292 lines
8.2 KiB
Go
292 lines
8.2 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"net/http"
|
|
"strings"
|
|
"time"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
iop "iop/proto/gen/iop"
|
|
)
|
|
|
|
var anthropicResponseHeaderAllowlist = map[string]struct{}{
|
|
"Cache-Control": {},
|
|
"Content-Length": {},
|
|
"Content-Type": {},
|
|
"Request-Id": {},
|
|
"Retry-After": {},
|
|
"X-Request-Id": {},
|
|
"X-Robots-Tag": {},
|
|
}
|
|
|
|
func (s *Server) writeAnthropicNativeTunnelResponse(w http.ResponseWriter, r *http.Request, handle edgeservice.ProviderTunnelResult, publicModelID string) {
|
|
frames := handle.Stream().Frames
|
|
if frames == nil {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel is unavailable")
|
|
return
|
|
}
|
|
flusher, _ := w.(http.Flusher)
|
|
timer := time.NewTimer(handle.WaitTimeout())
|
|
defer timer.Stop()
|
|
wroteHeader := false
|
|
receivedResponseStart := false
|
|
responseStatus := http.StatusOK
|
|
responseStreaming := false
|
|
rewriteResponse := strings.TrimSpace(publicModelID) != ""
|
|
var responseBody []byte
|
|
var streamRewriter *anthropicNativeModelRewriter
|
|
|
|
for {
|
|
select {
|
|
case <-r.Context().Done():
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), r.Context().Err())
|
|
return
|
|
case <-timer.C:
|
|
s.cancelRunOnHTTPGiveUp(handle.Dispatch(), errRunTimedOut)
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider response timed out")
|
|
}
|
|
return
|
|
case frame, ok := <-frames:
|
|
if !ok {
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel closed before a response")
|
|
}
|
|
return
|
|
}
|
|
switch frame.GetKind() {
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_RESPONSE_START:
|
|
if receivedResponseStart || wroteHeader {
|
|
continue
|
|
}
|
|
receivedResponseStart = true
|
|
copyAnthropicResponseHeaders(w.Header(), frame.GetHeaders())
|
|
status := int(frame.GetStatusCode())
|
|
if status == 0 {
|
|
status = http.StatusOK
|
|
}
|
|
responseStatus = status
|
|
if rewriteResponse && status >= http.StatusOK && status < http.StatusMultipleChoices {
|
|
w.Header().Del("Content-Length")
|
|
responseStreaming = strings.Contains(strings.ToLower(w.Header().Get("Content-Type")), "text/event-stream")
|
|
if responseStreaming {
|
|
streamRewriter = newAnthropicNativeModelRewriter(publicModelID)
|
|
w.WriteHeader(status)
|
|
wroteHeader = true
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
continue
|
|
}
|
|
w.WriteHeader(status)
|
|
wroteHeader = true
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_BODY:
|
|
if len(frame.GetBody()) == 0 {
|
|
continue
|
|
}
|
|
if rewriteResponse && receivedResponseStart && responseStatus >= http.StatusOK && responseStatus < http.StatusMultipleChoices {
|
|
if responseStreaming {
|
|
if err := writeAnthropicNativeBody(w, streamRewriter.Append(frame.GetBody()), flusher); err != nil {
|
|
s.sendCancelRun(handle.Dispatch())
|
|
return
|
|
}
|
|
} else {
|
|
responseBody = append(responseBody, frame.GetBody()...)
|
|
}
|
|
continue
|
|
}
|
|
if !wroteHeader {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
w.WriteHeader(http.StatusOK)
|
|
wroteHeader = true
|
|
}
|
|
if err := writeAnthropicNativeBody(w, frame.GetBody(), flusher); err != nil {
|
|
s.sendCancelRun(handle.Dispatch())
|
|
return
|
|
}
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_ERROR:
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel failed")
|
|
} else if strings.Contains(strings.ToLower(w.Header().Get("Content-Type")), "text/event-stream") {
|
|
_ = writeAnthropicSSEEvent(w, "error", anthropicErrorResponse{
|
|
Type: "error", Error: errorBody{Type: "api_error", Message: "provider tunnel failed"},
|
|
})
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_END:
|
|
if rewriteResponse && receivedResponseStart && responseStatus >= http.StatusOK && responseStatus < http.StatusMultipleChoices {
|
|
if responseStreaming {
|
|
if err := writeAnthropicNativeBody(w, streamRewriter.Flush(), flusher); err != nil {
|
|
s.sendCancelRun(handle.Dispatch())
|
|
}
|
|
return
|
|
}
|
|
if !wroteHeader {
|
|
w.WriteHeader(responseStatus)
|
|
wroteHeader = true
|
|
}
|
|
if err := writeAnthropicNativeBody(w, rewriteProviderJSONModel(responseBody, publicModelID), flusher); err != nil {
|
|
s.sendCancelRun(handle.Dispatch())
|
|
}
|
|
return
|
|
}
|
|
if !wroteHeader {
|
|
writeAnthropicError(w, http.StatusBadGateway, "api_error", "provider tunnel ended before a response")
|
|
}
|
|
return
|
|
case iop.ProviderTunnelFrameKind_PROVIDER_TUNNEL_FRAME_KIND_USAGE:
|
|
continue
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func writeAnthropicNativeBody(w http.ResponseWriter, body []byte, flusher http.Flusher) error {
|
|
if len(body) == 0 {
|
|
return nil
|
|
}
|
|
if _, err := w.Write(body); err != nil {
|
|
return err
|
|
}
|
|
if flusher != nil {
|
|
flusher.Flush()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
type anthropicNativeModelRewriter struct {
|
|
model string
|
|
pending []byte
|
|
messageStart bool
|
|
}
|
|
|
|
func newAnthropicNativeModelRewriter(model string) *anthropicNativeModelRewriter {
|
|
model = strings.TrimSpace(model)
|
|
if model == "" {
|
|
return nil
|
|
}
|
|
return &anthropicNativeModelRewriter{model: model}
|
|
}
|
|
|
|
func (r *anthropicNativeModelRewriter) Append(chunk []byte) []byte {
|
|
if r == nil || len(chunk) == 0 {
|
|
return chunk
|
|
}
|
|
r.pending = append(r.pending, chunk...)
|
|
var out bytes.Buffer
|
|
for {
|
|
index := bytes.IndexByte(r.pending, '\n')
|
|
if index < 0 {
|
|
break
|
|
}
|
|
line := r.pending[:index+1]
|
|
out.Write(r.rewriteLine(line))
|
|
r.pending = r.pending[index+1:]
|
|
}
|
|
return out.Bytes()
|
|
}
|
|
|
|
func (r *anthropicNativeModelRewriter) Flush() []byte {
|
|
if r == nil || len(r.pending) == 0 {
|
|
return nil
|
|
}
|
|
pending := r.pending
|
|
r.pending = nil
|
|
return r.rewriteLine(pending)
|
|
}
|
|
|
|
func (r *anthropicNativeModelRewriter) rewriteLine(line []byte) []byte {
|
|
body, ending := splitLineEnding(line)
|
|
prefix, payload, ok := bytes.Cut(body, []byte(":"))
|
|
if !ok {
|
|
return line
|
|
}
|
|
|
|
switch strings.TrimSpace(string(prefix)) {
|
|
case "event":
|
|
r.messageStart = strings.TrimSpace(string(payload)) == "message_start"
|
|
return line
|
|
case "data":
|
|
if !r.messageStart {
|
|
return line
|
|
}
|
|
r.messageStart = false
|
|
default:
|
|
return line
|
|
}
|
|
|
|
leading := len(payload) - len(bytes.TrimLeft(payload, " \t"))
|
|
trailing := len(payload) - len(bytes.TrimRight(payload, " \t"))
|
|
if leading+trailing >= len(payload) {
|
|
return line
|
|
}
|
|
rewritten := rewriteAnthropicMessageStartModel(payload[leading:len(payload)-trailing], r.model)
|
|
if bytes.Equal(rewritten, payload[leading:len(payload)-trailing]) {
|
|
return line
|
|
}
|
|
out := make([]byte, 0, len(body)+len(rewritten)-len(payload)+len(ending))
|
|
out = append(out, body[:len(prefix)+1+leading]...)
|
|
out = append(out, rewritten...)
|
|
out = append(out, payload[len(payload)-trailing:]...)
|
|
out = append(out, ending...)
|
|
return out
|
|
}
|
|
|
|
func rewriteAnthropicMessageStartModel(body []byte, model string) []byte {
|
|
modelJSON, err := json.Marshal(model)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
fields, _, err := scanTopLevelJSONObject(body)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
for _, field := range fields {
|
|
if field.name != "message" {
|
|
continue
|
|
}
|
|
message := body[field.valueFrom:field.valueTo]
|
|
messageFields, _, err := scanTopLevelJSONObject(message)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
for _, messageField := range messageFields {
|
|
if messageField.name != "model" {
|
|
continue
|
|
}
|
|
plan, err := planTopLevelJSONPatches(message, []topLevelJSONPatch{{name: "model", value: modelJSON}})
|
|
if err != nil {
|
|
return body
|
|
}
|
|
return topLevelJSONPatchPlan{
|
|
body: body,
|
|
edits: []jsonByteEdit{{
|
|
from: field.valueFrom, to: field.valueTo, replacement: plan.apply(),
|
|
}},
|
|
outputSize: len(body) + plan.outputSize - len(message),
|
|
}.apply()
|
|
}
|
|
}
|
|
return body
|
|
}
|
|
|
|
func copyAnthropicResponseHeaders(dst http.Header, headers map[string]string) {
|
|
for key, value := range headers {
|
|
canonical := http.CanonicalHeaderKey(key)
|
|
_, exact := anthropicResponseHeaderAllowlist[canonical]
|
|
lower := strings.ToLower(key)
|
|
if !exact && !strings.HasPrefix(lower, "anthropic-ratelimit-") && !strings.HasPrefix(lower, "ratelimit-") {
|
|
continue
|
|
}
|
|
dst.Set(canonical, value)
|
|
}
|
|
}
|