Enhance response and request handling in translators

- Refactored response handling to process multiple content parts effectively.
- Improved `tool_calls` structure with unique ID generation and enhanced mapping logic.
- Simplified `SystemInstruction` and tool message parsing in requests for better accuracy.
- Enhanced handling of function calls and tool responses with improved data integration.

README.md

@@ -11,6 +11,7 @@ A proxy server that provides an OpenAI-compatible API interface for CLI. This al
 - Multiple account support with load balancing
 - Simple CLI authentication flow
 - Support for Generative Language API Key
+- Support Gemini CLI with multiple account load balancing
 
 ## Installation
 
@@ -172,6 +173,13 @@ debug: false
 api-keys:
   - "your-api-key-1"
   - "your-api-key-2"
+
+# API keys for official Generative Language API
+generative-language-api-key:
+  - "AIzaSy...01"
+  - "AIzaSy...02"
+  - "AIzaSy...03"
+  - "AIzaSy...04"
 ```
 
 ### Authentication Directory
@@ -186,6 +194,20 @@ The `api-keys` parameter allows you to define a list of API keys that can be use
 Authorization: Bearer your-api-key-1
 ```
 
+### Official Generative Language API
+
+The `generative-language-api-key` parameter allows you to define a list of API keys that can be used to authenticate requests to the official Generative Language API.
+
+## Gemini CLI with multiple account load balancing
+
+Start CLI Proxy API server, and then set the `CODE_ASSIST_ENDPOINT` environment variable to the URL of the CLI Proxy API server.
+
+```bash
+export CODE_ASSIST_ENDPOINT="http://127.0.0.1:8317"
+```
+
+The server will relay the `loadCodeAssist`, `onboardUser`, and `countTokens` requests. And automatically load balance the text generation requests between the multiple accounts.
+
 ## Contributing
 
 Contributions are welcome! Please feel free to submit a Pull Request.

internal/api/cli-handlers.go

@@ -0,0 +1,228 @@
+package api
+
+import (
+	"bytes"
+	"context"
+	"fmt"
+	"github.com/gin-gonic/gin"
+	"github.com/luispater/CLIProxyAPI/internal/client"
+	"github.com/luispater/CLIProxyAPI/internal/util"
+	log "github.com/sirupsen/logrus"
+	"github.com/tidwall/gjson"
+	"github.com/tidwall/sjson"
+	"io"
+	"net/http"
+	"time"
+)
+
+func (h *APIHandlers) CLIHandler(c *gin.Context) {
+	rawJson, _ := c.GetRawData()
+	requestRawURI := c.Request.URL.Path
+	if requestRawURI == "/v1internal:generateContent" {
+		h.internalGenerateContent(c, rawJson)
+	} else if requestRawURI == "/v1internal:streamGenerateContent" {
+		h.internalStreamGenerateContent(c, rawJson)
+	} else {
+		reqBody := bytes.NewBuffer(rawJson)
+		req, err := http.NewRequest("POST", fmt.Sprintf("https://cloudcode-pa.googleapis.com%s", c.Request.URL.RequestURI()), reqBody)
+		if err != nil {
+			c.JSON(http.StatusBadRequest, ErrorResponse{
+				Error: ErrorDetail{
+					Message: fmt.Sprintf("Invalid request: %v", err),
+					Type:    "invalid_request_error",
+				},
+			})
+			return
+		}
+		for key, value := range c.Request.Header {
+			req.Header[key] = value
+		}
+
+		httpClient, err := util.SetProxy(h.cfg, &http.Client{})
+		if err != nil {
+			log.Fatalf("set proxy failed: %v", err)
+		}
+
+		resp, err := httpClient.Do(req)
+		if err != nil {
+			c.JSON(http.StatusBadRequest, ErrorResponse{
+				Error: ErrorDetail{
+					Message: fmt.Sprintf("Invalid request: %v", err),
+					Type:    "invalid_request_error",
+				},
+			})
+			return
+		}
+
+		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)
+
+			c.JSON(http.StatusBadRequest, ErrorResponse{
+				Error: ErrorDetail{
+					Message: string(bodyBytes),
+					Type:    "invalid_request_error",
+				},
+			})
+			return
+		}
+
+		defer func() {
+			_ = resp.Body.Close()
+		}()
+
+		for key, value := range resp.Header {
+			c.Header(key, value[0])
+		}
+		output, err := io.ReadAll(resp.Body)
+		if err != nil {
+			log.Errorf("Failed to read response body: %v", err)
+			return
+		}
+		_, _ = c.Writer.Write(output)
+	}
+}
+
+func (h *APIHandlers) internalStreamGenerateContent(c *gin.Context, rawJson []byte) {
+	// Get the http.Flusher interface to manually flush the response.
+	flusher, ok := c.Writer.(http.Flusher)
+	if !ok {
+		c.JSON(http.StatusInternalServerError, ErrorResponse{
+			Error: ErrorDetail{
+				Message: "Streaming not supported",
+				Type:    "server_error",
+			},
+		})
+		return
+	}
+
+	modelResult := gjson.GetBytes(rawJson, "model")
+	modelName := modelResult.String()
+
+	cliCtx, cliCancel := context.WithCancel(context.Background())
+	var cliClient *client.Client
+	defer func() {
+		// Ensure the client's mutex is unlocked on function exit.
+		if cliClient != nil {
+			cliClient.RequestMutex.Unlock()
+		}
+	}()
+
+outLoop:
+	for {
+		var errorResponse *client.ErrorMessage
+		cliClient, errorResponse = h.getClient(modelName)
+		if errorResponse != nil {
+			c.Status(errorResponse.StatusCode)
+			_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
+			flusher.Flush()
+			cliCancel()
+			return
+		}
+
+		if glAPIKey := cliClient.GetGenerativeLanguageAPIKey(); glAPIKey != "" {
+			log.Debugf("Request use generative language API Key: %s", glAPIKey)
+		} else {
+			log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
+		}
+		// Send the message and receive response chunks and errors via channels.
+		respChan, errChan := cliClient.SendRawMessageStream(cliCtx, rawJson)
+		hasFirstResponse := false
+		for {
+			select {
+			// Handle client disconnection.
+			case <-c.Request.Context().Done():
+				if c.Request.Context().Err().Error() == "context canceled" {
+					log.Debugf("Client disconnected: %v", c.Request.Context().Err())
+					cliCancel() // Cancel the backend request.
+					return
+				}
+			// Process incoming response chunks.
+			case chunk, okStream := <-respChan:
+				if !okStream {
+					cliCancel()
+					return
+				} else {
+					hasFirstResponse = true
+					if cliClient.GetGenerativeLanguageAPIKey() != "" {
+						chunk, _ = sjson.SetRawBytes(chunk, "response", chunk)
+					}
+					_, _ = c.Writer.Write([]byte("data: "))
+					_, _ = c.Writer.Write(chunk)
+					_, _ = c.Writer.Write([]byte("\n\n"))
+					flusher.Flush()
+				}
+			// Handle errors from the backend.
+			case err, okError := <-errChan:
+				if okError {
+					if err.StatusCode == 429 && h.cfg.QuotaExceeded.SwitchProject {
+						continue outLoop
+					} else {
+						c.Status(err.StatusCode)
+						_, _ = fmt.Fprint(c.Writer, err.Error.Error())
+						flusher.Flush()
+						cliCancel()
+					}
+					return
+				}
+			// Send a keep-alive signal to the client.
+			case <-time.After(500 * time.Millisecond):
+				if hasFirstResponse {
+					_, _ = c.Writer.Write([]byte("\n"))
+					flusher.Flush()
+				}
+			}
+		}
+	}
+}
+
+func (h *APIHandlers) internalGenerateContent(c *gin.Context, rawJson []byte) {
+	c.Header("Content-Type", "application/json")
+
+	modelResult := gjson.GetBytes(rawJson, "model")
+	modelName := modelResult.String()
+	cliCtx, cliCancel := context.WithCancel(context.Background())
+	var cliClient *client.Client
+	defer func() {
+		if cliClient != nil {
+			cliClient.RequestMutex.Unlock()
+		}
+	}()
+
+	for {
+		var errorResponse *client.ErrorMessage
+		cliClient, errorResponse = h.getClient(modelName)
+		if errorResponse != nil {
+			c.Status(errorResponse.StatusCode)
+			_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
+			cliCancel()
+			return
+		}
+
+		if glAPIKey := cliClient.GetGenerativeLanguageAPIKey(); glAPIKey != "" {
+			log.Debugf("Request use generative language API Key: %s", glAPIKey)
+		} else {
+			log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
+		}
+
+		resp, err := cliClient.SendRawMessage(cliCtx, rawJson)
+		if err != nil {
+			if err.StatusCode == 429 && h.cfg.QuotaExceeded.SwitchProject {
+				continue
+			} else {
+				c.Status(err.StatusCode)
+				_, _ = c.Writer.Write([]byte(err.Error.Error()))
+				cliCancel()
+			}
+			break
+		} else {
+			_, _ = c.Writer.Write(resp)
+			cliCancel()
+			break
+		}
+	}
+}

internal/api/handlers.go

@@ -165,6 +165,48 @@ func (h *APIHandlers) Models(c *gin.Context) {
 	})
 }
 
+func (h *APIHandlers) getClient(modelName string) (*client.Client, *client.ErrorMessage) {
+	var cliClient *client.Client
+
+	// Lock the mutex to update the last used client index
+	mutex.Lock()
+	startIndex := lastUsedClientIndex
+	currentIndex := (startIndex + 1) % len(h.cliClients)
+	lastUsedClientIndex = currentIndex
+	mutex.Unlock()
+
+	// Reorder the client to start from the last used index
+	reorderedClients := make([]*client.Client, 0)
+	for i := 0; i < len(h.cliClients); i++ {
+		cliClient = h.cliClients[(startIndex+1+i)%len(h.cliClients)]
+		if cliClient.IsModelQuotaExceeded(modelName) {
+			log.Debugf("Model %s is quota exceeded for account %s, project id: %s", modelName, cliClient.GetEmail(), cliClient.GetProjectID())
+			cliClient = nil
+			continue
+		}
+		reorderedClients = append(reorderedClients, cliClient)
+	}
+
+	if len(reorderedClients) == 0 {
+		return nil, &client.ErrorMessage{StatusCode: 429, Error: fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName)}
+	}
+
+	locked := false
+	for i := 0; i < len(reorderedClients); i++ {
+		cliClient = reorderedClients[i]
+		if cliClient.RequestMutex.TryLock() {
+			locked = true
+			break
+		}
+	}
+	if !locked {
+		cliClient = h.cliClients[0]
+		cliClient.RequestMutex.Lock()
+	}
+
+	return cliClient, nil
+}
+
 // ChatCompletions handles the /v1/chat/completions endpoint.
 // It determines whether the request is for a streaming or non-streaming response
 // and calls the appropriate handler.
@@ -196,19 +238,7 @@ func (h *APIHandlers) ChatCompletions(c *gin.Context) {
 func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte) {
 	c.Header("Content-Type", "application/json")
 
-	// Handle streaming manually
+	modelName, systemInstruction, contents, tools := translator.PrepareRequest(rawJson)
-	flusher, ok := c.Writer.(http.Flusher)
-	if !ok {
-		c.JSON(http.StatusInternalServerError, ErrorResponse{
-			Error: ErrorDetail{
-				Message: "Streaming not supported",
-				Type:    "server_error",
-			},
-		})
-		return
-	}
-
-	modelName, contents, tools := translator.PrepareRequest(rawJson)
 	cliCtx, cliCancel := context.WithCancel(context.Background())
 	var cliClient *client.Client
 	defer func() {
@@ -218,46 +248,15 @@ func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte)
 	}()
 
 	for {
-		// Lock the mutex to update the last used client index
+		var errorResponse *client.ErrorMessage
-		mutex.Lock()
+		cliClient, errorResponse = h.getClient(modelName)
-		startIndex := lastUsedClientIndex
+		if errorResponse != nil {
-		currentIndex := (startIndex + 1) % len(h.cliClients)
+			c.Status(errorResponse.StatusCode)
-		lastUsedClientIndex = currentIndex
+			_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
-		mutex.Unlock()
-
-		// Reorder the client to start from the last used index
-		reorderedClients := make([]*client.Client, 0)
-		for i := 0; i < len(h.cliClients); i++ {
-			cliClient = h.cliClients[(startIndex+1+i)%len(h.cliClients)]
-			if cliClient.IsModelQuotaExceeded(modelName) {
-				log.Debugf("Model %s is quota exceeded for account %s, project id: %s", modelName, cliClient.GetEmail(), cliClient.GetProjectID())
-				cliClient = nil
-				continue
-			}
-			reorderedClients = append(reorderedClients, cliClient)
-		}
-
-		if len(reorderedClients) == 0 {
-			c.Status(429)
-			_, _ = fmt.Fprint(c.Writer, fmt.Sprintf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName))
-			flusher.Flush()
 			cliCancel()
 			return
 		}
 
-		locked := false
-		for i := 0; i < len(reorderedClients); i++ {
-			cliClient = reorderedClients[i]
-			if cliClient.RequestMutex.TryLock() {
-				locked = true
-				break
-			}
-		}
-		if !locked {
-			cliClient = h.cliClients[0]
-			cliClient.RequestMutex.Lock()
-		}
-
 		isGlAPIKey := false
 		if glAPIKey := cliClient.GetGenerativeLanguageAPIKey(); glAPIKey != "" {
 			log.Debugf("Request use generative language API Key: %s", glAPIKey)
@@ -266,22 +265,20 @@ func (h *APIHandlers) handleNonStreamingResponse(c *gin.Context, rawJson []byte)
 			log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
 		}
 
-		resp, err := cliClient.SendMessage(cliCtx, rawJson, modelName, contents, tools)
+		resp, err := cliClient.SendMessage(cliCtx, rawJson, modelName, systemInstruction, contents, tools)
 		if err != nil {
 			if err.StatusCode == 429 && h.cfg.QuotaExceeded.SwitchProject {
 				continue
 			} else {
 				c.Status(err.StatusCode)
-				_, _ = fmt.Fprint(c.Writer, err.Error.Error())
+				_, _ = c.Writer.Write([]byte(err.Error.Error()))
-				flusher.Flush()
 				cliCancel()
 			}
 			break
 		} else {
 			openAIFormat := translator.ConvertCliToOpenAINonStream(resp, time.Now().Unix(), isGlAPIKey)
 			if openAIFormat != "" {
-				_, _ = fmt.Fprintf(c.Writer, "data: %s\n\n", openAIFormat)
+				_, _ = c.Writer.Write([]byte(openAIFormat))
-				flusher.Flush()
 			}
 			cliCancel()
 			break
@@ -309,7 +306,7 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) {
 	}
 
 	// Prepare the request for the backend client.
-	modelName, contents, tools := translator.PrepareRequest(rawJson)
+	modelName, systemInstruction, contents, tools := translator.PrepareRequest(rawJson)
 	cliCtx, cliCancel := context.WithCancel(context.Background())
 	var cliClient *client.Client
 	defer func() {
@@ -321,46 +318,16 @@ func (h *APIHandlers) handleStreamingResponse(c *gin.Context, rawJson []byte) {
 
 outLoop:
 	for {
-		// Lock the mutex to update the last used client index
+		var errorResponse *client.ErrorMessage
-		mutex.Lock()
+		cliClient, errorResponse = h.getClient(modelName)
-		startIndex := lastUsedClientIndex
+		if errorResponse != nil {
-		currentIndex := (startIndex + 1) % len(h.cliClients)
+			c.Status(errorResponse.StatusCode)
-		lastUsedClientIndex = currentIndex
+			_, _ = fmt.Fprint(c.Writer, errorResponse.Error)
-		mutex.Unlock()
-
-		// Reorder the client to start from the last used index
-		reorderedClients := make([]*client.Client, 0)
-		for i := 0; i < len(h.cliClients); i++ {
-			cliClient = h.cliClients[(startIndex+1+i)%len(h.cliClients)]
-			if cliClient.IsModelQuotaExceeded(modelName) {
-				log.Debugf("Model %s is quota exceeded for account %s, project id: %s", modelName, cliClient.GetEmail(), cliClient.GetProjectID())
-				cliClient = nil
-				continue
-			}
-			reorderedClients = append(reorderedClients, cliClient)
-		}
-
-		if len(reorderedClients) == 0 {
-			c.Status(429)
-			_, _ = fmt.Fprint(c.Writer, fmt.Sprintf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, modelName))
 			flusher.Flush()
 			cliCancel()
 			return
 		}
 
-		locked := false
-		for i := 0; i < len(reorderedClients); i++ {
-			cliClient = reorderedClients[i]
-			if cliClient.RequestMutex.TryLock() {
-				locked = true
-				break
-			}
-		}
-		if !locked {
-			cliClient = h.cliClients[0]
-			cliClient.RequestMutex.Lock()
-		}
-
 		isGlAPIKey := false
 		if glAPIKey := cliClient.GetGenerativeLanguageAPIKey(); glAPIKey != "" {
 			log.Debugf("Request use generative language API Key: %s", glAPIKey)
@@ -369,7 +336,7 @@ outLoop:
 			log.Debugf("Request use account: %s, project id: %s", cliClient.GetEmail(), cliClient.GetProjectID())
 		}
 		// Send the message and receive response chunks and errors via channels.
-		respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, contents, tools)
+		respChan, errChan := cliClient.SendMessageStream(cliCtx, rawJson, modelName, systemInstruction, contents, tools)
 		hasFirstResponse := false
 		for {
 			select {

internal/api/server.go

@@ -81,6 +81,8 @@ func (s *Server) setupRoutes() {
 			},
 		})
 	})
+	s.engine.POST("/v1internal:method", s.handlers.CLIHandler)
+
 }
 
 // Start begins listening for and serving HTTP requests.

internal/api/translator/request.go

@@ -12,7 +12,7 @@ import (
 // PrepareRequest translates a raw JSON request from an OpenAI-compatible format
 // to the internal format expected by the backend client. It parses messages,
 // roles, content types (text, image, file), and tool calls.
-func PrepareRequest(rawJson []byte) (string, []client.Content, []client.ToolDeclaration) {
+func PrepareRequest(rawJson []byte) (string, *client.Content, []client.Content, []client.ToolDeclaration) {
 	// Extract the model name from the request, defaulting to "gemini-2.5-pro".
 	modelName := "gemini-2.5-pro"
 	modelResult := gjson.GetBytes(rawJson, "model")
@@ -22,7 +22,41 @@ func PrepareRequest(rawJson []byte) (string, []client.Content, []client.ToolDecl
 
 	// Process the array of messages.
 	contents := make([]client.Content, 0)
+	var systemInstruction *client.Content
 	messagesResult := gjson.GetBytes(rawJson, "messages")
+
+	toolItems := make(map[string]*client.FunctionResponse)
+	if messagesResult.IsArray() {
+		messagesResults := messagesResult.Array()
+		for i := 0; i < len(messagesResults); i++ {
+			messageResult := messagesResults[i]
+			roleResult := messageResult.Get("role")
+			if roleResult.Type != gjson.String {
+				continue
+			}
+			contentResult := messageResult.Get("content")
+			if roleResult.String() == "tool" {
+				toolCallID := messageResult.Get("tool_call_id").String()
+				if toolCallID != "" {
+					var responseData string
+					if contentResult.Type == gjson.String {
+						responseData = contentResult.String()
+					} else if contentResult.IsObject() && contentResult.Get("type").String() == "text" {
+						responseData = contentResult.Get("text").String()
+					}
+
+					// drop the timestamp from the tool call ID
+					toolCallIDs := strings.Split(toolCallID, "-")
+					strings.Join(toolCallIDs, "-")
+					newToolCallID := strings.Join(toolCallIDs[:len(toolCallIDs)-1], "-")
+
+					functionResponse := client.FunctionResponse{Name: newToolCallID, Response: map[string]interface{}{"result": responseData}}
+					toolItems[toolCallID] = &functionResponse
+				}
+			}
+		}
+	}
+
 	if messagesResult.IsArray() {
 		messagesResults := messagesResult.Array()
 		for i := 0; i < len(messagesResults); i++ {
@@ -37,13 +71,11 @@ func PrepareRequest(rawJson []byte) (string, []client.Content, []client.ToolDecl
 			// System messages are converted to a user message followed by a model's acknowledgment.
 			case "system":
 				if contentResult.Type == gjson.String {
-					contents = append(contents, client.Content{Role: "user", Parts: []client.Part{{Text: contentResult.String()}}})
+					systemInstruction = &client.Content{Role: "user", Parts: []client.Part{{Text: contentResult.String()}}}
-					contents = append(contents, client.Content{Role: "model", Parts: []client.Part{{Text: "Understood. I will follow these instructions and use my tools to assist you."}}})
 				} else if contentResult.IsObject() {
 					// Handle object-based system messages.
 					if contentResult.Get("type").String() == "text" {
-						contents = append(contents, client.Content{Role: "user", Parts: []client.Part{{Text: contentResult.Get("text").String()}}})
+						systemInstruction = &client.Content{Role: "user", Parts: []client.Part{{Text: contentResult.Get("text").String()}}}
-						contents = append(contents, client.Content{Role: "model", Parts: []client.Part{{Text: "Understood. I will follow these instructions and use my tools to assist you."}}})
 					}
 				}
 			// User messages can contain simple text or a multi-part body.
@@ -98,40 +130,44 @@ func PrepareRequest(rawJson []byte) (string, []client.Content, []client.ToolDecl
 					contents = append(contents, client.Content{Role: "model", Parts: []client.Part{{Text: contentResult.String()}}})
 				} else if !contentResult.Exists() || contentResult.Type == gjson.Null {
 					// Handle tool calls made by the assistant.
+					functionIDs := make([]string, 0)
 					toolCallsResult := messageResult.Get("tool_calls")
 					if toolCallsResult.IsArray() {
+						parts := make([]client.Part, 0)
 						tcsResult := toolCallsResult.Array()
 						for j := 0; j < len(tcsResult); j++ {
 							tcResult := tcsResult[j]
+
+							functionID := tcResult.Get("id").String()
+							functionIDs = append(functionIDs, functionID)
+
 							functionName := tcResult.Get("function.name").String()
 							functionArgs := tcResult.Get("function.arguments").String()
 							var args map[string]any
 							if err := json.Unmarshal([]byte(functionArgs), &args); err == nil {
-								contents = append(contents, client.Content{
+								parts = append(parts, client.Part{
-									Role: "model", Parts: []client.Part{{
+									FunctionCall: &client.FunctionCall{
-										FunctionCall: &client.FunctionCall{
+										Name: functionName,
-											Name: functionName,
+										Args: args,
-											Args: args,
+									},
-										},
-									}},
 								})
 							}
 						}
+						if len(parts) > 0 {
+							contents = append(contents, client.Content{
+								Role: "model", Parts: parts,
+							})
+
+							toolParts := make([]client.Part, 0)
+							for _, functionID := range functionIDs {
+								if functionResponse, ok := toolItems[functionID]; ok {
+									toolParts = append(toolParts, client.Part{FunctionResponse: functionResponse})
+								}
+							}
+							contents = append(contents, client.Content{Role: "tool", Parts: toolParts})
+						}
 					}
 				}
-			// Tool messages contain the output of a tool call.
-			case "tool":
-				toolCallID := messageResult.Get("tool_call_id").String()
-				if toolCallID != "" {
-					var responseData string
-					if contentResult.Type == gjson.String {
-						responseData = contentResult.String()
-					} else if contentResult.IsObject() && contentResult.Get("type").String() == "text" {
-						responseData = contentResult.Get("text").String()
-					}
-					functionResponse := client.FunctionResponse{Name: toolCallID, Response: map[string]interface{}{"result": responseData}}
-					contents = append(contents, client.Content{Role: "tool", Parts: []client.Part{{FunctionResponse: &functionResponse}}})
-				}
 			}
 		}
 	}
@@ -159,5 +195,5 @@ func PrepareRequest(rawJson []byte) (string, []client.Content, []client.ToolDecl
 		tools = make([]client.ToolDeclaration, 0)
 	}
 
-	return modelName, contents, tools
+	return modelName, systemInstruction, contents, tools
 }

internal/api/translator/response.go

@@ -1,6 +1,7 @@
 package translator
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/tidwall/gjson"
@@ -62,32 +63,40 @@ func ConvertCliToOpenAI(rawJson []byte, unixTimestamp int64, isGlAPIKey bool) st
 	}
 
 	// Process the main content part of the response.
-	partResult := gjson.GetBytes(rawJson, "response.candidates.0.content.parts.0")
+	partsResult := gjson.GetBytes(rawJson, "response.candidates.0.content.parts")
-	partTextResult := partResult.Get("text")
+	if partsResult.IsArray() {
-	functionCallResult := partResult.Get("functionCall")
+		partResults := partsResult.Array()
+		for i := 0; i < len(partResults); i++ {
+			partResult := partResults[i]
+			partTextResult := partResult.Get("text")
+			functionCallResult := partResult.Get("functionCall")
 
-	if partTextResult.Exists() {
+			if partTextResult.Exists() {
-		// Handle text content, distinguishing between regular content and reasoning/thoughts.
+				// Handle text content, distinguishing between regular content and reasoning/thoughts.
-		if partResult.Get("thought").Bool() {
+				if partResult.Get("thought").Bool() {
-			template, _ = sjson.Set(template, "choices.0.delta.reasoning_content", partTextResult.String())
+					template, _ = sjson.Set(template, "choices.0.delta.reasoning_content", partTextResult.String())
-		} else {
+				} else {
-			template, _ = sjson.Set(template, "choices.0.delta.content", partTextResult.String())
+					template, _ = sjson.Set(template, "choices.0.delta.content", partTextResult.String())
+				}
+				template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
+			} else if functionCallResult.Exists() {
+				// Handle function call content.
+				toolCallsResult := gjson.Get(template, "choices.0.delta.tool_calls")
+				if !toolCallsResult.Exists() || !toolCallsResult.IsArray() {
+					template, _ = sjson.SetRaw(template, "choices.0.delta.tool_calls", `[]`)
+				}
+
+				functionCallTemplate := `{"id": "","type": "function","function": {"name": "","arguments": ""}}`
+				fcName := functionCallResult.Get("name").String()
+				functionCallTemplate, _ = sjson.Set(functionCallTemplate, "id", fmt.Sprintf("%s-%d", fcName, time.Now().UnixNano()))
+				functionCallTemplate, _ = sjson.Set(functionCallTemplate, "function.name", fcName)
+				if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
+					functionCallTemplate, _ = sjson.Set(functionCallTemplate, "function.arguments", fcArgsResult.Raw)
+				}
+				template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
+				template, _ = sjson.SetRaw(template, "choices.0.message.tool_calls.-1", functionCallTemplate)
+			}
 		}
-		template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
-	} else if functionCallResult.Exists() {
-		// Handle function call content.
-		functionCallTemplate := `[{"id": "","type": "function","function": {"name": "","arguments": ""}}]`
-		fcName := functionCallResult.Get("name").String()
-		functionCallTemplate, _ = sjson.Set(functionCallTemplate, "0.id", fcName)
-		functionCallTemplate, _ = sjson.Set(functionCallTemplate, "0.function.name", fcName)
-		if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
-			functionCallTemplate, _ = sjson.Set(functionCallTemplate, "0.function.arguments", fcArgsResult.Raw)
-		}
-		template, _ = sjson.Set(template, "choices.0.delta.role", "assistant")
-		template, _ = sjson.SetRaw(template, "choices.0.delta.tool_calls", functionCallTemplate)
-	} else {
-		// If no usable content is found, return an empty string.
-		return ""
 	}
 
 	return template
@@ -163,7 +172,7 @@ func ConvertCliToOpenAINonStream(rawJson []byte, unixTimestamp int64, isGlAPIKey
 				}
 				functionCallItemTemplate := `{"id": "","type": "function","function": {"name": "","arguments": ""}}`
 				fcName := functionCallResult.Get("name").String()
-				functionCallItemTemplate, _ = sjson.Set(functionCallItemTemplate, "id", fcName)
+				functionCallItemTemplate, _ = sjson.Set(functionCallItemTemplate, "id", fmt.Sprintf("%s-%d", fcName, time.Now().UnixNano()))
 				functionCallItemTemplate, _ = sjson.Set(functionCallItemTemplate, "function.name", fcName)
 				if fcArgsResult := functionCallResult.Get("args"); fcArgsResult.Exists() {
 					functionCallItemTemplate, _ = sjson.Set(functionCallItemTemplate, "function.arguments", fcArgsResult.Raw)

internal/client/client.go

@@ -212,6 +212,7 @@ func (c *Client) makeAPIRequest(ctx context.Context, endpoint, method string, bo
 	metadataStr := getClientMetadataString()
 	req.Header.Set("Content-Type", "application/json")
 	req.Header.Set("User-Agent", 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))
 
@@ -266,6 +267,12 @@ func (c *Client) APIRequest(ctx context.Context, endpoint string, body interface
 			url = url + "?alt=sse"
 		}
 		jsonBody = []byte(gjson.GetBytes(jsonBody, "request").Raw)
+		systemInstructionResult := gjson.GetBytes(jsonBody, "systemInstruction")
+		if systemInstructionResult.Exists() {
+			jsonBody, _ = sjson.SetRawBytes(jsonBody, "system_instruction", []byte(systemInstructionResult.Raw))
+			jsonBody, _ = sjson.DeleteBytes(jsonBody, "systemInstruction")
+			jsonBody, _ = sjson.DeleteBytes(jsonBody, "session_id")
+		}
 	}
 
 	// log.Debug(string(jsonBody))
@@ -285,6 +292,7 @@ func (c *Client) APIRequest(ctx context.Context, endpoint string, body interface
 			return nil, &ErrorMessage{500, fmt.Errorf("failed to get token: %v", errToken)}
 		}
 		req.Header.Set("User-Agent", 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))
 	} else {
@@ -303,15 +311,14 @@ func (c *Client) APIRequest(ctx context.Context, endpoint string, body interface
 			}
 		}()
 		bodyBytes, _ := io.ReadAll(resp.Body)
-
 		return nil, &ErrorMessage{resp.StatusCode, fmt.Errorf(string(bodyBytes))}
 	}
 
 	return resp.Body, nil
 }
 
-// SendMessageStream handles a single conversational turn, including tool calls.
+// SendMessage handles a single conversational turn, including tool calls.
-func (c *Client) SendMessage(ctx context.Context, rawJson []byte, model string, contents []Content, tools []ToolDeclaration) ([]byte, *ErrorMessage) {
+func (c *Client) SendMessage(ctx context.Context, rawJson []byte, model string, systemInstruction *Content, contents []Content, tools []ToolDeclaration) ([]byte, *ErrorMessage) {
 	request := GenerateContentRequest{
 		Contents: contents,
 		GenerationConfig: GenerationConfig{
@@ -320,6 +327,9 @@ func (c *Client) SendMessage(ctx context.Context, rawJson []byte, model string,
 			},
 		},
 	}
+
+	request.SystemInstruction = systemInstruction
+
 	request.Tools = tools
 
 	requestBody := map[string]interface{}{
@@ -402,7 +412,7 @@ func (c *Client) SendMessage(ctx context.Context, rawJson []byte, model string,
 }
 
 // SendMessageStream handles a single conversational turn, including tool calls.
-func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model string, contents []Content, tools []ToolDeclaration) (<-chan []byte, <-chan *ErrorMessage) {
+func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model string, systemInstruction *Content, contents []Content, tools []ToolDeclaration) (<-chan []byte, <-chan *ErrorMessage) {
 	dataTag := []byte("data: ")
 	errChan := make(chan *ErrorMessage)
 	dataChan := make(chan []byte)
@@ -418,6 +428,9 @@ func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model st
 				},
 			},
 		}
+
+		request.SystemInstruction = systemInstruction
+
 		request.Tools = tools
 
 		requestBody := map[string]interface{}{
@@ -519,6 +532,117 @@ func (c *Client) SendMessageStream(ctx context.Context, rawJson []byte, model st
 	return dataChan, errChan
 }
 
+// SendRawMessage handles a single conversational turn, including tool calls.
+func (c *Client) SendRawMessage(ctx context.Context, rawJson []byte) ([]byte, *ErrorMessage) {
+	rawJson, _ = sjson.SetBytes(rawJson, "project", c.GetProjectID())
+
+	modelResult := gjson.GetBytes(rawJson, "model")
+	model := modelResult.String()
+	modelName := model
+	for {
+		if c.isModelQuotaExceeded(modelName) {
+			if c.cfg.QuotaExceeded.SwitchPreviewModel && c.glAPIKey == "" {
+				modelName = c.getPreviewModel(model)
+				if modelName != "" {
+					log.Debugf("Model %s is quota exceeded. Switch to preview model %s", model, modelName)
+					rawJson, _ = sjson.SetBytes(rawJson, "model", modelName)
+					continue
+				}
+			}
+			return nil, &ErrorMessage{
+				StatusCode: 429,
+				Error:      fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, model),
+			}
+		}
+
+		respBody, err := c.APIRequest(ctx, "generateContent", rawJson, false)
+		if err != nil {
+			if err.StatusCode == 429 {
+				now := time.Now()
+				c.modelQuotaExceeded[modelName] = &now
+				if c.cfg.QuotaExceeded.SwitchPreviewModel && c.glAPIKey == "" {
+					continue
+				}
+			}
+			return nil, err
+		}
+		delete(c.modelQuotaExceeded, modelName)
+		bodyBytes, errReadAll := io.ReadAll(respBody)
+		if errReadAll != nil {
+			return nil, &ErrorMessage{StatusCode: 500, Error: errReadAll}
+		}
+		return bodyBytes, nil
+	}
+}
+
+// SendRawMessageStream handles a single conversational turn, including tool calls.
+func (c *Client) SendRawMessageStream(ctx context.Context, rawJson []byte) (<-chan []byte, <-chan *ErrorMessage) {
+	dataTag := []byte("data: ")
+	errChan := make(chan *ErrorMessage)
+	dataChan := make(chan []byte)
+	go func() {
+		defer close(errChan)
+		defer close(dataChan)
+
+		rawJson, _ = sjson.SetBytes(rawJson, "project", c.GetProjectID())
+
+		modelResult := gjson.GetBytes(rawJson, "model")
+		model := modelResult.String()
+		modelName := model
+		var stream io.ReadCloser
+		for {
+			if c.isModelQuotaExceeded(modelName) {
+				if c.cfg.QuotaExceeded.SwitchPreviewModel && c.glAPIKey == "" {
+					modelName = c.getPreviewModel(model)
+					if modelName != "" {
+						log.Debugf("Model %s is quota exceeded. Switch to preview model %s", model, modelName)
+						rawJson, _ = sjson.SetBytes(rawJson, "model", modelName)
+						continue
+					}
+				}
+				errChan <- &ErrorMessage{
+					StatusCode: 429,
+					Error:      fmt.Errorf(`{"error":{"code":429,"message":"All the models of '%s' are quota exceeded","status":"RESOURCE_EXHAUSTED"}}`, model),
+				}
+				return
+			}
+			var err *ErrorMessage
+			stream, err = c.APIRequest(ctx, "streamGenerateContent", rawJson, true)
+			if err != nil {
+				if err.StatusCode == 429 {
+					now := time.Now()
+					c.modelQuotaExceeded[modelName] = &now
+					if c.cfg.QuotaExceeded.SwitchPreviewModel && c.glAPIKey == "" {
+						continue
+					}
+				}
+				errChan <- err
+				return
+			}
+			delete(c.modelQuotaExceeded, modelName)
+			break
+		}
+
+		scanner := bufio.NewScanner(stream)
+		for scanner.Scan() {
+			line := scanner.Bytes()
+			if bytes.HasPrefix(line, dataTag) {
+				dataChan <- line[6:]
+			}
+		}
+
+		if errScanner := scanner.Err(); errScanner != nil {
+			errChan <- &ErrorMessage{500, errScanner}
+			_ = stream.Close()
+			return
+		}
+
+		_ = stream.Close()
+	}()
+
+	return dataChan, errChan
+}
+
 func (c *Client) isModelQuotaExceeded(model string) bool {
 	if lastExceededTime, hasKey := c.modelQuotaExceeded[model]; hasKey {
 		duration := time.Now().Sub(*lastExceededTime)
@@ -659,10 +783,10 @@ func (c *Client) SaveTokenToFile() error {
 // such as IDE type, platform, and plugin version.
 func getClientMetadata() map[string]string {
 	return map[string]string{
-		"ideType":       "IDE_UNSPECIFIED",
+		"ideType":    "IDE_UNSPECIFIED",
-		"platform":      getPlatform(),
+		"platform":   "PLATFORM_UNSPECIFIED",
-		"pluginType":    "GEMINI",
+		"pluginType": "GEMINI",
-		"pluginVersion": pluginVersion,
+		// "pluginVersion": pluginVersion,
 	}
 }
 
@@ -679,7 +803,8 @@ func getClientMetadataString() string {
 
 // getUserAgent constructs the User-Agent string for HTTP requests.
 func getUserAgent() string {
-	return fmt.Sprintf("GeminiCLI/%s (%s; %s)", pluginVersion, runtime.GOOS, runtime.GOARCH)
+	// return fmt.Sprintf("GeminiCLI/%s (%s; %s)", pluginVersion, runtime.GOOS, runtime.GOARCH)
+	return "google-api-nodejs-client/9.15.1"
 }
 
 // getPlatform determines the operating system and architecture and formats

internal/client/models.go

@@ -64,9 +64,10 @@ type FunctionResponse struct {
 
 // GenerateContentRequest is the top-level request structure for the streamGenerateContent endpoint.
 type GenerateContentRequest struct {
-	Contents         []Content         `json:"contents"`
+	SystemInstruction *Content          `json:"systemInstruction,omitempty"`
-	Tools            []ToolDeclaration `json:"tools,omitempty"`
+	Contents          []Content         `json:"contents"`
-	GenerationConfig `json:"generationConfig"`
+	Tools             []ToolDeclaration `json:"tools,omitempty"`
+	GenerationConfig  `json:"generationConfig"`
 }
 
 // GenerationConfig defines parameters that control the model's generation behavior.

internal/cmd/run.go

@@ -7,12 +7,10 @@ import (
 	"github.com/luispater/CLIProxyAPI/internal/auth"
 	"github.com/luispater/CLIProxyAPI/internal/client"
 	"github.com/luispater/CLIProxyAPI/internal/config"
+	"github.com/luispater/CLIProxyAPI/internal/util"
 	log "github.com/sirupsen/logrus"
-	"golang.org/x/net/proxy"
 	"io/fs"
-	"net"
 	"net/http"
-	"net/url"
 	"os"
 	"os/signal"
 	"path/filepath"
@@ -69,33 +67,12 @@ func StartService(cfg *config.Config) {
 	}
 
 	if len(cfg.GlAPIKey) > 0 {
-		var transport *http.Transport
-		proxyURL, errParse := url.Parse(cfg.ProxyUrl)
-		if errParse == nil {
-			if proxyURL.Scheme == "socks5" {
-				username := proxyURL.User.Username()
-				password, _ := proxyURL.User.Password()
-				proxyAuth := &proxy.Auth{User: username, Password: password}
-				dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, proxyAuth, proxy.Direct)
-				if errSOCKS5 != nil {
-					log.Fatalf("create SOCKS5 dialer failed: %v", errSOCKS5)
-				}
-				transport = &http.Transport{
-					DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
-						return dialer.Dial(network, addr)
-					},
-				}
-			} else if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
-				// Handle HTTP/HTTPS proxy.
-				transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
-			}
-		}
-
 		for i := 0; i < len(cfg.GlAPIKey); i++ {
-			httpClient := &http.Client{}
+			httpClient, errSetProxy := util.SetProxy(cfg, &http.Client{})
-			if transport != nil {
+			if errSetProxy != nil {
-				httpClient.Transport = transport
+				log.Fatalf("set proxy failed: %v", errSetProxy)
 			}
+
 			log.Debug("Initializing with Generative Language API key...")
 			cliClient := client.NewClient(httpClient, nil, cfg, cfg.GlAPIKey[i])
 			cliClients = append(cliClients, cliClient)

internal/util/proxy.go

@@ -0,0 +1,37 @@
+package util
+
+import (
+	"context"
+	"github.com/luispater/CLIProxyAPI/internal/config"
+	"golang.org/x/net/proxy"
+	"net"
+	"net/http"
+	"net/url"
+)
+
+func SetProxy(cfg *config.Config, httpClient *http.Client) (*http.Client, error) {
+	var transport *http.Transport
+	proxyURL, errParse := url.Parse(cfg.ProxyUrl)
+	if errParse == nil {
+		if proxyURL.Scheme == "socks5" {
+			username := proxyURL.User.Username()
+			password, _ := proxyURL.User.Password()
+			proxyAuth := &proxy.Auth{User: username, Password: password}
+			dialer, errSOCKS5 := proxy.SOCKS5("tcp", proxyURL.Host, proxyAuth, proxy.Direct)
+			if errSOCKS5 != nil {
+				return nil, errSOCKS5
+			}
+			transport = &http.Transport{
+				DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
+					return dialer.Dial(network, addr)
+				},
+			}
+		} else if proxyURL.Scheme == "http" || proxyURL.Scheme == "https" {
+			transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
+		}
+	}
+	if transport != nil {
+		httpClient.Transport = transport
+	}
+	return httpClient, nil
+}