chuan/CLIProxyAPI
1
0
Fork 0
mirror of https://github.com/router-for-me/CLIProxyAPI.git synced 2026-02-03 13:00:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor codebase

Browse Source
This commit is contained in:
Luis Pater
2025-08-22 01:31:12 +08:00
parent 2b1762be16
commit 8c555c4e69
109 changed files with 7319 additions and 5735 deletions
Download Patch File Download Diff File Expand all files Collapse all files

302
internal/client/claude_client.go
View File

@@ -1,3 +1,6 @@
// Package client provides HTTP client functionality for interacting with Anthropic's Claude API.
// It handles authentication, request/response translation, streaming communication,
// and quota management for Claude models.
package client
import (
@@ -17,7 +20,10 @@ import (
"github.com/luispater/CLIProxyAPI/internal/auth/claude"
"github.com/luispater/CLIProxyAPI/internal/auth/empty"
"github.com/luispater/CLIProxyAPI/internal/config"
. "github.com/luispater/CLIProxyAPI/internal/constant"
"github.com/luispater/CLIProxyAPI/internal/interfaces"
"github.com/luispater/CLIProxyAPI/internal/misc"
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
"github.com/luispater/CLIProxyAPI/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
@@ -28,14 +34,25 @@ const (
claudeEndpoint = "https://api.anthropic.com"
)
// ClaudeClient implements the Client interface for OpenAI API
// ClaudeClient implements the Client interface for Anthropic's Claude API.
// It provides methods for authenticating with Claude and sending requests to Claude models.
type ClaudeClient struct {
ClientBase
claudeAuth *claude.ClaudeAuth
// claudeAuth handles authentication with Claude API
claudeAuth *claude.ClaudeAuth
// apiKeyIndex is the index of the API key to use from the config, -1 if not using API keys
apiKeyIndex int
}
// NewClaudeClient creates a new OpenAI client instance
// NewClaudeClient creates a new Claude client instance using token-based authentication.
// It initializes the client with the provided configuration and token storage.
//
// Parameters:
// - cfg: The application configuration.
// - ts: The token storage for Claude authentication.
//
// Returns:
// - *ClaudeClient: A new Claude client instance.
func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeClient {
httpClient := util.SetProxy(cfg, &http.Client{})
client := &ClaudeClient{
@@ -53,7 +70,16 @@ func NewClaudeClient(cfg *config.Config, ts *claude.ClaudeTokenStorage) *ClaudeC
return client
}
// NewClaudeClientWithKey creates a new OpenAI client instance with api key
// NewClaudeClientWithKey creates a new Claude client instance using API key authentication.
// It initializes the client with the provided configuration and selects the API key
// at the specified index from the configuration.
//
// Parameters:
// - cfg: The application configuration.
// - apiKeyIndex: The index of the API key to use from the configuration.
//
// Returns:
// - *ClaudeClient: A new Claude client instance.
func NewClaudeClientWithKey(cfg *config.Config, apiKeyIndex int) *ClaudeClient {
httpClient := util.SetProxy(cfg, &http.Client{})
client := &ClaudeClient{
@@ -71,7 +97,41 @@ func NewClaudeClientWithKey(cfg *config.Config, apiKeyIndex int) *ClaudeClient {
return client
}
// GetAPIKey returns the api key index
// Type returns the client type identifier.
// This method returns "claude" to identify this client as a Claude API client.
func (c *ClaudeClient) Type() string {
return CLAUDE
}
// Provider returns the provider name for this client.
// This method returns "claude" to identify Anthropic's Claude as the provider.
func (c *ClaudeClient) Provider() string {
return CLAUDE
}
// CanProvideModel checks if this client can provide the specified model.
// It returns true if the model is supported by Claude, false otherwise.
//
// Parameters:
// - modelName: The name of the model to check.
//
// Returns:
// - bool: True if the model is supported, false otherwise.
func (c *ClaudeClient) CanProvideModel(modelName string) bool {
// List of Claude models supported by this client
models := []string{
"claude-opus-4-1-20250805",
"claude-opus-4-20250514",
"claude-sonnet-4-20250514",
"claude-3-7-sonnet-20250219",
"claude-3-5-haiku-20241022",
}
return util.InArray(models, modelName)
}
// GetAPIKey returns the API key for Claude API requests.
// If an API key index is specified, it returns the corresponding key from the configuration.
// Otherwise, it returns an empty string, indicating token-based authentication should be used.
func (c *ClaudeClient) GetAPIKey() string {
if c.apiKeyIndex != -1 {
return c.cfg.ClaudeKey[c.apiKeyIndex].APIKey
@@ -79,43 +139,37 @@ func (c *ClaudeClient) GetAPIKey() string {
return ""
}
// GetUserAgent returns the user agent string for OpenAI API requests
// GetUserAgent returns the user agent string for Claude API requests.
// This identifies the client as the Claude CLI to the Anthropic API.
func (c *ClaudeClient) GetUserAgent() string {
return "claude-cli/1.0.83 (external, cli)"
}
// TokenStorage returns the token storage interface used by this client.
// This provides access to the authentication token management system.
func (c *ClaudeClient) TokenStorage() auth.TokenStorage {
return c.tokenStorage
}
// SendMessage sends a message to OpenAI API (non-streaming)
func (c *ClaudeClient) SendMessage(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration) ([]byte, *ErrorMessage) {
// For now, return an error as OpenAI integration is not fully implemented
return nil, &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("claude message sending not yet implemented"),
}
}
// SendRawMessage sends a raw message to Claude API and returns the response.
// It handles request translation, API communication, error handling, and response translation.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: The response body.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *ClaudeClient) SendRawMessage(ctx context.Context, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, false)
rawJSON, _ = sjson.SetBytes(rawJSON, "stream", true)
// SendMessageStream sends a streaming message to OpenAI API
func (c *ClaudeClient) SendMessageStream(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration, _ ...bool) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage, 1)
errChan <- &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("claude streaming not yet implemented"),
}
close(errChan)
return nil, errChan
}
// SendRawMessage sends a raw message to OpenAI API
func (c *ClaudeClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt string) ([]byte, *ErrorMessage) {
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
respBody, err := c.APIRequest(ctx, "/v1/messages?beta=true", rawJSON, alt, false)
respBody, err := c.APIRequest(ctx, modelName, "/v1/messages?beta=true", rawJSON, alt, false)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
@@ -126,50 +180,88 @@ func (c *ClaudeClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt s
delete(c.modelQuotaExceeded, modelName)
bodyBytes, errReadAll := io.ReadAll(respBody)
if errReadAll != nil {
return nil, &ErrorMessage{StatusCode: 500, Error: errReadAll}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
}
return bodyBytes, nil
c.AddAPIResponseData(ctx, bodyBytes)
var param any
bodyBytes = []byte(translator.ResponseNonStream(handlerType, c.Type(), ctx, modelName, bodyBytes, &param))
return bodyBytes, nil
}
// SendRawMessageStream sends a raw streaming message to OpenAI API
func (c *ClaudeClient) SendRawMessageStream(ctx context.Context, rawJSON []byte, alt string) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage)
// SendRawMessageStream sends a raw streaming message to Claude API.
// It returns two channels: one for receiving response data chunks and one for errors.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - <-chan []byte: A channel for receiving response data chunks.
// - <-chan *interfaces.ErrorMessage: A channel for receiving error messages.
func (c *ClaudeClient) SendRawMessageStream(ctx context.Context, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
// return dataChan, errChan
go func() {
defer close(errChan)
defer close(dataChan)
rawJSON, _ = sjson.SetBytes(rawJSON, "stream", true)
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
var stream io.ReadCloser
for {
var err *ErrorMessage
stream, err = c.APIRequest(ctx, "/v1/messages?beta=true", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
if c.IsModelQuotaExceeded(modelName) {
errChan <- &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
delete(c.modelQuotaExceeded, modelName)
break
return
}
var err *interfaces.ErrorMessage
stream, err = c.APIRequest(ctx, modelName, "/v1/messages?beta=true", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
}
delete(c.modelQuotaExceeded, modelName)
scanner := bufio.NewScanner(stream)
buffer := make([]byte, 10240*1024)
scanner.Buffer(buffer, 10240*1024)
for scanner.Scan() {
line := scanner.Bytes()
dataChan <- line
if translator.NeedConvert(handlerType, c.Type()) {
var param any
for scanner.Scan() {
line := scanner.Bytes()
lines := translator.Response(handlerType, c.Type(), ctx, modelName, line, &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
c.AddAPIResponseData(ctx, line)
}
} else {
for scanner.Scan() {
line := scanner.Bytes()
dataChan <- line
c.AddAPIResponseData(ctx, line)
}
}
if errScanner := scanner.Err(); errScanner != nil {
errChan <- &ErrorMessage{500, errScanner, nil}
errChan <- &interfaces.ErrorMessage{StatusCode: 500, Error: errScanner}
_ = stream.Close()
return
}
@@ -180,36 +272,62 @@ func (c *ClaudeClient) SendRawMessageStream(ctx context.Context, rawJSON []byte,
return dataChan, errChan
}
// SendRawTokenCount sends a token count request to OpenAI API
func (c *ClaudeClient) SendRawTokenCount(_ context.Context, _ []byte, _ string) ([]byte, *ErrorMessage) {
return nil, &ErrorMessage{
// SendRawTokenCount sends a token count request to Claude API.
// Currently, this functionality is not implemented for Claude models.
// It returns a NotImplemented error.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: Always nil for this implementation.
// - *interfaces.ErrorMessage: An error message indicating that the feature is not implemented.
func (c *ClaudeClient) SendRawTokenCount(_ context.Context, _ string, _ []byte, _ string) ([]byte, *interfaces.ErrorMessage) {
return nil, &interfaces.ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("claude token counting not yet implemented"),
}
}
// SaveTokenToFile persists the token storage to disk
// SaveTokenToFile persists the authentication tokens to disk.
// It saves the token data to a JSON file in the configured authentication directory,
// with a filename based on the user's email address.
//
// Returns:
// - error: An error if the save operation fails, nil otherwise.
func (c *ClaudeClient) SaveTokenToFile() error {
fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("claude-%s.json", c.tokenStorage.(*claude.ClaudeTokenStorage).Email))
return c.tokenStorage.SaveTokenToFile(fileName)
}
// RefreshTokens refreshes the access tokens if needed
// RefreshTokens refreshes the access tokens if they have expired.
// It uses the refresh token to obtain new access tokens from the Claude authentication service.
// If successful, it updates the token storage and persists the new tokens to disk.
//
// Parameters:
// - ctx: The context for the request.
//
// Returns:
// - error: An error if the refresh operation fails, nil otherwise.
func (c *ClaudeClient) RefreshTokens(ctx context.Context) error {
// Check if we have a valid refresh token
if c.tokenStorage == nil || c.tokenStorage.(*claude.ClaudeTokenStorage).RefreshToken == "" {
return fmt.Errorf("no refresh token available")
}
// Refresh tokens using the auth service
// Refresh tokens using the auth service with retry mechanism
newTokenData, err := c.claudeAuth.RefreshTokensWithRetry(ctx, c.tokenStorage.(*claude.ClaudeTokenStorage).RefreshToken, 3)
if err != nil {
return fmt.Errorf("failed to refresh tokens: %w", err)
}
// Update token storage
// Update token storage with new token data
c.claudeAuth.UpdateTokenStorage(c.tokenStorage.(*claude.ClaudeTokenStorage), newTokenData)
// Save updated tokens
// Save updated tokens to persistent storage
if err = c.SaveTokenToFile(); err != nil {
log.Warnf("Failed to save refreshed tokens: %v", err)
}
@@ -218,16 +336,30 @@ func (c *ClaudeClient) RefreshTokens(ctx context.Context) error {
return nil
}
// APIRequest handles making requests to the CLI API endpoints.
func (c *ClaudeClient) APIRequest(ctx context.Context, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *ErrorMessage) {
// APIRequest handles making HTTP requests to the Claude API endpoints.
// It manages authentication, request preparation, and response handling.
//
// Parameters:
// - ctx: The context for the request, which may contain additional request metadata.
// - modelName: The name of the model being requested.
// - endpoint: The API endpoint path to call (e.g., "/v1/messages").
// - body: The request body, either as a byte array or an object to be marshaled to JSON.
// - alt: An alternative response format parameter (unused in this implementation).
// - stream: A boolean indicating if the request is for a streaming response (unused in this implementation).
//
// Returns:
// - io.ReadCloser: The response body reader if successful.
// - *interfaces.ErrorMessage: Error information if the request fails.
func (c *ClaudeClient) APIRequest(ctx context.Context, modelName, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *interfaces.ErrorMessage) {
var jsonBody []byte
var err error
// Convert body to JSON bytes
if byteBody, ok := body.([]byte); ok {
jsonBody = byteBody
} else {
jsonBody, err = json.Marshal(body)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to marshal request body: %w", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to marshal request body: %w", err)}
}
}
@@ -268,7 +400,7 @@ func (c *ClaudeClient) APIRequest(ctx context.Context, endpoint string, body int
req, err := http.NewRequestWithContext(ctx, "POST", url, reqBody)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to create request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to create request: %v", err)}
}
// Set headers
@@ -294,13 +426,21 @@ func (c *ClaudeClient) APIRequest(ctx context.Context, endpoint string, body int
req.Header.Set("Accept-Encoding", "gzip, deflate, br, zstd")
req.Header.Set("Anthropic-Beta", "claude-code-20250219,oauth-2025-04-20,interleaved-thinking-2025-05-14,fine-grained-tool-streaming-2025-05-14")
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
if c.cfg.RequestLog {
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
}
}
if c.apiKeyIndex != -1 {
log.Debugf("Use Claude API key %s for model %s", util.HideAPIKey(c.cfg.ClaudeKey[c.apiKeyIndex].APIKey), modelName)
} else {
log.Debugf("Use Claude account %s for model %s", c.GetEmail(), modelName)
}
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to execute request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to execute request: %v", err)}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
@@ -314,12 +454,20 @@ func (c *ClaudeClient) APIRequest(ctx context.Context, endpoint string, body int
addon := c.createAddon(resp.Header)
// log.Debug(string(jsonBody))
return nil, &ErrorMessage{resp.StatusCode, fmt.Errorf(string(bodyBytes)), addon}
return nil, &interfaces.ErrorMessage{StatusCode: resp.StatusCode, Error: fmt.Errorf("%s", string(bodyBytes)), Addon: addon}
}
return resp.Body, nil
}
// createAddon creates a new http.Header containing selected headers from the original response.
// This is used to pass relevant rate limit and retry information back to the caller.
//
// Parameters:
// - header: The original http.Header from the API response.
//
// Returns:
// - http.Header: A new header containing the selected headers.
func (c *ClaudeClient) createAddon(header http.Header) http.Header {
addon := http.Header{}
if _, ok := header["X-Should-Retry"]; ok {
@@ -352,6 +500,8 @@ func (c *ClaudeClient) createAddon(header http.Header) http.Header {
return addon
}
// GetEmail returns the email address associated with the client's token storage.
// If the client is using API key authentication, it returns an empty string.
func (c *ClaudeClient) GetEmail() string {
if ts, ok := c.tokenStorage.(*claude.ClaudeTokenStorage); ok {
return ts.Email
@@ -362,6 +512,12 @@ func (c *ClaudeClient) GetEmail() string {
// IsModelQuotaExceeded returns true if the specified model has exceeded its quota
// and no fallback options are available.
//
// Parameters:
// - model: The name of the model to check.
//
// Returns:
// - bool: True if the model's quota is exceeded, false otherwise.
func (c *ClaudeClient) IsModelQuotaExceeded(model string) bool {
if lastExceededTime, hasKey := c.modelQuotaExceeded[model]; hasKey {
duration := time.Now().Sub(*lastExceededTime)

78
internal/client/client.go
View File

@@ -4,61 +4,17 @@
package client
import (
"bytes"
"context"
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
"github.com/luispater/CLIProxyAPI/internal/auth"
"github.com/luispater/CLIProxyAPI/internal/config"
)
// Client defines the interface that all AI API clients must implement.
// This interface provides methods for interacting with various AI services
// including sending messages, streaming responses, and managing authentication.
type Client interface {
// GetRequestMutex returns the mutex used to synchronize requests for this client.
// This ensures that only one request is processed at a time for quota management.
GetRequestMutex() *sync.Mutex
// GetUserAgent returns the User-Agent string used for HTTP requests.
GetUserAgent() string
// SendMessage sends a single message to the AI service and returns the response.
// It takes the raw JSON request, model name, system instructions, conversation contents,
// and tool declarations, then returns the response bytes and any error that occurred.
SendMessage(ctx context.Context, rawJSON []byte, model string, systemInstruction *Content, contents []Content, tools []ToolDeclaration) ([]byte, *ErrorMessage)
// SendMessageStream sends a message to the AI service and returns streaming responses.
// It takes similar parameters to SendMessage but returns channels for streaming data
// and errors, enabling real-time response processing.
SendMessageStream(ctx context.Context, rawJSON []byte, model string, systemInstruction *Content, contents []Content, tools []ToolDeclaration, includeThoughts ...bool) (<-chan []byte, <-chan *ErrorMessage)
// SendRawMessage sends a raw JSON message to the AI service without translation.
// This method is used when the request is already in the service's native format.
SendRawMessage(ctx context.Context, rawJSON []byte, alt string) ([]byte, *ErrorMessage)
// SendRawMessageStream sends a raw JSON message and returns streaming responses.
// Similar to SendRawMessage but for streaming responses.
SendRawMessageStream(ctx context.Context, rawJSON []byte, alt string) (<-chan []byte, <-chan *ErrorMessage)
// SendRawTokenCount sends a token count request to the AI service.
// This method is used to estimate the number of tokens in a given text.
SendRawTokenCount(ctx context.Context, rawJSON []byte, alt string) ([]byte, *ErrorMessage)
// SaveTokenToFile saves the client's authentication token to a file.
// This is used for persisting authentication state between sessions.
SaveTokenToFile() error
// IsModelQuotaExceeded checks if the specified model has exceeded its quota.
// This helps with load balancing and automatic failover to alternative models.
IsModelQuotaExceeded(model string) bool
// GetEmail returns the email associated with the client's authentication.
// This is used for logging and identification purposes.
GetEmail() string
}
// ClientBase provides a common base structure for all AI API clients.
// It implements shared functionality such as request synchronization, HTTP client management,
// configuration access, token storage, and quota tracking.
@@ -82,6 +38,36 @@ type ClientBase struct {
// GetRequestMutex returns the mutex used to synchronize requests for this client.
// This ensures that only one request is processed at a time for quota management.
//
// Returns:
// - *sync.Mutex: The mutex used for request synchronization
func (c *ClientBase) GetRequestMutex() *sync.Mutex {
return c.RequestMutex
}
// AddAPIResponseData adds API response data to the Gin context for logging purposes.
// This method appends the provided data to any existing response data in the context,
// or creates a new entry if none exists. It only performs this operation if request
// logging is enabled in the configuration.
//
// Parameters:
// - ctx: The context for the request
// - line: The response data to be added
func (c *ClientBase) AddAPIResponseData(ctx context.Context, line []byte) {
if c.cfg.RequestLog {
data := bytes.TrimSpace(bytes.Clone(line))
if ginContext, ok := ctx.Value("gin").(*gin.Context); len(data) > 0 && ok {
if apiResponseData, isExist := ginContext.Get("API_RESPONSE"); isExist {
if byteAPIResponseData, isOk := apiResponseData.([]byte); isOk {
// Append new data and separator to existing response data
byteAPIResponseData = append(byteAPIResponseData, data...)
byteAPIResponseData = append(byteAPIResponseData, []byte("\n\n")...)
ginContext.Set("API_RESPONSE", byteAPIResponseData)
}
} else {
// Create new response data entry
ginContext.Set("API_RESPONSE", data)
}
}
}
}

165
internal/client/client_models.go
View File

@@ -1,165 +0,0 @@
// Package client defines the data structures used across all AI API clients.
// These structures represent the common data models for requests, responses,
// and configuration parameters used when communicating with various AI services.
package client
import (
"net/http"
"time"
)
// ErrorMessage encapsulates an error with an associated HTTP status code.
// This structure is used to provide detailed error information including
// both the HTTP status and the underlying error.
type ErrorMessage struct {
// StatusCode is the HTTP status code returned by the API.
StatusCode int
// Error is the underlying error that occurred.
Error error
// Addon is the additional headers to be added to the response
Addon http.Header
}
// GCPProject represents the response structure for a Google Cloud project list request.
// This structure is used when fetching available projects for a Google Cloud account.
type GCPProject struct {
// Projects is a list of Google Cloud projects accessible by the user.
Projects []GCPProjectProjects `json:"projects"`
}
// GCPProjectLabels defines the labels associated with a GCP project.
// These labels can contain metadata about the project's purpose or configuration.
type GCPProjectLabels struct {
// GenerativeLanguage indicates if the project has generative language APIs enabled.
GenerativeLanguage string `json:"generative-language"`
}
// GCPProjectProjects contains details about a single Google Cloud project.
// This includes identifying information, metadata, and configuration details.
type GCPProjectProjects struct {
// ProjectNumber is the unique numeric identifier for the project.
ProjectNumber string `json:"projectNumber"`
// ProjectID is the unique string identifier for the project.
ProjectID string `json:"projectId"`
// LifecycleState indicates the current state of the project (e.g., "ACTIVE").
LifecycleState string `json:"lifecycleState"`
// Name is the human-readable name of the project.
Name string `json:"name"`
// Labels contains metadata labels associated with the project.
Labels GCPProjectLabels `json:"labels"`
// CreateTime is the timestamp when the project was created.
CreateTime time.Time `json:"createTime"`
}
// Content represents a single message in a conversation, with a role and parts.
// This structure models a message exchange between a user and an AI model.
type Content struct {
// Role indicates who sent the message ("user", "model", or "tool").
Role string `json:"role"`
// Parts is a collection of content parts that make up the message.
Parts []Part `json:"parts"`
}
// Part represents a distinct piece of content within a message.
// A part can be text, inline data (like an image), a function call, or a function response.
type Part struct {
// Text contains plain text content.
Text string `json:"text,omitempty"`
// InlineData contains base64-encoded data with its MIME type (e.g., images).
InlineData *InlineData `json:"inlineData,omitempty"`
// FunctionCall represents a tool call requested by the model.
FunctionCall *FunctionCall `json:"functionCall,omitempty"`
// FunctionResponse represents the result of a tool execution.
FunctionResponse *FunctionResponse `json:"functionResponse,omitempty"`
}
// InlineData represents base64-encoded data with its MIME type.
// This is typically used for embedding images or other binary data in requests.
type InlineData struct {
// MimeType specifies the media type of the embedded data (e.g., "image/png").
MimeType string `json:"mime_type,omitempty"`
// Data contains the base64-encoded binary data.
Data string `json:"data,omitempty"`
}
// FunctionCall represents a tool call requested by the model.
// It includes the function name and its arguments that the model wants to execute.
type FunctionCall struct {
// Name is the identifier of the function to be called.
Name string `json:"name"`
// Args contains the arguments to pass to the function.
Args map[string]interface{} `json:"args"`
}
// FunctionResponse represents the result of a tool execution.
// This is sent back to the model after a tool call has been processed.
type FunctionResponse struct {
// Name is the identifier of the function that was called.
Name string `json:"name"`
// Response contains the result data from the function execution.
Response map[string]interface{} `json:"response"`
}
// GenerateContentRequest is the top-level request structure for the streamGenerateContent endpoint.
// This structure defines all the parameters needed for generating content from an AI model.
type GenerateContentRequest struct {
// SystemInstruction provides system-level instructions that guide the model's behavior.
SystemInstruction *Content `json:"systemInstruction,omitempty"`
// Contents is the conversation history between the user and the model.
Contents []Content `json:"contents"`
// Tools defines the available tools/functions that the model can call.
Tools []ToolDeclaration `json:"tools,omitempty"`
// GenerationConfig contains parameters that control the model's generation behavior.
GenerationConfig `json:"generationConfig"`
}
// GenerationConfig defines parameters that control the model's generation behavior.
// These parameters affect the creativity, randomness, and reasoning of the model's responses.
type GenerationConfig struct {
// ThinkingConfig specifies configuration for the model's "thinking" process.
ThinkingConfig GenerationConfigThinkingConfig `json:"thinkingConfig,omitempty"`
// Temperature controls the randomness of the model's responses.
// Values closer to 0 make responses more deterministic, while values closer to 1 increase randomness.
Temperature float64 `json:"temperature,omitempty"`
// TopP controls nucleus sampling, which affects the diversity of responses.
// It limits the model to consider only the top P% of probability mass.
TopP float64 `json:"topP,omitempty"`
// TopK limits the model to consider only the top K most likely tokens.
// This can help control the quality and diversity of generated text.
TopK float64 `json:"topK,omitempty"`
}
// GenerationConfigThinkingConfig specifies configuration for the model's "thinking" process.
// This controls whether the model should output its reasoning process along with the final answer.
type GenerationConfigThinkingConfig struct {
// IncludeThoughts determines whether the model should output its reasoning process.
// When enabled, the model will include its step-by-step thinking in the response.
IncludeThoughts bool `json:"include_thoughts,omitempty"`
}
// ToolDeclaration defines the structure for declaring tools (like functions)
// that the model can call during content generation.
type ToolDeclaration struct {
// FunctionDeclarations is a list of available functions that the model can call.
FunctionDeclarations []interface{} `json:"functionDeclarations"`
}

242
internal/client/codex_client.go
View File

@@ -1,3 +1,6 @@
// Package client defines the interface and base structure for AI API clients.
// It provides a common interface that all supported AI service clients must implement,
// including methods for sending messages, handling streams, and managing authentication.
package client
import (
@@ -17,6 +20,9 @@ import (
"github.com/luispater/CLIProxyAPI/internal/auth"
"github.com/luispater/CLIProxyAPI/internal/auth/codex"
"github.com/luispater/CLIProxyAPI/internal/config"
. "github.com/luispater/CLIProxyAPI/internal/constant"
"github.com/luispater/CLIProxyAPI/internal/interfaces"
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
"github.com/luispater/CLIProxyAPI/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
@@ -34,6 +40,14 @@ type CodexClient struct {
}
// NewCodexClient creates a new OpenAI client instance
//
// Parameters:
// - cfg: The application configuration.
// - ts: The token storage for Codex authentication.
//
// Returns:
// - *CodexClient: A new Codex client instance.
// - error: An error if the client creation fails.
func NewCodexClient(cfg *config.Config, ts *codex.CodexTokenStorage) (*CodexClient, error) {
httpClient := util.SetProxy(cfg, &http.Client{})
client := &CodexClient{
@@ -50,43 +64,61 @@ func NewCodexClient(cfg *config.Config, ts *codex.CodexTokenStorage) (*CodexClie
return client, nil
}
// Type returns the client type
func (c *CodexClient) Type() string {
return CODEX
}
// Provider returns the provider name for this client.
func (c *CodexClient) Provider() string {
return CODEX
}
// CanProvideModel checks if this client can provide the specified model.
//
// Parameters:
// - modelName: The name of the model to check.
//
// Returns:
// - bool: True if the model is supported, false otherwise.
func (c *CodexClient) CanProvideModel(modelName string) bool {
models := []string{
"gpt-5",
"gpt-5-mini",
"gpt-5-nano",
"gpt-5-high",
"codex-mini-latest",
}
return util.InArray(models, modelName)
}
// GetUserAgent returns the user agent string for OpenAI API requests
func (c *CodexClient) GetUserAgent() string {
return "codex-cli"
}
// TokenStorage returns the token storage for this client.
func (c *CodexClient) TokenStorage() auth.TokenStorage {
return c.tokenStorage
}
// SendMessage sends a message to OpenAI API (non-streaming)
func (c *CodexClient) SendMessage(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration) ([]byte, *ErrorMessage) {
// For now, return an error as OpenAI integration is not fully implemented
return nil, &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("codex message sending not yet implemented"),
}
}
// SendMessageStream sends a streaming message to OpenAI API
func (c *CodexClient) SendMessageStream(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration, _ ...bool) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage, 1)
errChan <- &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("codex streaming not yet implemented"),
}
close(errChan)
return nil, errChan
}
// SendRawMessage sends a raw message to OpenAI API
func (c *CodexClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt string) ([]byte, *ErrorMessage) {
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: The response body.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *CodexClient) SendRawMessage(ctx context.Context, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, false)
respBody, err := c.APIRequest(ctx, "/codex/responses", rawJSON, alt, false)
respBody, err := c.APIRequest(ctx, modelName, "/codex/responses", rawJSON, alt, false)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
@@ -97,49 +129,89 @@ func (c *CodexClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt st
delete(c.modelQuotaExceeded, modelName)
bodyBytes, errReadAll := io.ReadAll(respBody)
if errReadAll != nil {
return nil, &ErrorMessage{StatusCode: 500, Error: errReadAll}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
}
c.AddAPIResponseData(ctx, bodyBytes)
var param any
bodyBytes = []byte(translator.ResponseNonStream(handlerType, c.Type(), ctx, modelName, bodyBytes, &param))
return bodyBytes, nil
}
// SendRawMessageStream sends a raw streaming message to OpenAI API
func (c *CodexClient) SendRawMessageStream(ctx context.Context, rawJSON []byte, alt string) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage)
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - <-chan []byte: A channel for receiving response data chunks.
// - <-chan *interfaces.ErrorMessage: A channel for receiving error messages.
func (c *CodexClient) SendRawMessageStream(ctx context.Context, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
// return dataChan, errChan
go func() {
defer close(errChan)
defer close(dataChan)
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
var stream io.ReadCloser
for {
var err *ErrorMessage
stream, err = c.APIRequest(ctx, "/codex/responses", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
if c.IsModelQuotaExceeded(modelName) {
errChan <- &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
delete(c.modelQuotaExceeded, modelName)
break
return
}
var err *interfaces.ErrorMessage
stream, err = c.APIRequest(ctx, modelName, "/codex/responses", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
}
delete(c.modelQuotaExceeded, modelName)
scanner := bufio.NewScanner(stream)
buffer := make([]byte, 10240*1024)
scanner.Buffer(buffer, 10240*1024)
for scanner.Scan() {
line := scanner.Bytes()
dataChan <- line
if translator.NeedConvert(handlerType, c.Type()) {
var param any
for scanner.Scan() {
line := scanner.Bytes()
lines := translator.Response(handlerType, c.Type(), ctx, modelName, line, &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
c.AddAPIResponseData(ctx, line)
}
} else {
for scanner.Scan() {
line := scanner.Bytes()
dataChan <- line
c.AddAPIResponseData(ctx, line)
}
}
if errScanner := scanner.Err(); errScanner != nil {
errChan <- &ErrorMessage{500, errScanner, nil}
errChan <- &interfaces.ErrorMessage{StatusCode: 500, Error: errScanner}
_ = stream.Close()
return
}
@@ -151,20 +223,39 @@ func (c *CodexClient) SendRawMessageStream(ctx context.Context, rawJSON []byte,
}
// SendRawTokenCount sends a token count request to OpenAI API
func (c *CodexClient) SendRawTokenCount(_ context.Context, _ []byte, _ string) ([]byte, *ErrorMessage) {
return nil, &ErrorMessage{
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: Always nil for this implementation.
// - *interfaces.ErrorMessage: An error message indicating that the feature is not implemented.
func (c *CodexClient) SendRawTokenCount(_ context.Context, _ string, _ []byte, _ string) ([]byte, *interfaces.ErrorMessage) {
return nil, &interfaces.ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("codex token counting not yet implemented"),
}
}
// SaveTokenToFile persists the token storage to disk
//
// Returns:
// - error: An error if the save operation fails, nil otherwise.
func (c *CodexClient) SaveTokenToFile() error {
fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("codex-%s.json", c.tokenStorage.(*codex.CodexTokenStorage).Email))
return c.tokenStorage.SaveTokenToFile(fileName)
}
// RefreshTokens refreshes the access tokens if needed
//
// Parameters:
// - ctx: The context for the request.
//
// Returns:
// - error: An error if the refresh operation fails, nil otherwise.
func (c *CodexClient) RefreshTokens(ctx context.Context) error {
if c.tokenStorage == nil || c.tokenStorage.(*codex.CodexTokenStorage).RefreshToken == "" {
return fmt.Errorf("no refresh token available")
@@ -189,7 +280,19 @@ func (c *CodexClient) RefreshTokens(ctx context.Context) error {
}
// APIRequest handles making requests to the CLI API endpoints.
func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *ErrorMessage) {
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - endpoint: The API endpoint to call.
// - body: The request body.
// - alt: An alternative response format parameter.
// - stream: A boolean indicating if the request is for a streaming response.
//
// Returns:
// - io.ReadCloser: The response body reader.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *CodexClient) APIRequest(ctx context.Context, modelName, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *interfaces.ErrorMessage) {
var jsonBody []byte
var err error
if byteBody, ok := body.([]byte); ok {
@@ -197,7 +300,7 @@ func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body inte
} else {
jsonBody, err = json.Marshal(body)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to marshal request body: %w", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to marshal request body: %w", err)}
}
}
@@ -220,6 +323,20 @@ func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body inte
// Stream must be set to true
jsonBody, _ = sjson.SetBytes(jsonBody, "stream", true)
if util.InArray([]string{"gpt-5-nano", "gpt-5-mini", "gpt-5", "gpt-5-high"}, modelName) {
jsonBody, _ = sjson.SetBytes(jsonBody, "model", "gpt-5")
switch modelName {
case "gpt-5-nano":
jsonBody, _ = sjson.SetBytes(jsonBody, "reasoning.effort", "minimal")
case "gpt-5-mini":
jsonBody, _ = sjson.SetBytes(jsonBody, "reasoning.effort", "low")
case "gpt-5":
jsonBody, _ = sjson.SetBytes(jsonBody, "reasoning.effort", "medium")
case "gpt-5-high":
jsonBody, _ = sjson.SetBytes(jsonBody, "reasoning.effort", "high")
}
}
url := fmt.Sprintf("%s%s", chatGPTEndpoint, endpoint)
// log.Debug(string(jsonBody))
@@ -228,7 +345,7 @@ func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body inte
req, err := http.NewRequestWithContext(ctx, "POST", url, reqBody)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to create request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to create request: %v", err)}
}
sessionID := uuid.New().String()
@@ -242,13 +359,17 @@ func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body inte
req.Header.Set("Originator", "codex_cli_rs")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.tokenStorage.(*codex.CodexTokenStorage).AccessToken))
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
if c.cfg.RequestLog {
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
}
}
log.Debugf("Use ChatGPT account %s for model %s", c.GetEmail(), modelName)
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to execute request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to execute request: %v", err)}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
@@ -259,18 +380,25 @@ func (c *CodexClient) APIRequest(ctx context.Context, endpoint string, body inte
}()
bodyBytes, _ := io.ReadAll(resp.Body)
// log.Debug(string(jsonBody))
return nil, &ErrorMessage{resp.StatusCode, fmt.Errorf(string(bodyBytes)), nil}
return nil, &interfaces.ErrorMessage{StatusCode: resp.StatusCode, Error: fmt.Errorf("%s", string(bodyBytes))}
}
return resp.Body, nil
}
// GetEmail returns the email associated with the client's token storage.
func (c *CodexClient) GetEmail() string {
return c.tokenStorage.(*codex.CodexTokenStorage).Email
}
// IsModelQuotaExceeded returns true if the specified model has exceeded its quota
// and no fallback options are available.
//
// Parameters:
// - model: The name of the model to check.
//
// Returns:
// - bool: True if the model's quota is exceeded, false otherwise.
func (c *CodexClient) IsModelQuotaExceeded(model string) bool {
if lastExceededTime, hasKey := c.modelQuotaExceeded[model]; hasKey {
duration := time.Now().Sub(*lastExceededTime)

826
internal/client/gemini-cli_client.go Normal file
View File

@@ -0,0 +1,826 @@
// Package client defines the interface and base structure for AI API clients.
// It provides a common interface that all supported AI service clients must implement,
// including methods for sending messages, handling streams, and managing authentication.
package client
import (
"bufio"
"bytes"
"context"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"strings"
"sync"
"time"
"github.com/gin-gonic/gin"
geminiAuth "github.com/luispater/CLIProxyAPI/internal/auth/gemini"
"github.com/luispater/CLIProxyAPI/internal/config"
. "github.com/luispater/CLIProxyAPI/internal/constant"
"github.com/luispater/CLIProxyAPI/internal/interfaces"
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
"github.com/luispater/CLIProxyAPI/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
"github.com/tidwall/sjson"
"golang.org/x/oauth2"
)
const (
codeAssistEndpoint = "https://cloudcode-pa.googleapis.com"
apiVersion = "v1internal"
)
var (
previewModels = map[string][]string{
"gemini-2.5-pro": {"gemini-2.5-pro-preview-05-06", "gemini-2.5-pro-preview-06-05"},
"gemini-2.5-flash": {"gemini-2.5-flash-preview-04-17", "gemini-2.5-flash-preview-05-20"},
}
)
// GeminiCLIClient is the main client for interacting with the CLI API.
type GeminiCLIClient struct {
ClientBase
}
// NewGeminiCLIClient creates a new CLI API client.
//
// Parameters:
// - httpClient: The HTTP client to use for requests.
// - ts: The token storage for Gemini authentication.
// - cfg: The application configuration.
//
// Returns:
// - *GeminiCLIClient: A new Gemini CLI client instance.
func NewGeminiCLIClient(httpClient *http.Client, ts *geminiAuth.GeminiTokenStorage, cfg *config.Config) *GeminiCLIClient {
client := &GeminiCLIClient{
ClientBase: ClientBase{
RequestMutex: &sync.Mutex{},
httpClient: httpClient,
cfg: cfg,
tokenStorage: ts,
modelQuotaExceeded: make(map[string]*time.Time),
},
}
return client
}
// Type returns the client type
func (c *GeminiCLIClient) Type() string {
return GEMINICLI
}
// Provider returns the provider name for this client.
func (c *GeminiCLIClient) Provider() string {
return GEMINICLI
}
// CanProvideModel checks if this client can provide the specified model.
//
// Parameters:
// - modelName: The name of the model to check.
//
// Returns:
// - bool: True if the model is supported, false otherwise.
func (c *GeminiCLIClient) CanProvideModel(modelName string) bool {
models := []string{
"gemini-2.5-pro",
"gemini-2.5-flash",
}
return util.InArray(models, modelName)
}
// SetProjectID updates the project ID for the client's token storage.
//
// Parameters:
// - projectID: The new project ID.
func (c *GeminiCLIClient) SetProjectID(projectID string) {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID = projectID
}
// SetIsAuto configures whether the client should operate in automatic mode.
//
// Parameters:
// - auto: A boolean indicating if automatic mode should be enabled.
func (c *GeminiCLIClient) SetIsAuto(auto bool) {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Auto = auto
}
// SetIsChecked sets the checked status for the client's token storage.
//
// Parameters:
// - checked: A boolean indicating if the token storage has been checked.
func (c *GeminiCLIClient) SetIsChecked(checked bool) {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Checked = checked
}
// IsChecked returns whether the client's token storage has been checked.
func (c *GeminiCLIClient) IsChecked() bool {
return c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Checked
}
// IsAuto returns whether the client is operating in automatic mode.
func (c *GeminiCLIClient) IsAuto() bool {
return c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Auto
}
// GetEmail returns the email address associated with the client's token storage.
func (c *GeminiCLIClient) GetEmail() string {
return c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Email
}
// GetProjectID returns the Google Cloud project ID from the client's token storage.
func (c *GeminiCLIClient) GetProjectID() string {
if c.tokenStorage != nil {
if ts, ok := c.tokenStorage.(*geminiAuth.GeminiTokenStorage); ok {
return ts.ProjectID
}
}
return ""
}
// SetupUser performs the initial user onboarding and setup.
//
// Parameters:
// - ctx: The context for the request.
// - email: The user's email address.
// - projectID: The Google Cloud project ID.
//
// Returns:
// - error: An error if the setup fails, nil otherwise.
func (c *GeminiCLIClient) SetupUser(ctx context.Context, email, projectID string) error {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Email = email
log.Info("Performing user onboarding...")
// 1. LoadCodeAssist
loadAssistReqBody := map[string]interface{}{
"metadata": c.getClientMetadata(),
}
if projectID != "" {
loadAssistReqBody["cloudaicompanionProject"] = projectID
}
var loadAssistResp map[string]interface{}
err := c.makeAPIRequest(ctx, "loadCodeAssist", "POST", loadAssistReqBody, &loadAssistResp)
if err != nil {
return fmt.Errorf("failed to load code assist: %w", err)
}
// 2. OnboardUser
var onboardTierID = "legacy-tier"
if tiers, ok := loadAssistResp["allowedTiers"].([]interface{}); ok {
for _, t := range tiers {
if tier, tierOk := t.(map[string]interface{}); tierOk {
if isDefault, isDefaultOk := tier["isDefault"].(bool); isDefaultOk && isDefault {
if id, idOk := tier["id"].(string); idOk {
onboardTierID = id
break
}
}
}
}
}
onboardProjectID := projectID
if p, ok := loadAssistResp["cloudaicompanionProject"].(string); ok && p != "" {
onboardProjectID = p
}
onboardReqBody := map[string]interface{}{
"tierId": onboardTierID,
"metadata": c.getClientMetadata(),
}
if onboardProjectID != "" {
onboardReqBody["cloudaicompanionProject"] = onboardProjectID
} else {
return fmt.Errorf("failed to start user onboarding, need define a project id")
}
for {
var lroResp map[string]interface{}
err = c.makeAPIRequest(ctx, "onboardUser", "POST", onboardReqBody, &lroResp)
if err != nil {
return fmt.Errorf("failed to start user onboarding: %w", err)
}
// a, _ := json.Marshal(&lroResp)
// log.Debug(string(a))
// 3. Poll Long-Running Operation (LRO)
done, doneOk := lroResp["done"].(bool)
if doneOk && done {
if project, projectOk := lroResp["response"].(map[string]interface{})["cloudaicompanionProject"].(map[string]interface{}); projectOk {
if projectID != "" {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID = projectID
} else {
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID = project["id"].(string)
}
log.Infof("Onboarding complete. Using Project ID: %s", c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID)
return nil
}
} else {
log.Println("Onboarding in progress, waiting 5 seconds...")
time.Sleep(5 * time.Second)
}
}
}
// makeAPIRequest handles making requests to the CLI API endpoints.
//
// Parameters:
// - ctx: The context for the request.
// - endpoint: The API endpoint to call.
// - method: The HTTP method to use.
// - body: The request body.
// - result: A pointer to a variable to store the response.
//
// Returns:
// - error: An error if the request fails, nil otherwise.
func (c *GeminiCLIClient) makeAPIRequest(ctx context.Context, endpoint, method string, body interface{}, result interface{}) error {
var reqBody io.Reader
var jsonBody []byte
var err error
if body != nil {
jsonBody, err = json.Marshal(body)
if err != nil {
return fmt.Errorf("failed to marshal request body: %w", err)
}
reqBody = bytes.NewBuffer(jsonBody)
}
url := fmt.Sprintf("%s/%s:%s", codeAssistEndpoint, apiVersion, endpoint)
if strings.HasPrefix(endpoint, "operations/") {
url = fmt.Sprintf("%s/%s", codeAssistEndpoint, endpoint)
}
req, err := http.NewRequestWithContext(ctx, method, url, reqBody)
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
token, err := c.httpClient.Transport.(*oauth2.Transport).Source.Token()
if err != nil {
return fmt.Errorf("failed to get token: %w", err)
}
// Set headers
metadataStr := c.getClientMetadataString()
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", c.GetUserAgent())
req.Header.Set("X-Goog-Api-Client", "gl-node/22.17.0")
req.Header.Set("Client-Metadata", metadataStr)
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
}
resp, err := c.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to execute request: %w", err)
}
defer func() {
if err = resp.Body.Close(); err != nil {
log.Printf("warn: failed to close response body: %v", err)
}
}()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
bodyBytes, _ := io.ReadAll(resp.Body)
return fmt.Errorf("api request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
}
if result != nil {
if err = json.NewDecoder(resp.Body).Decode(result); err != nil {
return fmt.Errorf("failed to decode response body: %w", err)
}
}
return nil
}
// APIRequest handles making requests to the CLI API endpoints.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - endpoint: The API endpoint to call.
// - body: The request body.
// - alt: An alternative response format parameter.
// - stream: A boolean indicating if the request is for a streaming response.
//
// Returns:
// - io.ReadCloser: The response body reader.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *GeminiCLIClient) APIRequest(ctx context.Context, modelName, endpoint string, body interface{}, alt string, stream bool) (io.ReadCloser, *interfaces.ErrorMessage) {
var jsonBody []byte
var err error
if byteBody, ok := body.([]byte); ok {
jsonBody = byteBody
} else {
jsonBody, err = json.Marshal(body)
if err != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to marshal request body: %w", err)}
}
}
var url string
// Add alt=sse for streaming
url = fmt.Sprintf("%s/%s:%s", codeAssistEndpoint, apiVersion, endpoint)
if alt == "" && stream {
url = url + "?alt=sse"
} else {
if alt != "" {
url = url + fmt.Sprintf("?$alt=%s", alt)
}
}
// log.Debug(string(jsonBody))
// log.Debug(url)
reqBody := bytes.NewBuffer(jsonBody)
req, err := http.NewRequestWithContext(ctx, "POST", url, reqBody)
if err != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to create request: %v", err)}
}
// Set headers
metadataStr := c.getClientMetadataString()
req.Header.Set("Content-Type", "application/json")
token, errToken := c.httpClient.Transport.(*oauth2.Transport).Source.Token()
if errToken != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to get token: %v", errToken)}
}
req.Header.Set("User-Agent", c.GetUserAgent())
req.Header.Set("X-Goog-Api-Client", "gl-node/22.17.0")
req.Header.Set("Client-Metadata", metadataStr)
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
if c.cfg.RequestLog {
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
}
}
log.Debugf("Use Gemini CLI account %s (project id: %s) for model %s", c.GetEmail(), c.GetProjectID(), modelName)
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to execute request: %v", err)}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
defer func() {
if err = resp.Body.Close(); err != nil {
log.Printf("warn: failed to close response body: %v", err)
}
}()
bodyBytes, _ := io.ReadAll(resp.Body)
// log.Debug(string(jsonBody))
return nil, &interfaces.ErrorMessage{StatusCode: resp.StatusCode, Error: fmt.Errorf("%s", string(bodyBytes))}
}
return resp.Body, nil
}
// SendRawTokenCount handles a token count.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: The response body.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *GeminiCLIClient) SendRawTokenCount(ctx context.Context, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
for {
if c.isModelQuotaExceeded(modelName) {
if c.cfg.QuotaExceeded.SwitchPreviewModel {
newModelName := c.getPreviewModel(modelName)
if newModelName != "" {
log.Debugf("Model %s is quota exceeded. Switch to preview model %s", modelName, newModelName)
rawJSON, _ = sjson.SetBytes(rawJSON, "model", newModelName)
continue
}
}
return nil, &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
}
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, false)
// Remove project and model from the request body
rawJSON, _ = sjson.DeleteBytes(rawJSON, "project")
rawJSON, _ = sjson.DeleteBytes(rawJSON, "model")
respBody, err := c.APIRequest(ctx, modelName, "countTokens", rawJSON, alt, false)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
if c.cfg.QuotaExceeded.SwitchPreviewModel {
continue
}
}
return nil, err
}
delete(c.modelQuotaExceeded, modelName)
bodyBytes, errReadAll := io.ReadAll(respBody)
if errReadAll != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
}
c.AddAPIResponseData(ctx, bodyBytes)
var param any
bodyBytes = []byte(translator.ResponseNonStream(handlerType, c.Type(), ctx, modelName, bodyBytes, &param))
return bodyBytes, nil
}
}
// SendRawMessage handles a single conversational turn, including tool calls.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: The response body.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *GeminiCLIClient) SendRawMessage(ctx context.Context, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, false)
rawJSON, _ = sjson.SetBytes(rawJSON, "project", c.GetProjectID())
rawJSON, _ = sjson.SetBytes(rawJSON, "model", modelName)
for {
if c.isModelQuotaExceeded(modelName) {
if c.cfg.QuotaExceeded.SwitchPreviewModel {
newModelName := c.getPreviewModel(modelName)
if newModelName != "" {
log.Debugf("Model %s is quota exceeded. Switch to preview model %s", modelName, newModelName)
rawJSON, _ = sjson.SetBytes(rawJSON, "model", newModelName)
continue
}
}
return nil, &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
}
respBody, err := c.APIRequest(ctx, modelName, "generateContent", rawJSON, alt, false)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
if c.cfg.QuotaExceeded.SwitchPreviewModel {
continue
}
}
return nil, err
}
delete(c.modelQuotaExceeded, modelName)
bodyBytes, errReadAll := io.ReadAll(respBody)
if errReadAll != nil {
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
}
c.AddAPIResponseData(ctx, bodyBytes)
newCtx := context.WithValue(ctx, "alt", alt)
var param any
bodyBytes = []byte(translator.ResponseNonStream(handlerType, c.Type(), newCtx, modelName, bodyBytes, &param))
return bodyBytes, nil
}
}
// SendRawMessageStream handles a single conversational turn, including tool calls.
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - <-chan []byte: A channel for receiving response data chunks.
// - <-chan *interfaces.ErrorMessage: A channel for receiving error messages.
func (c *GeminiCLIClient) SendRawMessageStream(ctx context.Context, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
rawJSON, _ = sjson.SetBytes(rawJSON, "project", c.GetProjectID())
rawJSON, _ = sjson.SetBytes(rawJSON, "model", modelName)
dataTag := []byte("data: ")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
// return dataChan, errChan
go func() {
defer close(errChan)
defer close(dataChan)
rawJSON, _ = sjson.SetBytes(rawJSON, "project", c.GetProjectID())
var stream io.ReadCloser
for {
if c.isModelQuotaExceeded(modelName) {
if c.cfg.QuotaExceeded.SwitchPreviewModel {
newModelName := c.getPreviewModel(modelName)
if newModelName != "" {
log.Debugf("Model %s is quota exceeded. Switch to preview model %s", modelName, newModelName)
rawJSON, _ = sjson.SetBytes(rawJSON, "model", newModelName)
continue
}
}
errChan <- &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
return
}
var err *interfaces.ErrorMessage
stream, err = c.APIRequest(ctx, modelName, "streamGenerateContent", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
if c.cfg.QuotaExceeded.SwitchPreviewModel {
continue
}
}
errChan <- err
return
}
delete(c.modelQuotaExceeded, modelName)
break
}
newCtx := context.WithValue(ctx, "alt", alt)
var param any
if alt == "" {
scanner := bufio.NewScanner(stream)
if translator.NeedConvert(handlerType, c.Type()) {
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, line[6:], &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
}
c.AddAPIResponseData(ctx, line)
}
} else {
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
dataChan <- line[6:]
}
c.AddAPIResponseData(ctx, line)
}
}
if errScanner := scanner.Err(); errScanner != nil {
errChan <- &interfaces.ErrorMessage{StatusCode: 500, Error: errScanner}
_ = stream.Close()
return
}
} else {
data, err := io.ReadAll(stream)
if err != nil {
errChan <- &interfaces.ErrorMessage{StatusCode: 500, Error: err}
_ = stream.Close()
return
}
if translator.NeedConvert(handlerType, c.Type()) {
lines := translator.Response(handlerType, c.Type(), newCtx, modelName, data, &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
} else {
dataChan <- data
}
c.AddAPIResponseData(ctx, data)
}
if translator.NeedConvert(handlerType, c.Type()) {
lines := translator.Response(handlerType, c.Type(), ctx, modelName, []byte("[DONE]"), &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
}
_ = stream.Close()
}()
return dataChan, errChan
}
// isModelQuotaExceeded checks if the specified model has exceeded its quota
// within the last 30 minutes.
//
// Parameters:
// - model: The name of the model to check.
//
// Returns:
// - bool: True if the model's quota is exceeded, false otherwise.
func (c *GeminiCLIClient) isModelQuotaExceeded(model string) bool {
if lastExceededTime, hasKey := c.modelQuotaExceeded[model]; hasKey {
duration := time.Now().Sub(*lastExceededTime)
if duration > 30*time.Minute {
return false
}
return true
}
return false
}
// getPreviewModel returns an available preview model for the given base model,
// or an empty string if no preview models are available or all are quota exceeded.
//
// Parameters:
// - model: The base model name.
//
// Returns:
// - string: The name of the preview model to use, or an empty string.
func (c *GeminiCLIClient) getPreviewModel(model string) string {
if models, hasKey := previewModels[model]; hasKey {
for i := 0; i < len(models); i++ {
if !c.isModelQuotaExceeded(models[i]) {
return models[i]
}
}
}
return ""
}
// IsModelQuotaExceeded returns true if the specified model has exceeded its quota
// and no fallback options are available.
//
// Parameters:
// - model: The name of the model to check.
//
// Returns:
// - bool: True if the model's quota is exceeded, false otherwise.
func (c *GeminiCLIClient) IsModelQuotaExceeded(model string) bool {
if c.isModelQuotaExceeded(model) {
if c.cfg.QuotaExceeded.SwitchPreviewModel {
return c.getPreviewModel(model) == ""
}
return true
}
return false
}
// CheckCloudAPIIsEnabled sends a simple test request to the API to verify
// that the Cloud AI API is enabled for the user's project. It provides
// an activation URL if the API is disabled.
//
// Returns:
// - bool: True if the API is enabled, false otherwise.
// - error: An error if the request fails, nil otherwise.
func (c *GeminiCLIClient) CheckCloudAPIIsEnabled() (bool, error) {
ctx, cancel := context.WithCancel(context.Background())
defer func() {
c.RequestMutex.Unlock()
cancel()
}()
c.RequestMutex.Lock()
// A simple request to test the API endpoint.
requestBody := fmt.Sprintf(`{"project":"%s","request":{"contents":[{"role":"user","parts":[{"text":"Be concise. What is the capital of France?"}]}],"generationConfig":{"thinkingConfig":{"include_thoughts":false,"thinkingBudget":0}}},"model":"gemini-2.5-flash"}`, c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID)
stream, err := c.APIRequest(ctx, "gemini-2.5-flash", "streamGenerateContent", []byte(requestBody), "", true)
if err != nil {
// If a 403 Forbidden error occurs, it likely means the API is not enabled.
if err.StatusCode == 403 {
errJSON := err.Error.Error()
// Check for a specific error code and extract the activation URL.
if gjson.Get(errJSON, "0.error.code").Int() == 403 {
activationURL := gjson.Get(errJSON, "0.error.details.0.metadata.activationUrl").String()
if activationURL != "" {
log.Warnf(
"\n\nPlease activate your account with this url:\n\n%s\n\n And execute this command again:\n%s --login --project_id %s",
activationURL,
os.Args[0],
c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID,
)
}
}
log.Warnf("\n\nPlease copy this message and create an issue.\n\n%s\n\n", errJSON)
return false, nil
}
return false, err.Error
}
defer func() {
_ = stream.Close()
}()
// We only need to know if the request was successful, so we can drain the stream.
scanner := bufio.NewScanner(stream)
for scanner.Scan() {
// Do nothing, just consume the stream.
}
return scanner.Err() == nil, scanner.Err()
}
// GetProjectList fetches a list of Google Cloud projects accessible by the user.
//
// Parameters:
// - ctx: The context for the request.
//
// Returns:
// - *interfaces.GCPProject: A list of GCP projects.
// - error: An error if the request fails, nil otherwise.
func (c *GeminiCLIClient) GetProjectList(ctx context.Context) (*interfaces.GCPProject, error) {
token, err := c.httpClient.Transport.(*oauth2.Transport).Source.Token()
if err != nil {
return nil, fmt.Errorf("failed to get token: %w", err)
}
req, err := http.NewRequestWithContext(ctx, "GET", "https://cloudresourcemanager.googleapis.com/v1/projects", nil)
if err != nil {
return nil, fmt.Errorf("could not create project list request: %v", err)
}
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token.AccessToken))
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to execute project list request: %w", err)
}
defer func() {
_ = resp.Body.Close()
}()
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
bodyBytes, _ := io.ReadAll(resp.Body)
return nil, fmt.Errorf("project list request failed with status %d: %s", resp.StatusCode, string(bodyBytes))
}
var project interfaces.GCPProject
if err = json.NewDecoder(resp.Body).Decode(&project); err != nil {
return nil, fmt.Errorf("failed to unmarshal project list: %w", err)
}
return &project, nil
}
// SaveTokenToFile serializes the client's current token storage to a JSON file.
// The filename is constructed from the user's email and project ID.
//
// Returns:
// - error: An error if the save operation fails, nil otherwise.
func (c *GeminiCLIClient) SaveTokenToFile() error {
fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("%s-%s.json", c.tokenStorage.(*geminiAuth.GeminiTokenStorage).Email, c.tokenStorage.(*geminiAuth.GeminiTokenStorage).ProjectID))
log.Infof("Saving credentials to %s", fileName)
return c.tokenStorage.SaveTokenToFile(fileName)
}
// getClientMetadata returns a map of metadata about the client environment,
// such as IDE type, platform, and plugin version.
func (c *GeminiCLIClient) getClientMetadata() map[string]string {
return map[string]string{
"ideType": "IDE_UNSPECIFIED",
"platform": "PLATFORM_UNSPECIFIED",
"pluginType": "GEMINI",
// "pluginVersion": pluginVersion,
}
}
// getClientMetadataString returns the client metadata as a single,
// comma-separated string, which is required for the 'GeminiClient-Metadata' header.
func (c *GeminiCLIClient) getClientMetadataString() string {
md := c.getClientMetadata()
parts := make([]string, 0, len(md))
for k, v := range md {
parts = append(parts, fmt.Sprintf("%s=%s", k, v))
}
return strings.Join(parts, ",")
}
// GetUserAgent constructs the User-Agent string for HTTP requests.
func (c *GeminiCLIClient) GetUserAgent() string {
// return fmt.Sprintf("GeminiCLI/%s (%s; %s)", pluginVersion, runtime.GOOS, runtime.GOARCH)
return "google-api-nodejs-client/9.15.1"
}

989
internal/client/gemini_client.go
View File

File diff suppressed because it is too large Load Diff

234
internal/client/qwen_client.go
View File

@@ -1,3 +1,6 @@
// Package client defines the interface and base structure for AI API clients.
// It provides a common interface that all supported AI service clients must implement,
// including methods for sending messages, handling streams, and managing authentication.
package client
import (
@@ -17,6 +20,9 @@ import (
"github.com/luispater/CLIProxyAPI/internal/auth"
"github.com/luispater/CLIProxyAPI/internal/auth/qwen"
"github.com/luispater/CLIProxyAPI/internal/config"
. "github.com/luispater/CLIProxyAPI/internal/constant"
"github.com/luispater/CLIProxyAPI/internal/interfaces"
"github.com/luispater/CLIProxyAPI/internal/translator/translator"
"github.com/luispater/CLIProxyAPI/internal/util"
log "github.com/sirupsen/logrus"
"github.com/tidwall/gjson"
@@ -34,6 +40,13 @@ type QwenClient struct {
}
// NewQwenClient creates a new OpenAI client instance
//
// Parameters:
// - cfg: The application configuration.
// - ts: The token storage for Qwen authentication.
//
// Returns:
// - *QwenClient: A new Qwen client instance.
func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient {
httpClient := util.SetProxy(cfg, &http.Client{})
client := &QwenClient{
@@ -50,43 +63,58 @@ func NewQwenClient(cfg *config.Config, ts *qwen.QwenTokenStorage) *QwenClient {
return client
}
// Type returns the client type
func (c *QwenClient) Type() string {
return OPENAI
}
// Provider returns the provider name for this client.
func (c *QwenClient) Provider() string {
return "qwen"
}
// CanProvideModel checks if this client can provide the specified model.
//
// Parameters:
// - modelName: The name of the model to check.
//
// Returns:
// - bool: True if the model is supported, false otherwise.
func (c *QwenClient) CanProvideModel(modelName string) bool {
models := []string{
"qwen3-coder-plus",
"qwen3-coder-flash",
}
return util.InArray(models, modelName)
}
// GetUserAgent returns the user agent string for OpenAI API requests
func (c *QwenClient) GetUserAgent() string {
return "google-api-nodejs-client/9.15.1"
}
// TokenStorage returns the token storage for this client.
func (c *QwenClient) TokenStorage() auth.TokenStorage {
return c.tokenStorage
}
// SendMessage sends a message to OpenAI API (non-streaming)
func (c *QwenClient) SendMessage(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration) ([]byte, *ErrorMessage) {
// For now, return an error as OpenAI integration is not fully implemented
return nil, &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("qwen message sending not yet implemented"),
}
}
// SendMessageStream sends a streaming message to OpenAI API
func (c *QwenClient) SendMessageStream(_ context.Context, _ []byte, _ string, _ *Content, _ []Content, _ []ToolDeclaration, _ ...bool) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage, 1)
errChan <- &ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("qwen streaming not yet implemented"),
}
close(errChan)
return nil, errChan
}
// SendRawMessage sends a raw message to OpenAI API
func (c *QwenClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt string) ([]byte, *ErrorMessage) {
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: The response body.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *QwenClient) SendRawMessage(ctx context.Context, modelName string, rawJSON []byte, alt string) ([]byte, *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, false)
respBody, err := c.APIRequest(ctx, "/chat/completions", rawJSON, alt, false)
respBody, err := c.APIRequest(ctx, modelName, "/chat/completions", rawJSON, alt, false)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
@@ -97,49 +125,97 @@ func (c *QwenClient) SendRawMessage(ctx context.Context, rawJSON []byte, alt str
delete(c.modelQuotaExceeded, modelName)
bodyBytes, errReadAll := io.ReadAll(respBody)
if errReadAll != nil {
return nil, &ErrorMessage{StatusCode: 500, Error: errReadAll}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: errReadAll}
}
c.AddAPIResponseData(ctx, bodyBytes)
var param any
bodyBytes = []byte(translator.ResponseNonStream(handlerType, c.Type(), ctx, modelName, bodyBytes, &param))
return bodyBytes, nil
}
// SendRawMessageStream sends a raw streaming message to OpenAI API
func (c *QwenClient) SendRawMessageStream(ctx context.Context, rawJSON []byte, alt string) (<-chan []byte, <-chan *ErrorMessage) {
errChan := make(chan *ErrorMessage)
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - <-chan []byte: A channel for receiving response data chunks.
// - <-chan *interfaces.ErrorMessage: A channel for receiving error messages.
func (c *QwenClient) SendRawMessageStream(ctx context.Context, modelName string, rawJSON []byte, alt string) (<-chan []byte, <-chan *interfaces.ErrorMessage) {
handler := ctx.Value("handler").(interfaces.APIHandler)
handlerType := handler.HandlerType()
rawJSON = translator.Request(handlerType, c.Type(), modelName, rawJSON, true)
dataTag := []byte("data: ")
doneTag := []byte("data: [DONE]")
errChan := make(chan *interfaces.ErrorMessage)
dataChan := make(chan []byte)
// log.Debugf(string(rawJSON))
// return dataChan, errChan
go func() {
defer close(errChan)
defer close(dataChan)
modelResult := gjson.GetBytes(rawJSON, "model")
model := modelResult.String()
modelName := model
var stream io.ReadCloser
for {
var err *ErrorMessage
stream, err = c.APIRequest(ctx, "/chat/completions", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
if c.IsModelQuotaExceeded(modelName) {
errChan <- &interfaces.ErrorMessage{
StatusCode: 429,
Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName),
}
delete(c.modelQuotaExceeded, modelName)
break
return
}
var err *interfaces.ErrorMessage
stream, err = c.APIRequest(ctx, modelName, "/chat/completions", rawJSON, alt, true)
if err != nil {
if err.StatusCode == 429 {
now := time.Now()
c.modelQuotaExceeded[modelName] = &now
}
errChan <- err
return
}
delete(c.modelQuotaExceeded, modelName)
scanner := bufio.NewScanner(stream)
buffer := make([]byte, 10240*1024)
scanner.Buffer(buffer, 10240*1024)
for scanner.Scan() {
line := scanner.Bytes()
dataChan <- line
if translator.NeedConvert(handlerType, c.Type()) {
var param any
for scanner.Scan() {
line := scanner.Bytes()
if bytes.HasPrefix(line, dataTag) {
lines := translator.Response(handlerType, c.Type(), ctx, modelName, line[6:], &param)
for i := 0; i < len(lines); i++ {
dataChan <- []byte(lines[i])
}
}
c.AddAPIResponseData(ctx, line)
}
} else {
for scanner.Scan() {
line := scanner.Bytes()
if !bytes.HasPrefix(line, doneTag) {
if bytes.HasPrefix(line, dataTag) {
dataChan <- line[6:]
}
}
c.AddAPIResponseData(ctx, line)
}
}
if errScanner := scanner.Err(); errScanner != nil {
errChan <- &ErrorMessage{500, errScanner, nil}
errChan <- &interfaces.ErrorMessage{StatusCode: 500, Error: errScanner}
_ = stream.Close()
return
}
@@ -151,20 +227,39 @@ func (c *QwenClient) SendRawMessageStream(ctx context.Context, rawJSON []byte, a
}
// SendRawTokenCount sends a token count request to OpenAI API
func (c *QwenClient) SendRawTokenCount(_ context.Context, _ []byte, _ string) ([]byte, *ErrorMessage) {
return nil, &ErrorMessage{
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - rawJSON: The raw JSON request body.
// - alt: An alternative response format parameter.
//
// Returns:
// - []byte: Always nil for this implementation.
// - *interfaces.ErrorMessage: An error message indicating that the feature is not implemented.
func (c *QwenClient) SendRawTokenCount(_ context.Context, _ string, _ []byte, _ string) ([]byte, *interfaces.ErrorMessage) {
return nil, &interfaces.ErrorMessage{
StatusCode: http.StatusNotImplemented,
Error: fmt.Errorf("qwen token counting not yet implemented"),
}
}
// SaveTokenToFile persists the token storage to disk
//
// Returns:
// - error: An error if the save operation fails, nil otherwise.
func (c *QwenClient) SaveTokenToFile() error {
fileName := filepath.Join(c.cfg.AuthDir, fmt.Sprintf("qwen-%s.json", c.tokenStorage.(*qwen.QwenTokenStorage).Email))
return c.tokenStorage.SaveTokenToFile(fileName)
}
// RefreshTokens refreshes the access tokens if needed
//
// Parameters:
// - ctx: The context for the request.
//
// Returns:
// - error: An error if the refresh operation fails, nil otherwise.
func (c *QwenClient) RefreshTokens(ctx context.Context) error {
if c.tokenStorage == nil || c.tokenStorage.(*qwen.QwenTokenStorage).RefreshToken == "" {
return fmt.Errorf("no refresh token available")
@@ -189,7 +284,19 @@ func (c *QwenClient) RefreshTokens(ctx context.Context) error {
}
// APIRequest handles making requests to the CLI API endpoints.
func (c *QwenClient) APIRequest(ctx context.Context, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *ErrorMessage) {
//
// Parameters:
// - ctx: The context for the request.
// - modelName: The name of the model to use.
// - endpoint: The API endpoint to call.
// - body: The request body.
// - alt: An alternative response format parameter.
// - stream: A boolean indicating if the request is for a streaming response.
//
// Returns:
// - io.ReadCloser: The response body reader.
// - *interfaces.ErrorMessage: An error message if the request fails.
func (c *QwenClient) APIRequest(ctx context.Context, modelName, endpoint string, body interface{}, _ string, _ bool) (io.ReadCloser, *interfaces.ErrorMessage) {
var jsonBody []byte
var err error
if byteBody, ok := body.([]byte); ok {
@@ -197,7 +304,7 @@ func (c *QwenClient) APIRequest(ctx context.Context, endpoint string, body inter
} else {
jsonBody, err = json.Marshal(body)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to marshal request body: %w", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to marshal request body: %w", err)}
}
}
@@ -219,7 +326,7 @@ func (c *QwenClient) APIRequest(ctx context.Context, endpoint string, body inter
req, err := http.NewRequestWithContext(ctx, "POST", url, reqBody)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to create request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to create request: %v", err)}
}
// Set headers
@@ -229,13 +336,17 @@ func (c *QwenClient) APIRequest(ctx context.Context, endpoint string, body inter
req.Header.Set("Client-Metadata", c.getClientMetadataString())
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.tokenStorage.(*qwen.QwenTokenStorage).AccessToken))
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
if c.cfg.RequestLog {
if ginContext, ok := ctx.Value("gin").(*gin.Context); ok {
ginContext.Set("API_REQUEST", jsonBody)
}
}
log.Debugf("Use Qwen Code account %s for model %s", c.GetEmail(), modelName)
resp, err := c.httpClient.Do(req)
if err != nil {
return nil, &ErrorMessage{500, fmt.Errorf("failed to execute request: %v", err), nil}
return nil, &interfaces.ErrorMessage{StatusCode: 500, Error: fmt.Errorf("failed to execute request: %v", err)}
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
@@ -246,12 +357,13 @@ func (c *QwenClient) APIRequest(ctx context.Context, endpoint string, body inter
}()
bodyBytes, _ := io.ReadAll(resp.Body)
// log.Debug(string(jsonBody))
return nil, &ErrorMessage{resp.StatusCode, fmt.Errorf(string(bodyBytes)), nil}
return nil, &interfaces.ErrorMessage{StatusCode: resp.StatusCode, Error: fmt.Errorf("%s", string(bodyBytes))}
}
return resp.Body, nil
}
// getClientMetadata returns a map of metadata about the client environment.
func (c *QwenClient) getClientMetadata() map[string]string {
return map[string]string{
"ideType": "IDE_UNSPECIFIED",
@@ -261,6 +373,7 @@ func (c *QwenClient) getClientMetadata() map[string]string {
}
}
// getClientMetadataString returns the client metadata as a single, comma-separated string.
func (c *QwenClient) getClientMetadataString() string {
md := c.getClientMetadata()
parts := make([]string, 0, len(md))
@@ -270,12 +383,19 @@ func (c *QwenClient) getClientMetadataString() string {
return strings.Join(parts, ",")
}
// GetEmail returns the email associated with the client's token storage.
func (c *QwenClient) GetEmail() string {
return c.tokenStorage.(*qwen.QwenTokenStorage).Email
}
// IsModelQuotaExceeded returns true if the specified model has exceeded its quota
// and no fallback options are available.
//
// Parameters:
// - model: The name of the model to check.
//
// Returns:
// - bool: True if the model's quota is exceeded, false otherwise.
func (c *QwenClient) IsModelQuotaExceeded(model string) bool {
if lastExceededTime, hasKey := c.modelQuotaExceeded[model]; hasKey {
duration := time.Now().Sub(*lastExceededTime)