e2037ad756
This commit introduces a major architectural refactoring to improve quota management, centralize logging, and streamline the relay handling logic. Key changes: - **Pre-consume Quota:** Implements a new mechanism to check and reserve user quota *before* making the request to the upstream provider. This ensures more accurate quota deduction and prevents users from exceeding their limits due to concurrent requests. - **Unified Relay Handlers:** Refactors the relay logic to use generic handlers (e.g., `ChatHandler`, `ImageHandler`) instead of provider-specific implementations. This significantly reduces code duplication and simplifies adding new channels. - **Centralized Logger:** A new dedicated `logger` package is introduced, and all system logging calls are migrated to use it, moving this responsibility out of the `common` package. - **Code Reorganization:** DTOs are generalized (e.g., `dalle.go` -> `openai_image.go`) and utility code is moved to more appropriate packages (e.g., `common/http.go` -> `service/http.go`) for better code structure.
60 lines
1.5 KiB
Go
60 lines
1.5 KiB
Go
package service
|
|
|
|
import (
|
|
"bytes"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"one-api/common"
|
|
"one-api/logger"
|
|
|
|
"github.com/gin-gonic/gin"
|
|
)
|
|
|
|
func CloseResponseBodyGracefully(httpResponse *http.Response) {
|
|
if httpResponse == nil || httpResponse.Body == nil {
|
|
return
|
|
}
|
|
err := httpResponse.Body.Close()
|
|
if err != nil {
|
|
common.SysError("failed to close response body: " + err.Error())
|
|
}
|
|
}
|
|
|
|
func IOCopyBytesGracefully(c *gin.Context, src *http.Response, data []byte) {
|
|
if c.Writer == nil {
|
|
return
|
|
}
|
|
|
|
body := io.NopCloser(bytes.NewBuffer(data))
|
|
|
|
// We shouldn't set the header before we parse the response body, because the parse part may fail.
|
|
// And then we will have to send an error response, but in this case, the header has already been set.
|
|
// So the httpClient will be confused by the response.
|
|
// For example, Postman will report error, and we cannot check the response at all.
|
|
if src != nil {
|
|
for k, v := range src.Header {
|
|
// avoid setting Content-Length
|
|
if k == "Content-Length" {
|
|
continue
|
|
}
|
|
c.Writer.Header().Set(k, v[0])
|
|
}
|
|
}
|
|
|
|
// set Content-Length header manually BEFORE calling WriteHeader
|
|
c.Writer.Header().Set("Content-Length", fmt.Sprintf("%d", len(data)))
|
|
|
|
// Write header with status code (this sends the headers)
|
|
if src != nil {
|
|
c.Writer.WriteHeader(src.StatusCode)
|
|
} else {
|
|
c.Writer.WriteHeader(http.StatusOK)
|
|
}
|
|
|
|
_, err := io.Copy(c.Writer, body)
|
|
if err != nil {
|
|
logger.LogError(c, fmt.Sprintf("failed to copy response body: %s", err.Error()))
|
|
}
|
|
}
|