iop/apps/edge/internal/openai/gemini_bridge.go

package openai

import (
	"bytes"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"
	"strings"
)

type geminiBridgeResponseWriter struct {
	target                 http.ResponseWriter
	header                 http.Header
	status                 int
	committed              bool
	providerStatusObserved bool
	observeProviderStatus  func(int)
	buffer                 bytes.Buffer
	stream                 *geminiBridgeStream
}

func newGeminiBridgeResponseWriter(target http.ResponseWriter, model string, observeProviderStatus func(int)) *geminiBridgeResponseWriter {
	return &geminiBridgeResponseWriter{
		target: target, header: make(http.Header), observeProviderStatus: observeProviderStatus, stream: newGeminiBridgeStream(target, model),
	}
}

func (w *geminiBridgeResponseWriter) Header() http.Header { return w.header }

func (w *geminiBridgeResponseWriter) observeActualProviderHTTPStatus(status int) {
	if status < http.StatusBadRequest || w.providerStatusObserved || w.observeProviderStatus == nil {
		return
	}
	w.providerStatusObserved = true
	w.observeProviderStatus(status)
}

func (w *geminiBridgeResponseWriter) WriteHeader(status int) {
	if w.status == 0 {
		w.status = status
	}
}

func (w *geminiBridgeResponseWriter) Write(payload []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	if w.status >= http.StatusBadRequest {
		return w.buffer.Write(payload)
	}
	w.commit()
	if err := w.stream.Feed(payload); err != nil {
		_ = w.stream.Error("upstream stream could not be translated")
		return 0, err
	}
	return len(payload), nil
}

func (w *geminiBridgeResponseWriter) Flush() {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	if w.status < http.StatusBadRequest {
		w.commit()
	}
	if flusher, ok := w.target.(http.Flusher); ok {
		flusher.Flush()
	}
}

func (w *geminiBridgeResponseWriter) commit() {
	if w.committed {
		return
	}
	w.target.Header().Set("Content-Type", "text/event-stream")
	w.target.Header().Set("Cache-Control", "no-cache")
	w.target.Header().Del("Content-Length")
	w.target.WriteHeader(http.StatusOK)
	w.committed = true
}

func (w *geminiBridgeResponseWriter) Finish() {
	if w.status == 0 {
		writeGeminiError(w.target, http.StatusBadGateway, "UNAVAILABLE", "runtime request failed")
		return
	}
	if w.status >= http.StatusBadRequest {
		status, code := geminiHTTPError(w.status)
		writeGeminiError(w.target, status, code, geminiSafeErrorMessage(status))
		return
	}
	w.commit()
	if err := w.stream.Finish(); err != nil {
		_ = w.stream.Error("upstream stream could not be translated")
	}
	if flusher, ok := w.target.(http.Flusher); ok {
		flusher.Flush()
	}
}

func geminiHTTPError(status int) (int, string) {
	switch {
	case status == http.StatusUnauthorized || status == http.StatusForbidden:
		return http.StatusUnauthorized, "UNAUTHENTICATED"
	case status >= 400 && status < 500:
		return http.StatusBadRequest, "INVALID_ARGUMENT"
	default:
		return http.StatusBadGateway, "UNAVAILABLE"
	}
}

func geminiSafeErrorMessage(status int) string {
	if status == http.StatusUnauthorized {
		return "authentication failed"
	}
	if status == http.StatusBadRequest {
		return "request is invalid"
	}
	return "runtime request failed"
}

type geminiBridgeToolState struct {
	id               string
	name             string
	thoughtSignature string
	arguments        strings.Builder
}

type geminiBridgeStream struct {
	w          http.ResponseWriter
	model      string
	pendingSSE []byte
	tools      map[int]*geminiBridgeToolState
	finish     string
	usage      map[string]int
	stopped    bool
	errored    bool
}

func newGeminiBridgeStream(w http.ResponseWriter, model string) *geminiBridgeStream {
	return &geminiBridgeStream{w: w, model: model, tools: make(map[int]*geminiBridgeToolState), usage: make(map[string]int)}
}

func (s *geminiBridgeStream) Feed(chunk []byte) error {
	if s.stopped {
		return nil
	}
	s.pendingSSE = append(s.pendingSSE, chunk...)
	s.pendingSSE = bytes.ReplaceAll(s.pendingSSE, []byte("\r\n"), []byte("\n"))
	for {
		index := bytes.Index(s.pendingSSE, []byte("\n\n"))
		if index < 0 {
			return nil
		}
		event := append([]byte(nil), s.pendingSSE[:index]...)
		s.pendingSSE = s.pendingSSE[index+2:]
		if err := s.consumeSSEEvent(event); err != nil {
			return err
		}
		if s.stopped {
			s.pendingSSE = nil
			return nil
		}
	}
}

func (s *geminiBridgeStream) consumeSSEEvent(event []byte) error {
	var lines [][]byte
	for _, line := range bytes.Split(event, []byte("\n")) {
		line = bytes.TrimSpace(line)
		if bytes.HasPrefix(line, []byte("data:")) {
			lines = append(lines, bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))))
		}
	}
	if len(lines) == 0 {
		return nil
	}
	payload := bytes.Join(lines, []byte("\n"))
	if bytes.Equal(payload, []byte("[DONE]")) {
		return s.Finish()
	}
	var chunk geminiChatStreamChunk
	if err := json.Unmarshal(payload, &chunk); err != nil {
		return fmt.Errorf("decode Chat SSE: %w", err)
	}
	if chunk.Error != nil {
		return s.Error("upstream request failed")
	}
	if chunk.Usage != nil {
		setGeminiUsage(s.usage, "promptTokenCount", chunk.Usage.PromptTokens)
		setGeminiUsage(s.usage, "candidatesTokenCount", chunk.Usage.CompletionTokens)
		setGeminiUsage(s.usage, "totalTokenCount", chunk.Usage.TotalTokens)
		if chunk.Usage.PromptTokensDetails != nil {
			setGeminiUsage(s.usage, "cachedContentTokenCount", chunk.Usage.PromptTokensDetails.CachedTokens)
		}
		if chunk.Usage.CompletionTokensDetails != nil {
			setGeminiUsage(s.usage, "thoughtsTokenCount", chunk.Usage.CompletionTokensDetails.ReasoningTokens)
		}
	}
	for _, choice := range chunk.Choices {
		reasoning := choice.Delta.ReasoningContent
		if reasoning == "" {
			reasoning = choice.Delta.Reasoning
		}
		if reasoning != "" {
			if err := s.emitParts([]any{map[string]any{"text": reasoning, "thought": true}}, ""); err != nil {
				return err
			}
		}
		if choice.Delta.Content != "" {
			if err := s.emitParts([]any{map[string]any{"text": choice.Delta.Content}}, ""); err != nil {
				return err
			}
		}
		for _, delta := range choice.Delta.ToolCalls {
			state := s.tools[delta.Index]
			if state == nil {
				state = &geminiBridgeToolState{}
				s.tools[delta.Index] = state
			}
			if delta.Function.Name != "" {
				state.name = delta.Function.Name
			}
			if delta.ID != "" {
				if !geminiToolCallID.MatchString(delta.ID) || (state.id != "" && state.id != delta.ID) {
					return fmt.Errorf("tool call id is invalid")
				}
				state.id = delta.ID
			}
			if delta.ExtraContent.Google != nil && delta.ExtraContent.Google.ThoughtSignature != "" {
				state.thoughtSignature = delta.ExtraContent.Google.ThoughtSignature
			}
			if state.arguments.Len()+len(delta.Function.Arguments) > geminiToolArgumentLimit {
				return fmt.Errorf("tool arguments exceed limit")
			}
			state.arguments.WriteString(delta.Function.Arguments)
		}
		if choice.FinishReason != nil {
			s.finish = *choice.FinishReason
		}
	}
	return nil
}

func setGeminiUsage(target map[string]int, key string, value *int) {
	if value != nil {
		target[key] = *value
	}
}

func (s *geminiBridgeStream) emitParts(parts []any, finish string) error {
	candidate := map[string]any{"content": map[string]any{"role": "model", "parts": parts}}
	if finish != "" {
		candidate["finishReason"] = finish
	}
	payload := map[string]any{"candidates": []any{candidate}, "modelVersion": s.model}
	if len(s.usage) > 0 {
		payload["usageMetadata"] = s.usage
	}
	return writeGeminiSSE(s.w, payload)
}

func (s *geminiBridgeStream) emitTools() error {
	if len(s.tools) == 0 {
		return nil
	}
	indices := make([]int, 0, len(s.tools))
	for index := range s.tools {
		indices = append(indices, index)
	}
	sort.Ints(indices)
	parts := make([]any, 0, len(indices))
	seenIDs := make(map[string]struct{})
	for _, index := range indices {
		state := s.tools[index]
		if !geminiPathToken.MatchString(state.name) {
			return fmt.Errorf("tool name is invalid")
		}
		var args map[string]any
		if json.Unmarshal([]byte(state.arguments.String()), &args) != nil {
			return fmt.Errorf("tool arguments are invalid")
		}
		call := map[string]any{"name": state.name, "args": args}
		if state.id != "" {
			if _, duplicate := seenIDs[state.id]; duplicate {
				return fmt.Errorf("duplicate tool call id")
			}
			seenIDs[state.id] = struct{}{}
			call["id"] = state.id
		}
		part := map[string]any{"functionCall": call}
		if state.thoughtSignature != "" {
			part["thoughtSignature"] = state.thoughtSignature
		}
		parts = append(parts, part)
	}
	return s.emitParts(parts, "")
}

func (s *geminiBridgeStream) Finish() error {
	if s.stopped {
		return nil
	}
	if len(bytes.TrimSpace(s.pendingSSE)) > 0 {
		return fmt.Errorf("truncated Chat SSE")
	}
	if err := s.emitTools(); err != nil {
		return err
	}
	finish := "STOP"
	if s.finish == "length" {
		finish = "MAX_TOKENS"
	}
	s.stopped = true
	return s.emitParts([]any{}, finish)
}

func (s *geminiBridgeStream) Error(message string) error {
	if s.errored || s.stopped {
		return nil
	}
	s.errored, s.stopped = true, true
	return writeGeminiSSE(s.w, geminiErrorResponse{Error: geminiErrorBody{
		Code: http.StatusBadGateway, Message: message, Status: "UNAVAILABLE",
	}})
}

func writeGeminiSSE(w http.ResponseWriter, value any) error {
	payload, err := json.Marshal(value)
	if err != nil {
		return err
	}
	_, err = fmt.Fprintf(w, "data: %s\n\n", payload)
	return err
}