mirror of
https://github.com/router-for-me/CLIProxyAPIPlus.git
synced 2026-03-09 23:33:24 +00:00
Compare commits
119 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
0f51e73baa | ||
|
|
3a436e116a | ||
|
|
d06e2dc83c | ||
|
|
336867853b | ||
|
|
6403ff4ec4 | ||
|
|
d222469b44 | ||
|
|
790a17ce98 | ||
|
|
d473c952fb | ||
|
|
7646a2b877 | ||
|
|
62090f2568 | ||
|
|
d35152bbef | ||
|
|
c281f4cbaf | ||
|
|
09455f9e85 | ||
|
|
c8e72ba0dc | ||
|
|
375ef252ab | ||
|
|
ee552f8720 | ||
|
|
2e88c4858e | ||
|
|
3f50da85c1 | ||
|
|
8be06255f7 | ||
|
|
60936b5185 | ||
|
|
72274099aa | ||
|
|
b7f7b3a1d8 | ||
|
|
dcae098e23 | ||
|
|
618606966f | ||
|
|
05f249d77f | ||
|
|
2eb05ec640 | ||
|
|
3ce0d76aa4 | ||
|
|
a00b79d9be | ||
|
|
9fe6a215e6 | ||
|
|
33e53a2a56 | ||
|
|
cd5b80785f | ||
|
|
54f71aa273 | ||
|
|
3f949b7f84 | ||
|
|
cf8b2dcc85 | ||
|
|
8e24d9dc34 | ||
|
|
443c4538bb | ||
|
|
a7fc2ee4cf | ||
|
|
8e749ac22d | ||
|
|
69e09d9bc7 | ||
|
|
ed57d82bc1 | ||
|
|
06ad527e8c | ||
|
|
7af5a90a0b | ||
|
|
7551faff79 | ||
|
|
b7409dd2de | ||
|
|
5ba325a8fc | ||
|
|
d502840f91 | ||
|
|
99238a4b59 | ||
|
|
6d43a2ff9a | ||
|
|
cdb9c2e6e8 | ||
|
|
3faa1ca9af | ||
|
|
9d975e0375 | ||
|
|
2a6d8b78d4 | ||
|
|
671558a822 | ||
|
|
6b80ec79a0 | ||
|
|
d3f4783a24 | ||
|
|
1cb6bdbc87 | ||
|
|
96ddfc1f24 | ||
|
|
c169b32570 | ||
|
|
36a512fdf2 | ||
|
|
26fbb77901 | ||
|
|
a277302262 | ||
|
|
969c1a5b72 | ||
|
|
872339bceb | ||
|
|
5dc0dbc7aa | ||
|
|
ee6fc4e8a1 | ||
|
|
8fee16aecd | ||
|
|
2b7ba54a2f | ||
|
|
007c3304f2 | ||
|
|
e76ba0ede9 | ||
|
|
c06ac07e23 | ||
|
|
e592a57458 | ||
|
|
66769ec657 | ||
|
|
f413feec61 | ||
|
|
2e538e3486 | ||
|
|
9617a7b0d6 | ||
|
|
7569320770 | ||
|
|
8d25cf0d75 | ||
|
|
64e85e7019 | ||
|
|
a862984dca | ||
|
|
f0365f0465 | ||
|
|
6d1e20e940 | ||
|
|
0c0aae1eac | ||
|
|
5dcf7cb846 | ||
|
|
349b2ba3af | ||
|
|
98db5aabd0 | ||
|
|
e52b542e22 | ||
|
|
8f6abb8a86 | ||
|
|
ed8eaae964 | ||
|
|
7fd98f3556 | ||
|
|
e8de87ee90 | ||
|
|
4e572ec8b9 | ||
|
|
6c7f18c448 | ||
|
|
24bc9cba67 | ||
|
|
97356b1a04 | ||
|
|
1084b53fba | ||
|
|
b1aecc2bf1 | ||
|
|
83b90e106f | ||
|
|
f52114dab2 | ||
|
|
5106caf641 | ||
|
|
12370ee84e | ||
|
|
b84ccc6e7a | ||
|
|
e19ddb53e7 | ||
|
|
5bf89dd757 | ||
|
|
2a0100b2d6 | ||
|
|
4442574e53 | ||
|
|
c020fa60d0 | ||
|
|
b078be4613 | ||
|
|
71a6dffbb6 | ||
|
|
5f65dd5bb4 | ||
|
|
27b43ed63f | ||
|
|
f6a3a1d0ba | ||
|
|
830fd8eac2 | ||
|
|
a86d501dc2 | ||
|
|
24e8e20b59 | ||
|
|
a87f09bad2 | ||
|
|
bc6c4cdbfc | ||
|
|
404546ce93 | ||
|
|
6dd1cf1dd6 | ||
|
|
9058d406a3 |
@@ -13,8 +13,6 @@ Dockerfile
|
||||
docs/*
|
||||
README.md
|
||||
README_CN.md
|
||||
MANAGEMENT_API.md
|
||||
MANAGEMENT_API_CN.md
|
||||
LICENSE
|
||||
|
||||
# Runtime data folders (should be mounted as volumes)
|
||||
@@ -32,3 +30,4 @@ bin/*
|
||||
.agent/*
|
||||
.bmad/*
|
||||
_bmad/*
|
||||
_bmad-output/*
|
||||
|
||||
7
.github/ISSUE_TEMPLATE/bug_report.md
vendored
7
.github/ISSUE_TEMPLATE/bug_report.md
vendored
@@ -7,6 +7,13 @@ assignees: ''
|
||||
|
||||
---
|
||||
|
||||
**Is it a request payload issue?**
|
||||
[ ] Yes, this is a request payload issue. I am using a client/cURL to send a request payload, but I received an unexpected error.
|
||||
[ ] No, it's another issue.
|
||||
|
||||
**If it's a request payload issue, you MUST know**
|
||||
Our team doesn't have any GODs or ORACLEs or MIND READERs. Please make sure to attach the request log or curl payload.
|
||||
|
||||
**Describe the bug**
|
||||
A clear and concise description of what the bug is.
|
||||
|
||||
|
||||
7
.gitignore
vendored
7
.gitignore
vendored
@@ -12,11 +12,15 @@ bin/*
|
||||
logs/*
|
||||
conv/*
|
||||
temp/*
|
||||
refs/*
|
||||
|
||||
# Storage backends
|
||||
pgstore/*
|
||||
gitstore/*
|
||||
objectstore/*
|
||||
|
||||
# Static assets
|
||||
static/*
|
||||
refs/*
|
||||
|
||||
# Authentication data
|
||||
auths/*
|
||||
@@ -36,6 +40,7 @@ GEMINI.md
|
||||
.agent/*
|
||||
.bmad/*
|
||||
_bmad/*
|
||||
_bmad-output/*
|
||||
.mcp/cache/
|
||||
|
||||
# macOS
|
||||
|
||||
BIN
assets/cubence.png
Normal file
BIN
assets/cubence.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 51 KiB |
BIN
assets/packycode.png
Normal file
BIN
assets/packycode.png
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 8.1 KiB |
@@ -39,6 +39,9 @@ api-keys:
|
||||
# Enable debug logging
|
||||
debug: false
|
||||
|
||||
# When true, disable high-overhead HTTP middleware features to reduce per-request memory usage under high concurrency.
|
||||
commercial-mode: false
|
||||
|
||||
# Open OAuth URLs in incognito/private browser mode.
|
||||
# Useful when you want to login with a different account without logging out from your current session.
|
||||
# Default: false (but Kiro auth defaults to true for multi-account support)
|
||||
@@ -71,9 +74,18 @@ quota-exceeded:
|
||||
switch-project: true # Whether to automatically switch to another project when a quota is exceeded
|
||||
switch-preview-model: true # Whether to automatically switch to a preview model when a quota is exceeded
|
||||
|
||||
# Routing strategy for selecting credentials when multiple match.
|
||||
routing:
|
||||
strategy: "round-robin" # round-robin (default), fill-first
|
||||
|
||||
# When true, enable authentication for the WebSocket API (/v1/ws).
|
||||
ws-auth: false
|
||||
|
||||
# Streaming behavior (SSE keep-alives + safe bootstrap retries).
|
||||
# streaming:
|
||||
# keepalive-seconds: 15 # Default: 0 (disabled). <= 0 disables keep-alives.
|
||||
# bootstrap-retries: 1 # Default: 0 (disabled). Retries before first byte is sent.
|
||||
|
||||
# Gemini API keys
|
||||
# gemini-api-key:
|
||||
# - api-key: "AIzaSy...01"
|
||||
@@ -97,6 +109,9 @@ ws-auth: false
|
||||
# headers:
|
||||
# X-Custom-Header: "custom-value"
|
||||
# proxy-url: "socks5://proxy.example.com:1080" # optional: per-key proxy override
|
||||
# models:
|
||||
# - name: "gpt-5-codex" # upstream model name
|
||||
# alias: "codex-latest" # client alias mapped to the upstream model
|
||||
# excluded-models:
|
||||
# - "gpt-5.1" # exclude specific models (exact match)
|
||||
# - "gpt-5-*" # wildcard matching prefix (e.g. gpt-5-medium, gpt-5-codex)
|
||||
|
||||
@@ -597,11 +597,7 @@ func (h *Handler) PutCodexKeys(c *gin.Context) {
|
||||
filtered := make([]config.CodexKey, 0, len(arr))
|
||||
for i := range arr {
|
||||
entry := arr[i]
|
||||
entry.APIKey = strings.TrimSpace(entry.APIKey)
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = config.NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(entry.ExcludedModels)
|
||||
normalizeCodexKey(&entry)
|
||||
if entry.BaseURL == "" {
|
||||
continue
|
||||
}
|
||||
@@ -613,12 +609,13 @@ func (h *Handler) PutCodexKeys(c *gin.Context) {
|
||||
}
|
||||
func (h *Handler) PatchCodexKey(c *gin.Context) {
|
||||
type codexKeyPatch struct {
|
||||
APIKey *string `json:"api-key"`
|
||||
Prefix *string `json:"prefix"`
|
||||
BaseURL *string `json:"base-url"`
|
||||
ProxyURL *string `json:"proxy-url"`
|
||||
Headers *map[string]string `json:"headers"`
|
||||
ExcludedModels *[]string `json:"excluded-models"`
|
||||
APIKey *string `json:"api-key"`
|
||||
Prefix *string `json:"prefix"`
|
||||
BaseURL *string `json:"base-url"`
|
||||
ProxyURL *string `json:"proxy-url"`
|
||||
Models *[]config.CodexModel `json:"models"`
|
||||
Headers *map[string]string `json:"headers"`
|
||||
ExcludedModels *[]string `json:"excluded-models"`
|
||||
}
|
||||
var body struct {
|
||||
Index *int `json:"index"`
|
||||
@@ -667,12 +664,16 @@ func (h *Handler) PatchCodexKey(c *gin.Context) {
|
||||
if body.Value.ProxyURL != nil {
|
||||
entry.ProxyURL = strings.TrimSpace(*body.Value.ProxyURL)
|
||||
}
|
||||
if body.Value.Models != nil {
|
||||
entry.Models = append([]config.CodexModel(nil), (*body.Value.Models)...)
|
||||
}
|
||||
if body.Value.Headers != nil {
|
||||
entry.Headers = config.NormalizeHeaders(*body.Value.Headers)
|
||||
}
|
||||
if body.Value.ExcludedModels != nil {
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(*body.Value.ExcludedModels)
|
||||
}
|
||||
normalizeCodexKey(&entry)
|
||||
h.cfg.CodexKey[targetIndex] = entry
|
||||
h.cfg.SanitizeCodexKeys()
|
||||
h.persist(c)
|
||||
@@ -762,6 +763,32 @@ func normalizeClaudeKey(entry *config.ClaudeKey) {
|
||||
entry.Models = normalized
|
||||
}
|
||||
|
||||
func normalizeCodexKey(entry *config.CodexKey) {
|
||||
if entry == nil {
|
||||
return
|
||||
}
|
||||
entry.APIKey = strings.TrimSpace(entry.APIKey)
|
||||
entry.Prefix = strings.TrimSpace(entry.Prefix)
|
||||
entry.BaseURL = strings.TrimSpace(entry.BaseURL)
|
||||
entry.ProxyURL = strings.TrimSpace(entry.ProxyURL)
|
||||
entry.Headers = config.NormalizeHeaders(entry.Headers)
|
||||
entry.ExcludedModels = config.NormalizeExcludedModels(entry.ExcludedModels)
|
||||
if len(entry.Models) == 0 {
|
||||
return
|
||||
}
|
||||
normalized := make([]config.CodexModel, 0, len(entry.Models))
|
||||
for i := range entry.Models {
|
||||
model := entry.Models[i]
|
||||
model.Name = strings.TrimSpace(model.Name)
|
||||
model.Alias = strings.TrimSpace(model.Alias)
|
||||
if model.Name == "" && model.Alias == "" {
|
||||
continue
|
||||
}
|
||||
normalized = append(normalized, model)
|
||||
}
|
||||
entry.Models = normalized
|
||||
}
|
||||
|
||||
// GetAmpCode returns the complete ampcode configuration.
|
||||
func (h *Handler) GetAmpCode(c *gin.Context) {
|
||||
if h == nil || h.cfg == nil {
|
||||
|
||||
@@ -209,6 +209,94 @@ func (h *Handler) GetRequestErrorLogs(c *gin.Context) {
|
||||
c.JSON(http.StatusOK, gin.H{"files": files})
|
||||
}
|
||||
|
||||
// GetRequestLogByID finds and downloads a request log file by its request ID.
|
||||
// The ID is matched against the suffix of log file names (format: *-{requestID}.log).
|
||||
func (h *Handler) GetRequestLogByID(c *gin.Context) {
|
||||
if h == nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "handler unavailable"})
|
||||
return
|
||||
}
|
||||
if h.cfg == nil {
|
||||
c.JSON(http.StatusServiceUnavailable, gin.H{"error": "configuration unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
dir := h.logDirectory()
|
||||
if strings.TrimSpace(dir) == "" {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "log directory not configured"})
|
||||
return
|
||||
}
|
||||
|
||||
requestID := strings.TrimSpace(c.Param("id"))
|
||||
if requestID == "" {
|
||||
requestID = strings.TrimSpace(c.Query("id"))
|
||||
}
|
||||
if requestID == "" {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "missing request ID"})
|
||||
return
|
||||
}
|
||||
if strings.ContainsAny(requestID, "/\\") {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid request ID"})
|
||||
return
|
||||
}
|
||||
|
||||
entries, err := os.ReadDir(dir)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "log directory not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to list log directory: %v", err)})
|
||||
return
|
||||
}
|
||||
|
||||
suffix := "-" + requestID + ".log"
|
||||
var matchedFile string
|
||||
for _, entry := range entries {
|
||||
if entry.IsDir() {
|
||||
continue
|
||||
}
|
||||
name := entry.Name()
|
||||
if strings.HasSuffix(name, suffix) {
|
||||
matchedFile = name
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if matchedFile == "" {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "log file not found for the given request ID"})
|
||||
return
|
||||
}
|
||||
|
||||
dirAbs, errAbs := filepath.Abs(dir)
|
||||
if errAbs != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to resolve log directory: %v", errAbs)})
|
||||
return
|
||||
}
|
||||
fullPath := filepath.Clean(filepath.Join(dirAbs, matchedFile))
|
||||
prefix := dirAbs + string(os.PathSeparator)
|
||||
if !strings.HasPrefix(fullPath, prefix) {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid log file path"})
|
||||
return
|
||||
}
|
||||
|
||||
info, errStat := os.Stat(fullPath)
|
||||
if errStat != nil {
|
||||
if os.IsNotExist(errStat) {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "log file not found"})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": fmt.Sprintf("failed to read log file: %v", errStat)})
|
||||
return
|
||||
}
|
||||
if info.IsDir() {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid log file"})
|
||||
return
|
||||
}
|
||||
|
||||
c.FileAttachment(fullPath, matchedFile)
|
||||
}
|
||||
|
||||
// DownloadRequestErrorLog downloads a specific error request log file by name.
|
||||
func (h *Handler) DownloadRequestErrorLog(c *gin.Context) {
|
||||
if h == nil {
|
||||
|
||||
@@ -1,12 +1,25 @@
|
||||
package management
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
|
||||
)
|
||||
|
||||
type usageExportPayload struct {
|
||||
Version int `json:"version"`
|
||||
ExportedAt time.Time `json:"exported_at"`
|
||||
Usage usage.StatisticsSnapshot `json:"usage"`
|
||||
}
|
||||
|
||||
type usageImportPayload struct {
|
||||
Version int `json:"version"`
|
||||
Usage usage.StatisticsSnapshot `json:"usage"`
|
||||
}
|
||||
|
||||
// GetUsageStatistics returns the in-memory request statistics snapshot.
|
||||
func (h *Handler) GetUsageStatistics(c *gin.Context) {
|
||||
var snapshot usage.StatisticsSnapshot
|
||||
@@ -18,3 +31,49 @@ func (h *Handler) GetUsageStatistics(c *gin.Context) {
|
||||
"failed_requests": snapshot.FailureCount,
|
||||
})
|
||||
}
|
||||
|
||||
// ExportUsageStatistics returns a complete usage snapshot for backup/migration.
|
||||
func (h *Handler) ExportUsageStatistics(c *gin.Context) {
|
||||
var snapshot usage.StatisticsSnapshot
|
||||
if h != nil && h.usageStats != nil {
|
||||
snapshot = h.usageStats.Snapshot()
|
||||
}
|
||||
c.JSON(http.StatusOK, usageExportPayload{
|
||||
Version: 1,
|
||||
ExportedAt: time.Now().UTC(),
|
||||
Usage: snapshot,
|
||||
})
|
||||
}
|
||||
|
||||
// ImportUsageStatistics merges a previously exported usage snapshot into memory.
|
||||
func (h *Handler) ImportUsageStatistics(c *gin.Context) {
|
||||
if h == nil || h.usageStats == nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "usage statistics unavailable"})
|
||||
return
|
||||
}
|
||||
|
||||
data, err := c.GetRawData()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "failed to read request body"})
|
||||
return
|
||||
}
|
||||
|
||||
var payload usageImportPayload
|
||||
if err := json.Unmarshal(data, &payload); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "invalid json"})
|
||||
return
|
||||
}
|
||||
if payload.Version != 0 && payload.Version != 1 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "unsupported version"})
|
||||
return
|
||||
}
|
||||
|
||||
result := h.usageStats.MergeSnapshot(payload.Usage)
|
||||
snapshot := h.usageStats.Snapshot()
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"added": result.Added,
|
||||
"skipped": result.Skipped,
|
||||
"total_requests": snapshot.TotalRequests,
|
||||
"failed_requests": snapshot.FailureCount,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -98,10 +98,11 @@ func captureRequestInfo(c *gin.Context) (*RequestInfo, error) {
|
||||
}
|
||||
|
||||
return &RequestInfo{
|
||||
URL: url,
|
||||
Method: method,
|
||||
Headers: headers,
|
||||
Body: body,
|
||||
URL: url,
|
||||
Method: method,
|
||||
Headers: headers,
|
||||
Body: body,
|
||||
RequestID: logging.GetGinRequestID(c),
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -15,10 +15,11 @@ import (
|
||||
|
||||
// RequestInfo holds essential details of an incoming HTTP request for logging purposes.
|
||||
type RequestInfo struct {
|
||||
URL string // URL is the request URL.
|
||||
Method string // Method is the HTTP method (e.g., GET, POST).
|
||||
Headers map[string][]string // Headers contains the request headers.
|
||||
Body []byte // Body is the raw request body.
|
||||
URL string // URL is the request URL.
|
||||
Method string // Method is the HTTP method (e.g., GET, POST).
|
||||
Headers map[string][]string // Headers contains the request headers.
|
||||
Body []byte // Body is the raw request body.
|
||||
RequestID string // RequestID is the unique identifier for the request.
|
||||
}
|
||||
|
||||
// ResponseWriterWrapper wraps the standard gin.ResponseWriter to intercept and log response data.
|
||||
@@ -149,6 +150,7 @@ func (w *ResponseWriterWrapper) WriteHeader(statusCode int) {
|
||||
w.requestInfo.Method,
|
||||
w.requestInfo.Headers,
|
||||
w.requestInfo.Body,
|
||||
w.requestInfo.RequestID,
|
||||
)
|
||||
if err == nil {
|
||||
w.streamWriter = streamWriter
|
||||
@@ -346,7 +348,7 @@ func (w *ResponseWriterWrapper) logRequest(statusCode int, headers map[string][]
|
||||
}
|
||||
|
||||
if loggerWithOptions, ok := w.logger.(interface {
|
||||
LogRequestWithOptions(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, []byte, []*interfaces.ErrorMessage, bool) error
|
||||
LogRequestWithOptions(string, string, map[string][]string, []byte, int, map[string][]string, []byte, []byte, []byte, []*interfaces.ErrorMessage, bool, string) error
|
||||
}); ok {
|
||||
return loggerWithOptions.LogRequestWithOptions(
|
||||
w.requestInfo.URL,
|
||||
@@ -360,6 +362,7 @@ func (w *ResponseWriterWrapper) logRequest(statusCode int, headers map[string][]
|
||||
apiResponseBody,
|
||||
apiResponseErrors,
|
||||
forceLog,
|
||||
w.requestInfo.RequestID,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -374,5 +377,6 @@ func (w *ResponseWriterWrapper) logRequest(statusCode int, headers map[string][]
|
||||
apiRequestBody,
|
||||
apiResponseBody,
|
||||
apiResponseErrors,
|
||||
w.requestInfo.RequestID,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -279,16 +279,23 @@ func (m *AmpModule) hasModelMappingsChanged(old *config.AmpCode, new *config.Amp
|
||||
return true
|
||||
}
|
||||
|
||||
// Build map for efficient comparison
|
||||
oldMap := make(map[string]string, len(old.ModelMappings))
|
||||
// Build map for efficient and robust comparison
|
||||
type mappingInfo struct {
|
||||
to string
|
||||
regex bool
|
||||
}
|
||||
oldMap := make(map[string]mappingInfo, len(old.ModelMappings))
|
||||
for _, mapping := range old.ModelMappings {
|
||||
oldMap[strings.TrimSpace(mapping.From)] = strings.TrimSpace(mapping.To)
|
||||
oldMap[strings.TrimSpace(mapping.From)] = mappingInfo{
|
||||
to: strings.TrimSpace(mapping.To),
|
||||
regex: mapping.Regex,
|
||||
}
|
||||
}
|
||||
|
||||
for _, mapping := range new.ModelMappings {
|
||||
from := strings.TrimSpace(mapping.From)
|
||||
to := strings.TrimSpace(mapping.To)
|
||||
if oldTo, exists := oldMap[from]; !exists || oldTo != to {
|
||||
if oldVal, exists := oldMap[from]; !exists || oldVal.to != to || oldVal.regex != mapping.Regex {
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
package amp
|
||||
|
||||
import (
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
|
||||
@@ -26,13 +27,15 @@ type ModelMapper interface {
|
||||
// DefaultModelMapper implements ModelMapper with thread-safe mapping storage.
|
||||
type DefaultModelMapper struct {
|
||||
mu sync.RWMutex
|
||||
mappings map[string]string // from -> to (normalized lowercase keys)
|
||||
mappings map[string]string // exact: from -> to (normalized lowercase keys)
|
||||
regexps []regexMapping // regex rules evaluated in order
|
||||
}
|
||||
|
||||
// NewModelMapper creates a new model mapper with the given initial mappings.
|
||||
func NewModelMapper(mappings []config.AmpModelMapping) *DefaultModelMapper {
|
||||
m := &DefaultModelMapper{
|
||||
mappings: make(map[string]string),
|
||||
regexps: nil,
|
||||
}
|
||||
m.UpdateMappings(mappings)
|
||||
return m
|
||||
@@ -55,7 +58,18 @@ func (m *DefaultModelMapper) MapModel(requestedModel string) string {
|
||||
// Check for direct mapping
|
||||
targetModel, exists := m.mappings[normalizedRequest]
|
||||
if !exists {
|
||||
return ""
|
||||
// Try regex mappings in order
|
||||
base, _ := util.NormalizeThinkingModel(requestedModel)
|
||||
for _, rm := range m.regexps {
|
||||
if rm.re.MatchString(requestedModel) || (base != "" && rm.re.MatchString(base)) {
|
||||
targetModel = rm.to
|
||||
exists = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !exists {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// Verify target model has available providers
|
||||
@@ -78,6 +92,7 @@ func (m *DefaultModelMapper) UpdateMappings(mappings []config.AmpModelMapping) {
|
||||
|
||||
// Clear and rebuild mappings
|
||||
m.mappings = make(map[string]string, len(mappings))
|
||||
m.regexps = make([]regexMapping, 0, len(mappings))
|
||||
|
||||
for _, mapping := range mappings {
|
||||
from := strings.TrimSpace(mapping.From)
|
||||
@@ -88,16 +103,30 @@ func (m *DefaultModelMapper) UpdateMappings(mappings []config.AmpModelMapping) {
|
||||
continue
|
||||
}
|
||||
|
||||
// Store with normalized lowercase key for case-insensitive lookup
|
||||
normalizedFrom := strings.ToLower(from)
|
||||
m.mappings[normalizedFrom] = to
|
||||
|
||||
log.Debugf("amp model mapping registered: %s -> %s", from, to)
|
||||
if mapping.Regex {
|
||||
// Compile case-insensitive regex; wrap with (?i) to match behavior of exact lookups
|
||||
pattern := "(?i)" + from
|
||||
re, err := regexp.Compile(pattern)
|
||||
if err != nil {
|
||||
log.Warnf("amp model mapping: invalid regex %q: %v", from, err)
|
||||
continue
|
||||
}
|
||||
m.regexps = append(m.regexps, regexMapping{re: re, to: to})
|
||||
log.Debugf("amp model regex mapping registered: /%s/ -> %s", from, to)
|
||||
} else {
|
||||
// Store with normalized lowercase key for case-insensitive lookup
|
||||
normalizedFrom := strings.ToLower(from)
|
||||
m.mappings[normalizedFrom] = to
|
||||
log.Debugf("amp model mapping registered: %s -> %s", from, to)
|
||||
}
|
||||
}
|
||||
|
||||
if len(m.mappings) > 0 {
|
||||
log.Infof("amp model mapping: loaded %d mapping(s)", len(m.mappings))
|
||||
}
|
||||
if n := len(m.regexps); n > 0 {
|
||||
log.Infof("amp model mapping: loaded %d regex mapping(s)", n)
|
||||
}
|
||||
}
|
||||
|
||||
// GetMappings returns a copy of current mappings (for debugging/status).
|
||||
@@ -111,3 +140,8 @@ func (m *DefaultModelMapper) GetMappings() map[string]string {
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
type regexMapping struct {
|
||||
re *regexp.Regexp
|
||||
to string
|
||||
}
|
||||
|
||||
@@ -203,3 +203,81 @@ func TestModelMapper_GetMappings_ReturnsCopy(t *testing.T) {
|
||||
t.Error("Original map was modified")
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_Regex_MatchBaseWithoutParens(t *testing.T) {
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client-regex-1", "gemini", []*registry.ModelInfo{
|
||||
{ID: "gemini-2.5-pro", OwnedBy: "google", Type: "gemini"},
|
||||
})
|
||||
defer reg.UnregisterClient("test-client-regex-1")
|
||||
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "^gpt-5$", To: "gemini-2.5-pro", Regex: true},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Incoming model has reasoning suffix but should match base via regex
|
||||
result := mapper.MapModel("gpt-5(high)")
|
||||
if result != "gemini-2.5-pro" {
|
||||
t.Errorf("Expected gemini-2.5-pro, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_Regex_ExactPrecedence(t *testing.T) {
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client-regex-2", "claude", []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
|
||||
})
|
||||
reg.RegisterClient("test-client-regex-3", "gemini", []*registry.ModelInfo{
|
||||
{ID: "gemini-2.5-pro", OwnedBy: "google", Type: "gemini"},
|
||||
})
|
||||
defer reg.UnregisterClient("test-client-regex-2")
|
||||
defer reg.UnregisterClient("test-client-regex-3")
|
||||
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "gpt-5", To: "claude-sonnet-4"}, // exact
|
||||
{From: "^gpt-5.*$", To: "gemini-2.5-pro", Regex: true}, // regex
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
// Exact match should win over regex
|
||||
result := mapper.MapModel("gpt-5")
|
||||
if result != "claude-sonnet-4" {
|
||||
t.Errorf("Expected claude-sonnet-4, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_Regex_InvalidPattern_Skipped(t *testing.T) {
|
||||
// Invalid regex should be skipped and not cause panic
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "(", To: "target", Regex: true},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
result := mapper.MapModel("anything")
|
||||
if result != "" {
|
||||
t.Errorf("Expected empty result due to invalid regex, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMapper_Regex_CaseInsensitive(t *testing.T) {
|
||||
reg := registry.GetGlobalRegistry()
|
||||
reg.RegisterClient("test-client-regex-4", "claude", []*registry.ModelInfo{
|
||||
{ID: "claude-sonnet-4", OwnedBy: "anthropic", Type: "claude"},
|
||||
})
|
||||
defer reg.UnregisterClient("test-client-regex-4")
|
||||
|
||||
mappings := []config.AmpModelMapping{
|
||||
{From: "^CLAUDE-OPUS-.*$", To: "claude-sonnet-4", Regex: true},
|
||||
}
|
||||
|
||||
mapper := NewModelMapper(mappings)
|
||||
|
||||
result := mapper.MapModel("claude-opus-4.5")
|
||||
if result != "claude-sonnet-4" {
|
||||
t.Errorf("Expected claude-sonnet-4, got %s", result)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -209,13 +209,15 @@ func NewServer(cfg *config.Config, authManager *auth.Manager, accessManager *sdk
|
||||
// Resolve logs directory relative to the configuration file directory.
|
||||
var requestLogger logging.RequestLogger
|
||||
var toggle func(bool)
|
||||
if optionState.requestLoggerFactory != nil {
|
||||
requestLogger = optionState.requestLoggerFactory(cfg, configFilePath)
|
||||
}
|
||||
if requestLogger != nil {
|
||||
engine.Use(middleware.RequestLoggingMiddleware(requestLogger))
|
||||
if setter, ok := requestLogger.(interface{ SetEnabled(bool) }); ok {
|
||||
toggle = setter.SetEnabled
|
||||
if !cfg.CommercialMode {
|
||||
if optionState.requestLoggerFactory != nil {
|
||||
requestLogger = optionState.requestLoggerFactory(cfg, configFilePath)
|
||||
}
|
||||
if requestLogger != nil {
|
||||
engine.Use(middleware.RequestLoggingMiddleware(requestLogger))
|
||||
if setter, ok := requestLogger.(interface{ SetEnabled(bool) }); ok {
|
||||
toggle = setter.SetEnabled
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -494,6 +496,8 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.Use(s.managementAvailabilityMiddleware(), s.mgmt.Middleware())
|
||||
{
|
||||
mgmt.GET("/usage", s.mgmt.GetUsageStatistics)
|
||||
mgmt.GET("/usage/export", s.mgmt.ExportUsageStatistics)
|
||||
mgmt.POST("/usage/import", s.mgmt.ImportUsageStatistics)
|
||||
mgmt.GET("/config", s.mgmt.GetConfig)
|
||||
mgmt.GET("/config.yaml", s.mgmt.GetConfigYAML)
|
||||
mgmt.PUT("/config.yaml", s.mgmt.PutConfigYAML)
|
||||
@@ -538,6 +542,7 @@ func (s *Server) registerManagementRoutes() {
|
||||
mgmt.DELETE("/logs", s.mgmt.DeleteLogs)
|
||||
mgmt.GET("/request-error-logs", s.mgmt.GetRequestErrorLogs)
|
||||
mgmt.GET("/request-error-logs/:name", s.mgmt.DownloadRequestErrorLog)
|
||||
mgmt.GET("/request-log-by-id/:id", s.mgmt.GetRequestLogByID)
|
||||
mgmt.GET("/request-log", s.mgmt.GetRequestLog)
|
||||
mgmt.PUT("/request-log", s.mgmt.PutRequestLog)
|
||||
mgmt.PATCH("/request-log", s.mgmt.PutRequestLog)
|
||||
|
||||
@@ -40,6 +40,10 @@ type KiroTokenData struct {
|
||||
ClientSecret string `json:"clientSecret,omitempty"`
|
||||
// Email is the user's email address (used for file naming)
|
||||
Email string `json:"email,omitempty"`
|
||||
// StartURL is the IDC/Identity Center start URL (only for IDC auth method)
|
||||
StartURL string `json:"startUrl,omitempty"`
|
||||
// Region is the AWS region for IDC authentication (only for IDC auth method)
|
||||
Region string `json:"region,omitempty"`
|
||||
}
|
||||
|
||||
// KiroAuthBundle aggregates authentication data after OAuth flow completion
|
||||
|
||||
@@ -2,16 +2,19 @@
|
||||
package kiro
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"html"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -24,19 +27,31 @@ import (
|
||||
const (
|
||||
// AWS SSO OIDC endpoints
|
||||
ssoOIDCEndpoint = "https://oidc.us-east-1.amazonaws.com"
|
||||
|
||||
|
||||
// Kiro's start URL for Builder ID
|
||||
builderIDStartURL = "https://view.awsapps.com/start"
|
||||
|
||||
|
||||
// Default region for IDC
|
||||
defaultIDCRegion = "us-east-1"
|
||||
|
||||
// Polling interval
|
||||
pollInterval = 5 * time.Second
|
||||
|
||||
|
||||
// Authorization code flow callback
|
||||
authCodeCallbackPath = "/oauth/callback"
|
||||
authCodeCallbackPort = 19877
|
||||
|
||||
|
||||
// User-Agent to match official Kiro IDE
|
||||
kiroUserAgent = "KiroIDE"
|
||||
|
||||
// IDC token refresh headers (matching Kiro IDE behavior)
|
||||
idcAmzUserAgent = "aws-sdk-js/3.738.0 ua/2.1 os/other lang/js md/browser#unknown_unknown api/sso-oidc#3.738.0 m/E KiroIDE"
|
||||
)
|
||||
|
||||
// Sentinel errors for OIDC token polling
|
||||
var (
|
||||
ErrAuthorizationPending = errors.New("authorization_pending")
|
||||
ErrSlowDown = errors.New("slow_down")
|
||||
)
|
||||
|
||||
// SSOOIDCClient handles AWS SSO OIDC authentication.
|
||||
@@ -83,6 +98,440 @@ type CreateTokenResponse struct {
|
||||
RefreshToken string `json:"refreshToken"`
|
||||
}
|
||||
|
||||
// getOIDCEndpoint returns the OIDC endpoint for the given region.
|
||||
func getOIDCEndpoint(region string) string {
|
||||
if region == "" {
|
||||
region = defaultIDCRegion
|
||||
}
|
||||
return fmt.Sprintf("https://oidc.%s.amazonaws.com", region)
|
||||
}
|
||||
|
||||
// promptInput prompts the user for input with an optional default value.
|
||||
func promptInput(prompt, defaultValue string) string {
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
if defaultValue != "" {
|
||||
fmt.Printf("%s [%s]: ", prompt, defaultValue)
|
||||
} else {
|
||||
fmt.Printf("%s: ", prompt)
|
||||
}
|
||||
input, err := reader.ReadString('\n')
|
||||
if err != nil {
|
||||
log.Warnf("Error reading input: %v", err)
|
||||
return defaultValue
|
||||
}
|
||||
input = strings.TrimSpace(input)
|
||||
if input == "" {
|
||||
return defaultValue
|
||||
}
|
||||
return input
|
||||
}
|
||||
|
||||
// promptSelect prompts the user to select from options using number input.
|
||||
func promptSelect(prompt string, options []string) int {
|
||||
reader := bufio.NewReader(os.Stdin)
|
||||
|
||||
for {
|
||||
fmt.Println(prompt)
|
||||
for i, opt := range options {
|
||||
fmt.Printf(" %d) %s\n", i+1, opt)
|
||||
}
|
||||
fmt.Printf("Enter selection (1-%d): ", len(options))
|
||||
|
||||
input, err := reader.ReadString('\n')
|
||||
if err != nil {
|
||||
log.Warnf("Error reading input: %v", err)
|
||||
return 0 // Default to first option on error
|
||||
}
|
||||
input = strings.TrimSpace(input)
|
||||
|
||||
// Parse the selection
|
||||
var selection int
|
||||
if _, err := fmt.Sscanf(input, "%d", &selection); err != nil || selection < 1 || selection > len(options) {
|
||||
fmt.Printf("Invalid selection '%s'. Please enter a number between 1 and %d.\n\n", input, len(options))
|
||||
continue
|
||||
}
|
||||
return selection - 1
|
||||
}
|
||||
}
|
||||
|
||||
// RegisterClientWithRegion registers a new OIDC client with AWS using a specific region.
|
||||
func (c *SSOOIDCClient) RegisterClientWithRegion(ctx context.Context, region string) (*RegisterClientResponse, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]interface{}{
|
||||
"clientName": "Kiro IDE",
|
||||
"clientType": "public",
|
||||
"scopes": []string{"codewhisperer:completions", "codewhisperer:analysis", "codewhisperer:conversations", "codewhisperer:transformations", "codewhisperer:taskassist"},
|
||||
"grantTypes": []string{"urn:ietf:params:oauth:grant-type:device_code", "refresh_token"},
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/client/register", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("register client failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("register client failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result RegisterClientResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// StartDeviceAuthorizationWithIDC starts the device authorization flow for IDC.
|
||||
func (c *SSOOIDCClient) StartDeviceAuthorizationWithIDC(ctx context.Context, clientID, clientSecret, startURL, region string) (*StartDeviceAuthResponse, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"startUrl": startURL,
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/device_authorization", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("start device auth failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("start device auth failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result StartDeviceAuthResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// CreateTokenWithRegion polls for the access token after user authorization using a specific region.
|
||||
func (c *SSOOIDCClient) CreateTokenWithRegion(ctx context.Context, clientID, clientSecret, deviceCode, region string) (*CreateTokenResponse, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"deviceCode": deviceCode,
|
||||
"grantType": "urn:ietf:params:oauth:grant-type:device_code",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("User-Agent", kiroUserAgent)
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Check for pending authorization
|
||||
if resp.StatusCode == http.StatusBadRequest {
|
||||
var errResp struct {
|
||||
Error string `json:"error"`
|
||||
}
|
||||
if json.Unmarshal(respBody, &errResp) == nil {
|
||||
if errResp.Error == "authorization_pending" {
|
||||
return nil, ErrAuthorizationPending
|
||||
}
|
||||
if errResp.Error == "slow_down" {
|
||||
return nil, ErrSlowDown
|
||||
}
|
||||
}
|
||||
log.Debugf("create token failed: %s", string(respBody))
|
||||
return nil, fmt.Errorf("create token failed")
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Debugf("create token failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("create token failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result CreateTokenResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return &result, nil
|
||||
}
|
||||
|
||||
// RefreshTokenWithRegion refreshes an access token using the refresh token with a specific region.
|
||||
func (c *SSOOIDCClient) RefreshTokenWithRegion(ctx context.Context, clientID, clientSecret, refreshToken, region, startURL string) (*KiroTokenData, error) {
|
||||
endpoint := getOIDCEndpoint(region)
|
||||
|
||||
payload := map[string]string{
|
||||
"clientId": clientID,
|
||||
"clientSecret": clientSecret,
|
||||
"refreshToken": refreshToken,
|
||||
"grantType": "refresh_token",
|
||||
}
|
||||
|
||||
body, err := json.Marshal(payload)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint+"/token", strings.NewReader(string(body)))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// Set headers matching kiro2api's IDC token refresh
|
||||
// These headers are required for successful IDC token refresh
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Host", fmt.Sprintf("oidc.%s.amazonaws.com", region))
|
||||
req.Header.Set("Connection", "keep-alive")
|
||||
req.Header.Set("x-amz-user-agent", idcAmzUserAgent)
|
||||
req.Header.Set("Accept", "*/*")
|
||||
req.Header.Set("Accept-Language", "*")
|
||||
req.Header.Set("sec-fetch-mode", "cors")
|
||||
req.Header.Set("User-Agent", "node")
|
||||
req.Header.Set("Accept-Encoding", "br, gzip, deflate")
|
||||
|
||||
resp, err := c.httpClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
respBody, err := io.ReadAll(resp.Body)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
log.Warnf("IDC token refresh failed (status %d): %s", resp.StatusCode, string(respBody))
|
||||
return nil, fmt.Errorf("token refresh failed (status %d)", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result CreateTokenResponse
|
||||
if err := json.Unmarshal(respBody, &result); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
expiresAt := time.Now().Add(time.Duration(result.ExpiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: result.AccessToken,
|
||||
RefreshToken: result.RefreshToken,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "idc",
|
||||
Provider: "AWS",
|
||||
ClientID: clientID,
|
||||
ClientSecret: clientSecret,
|
||||
StartURL: startURL,
|
||||
Region: region,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// LoginWithIDC performs the full device code flow for AWS Identity Center (IDC).
|
||||
func (c *SSOOIDCClient) LoginWithIDC(ctx context.Context, startURL, region string) (*KiroTokenData, error) {
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Authentication (AWS Identity Center) ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
// Step 1: Register client with the specified region
|
||||
fmt.Println("\nRegistering client...")
|
||||
regResp, err := c.RegisterClientWithRegion(ctx, region)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to register client: %w", err)
|
||||
}
|
||||
log.Debugf("Client registered: %s", regResp.ClientID)
|
||||
|
||||
// Step 2: Start device authorization with IDC start URL
|
||||
fmt.Println("Starting device authorization...")
|
||||
authResp, err := c.StartDeviceAuthorizationWithIDC(ctx, regResp.ClientID, regResp.ClientSecret, startURL, region)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to start device auth: %w", err)
|
||||
}
|
||||
|
||||
// Step 3: Show user the verification URL
|
||||
fmt.Printf("\n")
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf(" Confirm the following code in the browser:\n")
|
||||
fmt.Printf(" Code: %s\n", authResp.UserCode)
|
||||
fmt.Println("════════════════════════════════════════════════════════════")
|
||||
fmt.Printf("\n Open this URL: %s\n\n", authResp.VerificationURIComplete)
|
||||
|
||||
// Set incognito mode based on config
|
||||
if c.cfg != nil {
|
||||
browser.SetIncognitoMode(c.cfg.IncognitoBrowser)
|
||||
if !c.cfg.IncognitoBrowser {
|
||||
log.Info("kiro: using normal browser mode (--no-incognito). Note: You may not be able to select a different account.")
|
||||
} else {
|
||||
log.Debug("kiro: using incognito mode for multi-account support")
|
||||
}
|
||||
} else {
|
||||
browser.SetIncognitoMode(true)
|
||||
log.Debug("kiro: using incognito mode for multi-account support (default)")
|
||||
}
|
||||
|
||||
// Open browser
|
||||
if err := browser.OpenURL(authResp.VerificationURIComplete); err != nil {
|
||||
log.Warnf("Could not open browser automatically: %v", err)
|
||||
fmt.Println(" Please open the URL manually in your browser.")
|
||||
} else {
|
||||
fmt.Println(" (Browser opened automatically)")
|
||||
}
|
||||
|
||||
// Step 4: Poll for token
|
||||
fmt.Println("Waiting for authorization...")
|
||||
|
||||
interval := pollInterval
|
||||
if authResp.Interval > 0 {
|
||||
interval = time.Duration(authResp.Interval) * time.Second
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(time.Duration(authResp.ExpiresIn) * time.Second)
|
||||
|
||||
for time.Now().Before(deadline) {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
browser.CloseBrowser()
|
||||
return nil, ctx.Err()
|
||||
case <-time.After(interval):
|
||||
tokenResp, err := c.CreateTokenWithRegion(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode, region)
|
||||
if err != nil {
|
||||
if errors.Is(err, ErrAuthorizationPending) {
|
||||
fmt.Print(".")
|
||||
continue
|
||||
}
|
||||
if errors.Is(err, ErrSlowDown) {
|
||||
interval += 5 * time.Second
|
||||
continue
|
||||
}
|
||||
browser.CloseBrowser()
|
||||
return nil, fmt.Errorf("token creation failed: %w", err)
|
||||
}
|
||||
|
||||
fmt.Println("\n\n✓ Authorization successful!")
|
||||
|
||||
// Close the browser window
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser: %v", err)
|
||||
}
|
||||
|
||||
// Step 5: Get profile ARN from CodeWhisperer API
|
||||
fmt.Println("Fetching profile information...")
|
||||
profileArn := c.fetchProfileArn(ctx, tokenResp.AccessToken)
|
||||
|
||||
// Fetch user email
|
||||
email := FetchUserEmailWithFallback(ctx, c.cfg, tokenResp.AccessToken)
|
||||
if email != "" {
|
||||
fmt.Printf(" Logged in as: %s\n", email)
|
||||
}
|
||||
|
||||
expiresAt := time.Now().Add(time.Duration(tokenResp.ExpiresIn) * time.Second)
|
||||
|
||||
return &KiroTokenData{
|
||||
AccessToken: tokenResp.AccessToken,
|
||||
RefreshToken: tokenResp.RefreshToken,
|
||||
ProfileArn: profileArn,
|
||||
ExpiresAt: expiresAt.Format(time.RFC3339),
|
||||
AuthMethod: "idc",
|
||||
Provider: "AWS",
|
||||
ClientID: regResp.ClientID,
|
||||
ClientSecret: regResp.ClientSecret,
|
||||
Email: email,
|
||||
StartURL: startURL,
|
||||
Region: region,
|
||||
}, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Close browser on timeout
|
||||
if err := browser.CloseBrowser(); err != nil {
|
||||
log.Debugf("Failed to close browser on timeout: %v", err)
|
||||
}
|
||||
return nil, fmt.Errorf("authorization timed out")
|
||||
}
|
||||
|
||||
// LoginWithMethodSelection prompts the user to select between Builder ID and IDC, then performs the login.
|
||||
func (c *SSOOIDCClient) LoginWithMethodSelection(ctx context.Context) (*KiroTokenData, error) {
|
||||
fmt.Println("\n╔══════════════════════════════════════════════════════════╗")
|
||||
fmt.Println("║ Kiro Authentication (AWS) ║")
|
||||
fmt.Println("╚══════════════════════════════════════════════════════════╝")
|
||||
|
||||
// Prompt for login method
|
||||
options := []string{
|
||||
"Use with Builder ID (personal AWS account)",
|
||||
"Use with IDC Account (organization SSO)",
|
||||
}
|
||||
selection := promptSelect("\n? Select login method:", options)
|
||||
|
||||
if selection == 0 {
|
||||
// Builder ID flow - use existing implementation
|
||||
return c.LoginWithBuilderID(ctx)
|
||||
}
|
||||
|
||||
// IDC flow - prompt for start URL and region
|
||||
fmt.Println()
|
||||
startURL := promptInput("? Enter Start URL", "")
|
||||
if startURL == "" {
|
||||
return nil, fmt.Errorf("start URL is required for IDC login")
|
||||
}
|
||||
|
||||
region := promptInput("? Enter Region", defaultIDCRegion)
|
||||
|
||||
return c.LoginWithIDC(ctx, startURL, region)
|
||||
}
|
||||
|
||||
// RegisterClient registers a new OIDC client with AWS.
|
||||
func (c *SSOOIDCClient) RegisterClient(ctx context.Context) (*RegisterClientResponse, error) {
|
||||
payload := map[string]interface{}{
|
||||
@@ -211,10 +660,10 @@ func (c *SSOOIDCClient) CreateToken(ctx context.Context, clientID, clientSecret,
|
||||
}
|
||||
if json.Unmarshal(respBody, &errResp) == nil {
|
||||
if errResp.Error == "authorization_pending" {
|
||||
return nil, fmt.Errorf("authorization_pending")
|
||||
return nil, ErrAuthorizationPending
|
||||
}
|
||||
if errResp.Error == "slow_down" {
|
||||
return nil, fmt.Errorf("slow_down")
|
||||
return nil, ErrSlowDown
|
||||
}
|
||||
}
|
||||
log.Debugf("create token failed: %s", string(respBody))
|
||||
@@ -359,12 +808,11 @@ func (c *SSOOIDCClient) LoginWithBuilderID(ctx context.Context) (*KiroTokenData,
|
||||
case <-time.After(interval):
|
||||
tokenResp, err := c.CreateToken(ctx, regResp.ClientID, regResp.ClientSecret, authResp.DeviceCode)
|
||||
if err != nil {
|
||||
errStr := err.Error()
|
||||
if strings.Contains(errStr, "authorization_pending") {
|
||||
if errors.Is(err, ErrAuthorizationPending) {
|
||||
fmt.Print(".")
|
||||
continue
|
||||
}
|
||||
if strings.Contains(errStr, "slow_down") {
|
||||
if errors.Is(err, ErrSlowDown) {
|
||||
interval += 5 * time.Second
|
||||
continue
|
||||
}
|
||||
|
||||
@@ -39,6 +39,9 @@ type Config struct {
|
||||
// Debug enables or disables debug-level logging and other debug features.
|
||||
Debug bool `yaml:"debug" json:"debug"`
|
||||
|
||||
// CommercialMode disables high-overhead HTTP middleware features to minimize per-request memory usage.
|
||||
CommercialMode bool `yaml:"commercial-mode" json:"commercial-mode"`
|
||||
|
||||
// LoggingToFile controls whether application logs are written to rotating files or stdout.
|
||||
LoggingToFile bool `yaml:"logging-to-file" json:"logging-to-file"`
|
||||
|
||||
@@ -60,6 +63,9 @@ type Config struct {
|
||||
// QuotaExceeded defines the behavior when a quota is exceeded.
|
||||
QuotaExceeded QuotaExceeded `yaml:"quota-exceeded" json:"quota-exceeded"`
|
||||
|
||||
// Routing controls credential selection behavior.
|
||||
Routing RoutingConfig `yaml:"routing" json:"routing"`
|
||||
|
||||
// WebsocketAuth enables or disables authentication for the WebSocket API.
|
||||
WebsocketAuth bool `yaml:"ws-auth" json:"ws-auth"`
|
||||
|
||||
@@ -136,6 +142,13 @@ type QuotaExceeded struct {
|
||||
SwitchPreviewModel bool `yaml:"switch-preview-model" json:"switch-preview-model"`
|
||||
}
|
||||
|
||||
// RoutingConfig configures how credentials are selected for requests.
|
||||
type RoutingConfig struct {
|
||||
// Strategy selects the credential selection strategy.
|
||||
// Supported values: "round-robin" (default), "fill-first".
|
||||
Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"`
|
||||
}
|
||||
|
||||
// AmpModelMapping defines a model name mapping for Amp CLI requests.
|
||||
// When Amp requests a model that isn't available locally, this mapping
|
||||
// allows routing to an alternative model that IS available.
|
||||
@@ -146,6 +159,11 @@ type AmpModelMapping struct {
|
||||
// To is the target model name to route to (e.g., "claude-sonnet-4").
|
||||
// The target model must have available providers in the registry.
|
||||
To string `yaml:"to" json:"to"`
|
||||
|
||||
// Regex indicates whether the 'from' field should be interpreted as a regular
|
||||
// expression for matching model names. When true, this mapping is evaluated
|
||||
// after exact matches and in the order provided. Defaults to false (exact match).
|
||||
Regex bool `yaml:"regex,omitempty" json:"regex,omitempty"`
|
||||
}
|
||||
|
||||
// AmpCode groups Amp CLI integration settings including upstream routing,
|
||||
@@ -247,6 +265,9 @@ type CodexKey struct {
|
||||
// ProxyURL overrides the global proxy setting for this API key if provided.
|
||||
ProxyURL string `yaml:"proxy-url" json:"proxy-url"`
|
||||
|
||||
// Models defines upstream model names and aliases for request routing.
|
||||
Models []CodexModel `yaml:"models" json:"models"`
|
||||
|
||||
// Headers optionally adds extra HTTP headers for requests sent with this key.
|
||||
Headers map[string]string `yaml:"headers,omitempty" json:"headers,omitempty"`
|
||||
|
||||
@@ -254,6 +275,15 @@ type CodexKey struct {
|
||||
ExcludedModels []string `yaml:"excluded-models,omitempty" json:"excluded-models,omitempty"`
|
||||
}
|
||||
|
||||
// CodexModel describes a mapping between an alias and the actual upstream model name.
|
||||
type CodexModel struct {
|
||||
// Name is the upstream model identifier used when issuing requests.
|
||||
Name string `yaml:"name" json:"name"`
|
||||
|
||||
// Alias is the client-facing model name that maps to Name.
|
||||
Alias string `yaml:"alias" json:"alias"`
|
||||
}
|
||||
|
||||
// GeminiKey represents the configuration for a Gemini API key,
|
||||
// including optional overrides for upstream base URL, proxy routing, and headers.
|
||||
type GeminiKey struct {
|
||||
@@ -391,7 +421,7 @@ func LoadConfigOptional(configFile string, optional bool) (*Config, error) {
|
||||
cfg.DisableCooling = false
|
||||
cfg.AmpCode.RestrictManagementToLocalhost = false // Default to false: API key auth is sufficient
|
||||
cfg.RemoteManagement.PanelGitHubRepository = DefaultPanelGitHubRepository
|
||||
cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
|
||||
cfg.IncognitoBrowser = false // Default to normal browser (AWS uses incognito by force)
|
||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||
if optional {
|
||||
// In cloud deploy mode, if YAML parsing fails, return empty config instead of error.
|
||||
@@ -861,8 +891,8 @@ func getOrCreateMapValue(mapNode *yaml.Node, key string) *yaml.Node {
|
||||
}
|
||||
|
||||
// mergeMappingPreserve merges keys from src into dst mapping node while preserving
|
||||
// key order and comments of existing keys in dst. Unknown keys from src are appended
|
||||
// to dst at the end, copying their node structure from src.
|
||||
// key order and comments of existing keys in dst. New keys are only added if their
|
||||
// value is non-zero to avoid polluting the config with defaults.
|
||||
func mergeMappingPreserve(dst, src *yaml.Node) {
|
||||
if dst == nil || src == nil {
|
||||
return
|
||||
@@ -873,20 +903,19 @@ func mergeMappingPreserve(dst, src *yaml.Node) {
|
||||
copyNodeShallow(dst, src)
|
||||
return
|
||||
}
|
||||
// Build a lookup of existing keys in dst
|
||||
for i := 0; i+1 < len(src.Content); i += 2 {
|
||||
sk := src.Content[i]
|
||||
sv := src.Content[i+1]
|
||||
idx := findMapKeyIndex(dst, sk.Value)
|
||||
if idx >= 0 {
|
||||
// Merge into existing value node
|
||||
// Merge into existing value node (always update, even to zero values)
|
||||
dv := dst.Content[idx+1]
|
||||
mergeNodePreserve(dv, sv)
|
||||
} else {
|
||||
if shouldSkipEmptyCollectionOnPersist(sk.Value, sv) {
|
||||
// New key: only add if value is non-zero to avoid polluting config with defaults
|
||||
if isZeroValueNode(sv) {
|
||||
continue
|
||||
}
|
||||
// Append new key/value pair by deep-copying from src
|
||||
dst.Content = append(dst.Content, deepCopyNode(sk), deepCopyNode(sv))
|
||||
}
|
||||
}
|
||||
@@ -969,32 +998,49 @@ func findMapKeyIndex(mapNode *yaml.Node, key string) int {
|
||||
return -1
|
||||
}
|
||||
|
||||
func shouldSkipEmptyCollectionOnPersist(key string, node *yaml.Node) bool {
|
||||
switch key {
|
||||
case "generative-language-api-key",
|
||||
"gemini-api-key",
|
||||
"vertex-api-key",
|
||||
"claude-api-key",
|
||||
"codex-api-key",
|
||||
"openai-compatibility":
|
||||
return isEmptyCollectionNode(node)
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func isEmptyCollectionNode(node *yaml.Node) bool {
|
||||
// isZeroValueNode returns true if the YAML node represents a zero/default value
|
||||
// that should not be written as a new key to preserve config cleanliness.
|
||||
// For mappings and sequences, recursively checks if all children are zero values.
|
||||
func isZeroValueNode(node *yaml.Node) bool {
|
||||
if node == nil {
|
||||
return true
|
||||
}
|
||||
switch node.Kind {
|
||||
case yaml.SequenceNode:
|
||||
return len(node.Content) == 0
|
||||
case yaml.ScalarNode:
|
||||
return node.Tag == "!!null"
|
||||
default:
|
||||
return false
|
||||
switch node.Tag {
|
||||
case "!!bool":
|
||||
return node.Value == "false"
|
||||
case "!!int", "!!float":
|
||||
return node.Value == "0" || node.Value == "0.0"
|
||||
case "!!str":
|
||||
return node.Value == ""
|
||||
case "!!null":
|
||||
return true
|
||||
}
|
||||
case yaml.SequenceNode:
|
||||
if len(node.Content) == 0 {
|
||||
return true
|
||||
}
|
||||
// Check if all elements are zero values
|
||||
for _, child := range node.Content {
|
||||
if !isZeroValueNode(child) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
case yaml.MappingNode:
|
||||
if len(node.Content) == 0 {
|
||||
return true
|
||||
}
|
||||
// Check if all values are zero values (values are at odd indices)
|
||||
for i := 1; i < len(node.Content); i += 2 {
|
||||
if !isZeroValueNode(node.Content[i]) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// deepCopyNode creates a deep copy of a yaml.Node graph.
|
||||
|
||||
@@ -22,6 +22,21 @@ type SDKConfig struct {
|
||||
|
||||
// Access holds request authentication provider configuration.
|
||||
Access AccessConfig `yaml:"auth,omitempty" json:"auth,omitempty"`
|
||||
|
||||
// Streaming configures server-side streaming behavior (keep-alives and safe bootstrap retries).
|
||||
Streaming StreamingConfig `yaml:"streaming" json:"streaming"`
|
||||
}
|
||||
|
||||
// StreamingConfig holds server streaming behavior configuration.
|
||||
type StreamingConfig struct {
|
||||
// KeepAliveSeconds controls how often the server emits SSE heartbeats (": keep-alive\n\n").
|
||||
// <= 0 disables keep-alives. Default is 0.
|
||||
KeepAliveSeconds int `yaml:"keepalive-seconds,omitempty" json:"keepalive-seconds,omitempty"`
|
||||
|
||||
// BootstrapRetries controls how many times the server may retry a streaming request before any bytes are sent,
|
||||
// to allow auth rotation / transient recovery.
|
||||
// <= 0 disables bootstrap retries. Default is 0.
|
||||
BootstrapRetries int `yaml:"bootstrap-retries,omitempty" json:"bootstrap-retries,omitempty"`
|
||||
}
|
||||
|
||||
// AccessConfig groups request authentication providers.
|
||||
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"runtime/debug"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -14,11 +15,24 @@ import (
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// aiAPIPrefixes defines path prefixes for AI API requests that should have request ID tracking.
|
||||
var aiAPIPrefixes = []string{
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/v1/messages",
|
||||
"/v1/responses",
|
||||
"/v1beta/models/",
|
||||
"/api/provider/",
|
||||
}
|
||||
|
||||
const skipGinLogKey = "__gin_skip_request_logging__"
|
||||
|
||||
// GinLogrusLogger returns a Gin middleware handler that logs HTTP requests and responses
|
||||
// using logrus. It captures request details including method, path, status code, latency,
|
||||
// client IP, and any error messages, formatting them in a Gin-style log format.
|
||||
// client IP, and any error messages. Request ID is only added for AI API requests.
|
||||
//
|
||||
// Output format (AI API): [2025-12-23 20:14:10] [info ] | a1b2c3d4 | 200 | 23.559s | ...
|
||||
// Output format (others): [2025-12-23 20:14:10] [info ] | -------- | 200 | 23.559s | ...
|
||||
//
|
||||
// Returns:
|
||||
// - gin.HandlerFunc: A middleware handler for request logging
|
||||
@@ -28,6 +42,15 @@ func GinLogrusLogger() gin.HandlerFunc {
|
||||
path := c.Request.URL.Path
|
||||
raw := util.MaskSensitiveQuery(c.Request.URL.RawQuery)
|
||||
|
||||
// Only generate request ID for AI API paths
|
||||
var requestID string
|
||||
if isAIAPIPath(path) {
|
||||
requestID = GenerateRequestID()
|
||||
SetGinRequestID(c, requestID)
|
||||
ctx := WithRequestID(c.Request.Context(), requestID)
|
||||
c.Request = c.Request.WithContext(ctx)
|
||||
}
|
||||
|
||||
c.Next()
|
||||
|
||||
if shouldSkipGinRequestLogging(c) {
|
||||
@@ -49,23 +72,38 @@ func GinLogrusLogger() gin.HandlerFunc {
|
||||
clientIP := c.ClientIP()
|
||||
method := c.Request.Method
|
||||
errorMessage := c.Errors.ByType(gin.ErrorTypePrivate).String()
|
||||
timestamp := time.Now().Format("2006/01/02 - 15:04:05")
|
||||
logLine := fmt.Sprintf("[GIN] %s | %3d | %13v | %15s | %-7s \"%s\"", timestamp, statusCode, latency, clientIP, method, path)
|
||||
|
||||
if requestID == "" {
|
||||
requestID = "--------"
|
||||
}
|
||||
logLine := fmt.Sprintf("%3d | %13v | %15s | %-7s \"%s\"", statusCode, latency, clientIP, method, path)
|
||||
if errorMessage != "" {
|
||||
logLine = logLine + " | " + errorMessage
|
||||
}
|
||||
|
||||
entry := log.WithField("request_id", requestID)
|
||||
|
||||
switch {
|
||||
case statusCode >= http.StatusInternalServerError:
|
||||
log.Error(logLine)
|
||||
entry.Error(logLine)
|
||||
case statusCode >= http.StatusBadRequest:
|
||||
log.Warn(logLine)
|
||||
entry.Warn(logLine)
|
||||
default:
|
||||
log.Info(logLine)
|
||||
entry.Info(logLine)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// isAIAPIPath checks if the given path is an AI API endpoint that should have request ID tracking.
|
||||
func isAIAPIPath(path string) bool {
|
||||
for _, prefix := range aiAPIPrefixes {
|
||||
if strings.HasPrefix(path, prefix) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// GinLogrusRecovery returns a Gin middleware handler that recovers from panics and logs
|
||||
// them using logrus. When a panic occurs, it captures the panic value, stack trace,
|
||||
// and request path, then returns a 500 Internal Server Error response to the client.
|
||||
|
||||
@@ -24,7 +24,8 @@ var (
|
||||
)
|
||||
|
||||
// LogFormatter defines a custom log format for logrus.
|
||||
// This formatter adds timestamp, level, and source location to each log entry.
|
||||
// This formatter adds timestamp, level, request ID, and source location to each log entry.
|
||||
// Format: [2025-12-23 20:14:04] [debug] [manager.go:524] | a1b2c3d4 | Use API key sk-9...0RHO for model gpt-5.2
|
||||
type LogFormatter struct{}
|
||||
|
||||
// Format renders a single log entry with custom formatting.
|
||||
@@ -38,16 +39,24 @@ func (m *LogFormatter) Format(entry *log.Entry) ([]byte, error) {
|
||||
|
||||
timestamp := entry.Time.Format("2006-01-02 15:04:05")
|
||||
message := strings.TrimRight(entry.Message, "\r\n")
|
||||
|
||||
// Handle nil Caller (can happen with some log entries)
|
||||
callerFile := "unknown"
|
||||
callerLine := 0
|
||||
if entry.Caller != nil {
|
||||
callerFile = filepath.Base(entry.Caller.File)
|
||||
callerLine = entry.Caller.Line
|
||||
|
||||
reqID := "--------"
|
||||
if id, ok := entry.Data["request_id"].(string); ok && id != "" {
|
||||
reqID = id
|
||||
}
|
||||
|
||||
level := entry.Level.String()
|
||||
if level == "warning" {
|
||||
level = "warn"
|
||||
}
|
||||
levelStr := fmt.Sprintf("%-5s", level)
|
||||
|
||||
var formatted string
|
||||
if entry.Caller != nil {
|
||||
formatted = fmt.Sprintf("[%s] [%s] [%s] [%s:%d] %s\n", timestamp, reqID, levelStr, filepath.Base(entry.Caller.File), entry.Caller.Line, message)
|
||||
} else {
|
||||
formatted = fmt.Sprintf("[%s] [%s] [%s] %s\n", timestamp, reqID, levelStr, message)
|
||||
}
|
||||
|
||||
formatted := fmt.Sprintf("[%s] [%s] [%s:%d] %s\n", timestamp, entry.Level, callerFile, callerLine, message)
|
||||
buffer.WriteString(formatted)
|
||||
|
||||
return buffer.Bytes(), nil
|
||||
|
||||
@@ -43,10 +43,11 @@ type RequestLogger interface {
|
||||
// - response: The raw response data
|
||||
// - apiRequest: The API request data
|
||||
// - apiResponse: The API response data
|
||||
// - requestID: Optional request ID for log file naming
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if logging fails, nil otherwise
|
||||
LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage) error
|
||||
LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string) error
|
||||
|
||||
// LogStreamingRequest initiates logging for a streaming request and returns a writer for chunks.
|
||||
//
|
||||
@@ -55,11 +56,12 @@ type RequestLogger interface {
|
||||
// - method: The HTTP method
|
||||
// - headers: The request headers
|
||||
// - body: The request body
|
||||
// - requestID: Optional request ID for log file naming
|
||||
//
|
||||
// Returns:
|
||||
// - StreamingLogWriter: A writer for streaming response chunks
|
||||
// - error: An error if logging initialization fails, nil otherwise
|
||||
LogStreamingRequest(url, method string, headers map[string][]string, body []byte) (StreamingLogWriter, error)
|
||||
LogStreamingRequest(url, method string, headers map[string][]string, body []byte, requestID string) (StreamingLogWriter, error)
|
||||
|
||||
// IsEnabled returns whether request logging is currently enabled.
|
||||
//
|
||||
@@ -177,20 +179,21 @@ func (l *FileRequestLogger) SetEnabled(enabled bool) {
|
||||
// - response: The raw response data
|
||||
// - apiRequest: The API request data
|
||||
// - apiResponse: The API response data
|
||||
// - requestID: Optional request ID for log file naming
|
||||
//
|
||||
// Returns:
|
||||
// - error: An error if logging fails, nil otherwise
|
||||
func (l *FileRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage) error {
|
||||
return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, apiRequest, apiResponse, apiResponseErrors, false)
|
||||
func (l *FileRequestLogger) LogRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, requestID string) error {
|
||||
return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, apiRequest, apiResponse, apiResponseErrors, false, requestID)
|
||||
}
|
||||
|
||||
// LogRequestWithOptions logs a request with optional forced logging behavior.
|
||||
// The force flag allows writing error logs even when regular request logging is disabled.
|
||||
func (l *FileRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool) error {
|
||||
return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, apiRequest, apiResponse, apiResponseErrors, force)
|
||||
func (l *FileRequestLogger) LogRequestWithOptions(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string) error {
|
||||
return l.logRequest(url, method, requestHeaders, body, statusCode, responseHeaders, response, apiRequest, apiResponse, apiResponseErrors, force, requestID)
|
||||
}
|
||||
|
||||
func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool) error {
|
||||
func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[string][]string, body []byte, statusCode int, responseHeaders map[string][]string, response, apiRequest, apiResponse []byte, apiResponseErrors []*interfaces.ErrorMessage, force bool, requestID string) error {
|
||||
if !l.enabled && !force {
|
||||
return nil
|
||||
}
|
||||
@@ -200,10 +203,10 @@ func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[st
|
||||
return fmt.Errorf("failed to create logs directory: %w", errEnsure)
|
||||
}
|
||||
|
||||
// Generate filename
|
||||
filename := l.generateFilename(url)
|
||||
// Generate filename with request ID
|
||||
filename := l.generateFilename(url, requestID)
|
||||
if force && !l.enabled {
|
||||
filename = l.generateErrorFilename(url)
|
||||
filename = l.generateErrorFilename(url, requestID)
|
||||
}
|
||||
filePath := filepath.Join(l.logsDir, filename)
|
||||
|
||||
@@ -271,11 +274,12 @@ func (l *FileRequestLogger) logRequest(url, method string, requestHeaders map[st
|
||||
// - method: The HTTP method
|
||||
// - headers: The request headers
|
||||
// - body: The request body
|
||||
// - requestID: Optional request ID for log file naming
|
||||
//
|
||||
// Returns:
|
||||
// - StreamingLogWriter: A writer for streaming response chunks
|
||||
// - error: An error if logging initialization fails, nil otherwise
|
||||
func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[string][]string, body []byte) (StreamingLogWriter, error) {
|
||||
func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[string][]string, body []byte, requestID string) (StreamingLogWriter, error) {
|
||||
if !l.enabled {
|
||||
return &NoOpStreamingLogWriter{}, nil
|
||||
}
|
||||
@@ -285,8 +289,8 @@ func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[
|
||||
return nil, fmt.Errorf("failed to create logs directory: %w", err)
|
||||
}
|
||||
|
||||
// Generate filename
|
||||
filename := l.generateFilename(url)
|
||||
// Generate filename with request ID
|
||||
filename := l.generateFilename(url, requestID)
|
||||
filePath := filepath.Join(l.logsDir, filename)
|
||||
|
||||
requestHeaders := make(map[string][]string, len(headers))
|
||||
@@ -330,8 +334,8 @@ func (l *FileRequestLogger) LogStreamingRequest(url, method string, headers map[
|
||||
}
|
||||
|
||||
// generateErrorFilename creates a filename with an error prefix to differentiate forced error logs.
|
||||
func (l *FileRequestLogger) generateErrorFilename(url string) string {
|
||||
return fmt.Sprintf("error-%s", l.generateFilename(url))
|
||||
func (l *FileRequestLogger) generateErrorFilename(url string, requestID ...string) string {
|
||||
return fmt.Sprintf("error-%s", l.generateFilename(url, requestID...))
|
||||
}
|
||||
|
||||
// ensureLogsDir creates the logs directory if it doesn't exist.
|
||||
@@ -346,13 +350,15 @@ func (l *FileRequestLogger) ensureLogsDir() error {
|
||||
}
|
||||
|
||||
// generateFilename creates a sanitized filename from the URL path and current timestamp.
|
||||
// Format: v1-responses-2025-12-23T195811-a1b2c3d4.log
|
||||
//
|
||||
// Parameters:
|
||||
// - url: The request URL
|
||||
// - requestID: Optional request ID to include in filename
|
||||
//
|
||||
// Returns:
|
||||
// - string: A sanitized filename for the log file
|
||||
func (l *FileRequestLogger) generateFilename(url string) string {
|
||||
func (l *FileRequestLogger) generateFilename(url string, requestID ...string) string {
|
||||
// Extract path from URL
|
||||
path := url
|
||||
if strings.Contains(url, "?") {
|
||||
@@ -368,12 +374,18 @@ func (l *FileRequestLogger) generateFilename(url string) string {
|
||||
sanitized := l.sanitizeForFilename(path)
|
||||
|
||||
// Add timestamp
|
||||
timestamp := time.Now().Format("2006-01-02T150405-.000000000")
|
||||
timestamp = strings.Replace(timestamp, ".", "", -1)
|
||||
timestamp := time.Now().Format("2006-01-02T150405")
|
||||
|
||||
id := requestLogID.Add(1)
|
||||
// Use request ID if provided, otherwise use sequential ID
|
||||
var idPart string
|
||||
if len(requestID) > 0 && requestID[0] != "" {
|
||||
idPart = requestID[0]
|
||||
} else {
|
||||
id := requestLogID.Add(1)
|
||||
idPart = fmt.Sprintf("%d", id)
|
||||
}
|
||||
|
||||
return fmt.Sprintf("%s-%s-%d.log", sanitized, timestamp, id)
|
||||
return fmt.Sprintf("%s-%s-%s.log", sanitized, timestamp, idPart)
|
||||
}
|
||||
|
||||
// sanitizeForFilename replaces characters that are not safe for filenames.
|
||||
|
||||
61
internal/logging/requestid.go
Normal file
61
internal/logging/requestid.go
Normal file
@@ -0,0 +1,61 @@
|
||||
package logging
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// requestIDKey is the context key for storing/retrieving request IDs.
|
||||
type requestIDKey struct{}
|
||||
|
||||
// ginRequestIDKey is the Gin context key for request IDs.
|
||||
const ginRequestIDKey = "__request_id__"
|
||||
|
||||
// GenerateRequestID creates a new 8-character hex request ID.
|
||||
func GenerateRequestID() string {
|
||||
b := make([]byte, 4)
|
||||
if _, err := rand.Read(b); err != nil {
|
||||
return "00000000"
|
||||
}
|
||||
return hex.EncodeToString(b)
|
||||
}
|
||||
|
||||
// WithRequestID returns a new context with the request ID attached.
|
||||
func WithRequestID(ctx context.Context, requestID string) context.Context {
|
||||
return context.WithValue(ctx, requestIDKey{}, requestID)
|
||||
}
|
||||
|
||||
// GetRequestID retrieves the request ID from the context.
|
||||
// Returns empty string if not found.
|
||||
func GetRequestID(ctx context.Context) string {
|
||||
if ctx == nil {
|
||||
return ""
|
||||
}
|
||||
if id, ok := ctx.Value(requestIDKey{}).(string); ok {
|
||||
return id
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// SetGinRequestID stores the request ID in the Gin context.
|
||||
func SetGinRequestID(c *gin.Context, requestID string) {
|
||||
if c != nil {
|
||||
c.Set(ginRequestIDKey, requestID)
|
||||
}
|
||||
}
|
||||
|
||||
// GetGinRequestID retrieves the request ID from the Gin context.
|
||||
func GetGinRequestID(c *gin.Context) string {
|
||||
if c == nil {
|
||||
return ""
|
||||
}
|
||||
if id, exists := c.Get(ginRequestIDKey); exists {
|
||||
if s, ok := id.(string); ok {
|
||||
return s
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
@@ -727,6 +727,7 @@ func GetIFlowModels() []*ModelInfo {
|
||||
{ID: "qwen3-max-preview", DisplayName: "Qwen3-Max-Preview", Description: "Qwen3 Max preview build", Created: 1757030400},
|
||||
{ID: "kimi-k2-0905", DisplayName: "Kimi-K2-Instruct-0905", Description: "Moonshot Kimi K2 instruct 0905", Created: 1757030400},
|
||||
{ID: "glm-4.6", DisplayName: "GLM-4.6", Description: "Zhipu GLM 4.6 general model", Created: 1759190400, Thinking: iFlowThinkingSupport},
|
||||
{ID: "glm-4.7", DisplayName: "GLM-4.7", Description: "Zhipu GLM 4.7 general model", Created: 1766448000, Thinking: iFlowThinkingSupport},
|
||||
{ID: "kimi-k2", DisplayName: "Kimi-K2", Description: "Moonshot Kimi K2 general model", Created: 1752192000},
|
||||
{ID: "kimi-k2-thinking", DisplayName: "Kimi-K2-Thinking", Description: "Moonshot Kimi K2 thinking model", Created: 1762387200},
|
||||
{ID: "deepseek-v3.2-chat", DisplayName: "DeepSeek-V3.2", Description: "DeepSeek V3.2 Chat", Created: 1764576000},
|
||||
@@ -740,6 +741,7 @@ func GetIFlowModels() []*ModelInfo {
|
||||
{ID: "qwen3-235b-a22b-instruct", DisplayName: "Qwen3-235B-A22B-Instruct", Description: "Qwen3 235B A22B Instruct", Created: 1753401600},
|
||||
{ID: "qwen3-235b", DisplayName: "Qwen3-235B-A22B", Description: "Qwen3 235B A22B", Created: 1753401600},
|
||||
{ID: "minimax-m2", DisplayName: "MiniMax-M2", Description: "MiniMax M2", Created: 1758672000},
|
||||
{ID: "minimax-m2.1", DisplayName: "MiniMax-M2.1", Description: "MiniMax M2.1", Created: 1766448000, Thinking: iFlowThinkingSupport},
|
||||
}
|
||||
models := make([]*ModelInfo, 0, len(entries))
|
||||
for _, entry := range entries {
|
||||
|
||||
@@ -33,18 +33,18 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
antigravityBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
// antigravityBaseURLAutopush = "https://autopush-cloudcode-pa.sandbox.googleapis.com"
|
||||
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityCountTokensPath = "/v1internal:countTokens"
|
||||
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
||||
antigravityGeneratePath = "/v1internal:generateContent"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityClientSecret = "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"
|
||||
defaultAntigravityAgent = "antigravity/1.11.5 windows/amd64"
|
||||
antigravityAuthType = "antigravity"
|
||||
refreshSkew = 3000 * time.Second
|
||||
antigravityBaseURLDaily = "https://daily-cloudcode-pa.googleapis.com"
|
||||
antigravitySandboxBaseURLDaily = "https://daily-cloudcode-pa.sandbox.googleapis.com"
|
||||
antigravityBaseURLProd = "https://cloudcode-pa.googleapis.com"
|
||||
antigravityCountTokensPath = "/v1internal:countTokens"
|
||||
antigravityStreamPath = "/v1internal:streamGenerateContent"
|
||||
antigravityGeneratePath = "/v1internal:generateContent"
|
||||
antigravityModelsPath = "/v1internal:fetchAvailableModels"
|
||||
antigravityClientID = "1071006060591-tmhssin2h21lcre235vtolojh4g403ep.apps.googleusercontent.com"
|
||||
antigravityClientSecret = "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"
|
||||
defaultAntigravityAgent = "antigravity/1.104.0 darwin/arm64"
|
||||
antigravityAuthType = "antigravity"
|
||||
refreshSkew = 3000 * time.Second
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -1156,7 +1156,7 @@ func antigravityBaseURLFallbackOrder(auth *cliproxyauth.Auth) []string {
|
||||
}
|
||||
return []string{
|
||||
antigravityBaseURLDaily,
|
||||
// antigravityBaseURLAutopush,
|
||||
antigravitySandboxBaseURLDaily,
|
||||
antigravityBaseURLProd,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -662,7 +662,14 @@ func decodeResponseBody(body io.ReadCloser, contentEncoding string) (io.ReadClos
|
||||
}
|
||||
|
||||
func applyClaudeHeaders(r *http.Request, auth *cliproxyauth.Auth, apiKey string, stream bool, extraBetas []string) {
|
||||
r.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
useAPIKey := auth != nil && auth.Attributes != nil && strings.TrimSpace(auth.Attributes["api_key"]) != ""
|
||||
isAnthropicBase := r.URL != nil && strings.EqualFold(r.URL.Scheme, "https") && strings.EqualFold(r.URL.Host, "api.anthropic.com")
|
||||
if isAnthropicBase && useAPIKey {
|
||||
r.Header.Del("Authorization")
|
||||
r.Header.Set("x-api-key", apiKey)
|
||||
} else {
|
||||
r.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
r.Header.Set("Content-Type", "application/json")
|
||||
|
||||
var ginHeaders http.Header
|
||||
|
||||
@@ -50,6 +50,16 @@ func (e *CodexExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
||||
if upstreamModel == "" {
|
||||
upstreamModel = req.Model
|
||||
}
|
||||
if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
} else if !strings.EqualFold(upstreamModel, req.Model) {
|
||||
if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
}
|
||||
}
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("codex")
|
||||
@@ -147,6 +157,16 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
||||
defer reporter.trackFailure(ctx, &err)
|
||||
|
||||
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
||||
if upstreamModel == "" {
|
||||
upstreamModel = req.Model
|
||||
}
|
||||
if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
} else if !strings.EqualFold(upstreamModel, req.Model) {
|
||||
if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
}
|
||||
}
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("codex")
|
||||
@@ -247,12 +267,22 @@ func (e *CodexExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
||||
|
||||
func (e *CodexExecutor) CountTokens(ctx context.Context, auth *cliproxyauth.Auth, req cliproxyexecutor.Request, opts cliproxyexecutor.Options) (cliproxyexecutor.Response, error) {
|
||||
upstreamModel := util.ResolveOriginalModel(req.Model, req.Metadata)
|
||||
if upstreamModel == "" {
|
||||
upstreamModel = req.Model
|
||||
}
|
||||
if modelOverride := e.resolveUpstreamModel(upstreamModel, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
} else if !strings.EqualFold(upstreamModel, req.Model) {
|
||||
if modelOverride := e.resolveUpstreamModel(req.Model, auth); modelOverride != "" {
|
||||
upstreamModel = modelOverride
|
||||
}
|
||||
}
|
||||
|
||||
from := opts.SourceFormat
|
||||
to := sdktranslator.FromString("codex")
|
||||
body := sdktranslator.TranslateRequest(from, to, req.Model, bytes.Clone(req.Payload), false)
|
||||
|
||||
modelForCounting := req.Model
|
||||
modelForCounting := upstreamModel
|
||||
|
||||
body = ApplyReasoningEffortMetadata(body, req.Metadata, req.Model, "reasoning.effort", false)
|
||||
body, _ = sjson.SetBytes(body, "model", upstreamModel)
|
||||
@@ -520,3 +550,87 @@ func codexCreds(a *cliproxyauth.Auth) (apiKey, baseURL string) {
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (e *CodexExecutor) resolveUpstreamModel(alias string, auth *cliproxyauth.Auth) string {
|
||||
trimmed := strings.TrimSpace(alias)
|
||||
if trimmed == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
entry := e.resolveCodexConfig(auth)
|
||||
if entry == nil {
|
||||
return ""
|
||||
}
|
||||
|
||||
normalizedModel, metadata := util.NormalizeThinkingModel(trimmed)
|
||||
|
||||
// Candidate names to match against configured aliases/names.
|
||||
candidates := []string{strings.TrimSpace(normalizedModel)}
|
||||
if !strings.EqualFold(normalizedModel, trimmed) {
|
||||
candidates = append(candidates, trimmed)
|
||||
}
|
||||
if original := util.ResolveOriginalModel(normalizedModel, metadata); original != "" && !strings.EqualFold(original, normalizedModel) {
|
||||
candidates = append(candidates, original)
|
||||
}
|
||||
|
||||
for i := range entry.Models {
|
||||
model := entry.Models[i]
|
||||
name := strings.TrimSpace(model.Name)
|
||||
modelAlias := strings.TrimSpace(model.Alias)
|
||||
|
||||
for _, candidate := range candidates {
|
||||
if candidate == "" {
|
||||
continue
|
||||
}
|
||||
if modelAlias != "" && strings.EqualFold(modelAlias, candidate) {
|
||||
if name != "" {
|
||||
return name
|
||||
}
|
||||
return candidate
|
||||
}
|
||||
if name != "" && strings.EqualFold(name, candidate) {
|
||||
return name
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func (e *CodexExecutor) resolveCodexConfig(auth *cliproxyauth.Auth) *config.CodexKey {
|
||||
if auth == nil || e.cfg == nil {
|
||||
return nil
|
||||
}
|
||||
var attrKey, attrBase string
|
||||
if auth.Attributes != nil {
|
||||
attrKey = strings.TrimSpace(auth.Attributes["api_key"])
|
||||
attrBase = strings.TrimSpace(auth.Attributes["base_url"])
|
||||
}
|
||||
for i := range e.cfg.CodexKey {
|
||||
entry := &e.cfg.CodexKey[i]
|
||||
cfgKey := strings.TrimSpace(entry.APIKey)
|
||||
cfgBase := strings.TrimSpace(entry.BaseURL)
|
||||
if attrKey != "" && attrBase != "" {
|
||||
if strings.EqualFold(cfgKey, attrKey) && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
continue
|
||||
}
|
||||
if attrKey != "" && strings.EqualFold(cfgKey, attrKey) {
|
||||
if cfgBase == "" || strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
if attrKey == "" && attrBase != "" && strings.EqualFold(cfgBase, attrBase) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
if attrKey != "" {
|
||||
for i := range e.cfg.CodexKey {
|
||||
entry := &e.cfg.CodexKey[i]
|
||||
if strings.EqualFold(strings.TrimSpace(entry.APIKey), attrKey) {
|
||||
return entry
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -67,6 +67,7 @@ func (e *IFlowExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, re
|
||||
return resp, errValidate
|
||||
}
|
||||
body = applyIFlowThinkingConfig(body)
|
||||
body = preserveReasoningContentInMessages(body)
|
||||
body = applyPayloadConfig(e.cfg, req.Model, body)
|
||||
|
||||
endpoint := strings.TrimSuffix(baseURL, "/") + iflowDefaultEndpoint
|
||||
@@ -159,6 +160,7 @@ func (e *IFlowExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Au
|
||||
return nil, errValidate
|
||||
}
|
||||
body = applyIFlowThinkingConfig(body)
|
||||
body = preserveReasoningContentInMessages(body)
|
||||
// Ensure tools array exists to avoid provider quirks similar to Qwen's behaviour.
|
||||
toolsResult := gjson.GetBytes(body, "tools")
|
||||
if toolsResult.Exists() && toolsResult.IsArray() && len(toolsResult.Array()) == 0 {
|
||||
@@ -445,20 +447,98 @@ func ensureToolsArray(body []byte) []byte {
|
||||
return updated
|
||||
}
|
||||
|
||||
// applyIFlowThinkingConfig converts normalized reasoning_effort to iFlow chat_template_kwargs.enable_thinking.
|
||||
// This should be called after NormalizeThinkingConfig has processed the payload.
|
||||
// iFlow only supports boolean enable_thinking, so any non-"none" effort enables thinking.
|
||||
func applyIFlowThinkingConfig(body []byte) []byte {
|
||||
effort := gjson.GetBytes(body, "reasoning_effort")
|
||||
if !effort.Exists() {
|
||||
// preserveReasoningContentInMessages ensures reasoning_content from assistant messages in the
|
||||
// conversation history is preserved when sending to iFlow models that support thinking.
|
||||
// This is critical for multi-turn conversations where the model needs to see its previous
|
||||
// reasoning to maintain coherent thought chains across tool calls and conversation turns.
|
||||
//
|
||||
// For GLM-4.7 and MiniMax-M2.1, the full assistant response (including reasoning) must be
|
||||
// appended back into message history before the next call.
|
||||
func preserveReasoningContentInMessages(body []byte) []byte {
|
||||
model := strings.ToLower(gjson.GetBytes(body, "model").String())
|
||||
|
||||
// Only apply to models that support thinking with history preservation
|
||||
needsPreservation := strings.HasPrefix(model, "glm-4.7") ||
|
||||
strings.HasPrefix(model, "glm-4-7") ||
|
||||
strings.HasPrefix(model, "minimax-m2.1") ||
|
||||
strings.HasPrefix(model, "minimax-m2-1")
|
||||
|
||||
if !needsPreservation {
|
||||
return body
|
||||
}
|
||||
|
||||
val := strings.ToLower(strings.TrimSpace(effort.String()))
|
||||
enableThinking := val != "none" && val != ""
|
||||
messages := gjson.GetBytes(body, "messages")
|
||||
if !messages.Exists() || !messages.IsArray() {
|
||||
return body
|
||||
}
|
||||
|
||||
body, _ = sjson.DeleteBytes(body, "reasoning_effort")
|
||||
body, _ = sjson.SetBytes(body, "chat_template_kwargs.enable_thinking", enableThinking)
|
||||
// Check if any assistant message already has reasoning_content preserved
|
||||
hasReasoningContent := false
|
||||
messages.ForEach(func(_, msg gjson.Result) bool {
|
||||
role := msg.Get("role").String()
|
||||
if role == "assistant" {
|
||||
rc := msg.Get("reasoning_content")
|
||||
if rc.Exists() && rc.String() != "" {
|
||||
hasReasoningContent = true
|
||||
return false // stop iteration
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
// If reasoning content is already present, the messages are properly formatted
|
||||
// No need to modify - the client has correctly preserved reasoning in history
|
||||
if hasReasoningContent {
|
||||
log.Debugf("iflow executor: reasoning_content found in message history for %s", model)
|
||||
}
|
||||
|
||||
return body
|
||||
}
|
||||
|
||||
// applyIFlowThinkingConfig converts normalized reasoning_effort to model-specific thinking configurations.
|
||||
// This should be called after NormalizeThinkingConfig has processed the payload.
|
||||
//
|
||||
// Model-specific handling:
|
||||
// - GLM-4.7: Uses extra_body={"thinking": {"type": "enabled"}, "clear_thinking": false}
|
||||
// - MiniMax-M2.1: Uses reasoning_split=true for OpenAI-style reasoning separation
|
||||
// - Other iFlow models: Uses chat_template_kwargs.enable_thinking (boolean)
|
||||
func applyIFlowThinkingConfig(body []byte) []byte {
|
||||
effort := gjson.GetBytes(body, "reasoning_effort")
|
||||
model := strings.ToLower(gjson.GetBytes(body, "model").String())
|
||||
|
||||
// Check if thinking should be enabled
|
||||
val := ""
|
||||
if effort.Exists() {
|
||||
val = strings.ToLower(strings.TrimSpace(effort.String()))
|
||||
}
|
||||
enableThinking := effort.Exists() && val != "none" && val != ""
|
||||
|
||||
// Remove reasoning_effort as we'll convert to model-specific format
|
||||
if effort.Exists() {
|
||||
body, _ = sjson.DeleteBytes(body, "reasoning_effort")
|
||||
}
|
||||
|
||||
// GLM-4.7: Use extra_body with thinking config and clear_thinking: false
|
||||
if strings.HasPrefix(model, "glm-4.7") || strings.HasPrefix(model, "glm-4-7") {
|
||||
if enableThinking {
|
||||
body, _ = sjson.SetBytes(body, "extra_body.thinking.type", "enabled")
|
||||
body, _ = sjson.SetBytes(body, "extra_body.clear_thinking", false)
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
// MiniMax-M2.1: Use reasoning_split=true for interleaved thinking
|
||||
if strings.HasPrefix(model, "minimax-m2.1") || strings.HasPrefix(model, "minimax-m2-1") {
|
||||
if enableThinking {
|
||||
body, _ = sjson.SetBytes(body, "reasoning_split", true)
|
||||
}
|
||||
return body
|
||||
}
|
||||
|
||||
// Other iFlow models (including GLM-4.6): Use chat_template_kwargs.enable_thinking
|
||||
if effort.Exists() {
|
||||
body, _ = sjson.SetBytes(body, "chat_template_kwargs.enable_thinking", enableThinking)
|
||||
}
|
||||
|
||||
return body
|
||||
}
|
||||
|
||||
@@ -10,6 +10,8 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
@@ -43,10 +45,15 @@ const (
|
||||
// Event Stream error type constants
|
||||
ErrStreamFatal = "fatal" // Connection/authentication errors, not recoverable
|
||||
ErrStreamMalformed = "malformed" // Format errors, data cannot be parsed
|
||||
// kiroUserAgent matches amq2api format for User-Agent header
|
||||
// kiroUserAgent matches amq2api format for User-Agent header (Amazon Q CLI style)
|
||||
kiroUserAgent = "aws-sdk-rust/1.3.9 os/macos lang/rust/1.87.0"
|
||||
// kiroFullUserAgent is the complete x-amz-user-agent header matching amq2api
|
||||
// kiroFullUserAgent is the complete x-amz-user-agent header matching amq2api (Amazon Q CLI style)
|
||||
kiroFullUserAgent = "aws-sdk-rust/1.3.9 ua/2.1 api/ssooidc/1.88.0 os/macos lang/rust/1.87.0 m/E app/AmazonQ-For-CLI"
|
||||
|
||||
// Kiro IDE style headers (from kiro2api - for IDC auth)
|
||||
kiroIDEUserAgent = "aws-sdk-js/1.0.18 ua/2.1 os/darwin#25.0.0 lang/js md/nodejs#20.16.0 api/codewhispererstreaming#1.0.18 m/E KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1"
|
||||
kiroIDEAmzUserAgent = "aws-sdk-js/1.0.18 KiroIDE-0.2.13-66c23a8c5d15afabec89ef9954ef52a119f10d369df04d548fc6c1eac694b0d1"
|
||||
kiroIDEAgentModeSpec = "spec"
|
||||
)
|
||||
|
||||
// Real-time usage estimation configuration
|
||||
@@ -101,11 +108,24 @@ var kiroEndpointConfigs = []kiroEndpointConfig{
|
||||
|
||||
// getKiroEndpointConfigs returns the list of Kiro API endpoint configurations to try in order.
|
||||
// Supports reordering based on "preferred_endpoint" in auth metadata/attributes.
|
||||
// For IDC auth method, automatically uses CodeWhisperer endpoint with CLI origin.
|
||||
func getKiroEndpointConfigs(auth *cliproxyauth.Auth) []kiroEndpointConfig {
|
||||
if auth == nil {
|
||||
return kiroEndpointConfigs
|
||||
}
|
||||
|
||||
// For IDC auth, use CodeWhisperer endpoint with AI_EDITOR origin (same as Social auth)
|
||||
// Based on kiro2api analysis: IDC tokens work with CodeWhisperer endpoint using Bearer auth
|
||||
// The difference is only in how tokens are refreshed (OIDC with clientId/clientSecret for IDC)
|
||||
// NOT in how API calls are made - both Social and IDC use the same endpoint/origin
|
||||
if auth.Metadata != nil {
|
||||
authMethod, _ := auth.Metadata["auth_method"].(string)
|
||||
if authMethod == "idc" {
|
||||
log.Debugf("kiro: IDC auth, using CodeWhisperer endpoint")
|
||||
return kiroEndpointConfigs
|
||||
}
|
||||
}
|
||||
|
||||
// Check for preference
|
||||
var preference string
|
||||
if auth.Metadata != nil {
|
||||
@@ -162,6 +182,15 @@ type KiroExecutor struct {
|
||||
refreshMu sync.Mutex // Serializes token refresh operations to prevent race conditions
|
||||
}
|
||||
|
||||
// isIDCAuth checks if the auth uses IDC (Identity Center) authentication method.
|
||||
func isIDCAuth(auth *cliproxyauth.Auth) bool {
|
||||
if auth == nil || auth.Metadata == nil {
|
||||
return false
|
||||
}
|
||||
authMethod, _ := auth.Metadata["auth_method"].(string)
|
||||
return authMethod == "idc"
|
||||
}
|
||||
|
||||
// buildKiroPayloadForFormat builds the Kiro API payload based on the source format.
|
||||
// This is critical because OpenAI and Claude formats have different tool structures:
|
||||
// - OpenAI: tools[].function.name, tools[].function.description
|
||||
@@ -210,6 +239,10 @@ func (e *KiroExecutor) Execute(ctx context.Context, auth *cliproxyauth.Auth, req
|
||||
log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr)
|
||||
} else if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
log.Infof("kiro: token refreshed successfully before request")
|
||||
}
|
||||
@@ -262,15 +295,28 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", kiroContentType)
|
||||
httpReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
httpReq.Header.Set("Accept", kiroAcceptStream)
|
||||
// Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)
|
||||
httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget)
|
||||
httpReq.Header.Set("User-Agent", kiroUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
|
||||
|
||||
// Use different headers based on auth type
|
||||
// IDC auth uses Kiro IDE style headers (from kiro2api)
|
||||
// Other auth types use Amazon Q CLI style headers
|
||||
if isIDCAuth(auth) {
|
||||
httpReq.Header.Set("User-Agent", kiroIDEUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)
|
||||
httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)
|
||||
log.Debugf("kiro: using Kiro IDE headers for IDC auth")
|
||||
} else {
|
||||
httpReq.Header.Set("User-Agent", kiroUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
|
||||
}
|
||||
httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
|
||||
httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())
|
||||
|
||||
// Bearer token authentication for all auth types (Builder ID, IDC, social, etc.)
|
||||
httpReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
var attrs map[string]string
|
||||
if auth != nil {
|
||||
attrs = auth.Attributes
|
||||
@@ -358,6 +404,11 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.
|
||||
|
||||
if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
// Continue anyway - the token is valid for this request
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
// Rebuild payload with new profile ARN if changed
|
||||
kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
|
||||
@@ -416,6 +467,11 @@ func (e *KiroExecutor) executeWithRetry(ctx context.Context, auth *cliproxyauth.
|
||||
}
|
||||
if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
// Continue anyway - the token is valid for this request
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
|
||||
log.Infof("kiro: token refreshed for 403, retrying request")
|
||||
@@ -518,6 +574,10 @@ func (e *KiroExecutor) ExecuteStream(ctx context.Context, auth *cliproxyauth.Aut
|
||||
log.Warnf("kiro: pre-request token refresh failed: %v", refreshErr)
|
||||
} else if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
log.Infof("kiro: token refreshed successfully before stream request")
|
||||
}
|
||||
@@ -568,15 +628,28 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", kiroContentType)
|
||||
httpReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
httpReq.Header.Set("Accept", kiroAcceptStream)
|
||||
// Use endpoint-specific X-Amz-Target (critical for avoiding 403 errors)
|
||||
httpReq.Header.Set("X-Amz-Target", endpointConfig.AmzTarget)
|
||||
httpReq.Header.Set("User-Agent", kiroUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
|
||||
|
||||
// Use different headers based on auth type
|
||||
// IDC auth uses Kiro IDE style headers (from kiro2api)
|
||||
// Other auth types use Amazon Q CLI style headers
|
||||
if isIDCAuth(auth) {
|
||||
httpReq.Header.Set("User-Agent", kiroIDEUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroIDEAmzUserAgent)
|
||||
httpReq.Header.Set("x-amzn-kiro-agent-mode", kiroIDEAgentModeSpec)
|
||||
log.Debugf("kiro: using Kiro IDE headers for IDC auth")
|
||||
} else {
|
||||
httpReq.Header.Set("User-Agent", kiroUserAgent)
|
||||
httpReq.Header.Set("X-Amz-User-Agent", kiroFullUserAgent)
|
||||
}
|
||||
httpReq.Header.Set("Amz-Sdk-Request", "attempt=1; max=3")
|
||||
httpReq.Header.Set("Amz-Sdk-Invocation-Id", uuid.New().String())
|
||||
|
||||
// Bearer token authentication for all auth types (Builder ID, IDC, social, etc.)
|
||||
httpReq.Header.Set("Authorization", "Bearer "+accessToken)
|
||||
|
||||
var attrs map[string]string
|
||||
if auth != nil {
|
||||
attrs = auth.Attributes
|
||||
@@ -677,6 +750,11 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox
|
||||
|
||||
if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
// Continue anyway - the token is valid for this request
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
// Rebuild payload with new profile ARN if changed
|
||||
kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
|
||||
@@ -735,6 +813,11 @@ func (e *KiroExecutor) executeStreamWithRetry(ctx context.Context, auth *cliprox
|
||||
}
|
||||
if refreshedAuth != nil {
|
||||
auth = refreshedAuth
|
||||
// Persist the refreshed auth to file so subsequent requests use it
|
||||
if persistErr := e.persistRefreshedAuth(auth); persistErr != nil {
|
||||
log.Warnf("kiro: failed to persist refreshed auth: %v", persistErr)
|
||||
// Continue anyway - the token is valid for this request
|
||||
}
|
||||
accessToken, profileArn = kiroCredentials(auth)
|
||||
kiroPayload, _ = buildKiroPayloadForFormat(body, kiroModelID, profileArn, currentOrigin, isAgentic, isChatOnly, from, opts.Headers)
|
||||
log.Infof("kiro: token refreshed for 403, retrying stream request")
|
||||
@@ -1001,12 +1084,12 @@ func getEffectiveProfileArn(auth *cliproxyauth.Auth, profileArn string) string {
|
||||
// This consolidates the auth_method check that was previously done separately.
|
||||
func getEffectiveProfileArnWithWarning(auth *cliproxyauth.Auth, profileArn string) string {
|
||||
if auth != nil && auth.Metadata != nil {
|
||||
if authMethod, ok := auth.Metadata["auth_method"].(string); ok && authMethod == "builder-id" {
|
||||
// builder-id auth doesn't need profileArn
|
||||
if authMethod, ok := auth.Metadata["auth_method"].(string); ok && (authMethod == "builder-id" || authMethod == "idc") {
|
||||
// builder-id and idc auth don't need profileArn
|
||||
return ""
|
||||
}
|
||||
}
|
||||
// For non-builder-id auth (social auth), profileArn is required
|
||||
// For non-builder-id/idc auth (social auth), profileArn is required
|
||||
if profileArn == "" {
|
||||
log.Warnf("kiro: profile ARN not found in auth, API calls may fail")
|
||||
}
|
||||
@@ -3010,6 +3093,7 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c
|
||||
var refreshToken string
|
||||
var clientID, clientSecret string
|
||||
var authMethod string
|
||||
var region, startURL string
|
||||
|
||||
if auth.Metadata != nil {
|
||||
if rt, ok := auth.Metadata["refresh_token"].(string); ok {
|
||||
@@ -3024,6 +3108,12 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c
|
||||
if am, ok := auth.Metadata["auth_method"].(string); ok {
|
||||
authMethod = am
|
||||
}
|
||||
if r, ok := auth.Metadata["region"].(string); ok {
|
||||
region = r
|
||||
}
|
||||
if su, ok := auth.Metadata["start_url"].(string); ok {
|
||||
startURL = su
|
||||
}
|
||||
}
|
||||
|
||||
if refreshToken == "" {
|
||||
@@ -3033,12 +3123,20 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c
|
||||
var tokenData *kiroauth.KiroTokenData
|
||||
var err error
|
||||
|
||||
// Use SSO OIDC refresh for AWS Builder ID, otherwise use Kiro's OAuth refresh endpoint
|
||||
if clientID != "" && clientSecret != "" && authMethod == "builder-id" {
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(e.cfg)
|
||||
|
||||
// Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint
|
||||
switch {
|
||||
case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "":
|
||||
// IDC refresh with region-specific endpoint
|
||||
log.Debugf("kiro executor: using SSO OIDC refresh for IDC (region=%s)", region)
|
||||
tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL)
|
||||
case clientID != "" && clientSecret != "" && authMethod == "builder-id":
|
||||
// Builder ID refresh with default endpoint
|
||||
log.Debugf("kiro executor: using SSO OIDC refresh for AWS Builder ID")
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(e.cfg)
|
||||
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
|
||||
} else {
|
||||
default:
|
||||
// Fallback to Kiro's OAuth refresh endpoint (for social auth: Google/GitHub)
|
||||
log.Debugf("kiro executor: using Kiro OAuth refresh endpoint")
|
||||
oauth := kiroauth.NewKiroOAuth(e.cfg)
|
||||
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
|
||||
@@ -3094,6 +3192,53 @@ func (e *KiroExecutor) Refresh(ctx context.Context, auth *cliproxyauth.Auth) (*c
|
||||
return updated, nil
|
||||
}
|
||||
|
||||
// persistRefreshedAuth persists a refreshed auth record to disk.
|
||||
// This ensures token refreshes from inline retry are saved to the auth file.
|
||||
func (e *KiroExecutor) persistRefreshedAuth(auth *cliproxyauth.Auth) error {
|
||||
if auth == nil || auth.Metadata == nil {
|
||||
return fmt.Errorf("kiro executor: cannot persist nil auth or metadata")
|
||||
}
|
||||
|
||||
// Determine the file path from auth attributes or filename
|
||||
var authPath string
|
||||
if auth.Attributes != nil {
|
||||
if p := strings.TrimSpace(auth.Attributes["path"]); p != "" {
|
||||
authPath = p
|
||||
}
|
||||
}
|
||||
if authPath == "" {
|
||||
fileName := strings.TrimSpace(auth.FileName)
|
||||
if fileName == "" {
|
||||
return fmt.Errorf("kiro executor: auth has no file path or filename")
|
||||
}
|
||||
if filepath.IsAbs(fileName) {
|
||||
authPath = fileName
|
||||
} else if e.cfg != nil && e.cfg.AuthDir != "" {
|
||||
authPath = filepath.Join(e.cfg.AuthDir, fileName)
|
||||
} else {
|
||||
return fmt.Errorf("kiro executor: cannot determine auth file path")
|
||||
}
|
||||
}
|
||||
|
||||
// Marshal metadata to JSON
|
||||
raw, err := json.Marshal(auth.Metadata)
|
||||
if err != nil {
|
||||
return fmt.Errorf("kiro executor: marshal metadata failed: %w", err)
|
||||
}
|
||||
|
||||
// Write to temp file first, then rename (atomic write)
|
||||
tmp := authPath + ".tmp"
|
||||
if err := os.WriteFile(tmp, raw, 0o600); err != nil {
|
||||
return fmt.Errorf("kiro executor: write temp auth file failed: %w", err)
|
||||
}
|
||||
if err := os.Rename(tmp, authPath); err != nil {
|
||||
return fmt.Errorf("kiro executor: rename auth file failed: %w", err)
|
||||
}
|
||||
|
||||
log.Debugf("kiro executor: persisted refreshed auth to %s", authPath)
|
||||
return nil
|
||||
}
|
||||
|
||||
// isTokenExpired checks if a JWT access token has expired.
|
||||
// Returns true if the token is expired or cannot be parsed.
|
||||
func (e *KiroExecutor) isTokenExpired(accessToken string) bool {
|
||||
|
||||
@@ -19,7 +19,7 @@ type usageReporter struct {
|
||||
provider string
|
||||
model string
|
||||
authID string
|
||||
authIndex uint64
|
||||
authIndex string
|
||||
apiKey string
|
||||
source string
|
||||
requestedAt time.Time
|
||||
@@ -275,6 +275,20 @@ func parseClaudeStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
return detail, true
|
||||
}
|
||||
|
||||
func parseGeminiFamilyUsageDetail(node gjson.Result) usage.Detail {
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
CachedTokens: node.Get("cachedContentTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
}
|
||||
|
||||
func parseGeminiCLIUsage(data []byte) usage.Detail {
|
||||
usageNode := gjson.ParseBytes(data)
|
||||
node := usageNode.Get("response.usageMetadata")
|
||||
@@ -284,16 +298,7 @@ func parseGeminiCLIUsage(data []byte) usage.Detail {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
return parseGeminiFamilyUsageDetail(node)
|
||||
}
|
||||
|
||||
func parseGeminiUsage(data []byte) usage.Detail {
|
||||
@@ -305,16 +310,7 @@ func parseGeminiUsage(data []byte) usage.Detail {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
return parseGeminiFamilyUsageDetail(node)
|
||||
}
|
||||
|
||||
func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
@@ -329,16 +325,7 @@ func parseGeminiStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail, true
|
||||
return parseGeminiFamilyUsageDetail(node), true
|
||||
}
|
||||
|
||||
func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
@@ -353,16 +340,7 @@ func parseGeminiCLIStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail, true
|
||||
return parseGeminiFamilyUsageDetail(node), true
|
||||
}
|
||||
|
||||
func parseAntigravityUsage(data []byte) usage.Detail {
|
||||
@@ -377,16 +355,7 @@ func parseAntigravityUsage(data []byte) usage.Detail {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail
|
||||
return parseGeminiFamilyUsageDetail(node)
|
||||
}
|
||||
|
||||
func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
@@ -404,16 +373,7 @@ func parseAntigravityStreamUsage(line []byte) (usage.Detail, bool) {
|
||||
if !node.Exists() {
|
||||
return usage.Detail{}, false
|
||||
}
|
||||
detail := usage.Detail{
|
||||
InputTokens: node.Get("promptTokenCount").Int(),
|
||||
OutputTokens: node.Get("candidatesTokenCount").Int(),
|
||||
ReasoningTokens: node.Get("thoughtsTokenCount").Int(),
|
||||
TotalTokens: node.Get("totalTokenCount").Int(),
|
||||
}
|
||||
if detail.TotalTokens == 0 {
|
||||
detail.TotalTokens = detail.InputTokens + detail.OutputTokens + detail.ReasoningTokens
|
||||
}
|
||||
return detail, true
|
||||
return parseGeminiFamilyUsageDetail(node), true
|
||||
}
|
||||
|
||||
var stopChunkWithoutUsage sync.Map
|
||||
@@ -522,12 +482,16 @@ func StripUsageMetadataFromJSON(rawJSON []byte) ([]byte, bool) {
|
||||
cleaned := jsonBytes
|
||||
var changed bool
|
||||
|
||||
if gjson.GetBytes(cleaned, "usageMetadata").Exists() {
|
||||
if usageMetadata = gjson.GetBytes(cleaned, "usageMetadata"); usageMetadata.Exists() {
|
||||
// Rename usageMetadata to cpaUsageMetadata in the message_start event of Claude
|
||||
cleaned, _ = sjson.SetRawBytes(cleaned, "cpaUsageMetadata", []byte(usageMetadata.Raw))
|
||||
cleaned, _ = sjson.DeleteBytes(cleaned, "usageMetadata")
|
||||
changed = true
|
||||
}
|
||||
|
||||
if gjson.GetBytes(cleaned, "response.usageMetadata").Exists() {
|
||||
if usageMetadata = gjson.GetBytes(cleaned, "response.usageMetadata"); usageMetadata.Exists() {
|
||||
// Rename usageMetadata to cpaUsageMetadata in the message_start event of Claude
|
||||
cleaned, _ = sjson.SetRawBytes(cleaned, "response.cpaUsageMetadata", []byte(usageMetadata.Raw))
|
||||
cleaned, _ = sjson.DeleteBytes(cleaned, "response.usageMetadata")
|
||||
changed = true
|
||||
}
|
||||
|
||||
@@ -86,6 +86,10 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
hasSystemInstruction = true
|
||||
}
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
systemInstructionJSON = `{"role":"user","parts":[{"text":""}]}`
|
||||
systemInstructionJSON, _ = sjson.Set(systemInstructionJSON, "parts.0.text", systemResult.String())
|
||||
hasSystemInstruction = true
|
||||
}
|
||||
|
||||
// contents
|
||||
@@ -121,32 +125,31 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
// Use GetThinkingText to handle wrapped thinking objects
|
||||
thinkingText := util.GetThinkingText(contentResult)
|
||||
signatureResult := contentResult.Get("signature")
|
||||
clientSignature := ""
|
||||
if signatureResult.Exists() && signatureResult.String() != "" {
|
||||
clientSignature = signatureResult.String()
|
||||
}
|
||||
|
||||
// Always try cached signature first (more reliable than client-provided)
|
||||
// Client may send stale or invalid signatures from different sessions
|
||||
signature := ""
|
||||
if sessionID != "" && thinkingText != "" {
|
||||
if cachedSig := cache.GetCachedSignature(sessionID, thinkingText); cachedSig != "" {
|
||||
signature = cachedSig
|
||||
log.Debugf("Using cached signature for thinking block")
|
||||
clientSignature := ""
|
||||
if signatureResult.Exists() && signatureResult.String() != "" {
|
||||
clientSignature = signatureResult.String()
|
||||
}
|
||||
}
|
||||
|
||||
// Fallback to client signature only if cache miss and client signature is valid
|
||||
if signature == "" && cache.HasValidSignature(clientSignature) {
|
||||
signature = clientSignature
|
||||
log.Debugf("Using client-provided signature for thinking block")
|
||||
}
|
||||
// Always try cached signature first (more reliable than client-provided)
|
||||
// Client may send stale or invalid signatures from different sessions
|
||||
signature := ""
|
||||
if sessionID != "" && thinkingText != "" {
|
||||
if cachedSig := cache.GetCachedSignature(sessionID, thinkingText); cachedSig != "" {
|
||||
signature = cachedSig
|
||||
log.Debugf("Using cached signature for thinking block")
|
||||
}
|
||||
}
|
||||
|
||||
// Store for subsequent tool_use in the same message
|
||||
if cache.HasValidSignature(signature) {
|
||||
currentMessageThinkingSignature = signature
|
||||
}
|
||||
// Fallback to client signature only if cache miss and client signature is valid
|
||||
if signature == "" && cache.HasValidSignature(clientSignature) {
|
||||
signature = clientSignature
|
||||
log.Debugf("Using client-provided signature for thinking block")
|
||||
}
|
||||
|
||||
// Store for subsequent tool_use in the same message
|
||||
if cache.HasValidSignature(signature) {
|
||||
currentMessageThinkingSignature = signature
|
||||
}
|
||||
|
||||
// Skip trailing unsigned thinking blocks on last assistant message
|
||||
isUnsigned := !cache.HasValidSignature(signature)
|
||||
@@ -321,6 +324,7 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
// tools
|
||||
toolsJSON := ""
|
||||
toolDeclCount := 0
|
||||
allowedToolKeys := []string{"name", "description", "behavior", "parameters", "parametersJsonSchema", "response", "responseJsonSchema"}
|
||||
toolsResult := gjson.GetBytes(rawJSON, "tools")
|
||||
if toolsResult.IsArray() {
|
||||
toolsJSON = `[{"functionDeclarations":[]}]`
|
||||
@@ -333,10 +337,12 @@ func ConvertClaudeRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
inputSchema := util.CleanJSONSchemaForAntigravity(inputSchemaResult.Raw)
|
||||
tool, _ := sjson.Delete(toolResult.Raw, "input_schema")
|
||||
tool, _ = sjson.SetRaw(tool, "parametersJsonSchema", inputSchema)
|
||||
tool, _ = sjson.Delete(tool, "strict")
|
||||
tool, _ = sjson.Delete(tool, "input_examples")
|
||||
tool, _ = sjson.Delete(tool, "type")
|
||||
tool, _ = sjson.Delete(tool, "cache_control")
|
||||
for toolKey := range gjson.Parse(tool).Map() {
|
||||
if util.InArray(allowedToolKeys, toolKey) {
|
||||
continue
|
||||
}
|
||||
tool, _ = sjson.Delete(tool, toolKey)
|
||||
}
|
||||
toolsJSON, _ = sjson.SetRaw(toolsJSON, "0.functionDeclarations.-1", tool)
|
||||
toolDeclCount++
|
||||
}
|
||||
|
||||
@@ -35,6 +35,7 @@ type Params struct {
|
||||
CandidatesTokenCount int64 // Cached candidate token count from usage metadata
|
||||
ThoughtsTokenCount int64 // Cached thinking token count from usage metadata
|
||||
TotalTokenCount int64 // Cached total token count from usage metadata
|
||||
CachedTokenCount int64 // Cached content token count (indicates prompt caching)
|
||||
HasSentFinalEvents bool // Indicates if final content/message events have been sent
|
||||
HasToolUse bool // Indicates if tool use was observed in the stream
|
||||
HasContent bool // Tracks whether any content (text, thinking, or tool use) has been output
|
||||
@@ -98,6 +99,14 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
// This follows the Claude Code API specification for streaming message initialization
|
||||
messageStartTemplate := `{"type": "message_start", "message": {"id": "msg_1nZdL29xx5MUA1yADyHTEsnR8uuvGzszyY", "type": "message", "role": "assistant", "content": [], "model": "claude-3-5-sonnet-20241022", "stop_reason": null, "stop_sequence": null, "usage": {"input_tokens": 0, "output_tokens": 0}}}`
|
||||
|
||||
// Use cpaUsageMetadata within the message_start event for Claude.
|
||||
if promptTokenCount := gjson.GetBytes(rawJSON, "response.cpaUsageMetadata.promptTokenCount"); promptTokenCount.Exists() {
|
||||
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.usage.input_tokens", promptTokenCount.Int())
|
||||
}
|
||||
if candidatesTokenCount := gjson.GetBytes(rawJSON, "response.cpaUsageMetadata.candidatesTokenCount"); candidatesTokenCount.Exists() {
|
||||
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.usage.output_tokens", candidatesTokenCount.Int())
|
||||
}
|
||||
|
||||
// Override default values with actual response metadata if available from the Gemini CLI response
|
||||
if modelVersionResult := gjson.GetBytes(rawJSON, "response.modelVersion"); modelVersionResult.Exists() {
|
||||
messageStartTemplate, _ = sjson.Set(messageStartTemplate, "message.model", modelVersionResult.String())
|
||||
@@ -270,7 +279,8 @@ func ConvertAntigravityResponseToClaude(_ context.Context, _ string, originalReq
|
||||
|
||||
if usageResult := gjson.GetBytes(rawJSON, "response.usageMetadata"); usageResult.Exists() {
|
||||
params.HasUsageMetadata = true
|
||||
params.PromptTokenCount = usageResult.Get("promptTokenCount").Int()
|
||||
params.CachedTokenCount = usageResult.Get("cachedContentTokenCount").Int()
|
||||
params.PromptTokenCount = usageResult.Get("promptTokenCount").Int() - params.CachedTokenCount
|
||||
params.CandidatesTokenCount = usageResult.Get("candidatesTokenCount").Int()
|
||||
params.ThoughtsTokenCount = usageResult.Get("thoughtsTokenCount").Int()
|
||||
params.TotalTokenCount = usageResult.Get("totalTokenCount").Int()
|
||||
@@ -322,6 +332,14 @@ func appendFinalEvents(params *Params, output *string, force bool) {
|
||||
*output = *output + "event: message_delta\n"
|
||||
*output = *output + "data: "
|
||||
delta := fmt.Sprintf(`{"type":"message_delta","delta":{"stop_reason":"%s","stop_sequence":null},"usage":{"input_tokens":%d,"output_tokens":%d}}`, stopReason, params.PromptTokenCount, usageOutputTokens)
|
||||
// Add cache_read_input_tokens if cached tokens are present (indicates prompt caching is working)
|
||||
if params.CachedTokenCount > 0 {
|
||||
var err error
|
||||
delta, err = sjson.Set(delta, "usage.cache_read_input_tokens", params.CachedTokenCount)
|
||||
if err != nil {
|
||||
log.Warnf("antigravity claude response: failed to set cache_read_input_tokens: %v", err)
|
||||
}
|
||||
}
|
||||
*output = *output + delta + "\n\n\n"
|
||||
|
||||
params.HasSentFinalEvents = true
|
||||
@@ -361,6 +379,7 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
|
||||
candidateTokens := root.Get("response.usageMetadata.candidatesTokenCount").Int()
|
||||
thoughtTokens := root.Get("response.usageMetadata.thoughtsTokenCount").Int()
|
||||
totalTokens := root.Get("response.usageMetadata.totalTokenCount").Int()
|
||||
cachedTokens := root.Get("response.usageMetadata.cachedContentTokenCount").Int()
|
||||
outputTokens := candidateTokens + thoughtTokens
|
||||
if outputTokens == 0 && totalTokens > 0 {
|
||||
outputTokens = totalTokens - promptTokens
|
||||
@@ -374,6 +393,14 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
|
||||
responseJSON, _ = sjson.Set(responseJSON, "model", root.Get("response.modelVersion").String())
|
||||
responseJSON, _ = sjson.Set(responseJSON, "usage.input_tokens", promptTokens)
|
||||
responseJSON, _ = sjson.Set(responseJSON, "usage.output_tokens", outputTokens)
|
||||
// Add cache_read_input_tokens if cached tokens are present (indicates prompt caching is working)
|
||||
if cachedTokens > 0 {
|
||||
var err error
|
||||
responseJSON, err = sjson.Set(responseJSON, "usage.cache_read_input_tokens", cachedTokens)
|
||||
if err != nil {
|
||||
log.Warnf("antigravity claude response: failed to set cache_read_input_tokens: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
contentArrayInitialized := false
|
||||
ensureContentArray := func() {
|
||||
@@ -452,7 +479,7 @@ func ConvertAntigravityResponseToClaudeNonStream(_ context.Context, _ string, or
|
||||
toolBlock, _ = sjson.Set(toolBlock, "id", fmt.Sprintf("tool_%d", toolIDCounter))
|
||||
toolBlock, _ = sjson.Set(toolBlock, "name", name)
|
||||
|
||||
if args := functionCall.Get("args"); args.Exists() && args.Raw != "" && gjson.Valid(args.Raw) {
|
||||
if args := functionCall.Get("args"); args.Exists() && args.Raw != "" && gjson.Valid(args.Raw) && args.IsObject() {
|
||||
toolBlock, _ = sjson.SetRaw(toolBlock, "input", args.Raw)
|
||||
}
|
||||
|
||||
|
||||
@@ -7,7 +7,6 @@ package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/common"
|
||||
@@ -135,41 +134,31 @@ func ConvertGeminiRequestToAntigravity(_ string, inputRawJSON []byte, _ bool) []
|
||||
|
||||
// FunctionCallGroup represents a group of function calls and their responses
|
||||
type FunctionCallGroup struct {
|
||||
ModelContent map[string]interface{}
|
||||
FunctionCalls []gjson.Result
|
||||
ResponsesNeeded int
|
||||
}
|
||||
|
||||
// parseFunctionResponse attempts to unmarshal a function response part.
|
||||
// Falls back to gjson extraction if standard json.Unmarshal fails.
|
||||
func parseFunctionResponse(response gjson.Result) map[string]interface{} {
|
||||
var responseMap map[string]interface{}
|
||||
err := json.Unmarshal([]byte(response.Raw), &responseMap)
|
||||
if err == nil {
|
||||
return responseMap
|
||||
// parseFunctionResponseRaw attempts to normalize a function response part into a JSON object string.
|
||||
// Falls back to a minimal "functionResponse" object when parsing fails.
|
||||
func parseFunctionResponseRaw(response gjson.Result) string {
|
||||
if response.IsObject() && gjson.Valid(response.Raw) {
|
||||
return response.Raw
|
||||
}
|
||||
|
||||
log.Debugf("unmarshal function response failed, using fallback: %v", err)
|
||||
log.Debugf("parse function response failed, using fallback")
|
||||
funcResp := response.Get("functionResponse")
|
||||
if funcResp.Exists() {
|
||||
fr := map[string]interface{}{
|
||||
"name": funcResp.Get("name").String(),
|
||||
"response": map[string]interface{}{
|
||||
"result": funcResp.Get("response").String(),
|
||||
},
|
||||
}
|
||||
fr := `{"functionResponse":{"name":"","response":{"result":""}}}`
|
||||
fr, _ = sjson.Set(fr, "functionResponse.name", funcResp.Get("name").String())
|
||||
fr, _ = sjson.Set(fr, "functionResponse.response.result", funcResp.Get("response").String())
|
||||
if id := funcResp.Get("id").String(); id != "" {
|
||||
fr["id"] = id
|
||||
fr, _ = sjson.Set(fr, "functionResponse.id", id)
|
||||
}
|
||||
return map[string]interface{}{"functionResponse": fr}
|
||||
return fr
|
||||
}
|
||||
|
||||
return map[string]interface{}{
|
||||
"functionResponse": map[string]interface{}{
|
||||
"name": "unknown",
|
||||
"response": map[string]interface{}{"result": response.String()},
|
||||
},
|
||||
}
|
||||
fr := `{"functionResponse":{"name":"unknown","response":{"result":""}}}`
|
||||
fr, _ = sjson.Set(fr, "functionResponse.response.result", response.String())
|
||||
return fr
|
||||
}
|
||||
|
||||
// fixCLIToolResponse performs sophisticated tool response format conversion and grouping.
|
||||
@@ -196,7 +185,7 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
}
|
||||
|
||||
// Initialize data structures for processing and grouping
|
||||
var newContents []interface{} // Final processed contents array
|
||||
contentsWrapper := `{"contents":[]}`
|
||||
var pendingGroups []*FunctionCallGroup // Groups awaiting completion with responses
|
||||
var collectedResponses []gjson.Result // Standalone responses to be matched
|
||||
|
||||
@@ -228,17 +217,16 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
collectedResponses = collectedResponses[group.ResponsesNeeded:]
|
||||
|
||||
// Create merged function response content
|
||||
var responseParts []interface{}
|
||||
functionResponseContent := `{"parts":[],"role":"function"}`
|
||||
for _, response := range groupResponses {
|
||||
responseParts = append(responseParts, parseFunctionResponse(response))
|
||||
partRaw := parseFunctionResponseRaw(response)
|
||||
if partRaw != "" {
|
||||
functionResponseContent, _ = sjson.SetRaw(functionResponseContent, "parts.-1", partRaw)
|
||||
}
|
||||
}
|
||||
|
||||
if len(responseParts) > 0 {
|
||||
functionResponseContent := map[string]interface{}{
|
||||
"parts": responseParts,
|
||||
"role": "function",
|
||||
}
|
||||
newContents = append(newContents, functionResponseContent)
|
||||
if gjson.Get(functionResponseContent, "parts.#").Int() > 0 {
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", functionResponseContent)
|
||||
}
|
||||
|
||||
// Remove this group as it's been satisfied
|
||||
@@ -252,50 +240,42 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
|
||||
// If this is a model with function calls, create a new group
|
||||
if role == "model" {
|
||||
var functionCallsInThisModel []gjson.Result
|
||||
functionCallsCount := 0
|
||||
parts.ForEach(func(_, part gjson.Result) bool {
|
||||
if part.Get("functionCall").Exists() {
|
||||
functionCallsInThisModel = append(functionCallsInThisModel, part)
|
||||
functionCallsCount++
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if len(functionCallsInThisModel) > 0 {
|
||||
if functionCallsCount > 0 {
|
||||
// Add the model content
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal model content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse model content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
|
||||
// Create a new group for tracking responses
|
||||
group := &FunctionCallGroup{
|
||||
ModelContent: contentMap,
|
||||
FunctionCalls: functionCallsInThisModel,
|
||||
ResponsesNeeded: len(functionCallsInThisModel),
|
||||
ResponsesNeeded: functionCallsCount,
|
||||
}
|
||||
pendingGroups = append(pendingGroups, group)
|
||||
} else {
|
||||
// Regular model content without function calls
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
}
|
||||
} else {
|
||||
// Non-model content (user, etc.)
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
}
|
||||
|
||||
return true
|
||||
@@ -307,25 +287,23 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
groupResponses := collectedResponses[:group.ResponsesNeeded]
|
||||
collectedResponses = collectedResponses[group.ResponsesNeeded:]
|
||||
|
||||
var responseParts []interface{}
|
||||
functionResponseContent := `{"parts":[],"role":"function"}`
|
||||
for _, response := range groupResponses {
|
||||
responseParts = append(responseParts, parseFunctionResponse(response))
|
||||
partRaw := parseFunctionResponseRaw(response)
|
||||
if partRaw != "" {
|
||||
functionResponseContent, _ = sjson.SetRaw(functionResponseContent, "parts.-1", partRaw)
|
||||
}
|
||||
}
|
||||
|
||||
if len(responseParts) > 0 {
|
||||
functionResponseContent := map[string]interface{}{
|
||||
"parts": responseParts,
|
||||
"role": "function",
|
||||
}
|
||||
newContents = append(newContents, functionResponseContent)
|
||||
if gjson.Get(functionResponseContent, "parts.#").Int() > 0 {
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", functionResponseContent)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update the original JSON with the new contents
|
||||
result := input
|
||||
newContentsJSON, _ := json.Marshal(newContents)
|
||||
result, _ = sjson.Set(result, "request.contents", json.RawMessage(newContentsJSON))
|
||||
result, _ = sjson.SetRaw(result, "request.contents", gjson.Get(contentsWrapper, "contents").Raw)
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -192,6 +192,14 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
@@ -239,10 +247,30 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
} else if role == "assistant" {
|
||||
node := []byte(`{"role":"model","parts":[]}`)
|
||||
p := 0
|
||||
if content.Type == gjson.String {
|
||||
if content.Type == gjson.String && content.String() != "" {
|
||||
node, _ = sjson.SetBytes(node, "parts.-1.text", content.String())
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
|
||||
p++
|
||||
} else if content.IsArray() {
|
||||
// Assistant multimodal content (e.g. text + image) -> single model content with parts
|
||||
for _, item := range content.Array() {
|
||||
switch item.Get("type").String() {
|
||||
case "text":
|
||||
p++
|
||||
case "image_url":
|
||||
// If the assistant returned an inline data URL, preserve it for history fidelity.
|
||||
imageURL := item.Get("image_url.url").String()
|
||||
if len(imageURL) > 5 { // expect data:...
|
||||
pieces := strings.SplitN(imageURL[5:], ";", 2)
|
||||
if len(pieces) == 2 && len(pieces[1]) > 7 {
|
||||
mime := pieces[0]
|
||||
data := pieces[1][7:]
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".inlineData.mime_type", mime)
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".inlineData.data", data)
|
||||
p++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Tool calls -> single model content with functionCall parts
|
||||
@@ -258,7 +286,11 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
fargs := tc.Get("function.arguments").String()
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.id", fid)
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.name", fname)
|
||||
node, _ = sjson.SetRawBytes(node, "parts."+itoa(p)+".functionCall.args", []byte(fargs))
|
||||
if gjson.Valid(fargs) {
|
||||
node, _ = sjson.SetRawBytes(node, "parts."+itoa(p)+".functionCall.args", []byte(fargs))
|
||||
} else {
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".functionCall.args.params", []byte(fargs))
|
||||
}
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".thoughtSignature", geminiCLIFunctionThoughtSignature)
|
||||
p++
|
||||
if fid != "" {
|
||||
@@ -293,6 +325,8 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
if pp > 0 {
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", toolNode)
|
||||
}
|
||||
} else {
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -319,7 +353,7 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
@@ -334,7 +368,7 @@ func ConvertOpenAIRequestToAntigravity(modelName string, inputRawJSON []byte, _
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
|
||||
@@ -8,12 +8,13 @@ package chat_completions
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/openai/chat-completions"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
@@ -86,18 +87,27 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq
|
||||
|
||||
// Extract and set usage metadata (token counts).
|
||||
if usageResult := gjson.GetBytes(rawJSON, "response.usageMetadata"); usageResult.Exists() {
|
||||
cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int()
|
||||
if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens", candidatesTokenCountResult.Int())
|
||||
}
|
||||
if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() {
|
||||
template, _ = sjson.Set(template, "usage.total_tokens", totalTokenCountResult.Int())
|
||||
}
|
||||
promptTokenCount := usageResult.Get("promptTokenCount").Int()
|
||||
promptTokenCount := usageResult.Get("promptTokenCount").Int() - cachedTokenCount
|
||||
thoughtsTokenCount := usageResult.Get("thoughtsTokenCount").Int()
|
||||
template, _ = sjson.Set(template, "usage.prompt_tokens", promptTokenCount+thoughtsTokenCount)
|
||||
if thoughtsTokenCount > 0 {
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens_details.reasoning_tokens", thoughtsTokenCount)
|
||||
}
|
||||
// Include cached token count if present (indicates prompt caching is working)
|
||||
if cachedTokenCount > 0 {
|
||||
var err error
|
||||
template, err = sjson.Set(template, "usage.prompt_tokens_details.cached_tokens", cachedTokenCount)
|
||||
if err != nil {
|
||||
log.Warnf("antigravity openai response: failed to set cached_tokens: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process the main content part of the response.
|
||||
@@ -171,21 +181,16 @@ func ConvertAntigravityResponseToOpenAI(_ context.Context, _ string, originalReq
|
||||
mimeType = "image/png"
|
||||
}
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
imagePayload, err := json.Marshal(map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]string{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
imagesResult := gjson.Get(template, "choices.0.delta.images")
|
||||
if !imagesResult.Exists() || !imagesResult.IsArray() {
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images", `[]`)
|
||||
}
|
||||
imageIndex := len(gjson.Get(template, "choices.0.delta.images").Array())
|
||||
imagePayload := `{"type":"image_url","image_url":{"url":""}}`
|
||||
imagePayload, _ = sjson.Set(imagePayload, "index", imageIndex)
|
||||
imagePayload, _ = sjson.Set(imagePayload, "image_url.url", imageURL)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", string(imagePayload))
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", imagePayload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -194,7 +194,7 @@ func ConvertGeminiRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
if name := fc.Get("name"); name.Exists() {
|
||||
toolUse, _ = sjson.Set(toolUse, "name", name.String())
|
||||
}
|
||||
if args := fc.Get("args"); args.Exists() {
|
||||
if args := fc.Get("args"); args.Exists() && args.IsObject() {
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", args.Raw)
|
||||
}
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", toolUse)
|
||||
@@ -314,11 +314,11 @@ func ConvertGeminiRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
if mode := funcCalling.Get("mode"); mode.Exists() {
|
||||
switch mode.String() {
|
||||
case "AUTO":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "auto"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"auto"}`)
|
||||
case "NONE":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "none"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"none"}`)
|
||||
case "ANY":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "any"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"any"}`)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -263,51 +263,6 @@ func ConvertClaudeResponseToGemini(_ context.Context, modelName string, original
|
||||
}
|
||||
}
|
||||
|
||||
// convertArrayToJSON converts []interface{} to JSON array string
|
||||
func convertArrayToJSON(arr []interface{}) string {
|
||||
result := "[]"
|
||||
for _, item := range arr {
|
||||
switch itemData := item.(type) {
|
||||
case map[string]interface{}:
|
||||
itemJSON := convertMapToJSON(itemData)
|
||||
result, _ = sjson.SetRaw(result, "-1", itemJSON)
|
||||
case string:
|
||||
result, _ = sjson.Set(result, "-1", itemData)
|
||||
case bool:
|
||||
result, _ = sjson.Set(result, "-1", itemData)
|
||||
case float64, int, int64:
|
||||
result, _ = sjson.Set(result, "-1", itemData)
|
||||
default:
|
||||
result, _ = sjson.Set(result, "-1", itemData)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// convertMapToJSON converts map[string]interface{} to JSON object string
|
||||
func convertMapToJSON(m map[string]interface{}) string {
|
||||
result := "{}"
|
||||
for key, value := range m {
|
||||
switch val := value.(type) {
|
||||
case map[string]interface{}:
|
||||
nestedJSON := convertMapToJSON(val)
|
||||
result, _ = sjson.SetRaw(result, key, nestedJSON)
|
||||
case []interface{}:
|
||||
arrayJSON := convertArrayToJSON(val)
|
||||
result, _ = sjson.SetRaw(result, key, arrayJSON)
|
||||
case string:
|
||||
result, _ = sjson.Set(result, key, val)
|
||||
case bool:
|
||||
result, _ = sjson.Set(result, key, val)
|
||||
case float64, int, int64:
|
||||
result, _ = sjson.Set(result, key, val)
|
||||
default:
|
||||
result, _ = sjson.Set(result, key, val)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
// ConvertClaudeResponseToGeminiNonStream converts a non-streaming Claude Code response to a non-streaming Gemini response.
|
||||
// This function processes the complete Claude Code response and transforms it into a single Gemini-compatible
|
||||
// JSON response. It handles message content, tool calls, reasoning content, and usage metadata, combining all
|
||||
@@ -356,8 +311,8 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
}
|
||||
|
||||
// Process each streaming event and collect parts
|
||||
var allParts []interface{}
|
||||
var finalUsage map[string]interface{}
|
||||
var allParts []string
|
||||
var finalUsageJSON string
|
||||
var responseID string
|
||||
var createdAt int64
|
||||
|
||||
@@ -407,16 +362,14 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
if text := delta.Get("text"); text.Exists() && text.String() != "" {
|
||||
partJSON := `{"text":""}`
|
||||
partJSON, _ = sjson.Set(partJSON, "text", text.String())
|
||||
part := gjson.Parse(partJSON).Value().(map[string]interface{})
|
||||
allParts = append(allParts, part)
|
||||
allParts = append(allParts, partJSON)
|
||||
}
|
||||
case "thinking_delta":
|
||||
// Process reasoning/thinking content
|
||||
if text := delta.Get("thinking"); text.Exists() && text.String() != "" {
|
||||
partJSON := `{"thought":true,"text":""}`
|
||||
partJSON, _ = sjson.Set(partJSON, "text", text.String())
|
||||
part := gjson.Parse(partJSON).Value().(map[string]interface{})
|
||||
allParts = append(allParts, part)
|
||||
allParts = append(allParts, partJSON)
|
||||
}
|
||||
case "input_json_delta":
|
||||
// accumulate args partial_json for this index
|
||||
@@ -456,9 +409,7 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
if argsTrim != "" {
|
||||
functionCallJSON, _ = sjson.SetRaw(functionCallJSON, "functionCall.args", argsTrim)
|
||||
}
|
||||
// Parse back to interface{} for allParts
|
||||
functionCall := gjson.Parse(functionCallJSON).Value().(map[string]interface{})
|
||||
allParts = append(allParts, functionCall)
|
||||
allParts = append(allParts, functionCallJSON)
|
||||
// cleanup used state for this index
|
||||
if newParam.ToolUseArgs != nil {
|
||||
delete(newParam.ToolUseArgs, idx)
|
||||
@@ -501,8 +452,7 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
// Set traffic type (required by Gemini API)
|
||||
usageJSON, _ = sjson.Set(usageJSON, "trafficType", "PROVISIONED_THROUGHPUT")
|
||||
|
||||
// Convert to map[string]interface{} using gjson
|
||||
finalUsage = gjson.Parse(usageJSON).Value().(map[string]interface{})
|
||||
finalUsageJSON = usageJSON
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -520,12 +470,16 @@ func ConvertClaudeResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
|
||||
// Set the consolidated parts array
|
||||
if len(consolidatedParts) > 0 {
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", convertToJSONString(consolidatedParts))
|
||||
partsJSON := "[]"
|
||||
for _, partJSON := range consolidatedParts {
|
||||
partsJSON, _ = sjson.SetRaw(partsJSON, "-1", partJSON)
|
||||
}
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", partsJSON)
|
||||
}
|
||||
|
||||
// Set usage metadata
|
||||
if finalUsage != nil {
|
||||
template, _ = sjson.SetRaw(template, "usageMetadata", convertToJSONString(finalUsage))
|
||||
if finalUsageJSON != "" {
|
||||
template, _ = sjson.SetRaw(template, "usageMetadata", finalUsageJSON)
|
||||
}
|
||||
|
||||
return template
|
||||
@@ -539,12 +493,12 @@ func GeminiTokenCount(ctx context.Context, count int64) string {
|
||||
// This function processes the parts array to combine adjacent text elements and thinking elements
|
||||
// into single consolidated parts, which results in a more readable and efficient response structure.
|
||||
// Tool calls and other non-text parts are preserved as separate elements.
|
||||
func consolidateParts(parts []interface{}) []interface{} {
|
||||
func consolidateParts(parts []string) []string {
|
||||
if len(parts) == 0 {
|
||||
return parts
|
||||
}
|
||||
|
||||
var consolidated []interface{}
|
||||
var consolidated []string
|
||||
var currentTextPart strings.Builder
|
||||
var currentThoughtPart strings.Builder
|
||||
var hasText, hasThought bool
|
||||
@@ -554,8 +508,7 @@ func consolidateParts(parts []interface{}) []interface{} {
|
||||
if hasText && currentTextPart.Len() > 0 {
|
||||
textPartJSON := `{"text":""}`
|
||||
textPartJSON, _ = sjson.Set(textPartJSON, "text", currentTextPart.String())
|
||||
textPart := gjson.Parse(textPartJSON).Value().(map[string]interface{})
|
||||
consolidated = append(consolidated, textPart)
|
||||
consolidated = append(consolidated, textPartJSON)
|
||||
currentTextPart.Reset()
|
||||
hasText = false
|
||||
}
|
||||
@@ -566,42 +519,42 @@ func consolidateParts(parts []interface{}) []interface{} {
|
||||
if hasThought && currentThoughtPart.Len() > 0 {
|
||||
thoughtPartJSON := `{"thought":true,"text":""}`
|
||||
thoughtPartJSON, _ = sjson.Set(thoughtPartJSON, "text", currentThoughtPart.String())
|
||||
thoughtPart := gjson.Parse(thoughtPartJSON).Value().(map[string]interface{})
|
||||
consolidated = append(consolidated, thoughtPart)
|
||||
consolidated = append(consolidated, thoughtPartJSON)
|
||||
currentThoughtPart.Reset()
|
||||
hasThought = false
|
||||
}
|
||||
}
|
||||
|
||||
for _, part := range parts {
|
||||
partMap, ok := part.(map[string]interface{})
|
||||
if !ok {
|
||||
for _, partJSON := range parts {
|
||||
part := gjson.Parse(partJSON)
|
||||
if !part.Exists() || !part.IsObject() {
|
||||
// Flush any pending parts and add this non-text part
|
||||
flushText()
|
||||
flushThought()
|
||||
consolidated = append(consolidated, part)
|
||||
consolidated = append(consolidated, partJSON)
|
||||
continue
|
||||
}
|
||||
|
||||
if thought, isThought := partMap["thought"]; isThought && thought == true {
|
||||
thought := part.Get("thought")
|
||||
if thought.Exists() && thought.Type == gjson.True {
|
||||
// This is a thinking part - flush any pending text first
|
||||
flushText() // Flush any pending text first
|
||||
|
||||
if text, hasTextContent := partMap["text"].(string); hasTextContent {
|
||||
currentThoughtPart.WriteString(text)
|
||||
if text := part.Get("text"); text.Exists() && text.Type == gjson.String {
|
||||
currentThoughtPart.WriteString(text.String())
|
||||
hasThought = true
|
||||
}
|
||||
} else if text, hasTextContent := partMap["text"].(string); hasTextContent {
|
||||
} else if text := part.Get("text"); text.Exists() && text.Type == gjson.String {
|
||||
// This is a regular text part - flush any pending thought first
|
||||
flushThought() // Flush any pending thought first
|
||||
|
||||
currentTextPart.WriteString(text)
|
||||
currentTextPart.WriteString(text.String())
|
||||
hasText = true
|
||||
} else {
|
||||
// This is some other type of part (like function call) - flush both text and thought
|
||||
flushText()
|
||||
flushThought()
|
||||
consolidated = append(consolidated, part)
|
||||
consolidated = append(consolidated, partJSON)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -611,20 +564,3 @@ func consolidateParts(parts []interface{}) []interface{} {
|
||||
|
||||
return consolidated
|
||||
}
|
||||
|
||||
// convertToJSONString converts interface{} to JSON string using sjson/gjson.
|
||||
// This function provides a consistent way to serialize different data types to JSON strings
|
||||
// for inclusion in the Gemini API response structure.
|
||||
func convertToJSONString(v interface{}) string {
|
||||
switch val := v.(type) {
|
||||
case []interface{}:
|
||||
return convertArrayToJSON(val)
|
||||
case map[string]interface{}:
|
||||
return convertMapToJSON(val)
|
||||
default:
|
||||
// For simple types, create a temporary JSON and extract the value
|
||||
temp := `{"temp":null}`
|
||||
temp, _ = sjson.Set(temp, "temp", val)
|
||||
return gjson.Get(temp, "temp").Raw
|
||||
}
|
||||
}
|
||||
|
||||
@@ -10,7 +10,6 @@ import (
|
||||
"crypto/rand"
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"strings"
|
||||
@@ -137,9 +136,6 @@ func ConvertOpenAIRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
out, _ = sjson.Set(out, "stream", stream)
|
||||
|
||||
// Process messages and transform them to Claude Code format
|
||||
var anthropicMessages []interface{}
|
||||
var toolCallIDs []string // Track tool call IDs for matching with tool results
|
||||
|
||||
if messages := root.Get("messages"); messages.Exists() && messages.IsArray() {
|
||||
messages.ForEach(func(_, message gjson.Result) bool {
|
||||
role := message.Get("role").String()
|
||||
@@ -152,33 +148,23 @@ func ConvertOpenAIRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
role = "user"
|
||||
}
|
||||
|
||||
msg := map[string]interface{}{
|
||||
"role": role,
|
||||
"content": []interface{}{},
|
||||
}
|
||||
msg := `{"role":"","content":[]}`
|
||||
msg, _ = sjson.Set(msg, "role", role)
|
||||
|
||||
// Handle content based on its type (string or array)
|
||||
if contentResult.Exists() && contentResult.Type == gjson.String && contentResult.String() != "" {
|
||||
// Simple text content conversion
|
||||
msg["content"] = []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": contentResult.String(),
|
||||
},
|
||||
}
|
||||
part := `{"type":"text","text":""}`
|
||||
part, _ = sjson.Set(part, "text", contentResult.String())
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", part)
|
||||
} else if contentResult.Exists() && contentResult.IsArray() {
|
||||
// Array of content parts processing
|
||||
var contentParts []interface{}
|
||||
contentResult.ForEach(func(_, part gjson.Result) bool {
|
||||
partType := part.Get("type").String()
|
||||
|
||||
switch partType {
|
||||
case "text":
|
||||
// Text part conversion
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": part.Get("text").String(),
|
||||
})
|
||||
textPart := `{"type":"text","text":""}`
|
||||
textPart, _ = sjson.Set(textPart, "text", part.Get("text").String())
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", textPart)
|
||||
|
||||
case "image_url":
|
||||
// Convert OpenAI image format to Claude Code format
|
||||
@@ -191,132 +177,95 @@ func ConvertOpenAIRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
mediaType := strings.TrimPrefix(mediaTypePart, "data:")
|
||||
data := parts[1]
|
||||
|
||||
contentParts = append(contentParts, map[string]interface{}{
|
||||
"type": "image",
|
||||
"source": map[string]interface{}{
|
||||
"type": "base64",
|
||||
"media_type": mediaType,
|
||||
"data": data,
|
||||
},
|
||||
})
|
||||
imagePart := `{"type":"image","source":{"type":"base64","media_type":"","data":""}}`
|
||||
imagePart, _ = sjson.Set(imagePart, "source.media_type", mediaType)
|
||||
imagePart, _ = sjson.Set(imagePart, "source.data", data)
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", imagePart)
|
||||
}
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if len(contentParts) > 0 {
|
||||
msg["content"] = contentParts
|
||||
}
|
||||
} else {
|
||||
// Initialize empty content array for tool calls
|
||||
msg["content"] = []interface{}{}
|
||||
}
|
||||
|
||||
// Handle tool calls (for assistant messages)
|
||||
if toolCalls := message.Get("tool_calls"); toolCalls.Exists() && toolCalls.IsArray() && role == "assistant" {
|
||||
var contentParts []interface{}
|
||||
|
||||
// Add existing text content if any
|
||||
if existingContent, ok := msg["content"].([]interface{}); ok {
|
||||
contentParts = existingContent
|
||||
}
|
||||
|
||||
toolCalls.ForEach(func(_, toolCall gjson.Result) bool {
|
||||
if toolCall.Get("type").String() == "function" {
|
||||
toolCallID := toolCall.Get("id").String()
|
||||
if toolCallID == "" {
|
||||
toolCallID = genToolCallID()
|
||||
}
|
||||
toolCallIDs = append(toolCallIDs, toolCallID)
|
||||
|
||||
function := toolCall.Get("function")
|
||||
toolUse := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": toolCallID,
|
||||
"name": function.Get("name").String(),
|
||||
}
|
||||
toolUse := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolUse, _ = sjson.Set(toolUse, "id", toolCallID)
|
||||
toolUse, _ = sjson.Set(toolUse, "name", function.Get("name").String())
|
||||
|
||||
// Parse arguments for the tool call
|
||||
if args := function.Get("arguments"); args.Exists() {
|
||||
argsStr := args.String()
|
||||
if argsStr != "" {
|
||||
var argsMap map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &argsMap); err == nil {
|
||||
toolUse["input"] = argsMap
|
||||
if argsStr != "" && gjson.Valid(argsStr) {
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", argsJSON.Raw)
|
||||
} else {
|
||||
toolUse["input"] = map[string]interface{}{}
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", "{}")
|
||||
}
|
||||
} else {
|
||||
toolUse["input"] = map[string]interface{}{}
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", "{}")
|
||||
}
|
||||
} else {
|
||||
toolUse["input"] = map[string]interface{}{}
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", "{}")
|
||||
}
|
||||
|
||||
contentParts = append(contentParts, toolUse)
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", toolUse)
|
||||
}
|
||||
return true
|
||||
})
|
||||
msg["content"] = contentParts
|
||||
}
|
||||
|
||||
anthropicMessages = append(anthropicMessages, msg)
|
||||
out, _ = sjson.SetRaw(out, "messages.-1", msg)
|
||||
|
||||
case "tool":
|
||||
// Handle tool result messages conversion
|
||||
toolCallID := message.Get("tool_call_id").String()
|
||||
content := message.Get("content").String()
|
||||
|
||||
// Create tool result message in Claude Code format
|
||||
msg := map[string]interface{}{
|
||||
"role": "user",
|
||||
"content": []interface{}{
|
||||
map[string]interface{}{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": toolCallID,
|
||||
"content": content,
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
anthropicMessages = append(anthropicMessages, msg)
|
||||
msg := `{"role":"user","content":[{"type":"tool_result","tool_use_id":"","content":""}]}`
|
||||
msg, _ = sjson.Set(msg, "content.0.tool_use_id", toolCallID)
|
||||
msg, _ = sjson.Set(msg, "content.0.content", content)
|
||||
out, _ = sjson.SetRaw(out, "messages.-1", msg)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Set messages in the output template
|
||||
if len(anthropicMessages) > 0 {
|
||||
messagesJSON, _ := json.Marshal(anthropicMessages)
|
||||
out, _ = sjson.SetRaw(out, "messages", string(messagesJSON))
|
||||
}
|
||||
|
||||
// Tools mapping: OpenAI tools -> Claude Code tools
|
||||
if tools := root.Get("tools"); tools.Exists() && tools.IsArray() && len(tools.Array()) > 0 {
|
||||
var anthropicTools []interface{}
|
||||
hasAnthropicTools := false
|
||||
tools.ForEach(func(_, tool gjson.Result) bool {
|
||||
if tool.Get("type").String() == "function" {
|
||||
function := tool.Get("function")
|
||||
anthropicTool := map[string]interface{}{
|
||||
"name": function.Get("name").String(),
|
||||
"description": function.Get("description").String(),
|
||||
}
|
||||
anthropicTool := `{"name":"","description":""}`
|
||||
anthropicTool, _ = sjson.Set(anthropicTool, "name", function.Get("name").String())
|
||||
anthropicTool, _ = sjson.Set(anthropicTool, "description", function.Get("description").String())
|
||||
|
||||
// Convert parameters schema for the tool
|
||||
if parameters := function.Get("parameters"); parameters.Exists() {
|
||||
anthropicTool["input_schema"] = parameters.Value()
|
||||
} else if parameters = function.Get("parametersJsonSchema"); parameters.Exists() {
|
||||
anthropicTool["input_schema"] = parameters.Value()
|
||||
anthropicTool, _ = sjson.SetRaw(anthropicTool, "input_schema", parameters.Raw)
|
||||
} else if parameters := function.Get("parametersJsonSchema"); parameters.Exists() {
|
||||
anthropicTool, _ = sjson.SetRaw(anthropicTool, "input_schema", parameters.Raw)
|
||||
}
|
||||
|
||||
anthropicTools = append(anthropicTools, anthropicTool)
|
||||
out, _ = sjson.SetRaw(out, "tools.-1", anthropicTool)
|
||||
hasAnthropicTools = true
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if len(anthropicTools) > 0 {
|
||||
toolsJSON, _ := json.Marshal(anthropicTools)
|
||||
out, _ = sjson.SetRaw(out, "tools", string(toolsJSON))
|
||||
if !hasAnthropicTools {
|
||||
out, _ = sjson.Delete(out, "tools")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -329,18 +278,17 @@ func ConvertOpenAIRequestToClaude(modelName string, inputRawJSON []byte, stream
|
||||
case "none":
|
||||
// Don't set tool_choice, Claude Code will not use tools
|
||||
case "auto":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "auto"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"auto"}`)
|
||||
case "required":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "any"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"any"}`)
|
||||
}
|
||||
case gjson.JSON:
|
||||
// Specific tool choice mapping
|
||||
if toolChoice.Get("type").String() == "function" {
|
||||
functionName := toolChoice.Get("function.name").String()
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{
|
||||
"type": "tool",
|
||||
"name": functionName,
|
||||
})
|
||||
toolChoiceJSON := `{"type":"tool","name":""}`
|
||||
toolChoiceJSON, _ = sjson.Set(toolChoiceJSON, "name", functionName)
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", toolChoiceJSON)
|
||||
}
|
||||
default:
|
||||
}
|
||||
|
||||
@@ -8,7 +8,7 @@ package chat_completions
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
@@ -182,18 +182,11 @@ func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, original
|
||||
if arguments == "" {
|
||||
arguments = "{}"
|
||||
}
|
||||
|
||||
toolCall := map[string]interface{}{
|
||||
"index": index,
|
||||
"id": accumulator.ID,
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": accumulator.Name,
|
||||
"arguments": arguments,
|
||||
},
|
||||
}
|
||||
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls", []interface{}{toolCall})
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls.0.index", index)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls.0.id", accumulator.ID)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls.0.type", "function")
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls.0.function.name", accumulator.Name)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.tool_calls.0.function.arguments", arguments)
|
||||
|
||||
// Clean up the accumulator for this index
|
||||
delete((*param).(*ConvertAnthropicResponseToOpenAIParams).ToolCallsAccumulator, index)
|
||||
@@ -214,12 +207,14 @@ func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, original
|
||||
|
||||
// Handle usage information for token counts
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
usageObj := map[string]interface{}{
|
||||
"prompt_tokens": usage.Get("input_tokens").Int(),
|
||||
"completion_tokens": usage.Get("output_tokens").Int(),
|
||||
"total_tokens": usage.Get("input_tokens").Int() + usage.Get("output_tokens").Int(),
|
||||
}
|
||||
template, _ = sjson.Set(template, "usage", usageObj)
|
||||
inputTokens := usage.Get("input_tokens").Int()
|
||||
outputTokens := usage.Get("output_tokens").Int()
|
||||
cacheReadInputTokens := usage.Get("cache_read_input_tokens").Int()
|
||||
cacheCreationInputTokens := usage.Get("cache_creation_input_tokens").Int()
|
||||
template, _ = sjson.Set(template, "usage.prompt_tokens", inputTokens+cacheCreationInputTokens)
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens", outputTokens)
|
||||
template, _ = sjson.Set(template, "usage.total_tokens", inputTokens+outputTokens)
|
||||
template, _ = sjson.Set(template, "usage.prompt_tokens_details.cached_tokens", cacheReadInputTokens)
|
||||
}
|
||||
return []string{template}
|
||||
|
||||
@@ -234,14 +229,10 @@ func ConvertClaudeResponseToOpenAI(_ context.Context, modelName string, original
|
||||
case "error":
|
||||
// Error event - format and return error response
|
||||
if errorData := root.Get("error"); errorData.Exists() {
|
||||
errorResponse := map[string]interface{}{
|
||||
"error": map[string]interface{}{
|
||||
"message": errorData.Get("message").String(),
|
||||
"type": errorData.Get("type").String(),
|
||||
},
|
||||
}
|
||||
errorJSON, _ := json.Marshal(errorResponse)
|
||||
return []string{string(errorJSON)}
|
||||
errorJSON := `{"error":{"message":"","type":""}}`
|
||||
errorJSON, _ = sjson.Set(errorJSON, "error.message", errorData.Get("message").String())
|
||||
errorJSON, _ = sjson.Set(errorJSON, "error.type", errorData.Get("type").String())
|
||||
return []string{errorJSON}
|
||||
}
|
||||
return []string{}
|
||||
|
||||
@@ -297,15 +288,10 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
var messageID string
|
||||
var model string
|
||||
var createdAt int64
|
||||
var inputTokens, outputTokens int64
|
||||
var reasoningTokens int64
|
||||
var stopReason string
|
||||
var contentParts []string
|
||||
var reasoningParts []string
|
||||
// Use map to track tool calls by index for proper merging
|
||||
toolCallsMap := make(map[int]map[string]interface{})
|
||||
// Track tool call arguments accumulation
|
||||
toolCallArgsMap := make(map[int]strings.Builder)
|
||||
toolCallsAccumulator := make(map[int]*ToolCallAccumulator)
|
||||
|
||||
for _, chunk := range chunks {
|
||||
root := gjson.ParseBytes(chunk)
|
||||
@@ -318,9 +304,6 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
messageID = message.Get("id").String()
|
||||
model = message.Get("model").String()
|
||||
createdAt = time.Now().Unix()
|
||||
if usage := message.Get("usage"); usage.Exists() {
|
||||
inputTokens = usage.Get("input_tokens").Int()
|
||||
}
|
||||
}
|
||||
|
||||
case "content_block_start":
|
||||
@@ -331,18 +314,12 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
// Start of thinking/reasoning content - skip for now as it's handled in delta
|
||||
continue
|
||||
} else if blockType == "tool_use" {
|
||||
// Initialize tool call tracking for this index
|
||||
// Initialize tool call accumulator for this index
|
||||
index := int(root.Get("index").Int())
|
||||
toolCallsMap[index] = map[string]interface{}{
|
||||
"id": contentBlock.Get("id").String(),
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": contentBlock.Get("name").String(),
|
||||
"arguments": "",
|
||||
},
|
||||
toolCallsAccumulator[index] = &ToolCallAccumulator{
|
||||
ID: contentBlock.Get("id").String(),
|
||||
Name: contentBlock.Get("name").String(),
|
||||
}
|
||||
// Initialize arguments builder for this tool call
|
||||
toolCallArgsMap[index] = strings.Builder{}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -365,9 +342,8 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
// Accumulate tool call arguments
|
||||
if partialJSON := delta.Get("partial_json"); partialJSON.Exists() {
|
||||
index := int(root.Get("index").Int())
|
||||
if builder, exists := toolCallArgsMap[index]; exists {
|
||||
builder.WriteString(partialJSON.String())
|
||||
toolCallArgsMap[index] = builder
|
||||
if accumulator, exists := toolCallsAccumulator[index]; exists {
|
||||
accumulator.Arguments.WriteString(partialJSON.String())
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -376,14 +352,9 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
case "content_block_stop":
|
||||
// Finalize tool call arguments for this index when content block ends
|
||||
index := int(root.Get("index").Int())
|
||||
if toolCall, exists := toolCallsMap[index]; exists {
|
||||
if builder, argsExists := toolCallArgsMap[index]; argsExists {
|
||||
// Set the accumulated arguments for the tool call
|
||||
arguments := builder.String()
|
||||
if arguments == "" {
|
||||
arguments = "{}"
|
||||
}
|
||||
toolCall["function"].(map[string]interface{})["arguments"] = arguments
|
||||
if accumulator, exists := toolCallsAccumulator[index]; exists {
|
||||
if accumulator.Arguments.Len() == 0 {
|
||||
accumulator.Arguments.WriteString("{}")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -395,11 +366,14 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
}
|
||||
}
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
outputTokens = usage.Get("output_tokens").Int()
|
||||
// Estimate reasoning tokens from accumulated thinking content
|
||||
if len(reasoningParts) > 0 {
|
||||
reasoningTokens = int64(len(strings.Join(reasoningParts, "")) / 4) // Rough estimation
|
||||
}
|
||||
inputTokens := usage.Get("input_tokens").Int()
|
||||
outputTokens := usage.Get("output_tokens").Int()
|
||||
cacheReadInputTokens := usage.Get("cache_read_input_tokens").Int()
|
||||
cacheCreationInputTokens := usage.Get("cache_creation_input_tokens").Int()
|
||||
out, _ = sjson.Set(out, "usage.prompt_tokens", inputTokens+cacheCreationInputTokens)
|
||||
out, _ = sjson.Set(out, "usage.completion_tokens", outputTokens)
|
||||
out, _ = sjson.Set(out, "usage.total_tokens", inputTokens+outputTokens)
|
||||
out, _ = sjson.Set(out, "usage.prompt_tokens_details.cached_tokens", cacheReadInputTokens)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -421,24 +395,35 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
}
|
||||
|
||||
// Set tool calls if any were accumulated during processing
|
||||
if len(toolCallsMap) > 0 {
|
||||
// Convert tool calls map to array, preserving order by index
|
||||
var toolCallsArray []interface{}
|
||||
// Find the maximum index to determine the range
|
||||
if len(toolCallsAccumulator) > 0 {
|
||||
toolCallsCount := 0
|
||||
maxIndex := -1
|
||||
for index := range toolCallsMap {
|
||||
for index := range toolCallsAccumulator {
|
||||
if index > maxIndex {
|
||||
maxIndex = index
|
||||
}
|
||||
}
|
||||
// Iterate through all possible indices up to maxIndex
|
||||
|
||||
for i := 0; i <= maxIndex; i++ {
|
||||
if toolCall, exists := toolCallsMap[i]; exists {
|
||||
toolCallsArray = append(toolCallsArray, toolCall)
|
||||
accumulator, exists := toolCallsAccumulator[i]
|
||||
if !exists {
|
||||
continue
|
||||
}
|
||||
|
||||
arguments := accumulator.Arguments.String()
|
||||
|
||||
idPath := fmt.Sprintf("choices.0.message.tool_calls.%d.id", toolCallsCount)
|
||||
typePath := fmt.Sprintf("choices.0.message.tool_calls.%d.type", toolCallsCount)
|
||||
namePath := fmt.Sprintf("choices.0.message.tool_calls.%d.function.name", toolCallsCount)
|
||||
argumentsPath := fmt.Sprintf("choices.0.message.tool_calls.%d.function.arguments", toolCallsCount)
|
||||
|
||||
out, _ = sjson.Set(out, idPath, accumulator.ID)
|
||||
out, _ = sjson.Set(out, typePath, "function")
|
||||
out, _ = sjson.Set(out, namePath, accumulator.Name)
|
||||
out, _ = sjson.Set(out, argumentsPath, arguments)
|
||||
toolCallsCount++
|
||||
}
|
||||
if len(toolCallsArray) > 0 {
|
||||
out, _ = sjson.Set(out, "choices.0.message.tool_calls", toolCallsArray)
|
||||
if toolCallsCount > 0 {
|
||||
out, _ = sjson.Set(out, "choices.0.finish_reason", "tool_calls")
|
||||
} else {
|
||||
out, _ = sjson.Set(out, "choices.0.finish_reason", mapAnthropicStopReasonToOpenAI(stopReason))
|
||||
@@ -447,16 +432,5 @@ func ConvertClaudeResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
out, _ = sjson.Set(out, "choices.0.finish_reason", mapAnthropicStopReasonToOpenAI(stopReason))
|
||||
}
|
||||
|
||||
// Set usage information including prompt tokens, completion tokens, and total tokens
|
||||
totalTokens := inputTokens + outputTokens
|
||||
out, _ = sjson.Set(out, "usage.prompt_tokens", inputTokens)
|
||||
out, _ = sjson.Set(out, "usage.completion_tokens", outputTokens)
|
||||
out, _ = sjson.Set(out, "usage.total_tokens", totalTokens)
|
||||
|
||||
// Add reasoning tokens to usage details if any reasoning content was processed
|
||||
if reasoningTokens > 0 {
|
||||
out, _ = sjson.Set(out, "usage.completion_tokens_details.reasoning_tokens", reasoningTokens)
|
||||
}
|
||||
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -114,13 +114,16 @@ func ConvertOpenAIResponsesRequestToClaude(modelName string, inputRawJSON []byte
|
||||
var builder strings.Builder
|
||||
if parts := item.Get("content"); parts.Exists() && parts.IsArray() {
|
||||
parts.ForEach(func(_, part gjson.Result) bool {
|
||||
text := part.Get("text").String()
|
||||
textResult := part.Get("text")
|
||||
text := textResult.String()
|
||||
if builder.Len() > 0 && text != "" {
|
||||
builder.WriteByte('\n')
|
||||
}
|
||||
builder.WriteString(text)
|
||||
return true
|
||||
})
|
||||
} else if parts.Type == gjson.String {
|
||||
builder.WriteString(parts.String())
|
||||
}
|
||||
instructionsText = builder.String()
|
||||
if instructionsText != "" {
|
||||
@@ -207,6 +210,8 @@ func ConvertOpenAIResponsesRequestToClaude(modelName string, inputRawJSON []byte
|
||||
}
|
||||
return true
|
||||
})
|
||||
} else if parts.Type == gjson.String {
|
||||
textAggregate.WriteString(parts.String())
|
||||
}
|
||||
|
||||
// Fallback to given role if content types not decisive
|
||||
@@ -254,7 +259,10 @@ func ConvertOpenAIResponsesRequestToClaude(modelName string, inputRawJSON []byte
|
||||
toolUse, _ = sjson.Set(toolUse, "id", callID)
|
||||
toolUse, _ = sjson.Set(toolUse, "name", name)
|
||||
if argsStr != "" && gjson.Valid(argsStr) {
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", argsStr)
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", argsJSON.Raw)
|
||||
}
|
||||
}
|
||||
|
||||
asst := `{"role":"assistant","content":[]}`
|
||||
@@ -309,16 +317,18 @@ func ConvertOpenAIResponsesRequestToClaude(modelName string, inputRawJSON []byte
|
||||
case gjson.String:
|
||||
switch toolChoice.String() {
|
||||
case "auto":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "auto"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"auto"}`)
|
||||
case "none":
|
||||
// Leave unset; implies no tools
|
||||
case "required":
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "any"})
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", `{"type":"any"}`)
|
||||
}
|
||||
case gjson.JSON:
|
||||
if toolChoice.Get("type").String() == "function" {
|
||||
fn := toolChoice.Get("function.name").String()
|
||||
out, _ = sjson.Set(out, "tool_choice", map[string]interface{}{"type": "tool", "name": fn})
|
||||
toolChoiceJSON := `{"name":"","type":"tool"}`
|
||||
toolChoiceJSON, _ = sjson.Set(toolChoiceJSON, "name", fn)
|
||||
out, _ = sjson.SetRaw(out, "tool_choice", toolChoiceJSON)
|
||||
}
|
||||
default:
|
||||
|
||||
|
||||
@@ -344,31 +344,20 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
|
||||
}
|
||||
|
||||
// Build response.output from aggregated state
|
||||
var outputs []interface{}
|
||||
outputsWrapper := `{"arr":[]}`
|
||||
// reasoning item (if any)
|
||||
if st.ReasoningBuf.Len() > 0 || st.ReasoningPartAdded {
|
||||
r := map[string]interface{}{
|
||||
"id": st.ReasoningItemID,
|
||||
"type": "reasoning",
|
||||
"summary": []interface{}{map[string]interface{}{"type": "summary_text", "text": st.ReasoningBuf.String()}},
|
||||
}
|
||||
outputs = append(outputs, r)
|
||||
item := `{"id":"","type":"reasoning","summary":[{"type":"summary_text","text":""}]}`
|
||||
item, _ = sjson.Set(item, "id", st.ReasoningItemID)
|
||||
item, _ = sjson.Set(item, "summary.0.text", st.ReasoningBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
// assistant message item (if any text)
|
||||
if st.TextBuf.Len() > 0 || st.InTextBlock || st.CurrentMsgID != "" {
|
||||
m := map[string]interface{}{
|
||||
"id": st.CurrentMsgID,
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": st.TextBuf.String(),
|
||||
}},
|
||||
"role": "assistant",
|
||||
}
|
||||
outputs = append(outputs, m)
|
||||
item := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
item, _ = sjson.Set(item, "id", st.CurrentMsgID)
|
||||
item, _ = sjson.Set(item, "content.0.text", st.TextBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
// function_call items (in ascending index order for determinism)
|
||||
if len(st.FuncArgsBuf) > 0 {
|
||||
@@ -395,19 +384,16 @@ func ConvertClaudeResponseToOpenAIResponses(ctx context.Context, modelName strin
|
||||
if callID == "" && st.CurrentFCID != "" {
|
||||
callID = st.CurrentFCID
|
||||
}
|
||||
item := map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", callID),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": args,
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
}
|
||||
outputs = append(outputs, item)
|
||||
item := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("fc_%s", callID))
|
||||
item, _ = sjson.Set(item, "arguments", args)
|
||||
item, _ = sjson.Set(item, "call_id", callID)
|
||||
item, _ = sjson.Set(item, "name", name)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
}
|
||||
if len(outputs) > 0 {
|
||||
completed, _ = sjson.Set(completed, "response.output", outputs)
|
||||
if gjson.Get(outputsWrapper, "arr.#").Int() > 0 {
|
||||
completed, _ = sjson.SetRaw(completed, "response.output", gjson.Get(outputsWrapper, "arr").Raw)
|
||||
}
|
||||
|
||||
reasoningTokens := int64(0)
|
||||
@@ -628,27 +614,18 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(_ context.Context, _ string
|
||||
}
|
||||
|
||||
// Build output array
|
||||
var outputs []interface{}
|
||||
outputsWrapper := `{"arr":[]}`
|
||||
if reasoningBuf.Len() > 0 {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": reasoningItemID,
|
||||
"type": "reasoning",
|
||||
"summary": []interface{}{map[string]interface{}{"type": "summary_text", "text": reasoningBuf.String()}},
|
||||
})
|
||||
item := `{"id":"","type":"reasoning","summary":[{"type":"summary_text","text":""}]}`
|
||||
item, _ = sjson.Set(item, "id", reasoningItemID)
|
||||
item, _ = sjson.Set(item, "summary.0.text", reasoningBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
if currentMsgID != "" || textBuf.Len() > 0 {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": currentMsgID,
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": textBuf.String(),
|
||||
}},
|
||||
"role": "assistant",
|
||||
})
|
||||
item := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
item, _ = sjson.Set(item, "id", currentMsgID)
|
||||
item, _ = sjson.Set(item, "content.0.text", textBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
// Preserve index order
|
||||
@@ -669,18 +646,16 @@ func ConvertClaudeResponseToOpenAIResponsesNonStream(_ context.Context, _ string
|
||||
if args == "" {
|
||||
args = "{}"
|
||||
}
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", st.id),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": args,
|
||||
"call_id": st.id,
|
||||
"name": st.name,
|
||||
})
|
||||
item := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("fc_%s", st.id))
|
||||
item, _ = sjson.Set(item, "arguments", args)
|
||||
item, _ = sjson.Set(item, "call_id", st.id)
|
||||
item, _ = sjson.Set(item, "name", st.name)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
}
|
||||
if len(outputs) > 0 {
|
||||
out, _ = sjson.Set(out, "output", outputs)
|
||||
if gjson.Get(outputsWrapper, "arr.#").Int() > 0 {
|
||||
out, _ = sjson.SetRaw(out, "output", gjson.Get(outputsWrapper, "arr").Raw)
|
||||
}
|
||||
|
||||
// Usage
|
||||
|
||||
@@ -9,7 +9,6 @@ package claude
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
@@ -191,21 +190,12 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original
|
||||
return ""
|
||||
}
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": responseData.Get("id").String(),
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": responseData.Get("model").String(),
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": responseData.Get("usage.input_tokens").Int(),
|
||||
"output_tokens": responseData.Get("usage.output_tokens").Int(),
|
||||
},
|
||||
}
|
||||
out := `{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
out, _ = sjson.Set(out, "id", responseData.Get("id").String())
|
||||
out, _ = sjson.Set(out, "model", responseData.Get("model").String())
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", responseData.Get("usage.input_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", responseData.Get("usage.output_tokens").Int())
|
||||
|
||||
var contentBlocks []interface{}
|
||||
hasToolCall := false
|
||||
|
||||
if output := responseData.Get("output"); output.Exists() && output.IsArray() {
|
||||
@@ -244,10 +234,9 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original
|
||||
}
|
||||
}
|
||||
if thinkingBuilder.Len() > 0 {
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": thinkingBuilder.String(),
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", thinkingBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
case "message":
|
||||
if content := item.Get("content"); content.Exists() {
|
||||
@@ -256,10 +245,9 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original
|
||||
if part.Get("type").String() == "output_text" {
|
||||
text := part.Get("text").String()
|
||||
if text != "" {
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": text,
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", text)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
}
|
||||
return true
|
||||
@@ -267,10 +255,9 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original
|
||||
} else {
|
||||
text := content.String()
|
||||
if text != "" {
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": text,
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", text)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -281,54 +268,41 @@ func ConvertCodexResponseToClaudeNonStream(_ context.Context, _ string, original
|
||||
name = original
|
||||
}
|
||||
|
||||
toolBlock := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": item.Get("call_id").String(),
|
||||
"name": name,
|
||||
"input": map[string]interface{}{},
|
||||
}
|
||||
|
||||
if argsStr := item.Get("arguments").String(); argsStr != "" {
|
||||
var args interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
|
||||
toolBlock["input"] = args
|
||||
toolBlock := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolBlock, _ = sjson.Set(toolBlock, "id", item.Get("call_id").String())
|
||||
toolBlock, _ = sjson.Set(toolBlock, "name", name)
|
||||
inputRaw := "{}"
|
||||
if argsStr := item.Get("arguments").String(); argsStr != "" && gjson.Valid(argsStr) {
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
inputRaw = argsJSON.Raw
|
||||
}
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolBlock)
|
||||
toolBlock, _ = sjson.SetRaw(toolBlock, "input", inputRaw)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolBlock)
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
if len(contentBlocks) > 0 {
|
||||
response["content"] = contentBlocks
|
||||
}
|
||||
|
||||
if stopReason := responseData.Get("stop_reason"); stopReason.Exists() && stopReason.String() != "" {
|
||||
response["stop_reason"] = stopReason.String()
|
||||
out, _ = sjson.Set(out, "stop_reason", stopReason.String())
|
||||
} else if hasToolCall {
|
||||
response["stop_reason"] = "tool_use"
|
||||
out, _ = sjson.Set(out, "stop_reason", "tool_use")
|
||||
} else {
|
||||
response["stop_reason"] = "end_turn"
|
||||
out, _ = sjson.Set(out, "stop_reason", "end_turn")
|
||||
}
|
||||
|
||||
if stopSequence := responseData.Get("stop_sequence"); stopSequence.Exists() && stopSequence.String() != "" {
|
||||
response["stop_sequence"] = stopSequence.Value()
|
||||
out, _ = sjson.SetRaw(out, "stop_sequence", stopSequence.Raw)
|
||||
}
|
||||
|
||||
if responseData.Get("usage.input_tokens").Exists() || responseData.Get("usage.output_tokens").Exists() {
|
||||
response["usage"] = map[string]interface{}{
|
||||
"input_tokens": responseData.Get("usage.input_tokens").Int(),
|
||||
"output_tokens": responseData.Get("usage.output_tokens").Int(),
|
||||
}
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", responseData.Get("usage.input_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", responseData.Get("usage.output_tokens").Int())
|
||||
}
|
||||
|
||||
responseJSON, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(responseJSON)
|
||||
return out
|
||||
}
|
||||
|
||||
// buildReverseMapFromClaudeOriginalShortToOriginal builds a map[short]original from original Claude request tools.
|
||||
|
||||
@@ -7,7 +7,6 @@ package gemini
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"time"
|
||||
|
||||
@@ -190,19 +189,19 @@ func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
}
|
||||
|
||||
// Process output content to build parts array
|
||||
var parts []interface{}
|
||||
hasToolCall := false
|
||||
var pendingFunctionCalls []interface{}
|
||||
var pendingFunctionCalls []string
|
||||
|
||||
flushPendingFunctionCalls := func() {
|
||||
if len(pendingFunctionCalls) > 0 {
|
||||
// Add all pending function calls as individual parts
|
||||
// This maintains the original Gemini API format while ensuring consecutive calls are grouped together
|
||||
for _, fc := range pendingFunctionCalls {
|
||||
parts = append(parts, fc)
|
||||
}
|
||||
pendingFunctionCalls = nil
|
||||
if len(pendingFunctionCalls) == 0 {
|
||||
return
|
||||
}
|
||||
// Add all pending function calls as individual parts
|
||||
// This maintains the original Gemini API format while ensuring consecutive calls are grouped together
|
||||
for _, fc := range pendingFunctionCalls {
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", fc)
|
||||
}
|
||||
pendingFunctionCalls = nil
|
||||
}
|
||||
|
||||
if output := responseData.Get("output"); output.Exists() && output.IsArray() {
|
||||
@@ -216,11 +215,9 @@ func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
|
||||
// Add thinking content
|
||||
if content := value.Get("content"); content.Exists() {
|
||||
part := map[string]interface{}{
|
||||
"thought": true,
|
||||
"text": content.String(),
|
||||
}
|
||||
parts = append(parts, part)
|
||||
part := `{"text":"","thought":true}`
|
||||
part, _ = sjson.Set(part, "text", content.String())
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", part)
|
||||
}
|
||||
|
||||
case "message":
|
||||
@@ -232,10 +229,9 @@ func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
content.ForEach(func(_, contentItem gjson.Result) bool {
|
||||
if contentItem.Get("type").String() == "output_text" {
|
||||
if text := contentItem.Get("text"); text.Exists() {
|
||||
part := map[string]interface{}{
|
||||
"text": text.String(),
|
||||
}
|
||||
parts = append(parts, part)
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", text.String())
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts.-1", part)
|
||||
}
|
||||
}
|
||||
return true
|
||||
@@ -245,28 +241,21 @@ func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
case "function_call":
|
||||
// Collect function call for potential merging with consecutive ones
|
||||
hasToolCall = true
|
||||
functionCall := map[string]interface{}{
|
||||
"functionCall": map[string]interface{}{
|
||||
"name": func() string {
|
||||
n := value.Get("name").String()
|
||||
rev := buildReverseMapFromGeminiOriginal(originalRequestRawJSON)
|
||||
if orig, ok := rev[n]; ok {
|
||||
return orig
|
||||
}
|
||||
return n
|
||||
}(),
|
||||
"args": map[string]interface{}{},
|
||||
},
|
||||
functionCall := `{"functionCall":{"args":{},"name":""}}`
|
||||
{
|
||||
n := value.Get("name").String()
|
||||
rev := buildReverseMapFromGeminiOriginal(originalRequestRawJSON)
|
||||
if orig, ok := rev[n]; ok {
|
||||
n = orig
|
||||
}
|
||||
functionCall, _ = sjson.Set(functionCall, "functionCall.name", n)
|
||||
}
|
||||
|
||||
// Parse and set arguments
|
||||
if argsStr := value.Get("arguments").String(); argsStr != "" {
|
||||
argsResult := gjson.Parse(argsStr)
|
||||
if argsResult.IsObject() {
|
||||
var args map[string]interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
|
||||
functionCall["functionCall"].(map[string]interface{})["args"] = args
|
||||
}
|
||||
functionCall, _ = sjson.SetRaw(functionCall, "functionCall.args", argsStr)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -279,11 +268,6 @@ func ConvertCodexResponseToGeminiNonStream(_ context.Context, modelName string,
|
||||
flushPendingFunctionCalls()
|
||||
}
|
||||
|
||||
// Set the parts array
|
||||
if len(parts) > 0 {
|
||||
template, _ = sjson.SetRaw(template, "candidates.0.content.parts", mustMarshalJSON(parts))
|
||||
}
|
||||
|
||||
// Set finish reason based on whether there were tool calls
|
||||
if hasToolCall {
|
||||
template, _ = sjson.Set(template, "candidates.0.finishReason", "STOP")
|
||||
@@ -323,15 +307,6 @@ func buildReverseMapFromGeminiOriginal(original []byte) map[string]string {
|
||||
return rev
|
||||
}
|
||||
|
||||
// mustMarshalJSON marshals a value to JSON, panicking on error.
|
||||
func mustMarshalJSON(v interface{}) string {
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(data)
|
||||
}
|
||||
|
||||
func GeminiTokenCount(ctx context.Context, count int64) string {
|
||||
return fmt.Sprintf(`{"totalTokens":%d,"promptTokensDetails":[{"modality":"TEXT","tokenCount":%d}]}`, count, count)
|
||||
}
|
||||
|
||||
@@ -7,10 +7,8 @@ package claude
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
client "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/common"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -41,92 +39,102 @@ func ConvertClaudeRequestToCLI(modelName string, inputRawJSON []byte, _ bool) []
|
||||
rawJSON := bytes.Clone(inputRawJSON)
|
||||
rawJSON = bytes.Replace(rawJSON, []byte(`"url":{"type":"string","format":"uri",`), []byte(`"url":{"type":"string",`), -1)
|
||||
|
||||
// Build output Gemini CLI request JSON
|
||||
out := `{"model":"","request":{"contents":[]}}`
|
||||
out, _ = sjson.Set(out, "model", modelName)
|
||||
|
||||
// system instruction
|
||||
var systemInstruction *client.Content
|
||||
systemResult := gjson.GetBytes(rawJSON, "system")
|
||||
if systemResult.IsArray() {
|
||||
systemResults := systemResult.Array()
|
||||
systemInstruction = &client.Content{Role: "user", Parts: []client.Part{}}
|
||||
for i := 0; i < len(systemResults); i++ {
|
||||
systemPromptResult := systemResults[i]
|
||||
systemTypePromptResult := systemPromptResult.Get("type")
|
||||
if systemTypePromptResult.Type == gjson.String && systemTypePromptResult.String() == "text" {
|
||||
systemPrompt := systemPromptResult.Get("text").String()
|
||||
systemPart := client.Part{Text: systemPrompt}
|
||||
systemInstruction.Parts = append(systemInstruction.Parts, systemPart)
|
||||
if systemResult := gjson.GetBytes(rawJSON, "system"); systemResult.IsArray() {
|
||||
systemInstruction := `{"role":"user","parts":[]}`
|
||||
hasSystemParts := false
|
||||
systemResult.ForEach(func(_, systemPromptResult gjson.Result) bool {
|
||||
if systemPromptResult.Get("type").String() == "text" {
|
||||
textResult := systemPromptResult.Get("text")
|
||||
if textResult.Type == gjson.String {
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", textResult.String())
|
||||
systemInstruction, _ = sjson.SetRaw(systemInstruction, "parts.-1", part)
|
||||
hasSystemParts = true
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if hasSystemParts {
|
||||
out, _ = sjson.SetRaw(out, "request.systemInstruction", systemInstruction)
|
||||
}
|
||||
if len(systemInstruction.Parts) == 0 {
|
||||
systemInstruction = nil
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
out, _ = sjson.Set(out, "request.systemInstruction.parts.-1.text", systemResult.String())
|
||||
}
|
||||
|
||||
// contents
|
||||
contents := make([]client.Content, 0)
|
||||
messagesResult := gjson.GetBytes(rawJSON, "messages")
|
||||
if messagesResult.IsArray() {
|
||||
messageResults := messagesResult.Array()
|
||||
for i := 0; i < len(messageResults); i++ {
|
||||
messageResult := messageResults[i]
|
||||
if messagesResult := gjson.GetBytes(rawJSON, "messages"); messagesResult.IsArray() {
|
||||
messagesResult.ForEach(func(_, messageResult gjson.Result) bool {
|
||||
roleResult := messageResult.Get("role")
|
||||
if roleResult.Type != gjson.String {
|
||||
continue
|
||||
return true
|
||||
}
|
||||
role := roleResult.String()
|
||||
if role == "assistant" {
|
||||
role = "model"
|
||||
}
|
||||
clientContent := client.Content{Role: role, Parts: []client.Part{}}
|
||||
|
||||
contentJSON := `{"role":"","parts":[]}`
|
||||
contentJSON, _ = sjson.Set(contentJSON, "role", role)
|
||||
|
||||
contentsResult := messageResult.Get("content")
|
||||
if contentsResult.IsArray() {
|
||||
contentResults := contentsResult.Array()
|
||||
for j := 0; j < len(contentResults); j++ {
|
||||
contentResult := contentResults[j]
|
||||
contentTypeResult := contentResult.Get("type")
|
||||
if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "text" {
|
||||
prompt := contentResult.Get("text").String()
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{Text: prompt})
|
||||
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_use" {
|
||||
contentsResult.ForEach(func(_, contentResult gjson.Result) bool {
|
||||
switch contentResult.Get("type").String() {
|
||||
case "text":
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", contentResult.Get("text").String())
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
|
||||
case "tool_use":
|
||||
functionName := contentResult.Get("name").String()
|
||||
functionArgs := contentResult.Get("input").String()
|
||||
var args map[string]any
|
||||
if err := json.Unmarshal([]byte(functionArgs), &args); err == nil {
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{
|
||||
FunctionCall: &client.FunctionCall{Name: functionName, Args: args},
|
||||
ThoughtSignature: geminiCLIClaudeThoughtSignature,
|
||||
})
|
||||
argsResult := gjson.Parse(functionArgs)
|
||||
if argsResult.IsObject() && gjson.Valid(functionArgs) {
|
||||
part := `{"thoughtSignature":"","functionCall":{"name":"","args":{}}}`
|
||||
part, _ = sjson.Set(part, "thoughtSignature", geminiCLIClaudeThoughtSignature)
|
||||
part, _ = sjson.Set(part, "functionCall.name", functionName)
|
||||
part, _ = sjson.SetRaw(part, "functionCall.args", functionArgs)
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
}
|
||||
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_result" {
|
||||
|
||||
case "tool_result":
|
||||
toolCallID := contentResult.Get("tool_use_id").String()
|
||||
if toolCallID != "" {
|
||||
funcName := toolCallID
|
||||
toolCallIDs := strings.Split(toolCallID, "-")
|
||||
if len(toolCallIDs) > 1 {
|
||||
funcName = strings.Join(toolCallIDs[0:len(toolCallIDs)-1], "-")
|
||||
}
|
||||
responseData := contentResult.Get("content").Raw
|
||||
functionResponse := client.FunctionResponse{Name: funcName, Response: map[string]interface{}{"result": responseData}}
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{FunctionResponse: &functionResponse})
|
||||
if toolCallID == "" {
|
||||
return true
|
||||
}
|
||||
funcName := toolCallID
|
||||
toolCallIDs := strings.Split(toolCallID, "-")
|
||||
if len(toolCallIDs) > 1 {
|
||||
funcName = strings.Join(toolCallIDs[0:len(toolCallIDs)-1], "-")
|
||||
}
|
||||
responseData := contentResult.Get("content").Raw
|
||||
part := `{"functionResponse":{"name":"","response":{"result":""}}}`
|
||||
part, _ = sjson.Set(part, "functionResponse.name", funcName)
|
||||
part, _ = sjson.Set(part, "functionResponse.response.result", responseData)
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
}
|
||||
}
|
||||
contents = append(contents, clientContent)
|
||||
return true
|
||||
})
|
||||
out, _ = sjson.SetRaw(out, "request.contents.-1", contentJSON)
|
||||
} else if contentsResult.Type == gjson.String {
|
||||
prompt := contentsResult.String()
|
||||
contents = append(contents, client.Content{Role: role, Parts: []client.Part{{Text: prompt}}})
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", contentsResult.String())
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
out, _ = sjson.SetRaw(out, "request.contents.-1", contentJSON)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// tools
|
||||
var tools []client.ToolDeclaration
|
||||
toolsResult := gjson.GetBytes(rawJSON, "tools")
|
||||
if toolsResult.IsArray() {
|
||||
tools = make([]client.ToolDeclaration, 1)
|
||||
tools[0].FunctionDeclarations = make([]any, 0)
|
||||
toolsResults := toolsResult.Array()
|
||||
for i := 0; i < len(toolsResults); i++ {
|
||||
toolResult := toolsResults[i]
|
||||
if toolsResult := gjson.GetBytes(rawJSON, "tools"); toolsResult.IsArray() {
|
||||
hasTools := false
|
||||
toolsResult.ForEach(func(_, toolResult gjson.Result) bool {
|
||||
inputSchemaResult := toolResult.Get("input_schema")
|
||||
if inputSchemaResult.Exists() && inputSchemaResult.IsObject() {
|
||||
inputSchema := inputSchemaResult.Raw
|
||||
@@ -136,30 +144,19 @@ func ConvertClaudeRequestToCLI(modelName string, inputRawJSON []byte, _ bool) []
|
||||
tool, _ = sjson.Delete(tool, "input_examples")
|
||||
tool, _ = sjson.Delete(tool, "type")
|
||||
tool, _ = sjson.Delete(tool, "cache_control")
|
||||
var toolDeclaration any
|
||||
if err := json.Unmarshal([]byte(tool), &toolDeclaration); err == nil {
|
||||
tools[0].FunctionDeclarations = append(tools[0].FunctionDeclarations, toolDeclaration)
|
||||
if gjson.Valid(tool) && gjson.Parse(tool).IsObject() {
|
||||
if !hasTools {
|
||||
out, _ = sjson.SetRaw(out, "request.tools", `[{"functionDeclarations":[]}]`)
|
||||
hasTools = true
|
||||
}
|
||||
out, _ = sjson.SetRaw(out, "request.tools.0.functionDeclarations.-1", tool)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if !hasTools {
|
||||
out, _ = sjson.Delete(out, "request.tools")
|
||||
}
|
||||
} else {
|
||||
tools = make([]client.ToolDeclaration, 0)
|
||||
}
|
||||
|
||||
// Build output Gemini CLI request JSON
|
||||
out := `{"model":"","request":{"contents":[]}}`
|
||||
out, _ = sjson.Set(out, "model", modelName)
|
||||
if systemInstruction != nil {
|
||||
b, _ := json.Marshal(systemInstruction)
|
||||
out, _ = sjson.SetRaw(out, "request.systemInstruction", string(b))
|
||||
}
|
||||
if len(contents) > 0 {
|
||||
b, _ := json.Marshal(contents)
|
||||
out, _ = sjson.SetRaw(out, "request.contents", string(b))
|
||||
}
|
||||
if len(tools) > 0 && len(tools[0].FunctionDeclarations) > 0 {
|
||||
b, _ := json.Marshal(tools)
|
||||
out, _ = sjson.SetRaw(out, "request.tools", string(b))
|
||||
}
|
||||
|
||||
// Map Anthropic thinking -> Gemini thinkingBudget/include_thoughts when type==enabled
|
||||
|
||||
@@ -9,7 +9,6 @@ package claude
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
@@ -276,22 +275,16 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": root.Get("response.responseId").String(),
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": root.Get("response.modelVersion").String(),
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": root.Get("response.usageMetadata.promptTokenCount").Int(),
|
||||
"output_tokens": root.Get("response.usageMetadata.candidatesTokenCount").Int() + root.Get("response.usageMetadata.thoughtsTokenCount").Int(),
|
||||
},
|
||||
}
|
||||
out := `{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
out, _ = sjson.Set(out, "id", root.Get("response.responseId").String())
|
||||
out, _ = sjson.Set(out, "model", root.Get("response.modelVersion").String())
|
||||
|
||||
inputTokens := root.Get("response.usageMetadata.promptTokenCount").Int()
|
||||
outputTokens := root.Get("response.usageMetadata.candidatesTokenCount").Int() + root.Get("response.usageMetadata.thoughtsTokenCount").Int()
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", inputTokens)
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", outputTokens)
|
||||
|
||||
parts := root.Get("response.candidates.0.content.parts")
|
||||
var contentBlocks []interface{}
|
||||
textBuilder := strings.Builder{}
|
||||
thinkingBuilder := strings.Builder{}
|
||||
toolIDCounter := 0
|
||||
@@ -301,10 +294,9 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
if textBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": textBuilder.String(),
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", textBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
textBuilder.Reset()
|
||||
}
|
||||
|
||||
@@ -312,10 +304,9 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
if thinkingBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": thinkingBuilder.String(),
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", thinkingBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
thinkingBuilder.Reset()
|
||||
}
|
||||
|
||||
@@ -339,21 +330,15 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
|
||||
name := functionCall.Get("name").String()
|
||||
toolIDCounter++
|
||||
toolBlock := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": fmt.Sprintf("tool_%d", toolIDCounter),
|
||||
"name": name,
|
||||
"input": map[string]interface{}{},
|
||||
toolBlock := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolBlock, _ = sjson.Set(toolBlock, "id", fmt.Sprintf("tool_%d", toolIDCounter))
|
||||
toolBlock, _ = sjson.Set(toolBlock, "name", name)
|
||||
inputRaw := "{}"
|
||||
if args := functionCall.Get("args"); args.Exists() && gjson.Valid(args.Raw) && args.IsObject() {
|
||||
inputRaw = args.Raw
|
||||
}
|
||||
|
||||
if args := functionCall.Get("args"); args.Exists() {
|
||||
var parsed interface{}
|
||||
if err := json.Unmarshal([]byte(args.Raw), &parsed); err == nil {
|
||||
toolBlock["input"] = parsed
|
||||
}
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolBlock)
|
||||
toolBlock, _ = sjson.SetRaw(toolBlock, "input", inputRaw)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolBlock)
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -362,8 +347,6 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
flushThinking()
|
||||
flushText()
|
||||
|
||||
response["content"] = contentBlocks
|
||||
|
||||
stopReason := "end_turn"
|
||||
if hasToolCall {
|
||||
stopReason = "tool_use"
|
||||
@@ -379,19 +362,13 @@ func ConvertGeminiCLIResponseToClaudeNonStream(_ context.Context, _ string, orig
|
||||
}
|
||||
}
|
||||
}
|
||||
response["stop_reason"] = stopReason
|
||||
out, _ = sjson.Set(out, "stop_reason", stopReason)
|
||||
|
||||
if usage := response["usage"].(map[string]interface{}); usage["input_tokens"] == int64(0) && usage["output_tokens"] == int64(0) {
|
||||
if usageMeta := root.Get("response.usageMetadata"); !usageMeta.Exists() {
|
||||
delete(response, "usage")
|
||||
}
|
||||
if inputTokens == int64(0) && outputTokens == int64(0) && !root.Get("response.usageMetadata").Exists() {
|
||||
out, _ = sjson.Delete(out, "usage")
|
||||
}
|
||||
|
||||
encoded, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(encoded)
|
||||
return out
|
||||
}
|
||||
|
||||
func ClaudeTokenCount(ctx context.Context, count int64) string {
|
||||
|
||||
@@ -7,7 +7,6 @@ package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/common"
|
||||
@@ -117,8 +116,6 @@ func ConvertGeminiRequestToGeminiCLI(_ string, inputRawJSON []byte, _ bool) []by
|
||||
|
||||
// FunctionCallGroup represents a group of function calls and their responses
|
||||
type FunctionCallGroup struct {
|
||||
ModelContent map[string]interface{}
|
||||
FunctionCalls []gjson.Result
|
||||
ResponsesNeeded int
|
||||
}
|
||||
|
||||
@@ -146,7 +143,7 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
}
|
||||
|
||||
// Initialize data structures for processing and grouping
|
||||
var newContents []interface{} // Final processed contents array
|
||||
contentsWrapper := `{"contents":[]}`
|
||||
var pendingGroups []*FunctionCallGroup // Groups awaiting completion with responses
|
||||
var collectedResponses []gjson.Result // Standalone responses to be matched
|
||||
|
||||
@@ -178,23 +175,17 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
collectedResponses = collectedResponses[group.ResponsesNeeded:]
|
||||
|
||||
// Create merged function response content
|
||||
var responseParts []interface{}
|
||||
functionResponseContent := `{"parts":[],"role":"function"}`
|
||||
for _, response := range groupResponses {
|
||||
var responseMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(response.Raw), &responseMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal function response: %v\n", errUnmarshal)
|
||||
if !response.IsObject() {
|
||||
log.Warnf("failed to parse function response")
|
||||
continue
|
||||
}
|
||||
responseParts = append(responseParts, responseMap)
|
||||
functionResponseContent, _ = sjson.SetRaw(functionResponseContent, "parts.-1", response.Raw)
|
||||
}
|
||||
|
||||
if len(responseParts) > 0 {
|
||||
functionResponseContent := map[string]interface{}{
|
||||
"parts": responseParts,
|
||||
"role": "function",
|
||||
}
|
||||
newContents = append(newContents, functionResponseContent)
|
||||
if gjson.Get(functionResponseContent, "parts.#").Int() > 0 {
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", functionResponseContent)
|
||||
}
|
||||
|
||||
// Remove this group as it's been satisfied
|
||||
@@ -208,50 +199,42 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
|
||||
// If this is a model with function calls, create a new group
|
||||
if role == "model" {
|
||||
var functionCallsInThisModel []gjson.Result
|
||||
functionCallsCount := 0
|
||||
parts.ForEach(func(_, part gjson.Result) bool {
|
||||
if part.Get("functionCall").Exists() {
|
||||
functionCallsInThisModel = append(functionCallsInThisModel, part)
|
||||
functionCallsCount++
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if len(functionCallsInThisModel) > 0 {
|
||||
if functionCallsCount > 0 {
|
||||
// Add the model content
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal model content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse model content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
|
||||
// Create a new group for tracking responses
|
||||
group := &FunctionCallGroup{
|
||||
ModelContent: contentMap,
|
||||
FunctionCalls: functionCallsInThisModel,
|
||||
ResponsesNeeded: len(functionCallsInThisModel),
|
||||
ResponsesNeeded: functionCallsCount,
|
||||
}
|
||||
pendingGroups = append(pendingGroups, group)
|
||||
} else {
|
||||
// Regular model content without function calls
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
}
|
||||
} else {
|
||||
// Non-model content (user, etc.)
|
||||
var contentMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(value.Raw), &contentMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal content: %v\n", errUnmarshal)
|
||||
if !value.IsObject() {
|
||||
log.Warnf("failed to parse content")
|
||||
return true
|
||||
}
|
||||
newContents = append(newContents, contentMap)
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", value.Raw)
|
||||
}
|
||||
|
||||
return true
|
||||
@@ -263,31 +246,24 @@ func fixCLIToolResponse(input string) (string, error) {
|
||||
groupResponses := collectedResponses[:group.ResponsesNeeded]
|
||||
collectedResponses = collectedResponses[group.ResponsesNeeded:]
|
||||
|
||||
var responseParts []interface{}
|
||||
functionResponseContent := `{"parts":[],"role":"function"}`
|
||||
for _, response := range groupResponses {
|
||||
var responseMap map[string]interface{}
|
||||
errUnmarshal := json.Unmarshal([]byte(response.Raw), &responseMap)
|
||||
if errUnmarshal != nil {
|
||||
log.Warnf("failed to unmarshal function response: %v\n", errUnmarshal)
|
||||
if !response.IsObject() {
|
||||
log.Warnf("failed to parse function response")
|
||||
continue
|
||||
}
|
||||
responseParts = append(responseParts, responseMap)
|
||||
functionResponseContent, _ = sjson.SetRaw(functionResponseContent, "parts.-1", response.Raw)
|
||||
}
|
||||
|
||||
if len(responseParts) > 0 {
|
||||
functionResponseContent := map[string]interface{}{
|
||||
"parts": responseParts,
|
||||
"role": "function",
|
||||
}
|
||||
newContents = append(newContents, functionResponseContent)
|
||||
if gjson.Get(functionResponseContent, "parts.#").Int() > 0 {
|
||||
contentsWrapper, _ = sjson.SetRaw(contentsWrapper, "contents.-1", functionResponseContent)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Update the original JSON with the new contents
|
||||
result := input
|
||||
newContentsJSON, _ := json.Marshal(newContents)
|
||||
result, _ = sjson.Set(result, "request.contents", json.RawMessage(newContentsJSON))
|
||||
result, _ = sjson.SetRaw(result, "request.contents", gjson.Get(contentsWrapper, "contents").Raw)
|
||||
|
||||
return result, nil
|
||||
}
|
||||
|
||||
@@ -160,6 +160,14 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
@@ -210,8 +218,29 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
if content.Type == gjson.String {
|
||||
// Assistant text -> single model content
|
||||
node, _ = sjson.SetBytes(node, "parts.-1.text", content.String())
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
|
||||
p++
|
||||
} else if content.IsArray() {
|
||||
// Assistant multimodal content (e.g. text + image) -> single model content with parts
|
||||
for _, item := range content.Array() {
|
||||
switch item.Get("type").String() {
|
||||
case "text":
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".text", item.Get("text").String())
|
||||
p++
|
||||
case "image_url":
|
||||
// If the assistant returned an inline data URL, preserve it for history fidelity.
|
||||
imageURL := item.Get("image_url.url").String()
|
||||
if len(imageURL) > 5 { // expect data:...
|
||||
pieces := strings.SplitN(imageURL[5:], ";", 2)
|
||||
if len(pieces) == 2 && len(pieces[1]) > 7 {
|
||||
mime := pieces[0]
|
||||
data := pieces[1][7:]
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".inlineData.mime_type", mime)
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".inlineData.data", data)
|
||||
p++
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Tool calls -> single model content with functionCall parts
|
||||
@@ -236,7 +265,7 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
|
||||
|
||||
// Append a single tool content combining name + response per function
|
||||
toolNode := []byte(`{"role":"tool","parts":[]}`)
|
||||
toolNode := []byte(`{"role":"user","parts":[]}`)
|
||||
pp := 0
|
||||
for _, fid := range fIDs {
|
||||
if name, ok := tcID2Name[fid]; ok {
|
||||
@@ -252,6 +281,8 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
if pp > 0 {
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", toolNode)
|
||||
}
|
||||
} else {
|
||||
out, _ = sjson.SetRawBytes(out, "request.contents.-1", node)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -278,7 +309,7 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
@@ -293,7 +324,7 @@ func ConvertOpenAIRequestToGeminiCLI(modelName string, inputRawJSON []byte, _ bo
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
|
||||
@@ -8,7 +8,6 @@ package chat_completions
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
@@ -171,21 +170,16 @@ func ConvertCliResponseToOpenAI(_ context.Context, _ string, originalRequestRawJ
|
||||
mimeType = "image/png"
|
||||
}
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
imagePayload, err := json.Marshal(map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]string{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
imagesResult := gjson.Get(template, "choices.0.delta.images")
|
||||
if !imagesResult.Exists() || !imagesResult.IsArray() {
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images", `[]`)
|
||||
}
|
||||
imageIndex := len(gjson.Get(template, "choices.0.delta.images").Array())
|
||||
imagePayload := `{"type":"image_url","image_url":{"url":""}}`
|
||||
imagePayload, _ = sjson.Set(imagePayload, "index", imageIndex)
|
||||
imagePayload, _ = sjson.Set(imagePayload, "image_url.url", imageURL)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", string(imagePayload))
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", imagePayload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -7,10 +7,8 @@ package claude
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
client "github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/translator/gemini/common"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
"github.com/tidwall/gjson"
|
||||
@@ -34,92 +32,102 @@ func ConvertClaudeRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
rawJSON := bytes.Clone(inputRawJSON)
|
||||
rawJSON = bytes.Replace(rawJSON, []byte(`"url":{"type":"string","format":"uri",`), []byte(`"url":{"type":"string",`), -1)
|
||||
|
||||
// Build output Gemini CLI request JSON
|
||||
out := `{"contents":[]}`
|
||||
out, _ = sjson.Set(out, "model", modelName)
|
||||
|
||||
// system instruction
|
||||
var systemInstruction *client.Content
|
||||
systemResult := gjson.GetBytes(rawJSON, "system")
|
||||
if systemResult.IsArray() {
|
||||
systemResults := systemResult.Array()
|
||||
systemInstruction = &client.Content{Role: "user", Parts: []client.Part{}}
|
||||
for i := 0; i < len(systemResults); i++ {
|
||||
systemPromptResult := systemResults[i]
|
||||
systemTypePromptResult := systemPromptResult.Get("type")
|
||||
if systemTypePromptResult.Type == gjson.String && systemTypePromptResult.String() == "text" {
|
||||
systemPrompt := systemPromptResult.Get("text").String()
|
||||
systemPart := client.Part{Text: systemPrompt}
|
||||
systemInstruction.Parts = append(systemInstruction.Parts, systemPart)
|
||||
if systemResult := gjson.GetBytes(rawJSON, "system"); systemResult.IsArray() {
|
||||
systemInstruction := `{"role":"user","parts":[]}`
|
||||
hasSystemParts := false
|
||||
systemResult.ForEach(func(_, systemPromptResult gjson.Result) bool {
|
||||
if systemPromptResult.Get("type").String() == "text" {
|
||||
textResult := systemPromptResult.Get("text")
|
||||
if textResult.Type == gjson.String {
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", textResult.String())
|
||||
systemInstruction, _ = sjson.SetRaw(systemInstruction, "parts.-1", part)
|
||||
hasSystemParts = true
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if hasSystemParts {
|
||||
out, _ = sjson.SetRaw(out, "system_instruction", systemInstruction)
|
||||
}
|
||||
if len(systemInstruction.Parts) == 0 {
|
||||
systemInstruction = nil
|
||||
}
|
||||
} else if systemResult.Type == gjson.String {
|
||||
out, _ = sjson.Set(out, "request.system_instruction.parts.-1.text", systemResult.String())
|
||||
}
|
||||
|
||||
// contents
|
||||
contents := make([]client.Content, 0)
|
||||
messagesResult := gjson.GetBytes(rawJSON, "messages")
|
||||
if messagesResult.IsArray() {
|
||||
messageResults := messagesResult.Array()
|
||||
for i := 0; i < len(messageResults); i++ {
|
||||
messageResult := messageResults[i]
|
||||
if messagesResult := gjson.GetBytes(rawJSON, "messages"); messagesResult.IsArray() {
|
||||
messagesResult.ForEach(func(_, messageResult gjson.Result) bool {
|
||||
roleResult := messageResult.Get("role")
|
||||
if roleResult.Type != gjson.String {
|
||||
continue
|
||||
return true
|
||||
}
|
||||
role := roleResult.String()
|
||||
if role == "assistant" {
|
||||
role = "model"
|
||||
}
|
||||
clientContent := client.Content{Role: role, Parts: []client.Part{}}
|
||||
|
||||
contentJSON := `{"role":"","parts":[]}`
|
||||
contentJSON, _ = sjson.Set(contentJSON, "role", role)
|
||||
|
||||
contentsResult := messageResult.Get("content")
|
||||
if contentsResult.IsArray() {
|
||||
contentResults := contentsResult.Array()
|
||||
for j := 0; j < len(contentResults); j++ {
|
||||
contentResult := contentResults[j]
|
||||
contentTypeResult := contentResult.Get("type")
|
||||
if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "text" {
|
||||
prompt := contentResult.Get("text").String()
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{Text: prompt})
|
||||
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_use" {
|
||||
contentsResult.ForEach(func(_, contentResult gjson.Result) bool {
|
||||
switch contentResult.Get("type").String() {
|
||||
case "text":
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", contentResult.Get("text").String())
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
|
||||
case "tool_use":
|
||||
functionName := contentResult.Get("name").String()
|
||||
functionArgs := contentResult.Get("input").String()
|
||||
var args map[string]any
|
||||
if err := json.Unmarshal([]byte(functionArgs), &args); err == nil {
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{
|
||||
FunctionCall: &client.FunctionCall{Name: functionName, Args: args},
|
||||
ThoughtSignature: geminiClaudeThoughtSignature,
|
||||
})
|
||||
argsResult := gjson.Parse(functionArgs)
|
||||
if argsResult.IsObject() && gjson.Valid(functionArgs) {
|
||||
part := `{"thoughtSignature":"","functionCall":{"name":"","args":{}}}`
|
||||
part, _ = sjson.Set(part, "thoughtSignature", geminiClaudeThoughtSignature)
|
||||
part, _ = sjson.Set(part, "functionCall.name", functionName)
|
||||
part, _ = sjson.SetRaw(part, "functionCall.args", functionArgs)
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
}
|
||||
} else if contentTypeResult.Type == gjson.String && contentTypeResult.String() == "tool_result" {
|
||||
|
||||
case "tool_result":
|
||||
toolCallID := contentResult.Get("tool_use_id").String()
|
||||
if toolCallID != "" {
|
||||
funcName := toolCallID
|
||||
toolCallIDs := strings.Split(toolCallID, "-")
|
||||
if len(toolCallIDs) > 1 {
|
||||
funcName = strings.Join(toolCallIDs[0:len(toolCallIDs)-1], "-")
|
||||
}
|
||||
responseData := contentResult.Get("content").Raw
|
||||
functionResponse := client.FunctionResponse{Name: funcName, Response: map[string]interface{}{"result": responseData}}
|
||||
clientContent.Parts = append(clientContent.Parts, client.Part{FunctionResponse: &functionResponse})
|
||||
if toolCallID == "" {
|
||||
return true
|
||||
}
|
||||
funcName := toolCallID
|
||||
toolCallIDs := strings.Split(toolCallID, "-")
|
||||
if len(toolCallIDs) > 1 {
|
||||
funcName = strings.Join(toolCallIDs[0:len(toolCallIDs)-1], "-")
|
||||
}
|
||||
responseData := contentResult.Get("content").Raw
|
||||
part := `{"functionResponse":{"name":"","response":{"result":""}}}`
|
||||
part, _ = sjson.Set(part, "functionResponse.name", funcName)
|
||||
part, _ = sjson.Set(part, "functionResponse.response.result", responseData)
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
}
|
||||
}
|
||||
contents = append(contents, clientContent)
|
||||
return true
|
||||
})
|
||||
out, _ = sjson.SetRaw(out, "contents.-1", contentJSON)
|
||||
} else if contentsResult.Type == gjson.String {
|
||||
prompt := contentsResult.String()
|
||||
contents = append(contents, client.Content{Role: role, Parts: []client.Part{{Text: prompt}}})
|
||||
part := `{"text":""}`
|
||||
part, _ = sjson.Set(part, "text", contentsResult.String())
|
||||
contentJSON, _ = sjson.SetRaw(contentJSON, "parts.-1", part)
|
||||
out, _ = sjson.SetRaw(out, "contents.-1", contentJSON)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// tools
|
||||
var tools []client.ToolDeclaration
|
||||
toolsResult := gjson.GetBytes(rawJSON, "tools")
|
||||
if toolsResult.IsArray() {
|
||||
tools = make([]client.ToolDeclaration, 1)
|
||||
tools[0].FunctionDeclarations = make([]any, 0)
|
||||
toolsResults := toolsResult.Array()
|
||||
for i := 0; i < len(toolsResults); i++ {
|
||||
toolResult := toolsResults[i]
|
||||
if toolsResult := gjson.GetBytes(rawJSON, "tools"); toolsResult.IsArray() {
|
||||
hasTools := false
|
||||
toolsResult.ForEach(func(_, toolResult gjson.Result) bool {
|
||||
inputSchemaResult := toolResult.Get("input_schema")
|
||||
if inputSchemaResult.Exists() && inputSchemaResult.IsObject() {
|
||||
inputSchema := inputSchemaResult.Raw
|
||||
@@ -129,30 +137,19 @@ func ConvertClaudeRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
tool, _ = sjson.Delete(tool, "input_examples")
|
||||
tool, _ = sjson.Delete(tool, "type")
|
||||
tool, _ = sjson.Delete(tool, "cache_control")
|
||||
var toolDeclaration any
|
||||
if err := json.Unmarshal([]byte(tool), &toolDeclaration); err == nil {
|
||||
tools[0].FunctionDeclarations = append(tools[0].FunctionDeclarations, toolDeclaration)
|
||||
if gjson.Valid(tool) && gjson.Parse(tool).IsObject() {
|
||||
if !hasTools {
|
||||
out, _ = sjson.SetRaw(out, "tools", `[{"functionDeclarations":[]}]`)
|
||||
hasTools = true
|
||||
}
|
||||
out, _ = sjson.SetRaw(out, "tools.0.functionDeclarations.-1", tool)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
if !hasTools {
|
||||
out, _ = sjson.Delete(out, "tools")
|
||||
}
|
||||
} else {
|
||||
tools = make([]client.ToolDeclaration, 0)
|
||||
}
|
||||
|
||||
// Build output Gemini CLI request JSON
|
||||
out := `{"contents":[]}`
|
||||
out, _ = sjson.Set(out, "model", modelName)
|
||||
if systemInstruction != nil {
|
||||
b, _ := json.Marshal(systemInstruction)
|
||||
out, _ = sjson.SetRaw(out, "system_instruction", string(b))
|
||||
}
|
||||
if len(contents) > 0 {
|
||||
b, _ := json.Marshal(contents)
|
||||
out, _ = sjson.SetRaw(out, "contents", string(b))
|
||||
}
|
||||
if len(tools) > 0 && len(tools[0].FunctionDeclarations) > 0 {
|
||||
b, _ := json.Marshal(tools)
|
||||
out, _ = sjson.SetRaw(out, "tools", string(b))
|
||||
}
|
||||
|
||||
// Map Anthropic thinking -> Gemini thinkingBudget/include_thoughts when enabled
|
||||
|
||||
@@ -9,7 +9,6 @@ package claude
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
@@ -282,22 +281,16 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": root.Get("responseId").String(),
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": root.Get("modelVersion").String(),
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": root.Get("usageMetadata.promptTokenCount").Int(),
|
||||
"output_tokens": root.Get("usageMetadata.candidatesTokenCount").Int() + root.Get("usageMetadata.thoughtsTokenCount").Int(),
|
||||
},
|
||||
}
|
||||
out := `{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
out, _ = sjson.Set(out, "id", root.Get("responseId").String())
|
||||
out, _ = sjson.Set(out, "model", root.Get("modelVersion").String())
|
||||
|
||||
inputTokens := root.Get("usageMetadata.promptTokenCount").Int()
|
||||
outputTokens := root.Get("usageMetadata.candidatesTokenCount").Int() + root.Get("usageMetadata.thoughtsTokenCount").Int()
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", inputTokens)
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", outputTokens)
|
||||
|
||||
parts := root.Get("candidates.0.content.parts")
|
||||
var contentBlocks []interface{}
|
||||
textBuilder := strings.Builder{}
|
||||
thinkingBuilder := strings.Builder{}
|
||||
toolIDCounter := 0
|
||||
@@ -307,10 +300,9 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if textBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": textBuilder.String(),
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", textBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
textBuilder.Reset()
|
||||
}
|
||||
|
||||
@@ -318,10 +310,9 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if thinkingBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": thinkingBuilder.String(),
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", thinkingBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
thinkingBuilder.Reset()
|
||||
}
|
||||
|
||||
@@ -345,21 +336,15 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
|
||||
name := functionCall.Get("name").String()
|
||||
toolIDCounter++
|
||||
toolBlock := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": fmt.Sprintf("tool_%d", toolIDCounter),
|
||||
"name": name,
|
||||
"input": map[string]interface{}{},
|
||||
toolBlock := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolBlock, _ = sjson.Set(toolBlock, "id", fmt.Sprintf("tool_%d", toolIDCounter))
|
||||
toolBlock, _ = sjson.Set(toolBlock, "name", name)
|
||||
inputRaw := "{}"
|
||||
if args := functionCall.Get("args"); args.Exists() && gjson.Valid(args.Raw) && args.IsObject() {
|
||||
inputRaw = args.Raw
|
||||
}
|
||||
|
||||
if args := functionCall.Get("args"); args.Exists() {
|
||||
var parsed interface{}
|
||||
if err := json.Unmarshal([]byte(args.Raw), &parsed); err == nil {
|
||||
toolBlock["input"] = parsed
|
||||
}
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolBlock)
|
||||
toolBlock, _ = sjson.SetRaw(toolBlock, "input", inputRaw)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolBlock)
|
||||
continue
|
||||
}
|
||||
}
|
||||
@@ -368,8 +353,6 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
flushThinking()
|
||||
flushText()
|
||||
|
||||
response["content"] = contentBlocks
|
||||
|
||||
stopReason := "end_turn"
|
||||
if hasToolCall {
|
||||
stopReason = "tool_use"
|
||||
@@ -385,19 +368,13 @@ func ConvertGeminiResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
}
|
||||
}
|
||||
}
|
||||
response["stop_reason"] = stopReason
|
||||
out, _ = sjson.Set(out, "stop_reason", stopReason)
|
||||
|
||||
if usage := response["usage"].(map[string]interface{}); usage["input_tokens"] == int64(0) && usage["output_tokens"] == int64(0) {
|
||||
if usageMeta := root.Get("usageMetadata"); !usageMeta.Exists() {
|
||||
delete(response, "usage")
|
||||
}
|
||||
if inputTokens == int64(0) && outputTokens == int64(0) && !root.Get("usageMetadata").Exists() {
|
||||
out, _ = sjson.Delete(out, "usage")
|
||||
}
|
||||
|
||||
encoded, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(encoded)
|
||||
return out
|
||||
}
|
||||
|
||||
func ClaudeTokenCount(ctx context.Context, count int64) string {
|
||||
|
||||
@@ -178,6 +178,14 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
} else if content.IsObject() && content.Get("type").String() == "text" {
|
||||
out, _ = sjson.SetBytes(out, "system_instruction.role", "user")
|
||||
out, _ = sjson.SetBytes(out, "system_instruction.parts.0.text", content.Get("text").String())
|
||||
} else if content.IsArray() {
|
||||
contents := content.Array()
|
||||
if len(contents) > 0 {
|
||||
out, _ = sjson.SetBytes(out, "request.systemInstruction.role", "user")
|
||||
for j := 0; j < len(contents); j++ {
|
||||
out, _ = sjson.SetBytes(out, fmt.Sprintf("request.systemInstruction.parts.%d.text", j), contents[j].Get("text").String())
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if role == "user" || (role == "system" && len(arr) == 1) {
|
||||
// Build single user content node to avoid splitting into multiple contents
|
||||
@@ -225,18 +233,15 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
} else if role == "assistant" {
|
||||
node := []byte(`{"role":"model","parts":[]}`)
|
||||
p := 0
|
||||
|
||||
if content.Type == gjson.String {
|
||||
// Assistant text -> single model content
|
||||
node, _ = sjson.SetBytes(node, "parts.-1.text", content.String())
|
||||
out, _ = sjson.SetRawBytes(out, "contents.-1", node)
|
||||
p++
|
||||
} else if content.IsArray() {
|
||||
// Assistant multimodal content (e.g. text + image) -> single model content with parts
|
||||
for _, item := range content.Array() {
|
||||
switch item.Get("type").String() {
|
||||
case "text":
|
||||
node, _ = sjson.SetBytes(node, "parts."+itoa(p)+".text", item.Get("text").String())
|
||||
p++
|
||||
case "image_url":
|
||||
// If the assistant returned an inline data URL, preserve it for history fidelity.
|
||||
@@ -253,7 +258,6 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
}
|
||||
}
|
||||
}
|
||||
out, _ = sjson.SetRawBytes(out, "contents.-1", node)
|
||||
}
|
||||
|
||||
// Tool calls -> single model content with functionCall parts
|
||||
@@ -278,7 +282,7 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
out, _ = sjson.SetRawBytes(out, "contents.-1", node)
|
||||
|
||||
// Append a single tool content combining name + response per function
|
||||
toolNode := []byte(`{"role":"tool","parts":[]}`)
|
||||
toolNode := []byte(`{"role":"user","parts":[]}`)
|
||||
pp := 0
|
||||
for _, fid := range fIDs {
|
||||
if name, ok := tcID2Name[fid]; ok {
|
||||
@@ -294,6 +298,8 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
if pp > 0 {
|
||||
out, _ = sjson.SetRawBytes(out, "contents.-1", toolNode)
|
||||
}
|
||||
} else {
|
||||
out, _ = sjson.SetRawBytes(out, "contents.-1", node)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -320,7 +326,7 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
@@ -335,7 +341,7 @@ func ConvertOpenAIRequestToGemini(modelName string, inputRawJSON []byte, _ bool)
|
||||
log.Warnf("Failed to set default schema type for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
}
|
||||
fnRaw, errSet = sjson.Set(fnRaw, "parametersJsonSchema.properties", map[string]interface{}{})
|
||||
fnRaw, errSet = sjson.SetRaw(fnRaw, "parametersJsonSchema.properties", `{}`)
|
||||
if errSet != nil {
|
||||
log.Warnf("Failed to set default schema properties for tool '%s': %v", fn.Get("name").String(), errSet)
|
||||
continue
|
||||
|
||||
@@ -8,12 +8,12 @@ package chat_completions
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
@@ -89,18 +89,27 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
|
||||
|
||||
// Extract and set usage metadata (token counts).
|
||||
if usageResult := gjson.GetBytes(rawJSON, "usageMetadata"); usageResult.Exists() {
|
||||
cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int()
|
||||
if candidatesTokenCountResult := usageResult.Get("candidatesTokenCount"); candidatesTokenCountResult.Exists() {
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens", candidatesTokenCountResult.Int())
|
||||
}
|
||||
if totalTokenCountResult := usageResult.Get("totalTokenCount"); totalTokenCountResult.Exists() {
|
||||
template, _ = sjson.Set(template, "usage.total_tokens", totalTokenCountResult.Int())
|
||||
}
|
||||
promptTokenCount := usageResult.Get("promptTokenCount").Int()
|
||||
promptTokenCount := usageResult.Get("promptTokenCount").Int() - cachedTokenCount
|
||||
thoughtsTokenCount := usageResult.Get("thoughtsTokenCount").Int()
|
||||
template, _ = sjson.Set(template, "usage.prompt_tokens", promptTokenCount+thoughtsTokenCount)
|
||||
if thoughtsTokenCount > 0 {
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens_details.reasoning_tokens", thoughtsTokenCount)
|
||||
}
|
||||
// Include cached token count if present (indicates prompt caching is working)
|
||||
if cachedTokenCount > 0 {
|
||||
var err error
|
||||
template, err = sjson.Set(template, "usage.prompt_tokens_details.cached_tokens", cachedTokenCount)
|
||||
if err != nil {
|
||||
log.Warnf("gemini openai response: failed to set cached_tokens in streaming: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process the main content part of the response.
|
||||
@@ -173,21 +182,16 @@ func ConvertGeminiResponseToOpenAI(_ context.Context, _ string, originalRequestR
|
||||
mimeType = "image/png"
|
||||
}
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
imagePayload, err := json.Marshal(map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]string{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
imagesResult := gjson.Get(template, "choices.0.delta.images")
|
||||
if !imagesResult.Exists() || !imagesResult.IsArray() {
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images", `[]`)
|
||||
}
|
||||
imageIndex := len(gjson.Get(template, "choices.0.delta.images").Array())
|
||||
imagePayload := `{"type":"image_url","image_url":{"url":""}}`
|
||||
imagePayload, _ = sjson.Set(imagePayload, "index", imageIndex)
|
||||
imagePayload, _ = sjson.Set(imagePayload, "image_url.url", imageURL)
|
||||
template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", string(imagePayload))
|
||||
template, _ = sjson.SetRaw(template, "choices.0.delta.images.-1", imagePayload)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -248,10 +252,19 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
}
|
||||
promptTokenCount := usageResult.Get("promptTokenCount").Int()
|
||||
thoughtsTokenCount := usageResult.Get("thoughtsTokenCount").Int()
|
||||
cachedTokenCount := usageResult.Get("cachedContentTokenCount").Int()
|
||||
template, _ = sjson.Set(template, "usage.prompt_tokens", promptTokenCount+thoughtsTokenCount)
|
||||
if thoughtsTokenCount > 0 {
|
||||
template, _ = sjson.Set(template, "usage.completion_tokens_details.reasoning_tokens", thoughtsTokenCount)
|
||||
}
|
||||
// Include cached token count if present (indicates prompt caching is working)
|
||||
if cachedTokenCount > 0 {
|
||||
var err error
|
||||
template, err = sjson.Set(template, "usage.prompt_tokens_details.cached_tokens", cachedTokenCount)
|
||||
if err != nil {
|
||||
log.Warnf("gemini openai response: failed to set cached_tokens in non-streaming: %v", err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Process the main content part of the response.
|
||||
@@ -305,21 +318,16 @@ func ConvertGeminiResponseToOpenAINonStream(_ context.Context, _ string, origina
|
||||
mimeType = "image/png"
|
||||
}
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
imagePayload, err := json.Marshal(map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]string{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
imagesResult := gjson.Get(template, "choices.0.message.images")
|
||||
if !imagesResult.Exists() || !imagesResult.IsArray() {
|
||||
template, _ = sjson.SetRaw(template, "choices.0.message.images", `[]`)
|
||||
}
|
||||
imageIndex := len(gjson.Get(template, "choices.0.message.images").Array())
|
||||
imagePayload := `{"type":"image_url","image_url":{"url":""}}`
|
||||
imagePayload, _ = sjson.Set(imagePayload, "index", imageIndex)
|
||||
imagePayload, _ = sjson.Set(imagePayload, "image_url.url", imageURL)
|
||||
template, _ = sjson.Set(template, "choices.0.message.role", "assistant")
|
||||
template, _ = sjson.SetRaw(template, "choices.0.message.images.-1", string(imagePayload))
|
||||
template, _ = sjson.SetRaw(template, "choices.0.message.images.-1", imagePayload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -377,27 +377,18 @@ func ConvertGeminiResponseToOpenAIResponses(_ context.Context, modelName string,
|
||||
}
|
||||
|
||||
// Compose outputs in encountered order: reasoning, message, function_calls
|
||||
var outputs []interface{}
|
||||
outputsWrapper := `{"arr":[]}`
|
||||
if st.ReasoningOpened {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": st.ReasoningItemID,
|
||||
"type": "reasoning",
|
||||
"summary": []interface{}{map[string]interface{}{"type": "summary_text", "text": st.ReasoningBuf.String()}},
|
||||
})
|
||||
item := `{"id":"","type":"reasoning","summary":[{"type":"summary_text","text":""}]}`
|
||||
item, _ = sjson.Set(item, "id", st.ReasoningItemID)
|
||||
item, _ = sjson.Set(item, "summary.0.text", st.ReasoningBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
if st.MsgOpened {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": st.CurrentMsgID,
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": st.TextBuf.String(),
|
||||
}},
|
||||
"role": "assistant",
|
||||
})
|
||||
item := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
item, _ = sjson.Set(item, "id", st.CurrentMsgID)
|
||||
item, _ = sjson.Set(item, "content.0.text", st.TextBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
if len(st.FuncArgsBuf) > 0 {
|
||||
idxs := make([]int, 0, len(st.FuncArgsBuf))
|
||||
@@ -416,18 +407,16 @@ func ConvertGeminiResponseToOpenAIResponses(_ context.Context, modelName string,
|
||||
if b := st.FuncArgsBuf[idx]; b != nil {
|
||||
args = b.String()
|
||||
}
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", st.FuncCallIDs[idx]),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": args,
|
||||
"call_id": st.FuncCallIDs[idx],
|
||||
"name": st.FuncNames[idx],
|
||||
})
|
||||
item := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("fc_%s", st.FuncCallIDs[idx]))
|
||||
item, _ = sjson.Set(item, "arguments", args)
|
||||
item, _ = sjson.Set(item, "call_id", st.FuncCallIDs[idx])
|
||||
item, _ = sjson.Set(item, "name", st.FuncNames[idx])
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
}
|
||||
if len(outputs) > 0 {
|
||||
completed, _ = sjson.Set(completed, "response.output", outputs)
|
||||
if gjson.Get(outputsWrapper, "arr.#").Int() > 0 {
|
||||
completed, _ = sjson.SetRaw(completed, "response.output", gjson.Get(outputsWrapper, "arr").Raw)
|
||||
}
|
||||
|
||||
// usage mapping
|
||||
@@ -558,11 +547,24 @@ func ConvertGeminiResponseToOpenAIResponsesNonStream(_ context.Context, _ string
|
||||
}
|
||||
|
||||
// Build outputs from candidates[0].content.parts
|
||||
var outputs []interface{}
|
||||
var reasoningText strings.Builder
|
||||
var reasoningEncrypted string
|
||||
var messageText strings.Builder
|
||||
var haveMessage bool
|
||||
|
||||
haveOutput := false
|
||||
ensureOutput := func() {
|
||||
if haveOutput {
|
||||
return
|
||||
}
|
||||
resp, _ = sjson.SetRaw(resp, "output", "[]")
|
||||
haveOutput = true
|
||||
}
|
||||
appendOutput := func(itemJSON string) {
|
||||
ensureOutput()
|
||||
resp, _ = sjson.SetRaw(resp, "output.-1", itemJSON)
|
||||
}
|
||||
|
||||
if parts := root.Get("candidates.0.content.parts"); parts.Exists() && parts.IsArray() {
|
||||
parts.ForEach(func(_, p gjson.Result) bool {
|
||||
if p.Get("thought").Bool() {
|
||||
@@ -583,19 +585,16 @@ func ConvertGeminiResponseToOpenAIResponsesNonStream(_ context.Context, _ string
|
||||
name := fc.Get("name").String()
|
||||
args := fc.Get("args")
|
||||
callID := fmt.Sprintf("call_%x_%d", time.Now().UnixNano(), atomic.AddUint64(&funcCallIDCounter, 1))
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", callID),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": func() string {
|
||||
if args.Exists() {
|
||||
return args.Raw
|
||||
}
|
||||
return ""
|
||||
}(),
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
})
|
||||
itemJSON := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
itemJSON, _ = sjson.Set(itemJSON, "id", fmt.Sprintf("fc_%s", callID))
|
||||
itemJSON, _ = sjson.Set(itemJSON, "call_id", callID)
|
||||
itemJSON, _ = sjson.Set(itemJSON, "name", name)
|
||||
argsStr := ""
|
||||
if args.Exists() {
|
||||
argsStr = args.Raw
|
||||
}
|
||||
itemJSON, _ = sjson.Set(itemJSON, "arguments", argsStr)
|
||||
appendOutput(itemJSON)
|
||||
return true
|
||||
}
|
||||
return true
|
||||
@@ -605,42 +604,24 @@ func ConvertGeminiResponseToOpenAIResponsesNonStream(_ context.Context, _ string
|
||||
// Reasoning output item
|
||||
if reasoningText.Len() > 0 || reasoningEncrypted != "" {
|
||||
rid := strings.TrimPrefix(id, "resp_")
|
||||
item := map[string]interface{}{
|
||||
"id": fmt.Sprintf("rs_%s", rid),
|
||||
"type": "reasoning",
|
||||
"encrypted_content": reasoningEncrypted,
|
||||
}
|
||||
var summaries []interface{}
|
||||
itemJSON := `{"id":"","type":"reasoning","encrypted_content":""}`
|
||||
itemJSON, _ = sjson.Set(itemJSON, "id", fmt.Sprintf("rs_%s", rid))
|
||||
itemJSON, _ = sjson.Set(itemJSON, "encrypted_content", reasoningEncrypted)
|
||||
if reasoningText.Len() > 0 {
|
||||
summaries = append(summaries, map[string]interface{}{
|
||||
"type": "summary_text",
|
||||
"text": reasoningText.String(),
|
||||
})
|
||||
summaryJSON := `{"type":"summary_text","text":""}`
|
||||
summaryJSON, _ = sjson.Set(summaryJSON, "text", reasoningText.String())
|
||||
itemJSON, _ = sjson.SetRaw(itemJSON, "summary", "[]")
|
||||
itemJSON, _ = sjson.SetRaw(itemJSON, "summary.-1", summaryJSON)
|
||||
}
|
||||
if summaries != nil {
|
||||
item["summary"] = summaries
|
||||
}
|
||||
outputs = append(outputs, item)
|
||||
appendOutput(itemJSON)
|
||||
}
|
||||
|
||||
// Assistant message output item
|
||||
if haveMessage {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("msg_%s_0", strings.TrimPrefix(id, "resp_")),
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": messageText.String(),
|
||||
}},
|
||||
"role": "assistant",
|
||||
})
|
||||
}
|
||||
|
||||
if len(outputs) > 0 {
|
||||
resp, _ = sjson.Set(resp, "output", outputs)
|
||||
itemJSON := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
itemJSON, _ = sjson.Set(itemJSON, "id", fmt.Sprintf("msg_%s_0", strings.TrimPrefix(id, "resp_")))
|
||||
itemJSON, _ = sjson.Set(itemJSON, "content.0.text", messageText.String())
|
||||
appendOutput(itemJSON)
|
||||
}
|
||||
|
||||
// usage mapping
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
// MergeAdjacentMessages merges adjacent messages with the same role.
|
||||
// This reduces API call complexity and improves compatibility.
|
||||
// Based on AIClient-2-API implementation.
|
||||
// NOTE: Tool messages are NOT merged because each has a unique tool_call_id that must be preserved.
|
||||
func MergeAdjacentMessages(messages []gjson.Result) []gjson.Result {
|
||||
if len(messages) <= 1 {
|
||||
return messages
|
||||
@@ -26,6 +27,12 @@ func MergeAdjacentMessages(messages []gjson.Result) []gjson.Result {
|
||||
currentRole := msg.Get("role").String()
|
||||
lastRole := lastMsg.Get("role").String()
|
||||
|
||||
// Don't merge tool messages - each has a unique tool_call_id
|
||||
if currentRole == "tool" || lastRole == "tool" {
|
||||
merged = append(merged, msg)
|
||||
continue
|
||||
}
|
||||
|
||||
if currentRole == lastRole {
|
||||
// Merge content from current message into last message
|
||||
mergedContent := mergeMessageContent(lastMsg, msg)
|
||||
|
||||
@@ -450,24 +450,10 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
|
||||
// Merge adjacent messages with the same role
|
||||
messagesArray := kirocommon.MergeAdjacentMessages(messages.Array())
|
||||
|
||||
// Build tool_call_id to name mapping from assistant messages
|
||||
toolCallIDToName := make(map[string]string)
|
||||
for _, msg := range messagesArray {
|
||||
if msg.Get("role").String() == "assistant" {
|
||||
toolCalls := msg.Get("tool_calls")
|
||||
if toolCalls.IsArray() {
|
||||
for _, tc := range toolCalls.Array() {
|
||||
if tc.Get("type").String() == "function" {
|
||||
id := tc.Get("id").String()
|
||||
name := tc.Get("function.name").String()
|
||||
if id != "" && name != "" {
|
||||
toolCallIDToName[id] = name
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
// Track pending tool results that should be attached to the next user message
|
||||
// This is critical for LiteLLM-translated requests where tool results appear
|
||||
// as separate "tool" role messages between assistant and user messages
|
||||
var pendingToolResults []KiroToolResult
|
||||
|
||||
for i, msg := range messagesArray {
|
||||
role := msg.Get("role").String()
|
||||
@@ -480,6 +466,10 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
|
||||
|
||||
case "user":
|
||||
userMsg, toolResults := buildUserMessageFromOpenAI(msg, modelID, origin)
|
||||
// Merge any pending tool results from preceding "tool" role messages
|
||||
toolResults = append(pendingToolResults, toolResults...)
|
||||
pendingToolResults = nil // Reset pending tool results
|
||||
|
||||
if isLastMessage {
|
||||
currentUserMsg = &userMsg
|
||||
currentToolResults = toolResults
|
||||
@@ -505,6 +495,24 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
|
||||
|
||||
case "assistant":
|
||||
assistantMsg := buildAssistantMessageFromOpenAI(msg)
|
||||
|
||||
// If there are pending tool results, we need to insert a synthetic user message
|
||||
// before this assistant message to maintain proper conversation structure
|
||||
if len(pendingToolResults) > 0 {
|
||||
syntheticUserMsg := KiroUserInputMessage{
|
||||
Content: "Tool results provided.",
|
||||
ModelID: modelID,
|
||||
Origin: origin,
|
||||
UserInputMessageContext: &KiroUserInputMessageContext{
|
||||
ToolResults: pendingToolResults,
|
||||
},
|
||||
}
|
||||
history = append(history, KiroHistoryMessage{
|
||||
UserInputMessage: &syntheticUserMsg,
|
||||
})
|
||||
pendingToolResults = nil
|
||||
}
|
||||
|
||||
if isLastMessage {
|
||||
history = append(history, KiroHistoryMessage{
|
||||
AssistantResponseMessage: &assistantMsg,
|
||||
@@ -524,7 +532,7 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
|
||||
case "tool":
|
||||
// Tool messages in OpenAI format provide results for tool_calls
|
||||
// These are typically followed by user or assistant messages
|
||||
// Process them and merge into the next user message's tool results
|
||||
// Collect them as pending and attach to the next user message
|
||||
toolCallID := msg.Get("tool_call_id").String()
|
||||
content := msg.Get("content").String()
|
||||
|
||||
@@ -534,9 +542,21 @@ func processOpenAIMessages(messages gjson.Result, modelID, origin string) ([]Kir
|
||||
Content: []KiroTextContent{{Text: content}},
|
||||
Status: "success",
|
||||
}
|
||||
// Tool results should be included in the next user message
|
||||
// For now, collect them and they'll be handled when we build the current message
|
||||
currentToolResults = append(currentToolResults, toolResult)
|
||||
// Collect pending tool results to attach to the next user message
|
||||
pendingToolResults = append(pendingToolResults, toolResult)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Handle case where tool results are at the end with no following user message
|
||||
if len(pendingToolResults) > 0 {
|
||||
currentToolResults = append(currentToolResults, pendingToolResults...)
|
||||
// If there's no current user message, create a synthetic one for the tool results
|
||||
if currentUserMsg == nil {
|
||||
currentUserMsg = &KiroUserInputMessage{
|
||||
Content: "Tool results provided.",
|
||||
ModelID: modelID,
|
||||
Origin: origin,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -551,9 +571,6 @@ func buildUserMessageFromOpenAI(msg gjson.Result, modelID, origin string) (KiroU
|
||||
var toolResults []KiroToolResult
|
||||
var images []KiroImage
|
||||
|
||||
// Track seen toolCallIds to deduplicate
|
||||
seenToolCallIDs := make(map[string]bool)
|
||||
|
||||
if content.IsArray() {
|
||||
for _, part := range content.Array() {
|
||||
partType := part.Get("type").String()
|
||||
@@ -589,9 +606,6 @@ func buildUserMessageFromOpenAI(msg gjson.Result, modelID, origin string) (KiroU
|
||||
contentBuilder.WriteString(content.String())
|
||||
}
|
||||
|
||||
// Check for tool_calls in the message (shouldn't be in user messages, but handle edge cases)
|
||||
_ = seenToolCallIDs // Used for deduplication if needed
|
||||
|
||||
userMsg := KiroUserInputMessage{
|
||||
Content: contentBuilder.String(),
|
||||
ModelID: modelID,
|
||||
|
||||
386
internal/translator/kiro/openai/kiro_openai_request_test.go
Normal file
386
internal/translator/kiro/openai/kiro_openai_request_test.go
Normal file
@@ -0,0 +1,386 @@
|
||||
package openai
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestToolResultsAttachedToCurrentMessage verifies that tool results from "tool" role messages
|
||||
// are properly attached to the current user message (the last message in the conversation).
|
||||
// This is critical for LiteLLM-translated requests where tool results appear as separate messages.
|
||||
func TestToolResultsAttachedToCurrentMessage(t *testing.T) {
|
||||
// OpenAI format request simulating LiteLLM's translation from Anthropic format
|
||||
// Sequence: user -> assistant (with tool_calls) -> tool (result) -> user
|
||||
// The last user message should have the tool results attached
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello, can you read a file for me?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I'll read that file for you.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_abc123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/test.txt\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_abc123",
|
||||
"content": "File contents: Hello World!"
|
||||
},
|
||||
{"role": "user", "content": "What did the file say?"}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
// The last user message becomes currentMessage
|
||||
// History should have: user (first), assistant (with tool_calls)
|
||||
t.Logf("History count: %d", len(payload.ConversationState.History))
|
||||
if len(payload.ConversationState.History) != 2 {
|
||||
t.Errorf("Expected 2 history entries (user + assistant), got %d", len(payload.ConversationState.History))
|
||||
}
|
||||
|
||||
// Tool results should be attached to currentMessage (the last user message)
|
||||
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
|
||||
if ctx == nil {
|
||||
t.Fatal("Expected currentMessage to have UserInputMessageContext with tool results")
|
||||
}
|
||||
|
||||
if len(ctx.ToolResults) != 1 {
|
||||
t.Fatalf("Expected 1 tool result in currentMessage, got %d", len(ctx.ToolResults))
|
||||
}
|
||||
|
||||
tr := ctx.ToolResults[0]
|
||||
if tr.ToolUseID != "call_abc123" {
|
||||
t.Errorf("Expected toolUseId 'call_abc123', got '%s'", tr.ToolUseID)
|
||||
}
|
||||
if len(tr.Content) == 0 || tr.Content[0].Text != "File contents: Hello World!" {
|
||||
t.Errorf("Tool result content mismatch, got: %+v", tr.Content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolResultsInHistoryUserMessage verifies that when there are multiple user messages
|
||||
// after tool results, the tool results are attached to the correct user message in history.
|
||||
func TestToolResultsInHistoryUserMessage(t *testing.T) {
|
||||
// Sequence: user -> assistant (with tool_calls) -> tool (result) -> user -> assistant -> user
|
||||
// The first user after tool should have tool results in history
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I'll read the file.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{}"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": "File result"
|
||||
},
|
||||
{"role": "user", "content": "Thanks for the file"},
|
||||
{"role": "assistant", "content": "You're welcome"},
|
||||
{"role": "user", "content": "Bye"}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
// History should have: user, assistant, user (with tool results), assistant
|
||||
// CurrentMessage should be: last user "Bye"
|
||||
t.Logf("History count: %d", len(payload.ConversationState.History))
|
||||
|
||||
// Find the user message in history with tool results
|
||||
foundToolResults := false
|
||||
for i, h := range payload.ConversationState.History {
|
||||
if h.UserInputMessage != nil {
|
||||
t.Logf("History[%d]: user message content=%q", i, h.UserInputMessage.Content)
|
||||
if h.UserInputMessage.UserInputMessageContext != nil {
|
||||
if len(h.UserInputMessage.UserInputMessageContext.ToolResults) > 0 {
|
||||
foundToolResults = true
|
||||
t.Logf(" Found %d tool results", len(h.UserInputMessage.UserInputMessageContext.ToolResults))
|
||||
tr := h.UserInputMessage.UserInputMessageContext.ToolResults[0]
|
||||
if tr.ToolUseID != "call_1" {
|
||||
t.Errorf("Expected toolUseId 'call_1', got '%s'", tr.ToolUseID)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
if h.AssistantResponseMessage != nil {
|
||||
t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content)
|
||||
}
|
||||
}
|
||||
|
||||
if !foundToolResults {
|
||||
t.Error("Tool results were not attached to any user message in history")
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolResultsWithMultipleToolCalls verifies handling of multiple tool calls
|
||||
func TestToolResultsWithMultipleToolCalls(t *testing.T) {
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Read two files for me"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I'll read both files.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/file1.txt\"}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "call_2",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/file2.txt\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": "Content of file 1"
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_2",
|
||||
"content": "Content of file 2"
|
||||
},
|
||||
{"role": "user", "content": "What do they say?"}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
t.Logf("History count: %d", len(payload.ConversationState.History))
|
||||
t.Logf("CurrentMessage content: %q", payload.ConversationState.CurrentMessage.UserInputMessage.Content)
|
||||
|
||||
// Check if there are any tool results anywhere
|
||||
var totalToolResults int
|
||||
for i, h := range payload.ConversationState.History {
|
||||
if h.UserInputMessage != nil && h.UserInputMessage.UserInputMessageContext != nil {
|
||||
count := len(h.UserInputMessage.UserInputMessageContext.ToolResults)
|
||||
t.Logf("History[%d] user message has %d tool results", i, count)
|
||||
totalToolResults += count
|
||||
}
|
||||
}
|
||||
|
||||
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
|
||||
if ctx != nil {
|
||||
t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults))
|
||||
totalToolResults += len(ctx.ToolResults)
|
||||
} else {
|
||||
t.Logf("CurrentMessage has no UserInputMessageContext")
|
||||
}
|
||||
|
||||
if totalToolResults != 2 {
|
||||
t.Errorf("Expected 2 tool results total, got %d", totalToolResults)
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolResultsAtEndOfConversation verifies tool results are handled when
|
||||
// the conversation ends with tool results (no following user message)
|
||||
func TestToolResultsAtEndOfConversation(t *testing.T) {
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Read a file"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Reading the file.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_end",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/test.txt\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_end",
|
||||
"content": "File contents here"
|
||||
}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
// When the last message is a tool result, a synthetic user message is created
|
||||
// and tool results should be attached to it
|
||||
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
|
||||
if ctx == nil || len(ctx.ToolResults) == 0 {
|
||||
t.Error("Expected tool results to be attached to current message when conversation ends with tool result")
|
||||
} else {
|
||||
if ctx.ToolResults[0].ToolUseID != "call_end" {
|
||||
t.Errorf("Expected toolUseId 'call_end', got '%s'", ctx.ToolResults[0].ToolUseID)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestToolResultsFollowedByAssistant verifies handling when tool results are followed
|
||||
// by an assistant message (no intermediate user message).
|
||||
// This is the pattern from LiteLLM translation of Anthropic format where:
|
||||
// user message has ONLY tool_result blocks -> LiteLLM creates tool messages
|
||||
// then the next message is assistant
|
||||
func TestToolResultsFollowedByAssistant(t *testing.T) {
|
||||
// Sequence: user -> assistant (with tool_calls) -> tool -> tool -> assistant -> user
|
||||
// This simulates LiteLLM's translation of:
|
||||
// user: "Read files"
|
||||
// assistant: [tool_use, tool_use]
|
||||
// user: [tool_result, tool_result] <- becomes multiple "tool" role messages
|
||||
// assistant: "I've read them"
|
||||
// user: "What did they say?"
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Read two files for me"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I'll read both files.",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/a.txt\"}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": "call_2",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Read",
|
||||
"arguments": "{\"file_path\": \"/tmp/b.txt\"}"
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_1",
|
||||
"content": "Contents of file A"
|
||||
},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_2",
|
||||
"content": "Contents of file B"
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I've read both files."
|
||||
},
|
||||
{"role": "user", "content": "What did they say?"}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
t.Logf("History count: %d", len(payload.ConversationState.History))
|
||||
|
||||
// Tool results should be attached to a synthetic user message or the history should be valid
|
||||
var totalToolResults int
|
||||
for i, h := range payload.ConversationState.History {
|
||||
if h.UserInputMessage != nil {
|
||||
t.Logf("History[%d]: user message content=%q", i, h.UserInputMessage.Content)
|
||||
if h.UserInputMessage.UserInputMessageContext != nil {
|
||||
count := len(h.UserInputMessage.UserInputMessageContext.ToolResults)
|
||||
t.Logf(" Has %d tool results", count)
|
||||
totalToolResults += count
|
||||
}
|
||||
}
|
||||
if h.AssistantResponseMessage != nil {
|
||||
t.Logf("History[%d]: assistant message content=%q", i, h.AssistantResponseMessage.Content)
|
||||
}
|
||||
}
|
||||
|
||||
ctx := payload.ConversationState.CurrentMessage.UserInputMessage.UserInputMessageContext
|
||||
if ctx != nil {
|
||||
t.Logf("CurrentMessage has %d tool results", len(ctx.ToolResults))
|
||||
totalToolResults += len(ctx.ToolResults)
|
||||
}
|
||||
|
||||
if totalToolResults != 2 {
|
||||
t.Errorf("Expected 2 tool results total, got %d", totalToolResults)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAssistantEndsConversation verifies handling when assistant is the last message
|
||||
func TestAssistantEndsConversation(t *testing.T) {
|
||||
input := []byte(`{
|
||||
"model": "kiro-claude-opus-4-5-agentic",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hi there!"
|
||||
}
|
||||
]
|
||||
}`)
|
||||
|
||||
result, _ := BuildKiroPayloadFromOpenAI(input, "kiro-model", "", "CLI", false, false, nil, nil)
|
||||
|
||||
var payload KiroPayload
|
||||
if err := json.Unmarshal(result, &payload); err != nil {
|
||||
t.Fatalf("Failed to unmarshal result: %v", err)
|
||||
}
|
||||
|
||||
// When assistant is last, a "Continue" user message should be created
|
||||
if payload.ConversationState.CurrentMessage.UserInputMessage.Content == "" {
|
||||
t.Error("Expected a 'Continue' message to be created when assistant is last")
|
||||
}
|
||||
}
|
||||
@@ -7,7 +7,6 @@ package claude
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
@@ -138,11 +137,7 @@ func ConvertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
|
||||
// Convert input to arguments JSON string
|
||||
if input := part.Get("input"); input.Exists() {
|
||||
if inputJSON, err := json.Marshal(input.Value()); err == nil {
|
||||
toolCallJSON, _ = sjson.Set(toolCallJSON, "function.arguments", string(inputJSON))
|
||||
} else {
|
||||
toolCallJSON, _ = sjson.Set(toolCallJSON, "function.arguments", "{}")
|
||||
}
|
||||
toolCallJSON, _ = sjson.Set(toolCallJSON, "function.arguments", input.Raw)
|
||||
} else {
|
||||
toolCallJSON, _ = sjson.Set(toolCallJSON, "function.arguments", "{}")
|
||||
}
|
||||
@@ -191,8 +186,7 @@ func ConvertClaudeRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
// Emit tool calls in a separate assistant message
|
||||
if role == "assistant" && len(toolCalls) > 0 {
|
||||
toolCallMsgJSON := `{"role":"assistant","tool_calls":[]}`
|
||||
toolCallsJSON, _ := json.Marshal(toolCalls)
|
||||
toolCallMsgJSON, _ = sjson.SetRaw(toolCallMsgJSON, "tool_calls", string(toolCallsJSON))
|
||||
toolCallMsgJSON, _ = sjson.Set(toolCallMsgJSON, "tool_calls", toolCalls)
|
||||
messagesJSON, _ = sjson.Set(messagesJSON, "-1", gjson.Parse(toolCallMsgJSON).Value())
|
||||
}
|
||||
|
||||
|
||||
@@ -8,7 +8,6 @@ package claude
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
@@ -133,24 +132,10 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
if delta := root.Get("choices.0.delta"); delta.Exists() {
|
||||
if !param.MessageStarted {
|
||||
// Send message_start event
|
||||
messageStart := map[string]interface{}{
|
||||
"type": "message_start",
|
||||
"message": map[string]interface{}{
|
||||
"id": param.MessageID,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": param.Model,
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
},
|
||||
}
|
||||
messageStartJSON, _ := json.Marshal(messageStart)
|
||||
results = append(results, "event: message_start\ndata: "+string(messageStartJSON)+"\n\n")
|
||||
messageStartJSON := `{"type":"message_start","message":{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}}`
|
||||
messageStartJSON, _ = sjson.Set(messageStartJSON, "message.id", param.MessageID)
|
||||
messageStartJSON, _ = sjson.Set(messageStartJSON, "message.model", param.Model)
|
||||
results = append(results, "event: message_start\ndata: "+messageStartJSON+"\n\n")
|
||||
param.MessageStarted = true
|
||||
|
||||
// Don't send content_block_start for text here - wait for actual content
|
||||
@@ -168,29 +153,16 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
param.ThinkingContentBlockIndex = param.NextContentBlockIndex
|
||||
param.NextContentBlockIndex++
|
||||
}
|
||||
contentBlockStart := map[string]interface{}{
|
||||
"type": "content_block_start",
|
||||
"index": param.ThinkingContentBlockIndex,
|
||||
"content_block": map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": "",
|
||||
},
|
||||
}
|
||||
contentBlockStartJSON, _ := json.Marshal(contentBlockStart)
|
||||
results = append(results, "event: content_block_start\ndata: "+string(contentBlockStartJSON)+"\n\n")
|
||||
contentBlockStartJSON := `{"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}`
|
||||
contentBlockStartJSON, _ = sjson.Set(contentBlockStartJSON, "index", param.ThinkingContentBlockIndex)
|
||||
results = append(results, "event: content_block_start\ndata: "+contentBlockStartJSON+"\n\n")
|
||||
param.ThinkingContentBlockStarted = true
|
||||
}
|
||||
|
||||
thinkingDelta := map[string]interface{}{
|
||||
"type": "content_block_delta",
|
||||
"index": param.ThinkingContentBlockIndex,
|
||||
"delta": map[string]interface{}{
|
||||
"type": "thinking_delta",
|
||||
"thinking": reasoningText,
|
||||
},
|
||||
}
|
||||
thinkingDeltaJSON, _ := json.Marshal(thinkingDelta)
|
||||
results = append(results, "event: content_block_delta\ndata: "+string(thinkingDeltaJSON)+"\n\n")
|
||||
thinkingDeltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":""}}`
|
||||
thinkingDeltaJSON, _ = sjson.Set(thinkingDeltaJSON, "index", param.ThinkingContentBlockIndex)
|
||||
thinkingDeltaJSON, _ = sjson.Set(thinkingDeltaJSON, "delta.thinking", reasoningText)
|
||||
results = append(results, "event: content_block_delta\ndata: "+thinkingDeltaJSON+"\n\n")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -203,29 +175,16 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
param.TextContentBlockIndex = param.NextContentBlockIndex
|
||||
param.NextContentBlockIndex++
|
||||
}
|
||||
contentBlockStart := map[string]interface{}{
|
||||
"type": "content_block_start",
|
||||
"index": param.TextContentBlockIndex,
|
||||
"content_block": map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": "",
|
||||
},
|
||||
}
|
||||
contentBlockStartJSON, _ := json.Marshal(contentBlockStart)
|
||||
results = append(results, "event: content_block_start\ndata: "+string(contentBlockStartJSON)+"\n\n")
|
||||
contentBlockStartJSON := `{"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`
|
||||
contentBlockStartJSON, _ = sjson.Set(contentBlockStartJSON, "index", param.TextContentBlockIndex)
|
||||
results = append(results, "event: content_block_start\ndata: "+contentBlockStartJSON+"\n\n")
|
||||
param.TextContentBlockStarted = true
|
||||
}
|
||||
|
||||
contentDelta := map[string]interface{}{
|
||||
"type": "content_block_delta",
|
||||
"index": param.TextContentBlockIndex,
|
||||
"delta": map[string]interface{}{
|
||||
"type": "text_delta",
|
||||
"text": content.String(),
|
||||
},
|
||||
}
|
||||
contentDeltaJSON, _ := json.Marshal(contentDelta)
|
||||
results = append(results, "event: content_block_delta\ndata: "+string(contentDeltaJSON)+"\n\n")
|
||||
contentDeltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":""}}`
|
||||
contentDeltaJSON, _ = sjson.Set(contentDeltaJSON, "index", param.TextContentBlockIndex)
|
||||
contentDeltaJSON, _ = sjson.Set(contentDeltaJSON, "delta.text", content.String())
|
||||
results = append(results, "event: content_block_delta\ndata: "+contentDeltaJSON+"\n\n")
|
||||
|
||||
// Accumulate content
|
||||
param.ContentAccumulator.WriteString(content.String())
|
||||
@@ -263,18 +222,11 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
stopTextContentBlock(param, &results)
|
||||
|
||||
// Send content_block_start for tool_use
|
||||
contentBlockStart := map[string]interface{}{
|
||||
"type": "content_block_start",
|
||||
"index": blockIndex,
|
||||
"content_block": map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": accumulator.ID,
|
||||
"name": accumulator.Name,
|
||||
"input": map[string]interface{}{},
|
||||
},
|
||||
}
|
||||
contentBlockStartJSON, _ := json.Marshal(contentBlockStart)
|
||||
results = append(results, "event: content_block_start\ndata: "+string(contentBlockStartJSON)+"\n\n")
|
||||
contentBlockStartJSON := `{"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"","name":"","input":{}}}`
|
||||
contentBlockStartJSON, _ = sjson.Set(contentBlockStartJSON, "index", blockIndex)
|
||||
contentBlockStartJSON, _ = sjson.Set(contentBlockStartJSON, "content_block.id", accumulator.ID)
|
||||
contentBlockStartJSON, _ = sjson.Set(contentBlockStartJSON, "content_block.name", accumulator.Name)
|
||||
results = append(results, "event: content_block_start\ndata: "+contentBlockStartJSON+"\n\n")
|
||||
}
|
||||
|
||||
// Handle function arguments
|
||||
@@ -298,12 +250,9 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
|
||||
// Send content_block_stop for thinking content if needed
|
||||
if param.ThinkingContentBlockStarted {
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": param.ThinkingContentBlockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
results = append(results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", param.ThinkingContentBlockIndex)
|
||||
results = append(results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
param.ThinkingContentBlockStarted = false
|
||||
param.ThinkingContentBlockIndex = -1
|
||||
}
|
||||
@@ -319,24 +268,15 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
|
||||
// Send complete input_json_delta with all accumulated arguments
|
||||
if accumulator.Arguments.Len() > 0 {
|
||||
inputDelta := map[string]interface{}{
|
||||
"type": "content_block_delta",
|
||||
"index": blockIndex,
|
||||
"delta": map[string]interface{}{
|
||||
"type": "input_json_delta",
|
||||
"partial_json": util.FixJSON(accumulator.Arguments.String()),
|
||||
},
|
||||
}
|
||||
inputDeltaJSON, _ := json.Marshal(inputDelta)
|
||||
results = append(results, "event: content_block_delta\ndata: "+string(inputDeltaJSON)+"\n\n")
|
||||
inputDeltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":""}}`
|
||||
inputDeltaJSON, _ = sjson.Set(inputDeltaJSON, "index", blockIndex)
|
||||
inputDeltaJSON, _ = sjson.Set(inputDeltaJSON, "delta.partial_json", util.FixJSON(accumulator.Arguments.String()))
|
||||
results = append(results, "event: content_block_delta\ndata: "+inputDeltaJSON+"\n\n")
|
||||
}
|
||||
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": blockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
results = append(results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", blockIndex)
|
||||
results = append(results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
delete(param.ToolCallBlockIndexes, index)
|
||||
}
|
||||
param.ContentBlocksStopped = true
|
||||
@@ -361,20 +301,11 @@ func convertOpenAIStreamingChunkToAnthropic(rawJSON []byte, param *ConvertOpenAI
|
||||
}
|
||||
}
|
||||
// Send message_delta with usage
|
||||
messageDelta := map[string]interface{}{
|
||||
"type": "message_delta",
|
||||
"delta": map[string]interface{}{
|
||||
"stop_reason": mapOpenAIFinishReasonToAnthropic(param.FinishReason),
|
||||
"stop_sequence": nil,
|
||||
},
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": inputTokens,
|
||||
"output_tokens": outputTokens,
|
||||
},
|
||||
}
|
||||
|
||||
messageDeltaJSON, _ := json.Marshal(messageDelta)
|
||||
results = append(results, "event: message_delta\ndata: "+string(messageDeltaJSON)+"\n\n")
|
||||
messageDeltaJSON := `{"type":"message_delta","delta":{"stop_reason":"","stop_sequence":null},"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
messageDeltaJSON, _ = sjson.Set(messageDeltaJSON, "delta.stop_reason", mapOpenAIFinishReasonToAnthropic(param.FinishReason))
|
||||
messageDeltaJSON, _ = sjson.Set(messageDeltaJSON, "usage.input_tokens", inputTokens)
|
||||
messageDeltaJSON, _ = sjson.Set(messageDeltaJSON, "usage.output_tokens", outputTokens)
|
||||
results = append(results, "event: message_delta\ndata: "+messageDeltaJSON+"\n\n")
|
||||
param.MessageDeltaSent = true
|
||||
|
||||
emitMessageStopIfNeeded(param, &results)
|
||||
@@ -390,12 +321,9 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams)
|
||||
|
||||
// Ensure all content blocks are stopped before final events
|
||||
if param.ThinkingContentBlockStarted {
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": param.ThinkingContentBlockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
results = append(results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", param.ThinkingContentBlockIndex)
|
||||
results = append(results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
param.ThinkingContentBlockStarted = false
|
||||
param.ThinkingContentBlockIndex = -1
|
||||
}
|
||||
@@ -408,24 +336,15 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams)
|
||||
blockIndex := param.toolContentBlockIndex(index)
|
||||
|
||||
if accumulator.Arguments.Len() > 0 {
|
||||
inputDelta := map[string]interface{}{
|
||||
"type": "content_block_delta",
|
||||
"index": blockIndex,
|
||||
"delta": map[string]interface{}{
|
||||
"type": "input_json_delta",
|
||||
"partial_json": util.FixJSON(accumulator.Arguments.String()),
|
||||
},
|
||||
}
|
||||
inputDeltaJSON, _ := json.Marshal(inputDelta)
|
||||
results = append(results, "event: content_block_delta\ndata: "+string(inputDeltaJSON)+"\n\n")
|
||||
inputDeltaJSON := `{"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":""}}`
|
||||
inputDeltaJSON, _ = sjson.Set(inputDeltaJSON, "index", blockIndex)
|
||||
inputDeltaJSON, _ = sjson.Set(inputDeltaJSON, "delta.partial_json", util.FixJSON(accumulator.Arguments.String()))
|
||||
results = append(results, "event: content_block_delta\ndata: "+inputDeltaJSON+"\n\n")
|
||||
}
|
||||
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": blockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
results = append(results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", blockIndex)
|
||||
results = append(results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
delete(param.ToolCallBlockIndexes, index)
|
||||
}
|
||||
param.ContentBlocksStopped = true
|
||||
@@ -433,16 +352,9 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams)
|
||||
|
||||
// If we haven't sent message_delta yet (no usage info was received), send it now
|
||||
if param.FinishReason != "" && !param.MessageDeltaSent {
|
||||
messageDelta := map[string]interface{}{
|
||||
"type": "message_delta",
|
||||
"delta": map[string]interface{}{
|
||||
"stop_reason": mapOpenAIFinishReasonToAnthropic(param.FinishReason),
|
||||
"stop_sequence": nil,
|
||||
},
|
||||
}
|
||||
|
||||
messageDeltaJSON, _ := json.Marshal(messageDelta)
|
||||
results = append(results, "event: message_delta\ndata: "+string(messageDeltaJSON)+"\n\n")
|
||||
messageDeltaJSON := `{"type":"message_delta","delta":{"stop_reason":"","stop_sequence":null}}`
|
||||
messageDeltaJSON, _ = sjson.Set(messageDeltaJSON, "delta.stop_reason", mapOpenAIFinishReasonToAnthropic(param.FinishReason))
|
||||
results = append(results, "event: message_delta\ndata: "+messageDeltaJSON+"\n\n")
|
||||
param.MessageDeltaSent = true
|
||||
}
|
||||
|
||||
@@ -455,105 +367,73 @@ func convertOpenAIDoneToAnthropic(param *ConvertOpenAIResponseToAnthropicParams)
|
||||
func convertOpenAINonStreamingToAnthropic(rawJSON []byte) []string {
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
|
||||
// Build Anthropic response
|
||||
response := map[string]interface{}{
|
||||
"id": root.Get("id").String(),
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": root.Get("model").String(),
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
}
|
||||
out := `{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
out, _ = sjson.Set(out, "id", root.Get("id").String())
|
||||
out, _ = sjson.Set(out, "model", root.Get("model").String())
|
||||
|
||||
// Process message content and tool calls
|
||||
var contentBlocks []interface{}
|
||||
|
||||
if choices := root.Get("choices"); choices.Exists() && choices.IsArray() {
|
||||
if choices := root.Get("choices"); choices.Exists() && choices.IsArray() && len(choices.Array()) > 0 {
|
||||
choice := choices.Array()[0] // Take first choice
|
||||
reasoningNode := choice.Get("message.reasoning_content")
|
||||
allReasoning := collectOpenAIReasoningTexts(reasoningNode)
|
||||
|
||||
for _, reasoningText := range allReasoning {
|
||||
reasoningNode := choice.Get("message.reasoning_content")
|
||||
for _, reasoningText := range collectOpenAIReasoningTexts(reasoningNode) {
|
||||
if reasoningText == "" {
|
||||
continue
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": reasoningText,
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", reasoningText)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
|
||||
// Handle text content
|
||||
if content := choice.Get("message.content"); content.Exists() && content.String() != "" {
|
||||
textBlock := map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": content.String(),
|
||||
}
|
||||
contentBlocks = append(contentBlocks, textBlock)
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", content.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
|
||||
// Handle tool calls
|
||||
if toolCalls := choice.Get("message.tool_calls"); toolCalls.Exists() && toolCalls.IsArray() {
|
||||
toolCalls.ForEach(func(_, toolCall gjson.Result) bool {
|
||||
toolUseBlock := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": toolCall.Get("id").String(),
|
||||
"name": toolCall.Get("function.name").String(),
|
||||
}
|
||||
toolUseBlock := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolUseBlock, _ = sjson.Set(toolUseBlock, "id", toolCall.Get("id").String())
|
||||
toolUseBlock, _ = sjson.Set(toolUseBlock, "name", toolCall.Get("function.name").String())
|
||||
|
||||
// Parse arguments
|
||||
argsStr := toolCall.Get("function.arguments").String()
|
||||
argsStr = util.FixJSON(argsStr)
|
||||
if argsStr != "" {
|
||||
var args interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
|
||||
toolUseBlock["input"] = args
|
||||
argsStr := util.FixJSON(toolCall.Get("function.arguments").String())
|
||||
if argsStr != "" && gjson.Valid(argsStr) {
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", argsJSON.Raw)
|
||||
} else {
|
||||
toolUseBlock["input"] = map[string]interface{}{}
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", "{}")
|
||||
}
|
||||
} else {
|
||||
toolUseBlock["input"] = map[string]interface{}{}
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", "{}")
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolUseBlock)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolUseBlock)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Set stop reason
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() {
|
||||
response["stop_reason"] = mapOpenAIFinishReasonToAnthropic(finishReason.String())
|
||||
out, _ = sjson.Set(out, "stop_reason", mapOpenAIFinishReasonToAnthropic(finishReason.String()))
|
||||
}
|
||||
}
|
||||
|
||||
response["content"] = contentBlocks
|
||||
|
||||
// Set usage information
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
response["usage"] = map[string]interface{}{
|
||||
"input_tokens": usage.Get("prompt_tokens").Int(),
|
||||
"output_tokens": usage.Get("completion_tokens").Int(),
|
||||
"reasoning_tokens": func() int64 {
|
||||
if v := usage.Get("completion_tokens_details.reasoning_tokens"); v.Exists() {
|
||||
return v.Int()
|
||||
}
|
||||
return 0
|
||||
}(),
|
||||
}
|
||||
} else {
|
||||
response["usage"] = map[string]interface{}{
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", usage.Get("prompt_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", usage.Get("completion_tokens").Int())
|
||||
reasoningTokens := int64(0)
|
||||
if v := usage.Get("completion_tokens_details.reasoning_tokens"); v.Exists() {
|
||||
reasoningTokens = v.Int()
|
||||
}
|
||||
out, _ = sjson.Set(out, "usage.reasoning_tokens", reasoningTokens)
|
||||
}
|
||||
|
||||
responseJSON, _ := json.Marshal(response)
|
||||
return []string{string(responseJSON)}
|
||||
return []string{out}
|
||||
}
|
||||
|
||||
// mapOpenAIFinishReasonToAnthropic maps OpenAI finish reasons to Anthropic equivalents
|
||||
@@ -620,12 +500,9 @@ func stopThinkingContentBlock(param *ConvertOpenAIResponseToAnthropicParams, res
|
||||
if !param.ThinkingContentBlockStarted {
|
||||
return
|
||||
}
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": param.ThinkingContentBlockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
*results = append(*results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", param.ThinkingContentBlockIndex)
|
||||
*results = append(*results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
param.ThinkingContentBlockStarted = false
|
||||
param.ThinkingContentBlockIndex = -1
|
||||
}
|
||||
@@ -642,12 +519,9 @@ func stopTextContentBlock(param *ConvertOpenAIResponseToAnthropicParams, results
|
||||
if !param.TextContentBlockStarted {
|
||||
return
|
||||
}
|
||||
contentBlockStop := map[string]interface{}{
|
||||
"type": "content_block_stop",
|
||||
"index": param.TextContentBlockIndex,
|
||||
}
|
||||
contentBlockStopJSON, _ := json.Marshal(contentBlockStop)
|
||||
*results = append(*results, "event: content_block_stop\ndata: "+string(contentBlockStopJSON)+"\n\n")
|
||||
contentBlockStopJSON := `{"type":"content_block_stop","index":0}`
|
||||
contentBlockStopJSON, _ = sjson.Set(contentBlockStopJSON, "index", param.TextContentBlockIndex)
|
||||
*results = append(*results, "event: content_block_stop\ndata: "+contentBlockStopJSON+"\n\n")
|
||||
param.TextContentBlockStarted = false
|
||||
param.TextContentBlockIndex = -1
|
||||
}
|
||||
@@ -667,29 +541,19 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
_ = requestRawJSON
|
||||
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
out := `{"id":"","type":"message","role":"assistant","model":"","content":[],"stop_reason":null,"stop_sequence":null,"usage":{"input_tokens":0,"output_tokens":0}}`
|
||||
out, _ = sjson.Set(out, "id", root.Get("id").String())
|
||||
out, _ = sjson.Set(out, "model", root.Get("model").String())
|
||||
|
||||
response := map[string]interface{}{
|
||||
"id": root.Get("id").String(),
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": root.Get("model").String(),
|
||||
"content": []interface{}{},
|
||||
"stop_reason": nil,
|
||||
"stop_sequence": nil,
|
||||
"usage": map[string]interface{}{
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
}
|
||||
|
||||
contentBlocks := make([]interface{}, 0)
|
||||
hasToolCall := false
|
||||
stopReasonSet := false
|
||||
|
||||
if choices := root.Get("choices"); choices.Exists() && choices.IsArray() && len(choices.Array()) > 0 {
|
||||
choice := choices.Array()[0]
|
||||
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() {
|
||||
response["stop_reason"] = mapOpenAIFinishReasonToAnthropic(finishReason.String())
|
||||
out, _ = sjson.Set(out, "stop_reason", mapOpenAIFinishReasonToAnthropic(finishReason.String()))
|
||||
stopReasonSet = true
|
||||
}
|
||||
|
||||
if message := choice.Get("message"); message.Exists() {
|
||||
@@ -702,10 +566,9 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if textBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": textBuilder.String(),
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", textBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
textBuilder.Reset()
|
||||
}
|
||||
|
||||
@@ -713,16 +576,14 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if thinkingBuilder.Len() == 0 {
|
||||
return
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": thinkingBuilder.String(),
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", thinkingBuilder.String())
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
thinkingBuilder.Reset()
|
||||
}
|
||||
|
||||
for _, item := range contentResult.Array() {
|
||||
typeStr := item.Get("type").String()
|
||||
switch typeStr {
|
||||
switch item.Get("type").String() {
|
||||
case "text":
|
||||
flushThinking()
|
||||
textBuilder.WriteString(item.Get("text").String())
|
||||
@@ -733,25 +594,23 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if toolCalls.IsArray() {
|
||||
toolCalls.ForEach(func(_, tc gjson.Result) bool {
|
||||
hasToolCall = true
|
||||
toolUse := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": tc.Get("id").String(),
|
||||
"name": tc.Get("function.name").String(),
|
||||
}
|
||||
toolUse := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolUse, _ = sjson.Set(toolUse, "id", tc.Get("id").String())
|
||||
toolUse, _ = sjson.Set(toolUse, "name", tc.Get("function.name").String())
|
||||
|
||||
argsStr := util.FixJSON(tc.Get("function.arguments").String())
|
||||
if argsStr != "" {
|
||||
var parsed interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &parsed); err == nil {
|
||||
toolUse["input"] = parsed
|
||||
if argsStr != "" && gjson.Valid(argsStr) {
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", argsJSON.Raw)
|
||||
} else {
|
||||
toolUse["input"] = map[string]interface{}{}
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", "{}")
|
||||
}
|
||||
} else {
|
||||
toolUse["input"] = map[string]interface{}{}
|
||||
toolUse, _ = sjson.SetRaw(toolUse, "input", "{}")
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolUse)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolUse)
|
||||
return true
|
||||
})
|
||||
}
|
||||
@@ -771,10 +630,9 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
} else if contentResult.Type == gjson.String {
|
||||
textContent := contentResult.String()
|
||||
if textContent != "" {
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": textContent,
|
||||
})
|
||||
block := `{"type":"text","text":""}`
|
||||
block, _ = sjson.Set(block, "text", textContent)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -784,81 +642,52 @@ func ConvertOpenAIResponseToClaudeNonStream(_ context.Context, _ string, origina
|
||||
if reasoningText == "" {
|
||||
continue
|
||||
}
|
||||
contentBlocks = append(contentBlocks, map[string]interface{}{
|
||||
"type": "thinking",
|
||||
"thinking": reasoningText,
|
||||
})
|
||||
block := `{"type":"thinking","thinking":""}`
|
||||
block, _ = sjson.Set(block, "thinking", reasoningText)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", block)
|
||||
}
|
||||
}
|
||||
|
||||
if toolCalls := message.Get("tool_calls"); toolCalls.Exists() && toolCalls.IsArray() {
|
||||
toolCalls.ForEach(func(_, toolCall gjson.Result) bool {
|
||||
hasToolCall = true
|
||||
toolUseBlock := map[string]interface{}{
|
||||
"type": "tool_use",
|
||||
"id": toolCall.Get("id").String(),
|
||||
"name": toolCall.Get("function.name").String(),
|
||||
}
|
||||
toolUseBlock := `{"type":"tool_use","id":"","name":"","input":{}}`
|
||||
toolUseBlock, _ = sjson.Set(toolUseBlock, "id", toolCall.Get("id").String())
|
||||
toolUseBlock, _ = sjson.Set(toolUseBlock, "name", toolCall.Get("function.name").String())
|
||||
|
||||
argsStr := toolCall.Get("function.arguments").String()
|
||||
argsStr = util.FixJSON(argsStr)
|
||||
if argsStr != "" {
|
||||
var args interface{}
|
||||
if err := json.Unmarshal([]byte(argsStr), &args); err == nil {
|
||||
toolUseBlock["input"] = args
|
||||
argsStr := util.FixJSON(toolCall.Get("function.arguments").String())
|
||||
if argsStr != "" && gjson.Valid(argsStr) {
|
||||
argsJSON := gjson.Parse(argsStr)
|
||||
if argsJSON.IsObject() {
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", argsJSON.Raw)
|
||||
} else {
|
||||
toolUseBlock["input"] = map[string]interface{}{}
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", "{}")
|
||||
}
|
||||
} else {
|
||||
toolUseBlock["input"] = map[string]interface{}{}
|
||||
toolUseBlock, _ = sjson.SetRaw(toolUseBlock, "input", "{}")
|
||||
}
|
||||
|
||||
contentBlocks = append(contentBlocks, toolUseBlock)
|
||||
out, _ = sjson.SetRaw(out, "content.-1", toolUseBlock)
|
||||
return true
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
response["content"] = contentBlocks
|
||||
|
||||
if respUsage := root.Get("usage"); respUsage.Exists() {
|
||||
usageJSON := `{}`
|
||||
usageJSON, _ = sjson.Set(usageJSON, "input_tokens", respUsage.Get("prompt_tokens").Int())
|
||||
usageJSON, _ = sjson.Set(usageJSON, "output_tokens", respUsage.Get("completion_tokens").Int())
|
||||
parsedUsage := gjson.Parse(usageJSON).Value().(map[string]interface{})
|
||||
response["usage"] = parsedUsage
|
||||
} else {
|
||||
response["usage"] = `{"input_tokens":0,"output_tokens":0}`
|
||||
out, _ = sjson.Set(out, "usage.input_tokens", respUsage.Get("prompt_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usage.output_tokens", respUsage.Get("completion_tokens").Int())
|
||||
}
|
||||
|
||||
if response["stop_reason"] == nil {
|
||||
if !stopReasonSet {
|
||||
if hasToolCall {
|
||||
response["stop_reason"] = "tool_use"
|
||||
out, _ = sjson.Set(out, "stop_reason", "tool_use")
|
||||
} else {
|
||||
response["stop_reason"] = "end_turn"
|
||||
out, _ = sjson.Set(out, "stop_reason", "end_turn")
|
||||
}
|
||||
}
|
||||
|
||||
if !hasToolCall {
|
||||
if toolBlocks := response["content"].([]interface{}); len(toolBlocks) > 0 {
|
||||
for _, block := range toolBlocks {
|
||||
if m, ok := block.(map[string]interface{}); ok && m["type"] == "tool_use" {
|
||||
hasToolCall = true
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
if hasToolCall {
|
||||
response["stop_reason"] = "tool_use"
|
||||
}
|
||||
}
|
||||
|
||||
responseJSON, err := json.Marshal(response)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
return string(responseJSON)
|
||||
return out
|
||||
}
|
||||
|
||||
func ClaudeTokenCount(ctx context.Context, count int64) string {
|
||||
|
||||
@@ -8,7 +8,6 @@ package gemini
|
||||
import (
|
||||
"bytes"
|
||||
"crypto/rand"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math/big"
|
||||
"strings"
|
||||
@@ -94,7 +93,6 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
out, _ = sjson.Set(out, "stream", stream)
|
||||
|
||||
// Process contents (Gemini messages) -> OpenAI messages
|
||||
var openAIMessages []interface{}
|
||||
var toolCallIDs []string // Track tool call IDs for matching with tool results
|
||||
|
||||
// System instruction -> OpenAI system message
|
||||
@@ -105,22 +103,17 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
}
|
||||
if systemInstruction.Exists() {
|
||||
parts := systemInstruction.Get("parts")
|
||||
msg := map[string]interface{}{
|
||||
"role": "system",
|
||||
"content": []interface{}{},
|
||||
}
|
||||
|
||||
var aggregatedParts []interface{}
|
||||
msg := `{"role":"system","content":[]}`
|
||||
hasContent := false
|
||||
|
||||
if parts.Exists() && parts.IsArray() {
|
||||
parts.ForEach(func(_, part gjson.Result) bool {
|
||||
// Handle text parts
|
||||
if text := part.Get("text"); text.Exists() {
|
||||
formattedText := text.String()
|
||||
aggregatedParts = append(aggregatedParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": formattedText,
|
||||
})
|
||||
contentPart := `{"type":"text","text":""}`
|
||||
contentPart, _ = sjson.Set(contentPart, "text", text.String())
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", contentPart)
|
||||
hasContent = true
|
||||
}
|
||||
|
||||
// Handle inline data (e.g., images)
|
||||
@@ -132,20 +125,17 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
data := inlineData.Get("data").String()
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
|
||||
aggregatedParts = append(aggregatedParts, map[string]interface{}{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]interface{}{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
contentPart := `{"type":"image_url","image_url":{"url":""}}`
|
||||
contentPart, _ = sjson.Set(contentPart, "image_url.url", imageURL)
|
||||
msg, _ = sjson.SetRaw(msg, "content.-1", contentPart)
|
||||
hasContent = true
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
if len(aggregatedParts) > 0 {
|
||||
msg["content"] = aggregatedParts
|
||||
openAIMessages = append(openAIMessages, msg)
|
||||
if hasContent {
|
||||
out, _ = sjson.SetRaw(out, "messages.-1", msg)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -159,16 +149,15 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
role = "assistant"
|
||||
}
|
||||
|
||||
// Create OpenAI message
|
||||
msg := map[string]interface{}{
|
||||
"role": role,
|
||||
"content": "",
|
||||
}
|
||||
msg := `{"role":"","content":""}`
|
||||
msg, _ = sjson.Set(msg, "role", role)
|
||||
|
||||
var textBuilder strings.Builder
|
||||
var aggregatedParts []interface{}
|
||||
contentWrapper := `{"arr":[]}`
|
||||
contentPartsCount := 0
|
||||
onlyTextContent := true
|
||||
var toolCalls []interface{}
|
||||
toolCallsWrapper := `{"arr":[]}`
|
||||
toolCallsCount := 0
|
||||
|
||||
if parts.Exists() && parts.IsArray() {
|
||||
parts.ForEach(func(_, part gjson.Result) bool {
|
||||
@@ -176,10 +165,10 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
if text := part.Get("text"); text.Exists() {
|
||||
formattedText := text.String()
|
||||
textBuilder.WriteString(formattedText)
|
||||
aggregatedParts = append(aggregatedParts, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": formattedText,
|
||||
})
|
||||
contentPart := `{"type":"text","text":""}`
|
||||
contentPart, _ = sjson.Set(contentPart, "text", formattedText)
|
||||
contentWrapper, _ = sjson.SetRaw(contentWrapper, "arr.-1", contentPart)
|
||||
contentPartsCount++
|
||||
}
|
||||
|
||||
// Handle inline data (e.g., images)
|
||||
@@ -193,12 +182,10 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
data := inlineData.Get("data").String()
|
||||
imageURL := fmt.Sprintf("data:%s;base64,%s", mimeType, data)
|
||||
|
||||
aggregatedParts = append(aggregatedParts, map[string]interface{}{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]interface{}{
|
||||
"url": imageURL,
|
||||
},
|
||||
})
|
||||
contentPart := `{"type":"image_url","image_url":{"url":""}}`
|
||||
contentPart, _ = sjson.Set(contentPart, "image_url.url", imageURL)
|
||||
contentWrapper, _ = sjson.SetRaw(contentWrapper, "arr.-1", contentPart)
|
||||
contentPartsCount++
|
||||
}
|
||||
|
||||
// Handle function calls (Gemini) -> tool calls (OpenAI)
|
||||
@@ -206,44 +193,32 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
toolCallID := genToolCallID()
|
||||
toolCallIDs = append(toolCallIDs, toolCallID)
|
||||
|
||||
toolCall := map[string]interface{}{
|
||||
"id": toolCallID,
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": functionCall.Get("name").String(),
|
||||
},
|
||||
}
|
||||
toolCall := `{"id":"","type":"function","function":{"name":"","arguments":""}}`
|
||||
toolCall, _ = sjson.Set(toolCall, "id", toolCallID)
|
||||
toolCall, _ = sjson.Set(toolCall, "function.name", functionCall.Get("name").String())
|
||||
|
||||
// Convert args to arguments JSON string
|
||||
if args := functionCall.Get("args"); args.Exists() {
|
||||
argsJSON, _ := json.Marshal(args.Value())
|
||||
toolCall["function"].(map[string]interface{})["arguments"] = string(argsJSON)
|
||||
toolCall, _ = sjson.Set(toolCall, "function.arguments", args.Raw)
|
||||
} else {
|
||||
toolCall["function"].(map[string]interface{})["arguments"] = "{}"
|
||||
toolCall, _ = sjson.Set(toolCall, "function.arguments", "{}")
|
||||
}
|
||||
|
||||
toolCalls = append(toolCalls, toolCall)
|
||||
toolCallsWrapper, _ = sjson.SetRaw(toolCallsWrapper, "arr.-1", toolCall)
|
||||
toolCallsCount++
|
||||
}
|
||||
|
||||
// Handle function responses (Gemini) -> tool role messages (OpenAI)
|
||||
if functionResponse := part.Get("functionResponse"); functionResponse.Exists() {
|
||||
// Create tool message for function response
|
||||
toolMsg := map[string]interface{}{
|
||||
"role": "tool",
|
||||
"tool_call_id": "", // Will be set based on context
|
||||
"content": "",
|
||||
}
|
||||
toolMsg := `{"role":"tool","tool_call_id":"","content":""}`
|
||||
|
||||
// Convert response.content to JSON string
|
||||
if response := functionResponse.Get("response"); response.Exists() {
|
||||
if content = response.Get("content"); content.Exists() {
|
||||
// Use the content field from the response
|
||||
contentJSON, _ := json.Marshal(content.Value())
|
||||
toolMsg["content"] = string(contentJSON)
|
||||
if contentField := response.Get("content"); contentField.Exists() {
|
||||
toolMsg, _ = sjson.Set(toolMsg, "content", contentField.Raw)
|
||||
} else {
|
||||
// Fallback to entire response
|
||||
responseJSON, _ := json.Marshal(response.Value())
|
||||
toolMsg["content"] = string(responseJSON)
|
||||
toolMsg, _ = sjson.Set(toolMsg, "content", response.Raw)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -252,13 +227,13 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
if len(toolCallIDs) > 0 {
|
||||
// Use the last tool call ID (simple matching by function name)
|
||||
// In a real implementation, you might want more sophisticated matching
|
||||
toolMsg["tool_call_id"] = toolCallIDs[len(toolCallIDs)-1]
|
||||
toolMsg, _ = sjson.Set(toolMsg, "tool_call_id", toolCallIDs[len(toolCallIDs)-1])
|
||||
} else {
|
||||
// Generate a tool call ID if none available
|
||||
toolMsg["tool_call_id"] = genToolCallID()
|
||||
toolMsg, _ = sjson.Set(toolMsg, "tool_call_id", genToolCallID())
|
||||
}
|
||||
|
||||
openAIMessages = append(openAIMessages, toolMsg)
|
||||
out, _ = sjson.SetRaw(out, "messages.-1", toolMsg)
|
||||
}
|
||||
|
||||
return true
|
||||
@@ -266,170 +241,46 @@ func ConvertGeminiRequestToOpenAI(modelName string, inputRawJSON []byte, stream
|
||||
}
|
||||
|
||||
// Set content
|
||||
if len(aggregatedParts) > 0 {
|
||||
if contentPartsCount > 0 {
|
||||
if onlyTextContent {
|
||||
msg["content"] = textBuilder.String()
|
||||
msg, _ = sjson.Set(msg, "content", textBuilder.String())
|
||||
} else {
|
||||
msg["content"] = aggregatedParts
|
||||
msg, _ = sjson.SetRaw(msg, "content", gjson.Get(contentWrapper, "arr").Raw)
|
||||
}
|
||||
}
|
||||
|
||||
// Set tool calls if any
|
||||
if len(toolCalls) > 0 {
|
||||
msg["tool_calls"] = toolCalls
|
||||
if toolCallsCount > 0 {
|
||||
msg, _ = sjson.SetRaw(msg, "tool_calls", gjson.Get(toolCallsWrapper, "arr").Raw)
|
||||
}
|
||||
|
||||
openAIMessages = append(openAIMessages, msg)
|
||||
|
||||
// switch role {
|
||||
// case "user", "model":
|
||||
// // Convert role: model -> assistant
|
||||
// if role == "model" {
|
||||
// role = "assistant"
|
||||
// }
|
||||
//
|
||||
// // Create OpenAI message
|
||||
// msg := map[string]interface{}{
|
||||
// "role": role,
|
||||
// "content": "",
|
||||
// }
|
||||
//
|
||||
// var contentParts []string
|
||||
// var toolCalls []interface{}
|
||||
//
|
||||
// if parts.Exists() && parts.IsArray() {
|
||||
// parts.ForEach(func(_, part gjson.Result) bool {
|
||||
// // Handle text parts
|
||||
// if text := part.Get("text"); text.Exists() {
|
||||
// contentParts = append(contentParts, text.String())
|
||||
// }
|
||||
//
|
||||
// // Handle function calls (Gemini) -> tool calls (OpenAI)
|
||||
// if functionCall := part.Get("functionCall"); functionCall.Exists() {
|
||||
// toolCallID := genToolCallID()
|
||||
// toolCallIDs = append(toolCallIDs, toolCallID)
|
||||
//
|
||||
// toolCall := map[string]interface{}{
|
||||
// "id": toolCallID,
|
||||
// "type": "function",
|
||||
// "function": map[string]interface{}{
|
||||
// "name": functionCall.Get("name").String(),
|
||||
// },
|
||||
// }
|
||||
//
|
||||
// // Convert args to arguments JSON string
|
||||
// if args := functionCall.Get("args"); args.Exists() {
|
||||
// argsJSON, _ := json.Marshal(args.Value())
|
||||
// toolCall["function"].(map[string]interface{})["arguments"] = string(argsJSON)
|
||||
// } else {
|
||||
// toolCall["function"].(map[string]interface{})["arguments"] = "{}"
|
||||
// }
|
||||
//
|
||||
// toolCalls = append(toolCalls, toolCall)
|
||||
// }
|
||||
//
|
||||
// return true
|
||||
// })
|
||||
// }
|
||||
//
|
||||
// // Set content
|
||||
// if len(contentParts) > 0 {
|
||||
// msg["content"] = strings.Join(contentParts, "")
|
||||
// }
|
||||
//
|
||||
// // Set tool calls if any
|
||||
// if len(toolCalls) > 0 {
|
||||
// msg["tool_calls"] = toolCalls
|
||||
// }
|
||||
//
|
||||
// openAIMessages = append(openAIMessages, msg)
|
||||
//
|
||||
// case "function":
|
||||
// // Handle Gemini function role -> OpenAI tool role
|
||||
// if parts.Exists() && parts.IsArray() {
|
||||
// parts.ForEach(func(_, part gjson.Result) bool {
|
||||
// // Handle function responses (Gemini) -> tool role messages (OpenAI)
|
||||
// if functionResponse := part.Get("functionResponse"); functionResponse.Exists() {
|
||||
// // Create tool message for function response
|
||||
// toolMsg := map[string]interface{}{
|
||||
// "role": "tool",
|
||||
// "tool_call_id": "", // Will be set based on context
|
||||
// "content": "",
|
||||
// }
|
||||
//
|
||||
// // Convert response.content to JSON string
|
||||
// if response := functionResponse.Get("response"); response.Exists() {
|
||||
// if content = response.Get("content"); content.Exists() {
|
||||
// // Use the content field from the response
|
||||
// contentJSON, _ := json.Marshal(content.Value())
|
||||
// toolMsg["content"] = string(contentJSON)
|
||||
// } else {
|
||||
// // Fallback to entire response
|
||||
// responseJSON, _ := json.Marshal(response.Value())
|
||||
// toolMsg["content"] = string(responseJSON)
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// // Try to match with previous tool call ID
|
||||
// _ = functionResponse.Get("name").String() // functionName not used for now
|
||||
// if len(toolCallIDs) > 0 {
|
||||
// // Use the last tool call ID (simple matching by function name)
|
||||
// // In a real implementation, you might want more sophisticated matching
|
||||
// toolMsg["tool_call_id"] = toolCallIDs[len(toolCallIDs)-1]
|
||||
// } else {
|
||||
// // Generate a tool call ID if none available
|
||||
// toolMsg["tool_call_id"] = genToolCallID()
|
||||
// }
|
||||
//
|
||||
// openAIMessages = append(openAIMessages, toolMsg)
|
||||
// }
|
||||
//
|
||||
// return true
|
||||
// })
|
||||
// }
|
||||
// }
|
||||
out, _ = sjson.SetRaw(out, "messages.-1", msg)
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Set messages
|
||||
if len(openAIMessages) > 0 {
|
||||
messagesJSON, _ := json.Marshal(openAIMessages)
|
||||
out, _ = sjson.SetRaw(out, "messages", string(messagesJSON))
|
||||
}
|
||||
|
||||
// Tools mapping: Gemini tools -> OpenAI tools
|
||||
if tools := root.Get("tools"); tools.Exists() && tools.IsArray() {
|
||||
var openAITools []interface{}
|
||||
tools.ForEach(func(_, tool gjson.Result) bool {
|
||||
if functionDeclarations := tool.Get("functionDeclarations"); functionDeclarations.Exists() && functionDeclarations.IsArray() {
|
||||
functionDeclarations.ForEach(func(_, funcDecl gjson.Result) bool {
|
||||
openAITool := map[string]interface{}{
|
||||
"type": "function",
|
||||
"function": map[string]interface{}{
|
||||
"name": funcDecl.Get("name").String(),
|
||||
"description": funcDecl.Get("description").String(),
|
||||
},
|
||||
}
|
||||
openAITool := `{"type":"function","function":{"name":"","description":""}}`
|
||||
openAITool, _ = sjson.Set(openAITool, "function.name", funcDecl.Get("name").String())
|
||||
openAITool, _ = sjson.Set(openAITool, "function.description", funcDecl.Get("description").String())
|
||||
|
||||
// Convert parameters schema
|
||||
if parameters := funcDecl.Get("parameters"); parameters.Exists() {
|
||||
openAITool["function"].(map[string]interface{})["parameters"] = parameters.Value()
|
||||
} else if parameters = funcDecl.Get("parametersJsonSchema"); parameters.Exists() {
|
||||
openAITool["function"].(map[string]interface{})["parameters"] = parameters.Value()
|
||||
openAITool, _ = sjson.SetRaw(openAITool, "function.parameters", parameters.Raw)
|
||||
} else if parameters := funcDecl.Get("parametersJsonSchema"); parameters.Exists() {
|
||||
openAITool, _ = sjson.SetRaw(openAITool, "function.parameters", parameters.Raw)
|
||||
}
|
||||
|
||||
openAITools = append(openAITools, openAITool)
|
||||
out, _ = sjson.SetRaw(out, "tools.-1", openAITool)
|
||||
return true
|
||||
})
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if len(openAITools) > 0 {
|
||||
toolsJSON, _ := json.Marshal(openAITools)
|
||||
out, _ = sjson.SetRaw(out, "tools", string(toolsJSON))
|
||||
}
|
||||
}
|
||||
|
||||
// Tool choice mapping (Gemini doesn't have direct equivalent, but we can handle it)
|
||||
|
||||
@@ -8,7 +8,6 @@ package gemini
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
@@ -84,15 +83,12 @@ func ConvertOpenAIResponseToGemini(_ context.Context, _ string, originalRequestR
|
||||
template, _ = sjson.Set(template, "model", model.String())
|
||||
}
|
||||
|
||||
usageObj := map[string]interface{}{
|
||||
"promptTokenCount": usage.Get("prompt_tokens").Int(),
|
||||
"candidatesTokenCount": usage.Get("completion_tokens").Int(),
|
||||
"totalTokenCount": usage.Get("total_tokens").Int(),
|
||||
}
|
||||
template, _ = sjson.Set(template, "usageMetadata.promptTokenCount", usage.Get("prompt_tokens").Int())
|
||||
template, _ = sjson.Set(template, "usageMetadata.candidatesTokenCount", usage.Get("completion_tokens").Int())
|
||||
template, _ = sjson.Set(template, "usageMetadata.totalTokenCount", usage.Get("total_tokens").Int())
|
||||
if reasoningTokens := reasoningTokensFromUsage(usage); reasoningTokens > 0 {
|
||||
usageObj["thoughtsTokenCount"] = reasoningTokens
|
||||
template, _ = sjson.Set(template, "usageMetadata.thoughtsTokenCount", reasoningTokens)
|
||||
}
|
||||
template, _ = sjson.Set(template, "usageMetadata", usageObj)
|
||||
return []string{template}
|
||||
}
|
||||
return []string{}
|
||||
@@ -133,13 +129,8 @@ func ConvertOpenAIResponseToGemini(_ context.Context, _ string, originalRequestR
|
||||
continue
|
||||
}
|
||||
reasoningTemplate := baseTemplate
|
||||
parts := []interface{}{
|
||||
map[string]interface{}{
|
||||
"thought": true,
|
||||
"text": reasoningText,
|
||||
},
|
||||
}
|
||||
reasoningTemplate, _ = sjson.Set(reasoningTemplate, "candidates.0.content.parts", parts)
|
||||
reasoningTemplate, _ = sjson.Set(reasoningTemplate, "candidates.0.content.parts.0.thought", true)
|
||||
reasoningTemplate, _ = sjson.Set(reasoningTemplate, "candidates.0.content.parts.0.text", reasoningText)
|
||||
chunkOutputs = append(chunkOutputs, reasoningTemplate)
|
||||
}
|
||||
}
|
||||
@@ -150,13 +141,8 @@ func ConvertOpenAIResponseToGemini(_ context.Context, _ string, originalRequestR
|
||||
(*param).(*ConvertOpenAIResponseToGeminiParams).ContentAccumulator.WriteString(contentText)
|
||||
|
||||
// Create text part for this delta
|
||||
parts := []interface{}{
|
||||
map[string]interface{}{
|
||||
"text": contentText,
|
||||
},
|
||||
}
|
||||
contentTemplate := baseTemplate
|
||||
contentTemplate, _ = sjson.Set(contentTemplate, "candidates.0.content.parts", parts)
|
||||
contentTemplate, _ = sjson.Set(contentTemplate, "candidates.0.content.parts.0.text", contentText)
|
||||
chunkOutputs = append(chunkOutputs, contentTemplate)
|
||||
}
|
||||
|
||||
@@ -225,24 +211,13 @@ func ConvertOpenAIResponseToGemini(_ context.Context, _ string, originalRequestR
|
||||
|
||||
// If we have accumulated tool calls, output them now
|
||||
if len((*param).(*ConvertOpenAIResponseToGeminiParams).ToolCallsAccumulator) > 0 {
|
||||
var parts []interface{}
|
||||
partIndex := 0
|
||||
for _, accumulator := range (*param).(*ConvertOpenAIResponseToGeminiParams).ToolCallsAccumulator {
|
||||
argsStr := accumulator.Arguments.String()
|
||||
var argsMap map[string]interface{}
|
||||
|
||||
argsMap = parseArgsToMap(argsStr)
|
||||
|
||||
functionCallPart := map[string]interface{}{
|
||||
"functionCall": map[string]interface{}{
|
||||
"name": accumulator.Name,
|
||||
"args": argsMap,
|
||||
},
|
||||
}
|
||||
parts = append(parts, functionCallPart)
|
||||
}
|
||||
|
||||
if len(parts) > 0 {
|
||||
template, _ = sjson.Set(template, "candidates.0.content.parts", parts)
|
||||
namePath := fmt.Sprintf("candidates.0.content.parts.%d.functionCall.name", partIndex)
|
||||
argsPath := fmt.Sprintf("candidates.0.content.parts.%d.functionCall.args", partIndex)
|
||||
template, _ = sjson.Set(template, namePath, accumulator.Name)
|
||||
template, _ = sjson.SetRaw(template, argsPath, parseArgsToObjectRaw(accumulator.Arguments.String()))
|
||||
partIndex++
|
||||
}
|
||||
|
||||
// Clear accumulators
|
||||
@@ -255,15 +230,12 @@ func ConvertOpenAIResponseToGemini(_ context.Context, _ string, originalRequestR
|
||||
|
||||
// Handle usage information
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
usageObj := map[string]interface{}{
|
||||
"promptTokenCount": usage.Get("prompt_tokens").Int(),
|
||||
"candidatesTokenCount": usage.Get("completion_tokens").Int(),
|
||||
"totalTokenCount": usage.Get("total_tokens").Int(),
|
||||
}
|
||||
template, _ = sjson.Set(template, "usageMetadata.promptTokenCount", usage.Get("prompt_tokens").Int())
|
||||
template, _ = sjson.Set(template, "usageMetadata.candidatesTokenCount", usage.Get("completion_tokens").Int())
|
||||
template, _ = sjson.Set(template, "usageMetadata.totalTokenCount", usage.Get("total_tokens").Int())
|
||||
if reasoningTokens := reasoningTokensFromUsage(usage); reasoningTokens > 0 {
|
||||
usageObj["thoughtsTokenCount"] = reasoningTokens
|
||||
template, _ = sjson.Set(template, "usageMetadata.thoughtsTokenCount", reasoningTokens)
|
||||
}
|
||||
template, _ = sjson.Set(template, "usageMetadata", usageObj)
|
||||
results = append(results, template)
|
||||
return true
|
||||
}
|
||||
@@ -291,46 +263,54 @@ func mapOpenAIFinishReasonToGemini(openAIReason string) string {
|
||||
}
|
||||
}
|
||||
|
||||
// parseArgsToMap safely parses a JSON string of function arguments into a map.
|
||||
// It returns an empty map if the input is empty or cannot be parsed as a JSON object.
|
||||
func parseArgsToMap(argsStr string) map[string]interface{} {
|
||||
// parseArgsToObjectRaw safely parses a JSON string of function arguments into an object JSON string.
|
||||
// It returns "{}" if the input is empty or cannot be parsed as a JSON object.
|
||||
func parseArgsToObjectRaw(argsStr string) string {
|
||||
trimmed := strings.TrimSpace(argsStr)
|
||||
if trimmed == "" || trimmed == "{}" {
|
||||
return map[string]interface{}{}
|
||||
return "{}"
|
||||
}
|
||||
|
||||
// First try strict JSON
|
||||
var out map[string]interface{}
|
||||
if errUnmarshal := json.Unmarshal([]byte(trimmed), &out); errUnmarshal == nil {
|
||||
return out
|
||||
if gjson.Valid(trimmed) {
|
||||
strict := gjson.Parse(trimmed)
|
||||
if strict.IsObject() {
|
||||
return strict.Raw
|
||||
}
|
||||
}
|
||||
|
||||
// Tolerant parse: handle streams where values are barewords (e.g., 北京, celsius)
|
||||
tolerant := tolerantParseJSONMap(trimmed)
|
||||
if len(tolerant) > 0 {
|
||||
tolerant := tolerantParseJSONObjectRaw(trimmed)
|
||||
if tolerant != "{}" {
|
||||
return tolerant
|
||||
}
|
||||
|
||||
// Fallback: return empty object when parsing fails
|
||||
return map[string]interface{}{}
|
||||
return "{}"
|
||||
}
|
||||
|
||||
// tolerantParseJSONMap attempts to parse a JSON-like object string into a map, tolerating
|
||||
func escapeSjsonPathKey(key string) string {
|
||||
key = strings.ReplaceAll(key, `\`, `\\`)
|
||||
key = strings.ReplaceAll(key, `.`, `\.`)
|
||||
return key
|
||||
}
|
||||
|
||||
// tolerantParseJSONObjectRaw attempts to parse a JSON-like object string into a JSON object string, tolerating
|
||||
// bareword values (unquoted strings) commonly seen during streamed tool calls.
|
||||
// Example input: {"location": 北京, "unit": celsius}
|
||||
func tolerantParseJSONMap(s string) map[string]interface{} {
|
||||
func tolerantParseJSONObjectRaw(s string) string {
|
||||
// Ensure we operate within the outermost braces if present
|
||||
start := strings.Index(s, "{")
|
||||
end := strings.LastIndex(s, "}")
|
||||
if start == -1 || end == -1 || start >= end {
|
||||
return map[string]interface{}{}
|
||||
return "{}"
|
||||
}
|
||||
content := s[start+1 : end]
|
||||
|
||||
runes := []rune(content)
|
||||
n := len(runes)
|
||||
i := 0
|
||||
result := make(map[string]interface{})
|
||||
result := "{}"
|
||||
|
||||
for i < n {
|
||||
// Skip whitespace and commas
|
||||
@@ -356,6 +336,7 @@ func tolerantParseJSONMap(s string) map[string]interface{} {
|
||||
break
|
||||
}
|
||||
keyName := jsonStringTokenToRawString(keyToken)
|
||||
sjsonKey := escapeSjsonPathKey(keyName)
|
||||
i = nextIdx
|
||||
|
||||
// Skip whitespace
|
||||
@@ -375,17 +356,16 @@ func tolerantParseJSONMap(s string) map[string]interface{} {
|
||||
}
|
||||
|
||||
// Parse value (string, number, object/array, bareword)
|
||||
var value interface{}
|
||||
switch runes[i] {
|
||||
case '"':
|
||||
// JSON string
|
||||
valToken, ni := parseJSONStringRunes(runes, i)
|
||||
if ni == -1 {
|
||||
// Malformed; treat as empty string
|
||||
value = ""
|
||||
result, _ = sjson.Set(result, sjsonKey, "")
|
||||
i = n
|
||||
} else {
|
||||
value = jsonStringTokenToRawString(valToken)
|
||||
result, _ = sjson.Set(result, sjsonKey, jsonStringTokenToRawString(valToken))
|
||||
i = ni
|
||||
}
|
||||
case '{', '[':
|
||||
@@ -394,11 +374,10 @@ func tolerantParseJSONMap(s string) map[string]interface{} {
|
||||
if ni == -1 {
|
||||
i = n
|
||||
} else {
|
||||
var anyVal interface{}
|
||||
if errUnmarshal := json.Unmarshal([]byte(seg), &anyVal); errUnmarshal == nil {
|
||||
value = anyVal
|
||||
if gjson.Valid(seg) {
|
||||
result, _ = sjson.SetRaw(result, sjsonKey, seg)
|
||||
} else {
|
||||
value = seg
|
||||
result, _ = sjson.Set(result, sjsonKey, seg)
|
||||
}
|
||||
i = ni
|
||||
}
|
||||
@@ -411,21 +390,19 @@ func tolerantParseJSONMap(s string) map[string]interface{} {
|
||||
token := strings.TrimSpace(string(runes[i:j]))
|
||||
// Interpret common JSON atoms and numbers; otherwise treat as string
|
||||
if token == "true" {
|
||||
value = true
|
||||
result, _ = sjson.Set(result, sjsonKey, true)
|
||||
} else if token == "false" {
|
||||
value = false
|
||||
result, _ = sjson.Set(result, sjsonKey, false)
|
||||
} else if token == "null" {
|
||||
value = nil
|
||||
result, _ = sjson.Set(result, sjsonKey, nil)
|
||||
} else if numVal, ok := tryParseNumber(token); ok {
|
||||
value = numVal
|
||||
result, _ = sjson.Set(result, sjsonKey, numVal)
|
||||
} else {
|
||||
value = token
|
||||
result, _ = sjson.Set(result, sjsonKey, token)
|
||||
}
|
||||
i = j
|
||||
}
|
||||
|
||||
result[keyName] = value
|
||||
|
||||
// Skip trailing whitespace and optional comma before next pair
|
||||
for i < n && (runes[i] == ' ' || runes[i] == '\n' || runes[i] == '\r' || runes[i] == '\t') {
|
||||
i++
|
||||
@@ -463,9 +440,9 @@ func parseJSONStringRunes(runes []rune, start int) (string, int) {
|
||||
|
||||
// jsonStringTokenToRawString converts a JSON string token (including quotes) to a raw Go string value.
|
||||
func jsonStringTokenToRawString(token string) string {
|
||||
var s string
|
||||
if errUnmarshal := json.Unmarshal([]byte(token), &s); errUnmarshal == nil {
|
||||
return s
|
||||
r := gjson.Parse(token)
|
||||
if r.Type == gjson.String {
|
||||
return r.String()
|
||||
}
|
||||
// Fallback: strip surrounding quotes if present
|
||||
if len(token) >= 2 && token[0] == '"' && token[len(token)-1] == '"' {
|
||||
@@ -579,7 +556,7 @@ func ConvertOpenAIResponseToGeminiNonStream(_ context.Context, _ string, origina
|
||||
}
|
||||
}
|
||||
|
||||
var parts []interface{}
|
||||
partIndex := 0
|
||||
|
||||
// Handle reasoning content before visible text
|
||||
if reasoning := message.Get("reasoning_content"); reasoning.Exists() {
|
||||
@@ -587,18 +564,16 @@ func ConvertOpenAIResponseToGeminiNonStream(_ context.Context, _ string, origina
|
||||
if reasoningText == "" {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, map[string]interface{}{
|
||||
"thought": true,
|
||||
"text": reasoningText,
|
||||
})
|
||||
out, _ = sjson.Set(out, fmt.Sprintf("candidates.0.content.parts.%d.thought", partIndex), true)
|
||||
out, _ = sjson.Set(out, fmt.Sprintf("candidates.0.content.parts.%d.text", partIndex), reasoningText)
|
||||
partIndex++
|
||||
}
|
||||
}
|
||||
|
||||
// Handle content first
|
||||
if content := message.Get("content"); content.Exists() && content.String() != "" {
|
||||
parts = append(parts, map[string]interface{}{
|
||||
"text": content.String(),
|
||||
})
|
||||
out, _ = sjson.Set(out, fmt.Sprintf("candidates.0.content.parts.%d.text", partIndex), content.String())
|
||||
partIndex++
|
||||
}
|
||||
|
||||
// Handle tool calls
|
||||
@@ -609,27 +584,16 @@ func ConvertOpenAIResponseToGeminiNonStream(_ context.Context, _ string, origina
|
||||
functionName := function.Get("name").String()
|
||||
functionArgs := function.Get("arguments").String()
|
||||
|
||||
// Parse arguments
|
||||
var argsMap map[string]interface{}
|
||||
argsMap = parseArgsToMap(functionArgs)
|
||||
|
||||
functionCallPart := map[string]interface{}{
|
||||
"functionCall": map[string]interface{}{
|
||||
"name": functionName,
|
||||
"args": argsMap,
|
||||
},
|
||||
}
|
||||
parts = append(parts, functionCallPart)
|
||||
namePath := fmt.Sprintf("candidates.0.content.parts.%d.functionCall.name", partIndex)
|
||||
argsPath := fmt.Sprintf("candidates.0.content.parts.%d.functionCall.args", partIndex)
|
||||
out, _ = sjson.Set(out, namePath, functionName)
|
||||
out, _ = sjson.SetRaw(out, argsPath, parseArgsToObjectRaw(functionArgs))
|
||||
partIndex++
|
||||
}
|
||||
return true
|
||||
})
|
||||
}
|
||||
|
||||
// Set parts
|
||||
if len(parts) > 0 {
|
||||
out, _ = sjson.Set(out, "candidates.0.content.parts", parts)
|
||||
}
|
||||
|
||||
// Handle finish reason
|
||||
if finishReason := choice.Get("finish_reason"); finishReason.Exists() {
|
||||
geminiFinishReason := mapOpenAIFinishReasonToGemini(finishReason.String())
|
||||
@@ -645,15 +609,12 @@ func ConvertOpenAIResponseToGeminiNonStream(_ context.Context, _ string, origina
|
||||
|
||||
// Handle usage information
|
||||
if usage := root.Get("usage"); usage.Exists() {
|
||||
usageObj := map[string]interface{}{
|
||||
"promptTokenCount": usage.Get("prompt_tokens").Int(),
|
||||
"candidatesTokenCount": usage.Get("completion_tokens").Int(),
|
||||
"totalTokenCount": usage.Get("total_tokens").Int(),
|
||||
}
|
||||
out, _ = sjson.Set(out, "usageMetadata.promptTokenCount", usage.Get("prompt_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usageMetadata.candidatesTokenCount", usage.Get("completion_tokens").Int())
|
||||
out, _ = sjson.Set(out, "usageMetadata.totalTokenCount", usage.Get("total_tokens").Int())
|
||||
if reasoningTokens := reasoningTokensFromUsage(usage); reasoningTokens > 0 {
|
||||
usageObj["thoughtsTokenCount"] = reasoningTokens
|
||||
out, _ = sjson.Set(out, "usageMetadata.thoughtsTokenCount", reasoningTokens)
|
||||
}
|
||||
out, _ = sjson.Set(out, "usageMetadata", usageObj)
|
||||
}
|
||||
|
||||
return out
|
||||
|
||||
@@ -484,16 +484,12 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponses(ctx context.Context,
|
||||
}
|
||||
}
|
||||
// Build response.output using aggregated buffers
|
||||
var outputs []interface{}
|
||||
outputsWrapper := `{"arr":[]}`
|
||||
if st.ReasoningBuf.Len() > 0 {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": st.ReasoningID,
|
||||
"type": "reasoning",
|
||||
"summary": []interface{}{map[string]interface{}{
|
||||
"type": "summary_text",
|
||||
"text": st.ReasoningBuf.String(),
|
||||
}},
|
||||
})
|
||||
item := `{"id":"","type":"reasoning","summary":[{"type":"summary_text","text":""}]}`
|
||||
item, _ = sjson.Set(item, "id", st.ReasoningID)
|
||||
item, _ = sjson.Set(item, "summary.0.text", st.ReasoningBuf.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
// Append message items in ascending index order
|
||||
if len(st.MsgItemAdded) > 0 {
|
||||
@@ -513,18 +509,10 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponses(ctx context.Context,
|
||||
if b := st.MsgTextBuf[i]; b != nil {
|
||||
txt = b.String()
|
||||
}
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("msg_%s_%d", st.ResponseID, i),
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": txt,
|
||||
}},
|
||||
"role": "assistant",
|
||||
})
|
||||
item := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("msg_%s_%d", st.ResponseID, i))
|
||||
item, _ = sjson.Set(item, "content.0.text", txt)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
}
|
||||
if len(st.FuncArgsBuf) > 0 {
|
||||
@@ -547,18 +535,16 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponses(ctx context.Context,
|
||||
}
|
||||
callID := st.FuncCallIDs[i]
|
||||
name := st.FuncNames[i]
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", callID),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": args,
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
})
|
||||
item := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("fc_%s", callID))
|
||||
item, _ = sjson.Set(item, "arguments", args)
|
||||
item, _ = sjson.Set(item, "call_id", callID)
|
||||
item, _ = sjson.Set(item, "name", name)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
}
|
||||
if len(outputs) > 0 {
|
||||
completed, _ = sjson.Set(completed, "response.output", outputs)
|
||||
if gjson.Get(outputsWrapper, "arr.#").Int() > 0 {
|
||||
completed, _ = sjson.SetRaw(completed, "response.output", gjson.Get(outputsWrapper, "arr").Raw)
|
||||
}
|
||||
if st.UsageSeen {
|
||||
completed, _ = sjson.Set(completed, "response.usage.input_tokens", st.PromptTokens)
|
||||
@@ -681,7 +667,7 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponsesNonStream(_ context.Co
|
||||
}
|
||||
|
||||
// Build output list from choices[...]
|
||||
var outputs []interface{}
|
||||
outputsWrapper := `{"arr":[]}`
|
||||
// Detect and capture reasoning content if present
|
||||
rcText := gjson.GetBytes(rawJSON, "choices.0.message.reasoning_content").String()
|
||||
includeReasoning := rcText != ""
|
||||
@@ -693,21 +679,14 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponsesNonStream(_ context.Co
|
||||
if strings.HasPrefix(rid, "resp_") {
|
||||
rid = strings.TrimPrefix(rid, "resp_")
|
||||
}
|
||||
reasoningItem := map[string]interface{}{
|
||||
"id": fmt.Sprintf("rs_%s", rid),
|
||||
"type": "reasoning",
|
||||
"encrypted_content": "",
|
||||
}
|
||||
// Prefer summary_text from reasoning_content; encrypted_content is optional
|
||||
var summaries []interface{}
|
||||
reasoningItem := `{"id":"","type":"reasoning","encrypted_content":"","summary":[]}`
|
||||
reasoningItem, _ = sjson.Set(reasoningItem, "id", fmt.Sprintf("rs_%s", rid))
|
||||
if rcText != "" {
|
||||
summaries = append(summaries, map[string]interface{}{
|
||||
"type": "summary_text",
|
||||
"text": rcText,
|
||||
})
|
||||
reasoningItem, _ = sjson.Set(reasoningItem, "summary.0.type", "summary_text")
|
||||
reasoningItem, _ = sjson.Set(reasoningItem, "summary.0.text", rcText)
|
||||
}
|
||||
reasoningItem["summary"] = summaries
|
||||
outputs = append(outputs, reasoningItem)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", reasoningItem)
|
||||
}
|
||||
|
||||
if choices := root.Get("choices"); choices.Exists() && choices.IsArray() {
|
||||
@@ -716,18 +695,10 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponsesNonStream(_ context.Co
|
||||
if msg.Exists() {
|
||||
// Text message part
|
||||
if c := msg.Get("content"); c.Exists() && c.String() != "" {
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("msg_%s_%d", id, int(choice.Get("index").Int())),
|
||||
"type": "message",
|
||||
"status": "completed",
|
||||
"content": []interface{}{map[string]interface{}{
|
||||
"type": "output_text",
|
||||
"annotations": []interface{}{},
|
||||
"logprobs": []interface{}{},
|
||||
"text": c.String(),
|
||||
}},
|
||||
"role": "assistant",
|
||||
})
|
||||
item := `{"id":"","type":"message","status":"completed","content":[{"type":"output_text","annotations":[],"logprobs":[],"text":""}],"role":"assistant"}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("msg_%s_%d", id, int(choice.Get("index").Int())))
|
||||
item, _ = sjson.Set(item, "content.0.text", c.String())
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
}
|
||||
|
||||
// Function/tool calls
|
||||
@@ -736,14 +707,12 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponsesNonStream(_ context.Co
|
||||
callID := tc.Get("id").String()
|
||||
name := tc.Get("function.name").String()
|
||||
args := tc.Get("function.arguments").String()
|
||||
outputs = append(outputs, map[string]interface{}{
|
||||
"id": fmt.Sprintf("fc_%s", callID),
|
||||
"type": "function_call",
|
||||
"status": "completed",
|
||||
"arguments": args,
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
})
|
||||
item := `{"id":"","type":"function_call","status":"completed","arguments":"","call_id":"","name":""}`
|
||||
item, _ = sjson.Set(item, "id", fmt.Sprintf("fc_%s", callID))
|
||||
item, _ = sjson.Set(item, "arguments", args)
|
||||
item, _ = sjson.Set(item, "call_id", callID)
|
||||
item, _ = sjson.Set(item, "name", name)
|
||||
outputsWrapper, _ = sjson.SetRaw(outputsWrapper, "arr.-1", item)
|
||||
return true
|
||||
})
|
||||
}
|
||||
@@ -751,8 +720,8 @@ func ConvertOpenAIChatCompletionsResponseToOpenAIResponsesNonStream(_ context.Co
|
||||
return true
|
||||
})
|
||||
}
|
||||
if len(outputs) > 0 {
|
||||
resp, _ = sjson.Set(resp, "output", outputs)
|
||||
if gjson.Get(outputsWrapper, "arr.#").Int() > 0 {
|
||||
resp, _ = sjson.SetRaw(resp, "output", gjson.Get(outputsWrapper, "arr").Raw)
|
||||
}
|
||||
|
||||
// usage mapping
|
||||
|
||||
@@ -6,6 +6,7 @@ package usage
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
@@ -90,7 +91,7 @@ type modelStats struct {
|
||||
type RequestDetail struct {
|
||||
Timestamp time.Time `json:"timestamp"`
|
||||
Source string `json:"source"`
|
||||
AuthIndex uint64 `json:"auth_index"`
|
||||
AuthIndex string `json:"auth_index"`
|
||||
Tokens TokenStats `json:"tokens"`
|
||||
Failed bool `json:"failed"`
|
||||
}
|
||||
@@ -281,6 +282,118 @@ func (s *RequestStatistics) Snapshot() StatisticsSnapshot {
|
||||
return result
|
||||
}
|
||||
|
||||
type MergeResult struct {
|
||||
Added int64 `json:"added"`
|
||||
Skipped int64 `json:"skipped"`
|
||||
}
|
||||
|
||||
// MergeSnapshot merges an exported statistics snapshot into the current store.
|
||||
// Existing data is preserved and duplicate request details are skipped.
|
||||
func (s *RequestStatistics) MergeSnapshot(snapshot StatisticsSnapshot) MergeResult {
|
||||
result := MergeResult{}
|
||||
if s == nil {
|
||||
return result
|
||||
}
|
||||
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
seen := make(map[string]struct{})
|
||||
for apiName, stats := range s.apis {
|
||||
if stats == nil {
|
||||
continue
|
||||
}
|
||||
for modelName, modelStatsValue := range stats.Models {
|
||||
if modelStatsValue == nil {
|
||||
continue
|
||||
}
|
||||
for _, detail := range modelStatsValue.Details {
|
||||
seen[dedupKey(apiName, modelName, detail)] = struct{}{}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
for apiName, apiSnapshot := range snapshot.APIs {
|
||||
apiName = strings.TrimSpace(apiName)
|
||||
if apiName == "" {
|
||||
continue
|
||||
}
|
||||
stats, ok := s.apis[apiName]
|
||||
if !ok || stats == nil {
|
||||
stats = &apiStats{Models: make(map[string]*modelStats)}
|
||||
s.apis[apiName] = stats
|
||||
} else if stats.Models == nil {
|
||||
stats.Models = make(map[string]*modelStats)
|
||||
}
|
||||
for modelName, modelSnapshot := range apiSnapshot.Models {
|
||||
modelName = strings.TrimSpace(modelName)
|
||||
if modelName == "" {
|
||||
modelName = "unknown"
|
||||
}
|
||||
for _, detail := range modelSnapshot.Details {
|
||||
detail.Tokens = normaliseTokenStats(detail.Tokens)
|
||||
if detail.Timestamp.IsZero() {
|
||||
detail.Timestamp = time.Now()
|
||||
}
|
||||
key := dedupKey(apiName, modelName, detail)
|
||||
if _, exists := seen[key]; exists {
|
||||
result.Skipped++
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
s.recordImported(apiName, modelName, stats, detail)
|
||||
result.Added++
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return result
|
||||
}
|
||||
|
||||
func (s *RequestStatistics) recordImported(apiName, modelName string, stats *apiStats, detail RequestDetail) {
|
||||
totalTokens := detail.Tokens.TotalTokens
|
||||
if totalTokens < 0 {
|
||||
totalTokens = 0
|
||||
}
|
||||
|
||||
s.totalRequests++
|
||||
if detail.Failed {
|
||||
s.failureCount++
|
||||
} else {
|
||||
s.successCount++
|
||||
}
|
||||
s.totalTokens += totalTokens
|
||||
|
||||
s.updateAPIStats(stats, modelName, detail)
|
||||
|
||||
dayKey := detail.Timestamp.Format("2006-01-02")
|
||||
hourKey := detail.Timestamp.Hour()
|
||||
|
||||
s.requestsByDay[dayKey]++
|
||||
s.requestsByHour[hourKey]++
|
||||
s.tokensByDay[dayKey] += totalTokens
|
||||
s.tokensByHour[hourKey] += totalTokens
|
||||
}
|
||||
|
||||
func dedupKey(apiName, modelName string, detail RequestDetail) string {
|
||||
timestamp := detail.Timestamp.UTC().Format(time.RFC3339Nano)
|
||||
tokens := normaliseTokenStats(detail.Tokens)
|
||||
return fmt.Sprintf(
|
||||
"%s|%s|%s|%s|%s|%t|%d|%d|%d|%d|%d",
|
||||
apiName,
|
||||
modelName,
|
||||
timestamp,
|
||||
detail.Source,
|
||||
detail.AuthIndex,
|
||||
detail.Failed,
|
||||
tokens.InputTokens,
|
||||
tokens.OutputTokens,
|
||||
tokens.ReasoningTokens,
|
||||
tokens.CachedTokens,
|
||||
tokens.TotalTokens,
|
||||
)
|
||||
}
|
||||
|
||||
func resolveAPIIdentifier(ctx context.Context, record coreusage.Record) string {
|
||||
if ctx != nil {
|
||||
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil {
|
||||
@@ -340,6 +453,16 @@ func normaliseDetail(detail coreusage.Detail) TokenStats {
|
||||
return tokens
|
||||
}
|
||||
|
||||
func normaliseTokenStats(tokens TokenStats) TokenStats {
|
||||
if tokens.TotalTokens == 0 {
|
||||
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens
|
||||
}
|
||||
if tokens.TotalTokens == 0 {
|
||||
tokens.TotalTokens = tokens.InputTokens + tokens.OutputTokens + tokens.ReasoningTokens + tokens.CachedTokens
|
||||
}
|
||||
return tokens
|
||||
}
|
||||
|
||||
func formatHour(hour int) string {
|
||||
if hour < 0 {
|
||||
hour = 0
|
||||
|
||||
@@ -56,6 +56,21 @@ func ComputeClaudeModelsHash(models []config.ClaudeModel) string {
|
||||
return hashJoined(keys)
|
||||
}
|
||||
|
||||
// ComputeCodexModelsHash returns a stable hash for Codex model aliases.
|
||||
func ComputeCodexModelsHash(models []config.CodexModel) string {
|
||||
keys := normalizeModelPairs(func(out func(key string)) {
|
||||
for _, model := range models {
|
||||
name := strings.TrimSpace(model.Name)
|
||||
alias := strings.TrimSpace(model.Alias)
|
||||
if name == "" && alias == "" {
|
||||
continue
|
||||
}
|
||||
out(strings.ToLower(name) + "|" + strings.ToLower(alias))
|
||||
}
|
||||
})
|
||||
return hashJoined(keys)
|
||||
}
|
||||
|
||||
// ComputeExcludedModelsHash returns a normalized hash for excluded model lists.
|
||||
func ComputeExcludedModelsHash(excluded []string) string {
|
||||
if len(excluded) == 0 {
|
||||
|
||||
@@ -81,6 +81,15 @@ func TestComputeClaudeModelsHash_Empty(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeCodexModelsHash_Empty(t *testing.T) {
|
||||
if got := ComputeCodexModelsHash(nil); got != "" {
|
||||
t.Fatalf("expected empty hash for nil models, got %q", got)
|
||||
}
|
||||
if got := ComputeCodexModelsHash([]config.CodexModel{}); got != "" {
|
||||
t.Fatalf("expected empty hash for empty slice, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeClaudeModelsHash_IgnoresBlankAndDedup(t *testing.T) {
|
||||
a := []config.ClaudeModel{
|
||||
{Name: "m1", Alias: "a1"},
|
||||
@@ -95,6 +104,20 @@ func TestComputeClaudeModelsHash_IgnoresBlankAndDedup(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeCodexModelsHash_IgnoresBlankAndDedup(t *testing.T) {
|
||||
a := []config.CodexModel{
|
||||
{Name: "m1", Alias: "a1"},
|
||||
{Name: " "},
|
||||
{Name: "M1", Alias: "A1"},
|
||||
}
|
||||
b := []config.CodexModel{
|
||||
{Name: "m1", Alias: "a1"},
|
||||
}
|
||||
if h1, h2 := ComputeCodexModelsHash(a), ComputeCodexModelsHash(b); h1 == "" || h1 != h2 {
|
||||
t.Fatalf("expected same hash ignoring blanks/dupes, got %q / %q", h1, h2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeExcludedModelsHash_Normalizes(t *testing.T) {
|
||||
hash1 := ComputeExcludedModelsHash([]string{" A ", "b", "a"})
|
||||
hash2 := ComputeExcludedModelsHash([]string{"a", " b", "A"})
|
||||
@@ -157,3 +180,15 @@ func TestComputeClaudeModelsHash_Deterministic(t *testing.T) {
|
||||
t.Fatalf("expected different hash when models change, got %s", h3)
|
||||
}
|
||||
}
|
||||
|
||||
func TestComputeCodexModelsHash_Deterministic(t *testing.T) {
|
||||
models := []config.CodexModel{{Name: "a", Alias: "A"}, {Name: "b"}}
|
||||
h1 := ComputeCodexModelsHash(models)
|
||||
h2 := ComputeCodexModelsHash(models)
|
||||
if h1 == "" || h1 != h2 {
|
||||
t.Fatalf("expected deterministic hash, got %s / %s", h1, h2)
|
||||
}
|
||||
if h3 := ComputeCodexModelsHash([]config.CodexModel{{Name: "a"}}); h3 == h1 {
|
||||
t.Fatalf("expected different hash when models change, got %s", h3)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -151,6 +151,9 @@ func (s *ConfigSynthesizer) synthesizeCodexKeys(ctx *SynthesisContext) []*coreau
|
||||
if ck.BaseURL != "" {
|
||||
attrs["base_url"] = ck.BaseURL
|
||||
}
|
||||
if hash := diff.ComputeCodexModelsHash(ck.Models); hash != "" {
|
||||
attrs["models_hash"] = hash
|
||||
}
|
||||
addConfigHeadersToAttrs(ck.Headers, attrs)
|
||||
proxyURL := strings.TrimSpace(ck.ProxyURL)
|
||||
a := &coreauth.Auth{
|
||||
|
||||
@@ -14,7 +14,6 @@ import (
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
|
||||
@@ -185,14 +184,6 @@ func (h *ClaudeCodeAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSO
|
||||
// - c: The Gin context for the request.
|
||||
// - rawJSON: The raw JSON request body.
|
||||
func (h *ClaudeCodeAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
// Set up Server-Sent Events (SSE) headers for streaming response
|
||||
// These headers are essential for maintaining a persistent connection
|
||||
// and enabling real-time streaming of chat completions
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
// This is crucial for streaming as it allows immediate sending of data chunks
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
@@ -213,58 +204,82 @@ func (h *ClaudeCodeAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON [
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
|
||||
dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
|
||||
h.forwardClaudeStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
func (h *ClaudeCodeAPIHandler) forwardClaudeStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
// OpenAI-style stream forwarding: write each SSE chunk and flush immediately.
|
||||
// This guarantees clients see incremental output even for small responses.
|
||||
// Peek at the first chunk to determine success or failure before setting headers
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
|
||||
case chunk, ok := <-data:
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
// Stream closed without data? Send DONE or just headers.
|
||||
setSSEHeaders()
|
||||
flusher.Flush()
|
||||
cancel(nil)
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers now.
|
||||
setSSEHeaders()
|
||||
|
||||
// Write the first chunk
|
||||
if len(chunk) > 0 {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
c.Status(status)
|
||||
|
||||
// An error occurred: emit as a proper SSE error event
|
||||
errorBytes, _ := json.Marshal(h.toClaudeError(errMsg))
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", errorBytes)
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
}
|
||||
cancel(execErr)
|
||||
// Continue streaming the rest
|
||||
h.forwardClaudeStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *ClaudeCodeAPIHandler) forwardClaudeStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if len(chunk) == 0 {
|
||||
return
|
||||
}
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
c.Status(status)
|
||||
|
||||
errorBytes, _ := json.Marshal(h.toClaudeError(errMsg))
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", errorBytes)
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
type claudeErrorDetail struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
|
||||
@@ -182,19 +182,18 @@ func (h *GeminiCLIAPIHandler) handleInternalGenerateContent(c *gin.Context, rawJ
|
||||
}
|
||||
|
||||
func (h *GeminiCLIAPIHandler) forwardCLIStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
cancel(nil)
|
||||
return
|
||||
}
|
||||
var keepAliveInterval *time.Duration
|
||||
if alt != "" {
|
||||
disabled := time.Duration(0)
|
||||
keepAliveInterval = &disabled
|
||||
}
|
||||
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
KeepAliveInterval: keepAliveInterval,
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if alt == "" {
|
||||
if bytes.Equal(chunk, []byte("data: [DONE]")) || bytes.Equal(chunk, []byte("[DONE]")) {
|
||||
continue
|
||||
return
|
||||
}
|
||||
|
||||
if !bytes.HasPrefix(chunk, []byte("data:")) {
|
||||
@@ -206,22 +205,25 @@ func (h *GeminiCLIAPIHandler) forwardCLIStream(c *gin.Context, flusher http.Flus
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
flusher.Flush()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
flusher.Flush()
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
cancel(execErr)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
if alt == "" {
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -226,13 +226,6 @@ func (h *GeminiAPIHandler) GeminiHandler(c *gin.Context) {
|
||||
func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
alt := h.GetAlt(c)
|
||||
|
||||
if alt == "" {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
@@ -247,8 +240,65 @@ func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName
|
||||
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
h.forwardGeminiStream(c, flusher, alt, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
// Closed without data
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
|
||||
// Write first chunk
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
flusher.Flush()
|
||||
|
||||
// Continue
|
||||
h.forwardGeminiStream(c, flusher, alt, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleCountTokens handles token counting requests for Gemini models.
|
||||
@@ -297,16 +347,15 @@ func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName strin
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) forwardGeminiStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
cancel(nil)
|
||||
return
|
||||
}
|
||||
var keepAliveInterval *time.Duration
|
||||
if alt != "" {
|
||||
disabled := time.Duration(0)
|
||||
keepAliveInterval = &disabled
|
||||
}
|
||||
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
KeepAliveInterval: keepAliveInterval,
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
@@ -314,22 +363,25 @@ func (h *GeminiAPIHandler) forwardGeminiStream(c *gin.Context, flusher http.Flus
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
flusher.Flush()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
flusher.Flush()
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
cancel(execErr)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
if alt == "" {
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -9,9 +9,12 @@ import (
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
@@ -40,6 +43,117 @@ type ErrorDetail struct {
|
||||
Code string `json:"code,omitempty"`
|
||||
}
|
||||
|
||||
const idempotencyKeyMetadataKey = "idempotency_key"
|
||||
|
||||
const (
|
||||
defaultStreamingKeepAliveSeconds = 0
|
||||
defaultStreamingBootstrapRetries = 0
|
||||
)
|
||||
|
||||
// BuildErrorResponseBody builds an OpenAI-compatible JSON error response body.
|
||||
// If errText is already valid JSON, it is returned as-is to preserve upstream error payloads.
|
||||
func BuildErrorResponseBody(status int, errText string) []byte {
|
||||
if status <= 0 {
|
||||
status = http.StatusInternalServerError
|
||||
}
|
||||
if strings.TrimSpace(errText) == "" {
|
||||
errText = http.StatusText(status)
|
||||
}
|
||||
|
||||
trimmed := strings.TrimSpace(errText)
|
||||
if trimmed != "" && json.Valid([]byte(trimmed)) {
|
||||
return []byte(trimmed)
|
||||
}
|
||||
|
||||
errType := "invalid_request_error"
|
||||
var code string
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
errType = "authentication_error"
|
||||
code = "invalid_api_key"
|
||||
case http.StatusForbidden:
|
||||
errType = "permission_error"
|
||||
code = "insufficient_quota"
|
||||
case http.StatusTooManyRequests:
|
||||
errType = "rate_limit_error"
|
||||
code = "rate_limit_exceeded"
|
||||
case http.StatusNotFound:
|
||||
errType = "invalid_request_error"
|
||||
code = "model_not_found"
|
||||
default:
|
||||
if status >= http.StatusInternalServerError {
|
||||
errType = "server_error"
|
||||
code = "internal_server_error"
|
||||
}
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(ErrorResponse{
|
||||
Error: ErrorDetail{
|
||||
Message: errText,
|
||||
Type: errType,
|
||||
Code: code,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return []byte(fmt.Sprintf(`{"error":{"message":%q,"type":"server_error","code":"internal_server_error"}}`, errText))
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
// StreamingKeepAliveInterval returns the SSE keep-alive interval for this server.
|
||||
// Returning 0 disables keep-alives (default when unset).
|
||||
func StreamingKeepAliveInterval(cfg *config.SDKConfig) time.Duration {
|
||||
seconds := defaultStreamingKeepAliveSeconds
|
||||
if cfg != nil {
|
||||
seconds = cfg.Streaming.KeepAliveSeconds
|
||||
}
|
||||
if seconds <= 0 {
|
||||
return 0
|
||||
}
|
||||
return time.Duration(seconds) * time.Second
|
||||
}
|
||||
|
||||
// StreamingBootstrapRetries returns how many times a streaming request may be retried before any bytes are sent.
|
||||
func StreamingBootstrapRetries(cfg *config.SDKConfig) int {
|
||||
retries := defaultStreamingBootstrapRetries
|
||||
if cfg != nil {
|
||||
retries = cfg.Streaming.BootstrapRetries
|
||||
}
|
||||
if retries < 0 {
|
||||
retries = 0
|
||||
}
|
||||
return retries
|
||||
}
|
||||
|
||||
func requestExecutionMetadata(ctx context.Context) map[string]any {
|
||||
// Idempotency-Key is an optional client-supplied header used to correlate retries.
|
||||
// It is forwarded as execution metadata; when absent we generate a UUID.
|
||||
key := ""
|
||||
if ctx != nil {
|
||||
if ginCtx, ok := ctx.Value("gin").(*gin.Context); ok && ginCtx != nil && ginCtx.Request != nil {
|
||||
key = strings.TrimSpace(ginCtx.GetHeader("Idempotency-Key"))
|
||||
}
|
||||
}
|
||||
if key == "" {
|
||||
key = uuid.NewString()
|
||||
}
|
||||
return map[string]any{idempotencyKeyMetadataKey: key}
|
||||
}
|
||||
|
||||
func mergeMetadata(base, overlay map[string]any) map[string]any {
|
||||
if len(base) == 0 && len(overlay) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(base)+len(overlay))
|
||||
for k, v := range base {
|
||||
out[k] = v
|
||||
}
|
||||
for k, v := range overlay {
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// BaseAPIHandler contains the handlers for API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service and manages
|
||||
// load balancing, client selection, and configuration.
|
||||
@@ -104,13 +218,39 @@ func (h *BaseAPIHandler) GetAlt(c *gin.Context) string {
|
||||
// Parameters:
|
||||
// - handler: The API handler associated with the request.
|
||||
// - c: The Gin context of the current request.
|
||||
// - ctx: The parent context.
|
||||
// - ctx: The parent context (caller values/deadlines are preserved; request context adds cancellation and request ID).
|
||||
//
|
||||
// Returns:
|
||||
// - context.Context: The new context with cancellation and embedded values.
|
||||
// - APIHandlerCancelFunc: A function to cancel the context and log the response.
|
||||
func (h *BaseAPIHandler) GetContextWithCancel(handler interfaces.APIHandler, c *gin.Context, ctx context.Context) (context.Context, APIHandlerCancelFunc) {
|
||||
newCtx, cancel := context.WithCancel(ctx)
|
||||
parentCtx := ctx
|
||||
if parentCtx == nil {
|
||||
parentCtx = context.Background()
|
||||
}
|
||||
|
||||
var requestCtx context.Context
|
||||
if c != nil && c.Request != nil {
|
||||
requestCtx = c.Request.Context()
|
||||
}
|
||||
|
||||
if requestCtx != nil && logging.GetRequestID(parentCtx) == "" {
|
||||
if requestID := logging.GetRequestID(requestCtx); requestID != "" {
|
||||
parentCtx = logging.WithRequestID(parentCtx, requestID)
|
||||
} else if requestID := logging.GetGinRequestID(c); requestID != "" {
|
||||
parentCtx = logging.WithRequestID(parentCtx, requestID)
|
||||
}
|
||||
}
|
||||
newCtx, cancel := context.WithCancel(parentCtx)
|
||||
if requestCtx != nil && requestCtx != parentCtx {
|
||||
go func() {
|
||||
select {
|
||||
case <-requestCtx.Done():
|
||||
cancel()
|
||||
case <-newCtx.Done():
|
||||
}
|
||||
}()
|
||||
}
|
||||
newCtx = context.WithValue(newCtx, "gin", c)
|
||||
newCtx = context.WithValue(newCtx, "handler", handler)
|
||||
return newCtx, func(params ...interface{}) {
|
||||
@@ -183,6 +323,7 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType
|
||||
if errMsg != nil {
|
||||
return nil, errMsg
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: cloneBytes(rawJSON),
|
||||
@@ -196,9 +337,7 @@ func (h *BaseAPIHandler) ExecuteWithAuthManager(ctx context.Context, handlerType
|
||||
OriginalRequest: cloneBytes(rawJSON),
|
||||
SourceFormat: sdktranslator.FromString(handlerType),
|
||||
}
|
||||
if cloned := cloneMetadata(metadata); cloned != nil {
|
||||
opts.Metadata = cloned
|
||||
}
|
||||
opts.Metadata = mergeMetadata(cloneMetadata(metadata), reqMeta)
|
||||
resp, err := h.AuthManager.Execute(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
@@ -225,6 +364,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle
|
||||
if errMsg != nil {
|
||||
return nil, errMsg
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: cloneBytes(rawJSON),
|
||||
@@ -238,9 +378,7 @@ func (h *BaseAPIHandler) ExecuteCountWithAuthManager(ctx context.Context, handle
|
||||
OriginalRequest: cloneBytes(rawJSON),
|
||||
SourceFormat: sdktranslator.FromString(handlerType),
|
||||
}
|
||||
if cloned := cloneMetadata(metadata); cloned != nil {
|
||||
opts.Metadata = cloned
|
||||
}
|
||||
opts.Metadata = mergeMetadata(cloneMetadata(metadata), reqMeta)
|
||||
resp, err := h.AuthManager.ExecuteCount(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
@@ -270,6 +408,7 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl
|
||||
close(errChan)
|
||||
return nil, errChan
|
||||
}
|
||||
reqMeta := requestExecutionMetadata(ctx)
|
||||
req := coreexecutor.Request{
|
||||
Model: normalizedModel,
|
||||
Payload: cloneBytes(rawJSON),
|
||||
@@ -283,9 +422,7 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl
|
||||
OriginalRequest: cloneBytes(rawJSON),
|
||||
SourceFormat: sdktranslator.FromString(handlerType),
|
||||
}
|
||||
if cloned := cloneMetadata(metadata); cloned != nil {
|
||||
opts.Metadata = cloned
|
||||
}
|
||||
opts.Metadata = mergeMetadata(cloneMetadata(metadata), reqMeta)
|
||||
chunks, err := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
|
||||
if err != nil {
|
||||
errChan := make(chan *interfaces.ErrorMessage, 1)
|
||||
@@ -310,31 +447,94 @@ func (h *BaseAPIHandler) ExecuteStreamWithAuthManager(ctx context.Context, handl
|
||||
go func() {
|
||||
defer close(dataChan)
|
||||
defer close(errChan)
|
||||
for chunk := range chunks {
|
||||
if chunk.Err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
if se, ok := chunk.Err.(interface{ StatusCode() int }); ok && se != nil {
|
||||
if code := se.StatusCode(); code > 0 {
|
||||
status = code
|
||||
}
|
||||
}
|
||||
var addon http.Header
|
||||
if he, ok := chunk.Err.(interface{ Headers() http.Header }); ok && he != nil {
|
||||
if hdr := he.Headers(); hdr != nil {
|
||||
addon = hdr.Clone()
|
||||
}
|
||||
}
|
||||
errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: chunk.Err, Addon: addon}
|
||||
return
|
||||
sentPayload := false
|
||||
bootstrapRetries := 0
|
||||
maxBootstrapRetries := StreamingBootstrapRetries(h.Cfg)
|
||||
|
||||
bootstrapEligible := func(err error) bool {
|
||||
status := statusFromError(err)
|
||||
if status == 0 {
|
||||
return true
|
||||
}
|
||||
if len(chunk.Payload) > 0 {
|
||||
dataChan <- cloneBytes(chunk.Payload)
|
||||
switch status {
|
||||
case http.StatusUnauthorized, http.StatusForbidden, http.StatusPaymentRequired,
|
||||
http.StatusRequestTimeout, http.StatusTooManyRequests:
|
||||
return true
|
||||
default:
|
||||
return status >= http.StatusInternalServerError
|
||||
}
|
||||
}
|
||||
|
||||
outer:
|
||||
for {
|
||||
for {
|
||||
var chunk coreexecutor.StreamChunk
|
||||
var ok bool
|
||||
if ctx != nil {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case chunk, ok = <-chunks:
|
||||
}
|
||||
} else {
|
||||
chunk, ok = <-chunks
|
||||
}
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
if chunk.Err != nil {
|
||||
streamErr := chunk.Err
|
||||
// Safe bootstrap recovery: if the upstream fails before any payload bytes are sent,
|
||||
// retry a few times (to allow auth rotation / transient recovery) and then attempt model fallback.
|
||||
if !sentPayload {
|
||||
if bootstrapRetries < maxBootstrapRetries && bootstrapEligible(streamErr) {
|
||||
bootstrapRetries++
|
||||
retryChunks, retryErr := h.AuthManager.ExecuteStream(ctx, providers, req, opts)
|
||||
if retryErr == nil {
|
||||
chunks = retryChunks
|
||||
continue outer
|
||||
}
|
||||
streamErr = retryErr
|
||||
}
|
||||
}
|
||||
|
||||
status := http.StatusInternalServerError
|
||||
if se, ok := streamErr.(interface{ StatusCode() int }); ok && se != nil {
|
||||
if code := se.StatusCode(); code > 0 {
|
||||
status = code
|
||||
}
|
||||
}
|
||||
var addon http.Header
|
||||
if he, ok := streamErr.(interface{ Headers() http.Header }); ok && he != nil {
|
||||
if hdr := he.Headers(); hdr != nil {
|
||||
addon = hdr.Clone()
|
||||
}
|
||||
}
|
||||
errChan <- &interfaces.ErrorMessage{StatusCode: status, Error: streamErr, Addon: addon}
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) > 0 {
|
||||
sentPayload = true
|
||||
dataChan <- cloneBytes(chunk.Payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
return dataChan, errChan
|
||||
}
|
||||
|
||||
func statusFromError(err error) int {
|
||||
if err == nil {
|
||||
return 0
|
||||
}
|
||||
if se, ok := err.(interface{ StatusCode() int }); ok && se != nil {
|
||||
if code := se.StatusCode(); code > 0 {
|
||||
return code
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) getRequestDetails(modelName string) (providers []string, normalizedModel string, metadata map[string]any, err *interfaces.ErrorMessage) {
|
||||
// Resolve "auto" model to an actual available model first
|
||||
resolvedModelName := util.ResolveAutoModel(modelName)
|
||||
@@ -418,38 +618,7 @@ func (h *BaseAPIHandler) WriteErrorResponse(c *gin.Context, msg *interfaces.Erro
|
||||
}
|
||||
}
|
||||
|
||||
// Prefer preserving upstream JSON error bodies when possible.
|
||||
buildJSONBody := func() []byte {
|
||||
trimmed := strings.TrimSpace(errText)
|
||||
if trimmed != "" && json.Valid([]byte(trimmed)) {
|
||||
return []byte(trimmed)
|
||||
}
|
||||
errType := "invalid_request_error"
|
||||
switch status {
|
||||
case http.StatusUnauthorized:
|
||||
errType = "authentication_error"
|
||||
case http.StatusForbidden:
|
||||
errType = "permission_error"
|
||||
case http.StatusTooManyRequests:
|
||||
errType = "rate_limit_error"
|
||||
default:
|
||||
if status >= http.StatusInternalServerError {
|
||||
errType = "server_error"
|
||||
}
|
||||
}
|
||||
payload, err := json.Marshal(ErrorResponse{
|
||||
Error: ErrorDetail{
|
||||
Message: errText,
|
||||
Type: errType,
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
return []byte(fmt.Sprintf(`{"error":{"message":%q,"type":"server_error"}}`, errText))
|
||||
}
|
||||
return payload
|
||||
}
|
||||
|
||||
body := buildJSONBody()
|
||||
body := BuildErrorResponseBody(status, errText)
|
||||
c.Set("API_RESPONSE", bytes.Clone(body))
|
||||
|
||||
if !c.Writer.Written() {
|
||||
|
||||
124
sdk/api/handlers/handlers_stream_bootstrap_test.go
Normal file
124
sdk/api/handlers/handlers_stream_bootstrap_test.go
Normal file
@@ -0,0 +1,124 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v6/sdk/config"
|
||||
)
|
||||
|
||||
type failOnceStreamExecutor struct {
|
||||
mu sync.Mutex
|
||||
calls int
|
||||
}
|
||||
|
||||
func (e *failOnceStreamExecutor) Identifier() string { return "codex" }
|
||||
|
||||
func (e *failOnceStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "Execute not implemented"}
|
||||
}
|
||||
|
||||
func (e *failOnceStreamExecutor) ExecuteStream(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (<-chan coreexecutor.StreamChunk, error) {
|
||||
e.mu.Lock()
|
||||
e.calls++
|
||||
call := e.calls
|
||||
e.mu.Unlock()
|
||||
|
||||
ch := make(chan coreexecutor.StreamChunk, 1)
|
||||
if call == 1 {
|
||||
ch <- coreexecutor.StreamChunk{
|
||||
Err: &coreauth.Error{
|
||||
Code: "unauthorized",
|
||||
Message: "unauthorized",
|
||||
Retryable: false,
|
||||
HTTPStatus: http.StatusUnauthorized,
|
||||
},
|
||||
}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
ch <- coreexecutor.StreamChunk{Payload: []byte("ok")}
|
||||
close(ch)
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
func (e *failOnceStreamExecutor) Refresh(ctx context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (e *failOnceStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, &coreauth.Error{Code: "not_implemented", Message: "CountTokens not implemented"}
|
||||
}
|
||||
|
||||
func (e *failOnceStreamExecutor) Calls() int {
|
||||
e.mu.Lock()
|
||||
defer e.mu.Unlock()
|
||||
return e.calls
|
||||
}
|
||||
|
||||
func TestExecuteStreamWithAuthManager_RetriesBeforeFirstByte(t *testing.T) {
|
||||
executor := &failOnceStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
|
||||
auth1 := &coreauth.Auth{
|
||||
ID: "auth1",
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Metadata: map[string]any{"email": "test1@example.com"},
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth1); err != nil {
|
||||
t.Fatalf("manager.Register(auth1): %v", err)
|
||||
}
|
||||
|
||||
auth2 := &coreauth.Auth{
|
||||
ID: "auth2",
|
||||
Provider: "codex",
|
||||
Status: coreauth.StatusActive,
|
||||
Metadata: map[string]any{"email": "test2@example.com"},
|
||||
}
|
||||
if _, err := manager.Register(context.Background(), auth2); err != nil {
|
||||
t.Fatalf("manager.Register(auth2): %v", err)
|
||||
}
|
||||
|
||||
registry.GetGlobalRegistry().RegisterClient(auth1.ID, auth1.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
registry.GetGlobalRegistry().RegisterClient(auth2.ID, auth2.Provider, []*registry.ModelInfo{{ID: "test-model"}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth1.ID)
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth2.ID)
|
||||
})
|
||||
|
||||
handler := NewBaseAPIHandlers(&sdkconfig.SDKConfig{
|
||||
Streaming: sdkconfig.StreamingConfig{
|
||||
BootstrapRetries: 1,
|
||||
},
|
||||
}, manager)
|
||||
dataChan, errChan := handler.ExecuteStreamWithAuthManager(context.Background(), "openai", "test-model", []byte(`{"model":"test-model"}`), "")
|
||||
if dataChan == nil || errChan == nil {
|
||||
t.Fatalf("expected non-nil channels")
|
||||
}
|
||||
|
||||
var got []byte
|
||||
for chunk := range dataChan {
|
||||
got = append(got, chunk...)
|
||||
}
|
||||
|
||||
for msg := range errChan {
|
||||
if msg != nil {
|
||||
t.Fatalf("unexpected error: %+v", msg)
|
||||
}
|
||||
}
|
||||
|
||||
if string(got) != "ok" {
|
||||
t.Fatalf("expected payload ok, got %q", string(got))
|
||||
}
|
||||
if executor.Calls() != 2 {
|
||||
t.Fatalf("expected 2 stream attempts, got %d", executor.Calls())
|
||||
}
|
||||
}
|
||||
@@ -11,12 +11,13 @@ import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
"sync"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
responsesconverter "github.com/router-for-me/CLIProxyAPI/v6/internal/translator/openai/openai/responses"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
@@ -109,7 +110,17 @@ func (h *OpenAIAPIHandler) ChatCompletions(c *gin.Context) {
|
||||
|
||||
// Check if the client requested a streaming response.
|
||||
streamResult := gjson.GetBytes(rawJSON, "stream")
|
||||
if streamResult.Type == gjson.True {
|
||||
stream := streamResult.Type == gjson.True
|
||||
|
||||
// Some clients send OpenAI Responses-format payloads to /v1/chat/completions.
|
||||
// Convert them to Chat Completions so downstream translators preserve tool metadata.
|
||||
if shouldTreatAsResponsesFormat(rawJSON) {
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
rawJSON = responsesconverter.ConvertOpenAIResponsesRequestToOpenAIChatCompletions(modelName, rawJSON, stream)
|
||||
stream = gjson.GetBytes(rawJSON, "stream").Bool()
|
||||
}
|
||||
|
||||
if stream {
|
||||
h.handleStreamingResponse(c, rawJSON)
|
||||
} else {
|
||||
h.handleNonStreamingResponse(c, rawJSON)
|
||||
@@ -117,6 +128,21 @@ func (h *OpenAIAPIHandler) ChatCompletions(c *gin.Context) {
|
||||
|
||||
}
|
||||
|
||||
// shouldTreatAsResponsesFormat detects OpenAI Responses-style payloads that are
|
||||
// accidentally sent to the Chat Completions endpoint.
|
||||
func shouldTreatAsResponsesFormat(rawJSON []byte) bool {
|
||||
if gjson.GetBytes(rawJSON, "messages").Exists() {
|
||||
return false
|
||||
}
|
||||
if gjson.GetBytes(rawJSON, "input").Exists() {
|
||||
return true
|
||||
}
|
||||
if gjson.GetBytes(rawJSON, "instructions").Exists() {
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Completions handles the /v1/completions endpoint.
|
||||
// It determines whether the request is for a streaming or non-streaming response
|
||||
// and calls the appropriate handler based on the model provider.
|
||||
@@ -417,11 +443,6 @@ func (h *OpenAIAPIHandler) handleNonStreamingResponse(c *gin.Context, rawJSON []
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible request
|
||||
func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
@@ -437,7 +458,55 @@ func (h *OpenAIAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byt
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, h.GetAlt(c))
|
||||
h.handleStreamResult(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk to determine success or failure before setting headers
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
// Stream closed without data? Send DONE or just headers.
|
||||
setSSEHeaders()
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Commit to streaming headers.
|
||||
setSSEHeaders()
|
||||
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
|
||||
flusher.Flush()
|
||||
|
||||
// Continue streaming the rest
|
||||
h.handleStreamResult(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleCompletionsNonStreamingResponse handles non-streaming completions responses.
|
||||
@@ -474,11 +543,6 @@ func (h *OpenAIAPIHandler) handleCompletionsNonStreamingResponse(c *gin.Context,
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAI-compatible completions request
|
||||
func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
@@ -498,71 +562,109 @@ func (h *OpenAIAPIHandler) handleCompletionsStreamingResponse(c *gin.Context, ra
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, chatCompletionsJSON, "")
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, isOk := <-dataChan:
|
||||
if !isOk {
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
setSSEHeaders()
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
cliCancel()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
setSSEHeaders()
|
||||
|
||||
// Write the first chunk
|
||||
converted := convertChatCompletionsStreamChunkToCompletions(chunk)
|
||||
if converted != nil {
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(converted))
|
||||
flusher.Flush()
|
||||
}
|
||||
case errMsg, isOk := <-errChan:
|
||||
if !isOk {
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
flusher.Flush()
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
}
|
||||
cliCancel(execErr)
|
||||
|
||||
done := make(chan struct{})
|
||||
var doneOnce sync.Once
|
||||
stop := func() { doneOnce.Do(func() { close(done) }) }
|
||||
|
||||
convertedChan := make(chan []byte)
|
||||
go func() {
|
||||
defer close(convertedChan)
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
converted := convertChatCompletionsStreamChunkToCompletions(chunk)
|
||||
if converted == nil {
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case <-done:
|
||||
return
|
||||
case convertedChan <- converted:
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
h.handleStreamResult(c, flusher, func(err error) {
|
||||
stop()
|
||||
cliCancel(err)
|
||||
}, convertedChan, errChan)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
func (h *OpenAIAPIHandler) handleStreamResult(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
cancel(nil)
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(chunk))
|
||||
flusher.Flush()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
flusher.Flush()
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
}
|
||||
cancel(execErr)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", string(body))
|
||||
},
|
||||
WriteDone: func() {
|
||||
_, _ = fmt.Fprint(c.Writer, "data: [DONE]\n\n")
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -11,7 +11,6 @@ import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v6/internal/constant"
|
||||
@@ -128,11 +127,6 @@ func (h *OpenAIResponsesAPIHandler) handleNonStreamingResponse(c *gin.Context, r
|
||||
// - c: The Gin context containing the HTTP request and response
|
||||
// - rawJSON: The raw JSON bytes of the OpenAIResponses-compatible request
|
||||
func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJSON []byte) {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
@@ -149,46 +143,88 @@ func (h *OpenAIResponsesAPIHandler) handleStreamingResponse(c *gin.Context, rawJ
|
||||
modelName := gjson.GetBytes(rawJSON, "model").String()
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, "")
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) forwardResponsesStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
// Stream closed without data? Send headers and done.
|
||||
setSSEHeaders()
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
flusher.Flush()
|
||||
cancel(nil)
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
setSSEHeaders()
|
||||
|
||||
// Write first chunk logic (matching forwardResponsesStream)
|
||||
if bytes.HasPrefix(chunk, []byte("event:")) {
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
}
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
|
||||
flusher.Flush()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
flusher.Flush()
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
}
|
||||
cancel(execErr)
|
||||
|
||||
// Continue
|
||||
h.forwardResponsesStream(c, flusher, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
case <-time.After(500 * time.Millisecond):
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (h *OpenAIResponsesAPIHandler) forwardResponsesStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if bytes.HasPrefix(chunk, []byte("event:")) {
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
}
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
_, _ = fmt.Fprintf(c.Writer, "\nevent: error\ndata: %s\n\n", string(body))
|
||||
},
|
||||
WriteDone: func() {
|
||||
_, _ = c.Writer.Write([]byte("\n"))
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
121
sdk/api/handlers/stream_forwarder.go
Normal file
121
sdk/api/handlers/stream_forwarder.go
Normal file
@@ -0,0 +1,121 @@
|
||||
package handlers
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/interfaces"
|
||||
)
|
||||
|
||||
type StreamForwardOptions struct {
|
||||
// KeepAliveInterval overrides the configured streaming keep-alive interval.
|
||||
// If nil, the configured default is used. If set to <= 0, keep-alives are disabled.
|
||||
KeepAliveInterval *time.Duration
|
||||
|
||||
// WriteChunk writes a single data chunk to the response body. It should not flush.
|
||||
WriteChunk func(chunk []byte)
|
||||
|
||||
// WriteTerminalError writes an error payload to the response body when streaming fails
|
||||
// after headers have already been committed. It should not flush.
|
||||
WriteTerminalError func(errMsg *interfaces.ErrorMessage)
|
||||
|
||||
// WriteDone optionally writes a terminal marker when the upstream data channel closes
|
||||
// without an error (e.g. OpenAI's `[DONE]`). It should not flush.
|
||||
WriteDone func()
|
||||
|
||||
// WriteKeepAlive optionally writes a keep-alive heartbeat. It should not flush.
|
||||
// When nil, a standard SSE comment heartbeat is used.
|
||||
WriteKeepAlive func()
|
||||
}
|
||||
|
||||
func (h *BaseAPIHandler) ForwardStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage, opts StreamForwardOptions) {
|
||||
if c == nil {
|
||||
return
|
||||
}
|
||||
if cancel == nil {
|
||||
return
|
||||
}
|
||||
|
||||
writeChunk := opts.WriteChunk
|
||||
if writeChunk == nil {
|
||||
writeChunk = func([]byte) {}
|
||||
}
|
||||
|
||||
writeKeepAlive := opts.WriteKeepAlive
|
||||
if writeKeepAlive == nil {
|
||||
writeKeepAlive = func() {
|
||||
_, _ = c.Writer.Write([]byte(": keep-alive\n\n"))
|
||||
}
|
||||
}
|
||||
|
||||
keepAliveInterval := StreamingKeepAliveInterval(h.Cfg)
|
||||
if opts.KeepAliveInterval != nil {
|
||||
keepAliveInterval = *opts.KeepAliveInterval
|
||||
}
|
||||
var keepAlive *time.Ticker
|
||||
var keepAliveC <-chan time.Time
|
||||
if keepAliveInterval > 0 {
|
||||
keepAlive = time.NewTicker(keepAliveInterval)
|
||||
defer keepAlive.Stop()
|
||||
keepAliveC = keepAlive.C
|
||||
}
|
||||
|
||||
var terminalErr *interfaces.ErrorMessage
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cancel(c.Request.Context().Err())
|
||||
return
|
||||
case chunk, ok := <-data:
|
||||
if !ok {
|
||||
// Prefer surfacing a terminal error if one is pending.
|
||||
if terminalErr == nil {
|
||||
select {
|
||||
case errMsg, ok := <-errs:
|
||||
if ok && errMsg != nil {
|
||||
terminalErr = errMsg
|
||||
}
|
||||
default:
|
||||
}
|
||||
}
|
||||
if terminalErr != nil {
|
||||
if opts.WriteTerminalError != nil {
|
||||
opts.WriteTerminalError(terminalErr)
|
||||
}
|
||||
flusher.Flush()
|
||||
cancel(terminalErr.Error)
|
||||
return
|
||||
}
|
||||
if opts.WriteDone != nil {
|
||||
opts.WriteDone()
|
||||
}
|
||||
flusher.Flush()
|
||||
cancel(nil)
|
||||
return
|
||||
}
|
||||
writeChunk(chunk)
|
||||
flusher.Flush()
|
||||
case errMsg, ok := <-errs:
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
if errMsg != nil {
|
||||
terminalErr = errMsg
|
||||
if opts.WriteTerminalError != nil {
|
||||
opts.WriteTerminalError(errMsg)
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
var execErr error
|
||||
if errMsg != nil {
|
||||
execErr = errMsg.Error
|
||||
}
|
||||
cancel(execErr)
|
||||
return
|
||||
case <-keepAliveC:
|
||||
writeKeepAlive()
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
116
sdk/auth/kiro.go
116
sdk/auth/kiro.go
@@ -53,20 +53,8 @@ func (a *KiroAuthenticator) RefreshLead() *time.Duration {
|
||||
return &d
|
||||
}
|
||||
|
||||
// Login performs OAuth login for Kiro with AWS Builder ID.
|
||||
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
||||
}
|
||||
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
|
||||
// Use AWS Builder ID device code flow
|
||||
tokenData, err := oauth.LoginWithBuilderID(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("login failed: %w", err)
|
||||
}
|
||||
|
||||
// createAuthRecord creates an auth record from token data.
|
||||
func (a *KiroAuthenticator) createAuthRecord(tokenData *kiroauth.KiroTokenData, source string) (*coreauth.Auth, error) {
|
||||
// Parse expires_at
|
||||
expiresAt, err := time.Parse(time.RFC3339, tokenData.ExpiresAt)
|
||||
if err != nil {
|
||||
@@ -76,34 +64,63 @@ func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts
|
||||
// Extract identifier for file naming
|
||||
idPart := extractKiroIdentifier(tokenData.Email, tokenData.ProfileArn)
|
||||
|
||||
// Determine label based on auth method
|
||||
label := fmt.Sprintf("kiro-%s", source)
|
||||
if tokenData.AuthMethod == "idc" {
|
||||
label = "kiro-idc"
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
fileName := fmt.Sprintf("kiro-aws-%s.json", idPart)
|
||||
fileName := fmt.Sprintf("%s-%s.json", label, idPart)
|
||||
|
||||
metadata := map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"client_id": tokenData.ClientID,
|
||||
"client_secret": tokenData.ClientSecret,
|
||||
"email": tokenData.Email,
|
||||
}
|
||||
|
||||
// Add IDC-specific fields if present
|
||||
if tokenData.StartURL != "" {
|
||||
metadata["start_url"] = tokenData.StartURL
|
||||
}
|
||||
if tokenData.Region != "" {
|
||||
metadata["region"] = tokenData.Region
|
||||
}
|
||||
|
||||
attributes := map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": source,
|
||||
"email": tokenData.Email,
|
||||
}
|
||||
|
||||
// Add IDC-specific attributes if present
|
||||
if tokenData.AuthMethod == "idc" {
|
||||
attributes["source"] = "aws-idc"
|
||||
if tokenData.StartURL != "" {
|
||||
attributes["start_url"] = tokenData.StartURL
|
||||
}
|
||||
if tokenData.Region != "" {
|
||||
attributes["region"] = tokenData.Region
|
||||
}
|
||||
}
|
||||
|
||||
record := &coreauth.Auth{
|
||||
ID: fileName,
|
||||
Provider: "kiro",
|
||||
FileName: fileName,
|
||||
Label: "kiro-aws",
|
||||
Label: label,
|
||||
Status: coreauth.StatusActive,
|
||||
CreatedAt: now,
|
||||
UpdatedAt: now,
|
||||
Metadata: map[string]any{
|
||||
"type": "kiro",
|
||||
"access_token": tokenData.AccessToken,
|
||||
"refresh_token": tokenData.RefreshToken,
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"expires_at": tokenData.ExpiresAt,
|
||||
"auth_method": tokenData.AuthMethod,
|
||||
"provider": tokenData.Provider,
|
||||
"client_id": tokenData.ClientID,
|
||||
"client_secret": tokenData.ClientSecret,
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Attributes: map[string]string{
|
||||
"profile_arn": tokenData.ProfileArn,
|
||||
"source": "aws-builder-id",
|
||||
"email": tokenData.Email,
|
||||
},
|
||||
Metadata: metadata,
|
||||
Attributes: attributes,
|
||||
// NextRefreshAfter is aligned with RefreshLead (5min)
|
||||
NextRefreshAfter: expiresAt.Add(-5 * time.Minute),
|
||||
}
|
||||
@@ -117,6 +134,23 @@ func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts
|
||||
return record, nil
|
||||
}
|
||||
|
||||
// Login performs OAuth login for Kiro with AWS (Builder ID or IDC).
|
||||
// This shows a method selection prompt and handles both flows.
|
||||
func (a *KiroAuthenticator) Login(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
if cfg == nil {
|
||||
return nil, fmt.Errorf("kiro auth: configuration is required")
|
||||
}
|
||||
|
||||
// Use the unified method selection flow (Builder ID or IDC)
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
||||
tokenData, err := ssoClient.LoginWithMethodSelection(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("login failed: %w", err)
|
||||
}
|
||||
|
||||
return a.createAuthRecord(tokenData, "aws")
|
||||
}
|
||||
|
||||
// LoginWithAuthCode performs OAuth login for Kiro with AWS Builder ID using authorization code flow.
|
||||
// This provides a better UX than device code flow as it uses automatic browser callback.
|
||||
func (a *KiroAuthenticator) LoginWithAuthCode(ctx context.Context, cfg *config.Config, opts *LoginOptions) (*coreauth.Auth, error) {
|
||||
@@ -388,15 +422,23 @@ func (a *KiroAuthenticator) Refresh(ctx context.Context, cfg *config.Config, aut
|
||||
clientID, _ := auth.Metadata["client_id"].(string)
|
||||
clientSecret, _ := auth.Metadata["client_secret"].(string)
|
||||
authMethod, _ := auth.Metadata["auth_method"].(string)
|
||||
startURL, _ := auth.Metadata["start_url"].(string)
|
||||
region, _ := auth.Metadata["region"].(string)
|
||||
|
||||
var tokenData *kiroauth.KiroTokenData
|
||||
var err error
|
||||
|
||||
// Use SSO OIDC refresh for AWS Builder ID, otherwise use Kiro's OAuth refresh endpoint
|
||||
if clientID != "" && clientSecret != "" && authMethod == "builder-id" {
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
||||
ssoClient := kiroauth.NewSSOOIDCClient(cfg)
|
||||
|
||||
// Use SSO OIDC refresh for AWS Builder ID or IDC, otherwise use Kiro's OAuth refresh endpoint
|
||||
switch {
|
||||
case clientID != "" && clientSecret != "" && authMethod == "idc" && region != "":
|
||||
// IDC refresh with region-specific endpoint
|
||||
tokenData, err = ssoClient.RefreshTokenWithRegion(ctx, clientID, clientSecret, refreshToken, region, startURL)
|
||||
case clientID != "" && clientSecret != "" && authMethod == "builder-id":
|
||||
// Builder ID refresh with default endpoint
|
||||
tokenData, err = ssoClient.RefreshToken(ctx, clientID, clientSecret, refreshToken)
|
||||
} else {
|
||||
default:
|
||||
// Fallback to Kiro's refresh endpoint (for social auth: Google/GitHub)
|
||||
oauth := kiroauth.NewKiroOAuth(cfg)
|
||||
tokenData, err = oauth.RefreshToken(ctx, refreshToken)
|
||||
|
||||
@@ -12,6 +12,7 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/google/uuid"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/logging"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/util"
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
@@ -135,6 +136,18 @@ func NewManager(store Store, selector Selector, hook Hook) *Manager {
|
||||
}
|
||||
}
|
||||
|
||||
func (m *Manager) SetSelector(selector Selector) {
|
||||
if m == nil {
|
||||
return
|
||||
}
|
||||
if selector == nil {
|
||||
selector = &RoundRobinSelector{}
|
||||
}
|
||||
m.mu.Lock()
|
||||
m.selector = selector
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
// SetStore swaps the underlying persistence store.
|
||||
func (m *Manager) SetStore(store Store) {
|
||||
m.mu.Lock()
|
||||
@@ -190,10 +203,10 @@ func (m *Manager) Register(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
if auth == nil {
|
||||
return nil, nil
|
||||
}
|
||||
auth.EnsureIndex()
|
||||
if auth.ID == "" {
|
||||
auth.ID = uuid.NewString()
|
||||
}
|
||||
auth.EnsureIndex()
|
||||
m.mu.Lock()
|
||||
m.auths[auth.ID] = auth.Clone()
|
||||
m.mu.Unlock()
|
||||
@@ -208,7 +221,7 @@ func (m *Manager) Update(ctx context.Context, auth *Auth) (*Auth, error) {
|
||||
return nil, nil
|
||||
}
|
||||
m.mu.Lock()
|
||||
if existing, ok := m.auths[auth.ID]; ok && existing != nil && !auth.indexAssigned && auth.Index == 0 {
|
||||
if existing, ok := m.auths[auth.ID]; ok && existing != nil && !auth.indexAssigned && auth.Index == "" {
|
||||
auth.Index = existing.Index
|
||||
auth.indexAssigned = existing.indexAssigned
|
||||
}
|
||||
@@ -250,7 +263,6 @@ func (m *Manager) Execute(ctx context.Context, providers []string, req cliproxye
|
||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||
}
|
||||
rotated := m.rotateProviders(req.Model, normalized)
|
||||
defer m.advanceProviderCursor(req.Model, normalized)
|
||||
|
||||
retryTimes, maxWait := m.retrySettings()
|
||||
attempts := retryTimes + 1
|
||||
@@ -289,7 +301,6 @@ func (m *Manager) ExecuteCount(ctx context.Context, providers []string, req clip
|
||||
return cliproxyexecutor.Response{}, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||
}
|
||||
rotated := m.rotateProviders(req.Model, normalized)
|
||||
defer m.advanceProviderCursor(req.Model, normalized)
|
||||
|
||||
retryTimes, maxWait := m.retrySettings()
|
||||
attempts := retryTimes + 1
|
||||
@@ -328,7 +339,6 @@ func (m *Manager) ExecuteStream(ctx context.Context, providers []string, req cli
|
||||
return nil, &Error{Code: "provider_not_found", Message: "no provider supplied"}
|
||||
}
|
||||
rotated := m.rotateProviders(req.Model, normalized)
|
||||
defer m.advanceProviderCursor(req.Model, normalized)
|
||||
|
||||
retryTimes, maxWait := m.retrySettings()
|
||||
attempts := retryTimes + 1
|
||||
@@ -377,17 +387,18 @@ func (m *Manager) executeWithProvider(ctx context.Context, provider string, req
|
||||
|
||||
accountType, accountInfo := auth.AccountInfo()
|
||||
proxyInfo := auth.ProxyInfo()
|
||||
entry := logEntryWithRequestID(ctx)
|
||||
if accountType == "api_key" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
entry.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
entry.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
}
|
||||
} else if accountType == "oauth" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
entry.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
entry.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -437,17 +448,18 @@ func (m *Manager) executeCountWithProvider(ctx context.Context, provider string,
|
||||
|
||||
accountType, accountInfo := auth.AccountInfo()
|
||||
proxyInfo := auth.ProxyInfo()
|
||||
entry := logEntryWithRequestID(ctx)
|
||||
if accountType == "api_key" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
entry.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
entry.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
}
|
||||
} else if accountType == "oauth" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
entry.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
entry.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -497,17 +509,18 @@ func (m *Manager) executeStreamWithProvider(ctx context.Context, provider string
|
||||
|
||||
accountType, accountInfo := auth.AccountInfo()
|
||||
proxyInfo := auth.ProxyInfo()
|
||||
entry := logEntryWithRequestID(ctx)
|
||||
if accountType == "api_key" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
entry.Debugf("Use API key %s for model %s %s", util.HideAPIKey(accountInfo), req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
entry.Debugf("Use API key %s for model %s", util.HideAPIKey(accountInfo), req.Model)
|
||||
}
|
||||
} else if accountType == "oauth" {
|
||||
if proxyInfo != "" {
|
||||
log.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
entry.Debugf("Use OAuth %s for model %s %s", accountInfo, req.Model, proxyInfo)
|
||||
} else {
|
||||
log.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
entry.Debugf("Use OAuth %s for model %s", accountInfo, req.Model)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -624,13 +637,20 @@ func (m *Manager) normalizeProviders(providers []string) []string {
|
||||
return result
|
||||
}
|
||||
|
||||
// rotateProviders returns a rotated view of the providers list starting from the
|
||||
// current offset for the model, and atomically increments the offset for the next call.
|
||||
// This ensures concurrent requests get different starting providers.
|
||||
func (m *Manager) rotateProviders(model string, providers []string) []string {
|
||||
if len(providers) == 0 {
|
||||
return nil
|
||||
}
|
||||
m.mu.RLock()
|
||||
|
||||
// Atomic read-and-increment: get current offset and advance cursor in one lock
|
||||
m.mu.Lock()
|
||||
offset := m.providerOffsets[model]
|
||||
m.mu.RUnlock()
|
||||
m.providerOffsets[model] = (offset + 1) % len(providers)
|
||||
m.mu.Unlock()
|
||||
|
||||
if len(providers) > 0 {
|
||||
offset %= len(providers)
|
||||
}
|
||||
@@ -646,19 +666,6 @@ func (m *Manager) rotateProviders(model string, providers []string) []string {
|
||||
return rotated
|
||||
}
|
||||
|
||||
func (m *Manager) advanceProviderCursor(model string, providers []string) {
|
||||
if len(providers) == 0 {
|
||||
m.mu.Lock()
|
||||
delete(m.providerOffsets, model)
|
||||
m.mu.Unlock()
|
||||
return
|
||||
}
|
||||
m.mu.Lock()
|
||||
current := m.providerOffsets[model]
|
||||
m.providerOffsets[model] = (current + 1) % len(providers)
|
||||
m.mu.Unlock()
|
||||
}
|
||||
|
||||
func (m *Manager) retrySettings() (int, time.Duration) {
|
||||
if m == nil {
|
||||
return 0, 0
|
||||
@@ -1594,6 +1601,17 @@ type RequestPreparer interface {
|
||||
PrepareRequest(req *http.Request, auth *Auth) error
|
||||
}
|
||||
|
||||
// logEntryWithRequestID returns a logrus entry with request_id field if available in context.
|
||||
func logEntryWithRequestID(ctx context.Context) *log.Entry {
|
||||
if ctx == nil {
|
||||
return log.NewEntry(log.StandardLogger())
|
||||
}
|
||||
if reqID := logging.GetRequestID(ctx); reqID != "" {
|
||||
return log.WithField("request_id", reqID)
|
||||
}
|
||||
return log.NewEntry(log.StandardLogger())
|
||||
}
|
||||
|
||||
// InjectCredentials delegates per-provider HTTP request preparation when supported.
|
||||
// If the registered executor for the auth provider implements RequestPreparer,
|
||||
// it will be invoked to modify the request (e.g., add headers).
|
||||
@@ -20,6 +20,11 @@ type RoundRobinSelector struct {
|
||||
cursors map[string]int
|
||||
}
|
||||
|
||||
// FillFirstSelector selects the first available credential (deterministic ordering).
|
||||
// This "burns" one account before moving to the next, which can help stagger
|
||||
// rolling-window subscription caps (e.g. chat message limits).
|
||||
type FillFirstSelector struct{}
|
||||
|
||||
type blockReason int
|
||||
|
||||
const (
|
||||
@@ -98,20 +103,8 @@ func (e *modelCooldownError) Headers() http.Header {
|
||||
return headers
|
||||
}
|
||||
|
||||
// Pick selects the next available auth for the provider in a round-robin manner.
|
||||
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
if len(auths) == 0 {
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
||||
}
|
||||
if s.cursors == nil {
|
||||
s.cursors = make(map[string]int)
|
||||
}
|
||||
available := make([]*Auth, 0, len(auths))
|
||||
now := time.Now()
|
||||
cooldownCount := 0
|
||||
var earliest time.Time
|
||||
func collectAvailable(auths []*Auth, model string, now time.Time) (available []*Auth, cooldownCount int, earliest time.Time) {
|
||||
available = make([]*Auth, 0, len(auths))
|
||||
for i := 0; i < len(auths); i++ {
|
||||
candidate := auths[i]
|
||||
blocked, reason, next := isAuthBlockedForModel(candidate, model, now)
|
||||
@@ -126,6 +119,18 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
}
|
||||
}
|
||||
}
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
}
|
||||
return available, cooldownCount, earliest
|
||||
}
|
||||
|
||||
func getAvailableAuths(auths []*Auth, provider, model string, now time.Time) ([]*Auth, error) {
|
||||
if len(auths) == 0 {
|
||||
return nil, &Error{Code: "auth_not_found", Message: "no auth candidates"}
|
||||
}
|
||||
|
||||
available, cooldownCount, earliest := collectAvailable(auths, model, now)
|
||||
if len(available) == 0 {
|
||||
if cooldownCount == len(auths) && !earliest.IsZero() {
|
||||
resetIn := earliest.Sub(now)
|
||||
@@ -136,12 +141,24 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
}
|
||||
return nil, &Error{Code: "auth_unavailable", Message: "no auth available"}
|
||||
}
|
||||
// Make round-robin deterministic even if caller's candidate order is unstable.
|
||||
if len(available) > 1 {
|
||||
sort.Slice(available, func(i, j int) bool { return available[i].ID < available[j].ID })
|
||||
|
||||
return available, nil
|
||||
}
|
||||
|
||||
// Pick selects the next available auth for the provider in a round-robin manner.
|
||||
func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
now := time.Now()
|
||||
available, err := getAvailableAuths(auths, provider, model, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
key := provider + ":" + model
|
||||
s.mu.Lock()
|
||||
if s.cursors == nil {
|
||||
s.cursors = make(map[string]int)
|
||||
}
|
||||
index := s.cursors[key]
|
||||
|
||||
if index >= 2_147_483_640 {
|
||||
@@ -154,6 +171,18 @@ func (s *RoundRobinSelector) Pick(ctx context.Context, provider, model string, o
|
||||
return available[index%len(available)], nil
|
||||
}
|
||||
|
||||
// Pick selects the first available auth for the provider in a deterministic manner.
|
||||
func (s *FillFirstSelector) Pick(ctx context.Context, provider, model string, opts cliproxyexecutor.Options, auths []*Auth) (*Auth, error) {
|
||||
_ = ctx
|
||||
_ = opts
|
||||
now := time.Now()
|
||||
available, err := getAvailableAuths(auths, provider, model, now)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return available[0], nil
|
||||
}
|
||||
|
||||
func isAuthBlockedForModel(auth *Auth, model string, now time.Time) (bool, blockReason, time.Time) {
|
||||
if auth == nil {
|
||||
return true, blockReasonOther, time.Time{}
|
||||
|
||||
113
sdk/cliproxy/auth/selector_test.go
Normal file
113
sdk/cliproxy/auth/selector_test.go
Normal file
@@ -0,0 +1,113 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
|
||||
)
|
||||
|
||||
func TestFillFirstSelectorPick_Deterministic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
selector := &FillFirstSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
t.Fatalf("Pick() error = %v", err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("Pick() auth = nil")
|
||||
}
|
||||
if got.ID != "a" {
|
||||
t.Fatalf("Pick() auth.ID = %q, want %q", got.ID, "a")
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinSelectorPick_CyclesDeterministic(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
selector := &RoundRobinSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
want := []string{"a", "b", "c", "a", "b"}
|
||||
for i, id := range want {
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
t.Fatalf("Pick() #%d error = %v", i, err)
|
||||
}
|
||||
if got == nil {
|
||||
t.Fatalf("Pick() #%d auth = nil", i)
|
||||
}
|
||||
if got.ID != id {
|
||||
t.Fatalf("Pick() #%d auth.ID = %q, want %q", i, got.ID, id)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestRoundRobinSelectorPick_Concurrent(t *testing.T) {
|
||||
selector := &RoundRobinSelector{}
|
||||
auths := []*Auth{
|
||||
{ID: "b"},
|
||||
{ID: "a"},
|
||||
{ID: "c"},
|
||||
}
|
||||
|
||||
start := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, 1)
|
||||
|
||||
goroutines := 32
|
||||
iterations := 100
|
||||
for i := 0; i < goroutines; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
for j := 0; j < iterations; j++ {
|
||||
got, err := selector.Pick(context.Background(), "gemini", "", cliproxyexecutor.Options{}, auths)
|
||||
if err != nil {
|
||||
select {
|
||||
case errCh <- err:
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
if got == nil {
|
||||
select {
|
||||
case errCh <- errors.New("Pick() returned nil auth"):
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
if got.ID == "" {
|
||||
select {
|
||||
case errCh <- errors.New("Pick() returned auth with empty ID"):
|
||||
default:
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
close(start)
|
||||
wg.Wait()
|
||||
|
||||
select {
|
||||
case err := <-errCh:
|
||||
t.Fatalf("concurrent Pick() error = %v", err)
|
||||
default:
|
||||
}
|
||||
}
|
||||
@@ -1,11 +1,12 @@
|
||||
package auth
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
baseauth "github.com/router-for-me/CLIProxyAPI/v6/internal/auth"
|
||||
@@ -15,8 +16,8 @@ import (
|
||||
type Auth struct {
|
||||
// ID uniquely identifies the auth record across restarts.
|
||||
ID string `json:"id"`
|
||||
// Index is a monotonically increasing runtime identifier used for diagnostics.
|
||||
Index uint64 `json:"-"`
|
||||
// Index is a stable runtime identifier derived from auth metadata (not persisted).
|
||||
Index string `json:"-"`
|
||||
// Provider is the upstream provider key (e.g. "gemini", "claude").
|
||||
Provider string `json:"provider"`
|
||||
// Prefix optionally namespaces models for routing (e.g., "teamA/gemini-3-pro-preview").
|
||||
@@ -94,12 +95,6 @@ type ModelState struct {
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
var authIndexCounter atomic.Uint64
|
||||
|
||||
func nextAuthIndex() uint64 {
|
||||
return authIndexCounter.Add(1) - 1
|
||||
}
|
||||
|
||||
// Clone shallow copies the Auth structure, duplicating maps to avoid accidental mutation.
|
||||
func (a *Auth) Clone() *Auth {
|
||||
if a == nil {
|
||||
@@ -128,15 +123,41 @@ func (a *Auth) Clone() *Auth {
|
||||
return ©Auth
|
||||
}
|
||||
|
||||
// EnsureIndex returns the global index, assigning one if it was not set yet.
|
||||
func (a *Auth) EnsureIndex() uint64 {
|
||||
if a == nil {
|
||||
return 0
|
||||
func stableAuthIndex(seed string) string {
|
||||
seed = strings.TrimSpace(seed)
|
||||
if seed == "" {
|
||||
return ""
|
||||
}
|
||||
if a.indexAssigned {
|
||||
sum := sha256.Sum256([]byte(seed))
|
||||
return hex.EncodeToString(sum[:8])
|
||||
}
|
||||
|
||||
// EnsureIndex returns a stable index derived from the auth file name or API key.
|
||||
func (a *Auth) EnsureIndex() string {
|
||||
if a == nil {
|
||||
return ""
|
||||
}
|
||||
if a.indexAssigned && a.Index != "" {
|
||||
return a.Index
|
||||
}
|
||||
idx := nextAuthIndex()
|
||||
|
||||
seed := strings.TrimSpace(a.FileName)
|
||||
if seed != "" {
|
||||
seed = "file:" + seed
|
||||
} else if a.Attributes != nil {
|
||||
if apiKey := strings.TrimSpace(a.Attributes["api_key"]); apiKey != "" {
|
||||
seed = "api_key:" + apiKey
|
||||
}
|
||||
}
|
||||
if seed == "" {
|
||||
if id := strings.TrimSpace(a.ID); id != "" {
|
||||
seed = "id:" + id
|
||||
} else {
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
idx := stableAuthIndex(seed)
|
||||
a.Index = idx
|
||||
a.indexAssigned = true
|
||||
return idx
|
||||
|
||||
@@ -5,6 +5,7 @@ package cliproxy
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/router-for-me/CLIProxyAPI/v6/internal/api"
|
||||
sdkaccess "github.com/router-for-me/CLIProxyAPI/v6/sdk/access"
|
||||
@@ -197,7 +198,20 @@ func (b *Builder) Build() (*Service, error) {
|
||||
if dirSetter, ok := tokenStore.(interface{ SetBaseDir(string) }); ok && b.cfg != nil {
|
||||
dirSetter.SetBaseDir(b.cfg.AuthDir)
|
||||
}
|
||||
coreManager = coreauth.NewManager(tokenStore, nil, nil)
|
||||
|
||||
strategy := ""
|
||||
if b.cfg != nil {
|
||||
strategy = strings.ToLower(strings.TrimSpace(b.cfg.Routing.Strategy))
|
||||
}
|
||||
var selector coreauth.Selector
|
||||
switch strategy {
|
||||
case "fill-first", "fillfirst", "ff":
|
||||
selector = &coreauth.FillFirstSelector{}
|
||||
default:
|
||||
selector = &coreauth.RoundRobinSelector{}
|
||||
}
|
||||
|
||||
coreManager = coreauth.NewManager(tokenStore, selector, nil)
|
||||
}
|
||||
// Attach a default RoundTripper provider so providers can opt-in per-auth transports.
|
||||
coreManager.SetRoundTripperProvider(newDefaultRoundTripperProvider())
|
||||
|
||||
@@ -510,6 +510,13 @@ func (s *Service) Run(ctx context.Context) error {
|
||||
|
||||
var watcherWrapper *WatcherWrapper
|
||||
reloadCallback := func(newCfg *config.Config) {
|
||||
previousStrategy := ""
|
||||
s.cfgMu.RLock()
|
||||
if s.cfg != nil {
|
||||
previousStrategy = strings.ToLower(strings.TrimSpace(s.cfg.Routing.Strategy))
|
||||
}
|
||||
s.cfgMu.RUnlock()
|
||||
|
||||
if newCfg == nil {
|
||||
s.cfgMu.RLock()
|
||||
newCfg = s.cfg
|
||||
@@ -518,6 +525,30 @@ func (s *Service) Run(ctx context.Context) error {
|
||||
if newCfg == nil {
|
||||
return
|
||||
}
|
||||
|
||||
nextStrategy := strings.ToLower(strings.TrimSpace(newCfg.Routing.Strategy))
|
||||
normalizeStrategy := func(strategy string) string {
|
||||
switch strategy {
|
||||
case "fill-first", "fillfirst", "ff":
|
||||
return "fill-first"
|
||||
default:
|
||||
return "round-robin"
|
||||
}
|
||||
}
|
||||
previousStrategy = normalizeStrategy(previousStrategy)
|
||||
nextStrategy = normalizeStrategy(nextStrategy)
|
||||
if s.coreManager != nil && previousStrategy != nextStrategy {
|
||||
var selector coreauth.Selector
|
||||
switch nextStrategy {
|
||||
case "fill-first":
|
||||
selector = &coreauth.FillFirstSelector{}
|
||||
default:
|
||||
selector = &coreauth.RoundRobinSelector{}
|
||||
}
|
||||
s.coreManager.SetSelector(selector)
|
||||
log.Infof("routing strategy updated to %s", nextStrategy)
|
||||
}
|
||||
|
||||
s.applyRetryConfig(newCfg)
|
||||
if s.server != nil {
|
||||
s.server.UpdateClients(newCfg)
|
||||
@@ -714,6 +745,9 @@ func (s *Service) registerModelsForAuth(a *coreauth.Auth) {
|
||||
case "codex":
|
||||
models = registry.GetOpenAIModels()
|
||||
if entry := s.resolveConfigCodexKey(a); entry != nil {
|
||||
if len(entry.Models) > 0 {
|
||||
models = buildCodexConfigModels(entry)
|
||||
}
|
||||
if authKind == "apikey" {
|
||||
excluded = entry.ExcludedModels
|
||||
}
|
||||
@@ -1157,3 +1191,41 @@ func buildClaudeConfigModels(entry *config.ClaudeKey) []*ModelInfo {
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func buildCodexConfigModels(entry *config.CodexKey) []*ModelInfo {
|
||||
if entry == nil || len(entry.Models) == 0 {
|
||||
return nil
|
||||
}
|
||||
now := time.Now().Unix()
|
||||
out := make([]*ModelInfo, 0, len(entry.Models))
|
||||
seen := make(map[string]struct{}, len(entry.Models))
|
||||
for i := range entry.Models {
|
||||
model := entry.Models[i]
|
||||
name := strings.TrimSpace(model.Name)
|
||||
alias := strings.TrimSpace(model.Alias)
|
||||
if alias == "" {
|
||||
alias = name
|
||||
}
|
||||
if alias == "" {
|
||||
continue
|
||||
}
|
||||
key := strings.ToLower(alias)
|
||||
if _, exists := seen[key]; exists {
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
display := name
|
||||
if display == "" {
|
||||
display = alias
|
||||
}
|
||||
out = append(out, &ModelInfo{
|
||||
ID: alias,
|
||||
Object: "model",
|
||||
Created: now,
|
||||
OwnedBy: "openai",
|
||||
Type: "openai",
|
||||
DisplayName: display,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
@@ -14,7 +14,7 @@ type Record struct {
|
||||
Model string
|
||||
APIKey string
|
||||
AuthID string
|
||||
AuthIndex uint64
|
||||
AuthIndex string
|
||||
Source string
|
||||
RequestedAt time.Time
|
||||
Failed bool
|
||||
|
||||
@@ -12,6 +12,7 @@ type AccessProvider = internalconfig.AccessProvider
|
||||
|
||||
type Config = internalconfig.Config
|
||||
|
||||
type StreamingConfig = internalconfig.StreamingConfig
|
||||
type TLSConfig = internalconfig.TLSConfig
|
||||
type RemoteManagement = internalconfig.RemoteManagement
|
||||
type AmpCode = internalconfig.AmpCode
|
||||
|
||||
Reference in New Issue
Block a user