mirror of
https://github.com/router-for-me/CLIProxyAPI.git
synced 2026-02-19 04:40:52 +08:00
refactor(executor): relocate gemini token counters
This commit is contained in:
hkfires
@@ -109,184 +109,6 @@ func (e *GeminiVertexExecutor) CountTokens(ctx context.Context, auth *cliproxyau
|
|||||||
return e.countTokensWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
|
return e.countTokensWithAPIKey(ctx, auth, req, opts, apiKey, baseURL)
|
||||||
}
|
}
|
||||||
|
|
||||||
// countTokensWithServiceAccount handles token counting using service account credentials.
|
|
||||||
func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (cliproxyexecutor.Response, error) {
|
|
||||||
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
|
||||||
|
|
||||||
from := opts.SourceFormat
|
|
||||||
to := sdktranslator.FromString("gemini")
|
|
||||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
|
||||||
if budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(req.Model, req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
|
||||||
if budgetOverride != nil {
|
|
||||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
|
||||||
budgetOverride = &norm
|
|
||||||
}
|
|
||||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
|
||||||
}
|
|
||||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
|
||||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
|
||||||
translatedReq, _ = sjson.SetBytes(translatedReq, "model", upstreamModel)
|
|
||||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
|
||||||
|
|
||||||
baseURL := vertexBaseURL(location)
|
|
||||||
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, upstreamModel, "countTokens")
|
|
||||||
|
|
||||||
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
|
||||||
if errNewReq != nil {
|
|
||||||
return cliproxyexecutor.Response{}, errNewReq
|
|
||||||
}
|
|
||||||
httpReq.Header.Set("Content-Type", "application/json")
|
|
||||||
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
|
|
||||||
httpReq.Header.Set("Authorization", "Bearer "+token)
|
|
||||||
} else if errTok != nil {
|
|
||||||
log.Errorf("vertex executor: access token error: %v", errTok)
|
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
|
|
||||||
}
|
|
||||||
applyGeminiHeaders(httpReq, auth)
|
|
||||||
|
|
||||||
var authID, authLabel, authType, authValue string
|
|
||||||
if auth != nil {
|
|
||||||
authID = auth.ID
|
|
||||||
authLabel = auth.Label
|
|
||||||
authType, authValue = auth.AccountInfo()
|
|
||||||
}
|
|
||||||
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
|
||||||
URL: url,
|
|
||||||
Method: http.MethodPost,
|
|
||||||
Headers: httpReq.Header.Clone(),
|
|
||||||
Body: translatedReq,
|
|
||||||
Provider: e.Identifier(),
|
|
||||||
AuthID: authID,
|
|
||||||
AuthLabel: authLabel,
|
|
||||||
AuthType: authType,
|
|
||||||
AuthValue: authValue,
|
|
||||||
})
|
|
||||||
|
|
||||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
|
||||||
httpResp, errDo := httpClient.Do(httpReq)
|
|
||||||
if errDo != nil {
|
|
||||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
|
||||||
return cliproxyexecutor.Response{}, errDo
|
|
||||||
}
|
|
||||||
defer func() {
|
|
||||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
|
||||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
||||||
b, _ := io.ReadAll(httpResp.Body)
|
|
||||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
|
||||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
|
||||||
}
|
|
||||||
data, errRead := io.ReadAll(httpResp.Body)
|
|
||||||
if errRead != nil {
|
|
||||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
|
||||||
return cliproxyexecutor.Response{}, errRead
|
|
||||||
}
|
|
||||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
||||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
|
||||||
}
|
|
||||||
count := gjson.GetBytes(data, "totalTokens").Int()
|
|
||||||
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
|
||||||
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// countTokensWithAPIKey handles token counting using API key credentials.
|
|
||||||
func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (cliproxyexecutor.Response, error) {
|
|
||||||
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
|
||||||
|
|
||||||
from := opts.SourceFormat
|
|
||||||
to := sdktranslator.FromString("gemini")
|
|
||||||
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
|
||||||
if budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(req.Model, req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
|
||||||
if budgetOverride != nil {
|
|
||||||
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
|
||||||
budgetOverride = &norm
|
|
||||||
}
|
|
||||||
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
|
||||||
}
|
|
||||||
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
|
||||||
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
|
||||||
translatedReq, _ = sjson.SetBytes(translatedReq, "model", upstreamModel)
|
|
||||||
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
|
||||||
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
|
||||||
|
|
||||||
// For API key auth, use simpler URL format without project/location
|
|
||||||
if baseURL == "" {
|
|
||||||
baseURL = "https://generativelanguage.googleapis.com"
|
|
||||||
}
|
|
||||||
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "countTokens")
|
|
||||||
|
|
||||||
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
|
||||||
if errNewReq != nil {
|
|
||||||
return cliproxyexecutor.Response{}, errNewReq
|
|
||||||
}
|
|
||||||
httpReq.Header.Set("Content-Type", "application/json")
|
|
||||||
if apiKey != "" {
|
|
||||||
httpReq.Header.Set("x-goog-api-key", apiKey)
|
|
||||||
}
|
|
||||||
applyGeminiHeaders(httpReq, auth)
|
|
||||||
|
|
||||||
var authID, authLabel, authType, authValue string
|
|
||||||
if auth != nil {
|
|
||||||
authID = auth.ID
|
|
||||||
authLabel = auth.Label
|
|
||||||
authType, authValue = auth.AccountInfo()
|
|
||||||
}
|
|
||||||
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
|
||||||
URL: url,
|
|
||||||
Method: http.MethodPost,
|
|
||||||
Headers: httpReq.Header.Clone(),
|
|
||||||
Body: translatedReq,
|
|
||||||
Provider: e.Identifier(),
|
|
||||||
AuthID: authID,
|
|
||||||
AuthLabel: authLabel,
|
|
||||||
AuthType: authType,
|
|
||||||
AuthValue: authValue,
|
|
||||||
})
|
|
||||||
|
|
||||||
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
|
||||||
httpResp, errDo := httpClient.Do(httpReq)
|
|
||||||
if errDo != nil {
|
|
||||||
recordAPIResponseError(ctx, e.cfg, errDo)
|
|
||||||
return cliproxyexecutor.Response{}, errDo
|
|
||||||
}
|
|
||||||
defer func() {
|
|
||||||
if errClose := httpResp.Body.Close(); errClose != nil {
|
|
||||||
log.Errorf("vertex executor: close response body error: %v", errClose)
|
|
||||||
}
|
|
||||||
}()
|
|
||||||
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
||||||
b, _ := io.ReadAll(httpResp.Body)
|
|
||||||
appendAPIResponseChunk(ctx, e.cfg, b)
|
|
||||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
|
||||||
}
|
|
||||||
data, errRead := io.ReadAll(httpResp.Body)
|
|
||||||
if errRead != nil {
|
|
||||||
recordAPIResponseError(ctx, e.cfg, errRead)
|
|
||||||
return cliproxyexecutor.Response{}, errRead
|
|
||||||
}
|
|
||||||
appendAPIResponseChunk(ctx, e.cfg, data)
|
|
||||||
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
|
||||||
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
|
||||||
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
|
||||||
}
|
|
||||||
count := gjson.GetBytes(data, "totalTokens").Int()
|
|
||||||
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
|
||||||
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
|
||||||
}
|
|
||||||
|
|
||||||
// Refresh refreshes the authentication credentials (no-op for Vertex).
|
// Refresh refreshes the authentication credentials (no-op for Vertex).
|
||||||
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
func (e *GeminiVertexExecutor) Refresh(_ context.Context, auth *cliproxyauth.Auth) (*cliproxyauth.Auth, error) {
|
||||||
return auth, nil
|
return auth, nil
|
||||||
@@ -728,6 +550,184 @@ func (e *GeminiVertexExecutor) executeStreamWithAPIKey(ctx context.Context, auth
|
|||||||
return stream, nil
|
return stream, nil
|
||||||
}
|
}
|
||||||
|
|
||||||
|
// countTokensWithServiceAccount counts tokens using service account credentials.
|
||||||
|
func (e *GeminiVertexExecutor) countTokensWithServiceAccount(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, projectID, location string, saJSON []byte) (cliproxyexecutor.Response, error) {
|
||||||
|
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
||||||
|
|
||||||
|
from := opts.SourceFormat
|
||||||
|
to := sdktranslator.FromString("gemini")
|
||||||
|
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||||
|
if budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(req.Model, req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||||
|
if budgetOverride != nil {
|
||||||
|
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||||
|
budgetOverride = &norm
|
||||||
|
}
|
||||||
|
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||||
|
}
|
||||||
|
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||||
|
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||||
|
translatedReq, _ = sjson.SetBytes(translatedReq, "model", upstreamModel)
|
||||||
|
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||||
|
|
||||||
|
baseURL := vertexBaseURL(location)
|
||||||
|
url := fmt.Sprintf("%s/%s/projects/%s/locations/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, projectID, location, upstreamModel, "countTokens")
|
||||||
|
|
||||||
|
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
||||||
|
if errNewReq != nil {
|
||||||
|
return cliproxyexecutor.Response{}, errNewReq
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
if token, errTok := vertexAccessToken(ctx, e.cfg, auth, saJSON); errTok == nil && token != "" {
|
||||||
|
httpReq.Header.Set("Authorization", "Bearer "+token)
|
||||||
|
} else if errTok != nil {
|
||||||
|
log.Errorf("vertex executor: access token error: %v", errTok)
|
||||||
|
return cliproxyexecutor.Response{}, statusErr{code: 500, msg: "internal server error"}
|
||||||
|
}
|
||||||
|
applyGeminiHeaders(httpReq, auth)
|
||||||
|
|
||||||
|
var authID, authLabel, authType, authValue string
|
||||||
|
if auth != nil {
|
||||||
|
authID = auth.ID
|
||||||
|
authLabel = auth.Label
|
||||||
|
authType, authValue = auth.AccountInfo()
|
||||||
|
}
|
||||||
|
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
||||||
|
URL: url,
|
||||||
|
Method: http.MethodPost,
|
||||||
|
Headers: httpReq.Header.Clone(),
|
||||||
|
Body: translatedReq,
|
||||||
|
Provider: e.Identifier(),
|
||||||
|
AuthID: authID,
|
||||||
|
AuthLabel: authLabel,
|
||||||
|
AuthType: authType,
|
||||||
|
AuthValue: authValue,
|
||||||
|
})
|
||||||
|
|
||||||
|
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||||
|
httpResp, errDo := httpClient.Do(httpReq)
|
||||||
|
if errDo != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||||
|
return cliproxyexecutor.Response{}, errDo
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||||
|
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
|
b, _ := io.ReadAll(httpResp.Body)
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||||
|
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||||
|
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||||
|
}
|
||||||
|
data, errRead := io.ReadAll(httpResp.Body)
|
||||||
|
if errRead != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||||
|
return cliproxyexecutor.Response{}, errRead
|
||||||
|
}
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||||
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
|
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||||
|
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
||||||
|
}
|
||||||
|
count := gjson.GetBytes(data, "totalTokens").Int()
|
||||||
|
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
||||||
|
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
|
// countTokensWithAPIKey handles token counting using API key credentials.
|
||||||
|
func (e *GeminiVertexExecutor) countTokensWithAPIKey(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options, apiKey, baseURL string) (cliproxyexecutor.Response, error) {
|
||||||
|
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
||||||
|
|
||||||
|
from := opts.SourceFormat
|
||||||
|
to := sdktranslator.FromString("gemini")
|
||||||
|
translatedReq := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||||
|
if budgetOverride, includeOverride, ok := util.ResolveThinkingConfigFromMetadata(req.Model, req.Metadata); ok && util.ModelSupportsThinking(req.Model) {
|
||||||
|
if budgetOverride != nil {
|
||||||
|
norm := util.NormalizeThinkingBudget(req.Model, *budgetOverride)
|
||||||
|
budgetOverride = &norm
|
||||||
|
}
|
||||||
|
translatedReq = util.ApplyGeminiThinkingConfig(translatedReq, budgetOverride, includeOverride)
|
||||||
|
}
|
||||||
|
translatedReq = util.StripThinkingConfigIfUnsupported(req.Model, translatedReq)
|
||||||
|
translatedReq = fixGeminiImageAspectRatio(req.Model, translatedReq)
|
||||||
|
translatedReq, _ = sjson.SetBytes(translatedReq, "model", upstreamModel)
|
||||||
|
respCtx := context.WithValue(ctx, "alt", opts.Alt)
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "tools")
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "generationConfig")
|
||||||
|
translatedReq, _ = sjson.DeleteBytes(translatedReq, "safetySettings")
|
||||||
|
|
||||||
|
// For API key auth, use simpler URL format without project/location
|
||||||
|
if baseURL == "" {
|
||||||
|
baseURL = "https://generativelanguage.googleapis.com"
|
||||||
|
}
|
||||||
|
url := fmt.Sprintf("%s/%s/publishers/google/models/%s:%s", baseURL, vertexAPIVersion, req.Model, "countTokens")
|
||||||
|
|
||||||
|
httpReq, errNewReq := http.NewRequestWithContext(respCtx, http.MethodPost, url, bytes.NewReader(translatedReq))
|
||||||
|
if errNewReq != nil {
|
||||||
|
return cliproxyexecutor.Response{}, errNewReq
|
||||||
|
}
|
||||||
|
httpReq.Header.Set("Content-Type", "application/json")
|
||||||
|
if apiKey != "" {
|
||||||
|
httpReq.Header.Set("x-goog-api-key", apiKey)
|
||||||
|
}
|
||||||
|
applyGeminiHeaders(httpReq, auth)
|
||||||
|
|
||||||
|
var authID, authLabel, authType, authValue string
|
||||||
|
if auth != nil {
|
||||||
|
authID = auth.ID
|
||||||
|
authLabel = auth.Label
|
||||||
|
authType, authValue = auth.AccountInfo()
|
||||||
|
}
|
||||||
|
recordAPIRequest(ctx, e.cfg, upstreamRequestLog{
|
||||||
|
URL: url,
|
||||||
|
Method: http.MethodPost,
|
||||||
|
Headers: httpReq.Header.Clone(),
|
||||||
|
Body: translatedReq,
|
||||||
|
Provider: e.Identifier(),
|
||||||
|
AuthID: authID,
|
||||||
|
AuthLabel: authLabel,
|
||||||
|
AuthType: authType,
|
||||||
|
AuthValue: authValue,
|
||||||
|
})
|
||||||
|
|
||||||
|
httpClient := newProxyAwareHTTPClient(ctx, e.cfg, auth, 0)
|
||||||
|
httpResp, errDo := httpClient.Do(httpReq)
|
||||||
|
if errDo != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errDo)
|
||||||
|
return cliproxyexecutor.Response{}, errDo
|
||||||
|
}
|
||||||
|
defer func() {
|
||||||
|
if errClose := httpResp.Body.Close(); errClose != nil {
|
||||||
|
log.Errorf("vertex executor: close response body error: %v", errClose)
|
||||||
|
}
|
||||||
|
}()
|
||||||
|
recordAPIResponseMetadata(ctx, e.cfg, httpResp.StatusCode, httpResp.Header.Clone())
|
||||||
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
|
b, _ := io.ReadAll(httpResp.Body)
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, b)
|
||||||
|
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), b))
|
||||||
|
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(b)}
|
||||||
|
}
|
||||||
|
data, errRead := io.ReadAll(httpResp.Body)
|
||||||
|
if errRead != nil {
|
||||||
|
recordAPIResponseError(ctx, e.cfg, errRead)
|
||||||
|
return cliproxyexecutor.Response{}, errRead
|
||||||
|
}
|
||||||
|
appendAPIResponseChunk(ctx, e.cfg, data)
|
||||||
|
if httpResp.StatusCode < 200 || httpResp.StatusCode >= 300 {
|
||||||
|
log.Debugf("request error, error status: %d, error body: %s", httpResp.StatusCode, summarizeErrorBody(httpResp.Header.Get("Content-Type"), data))
|
||||||
|
return cliproxyexecutor.Response{}, statusErr{code: httpResp.StatusCode, msg: string(data)}
|
||||||
|
}
|
||||||
|
count := gjson.GetBytes(data, "totalTokens").Int()
|
||||||
|
out := sdktranslator.TranslateTokenCount(ctx, to, from, count, data)
|
||||||
|
return cliproxyexecutor.Response{Payload: []byte(out)}, nil
|
||||||
|
}
|
||||||
|
|
||||||
// vertexCreds extracts project, location and raw service account JSON from auth metadata.
|
// vertexCreds extracts project, location and raw service account JSON from auth metadata.
|
||||||
func vertexCreds(a *cliproxyauth.Auth) (projectID, location string, serviceAccountJSON []byte, err error) {
|
func vertexCreds(a *cliproxyauth.Auth) (projectID, location string, serviceAccountJSON []byte, err error) {
|
||||||
if a == nil || a.Metadata == nil {
|
if a == nil || a.Metadata == nil {
|
||||||
|
|||||||
Reference in New Issue
Block a user