494 lines
14 KiB
Go
494 lines
14 KiB
Go
package openai
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"fmt"
|
|
"strings"
|
|
"sync"
|
|
|
|
edgeservice "iop/apps/edge/internal/service"
|
|
"iop/packages/go/config"
|
|
)
|
|
|
|
type openAIProviderBodyBuilder struct {
|
|
mu sync.Mutex
|
|
build func(string) (*openAIRebuiltLease, error)
|
|
lease *openAIRebuiltLease
|
|
built bool
|
|
closed bool
|
|
}
|
|
|
|
func newOpenAIProviderBodyBuilder(build func(string) (*openAIRebuiltLease, error)) *openAIProviderBodyBuilder {
|
|
return &openAIProviderBodyBuilder{build: build}
|
|
}
|
|
|
|
func (b *openAIProviderBodyBuilder) BuildBody(target string) ([]byte, error) {
|
|
if b == nil {
|
|
return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
|
|
}
|
|
b.mu.Lock()
|
|
if b.closed || b.built || b.build == nil {
|
|
b.mu.Unlock()
|
|
return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
|
|
}
|
|
b.built = true
|
|
build := b.build
|
|
b.mu.Unlock()
|
|
|
|
lease, err := build(target)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
b.mu.Lock()
|
|
if b.closed {
|
|
b.mu.Unlock()
|
|
lease.release()
|
|
return nil, fmt.Errorf("OpenAI provider body builder is unavailable")
|
|
}
|
|
b.lease = lease
|
|
b.mu.Unlock()
|
|
body, err := lease.body()
|
|
if err != nil {
|
|
b.Close()
|
|
return nil, err
|
|
}
|
|
return body, nil
|
|
}
|
|
|
|
func (b *openAIProviderBodyBuilder) Close() {
|
|
if b == nil {
|
|
return
|
|
}
|
|
b.mu.Lock()
|
|
b.closed = true
|
|
b.build = nil
|
|
lease := b.lease
|
|
b.lease = nil
|
|
b.mu.Unlock()
|
|
lease.release()
|
|
}
|
|
|
|
type providerModelRewriter struct {
|
|
streaming bool
|
|
model string
|
|
toolCallWire string
|
|
pending []byte
|
|
}
|
|
|
|
func newProviderModelRewriter(streaming bool, model string) *providerModelRewriter {
|
|
return newProviderModelRewriterWithToolCallWire(streaming, model, "")
|
|
}
|
|
|
|
func newProviderModelRewriterForDispatch(streaming bool, model string, dispatch edgeservice.RunDispatch) *providerModelRewriter {
|
|
return newProviderModelRewriterWithToolCallWire(streaming, model, dispatch.ProfileToolCallWire)
|
|
}
|
|
|
|
func newProviderModelRewriterWithToolCallWire(streaming bool, model, toolCallWire string) *providerModelRewriter {
|
|
model = strings.TrimSpace(model)
|
|
toolCallWire = strings.TrimSpace(toolCallWire)
|
|
if model == "" && toolCallWire == "" {
|
|
return nil
|
|
}
|
|
return &providerModelRewriter{streaming: streaming, model: model, toolCallWire: toolCallWire}
|
|
}
|
|
|
|
func (r *providerModelRewriter) setToolCallWire(toolCallWire string) {
|
|
if r != nil {
|
|
r.toolCallWire = strings.TrimSpace(toolCallWire)
|
|
}
|
|
}
|
|
|
|
func (r *providerModelRewriter) AppendStream(chunk []byte) []byte {
|
|
if r == nil || len(chunk) == 0 {
|
|
return chunk
|
|
}
|
|
if !r.streaming {
|
|
r.pending = append(r.pending, chunk...)
|
|
return nil
|
|
}
|
|
r.pending = append(r.pending, chunk...)
|
|
var out bytes.Buffer
|
|
for {
|
|
idx := bytes.IndexByte(r.pending, '\n')
|
|
if idx < 0 {
|
|
break
|
|
}
|
|
line := r.pending[:idx+1]
|
|
out.Write(rewriteProviderSSELine(line, r.model, r.toolCallWire))
|
|
r.pending = r.pending[idx+1:]
|
|
}
|
|
return out.Bytes()
|
|
}
|
|
|
|
func (r *providerModelRewriter) FlushStream() []byte {
|
|
if r == nil || len(r.pending) == 0 {
|
|
return nil
|
|
}
|
|
pending := r.pending
|
|
r.pending = nil
|
|
if !r.streaming {
|
|
return r.RewriteComplete(pending)
|
|
}
|
|
return rewriteProviderSSELine(pending, r.model, r.toolCallWire)
|
|
}
|
|
|
|
func (r *providerModelRewriter) RewriteComplete(body []byte) []byte {
|
|
if r == nil || len(body) == 0 {
|
|
return body
|
|
}
|
|
return rewriteProviderJSONResponse(body, r.model, r.toolCallWire)
|
|
}
|
|
|
|
func rewriteProviderSSEModelLine(line []byte, model string) []byte {
|
|
return rewriteProviderSSELine(line, model, "")
|
|
}
|
|
|
|
func rewriteProviderSSELine(line []byte, model, toolCallWire string) []byte {
|
|
body, ending := splitLineEnding(line)
|
|
prefix, payload, ok := bytes.Cut(body, []byte(":"))
|
|
if !ok || strings.TrimSpace(string(prefix)) != "data" {
|
|
return line
|
|
}
|
|
payload = bytes.TrimSpace(payload)
|
|
if len(payload) == 0 || bytes.Equal(payload, []byte("[DONE]")) {
|
|
return line
|
|
}
|
|
rewritten := rewriteProviderJSONResponse(payload, model, toolCallWire)
|
|
if bytes.Equal(rewritten, payload) {
|
|
return line
|
|
}
|
|
out := make([]byte, 0, len("data: ")+len(rewritten)+len(ending))
|
|
out = append(out, "data: "...)
|
|
out = append(out, rewritten...)
|
|
out = append(out, ending...)
|
|
return out
|
|
}
|
|
|
|
func rewriteProviderJSONResponse(body []byte, model, toolCallWire string) []byte {
|
|
rewritten := rewriteProviderJSONModel(body, model)
|
|
if toolCallWire != config.ProtocolToolCallWireGeminiChat {
|
|
return rewritten
|
|
}
|
|
return normalizeGeminiChatProviderResponse(rewritten)
|
|
}
|
|
|
|
func normalizeGeminiChatProviderResponse(body []byte) []byte {
|
|
if !bytes.Contains(body, []byte(`"thought_signature"`)) {
|
|
return body
|
|
}
|
|
decoder := json.NewDecoder(bytes.NewReader(body))
|
|
decoder.UseNumber()
|
|
var response map[string]any
|
|
if err := decoder.Decode(&response); err != nil {
|
|
return body
|
|
}
|
|
changed := false
|
|
for _, rawChoice := range anySlice(response["choices"]) {
|
|
choice, ok := rawChoice.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
for _, messageKey := range []string{"message", "delta"} {
|
|
message, ok := choice[messageKey].(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
for _, rawCall := range anySlice(message["tool_calls"]) {
|
|
call, ok := rawCall.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
id, idOK := call["id"].(string)
|
|
extra, extraOK := call["extra_content"].(map[string]any)
|
|
if !idOK || id == "" || !extraOK {
|
|
continue
|
|
}
|
|
google, googleOK := extra["google"].(map[string]any)
|
|
if !googleOK {
|
|
continue
|
|
}
|
|
signature, signatureOK := google["thought_signature"].(string)
|
|
if !signatureOK || signature == "" {
|
|
continue
|
|
}
|
|
call["id"] = encodeGeminiThoughtSignatureToolID(id, signature)
|
|
delete(google, "thought_signature")
|
|
if len(google) == 0 {
|
|
delete(extra, "google")
|
|
}
|
|
if len(extra) == 0 {
|
|
delete(call, "extra_content")
|
|
}
|
|
changed = true
|
|
}
|
|
}
|
|
}
|
|
if !changed {
|
|
return body
|
|
}
|
|
encoded, err := json.Marshal(response)
|
|
if err != nil {
|
|
return body
|
|
}
|
|
return encoded
|
|
}
|
|
|
|
func normalizeGeminiChatProviderRequest(body []byte) ([]byte, error) {
|
|
if !bytes.Contains(body, []byte(geminiThoughtSignatureToolIDPrefix)) {
|
|
return body, nil
|
|
}
|
|
decoder := json.NewDecoder(bytes.NewReader(body))
|
|
decoder.UseNumber()
|
|
var request map[string]any
|
|
if err := decoder.Decode(&request); err != nil {
|
|
return nil, fmt.Errorf("decode Gemini Chat provider request: %w", err)
|
|
}
|
|
changed := false
|
|
for _, rawMessage := range anySlice(request["messages"]) {
|
|
message, ok := rawMessage.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
for _, rawCall := range anySlice(message["tool_calls"]) {
|
|
call, ok := rawCall.(map[string]any)
|
|
if !ok {
|
|
continue
|
|
}
|
|
encodedID, ok := call["id"].(string)
|
|
if !ok {
|
|
continue
|
|
}
|
|
id, signature, encoded, err := decodeGeminiThoughtSignatureToolID(encodedID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if !encoded {
|
|
continue
|
|
}
|
|
if err := restoreGeminiThoughtSignature(call, signature); err != nil {
|
|
return nil, err
|
|
}
|
|
call["id"] = id
|
|
changed = true
|
|
}
|
|
if encodedID, ok := message["tool_call_id"].(string); ok {
|
|
id, _, encoded, err := decodeGeminiThoughtSignatureToolID(encodedID)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if encoded {
|
|
message["tool_call_id"] = id
|
|
changed = true
|
|
}
|
|
}
|
|
}
|
|
if !changed {
|
|
return body, nil
|
|
}
|
|
encoded, err := json.Marshal(request)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("encode Gemini Chat provider request: %w", err)
|
|
}
|
|
return encoded, nil
|
|
}
|
|
|
|
func restoreGeminiThoughtSignature(call map[string]any, signature string) error {
|
|
extra, ok := call["extra_content"].(map[string]any)
|
|
if !ok {
|
|
if call["extra_content"] != nil {
|
|
return fmt.Errorf("restore Gemini thought signature: extra_content is not an object")
|
|
}
|
|
extra = make(map[string]any)
|
|
call["extra_content"] = extra
|
|
}
|
|
google, ok := extra["google"].(map[string]any)
|
|
if !ok {
|
|
if extra["google"] != nil {
|
|
return fmt.Errorf("restore Gemini thought signature: extra_content.google is not an object")
|
|
}
|
|
google = make(map[string]any)
|
|
extra["google"] = google
|
|
}
|
|
if existing, exists := google["thought_signature"]; exists && existing != signature {
|
|
return fmt.Errorf("restore Gemini thought signature: conflicting signature")
|
|
}
|
|
google["thought_signature"] = signature
|
|
return nil
|
|
}
|
|
|
|
func prepareProviderChatToolCallNormalization(tunnel edgeservice.SubmitProviderTunnelRequest, selected edgeservice.ProviderPoolCandidate) (edgeservice.SubmitProviderTunnelRequest, error) {
|
|
if selected.ProtocolProfile == nil {
|
|
return tunnel, nil
|
|
}
|
|
mapping, ok := selected.ProtocolProfile.ToolCallMapping(config.OperationChatCompletions)
|
|
if !ok || mapping.Wire != config.ProtocolToolCallWireGeminiChat {
|
|
return tunnel, nil
|
|
}
|
|
build := tunnel.BuildBody
|
|
if build == nil {
|
|
return tunnel, fmt.Errorf("Gemini Chat provider body builder is unavailable")
|
|
}
|
|
tunnel.BuildBody = func(target string) ([]byte, error) {
|
|
body, err := build(target)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return normalizeGeminiChatProviderRequest(body)
|
|
}
|
|
return tunnel, nil
|
|
}
|
|
|
|
func splitLineEnding(line []byte) ([]byte, []byte) {
|
|
if len(line) == 0 || line[len(line)-1] != '\n' {
|
|
return line, nil
|
|
}
|
|
if len(line) >= 2 && line[len(line)-2] == '\r' {
|
|
return line[:len(line)-2], []byte("\r\n")
|
|
}
|
|
return line[:len(line)-1], []byte("\n")
|
|
}
|
|
|
|
func rewriteProviderJSONModel(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 == "model" {
|
|
plan, err := planTopLevelJSONPatches(body, []topLevelJSONPatch{{name: "model", value: modelJSON}})
|
|
if err == nil {
|
|
return plan.apply()
|
|
}
|
|
break
|
|
}
|
|
}
|
|
return body
|
|
}
|
|
|
|
// rewriteChatCompletionModel replaces only the model field of the caller's
|
|
// original request JSON so the provider receives its served model name. The
|
|
// rest of the caller payload is forwarded without IOP rewriting.
|
|
func rewriteChatCompletionModel(rawBody []byte, target string, req chatCompletionRequest) ([]byte, error) {
|
|
patches, err := chatCompletionProviderPatches(target, req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(patches) == 0 {
|
|
return rawBody, nil
|
|
}
|
|
plan, err := planTopLevelJSONPatches(rawBody, patches)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return plan.apply(), nil
|
|
}
|
|
|
|
func rewriteChatCompletionModelFromIngress(ingress *openAIIngressSnapshot, target string, req chatCompletionRequest) (*openAIRebuiltLease, error) {
|
|
patches, err := chatCompletionProviderPatches(target, req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return rewriteIngressTopLevelJSON(ingress, patches)
|
|
}
|
|
|
|
func chatCompletionProviderPatches(target string, req chatCompletionRequest) ([]topLevelJSONPatch, error) {
|
|
patches := make([]topLevelJSONPatch, 0, 5)
|
|
if strings.TrimSpace(target) != "" {
|
|
modelJSON, err := json.Marshal(target)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
patches = append(patches, topLevelJSONPatch{name: "model", value: modelJSON})
|
|
}
|
|
if req.MaxTokens != nil {
|
|
maxTokensJSON, err := json.Marshal(*req.MaxTokens)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
patches = append(patches,
|
|
topLevelJSONPatch{name: "max_tokens", value: maxTokensJSON},
|
|
topLevelJSONPatch{name: "max_completion_tokens"},
|
|
)
|
|
}
|
|
if req.ThinkingTokenBudget != nil {
|
|
budgetJSON, err := json.Marshal(*req.ThinkingTokenBudget)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
patches = append(patches, topLevelJSONPatch{name: "thinking_token_budget", value: budgetJSON})
|
|
}
|
|
if req.Think != nil {
|
|
thinkJSON, err := json.Marshal(*req.Think)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
patches = append(patches, topLevelJSONPatch{name: "think", value: thinkJSON})
|
|
}
|
|
return patches, nil
|
|
}
|
|
|
|
// rewriteResponsesModel replaces only the model field of the caller's original
|
|
// /v1/responses request JSON so the provider receives its served model name.
|
|
// Every other field (input, instructions, tools, max_output_tokens, and any
|
|
// Codex/Responses-specific field) is forwarded without IOP rewriting. An empty
|
|
// target leaves the body untouched; invalid JSON is rejected.
|
|
func rewriteResponsesModel(rawBody []byte, target string) ([]byte, error) {
|
|
if strings.TrimSpace(target) == "" {
|
|
return rawBody, nil
|
|
}
|
|
modelJSON, err := json.Marshal(target)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan, err := planTopLevelJSONPatches(rawBody, []topLevelJSONPatch{{name: "model", value: modelJSON}})
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return plan.apply(), nil
|
|
}
|
|
|
|
func rewriteResponsesModelFromIngress(ingress *openAIIngressSnapshot, target string) (*openAIRebuiltLease, error) {
|
|
if strings.TrimSpace(target) == "" {
|
|
return rewriteIngressTopLevelJSON(ingress, nil)
|
|
}
|
|
modelJSON, err := json.Marshal(target)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return rewriteIngressTopLevelJSON(ingress, []topLevelJSONPatch{{name: "model", value: modelJSON}})
|
|
}
|
|
|
|
// rewriteIngressTopLevelJSON places both allocation checks around the single
|
|
// patched output and returns a lease that stays live through synchronous
|
|
// provider submission. The canonical ingress owner remains unchanged.
|
|
func rewriteIngressTopLevelJSON(ingress *openAIIngressSnapshot, patches []topLevelJSONPatch) (*openAIRebuiltLease, error) {
|
|
body, err := ingress.canonicalBody()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
plan, err := planTopLevelJSONPatches(body, patches)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(plan.edits) == 0 {
|
|
return &openAIRebuiltLease{ingress: ingress, bodyAlias: body}, nil
|
|
}
|
|
guard, err := ingress.reserveRebuild(int64(plan.outputSize))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
output := plan.apply()
|
|
rebuilt, err := guard.CommitOwnedTyped(openAIRebuiltBodyViewName, output)
|
|
if err != nil {
|
|
guard.Close()
|
|
return nil, err
|
|
}
|
|
return &openAIRebuiltLease{
|
|
ingress: ingress, guard: guard, rebuilt: rebuilt, bodyAlias: output,
|
|
}, nil
|
|
}
|