310 lines
8.2 KiB
Go
310 lines
8.2 KiB
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
|
|
buffer bytes.Buffer
|
|
stream *geminiBridgeStream
|
|
}
|
|
|
|
func newGeminiBridgeResponseWriter(target http.ResponseWriter, model string) *geminiBridgeResponseWriter {
|
|
return &geminiBridgeResponseWriter{
|
|
target: target, header: make(http.Header), stream: newGeminiBridgeStream(target, model),
|
|
}
|
|
}
|
|
|
|
func (w *geminiBridgeResponseWriter) Header() http.Header { return w.header }
|
|
|
|
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 {
|
|
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.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))
|
|
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")
|
|
}
|
|
part := map[string]any{"functionCall": map[string]any{"name": state.name, "args": args}}
|
|
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
|
|
}
|