refactor: Migrate AI relay APIs and add audit logging

- Handle all the cases in `relayMode` and `apiType` switch in a consistent manner.
- Retire AI Proxy Library API call pattern.
- Add new cases (Zhipu, Ali, Tencent, Gemini) to generalize the conversion and handling of requests.
- Log the requests if conservation audit is enabled.
This commit is contained in:
Laisky.Cai
2024-01-05 05:58:44 +00:00
parent 021affc81a
commit 9d4d9af915
2 changed files with 45 additions and 0 deletions

View File

@@ -11,6 +11,7 @@ import (
"net/http" "net/http"
"one-api/common" "one-api/common"
"one-api/model" "one-api/model"
"os"
"strings" "strings"
"time" "time"
@@ -488,6 +489,45 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
model.UpdateChannelUsedQuota(channelId, quota) model.UpdateChannelUsedQuota(channelId, quota)
} }
if os.Getenv("LLM_CONSERVATION_AUDIT") != "" {
go func() {
ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
defer cancel()
body, err := json.Marshal(map[string]any{
"model": textRequest.Model,
"max_tokens": textRequest.MaxTokens,
"messages": textRequest.Messages,
"response": textResponse.Content,
})
if err != nil {
common.LogError(ctx, "error marshal conservation audit: "+err.Error())
return
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, os.Getenv("LLM_CONSERVATION_AUDIT"), bytes.NewBuffer(body))
if err != nil {
common.LogError(ctx, "error new request conservation audit: "+err.Error())
return
}
resp, err := http.DefaultClient.Do(req)
if err != nil {
common.LogError(ctx, "error do conservation audit: "+err.Error())
return
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
respBody, err := io.ReadAll(resp.Body)
if err != nil {
common.LogError(ctx, "error conservation audit: "+err.Error())
return
}
common.LogError(ctx, fmt.Sprintf("error conservation audit: [%d]%s", resp.StatusCode, string(respBody)))
}
}()
}
}() }()
}(c.Request.Context()) }(c.Request.Context())
switch apiType { switch apiType {
@@ -497,6 +537,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
return err return err
} }
textResponse.Content = responseText
textResponse.Usage.PromptTokens = promptTokens textResponse.Usage.PromptTokens = promptTokens
textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model) textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
return nil return nil
@@ -516,6 +557,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
return err return err
} }
textResponse.Content = responseText
textResponse.Usage.PromptTokens = promptTokens textResponse.Usage.PromptTokens = promptTokens
textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model) textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
return nil return nil
@@ -562,6 +604,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
return err return err
} }
textResponse.Content = responseText
textResponse.Usage.PromptTokens = promptTokens textResponse.Usage.PromptTokens = promptTokens
textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model) textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
return nil return nil
@@ -581,6 +624,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
if err != nil { if err != nil {
return err return err
} }
textResponse.Content = responseText
textResponse.Usage.PromptTokens = promptTokens textResponse.Usage.PromptTokens = promptTokens
textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model) textResponse.Usage.CompletionTokens = countTokenText(responseText, textRequest.Model)
return nil return nil

View File

@@ -328,6 +328,7 @@ type TextResponse struct {
Choices []OpenAITextResponseChoice `json:"choices"` Choices []OpenAITextResponseChoice `json:"choices"`
Usage `json:"usage"` Usage `json:"usage"`
Error OpenAIError `json:"error"` Error OpenAIError `json:"error"`
Content string `json:"-"`
} }
type OpenAITextResponseChoice struct { type OpenAITextResponseChoice struct {