Files
token_thief/proxy/sse.go
T

708 lines
20 KiB
Go

package proxy
import (
"bufio"
"bytes"
"encoding/json"
"strings"
)
func assembleSSEJSON(body []byte) ([]byte, bool) {
payloads := sseDataPayloads(body)
if len(payloads) == 0 {
return nil, false
}
if assembled, ok := assembleOpenAICompletionsSSE(payloads); ok {
return assembled, true
}
if assembled, ok := assembleOpenAIChatSSE(payloads); ok {
return assembled, true
}
if assembled, ok := assembleOpenAIResponsesSSE(payloads); ok {
return assembled, true
}
if assembled, ok := assembleAnthropicSSE(payloads); ok {
return assembled, true
}
if assembled, ok := assembleGeminiSSE(payloads); ok {
return assembled, true
}
return nil, false
}
func sseDataPayloads(body []byte) []string {
payloads := make([]string, 0)
scanner := bufio.NewScanner(bytes.NewReader(body))
scanner.Buffer(make([]byte, 0, 64*1024), len(body)+1)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
if !strings.HasPrefix(line, "data:") {
continue
}
payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if payload == "" || payload == "[DONE]" {
continue
}
payloads = append(payloads, payload)
}
return payloads
}
type openAICompletionChunk struct {
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int64 `json:"created,omitempty"`
Model string `json:"model,omitempty"`
Usage json.RawMessage `json:"usage"`
Choices []struct {
Index int `json:"index"`
Text *string `json:"text"`
FinishReason *string `json:"finish_reason"`
} `json:"choices"`
}
type openAICompletionChoice struct {
index int
text strings.Builder
finishReason string
}
func assembleOpenAICompletionsSSE(payloads []string) ([]byte, bool) {
choices := map[int]*openAICompletionChoice{}
order := make([]int, 0, 1)
var id, object, model string
var created int64
var usage json.RawMessage
matched := false
for _, payload := range payloads {
var chunk openAICompletionChunk
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
return nil, false
}
if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
return nil, false
}
for _, choice := range chunk.Choices {
if choice.Text == nil {
return nil, false
}
}
matched = true
if id == "" {
id = chunk.ID
}
if object == "" {
object = chunk.Object
}
if created == 0 {
created = chunk.Created
}
if model == "" {
model = chunk.Model
}
if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
usage = chunk.Usage
}
for _, choice := range chunk.Choices {
assembled := choices[choice.Index]
if assembled == nil {
assembled = &openAICompletionChoice{index: choice.Index}
choices[choice.Index] = assembled
order = append(order, choice.Index)
}
assembled.text.WriteString(*choice.Text)
if choice.FinishReason != nil {
assembled.finishReason = *choice.FinishReason
}
}
}
if !matched {
return nil, false
}
assembled := struct {
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int64 `json:"created,omitempty"`
Model string `json:"model,omitempty"`
Usage json.RawMessage `json:"usage,omitempty"`
Choices []struct {
Index int `json:"index"`
Text string `json:"text"`
FinishReason string `json:"finish_reason,omitempty"`
} `json:"choices"`
}{ID: id, Object: object, Created: created, Model: model, Usage: usage}
for _, index := range order {
choice := choices[index]
assembled.Choices = append(assembled.Choices, struct {
Index int `json:"index"`
Text string `json:"text"`
FinishReason string `json:"finish_reason,omitempty"`
}{Index: choice.index, Text: choice.text.String(), FinishReason: choice.finishReason})
}
data, err := json.Marshal(assembled)
return data, err == nil
}
type openAIChatChunk struct {
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int64 `json:"created,omitempty"`
Model string `json:"model,omitempty"`
SystemFingerprint string `json:"system_fingerprint,omitempty"`
Usage json.RawMessage `json:"usage"`
Choices []struct {
Index int `json:"index"`
Delta struct {
Role string `json:"role,omitempty"`
Content string `json:"content,omitempty"`
ReasoningContent string `json:"reasoning_content,omitempty"`
ToolCalls []struct {
Index int `json:"index"`
ID string `json:"id,omitempty"`
Type string `json:"type,omitempty"`
Function struct {
Name string `json:"name,omitempty"`
Arguments string `json:"arguments,omitempty"`
} `json:"function,omitempty"`
} `json:"tool_calls,omitempty"`
} `json:"delta"`
FinishReason *string `json:"finish_reason"`
NativeFinishReason *string `json:"native_finish_reason"`
} `json:"choices"`
}
type openAIChatChoice struct {
index int
role string
content strings.Builder
reasoning strings.Builder
toolCalls map[int]*openAIToolCall
toolOrder []int
finishReason string
nativeFinish string
}
type openAIToolCall struct {
id string
callType string
name string
arguments strings.Builder
}
type openAIChatResponse struct {
ID string `json:"id,omitempty"`
Object string `json:"object,omitempty"`
Created int64 `json:"created,omitempty"`
Model string `json:"model,omitempty"`
SystemFingerprint string `json:"system_fingerprint,omitempty"`
Usage json.RawMessage `json:"usage,omitempty"`
Choices []openAIResponseChoice `json:"choices"`
}
type openAIResponseChoice struct {
Index int `json:"index"`
Message openAIResponseMessage `json:"message"`
FinishReason string `json:"finish_reason,omitempty"`
NativeFinishReason string `json:"native_finish_reason,omitempty"`
}
type openAIResponseMessage struct {
Role string `json:"role,omitempty"`
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content,omitempty"`
ToolCalls []openAIResponseTool `json:"tool_calls,omitempty"`
}
type openAIResponseTool struct {
ID string `json:"id,omitempty"`
Type string `json:"type,omitempty"`
Function openAIResponseToolFunction `json:"function"`
}
type openAIResponseToolFunction struct {
Name string `json:"name,omitempty"`
Arguments string `json:"arguments"`
}
func assembleOpenAIChatSSE(payloads []string) ([]byte, bool) {
choices := map[int]*openAIChatChoice{}
order := make([]int, 0, 1)
var id, object string
var created int64
var model string
var systemFingerprint string
var usage json.RawMessage
matched := false
for _, payload := range payloads {
var chunk openAIChatChunk
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
return nil, false
}
if len(chunk.Choices) == 0 && len(chunk.Usage) == 0 {
return nil, false
}
matched = true
if id == "" {
id = chunk.ID
}
if object == "" {
object = strings.TrimSuffix(chunk.Object, ".chunk")
}
if created == 0 {
created = chunk.Created
}
if model == "" {
model = chunk.Model
}
if systemFingerprint == "" {
systemFingerprint = chunk.SystemFingerprint
}
if len(chunk.Usage) > 0 && string(chunk.Usage) != "null" {
usage = chunk.Usage
}
for _, choice := range chunk.Choices {
assembled := choices[choice.Index]
if assembled == nil {
assembled = &openAIChatChoice{index: choice.Index}
choices[choice.Index] = assembled
order = append(order, choice.Index)
}
if choice.Delta.Role != "" {
assembled.role = choice.Delta.Role
}
if choice.Delta.Content != "" {
assembled.content.WriteString(choice.Delta.Content)
}
if choice.Delta.ReasoningContent != "" {
assembled.reasoning.WriteString(choice.Delta.ReasoningContent)
}
for _, toolCall := range choice.Delta.ToolCalls {
if assembled.toolCalls == nil {
assembled.toolCalls = map[int]*openAIToolCall{}
}
assembledTool := assembled.toolCalls[toolCall.Index]
if assembledTool == nil {
assembledTool = &openAIToolCall{}
assembled.toolCalls[toolCall.Index] = assembledTool
assembled.toolOrder = append(assembled.toolOrder, toolCall.Index)
}
if toolCall.ID != "" {
assembledTool.id = toolCall.ID
}
if toolCall.Type != "" {
assembledTool.callType = toolCall.Type
}
if toolCall.Function.Name != "" {
assembledTool.name = toolCall.Function.Name
}
if toolCall.Function.Arguments != "" {
assembledTool.arguments.WriteString(toolCall.Function.Arguments)
}
}
if choice.FinishReason != nil {
assembled.finishReason = *choice.FinishReason
}
if choice.NativeFinishReason != nil {
assembled.nativeFinish = *choice.NativeFinishReason
}
}
}
if !matched {
return nil, false
}
assembled := openAIChatResponse{ID: id, Object: object, Created: created, Model: model, SystemFingerprint: systemFingerprint, Usage: usage}
for _, index := range order {
choice := choices[index]
out := openAIResponseChoice{Index: choice.index, FinishReason: choice.finishReason, NativeFinishReason: choice.nativeFinish}
out.Message.Role = choice.role
out.Message.Content = choice.content.String()
out.Message.ReasoningContent = choice.reasoning.String()
for _, toolIndex := range choice.toolOrder {
toolCall := choice.toolCalls[toolIndex]
outTool := openAIResponseTool{ID: toolCall.id, Type: toolCall.callType}
outTool.Function.Name = toolCall.name
outTool.Function.Arguments = toolCall.arguments.String()
out.Message.ToolCalls = append(out.Message.ToolCalls, outTool)
}
assembled.Choices = append(assembled.Choices, out)
}
data, err := json.Marshal(assembled)
return data, err == nil
}
type openAIResponsesEvent struct {
Type string `json:"type"`
Response json.RawMessage `json:"response"`
}
func assembleOpenAIResponsesSSE(payloads []string) ([]byte, bool) {
var completed json.RawMessage
matched := false
for _, payload := range payloads {
var event openAIResponsesEvent
if err := json.Unmarshal([]byte(payload), &event); err != nil {
return nil, false
}
if !strings.HasPrefix(event.Type, "response.") {
return nil, false
}
matched = true
if event.Type == "response.completed" && len(event.Response) > 0 {
completed = append(json.RawMessage(nil), event.Response...)
}
}
if !matched || len(completed) == 0 {
return nil, false
}
return completed, true
}
type anthropicEvent struct {
Type string `json:"type"`
Message *struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Model string `json:"model,omitempty"`
StopReason string `json:"stop_reason"`
StopSequence string `json:"stop_sequence"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
} `json:"usage"`
} `json:"message"`
Index int `json:"index"`
ContentBlock *struct {
Type string `json:"type"`
Text string `json:"text"`
ID string `json:"id"`
Name string `json:"name"`
} `json:"content_block"`
Delta *struct {
StopReason string `json:"stop_reason"`
StopSequence string `json:"stop_sequence"`
Text string `json:"text"`
PartialJSON string `json:"partial_json"`
} `json:"delta"`
Usage *struct {
OutputTokens int `json:"output_tokens"`
} `json:"usage"`
}
type anthropicContentBlock struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
ID string `json:"id,omitempty"`
Name string `json:"name,omitempty"`
Input map[string]any `json:"input,omitempty"`
partialJSON strings.Builder
}
func assembleAnthropicSSE(payloads []string) ([]byte, bool) {
var assembled struct {
ID string `json:"id"`
Type string `json:"type"`
Role string `json:"role"`
Content []anthropicContentBlock `json:"content"`
Model string `json:"model,omitempty"`
StopReason string `json:"stop_reason,omitempty"`
StopSequence string `json:"stop_sequence,omitempty"`
Usage struct {
InputTokens int `json:"input_tokens"`
OutputTokens int `json:"output_tokens"`
} `json:"usage"`
}
contents := map[int]*anthropicContentBlock{}
order := make([]int, 0, 1)
matched := false
for _, payload := range payloads {
var event anthropicEvent
if err := json.Unmarshal([]byte(payload), &event); err != nil {
return nil, false
}
if !strings.HasPrefix(event.Type, "message_") && !strings.HasPrefix(event.Type, "content_block_") {
return nil, false
}
matched = true
if event.Message != nil {
assembled.ID = event.Message.ID
assembled.Type = event.Message.Type
assembled.Role = event.Message.Role
assembled.Model = event.Message.Model
assembled.StopReason = event.Message.StopReason
assembled.StopSequence = event.Message.StopSequence
assembled.Usage.InputTokens = event.Message.Usage.InputTokens
assembled.Usage.OutputTokens = event.Message.Usage.OutputTokens
}
if event.ContentBlock != nil {
block := contents[event.Index]
if block == nil {
block = &anthropicContentBlock{Type: event.ContentBlock.Type, ID: event.ContentBlock.ID, Name: event.ContentBlock.Name}
contents[event.Index] = block
order = append(order, event.Index)
}
if event.ContentBlock.ID != "" {
block.ID = event.ContentBlock.ID
}
if event.ContentBlock.Name != "" {
block.Name = event.ContentBlock.Name
}
block.Text += event.ContentBlock.Text
}
if event.Delta != nil {
if event.Delta.Text != "" {
block := contents[event.Index]
if block == nil {
block = &anthropicContentBlock{Type: "text"}
contents[event.Index] = block
order = append(order, event.Index)
}
block.Text += event.Delta.Text
}
if event.Delta.PartialJSON != "" {
block := contents[event.Index]
if block == nil {
block = &anthropicContentBlock{Type: "tool_use"}
contents[event.Index] = block
order = append(order, event.Index)
}
block.partialJSON.WriteString(event.Delta.PartialJSON)
}
if event.Delta.StopReason != "" {
assembled.StopReason = event.Delta.StopReason
}
if event.Delta.StopSequence != "" {
assembled.StopSequence = event.Delta.StopSequence
}
}
if event.Usage != nil {
assembled.Usage.OutputTokens = event.Usage.OutputTokens
}
}
if !matched || assembled.Type == "" {
return nil, false
}
for _, index := range order {
block := contents[index]
if block.partialJSON.Len() > 0 {
var input map[string]any
if err := json.Unmarshal([]byte(block.partialJSON.String()), &input); err != nil {
return nil, false
}
block.Input = input
}
assembled.Content = append(assembled.Content, *block)
}
data, err := json.Marshal(assembled)
return data, err == nil
}
type geminiChunk struct {
Raw map[string]json.RawMessage `json:"-"`
Candidates []struct {
Raw map[string]json.RawMessage `json:"-"`
Content struct {
Parts []struct {
Text string `json:"text"`
} `json:"parts"`
Role string `json:"role"`
} `json:"content"`
FinishReason string `json:"finishReason"`
Index int `json:"index"`
} `json:"candidates"`
UsageMetadata json.RawMessage `json:"usageMetadata"`
}
func (g *geminiChunk) UnmarshalJSON(data []byte) error {
type alias geminiChunk
var decoded alias
if err := json.Unmarshal(data, &decoded); err != nil {
return err
}
var raw map[string]json.RawMessage
if err := json.Unmarshal(data, &raw); err != nil {
return err
}
*g = geminiChunk(decoded)
g.Raw = raw
if candidatesRaw, ok := raw["candidates"]; ok {
var rawCandidates []map[string]json.RawMessage
if err := json.Unmarshal(candidatesRaw, &rawCandidates); err != nil {
return err
}
for i := range g.Candidates {
if i < len(rawCandidates) {
g.Candidates[i].Raw = rawCandidates[i]
}
}
}
return nil
}
type geminiCandidate struct {
raw map[string]json.RawMessage
index int
role string
parts []*geminiPart
}
type geminiPart struct {
raw map[string]json.RawMessage
text strings.Builder
kind string
}
func assembleGeminiSSE(payloads []string) ([]byte, bool) {
candidates := map[int]*geminiCandidate{}
order := make([]int, 0, 1)
var usage json.RawMessage
matched := false
for _, payload := range payloads {
var chunk geminiChunk
if err := json.Unmarshal([]byte(payload), &chunk); err != nil {
return nil, false
}
if len(chunk.Candidates) == 0 {
return nil, false
}
matched = true
for _, candidate := range chunk.Candidates {
assembled := candidates[candidate.Index]
if assembled == nil {
assembled = &geminiCandidate{index: candidate.Index}
candidates[candidate.Index] = assembled
order = append(order, candidate.Index)
}
if candidate.Raw != nil {
assembled.raw = cloneRawMap(candidate.Raw)
}
if candidate.Content.Role != "" {
assembled.role = candidate.Content.Role
}
if len(candidate.Content.Parts) > 0 {
var contentRaw struct {
Parts []map[string]json.RawMessage `json:"parts"`
}
if rawContent, ok := candidate.Raw["content"]; ok {
_ = json.Unmarshal(rawContent, &contentRaw)
}
mergeByIndex := len(assembled.parts) == len(candidate.Content.Parts)
if mergeByIndex {
for i := range candidate.Content.Parts {
if i >= len(contentRaw.Parts) || assembled.parts[i].kind != geminiPartKind(contentRaw.Parts[i]) {
mergeByIndex = false
break
}
}
}
for i, part := range candidate.Content.Parts {
target := i
if !mergeByIndex {
target = len(assembled.parts)
assembled.parts = append(assembled.parts, &geminiPart{})
}
assembled.parts[target].text.WriteString(part.Text)
if i < len(contentRaw.Parts) {
assembled.parts[target].kind = geminiPartKind(contentRaw.Parts[i])
assembled.parts[target].raw = mergeRawMap(assembled.parts[target].raw, contentRaw.Parts[i])
}
}
}
}
if len(chunk.UsageMetadata) > 0 {
usage = chunk.UsageMetadata
}
}
if !matched {
return nil, false
}
assembled := map[string]any{"candidates": make([]any, 0, len(order))}
for _, index := range order {
candidate := candidates[index]
candidateRaw := cloneRawMap(candidate.raw)
candidateRaw["index"] = mustJSON(candidate.index)
parts := make([]any, 0, len(candidate.parts))
for _, part := range candidate.parts {
partRaw := cloneRawMap(part.raw)
if part.text.Len() > 0 || len(partRaw) == 0 {
partRaw["text"] = mustJSON(part.text.String())
}
parts = append(parts, rawMapToMap(partRaw))
}
contentRaw := map[string]any{"role": candidate.role, "parts": parts}
candidateRaw["content"] = mustJSON(contentRaw)
assembled["candidates"] = append(assembled["candidates"].([]any), rawMapToMap(candidateRaw))
}
if len(usage) > 0 {
assembled["usageMetadata"] = usage
}
data, err := json.Marshal(assembled)
return data, err == nil
}
func geminiPartKind(part map[string]json.RawMessage) string {
for _, key := range []string{"functionCall", "functionResponse", "inlineData", "fileData", "executableCode", "codeExecutionResult", "text"} {
if _, ok := part[key]; ok {
return key
}
}
return "unknown"
}
func mergeRawMap(dst, src map[string]json.RawMessage) map[string]json.RawMessage {
if dst == nil {
dst = make(map[string]json.RawMessage, len(src))
}
for key, value := range src {
if key == "text" {
continue
}
dst[key] = append(json.RawMessage(nil), value...)
}
return dst
}
func cloneRawMap(in map[string]json.RawMessage) map[string]json.RawMessage {
out := make(map[string]json.RawMessage, len(in))
for k, v := range in {
out[k] = append(json.RawMessage(nil), v...)
}
return out
}
func mustJSON(v any) json.RawMessage {
data, err := json.Marshal(v)
if err != nil {
return nil
}
return data
}
func rawMapToMap(in map[string]json.RawMessage) map[string]any {
out := make(map[string]any, len(in))
for k, v := range in {
var decoded any
if err := json.Unmarshal(v, &decoded); err != nil {
out[k] = string(v)
continue
}
out[k] = decoded
}
return out
}