Add projects
This commit is contained in:
parent
2d3a9ad623
commit
8b607dd700
1802 changed files with 503346 additions and 2 deletions
350
backend/sdk/api/handlers/gemini/gemini_handlers.go
Normal file
350
backend/sdk/api/handlers/gemini/gemini_handlers.go
Normal file
|
|
@ -0,0 +1,350 @@
|
|||
// Package gemini provides HTTP handlers for Gemini API endpoints.
|
||||
// This package implements handlers for managing Gemini model operations including
|
||||
// model listing, content generation, streaming content generation, and token counting.
|
||||
// It serves as a proxy layer between clients and the Gemini backend service,
|
||||
// handling request translation, client management, and response processing.
|
||||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
)
|
||||
|
||||
// GeminiAPIHandler contains the handlers for Gemini API endpoints.
|
||||
// It holds a pool of clients to interact with the backend service.
|
||||
type GeminiAPIHandler struct {
|
||||
*handlers.BaseAPIHandler
|
||||
}
|
||||
|
||||
// NewGeminiAPIHandler creates a new Gemini API handlers instance.
|
||||
// It takes an BaseAPIHandler instance as input and returns a GeminiAPIHandler.
|
||||
func NewGeminiAPIHandler(apiHandlers *handlers.BaseAPIHandler) *GeminiAPIHandler {
|
||||
return &GeminiAPIHandler{
|
||||
BaseAPIHandler: apiHandlers,
|
||||
}
|
||||
}
|
||||
|
||||
// HandlerType returns the identifier for this handler implementation.
|
||||
func (h *GeminiAPIHandler) HandlerType() string {
|
||||
return Gemini
|
||||
}
|
||||
|
||||
// Models returns the Gemini-compatible model metadata supported by this handler.
|
||||
func (h *GeminiAPIHandler) Models() []map[string]any {
|
||||
// Get dynamic models from the global registry
|
||||
modelRegistry := registry.GetGlobalRegistry()
|
||||
return modelRegistry.GetAvailableModels("gemini")
|
||||
}
|
||||
|
||||
// GeminiModels handles the Gemini models listing endpoint.
|
||||
// It returns a JSON response containing available Gemini models and their specifications.
|
||||
func (h *GeminiAPIHandler) GeminiModels(c *gin.Context) {
|
||||
rawModels := h.Models()
|
||||
normalizedModels := make([]map[string]any, 0, len(rawModels))
|
||||
defaultMethods := []string{"generateContent"}
|
||||
for _, model := range rawModels {
|
||||
normalizedModel := make(map[string]any, len(model))
|
||||
for k, v := range model {
|
||||
normalizedModel[k] = v
|
||||
}
|
||||
if name, ok := normalizedModel["name"].(string); ok && name != "" {
|
||||
if !strings.HasPrefix(name, "models/") {
|
||||
normalizedModel["name"] = "models/" + name
|
||||
}
|
||||
if displayName, _ := normalizedModel["displayName"].(string); displayName == "" {
|
||||
normalizedModel["displayName"] = name
|
||||
}
|
||||
if description, _ := normalizedModel["description"].(string); description == "" {
|
||||
normalizedModel["description"] = name
|
||||
}
|
||||
}
|
||||
if _, ok := normalizedModel["supportedGenerationMethods"]; !ok {
|
||||
normalizedModel["supportedGenerationMethods"] = defaultMethods
|
||||
}
|
||||
normalizedModels = append(normalizedModels, normalizedModel)
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"models": normalizedModels,
|
||||
})
|
||||
}
|
||||
|
||||
// GeminiGetHandler handles GET requests for specific Gemini model information.
|
||||
// It returns detailed information about a specific Gemini model based on the action parameter.
|
||||
func (h *GeminiAPIHandler) GeminiGetHandler(c *gin.Context) {
|
||||
var request struct {
|
||||
Action string `uri:"action" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindUri(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
action := strings.TrimPrefix(request.Action, "/")
|
||||
|
||||
// Get dynamic models from the global registry and find the matching one
|
||||
availableModels := h.Models()
|
||||
var targetModel map[string]any
|
||||
|
||||
for _, model := range availableModels {
|
||||
name, _ := model["name"].(string)
|
||||
// Match name with or without 'models/' prefix
|
||||
if name == action || name == "models/"+action {
|
||||
targetModel = model
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if targetModel != nil {
|
||||
// Ensure the name has 'models/' prefix in the output if it's a Gemini model
|
||||
if name, ok := targetModel["name"].(string); ok && name != "" && !strings.HasPrefix(name, "models/") {
|
||||
targetModel["name"] = "models/" + name
|
||||
}
|
||||
c.JSON(http.StatusOK, targetModel)
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusNotFound, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Not Found",
|
||||
Type: "not_found",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// GeminiHandler handles POST requests for Gemini API operations.
|
||||
// It routes requests to appropriate handlers based on the action parameter (model:method format).
|
||||
func (h *GeminiAPIHandler) GeminiHandler(c *gin.Context) {
|
||||
var request struct {
|
||||
Action string `uri:"action" binding:"required"`
|
||||
}
|
||||
if err := c.ShouldBindUri(&request); err != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("Invalid request: %v", err),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
action := strings.Split(strings.TrimPrefix(request.Action, "/"), ":")
|
||||
if len(action) != 2 {
|
||||
c.JSON(http.StatusNotFound, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: fmt.Sprintf("%s not found.", c.Request.URL.Path),
|
||||
Type: "invalid_request_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
method := action[1]
|
||||
rawJSON, _ := c.GetRawData()
|
||||
|
||||
switch method {
|
||||
case "generateContent":
|
||||
h.handleGenerateContent(c, action[0], rawJSON)
|
||||
case "streamGenerateContent":
|
||||
h.handleStreamGenerateContent(c, action[0], rawJSON)
|
||||
case "countTokens":
|
||||
h.handleCountTokens(c, action[0], rawJSON)
|
||||
}
|
||||
}
|
||||
|
||||
// handleStreamGenerateContent handles streaming content generation requests for Gemini models.
|
||||
// This function establishes a Server-Sent Events connection and streams the generated content
|
||||
// back to the client in real-time. It supports both SSE format and direct streaming based
|
||||
// on the 'alt' query parameter.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for content generation
|
||||
// - rawJSON: The raw JSON request body containing generation parameters
|
||||
func (h *GeminiAPIHandler) handleStreamGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
alt := h.GetAlt(c)
|
||||
|
||||
// Get the http.Flusher interface to manually flush the response.
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{
|
||||
Error: handlers.ErrorDetail{
|
||||
Message: "Streaming not supported",
|
||||
Type: "server_error",
|
||||
},
|
||||
})
|
||||
return
|
||||
}
|
||||
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
dataChan, upstreamHeaders, errChan := h.ExecuteStreamWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
|
||||
setSSEHeaders := func() {
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
}
|
||||
|
||||
// Peek at the first chunk
|
||||
for {
|
||||
select {
|
||||
case <-c.Request.Context().Done():
|
||||
cliCancel(c.Request.Context().Err())
|
||||
return
|
||||
case errMsg, ok := <-errChan:
|
||||
if !ok {
|
||||
// Err channel closed cleanly; wait for data channel.
|
||||
errChan = nil
|
||||
continue
|
||||
}
|
||||
// Upstream failed immediately. Return proper error status and JSON.
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
case chunk, ok := <-dataChan:
|
||||
if !ok {
|
||||
if errMsg, hasPendingError := handlers.PendingStreamError(errChan); hasPendingError {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
if errMsg != nil {
|
||||
cliCancel(errMsg.Error)
|
||||
} else {
|
||||
cliCancel(nil)
|
||||
}
|
||||
return
|
||||
}
|
||||
// Closed without data
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
flusher.Flush()
|
||||
cliCancel(nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Success! Set headers.
|
||||
if alt == "" {
|
||||
setSSEHeaders()
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
|
||||
// Write first chunk
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
flusher.Flush()
|
||||
|
||||
// Continue
|
||||
h.forwardGeminiStream(c, flusher, alt, func(err error) { cliCancel(err) }, dataChan, errChan)
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleCountTokens handles token counting requests for Gemini models.
|
||||
// This function counts the number of tokens in the provided content without
|
||||
// generating a response. It's useful for quota management and content validation.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for token counting
|
||||
// - rawJSON: The raw JSON request body containing the content to count
|
||||
func (h *GeminiAPIHandler) handleCountTokens(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteCountWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
// handleGenerateContent handles non-streaming content generation requests for Gemini models.
|
||||
// This function processes the request synchronously and returns the complete generated
|
||||
// response in a single API call. It supports various generation parameters and
|
||||
// response formats.
|
||||
//
|
||||
// Parameters:
|
||||
// - c: The Gin context for the request
|
||||
// - modelName: The name of the Gemini model to use for content generation
|
||||
// - rawJSON: The raw JSON request body containing generation parameters and content
|
||||
func (h *GeminiAPIHandler) handleGenerateContent(c *gin.Context, modelName string, rawJSON []byte) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, upstreamHeaders, errMsg := h.ExecuteWithAuthManager(cliCtx, h.HandlerType(), modelName, rawJSON, alt)
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), upstreamHeaders)
|
||||
_, _ = c.Writer.Write(resp)
|
||||
cliCancel()
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) forwardGeminiStream(c *gin.Context, flusher http.Flusher, alt string, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
var keepAliveInterval *time.Duration
|
||||
if alt != "" {
|
||||
keepAliveInterval = new(time.Duration(0))
|
||||
}
|
||||
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
KeepAliveInterval: keepAliveInterval,
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if alt == "" {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
if alt == "" {
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
} else {
|
||||
_, _ = c.Writer.Write(body)
|
||||
}
|
||||
},
|
||||
})
|
||||
}
|
||||
|
|
@ -0,0 +1,93 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
coreexecutor "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/executor"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
)
|
||||
|
||||
const (
|
||||
initialFailureGeminiModel = "initial-failure-gemini-model"
|
||||
)
|
||||
|
||||
type initialFailureGeminiStreamExecutor struct{}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Identifier() string {
|
||||
return "initial-failure-gemini-stream-executor"
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Execute(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) ExecuteStream(_ context.Context, _ *coreauth.Auth, _ coreexecutor.Request, _ coreexecutor.Options) (*coreexecutor.StreamResult, error) {
|
||||
chunks := make(chan coreexecutor.StreamChunk, 1)
|
||||
chunks <- coreexecutor.StreamChunk{Err: errors.New("upstream failed before first payload")}
|
||||
close(chunks)
|
||||
return &coreexecutor.StreamResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) Refresh(_ context.Context, auth *coreauth.Auth) (*coreauth.Auth, error) {
|
||||
return auth, nil
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) CountTokens(context.Context, *coreauth.Auth, coreexecutor.Request, coreexecutor.Options) (coreexecutor.Response, error) {
|
||||
return coreexecutor.Response{}, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (*initialFailureGeminiStreamExecutor) HttpRequest(context.Context, *coreauth.Auth, *http.Request) (*http.Response, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func TestGeminiStreamGenerateContentDoesNotLoseErrorBeforeFirstPayload(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 100; i++ {
|
||||
wg.Add(1)
|
||||
go func(idx int) {
|
||||
defer wg.Done()
|
||||
executor := &initialFailureGeminiStreamExecutor{}
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor)
|
||||
authID := fmt.Sprintf("initial-failure-gemini-auth-%d", idx)
|
||||
auth := &coreauth.Auth{ID: authID, Provider: executor.Identifier(), Status: coreauth.StatusActive}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Errorf("register auth %d: %v", idx, errRegister)
|
||||
return
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: initialFailureGeminiModel}})
|
||||
defer registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
|
||||
base := handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager)
|
||||
h := NewGeminiAPIHandler(base)
|
||||
router := gin.New()
|
||||
router.POST("/v1beta/models/*action", h.GeminiHandler)
|
||||
|
||||
request := httptest.NewRequest(http.MethodPost, "/v1beta/models/initial-failure-gemini-model:streamGenerateContent", strings.NewReader(`{"contents":[{"parts":[{"text":"hi"}]}]}`))
|
||||
request.Header.Set("Content-Type", "application/json")
|
||||
recorder := httptest.NewRecorder()
|
||||
router.ServeHTTP(recorder, request)
|
||||
|
||||
if recorder.Code == http.StatusOK {
|
||||
t.Errorf("request %d lost the buffered initial error and returned HTTP 200: %q", idx, recorder.Body.String())
|
||||
}
|
||||
if !strings.Contains(recorder.Body.String(), "upstream failed before first payload") {
|
||||
t.Errorf("request %d lost the initial upstream error: status=%d body=%q", idx, recorder.Code, recorder.Body.String())
|
||||
}
|
||||
}(i)
|
||||
}
|
||||
wg.Wait()
|
||||
}
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
)
|
||||
|
||||
func TestGeminiModelsResponseUsesConfiguredDisplayName(t *testing.T) {
|
||||
const clientID = "gemini-display-name-catalog-test"
|
||||
const modelID = "gemini-display-name-catalog-test"
|
||||
registryRef := registry.GetGlobalRegistry()
|
||||
registryRef.RegisterClient(clientID, "gemini", []*registry.ModelInfo{{
|
||||
ID: modelID, Name: modelID, DisplayName: "Configured Gemini Name",
|
||||
}})
|
||||
t.Cleanup(func() {
|
||||
registryRef.UnregisterClient(clientID)
|
||||
})
|
||||
|
||||
recorder := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(recorder)
|
||||
NewGeminiAPIHandler(&handlers.BaseAPIHandler{}).GeminiModels(ctx)
|
||||
|
||||
var response struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
DisplayName string `json:"displayName"`
|
||||
} `json:"models"`
|
||||
}
|
||||
if errUnmarshal := json.Unmarshal(recorder.Body.Bytes(), &response); errUnmarshal != nil {
|
||||
t.Fatalf("decode response: %v", errUnmarshal)
|
||||
}
|
||||
for _, model := range response.Models {
|
||||
if model.Name == "models/"+modelID {
|
||||
if model.DisplayName != "Configured Gemini Name" {
|
||||
t.Fatalf("displayName = %q, want Configured Gemini Name", model.DisplayName)
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
t.Fatalf("model %q not found in response", modelID)
|
||||
}
|
||||
202
backend/sdk/api/handlers/gemini/interactions_handlers.go
Normal file
202
backend/sdk/api/handlers/gemini/interactions_handlers.go
Normal file
|
|
@ -0,0 +1,202 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
. "github.com/router-for-me/CLIProxyAPI/v7/internal/constant"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
)
|
||||
|
||||
const interactionsAgentAuthSelectionModel = "gemini-2.5-flash"
|
||||
|
||||
type interactionsRequestTarget struct {
|
||||
Model string
|
||||
Agent string
|
||||
Stream bool
|
||||
}
|
||||
|
||||
func parseInteractionsRequestTarget(rawJSON []byte) (interactionsRequestTarget, error) {
|
||||
if !gjson.ValidBytes(rawJSON) {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("invalid JSON body")
|
||||
}
|
||||
root := gjson.ParseBytes(rawJSON)
|
||||
model := strings.TrimSpace(root.Get("model").String())
|
||||
agent := strings.TrimSpace(root.Get("agent").String())
|
||||
if model == "" && agent == "" {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("request requires exactly one of model or agent")
|
||||
}
|
||||
if model != "" && agent != "" {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("request requires exactly one of model or agent")
|
||||
}
|
||||
streamNode := root.Get("stream")
|
||||
stream := false
|
||||
if streamNode.Exists() {
|
||||
if !streamNode.IsBool() {
|
||||
return interactionsRequestTarget{}, fmt.Errorf("stream must be a boolean")
|
||||
}
|
||||
stream = streamNode.Bool()
|
||||
}
|
||||
return interactionsRequestTarget{Model: model, Agent: agent, Stream: stream}, nil
|
||||
}
|
||||
|
||||
func prepareInteractionsExecutionTarget(rawJSON []byte, target interactionsRequestTarget) (string, []byte) {
|
||||
if target.Agent != "" {
|
||||
return target.Agent, rawJSON
|
||||
}
|
||||
model := normalizeGeminiModelResourceName(target.Model)
|
||||
if model == target.Model {
|
||||
return model, rawJSON
|
||||
}
|
||||
updatedRawJSON, errSet := sjson.SetBytes(rawJSON, "model", model)
|
||||
if errSet != nil {
|
||||
return model, rawJSON
|
||||
}
|
||||
return model, updatedRawJSON
|
||||
}
|
||||
|
||||
func normalizeGeminiModelResourceName(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if strings.HasPrefix(model, "models/") && len(model) > len("models/") {
|
||||
return strings.TrimPrefix(model, "models/")
|
||||
}
|
||||
return model
|
||||
}
|
||||
|
||||
func buildInteractionsExecutionRequest(target interactionsRequestTarget, modelName string, rawJSON []byte, alt string) handlers.ProtocolExecutionRequest {
|
||||
forcedProvider := ""
|
||||
authSelectionModel := ""
|
||||
if target.Agent != "" {
|
||||
forcedProvider = GeminiInteractions
|
||||
authSelectionModel = interactionsAgentAuthSelectionModel
|
||||
}
|
||||
return handlers.ProtocolExecutionRequest{
|
||||
EntryProtocol: Interactions,
|
||||
ExitProtocol: Interactions,
|
||||
ForcedProvider: forcedProvider,
|
||||
AuthSelectionModel: authSelectionModel,
|
||||
Model: modelName,
|
||||
Stream: target.Stream,
|
||||
Body: rawJSON,
|
||||
Alt: alt,
|
||||
}
|
||||
}
|
||||
|
||||
// Interactions handles POST /v1beta/interactions.
|
||||
func (h *GeminiAPIHandler) Interactions(c *gin.Context) {
|
||||
rawJSON, errRead := c.GetRawData()
|
||||
if errRead != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: errRead.Error(), Type: "invalid_request_error"}})
|
||||
return
|
||||
}
|
||||
target, errParse := parseInteractionsRequestTarget(rawJSON)
|
||||
if errParse != nil {
|
||||
c.JSON(http.StatusBadRequest, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: errParse.Error(), Type: "invalid_request_error"}})
|
||||
return
|
||||
}
|
||||
|
||||
modelName, resolvedRawJSON := prepareInteractionsExecutionTarget(rawJSON, target)
|
||||
rawJSON = resolvedRawJSON
|
||||
|
||||
alt := h.GetAlt(c)
|
||||
cliCtx, cliCancel := h.GetContextWithCancel(h, c, context.Background())
|
||||
defer cliCancel(nil)
|
||||
|
||||
req := buildInteractionsExecutionRequest(target, modelName, rawJSON, alt)
|
||||
if target.Stream {
|
||||
h.handleInteractionsStream(c, cliCtx, cliCancel, req)
|
||||
return
|
||||
}
|
||||
h.handleInteractionsNonStream(c, cliCtx, cliCancel, req)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) handleInteractionsNonStream(c *gin.Context, cliCtx context.Context, cliCancel handlers.APIHandlerCancelFunc, req handlers.ProtocolExecutionRequest) {
|
||||
c.Header("Content-Type", "application/json")
|
||||
stopKeepAlive := h.StartNonStreamingKeepAlive(c, cliCtx)
|
||||
resp, errMsg := h.ExecuteProtocolWithAuthManager(cliCtx, req)
|
||||
stopKeepAlive()
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), resp.Headers)
|
||||
_, _ = c.Writer.Write(resp.Body)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) handleInteractionsStream(c *gin.Context, cliCtx context.Context, cliCancel handlers.APIHandlerCancelFunc, req handlers.ProtocolExecutionRequest) {
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, handlers.ErrorResponse{Error: handlers.ErrorDetail{Message: "Streaming not supported", Type: "server_error"}})
|
||||
return
|
||||
}
|
||||
stream, errMsg := h.ExecuteProtocolStreamWithAuthManager(cliCtx, req)
|
||||
if errMsg != nil {
|
||||
h.WriteErrorResponse(c, errMsg)
|
||||
cliCancel(errMsg.Error)
|
||||
return
|
||||
}
|
||||
c.Header("Content-Type", "text/event-stream")
|
||||
c.Header("Cache-Control", "no-cache")
|
||||
c.Header("Connection", "keep-alive")
|
||||
c.Header("Access-Control-Allow-Origin", "*")
|
||||
handlers.WriteUpstreamHeaders(c.Writer.Header(), stream.Headers)
|
||||
data := make(chan []byte)
|
||||
errs := make(chan *interfaces.ErrorMessage, 1)
|
||||
go func() {
|
||||
defer close(data)
|
||||
defer close(errs)
|
||||
for chunk := range stream.Chunks {
|
||||
if chunk.Err != nil {
|
||||
errs <- &interfaces.ErrorMessage{StatusCode: chunk.Err.StatusCode, Error: chunk.Err}
|
||||
return
|
||||
}
|
||||
if len(chunk.Payload) > 0 {
|
||||
data <- chunk.Payload
|
||||
}
|
||||
}
|
||||
}()
|
||||
h.forwardInteractionsStream(c, flusher, func(err error) { cliCancel(err) }, data, errs)
|
||||
}
|
||||
|
||||
func (h *GeminiAPIHandler) forwardInteractionsStream(c *gin.Context, flusher http.Flusher, cancel func(error), data <-chan []byte, errs <-chan *interfaces.ErrorMessage) {
|
||||
h.ForwardStream(c, flusher, cancel, data, errs, handlers.StreamForwardOptions{
|
||||
WriteChunk: func(chunk []byte) {
|
||||
if len(chunk) == 0 {
|
||||
return
|
||||
}
|
||||
trimmed := bytes.TrimSpace(chunk)
|
||||
if bytes.HasPrefix(trimmed, []byte("event:")) || bytes.HasPrefix(trimmed, []byte("data:")) {
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
} else {
|
||||
_, _ = c.Writer.Write([]byte("data: "))
|
||||
_, _ = c.Writer.Write(chunk)
|
||||
}
|
||||
if !bytes.HasSuffix(chunk, []byte("\n\n")) {
|
||||
_, _ = c.Writer.Write([]byte("\n\n"))
|
||||
}
|
||||
},
|
||||
WriteTerminalError: func(errMsg *interfaces.ErrorMessage) {
|
||||
if errMsg == nil {
|
||||
return
|
||||
}
|
||||
status := http.StatusInternalServerError
|
||||
if errMsg.StatusCode > 0 {
|
||||
status = errMsg.StatusCode
|
||||
}
|
||||
errText := http.StatusText(status)
|
||||
if errMsg.Error != nil && errMsg.Error.Error() != "" {
|
||||
errText = errMsg.Error.Error()
|
||||
}
|
||||
body := handlers.BuildErrorResponseBody(status, errText)
|
||||
_, _ = fmt.Fprintf(c.Writer, "event: error\ndata: %s\n\n", string(body))
|
||||
},
|
||||
})
|
||||
}
|
||||
320
backend/sdk/api/handlers/gemini/interactions_handlers_test.go
Normal file
320
backend/sdk/api/handlers/gemini/interactions_handlers_test.go
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
package gemini
|
||||
|
||||
import (
|
||||
"context"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/config"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/interfaces"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/registry"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/internal/runtime/executor"
|
||||
_ "github.com/router-for-me/CLIProxyAPI/v7/internal/translator"
|
||||
"github.com/router-for-me/CLIProxyAPI/v7/sdk/api/handlers"
|
||||
coreauth "github.com/router-for-me/CLIProxyAPI/v7/sdk/cliproxy/auth"
|
||||
sdkconfig "github.com/router-for-me/CLIProxyAPI/v7/sdk/config"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
func TestParseInteractionsRequestTarget(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
body string
|
||||
wantModel string
|
||||
wantAgent string
|
||||
wantErr bool
|
||||
}{
|
||||
{name: "model", body: `{"model":"gemini-3.5-flash","input":"hi"}`, wantModel: "gemini-3.5-flash"},
|
||||
{name: "model resource name", body: `{"model":"models/gemini-3.5-flash","input":"hi"}`, wantModel: "models/gemini-3.5-flash"},
|
||||
{name: "agent", body: `{"agent":"agents/test-agent","input":"hi"}`, wantAgent: "agents/test-agent"},
|
||||
{name: "missing", body: `{"input":"hi"}`, wantErr: true},
|
||||
{name: "both", body: `{"model":"gemini-3.5-flash","agent":"agents/test-agent","input":"hi"}`, wantErr: true},
|
||||
{name: "stream string", body: `{"model":"gemini-3.5-flash","stream":"true","input":"hi"}`, wantErr: true},
|
||||
{name: "stream true", body: `{"model":"gemini-3.5-flash","stream":true,"input":"hi"}`, wantModel: "gemini-3.5-flash"},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(tt.body))
|
||||
if tt.wantErr {
|
||||
if errParse == nil {
|
||||
t.Fatal("parseInteractionsRequestTarget() error = nil, want error")
|
||||
}
|
||||
return
|
||||
}
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
if target.Model != tt.wantModel || target.Agent != tt.wantAgent {
|
||||
t.Fatalf("target = %#v, want model %q agent %q", target, tt.wantModel, tt.wantAgent)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareInteractionsExecutionTargetNormalizesModelResourceName(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"model":"models/gemini-3.5-flash","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
model, body := prepareInteractionsExecutionTarget([]byte(`{"model":"models/gemini-3.5-flash","input":"hi"}`), target)
|
||||
if model != "gemini-3.5-flash" {
|
||||
t.Fatalf("model = %q, want gemini-3.5-flash", model)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "model").String(); got != "gemini-3.5-flash" {
|
||||
t.Fatalf("body model = %q, want gemini-3.5-flash. Body: %s", got, string(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrepareInteractionsExecutionTargetPreservesBareModel(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"model":"gemini-3.5-flash","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
model, body := prepareInteractionsExecutionTarget([]byte(`{"model":"gemini-3.5-flash","input":"hi"}`), target)
|
||||
if model != "gemini-3.5-flash" {
|
||||
t.Fatalf("model = %q, want gemini-3.5-flash", model)
|
||||
}
|
||||
if got := gjson.GetBytes(body, "model").String(); got != "gemini-3.5-flash" {
|
||||
t.Fatalf("body model = %q, want gemini-3.5-flash. Body: %s", got, string(body))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildInteractionsExecutionRequestUsesAgentAuthSelectionModel(t *testing.T) {
|
||||
target, errParse := parseInteractionsRequestTarget([]byte(`{"agent":"agents/test-agent","input":"hi"}`))
|
||||
if errParse != nil {
|
||||
t.Fatalf("parseInteractionsRequestTarget() error = %v", errParse)
|
||||
}
|
||||
req := buildInteractionsExecutionRequest(target, "agents/test-agent", []byte(`{"agent":"agents/test-agent","input":"hi"}`), "")
|
||||
if req.ForcedProvider != "gemini-interactions" {
|
||||
t.Fatalf("ForcedProvider = %q, want gemini-interactions", req.ForcedProvider)
|
||||
}
|
||||
if req.AuthSelectionModel != interactionsAgentAuthSelectionModel {
|
||||
t.Fatalf("AuthSelectionModel = %q, want %q", req.AuthSelectionModel, interactionsAgentAuthSelectionModel)
|
||||
}
|
||||
if req.Model != "agents/test-agent" {
|
||||
t.Fatalf("Model = %q, want agents/test-agent", req.Model)
|
||||
}
|
||||
if got := gjson.GetBytes(req.Body, "agent").String(); got != "agents/test-agent" {
|
||||
t.Fatalf("body agent = %q, want agents/test-agent", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsInvalidJSON(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "invalid_request_error") {
|
||||
t.Fatalf("body = %s, want invalid_request_error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsMissingModelAndAgent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "exactly one of model or agent") {
|
||||
t.Fatalf("body = %s, want model/agent validation error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsBothModelAndAgent(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"gemini-3.5-flash","agent":"agents/test-agent","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "exactly one of model or agent") {
|
||||
t.Fatalf("body = %s, want model/agent validation error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsRejectsNonBooleanStream(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"gemini-3.5-flash","stream":"true","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusBadRequest {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusBadRequest, rec.Body.String())
|
||||
}
|
||||
if !strings.Contains(rec.Body.String(), "invalid_request_error") {
|
||||
t.Fatalf("body = %s, want invalid_request_error", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsAgentUsesNativeInteractionsEndpoint(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
var gotPath string
|
||||
var upstreamBody []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
gotPath = r.URL.Path
|
||||
body, errRead := io.ReadAll(r.Body)
|
||||
if errRead != nil {
|
||||
http.Error(w, errRead.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
upstreamBody = body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"id":"interaction_1","object":"interaction","status":"completed","steps":[{"type":"model_output","content":[{"text":"ok"}]}],"usage":{"input_tokens":1,"output_tokens":1,"total_tokens":2}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor.NewGeminiInteractionsExecutor(&config.Config{RequestRetry: 1}))
|
||||
auth := &coreauth.Auth{
|
||||
ID: "interactions-agent-native-auth",
|
||||
Provider: "gemini-interactions",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"api_key": "test-key",
|
||||
"base_url": server.URL,
|
||||
},
|
||||
Metadata: map[string]any{"email": "interactions-agent@example.com"},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: interactionsAgentAuthSelectionModel}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"agent":"agents/test-agent","input":"hi"}`))
|
||||
h := NewGeminiAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager))
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if gotPath != "/v1beta/interactions" {
|
||||
t.Fatalf("path = %q, want /v1beta/interactions", gotPath)
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "agent").String(); got != "agents/test-agent" {
|
||||
t.Fatalf("upstream agent = %q, want agents/test-agent. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(rec.Body.Bytes(), "id").String(); got != "interaction_1" {
|
||||
t.Fatalf("response id = %q, want interaction_1. Body: %s", got, rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestInteractionsAntigravityModelUsesTranslatorBridge(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
model := "interactions-antigravity-bridge-model"
|
||||
var upstreamBody []byte
|
||||
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path != "/v1internal:generateContent" {
|
||||
http.Error(w, "unexpected path: "+r.URL.Path, http.StatusNotFound)
|
||||
return
|
||||
}
|
||||
body, errRead := io.ReadAll(r.Body)
|
||||
if errRead != nil {
|
||||
http.Error(w, errRead.Error(), http.StatusBadRequest)
|
||||
return
|
||||
}
|
||||
upstreamBody = body
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
_, _ = w.Write([]byte(`{"response":{"responseId":"resp_1","candidates":[{"content":{"role":"model","parts":[{"text":"translated-ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":2,"totalTokenCount":3}}}`))
|
||||
}))
|
||||
defer server.Close()
|
||||
|
||||
manager := coreauth.NewManager(nil, nil, nil)
|
||||
manager.RegisterExecutor(executor.NewAntigravityExecutor(&config.Config{RequestRetry: 1}))
|
||||
auth := &coreauth.Auth{
|
||||
ID: "interactions-antigravity-bridge-auth",
|
||||
Provider: "antigravity",
|
||||
Status: coreauth.StatusActive,
|
||||
Attributes: map[string]string{
|
||||
"base_url": server.URL,
|
||||
},
|
||||
Metadata: map[string]any{
|
||||
"access_token": "token",
|
||||
"project_id": "project-1",
|
||||
"expired": time.Now().Add(time.Hour).Format(time.RFC3339),
|
||||
},
|
||||
}
|
||||
if _, errRegister := manager.Register(context.Background(), auth); errRegister != nil {
|
||||
t.Fatalf("manager.Register(): %v", errRegister)
|
||||
}
|
||||
registry.GetGlobalRegistry().RegisterClient(auth.ID, auth.Provider, []*registry.ModelInfo{{ID: model}})
|
||||
t.Cleanup(func() {
|
||||
registry.GetGlobalRegistry().UnregisterClient(auth.ID)
|
||||
})
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{"model":"`+model+`","input":"hi","generation_config":{"top_p":0.8}}`))
|
||||
h := NewGeminiAPIHandler(handlers.NewBaseAPIHandlers(&sdkconfig.SDKConfig{}, manager))
|
||||
|
||||
h.Interactions(ctx)
|
||||
|
||||
if rec.Code != http.StatusOK {
|
||||
t.Fatalf("status = %d, want %d; body=%s", rec.Code, http.StatusOK, rec.Body.String())
|
||||
}
|
||||
if gjson.GetBytes(upstreamBody, "input").Exists() {
|
||||
t.Fatalf("upstream body still contains raw interactions input: %s", string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "request.contents.0.parts.0.text").String(); got != "hi" {
|
||||
t.Fatalf("upstream request text = %q, want hi. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(upstreamBody, "request.generationConfig.topP").Float(); got != 0.8 {
|
||||
t.Fatalf("upstream topP = %v, want 0.8. Body: %s", got, string(upstreamBody))
|
||||
}
|
||||
if got := gjson.GetBytes(rec.Body.Bytes(), "steps.0.content.0.text").String(); got != "translated-ok" {
|
||||
t.Fatalf("response text = %q, want translated-ok. Body: %s", got, rec.Body.String())
|
||||
}
|
||||
if gjson.GetBytes(rec.Body.Bytes(), "response").Exists() {
|
||||
t.Fatalf("response still contains raw antigravity response wrapper: %s", rec.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestForwardInteractionsStreamWrapsBareJSONAsSSEData(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
ctx, _ := gin.CreateTestContext(rec)
|
||||
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1beta/interactions", strings.NewReader(`{}`))
|
||||
data := make(chan []byte, 1)
|
||||
errs := make(chan *interfaces.ErrorMessage)
|
||||
data <- []byte(`{"type":"interaction.completed"}`)
|
||||
close(data)
|
||||
close(errs)
|
||||
h := NewGeminiAPIHandler(&handlers.BaseAPIHandler{})
|
||||
|
||||
h.forwardInteractionsStream(ctx, rec, func(error) {}, data, errs)
|
||||
|
||||
if got := rec.Body.String(); got != "data: {\"type\":\"interaction.completed\"}\n\n" {
|
||||
t.Fatalf("body = %q, want SSE data frame", got)
|
||||
}
|
||||
}
|
||||
Loading…
Reference in a new issue