446 lines
12 KiB
Go
446 lines
12 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
type openAIChatStreamChunk struct {
|
|
ID string `json:"id"`
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
Reasoning string `json:"reasoning"`
|
|
ToolCalls []struct {
|
|
Index int `json:"index"`
|
|
ID string `json:"id"`
|
|
Function struct {
|
|
Name string `json:"name"`
|
|
Arguments string `json:"arguments"`
|
|
} `json:"function"`
|
|
} `json:"tool_calls"`
|
|
} `json:"delta"`
|
|
FinishReason *string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
Usage *openAIChatBridgeUsage `json:"usage,omitempty"`
|
|
Error *struct {
|
|
Type string `json:"type"`
|
|
Message string `json:"message"`
|
|
} `json:"error,omitempty"`
|
|
}
|
|
|
|
type anthropicBridgeToolState struct {
|
|
id string
|
|
name string
|
|
arguments strings.Builder
|
|
}
|
|
|
|
type anthropicBridgeStream struct {
|
|
w http.ResponseWriter
|
|
model string
|
|
id string
|
|
started bool
|
|
stopped bool
|
|
nextBlock int
|
|
openBlock bool
|
|
openKind string
|
|
finish *string
|
|
usage anthropicUsage
|
|
tools map[int]*anthropicBridgeToolState
|
|
pendingSSE []byte
|
|
}
|
|
|
|
func newAnthropicBridgeStream(w http.ResponseWriter, model string) *anthropicBridgeStream {
|
|
return &anthropicBridgeStream{w: w, model: model, tools: make(map[int]*anthropicBridgeToolState)}
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) 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 *anthropicBridgeStream) consumeSSEEvent(event []byte) error {
|
|
var dataLines [][]byte
|
|
for _, line := range bytes.Split(event, []byte("\n")) {
|
|
line = bytes.TrimSpace(line)
|
|
if len(line) == 0 || line[0] == ':' {
|
|
continue
|
|
}
|
|
if bytes.HasPrefix(line, []byte("data:")) {
|
|
dataLines = append(dataLines, bytes.TrimSpace(bytes.TrimPrefix(line, []byte("data:"))))
|
|
}
|
|
}
|
|
if len(dataLines) == 0 {
|
|
return nil
|
|
}
|
|
payload := bytes.Join(dataLines, []byte("\n"))
|
|
if bytes.Equal(payload, []byte("[DONE]")) {
|
|
return s.Finish()
|
|
}
|
|
var chunk openAIChatStreamChunk
|
|
if err := json.Unmarshal(payload, &chunk); err != nil {
|
|
return fmt.Errorf("decode Chat SSE event: %w", err)
|
|
}
|
|
if chunk.Error != nil {
|
|
message := strings.TrimSpace(chunk.Error.Message)
|
|
if message == "" {
|
|
message = "upstream provider error"
|
|
}
|
|
return s.Error(chunk.Error.Type, message)
|
|
}
|
|
if chunk.ID != "" && s.id == "" {
|
|
s.id = chunk.ID
|
|
}
|
|
if chunk.Usage != nil {
|
|
s.usage.InputTokens = chunk.Usage.PromptTokens
|
|
s.usage.OutputTokens = chunk.Usage.CompletionTokens
|
|
s.usage.CacheReadInputTokens = chunk.Usage.PromptDetails.CachedTokens
|
|
}
|
|
if err := s.start(); err != nil {
|
|
return err
|
|
}
|
|
for _, choice := range chunk.Choices {
|
|
reasoning := choice.Delta.ReasoningContent
|
|
if reasoning == "" {
|
|
reasoning = choice.Delta.Reasoning
|
|
}
|
|
if reasoning != "" {
|
|
if err := s.delta("thinking", reasoning); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
if choice.Delta.Content != "" {
|
|
if err := s.delta("text", choice.Delta.Content); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
for _, delta := range choice.Delta.ToolCalls {
|
|
state := s.tools[delta.Index]
|
|
if state == nil {
|
|
state = &anthropicBridgeToolState{}
|
|
s.tools[delta.Index] = state
|
|
}
|
|
if delta.ID != "" {
|
|
state.id = delta.ID
|
|
}
|
|
if delta.Function.Name != "" {
|
|
state.name = delta.Function.Name
|
|
}
|
|
state.arguments.WriteString(delta.Function.Arguments)
|
|
}
|
|
if choice.FinishReason != nil {
|
|
s.finish = choice.FinishReason
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) start() error {
|
|
if s.started {
|
|
return nil
|
|
}
|
|
s.started = true
|
|
id := s.id
|
|
if id == "" {
|
|
id = "msg_iop"
|
|
}
|
|
return writeAnthropicSSEEvent(s.w, "message_start", map[string]any{
|
|
"type": "message_start",
|
|
"message": anthropicMessageResponse{
|
|
ID: id, Type: "message", Role: "assistant", Model: s.model,
|
|
Content: []map[string]any{}, StopReason: nil, StopSequence: nil,
|
|
Usage: anthropicUsage{InputTokens: s.usage.InputTokens},
|
|
},
|
|
})
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) delta(kind, value string) error {
|
|
if !s.openBlock || s.openKind != kind {
|
|
if err := s.closeBlock(); err != nil {
|
|
return err
|
|
}
|
|
block := map[string]any{"type": kind}
|
|
if kind == "thinking" {
|
|
block["thinking"] = ""
|
|
block["signature"] = ""
|
|
} else {
|
|
block["text"] = ""
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "content_block_start", map[string]any{
|
|
"type": "content_block_start", "index": s.nextBlock, "content_block": block,
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.openBlock, s.openKind = true, kind
|
|
}
|
|
deltaType, field := "text_delta", "text"
|
|
if kind == "thinking" {
|
|
deltaType, field = "thinking_delta", "thinking"
|
|
}
|
|
return writeAnthropicSSEEvent(s.w, "content_block_delta", map[string]any{
|
|
"type": "content_block_delta", "index": s.nextBlock,
|
|
"delta": map[string]any{"type": deltaType, field: value},
|
|
})
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) closeBlock() error {
|
|
if !s.openBlock {
|
|
return nil
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "content_block_stop", map[string]any{
|
|
"type": "content_block_stop", "index": s.nextBlock,
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.nextBlock++
|
|
s.openBlock = false
|
|
s.openKind = ""
|
|
return nil
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) 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)
|
|
for _, index := range indices {
|
|
tool := s.tools[index]
|
|
arguments := tool.arguments.String()
|
|
if tool.id == "" || tool.name == "" || !json.Valid([]byte(arguments)) {
|
|
return fmt.Errorf("Chat stream tool call has invalid id, name, or arguments")
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "content_block_start", map[string]any{
|
|
"type": "content_block_start", "index": s.nextBlock,
|
|
"content_block": map[string]any{"type": "tool_use", "id": tool.id, "name": tool.name, "input": map[string]any{}},
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "content_block_delta", map[string]any{
|
|
"type": "content_block_delta", "index": s.nextBlock,
|
|
"delta": map[string]any{"type": "input_json_delta", "partial_json": arguments},
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "content_block_stop", map[string]any{
|
|
"type": "content_block_stop", "index": s.nextBlock,
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
s.nextBlock++
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) Finish() error {
|
|
if s.stopped {
|
|
return nil
|
|
}
|
|
if err := s.start(); err != nil {
|
|
return err
|
|
}
|
|
if err := s.closeBlock(); err != nil {
|
|
return err
|
|
}
|
|
if err := s.emitTools(); err != nil {
|
|
return err
|
|
}
|
|
stopReason, err := anthropicStopReason(s.finish)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if stopReason == nil {
|
|
fallback := "end_turn"
|
|
stopReason = &fallback
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "message_delta", map[string]any{
|
|
"type": "message_delta", "delta": map[string]any{"stop_reason": *stopReason, "stop_sequence": nil},
|
|
"usage": s.usage,
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
if err := writeAnthropicSSEEvent(s.w, "message_stop", map[string]any{"type": "message_stop"}); err != nil {
|
|
return err
|
|
}
|
|
s.stopped = true
|
|
return nil
|
|
}
|
|
|
|
func (s *anthropicBridgeStream) Error(errorType, message string) error {
|
|
if s.stopped {
|
|
return nil
|
|
}
|
|
if strings.TrimSpace(errorType) == "" {
|
|
errorType = "api_error"
|
|
}
|
|
s.stopped = true
|
|
return writeAnthropicSSEEvent(s.w, "error", anthropicErrorResponse{
|
|
Type: "error", Error: errorBody{Type: errorType, Message: message},
|
|
})
|
|
}
|
|
|
|
func writeAnthropicSSEEvent(w http.ResponseWriter, event string, payload any) error {
|
|
encoded, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, encoded); err != nil {
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Server) writeAnthropicChatBridgeResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult, envelope anthropicRequestEnvelope) {
|
|
frames := handle.Stream().Frames
|
|
if frames == nil {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel is unavailable")
|
|
return
|
|
}
|
|
timer := time.NewTimer(handle.WaitTimeout())
|
|
defer timer.Stop()
|
|
status := http.StatusOK
|
|
headers := make(map[string]string)
|
|
started := false
|
|
committed := false
|
|
var body []byte
|
|
var stream *anthropicBridgeStream
|
|
if envelope.Stream {
|
|
stream = newAnthropicBridgeStream(w, envelope.Model)
|
|
}
|
|
flush := func() {
|
|
if flusher, ok := w.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
commitStream := func() {
|
|
if committed {
|
|
return
|
|
}
|
|
copyAnthropicResponseHeaders(w.Header(), headers)
|
|
w.Header().Del("Content-Length")
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.WriteHeader(status)
|
|
committed = true
|
|
flush()
|
|
}
|
|
|
|
for {
|
|
select {
|
|
case <-r.Context().Done():
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
|
|
return
|
|
case <-timer.C:
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
|
|
if !committed {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider response timed out")
|
|
} else {
|
|
_ = stream.Error("api_error", "provider response timed out")
|
|
flush()
|
|
}
|
|
return
|
|
case frame, ok := <-frames:
|
|
if !ok {
|
|
if !committed {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel closed before a response")
|
|
} else {
|
|
_ = stream.Error("api_error", "provider tunnel closed before a response")
|
|
flush()
|
|
}
|
|
return
|
|
}
|
|
switch frame.GetKind() {
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
|
|
if started {
|
|
continue
|
|
}
|
|
started = true
|
|
status = int(frame.GetStatusCode())
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
for key, value := range frame.GetHeaders() {
|
|
headers[key] = value
|
|
}
|
|
if envelope.Stream && status < http.StatusBadRequest {
|
|
commitStream()
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
|
|
if envelope.Stream && status < http.StatusBadRequest {
|
|
commitStream()
|
|
if err := stream.Feed(frame.GetBody()); err != nil {
|
|
_ = stream.Error("api_error", "upstream stream could not be translated")
|
|
flush()
|
|
return
|
|
}
|
|
flush()
|
|
} else {
|
|
body = append(body, frame.GetBody()...)
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
|
|
if committed {
|
|
_ = stream.Error("api_error", "provider tunnel failed")
|
|
flush()
|
|
} else {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel failed")
|
|
}
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
|
|
if envelope.Stream && status < http.StatusBadRequest {
|
|
commitStream()
|
|
if err := stream.Finish(); err != nil {
|
|
_ = stream.Error("api_error", "upstream stream could not be translated")
|
|
}
|
|
flush()
|
|
return
|
|
}
|
|
copyAnthropicResponseHeaders(w.Header(), headers)
|
|
w.Header().Del("Content-Length")
|
|
if status >= http.StatusBadRequest {
|
|
writeJSON(w, status, convertChatErrorToAnthropic(body))
|
|
return
|
|
}
|
|
converted, err := convertChatResponseToAnthropic(body, envelope.Model)
|
|
if err != nil {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "upstream response could not be translated")
|
|
return
|
|
}
|
|
writeJSON(w, http.StatusOK, converted)
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
}
|