572 lines
17 KiB
Go
572 lines
17 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"regexp"
|
|
"strings"
|
|
|
|
"go.uber.org/zap"
|
|
)
|
|
|
|
var geminiPathToken = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._-]*$`)
|
|
var geminiCallerModelLabel = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9 ._-]{0,127}$`)
|
|
var geminiToolCallID = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$`)
|
|
|
|
const geminiRejectionLogMessage = "edge_gemini_rejection"
|
|
|
|
type geminiRejectionClass string
|
|
|
|
const (
|
|
geminiRejectionPreIngress geminiRejectionClass = "pre_ingress"
|
|
geminiRejectionProviderHTTP geminiRejectionClass = "provider_http"
|
|
)
|
|
|
|
func isGeminiRequest(r *http.Request) bool {
|
|
return r != nil && strings.HasPrefix(r.URL.Path, geminiPathPrefix)
|
|
}
|
|
|
|
func writeGeminiError(w http.ResponseWriter, status int, code, message string) {
|
|
writeJSON(w, status, geminiErrorResponse{Error: geminiErrorBody{
|
|
Code: status, Message: message, Status: code,
|
|
}})
|
|
}
|
|
|
|
func (s *Server) observeGeminiRejection(class geminiRejectionClass, status int) {
|
|
s.logger.Info(
|
|
geminiRejectionLogMessage,
|
|
zap.String("surface", "gemini"),
|
|
zap.String("bridge", "chat"),
|
|
zap.String("rejection_class", string(class)),
|
|
zap.Int("http_status", status),
|
|
)
|
|
}
|
|
|
|
func (s *Server) writeGeminiPreIngressError(w http.ResponseWriter, status int, code, message string) {
|
|
s.observeGeminiRejection(geminiRejectionPreIngress, status)
|
|
writeGeminiError(w, status, code, message)
|
|
}
|
|
|
|
func (s *Server) handleGeminiStreamGenerateContent(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
s.writeGeminiPreIngressError(w, http.StatusMethodNotAllowed, "INVALID_ARGUMENT", "method not allowed")
|
|
return
|
|
}
|
|
routeID, callerModel, err := parseGeminiStreamPath(r)
|
|
if err != nil {
|
|
s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request path is invalid")
|
|
return
|
|
}
|
|
defer r.Body.Close()
|
|
body, err := readOpenAIIngressBody(w, r, s.maxIngressSnapshotBytes())
|
|
if err != nil {
|
|
s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request body is invalid")
|
|
return
|
|
}
|
|
chatBody, err := prepareGeminiChatBridge(body, routeID)
|
|
if err != nil {
|
|
s.writeGeminiPreIngressError(w, http.StatusBadRequest, "INVALID_ARGUMENT", "request body is invalid")
|
|
return
|
|
}
|
|
|
|
internal := r.Clone(r.Context())
|
|
internal.URL.Path = "/v1/chat/completions"
|
|
internal.URL.RawPath = ""
|
|
internal.URL.RawQuery = ""
|
|
internal.RequestURI = "/v1/chat/completions"
|
|
internal.Body = io.NopCloser(bytes.NewReader(chatBody))
|
|
internal.ContentLength = int64(len(chatBody))
|
|
internal.Header = r.Header.Clone()
|
|
internal.Header.Del("Authorization")
|
|
internal.Header.Del("X-Goog-Api-Key")
|
|
internal.Header.Set("Content-Type", "application/json")
|
|
|
|
bridge := newGeminiBridgeResponseWriter(w, callerModel, func(status int) {
|
|
s.observeGeminiRejection(geminiRejectionProviderHTTP, status)
|
|
})
|
|
s.handleChatCompletions(bridge, internal)
|
|
bridge.Finish()
|
|
}
|
|
|
|
func parseGeminiStreamPath(r *http.Request) (string, string, error) {
|
|
if r == nil || r.URL == nil {
|
|
return "", "", fmt.Errorf("missing URL")
|
|
}
|
|
parts := strings.Split(strings.TrimPrefix(r.URL.Path, geminiPathPrefix), "/")
|
|
if len(parts) != 4 || parts[1] != "v1beta" || parts[2] != "models" {
|
|
return "", "", fmt.Errorf("unexpected path")
|
|
}
|
|
const suffix = ":streamGenerateContent"
|
|
if !strings.HasSuffix(parts[3], suffix) {
|
|
return "", "", fmt.Errorf("unexpected method")
|
|
}
|
|
routeID := parts[0]
|
|
callerModel := strings.TrimSuffix(parts[3], suffix)
|
|
if !geminiPathToken.MatchString(routeID) || !geminiCallerModelLabel.MatchString(callerModel) || strings.TrimSpace(callerModel) != callerModel {
|
|
return "", "", fmt.Errorf("invalid path token")
|
|
}
|
|
query := r.URL.Query()
|
|
if len(query) != 1 || len(query["alt"]) != 1 || query.Get("alt") != "sse" {
|
|
return "", "", fmt.Errorf("alt=sse is required")
|
|
}
|
|
return routeID, callerModel, nil
|
|
}
|
|
|
|
func prepareGeminiChatBridge(body []byte, routeID string) ([]byte, error) {
|
|
if err := validateJSONMembers(body); err != nil {
|
|
return nil, err
|
|
}
|
|
decoder := json.NewDecoder(bytes.NewReader(body))
|
|
decoder.DisallowUnknownFields()
|
|
var req geminiRequest
|
|
if err := decoder.Decode(&req); err != nil {
|
|
return nil, err
|
|
}
|
|
if err := requireGeminiJSONEOF(decoder); err != nil {
|
|
return nil, err
|
|
}
|
|
if len(req.Contents) == 0 {
|
|
return nil, fmt.Errorf("contents are required")
|
|
}
|
|
chat := map[string]any{
|
|
"model": routeID, "stream": true,
|
|
"stream_options": map[string]any{"include_usage": true},
|
|
}
|
|
messages := make([]map[string]any, 0, len(req.Contents)+1)
|
|
if req.SystemInstruction != nil {
|
|
text, err := geminiTextOnly(*req.SystemInstruction)
|
|
if err != nil || strings.TrimSpace(text) == "" {
|
|
return nil, fmt.Errorf("systemInstruction is invalid")
|
|
}
|
|
messages = append(messages, map[string]any{"role": "system", "content": text})
|
|
}
|
|
pendingCalls := newGeminiPendingCalls()
|
|
for contentIndex, content := range req.Contents {
|
|
converted, err := geminiContentToChat(content, contentIndex, pendingCalls)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
messages = append(messages, converted...)
|
|
}
|
|
chat["messages"] = messages
|
|
if config := req.GenerationConfig; config != nil {
|
|
if config.CandidateCount != nil && *config.CandidateCount != 1 {
|
|
return nil, fmt.Errorf("candidateCount must be one")
|
|
}
|
|
if config.MaxOutputTokens != nil {
|
|
if *config.MaxOutputTokens <= 0 {
|
|
return nil, fmt.Errorf("maxOutputTokens must be positive")
|
|
}
|
|
chat["max_tokens"] = *config.MaxOutputTokens
|
|
}
|
|
if len(config.StopSequences) > 0 {
|
|
chat["stop"] = config.StopSequences
|
|
}
|
|
if config.Temperature != nil {
|
|
if *config.Temperature < 0 || *config.Temperature > 2 {
|
|
return nil, fmt.Errorf("temperature is invalid")
|
|
}
|
|
}
|
|
if config.TopK != nil {
|
|
if *config.TopK <= 0 {
|
|
return nil, fmt.Errorf("topK is invalid")
|
|
}
|
|
}
|
|
if config.TopP != nil {
|
|
if *config.TopP < 0 || *config.TopP > 1 {
|
|
return nil, fmt.Errorf("topP is invalid")
|
|
}
|
|
}
|
|
if thinking := config.ThinkingConfig; thinking != nil {
|
|
googleThinking := make(map[string]any)
|
|
if thinking.IncludeThoughts != nil {
|
|
googleThinking["include_thoughts"] = *thinking.IncludeThoughts
|
|
}
|
|
if thinking.ThinkingBudget != nil {
|
|
if *thinking.ThinkingBudget < -1 {
|
|
return nil, fmt.Errorf("thinkingBudget is invalid")
|
|
}
|
|
googleThinking["thinking_budget"] = *thinking.ThinkingBudget
|
|
}
|
|
if len(googleThinking) > 0 {
|
|
// Gemini's OpenAI-compatible endpoint accepts native options only
|
|
// below extra_body.google; generic think/include_reasoning fields
|
|
// are rejected by that endpoint.
|
|
chat["extra_body"] = map[string]any{
|
|
"google": map[string]any{"thinking_config": googleThinking},
|
|
}
|
|
}
|
|
}
|
|
responseSchema, err := geminiExclusiveSchema(config.ResponseSchema, config.ResponseJSONSchema)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("response schema is invalid")
|
|
}
|
|
switch config.ResponseMIMEType {
|
|
case "":
|
|
if responseSchema != nil {
|
|
return nil, fmt.Errorf("response MIME type is required")
|
|
}
|
|
case "text/plain":
|
|
if responseSchema != nil {
|
|
return nil, fmt.Errorf("text response schema is invalid")
|
|
}
|
|
case "application/json":
|
|
if responseSchema == nil {
|
|
chat["response_format"] = map[string]any{"type": "json_object"}
|
|
} else {
|
|
chat["response_format"] = map[string]any{
|
|
"type": "json_schema",
|
|
"json_schema": map[string]any{
|
|
"name": "agy_response", "strict": true, "schema": responseSchema,
|
|
},
|
|
}
|
|
}
|
|
default:
|
|
return nil, fmt.Errorf("response MIME type is invalid")
|
|
}
|
|
}
|
|
if len(req.Tools) > 0 {
|
|
tools := make([]map[string]any, 0)
|
|
for _, group := range req.Tools {
|
|
if len(group.FunctionDeclarations) == 0 {
|
|
return nil, fmt.Errorf("functionDeclarations are required")
|
|
}
|
|
for _, declaration := range group.FunctionDeclarations {
|
|
if !geminiPathToken.MatchString(declaration.Name) {
|
|
return nil, fmt.Errorf("function declaration is invalid")
|
|
}
|
|
schema, err := geminiExclusiveSchema(
|
|
declaration.Parameters, declaration.ParametersJSONSchema,
|
|
)
|
|
if err != nil || schema == nil {
|
|
return nil, fmt.Errorf("function schema is invalid")
|
|
}
|
|
if _, err := geminiExclusiveSchema(
|
|
declaration.Response, declaration.ResponseJSONSchema,
|
|
); err != nil {
|
|
return nil, fmt.Errorf("function response schema is invalid")
|
|
}
|
|
function := map[string]any{"name": declaration.Name, "parameters": schema}
|
|
if declaration.Description != "" {
|
|
function["description"] = declaration.Description
|
|
}
|
|
tools = append(tools, map[string]any{"type": "function", "function": function})
|
|
}
|
|
}
|
|
chat["tools"] = tools
|
|
}
|
|
if req.ToolConfig != nil && req.ToolConfig.FunctionCallingConfig != nil {
|
|
switch strings.ToUpper(req.ToolConfig.FunctionCallingConfig.Mode) {
|
|
case "", "AUTO":
|
|
chat["tool_choice"] = "auto"
|
|
case "ANY":
|
|
chat["tool_choice"] = "required"
|
|
case "NONE":
|
|
chat["tool_choice"] = "none"
|
|
default:
|
|
return nil, fmt.Errorf("function calling mode is invalid")
|
|
}
|
|
}
|
|
return json.Marshal(chat)
|
|
}
|
|
|
|
func geminiExclusiveSchema(first, second json.RawMessage) (map[string]any, error) {
|
|
if len(first) > 0 && len(second) > 0 {
|
|
return nil, fmt.Errorf("schema alternatives conflict")
|
|
}
|
|
raw := first
|
|
if len(raw) == 0 {
|
|
raw = second
|
|
}
|
|
if len(raw) == 0 {
|
|
return nil, nil
|
|
}
|
|
var schema map[string]any
|
|
if err := json.Unmarshal(raw, &schema); err != nil || schema == nil {
|
|
return nil, fmt.Errorf("schema is not an object")
|
|
}
|
|
if err := normalizeGeminiSchemaTypes(schema); err != nil {
|
|
return nil, err
|
|
}
|
|
return schema, nil
|
|
}
|
|
|
|
func normalizeGeminiSchemaTypes(value any) error {
|
|
switch typed := value.(type) {
|
|
case map[string]any:
|
|
for key, item := range typed {
|
|
if key == "type" {
|
|
name, ok := item.(string)
|
|
if !ok {
|
|
return fmt.Errorf("schema type is invalid")
|
|
}
|
|
normalized := strings.ToLower(name)
|
|
switch normalized {
|
|
case "null", "boolean", "object", "array", "number", "integer", "string":
|
|
typed[key] = normalized
|
|
default:
|
|
return fmt.Errorf("schema type is invalid")
|
|
}
|
|
continue
|
|
}
|
|
if err := normalizeGeminiSchemaTypes(item); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
case []any:
|
|
for _, item := range typed {
|
|
if err := normalizeGeminiSchemaTypes(item); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func geminiTextOnly(content geminiContent) (string, error) {
|
|
// Gemini represents systemInstruction as Content and official agy 1.1.12
|
|
// labels that Content with the API-native "user" role.
|
|
if content.Role != "" && content.Role != "user" {
|
|
return "", fmt.Errorf("role is invalid")
|
|
}
|
|
texts := make([]string, 0, len(content.Parts))
|
|
for _, part := range content.Parts {
|
|
if part.Text == nil || part.FunctionCall != nil || part.FunctionResponse != nil || part.Thought || part.ThoughtSignature != "" {
|
|
return "", fmt.Errorf("only text is supported")
|
|
}
|
|
texts = append(texts, *part.Text)
|
|
}
|
|
return strings.Join(texts, "\n"), nil
|
|
}
|
|
|
|
type geminiPendingCall struct {
|
|
id string
|
|
name string
|
|
explicit bool
|
|
}
|
|
|
|
type geminiPendingCalls struct {
|
|
byID map[string]geminiPendingCall
|
|
byName map[string][]string
|
|
}
|
|
|
|
func newGeminiPendingCalls() *geminiPendingCalls {
|
|
return &geminiPendingCalls{byID: make(map[string]geminiPendingCall), byName: make(map[string][]string)}
|
|
}
|
|
|
|
func (p *geminiPendingCalls) add(name, id string, explicit bool) error {
|
|
if _, duplicate := p.byID[id]; duplicate {
|
|
return fmt.Errorf("duplicate functionCall id")
|
|
}
|
|
p.byID[id] = geminiPendingCall{id: id, name: name, explicit: explicit}
|
|
p.byName[name] = append(p.byName[name], id)
|
|
return nil
|
|
}
|
|
|
|
func (p *geminiPendingCalls) consume(name, id string) (string, error) {
|
|
if id != "" {
|
|
if !geminiToolCallID.MatchString(id) {
|
|
return "", fmt.Errorf("functionResponse id is invalid")
|
|
}
|
|
call, ok := p.byID[id]
|
|
if !ok || call.name != name {
|
|
return "", fmt.Errorf("functionResponse has no matching call")
|
|
}
|
|
p.remove(call)
|
|
return id, nil
|
|
}
|
|
ids := p.byName[name]
|
|
if len(ids) == 0 {
|
|
return "", fmt.Errorf("functionResponse has no matching call")
|
|
}
|
|
call, ok := p.byID[ids[0]]
|
|
if !ok || call.explicit {
|
|
return "", fmt.Errorf("functionResponse id is required")
|
|
}
|
|
p.remove(call)
|
|
return call.id, nil
|
|
}
|
|
|
|
func (p *geminiPendingCalls) remove(call geminiPendingCall) {
|
|
delete(p.byID, call.id)
|
|
ids := p.byName[call.name]
|
|
for index, id := range ids {
|
|
if id != call.id {
|
|
continue
|
|
}
|
|
ids = append(ids[:index], ids[index+1:]...)
|
|
break
|
|
}
|
|
if len(ids) == 0 {
|
|
delete(p.byName, call.name)
|
|
} else {
|
|
p.byName[call.name] = ids
|
|
}
|
|
}
|
|
|
|
func geminiContentToChat(content geminiContent, contentIndex int, pending *geminiPendingCalls) ([]map[string]any, error) {
|
|
role := strings.ToLower(strings.TrimSpace(content.Role))
|
|
if role != "user" && role != "model" {
|
|
return nil, fmt.Errorf("content role is invalid")
|
|
}
|
|
if len(content.Parts) == 0 {
|
|
return nil, fmt.Errorf("content parts are required")
|
|
}
|
|
var texts []string
|
|
var reasoning []string
|
|
var toolCalls []any
|
|
var toolMessages []map[string]any
|
|
for partIndex, part := range content.Parts {
|
|
set := 0
|
|
if part.Text != nil {
|
|
set++
|
|
}
|
|
if part.FunctionCall != nil {
|
|
set++
|
|
}
|
|
if part.FunctionResponse != nil {
|
|
set++
|
|
}
|
|
if set != 1 {
|
|
return nil, fmt.Errorf("content part is invalid")
|
|
}
|
|
if part.Text != nil {
|
|
if role == "model" && part.Thought {
|
|
reasoning = append(reasoning, *part.Text)
|
|
} else if part.Thought {
|
|
return nil, fmt.Errorf("user thought is invalid")
|
|
} else {
|
|
texts = append(texts, *part.Text)
|
|
}
|
|
continue
|
|
}
|
|
if part.FunctionCall != nil {
|
|
if role != "model" || !geminiPathToken.MatchString(part.FunctionCall.Name) {
|
|
return nil, fmt.Errorf("functionCall is invalid")
|
|
}
|
|
var args map[string]any
|
|
if json.Unmarshal(part.FunctionCall.Args, &args) != nil {
|
|
return nil, fmt.Errorf("functionCall args are invalid")
|
|
}
|
|
callID := part.FunctionCall.ID
|
|
explicitID := callID != ""
|
|
if !explicitID {
|
|
callID = fmt.Sprintf("gemini_call_%d_%d", contentIndex, partIndex)
|
|
} else if !geminiToolCallID.MatchString(callID) {
|
|
return nil, fmt.Errorf("functionCall id is invalid")
|
|
}
|
|
if err := pending.add(part.FunctionCall.Name, callID, explicitID); err != nil {
|
|
return nil, err
|
|
}
|
|
call := map[string]any{"id": callID, "type": "function", "function": map[string]any{"name": part.FunctionCall.Name, "arguments": string(part.FunctionCall.Args)}}
|
|
if part.ThoughtSignature != "" {
|
|
call["extra_content"] = openAIChatThoughtSignature(part.ThoughtSignature)
|
|
}
|
|
toolCalls = append(toolCalls, call)
|
|
continue
|
|
}
|
|
response := part.FunctionResponse
|
|
if !geminiPathToken.MatchString(response.Name) {
|
|
return nil, fmt.Errorf("functionResponse is invalid")
|
|
}
|
|
if !json.Valid(response.Response) {
|
|
return nil, fmt.Errorf("functionResponse has no matching call")
|
|
}
|
|
callID, err := pending.consume(response.Name, response.ID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
toolMessages = append(toolMessages, map[string]any{"role": "tool", "tool_call_id": callID, "content": string(response.Response)})
|
|
}
|
|
if role == "model" {
|
|
if len(toolMessages) > 0 {
|
|
if len(texts) > 0 || len(reasoning) > 0 || len(toolCalls) > 0 {
|
|
return nil, fmt.Errorf("model functionResponse cannot be mixed with assistant content")
|
|
}
|
|
return toolMessages, nil
|
|
}
|
|
message := map[string]any{"role": "assistant", "content": strings.Join(texts, "\n")}
|
|
if len(reasoning) > 0 {
|
|
message["reasoning_content"] = strings.Join(reasoning, "")
|
|
}
|
|
if len(toolCalls) > 0 {
|
|
message["tool_calls"] = toolCalls
|
|
}
|
|
return []map[string]any{message}, nil
|
|
}
|
|
if len(texts) > 0 {
|
|
toolMessages = append(toolMessages, map[string]any{"role": "user", "content": strings.Join(texts, "\n")})
|
|
}
|
|
if len(toolMessages) == 0 {
|
|
return nil, fmt.Errorf("user content is empty")
|
|
}
|
|
return toolMessages, nil
|
|
}
|
|
|
|
func validateJSONMembers(raw []byte) error {
|
|
dec := json.NewDecoder(bytes.NewReader(raw))
|
|
if err := validateGeminiJSONValue(dec); err != nil {
|
|
return err
|
|
}
|
|
return requireGeminiJSONEOF(dec)
|
|
}
|
|
|
|
func validateGeminiJSONValue(dec *json.Decoder) error {
|
|
token, err := dec.Token()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
delim, ok := token.(json.Delim)
|
|
if !ok {
|
|
return nil
|
|
}
|
|
switch delim {
|
|
case '{':
|
|
seen := map[string]struct{}{}
|
|
for dec.More() {
|
|
keyToken, err := dec.Token()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
key, ok := keyToken.(string)
|
|
if !ok {
|
|
return fmt.Errorf("object key is invalid")
|
|
}
|
|
if _, duplicate := seen[key]; duplicate {
|
|
return fmt.Errorf("duplicate member")
|
|
}
|
|
seen[key] = struct{}{}
|
|
if err := validateGeminiJSONValue(dec); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
closeToken, err := dec.Token()
|
|
if err != nil || closeToken != json.Delim('}') {
|
|
return fmt.Errorf("object is invalid")
|
|
}
|
|
case '[':
|
|
for dec.More() {
|
|
if err := validateGeminiJSONValue(dec); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
closeToken, err := dec.Token()
|
|
if err != nil || closeToken != json.Delim(']') {
|
|
return fmt.Errorf("array is invalid")
|
|
}
|
|
default:
|
|
return fmt.Errorf("JSON delimiter is invalid")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func requireGeminiJSONEOF(dec *json.Decoder) error {
|
|
var extra any
|
|
if err := dec.Decode(&extra); err != io.EOF {
|
|
return fmt.Errorf("trailing JSON value")
|
|
}
|
|
return nil
|
|
}
|