hjjjj
5730b2a798
refactor: 重构 Gemini 适配器以支持图片编辑和生成 feat(relay): 添加图片编辑模式支持 feat(controller): 实现 UsageAPIURL 用于获取真实 token 用量 feat(web): 在渠道测试中添加模型选择功能 perf(token): 优化多模态 token 计算逻辑 fix(web): 修复日志分页组件显示问题 docs: 更新渠道配置中的 UsageAPIURL 说明 style: 清理调试日志和注释 feat(gemini): 支持 Imagen 3+ 图片生成模型 feat(openai): 添加生成 ID 捕获和元数据获取功能
95 lines
3.5 KiB
Go
95 lines
3.5 KiB
Go
package openai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/songquanpeng/one-api/common/logger"
|
|
"github.com/songquanpeng/one-api/relay/channeltype"
|
|
"github.com/songquanpeng/one-api/relay/model"
|
|
)
|
|
|
|
var metadataHTTPClient = &http.Client{Timeout: 8 * time.Second}
|
|
|
|
// usageMetadataResponse is a generic structure for upstream metadata endpoints
|
|
// that return nativeTokens (e.g. Zenmux /api/v1/generation).
|
|
type usageMetadataResponse struct {
|
|
NativeTokens struct {
|
|
PromptTokenCount int `json:"promptTokenCount"`
|
|
CandidatesTokenCount int `json:"candidatesTokenCount"`
|
|
TotalTokenCount int `json:"totalTokenCount"`
|
|
// ThoughtsTokenCount is billed at the completion rate, included in CompletionTokens.
|
|
ThoughtsTokenCount int `json:"thoughtsTokenCount"`
|
|
} `json:"nativeTokens"`
|
|
}
|
|
|
|
// FetchUsageFromMetadataURL fetches accurate token usage from an upstream metadata
|
|
// endpoint. urlTemplate must contain {id} which is replaced with generationId.
|
|
// Returns nil if the fetch fails or the response contains no usable token data.
|
|
func FetchUsageFromMetadataURL(urlTemplate, generationId string) (*model.Usage, error) {
|
|
url := strings.ReplaceAll(urlTemplate, "{id}", generationId)
|
|
resp, err := metadataHTTPClient.Get(url)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode != http.StatusOK {
|
|
return nil, fmt.Errorf("metadata API returned status %d", resp.StatusCode)
|
|
}
|
|
var meta usageMetadataResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&meta); err != nil {
|
|
return nil, err
|
|
}
|
|
if meta.NativeTokens.TotalTokenCount == 0 {
|
|
return nil, fmt.Errorf("metadata API returned zero total tokens")
|
|
}
|
|
prompt := meta.NativeTokens.PromptTokenCount
|
|
completion := meta.NativeTokens.CandidatesTokenCount + meta.NativeTokens.ThoughtsTokenCount
|
|
logger.SysLog(fmt.Sprintf("usage from metadata API (id=%s): prompt=%d completion=%d", generationId, prompt, completion))
|
|
return &model.Usage{
|
|
PromptTokens: prompt,
|
|
CompletionTokens: completion,
|
|
TotalTokens: prompt + completion,
|
|
}, nil
|
|
}
|
|
|
|
// countOutputMediaTokens returns a fixed token estimate for any embedded
|
|
// media data URIs found in text (image/video/audio), consistent with the
|
|
// fixed estimates used for input media in CountTokenMessages.
|
|
func countOutputMediaTokens(text string) int {
|
|
tokens := 0
|
|
tokens += strings.Count(text, "data:image/") * 2500
|
|
tokens += strings.Count(text, "data:video/") * 10000
|
|
tokens += strings.Count(text, "data:audio/") * 1500
|
|
return tokens
|
|
}
|
|
|
|
func ResponseText2Usage(responseText string, modelName string, promptTokens int) *model.Usage {
|
|
usage := &model.Usage{}
|
|
usage.PromptTokens = promptTokens
|
|
usage.CompletionTokens = CountTokenText(stripBase64Payloads(responseText), modelName) +
|
|
countOutputMediaTokens(responseText)
|
|
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
|
|
return usage
|
|
}
|
|
|
|
func GetFullRequestURL(baseURL string, requestURL string, channelType int) string {
|
|
if channelType == channeltype.OpenAICompatible {
|
|
return fmt.Sprintf("%s%s", strings.TrimSuffix(baseURL, "/"), strings.TrimPrefix(requestURL, "/v1"))
|
|
}
|
|
fullRequestURL := fmt.Sprintf("%s%s", baseURL, requestURL)
|
|
|
|
if strings.HasPrefix(baseURL, "https://gateway.ai.cloudflare.com") {
|
|
switch channelType {
|
|
case channeltype.OpenAI:
|
|
fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/v1"))
|
|
case channeltype.Azure:
|
|
fullRequestURL = fmt.Sprintf("%s%s", baseURL, strings.TrimPrefix(requestURL, "/openai/deployments"))
|
|
}
|
|
}
|
|
return fullRequestURL
|
|
}
|