530 lines
15 KiB
Go
530 lines
15 KiB
Go
package providers
|
|
|
|
import (
|
|
"git.gobha.me/xcaliber/chat-switchboard/capabilities"
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
)
|
|
|
|
// OpenAIProvider handles any OpenAI-compatible API.
|
|
type OpenAIProvider struct{}
|
|
|
|
func (p *OpenAIProvider) ID() string { return "openai" }
|
|
|
|
// ── Chat Completion (non-streaming) ─────────
|
|
|
|
func (p *OpenAIProvider) ChatCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (*CompletionResponse, error) {
|
|
req.Stream = false
|
|
|
|
body, err := p.doRequest(ctx, cfg, req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer body.Close()
|
|
|
|
var resp openaiChatResponse
|
|
if err := json.NewDecoder(body).Decode(&resp); err != nil {
|
|
return nil, fmt.Errorf("decode response: %w", err)
|
|
}
|
|
|
|
if len(resp.Choices) == 0 {
|
|
return nil, fmt.Errorf("no choices in response")
|
|
}
|
|
|
|
result := &CompletionResponse{
|
|
Content: resp.Choices[0].Message.Content,
|
|
Model: resp.Model,
|
|
FinishReason: resp.Choices[0].FinishReason,
|
|
InputTokens: resp.Usage.PromptTokens,
|
|
OutputTokens: resp.Usage.CompletionTokens,
|
|
}
|
|
|
|
// Cache tokens (OpenAI, Venice report these in prompt_tokens_details)
|
|
if resp.Usage.PromptTokensDetails != nil {
|
|
result.CacheReadTokens = resp.Usage.PromptTokensDetails.CachedTokens
|
|
}
|
|
|
|
// Extract tool calls if present
|
|
if len(resp.Choices[0].Message.ToolCalls) > 0 {
|
|
for _, tc := range resp.Choices[0].Message.ToolCalls {
|
|
result.ToolCalls = append(result.ToolCalls, ToolCall{
|
|
ID: tc.ID,
|
|
Type: tc.Type,
|
|
Function: FunctionCall{
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
},
|
|
})
|
|
}
|
|
if result.FinishReason == "" {
|
|
result.FinishReason = "tool_calls"
|
|
}
|
|
}
|
|
|
|
return result, nil
|
|
}
|
|
|
|
// ── Stream Completion ───────────────────────
|
|
|
|
func (p *OpenAIProvider) StreamCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (<-chan StreamEvent, error) {
|
|
req.Stream = true
|
|
|
|
body, err := p.doRequest(ctx, cfg, req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
ch := make(chan StreamEvent, 64)
|
|
|
|
go func() {
|
|
defer close(ch)
|
|
defer body.Close()
|
|
|
|
// Accumulate tool calls across stream chunks
|
|
toolCallMap := map[int]*ToolCall{} // index → accumulated call
|
|
|
|
// Usage arrives in a separate chunk (or on the final chunk).
|
|
// OpenAI protocol order: content → finish_reason → usage → [DONE]
|
|
// The usage chunk has choices=[] so we must hold the finish event
|
|
// until usage arrives, otherwise tokens are always 0.
|
|
var streamUsage *openaiUsage
|
|
var pendingFinish *StreamEvent // deferred until usage arrives
|
|
|
|
scanner := bufio.NewScanner(body)
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
|
|
if !strings.HasPrefix(line, "data: ") {
|
|
continue
|
|
}
|
|
|
|
data := strings.TrimPrefix(line, "data: ")
|
|
if data == "[DONE]" {
|
|
// Flush any pending finish event (usage may or may not have arrived)
|
|
if pendingFinish != nil {
|
|
ch <- *pendingFinish
|
|
} else {
|
|
ch <- StreamEvent{Done: true}
|
|
}
|
|
return
|
|
}
|
|
|
|
var chunk openaiStreamChunk
|
|
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
|
continue
|
|
}
|
|
|
|
// Capture usage — may arrive in a separate chunk with empty choices
|
|
if chunk.Usage != nil {
|
|
streamUsage = chunk.Usage
|
|
// If finish already arrived, attach usage and flush immediately
|
|
if pendingFinish != nil {
|
|
pendingFinish.InputTokens = streamUsage.PromptTokens
|
|
pendingFinish.OutputTokens = streamUsage.CompletionTokens
|
|
if streamUsage.PromptTokensDetails != nil {
|
|
pendingFinish.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
|
|
}
|
|
ch <- *pendingFinish
|
|
return // Done — [DONE] will be handled by defer body.Close()
|
|
}
|
|
}
|
|
|
|
if len(chunk.Choices) == 0 {
|
|
continue
|
|
}
|
|
|
|
choice := chunk.Choices[0]
|
|
|
|
// Accumulate tool call deltas
|
|
for _, tc := range choice.Delta.ToolCalls {
|
|
idx := 0
|
|
if tc.Index != nil {
|
|
idx = *tc.Index
|
|
}
|
|
|
|
if existing, ok := toolCallMap[idx]; ok {
|
|
// Append arguments fragment
|
|
existing.Function.Arguments += tc.Function.Arguments
|
|
} else {
|
|
// First chunk for this tool — has ID and name
|
|
toolCallMap[idx] = &ToolCall{
|
|
ID: tc.ID,
|
|
Type: "function",
|
|
Function: FunctionCall{
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
},
|
|
}
|
|
}
|
|
}
|
|
|
|
ev := StreamEvent{
|
|
Delta: choice.Delta.Content,
|
|
Reasoning: choice.Delta.ReasoningContent,
|
|
Model: chunk.Model,
|
|
FinishReason: choice.FinishReason,
|
|
}
|
|
|
|
if ev.FinishReason != "" {
|
|
ev.Done = true
|
|
// Attach accumulated tool calls on final event
|
|
if ev.FinishReason == "tool_calls" && len(toolCallMap) > 0 {
|
|
for _, tc := range toolCallMap {
|
|
ev.ToolCalls = append(ev.ToolCalls, *tc)
|
|
}
|
|
// Tool calls need to flush immediately — the tool loop
|
|
// needs the event to start executing tools
|
|
if streamUsage != nil {
|
|
ev.InputTokens = streamUsage.PromptTokens
|
|
ev.OutputTokens = streamUsage.CompletionTokens
|
|
if streamUsage.PromptTokensDetails != nil {
|
|
ev.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
|
|
}
|
|
}
|
|
ch <- ev
|
|
continue
|
|
}
|
|
// Normal finish — defer until usage chunk arrives
|
|
if streamUsage != nil {
|
|
// Usage already captured (rare: same chunk or earlier)
|
|
ev.InputTokens = streamUsage.PromptTokens
|
|
ev.OutputTokens = streamUsage.CompletionTokens
|
|
if streamUsage.PromptTokensDetails != nil {
|
|
ev.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
|
|
}
|
|
ch <- ev
|
|
continue
|
|
}
|
|
// Hold — usage chunk hasn't arrived yet
|
|
pendingFinish = &ev
|
|
continue
|
|
}
|
|
|
|
ch <- ev
|
|
}
|
|
|
|
if err := scanner.Err(); err != nil {
|
|
ch <- StreamEvent{Error: err}
|
|
} else if pendingFinish != nil {
|
|
// Scanner ended without [DONE] — flush pending finish
|
|
ch <- *pendingFinish
|
|
}
|
|
}()
|
|
|
|
return ch, nil
|
|
}
|
|
|
|
// ── List Models ─────────────────────────────
|
|
|
|
func (p *OpenAIProvider) ListModels(ctx context.Context, cfg ProviderConfig) ([]Model, error) {
|
|
endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/models"
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if cfg.APIKey != "" {
|
|
httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
|
|
}
|
|
for k, v := range cfg.CustomHeaders {
|
|
httpReq.Header.Set(k, v)
|
|
}
|
|
|
|
resp, err := http.DefaultClient.Do(httpReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("list models: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
b, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("list models: HTTP %d: %s", resp.StatusCode, string(b))
|
|
}
|
|
|
|
var result openaiModelsResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
|
return nil, fmt.Errorf("decode models: %w", err)
|
|
}
|
|
|
|
out := make([]Model, 0, len(result.Data))
|
|
for _, m := range result.Data {
|
|
// Heuristic inference from model ID
|
|
caps := capabilities.InferCapabilities(m.ID)
|
|
// Use context_length from API if available and we don't have it
|
|
if m.ContextLength > 0 && caps.MaxContext == 0 {
|
|
caps.MaxContext = m.ContextLength
|
|
}
|
|
|
|
out = append(out, Model{
|
|
ID: m.ID,
|
|
OwnedBy: m.OwnedBy,
|
|
Capabilities: caps,
|
|
})
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// ── Embeddings ─────────────────────────────
|
|
|
|
func (p *OpenAIProvider) Embed(ctx context.Context, cfg ProviderConfig, req EmbeddingRequest) (*EmbeddingResponse, error) {
|
|
endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/embeddings"
|
|
|
|
oaiReq := openaiEmbeddingRequest{
|
|
Model: req.Model,
|
|
Input: req.Input,
|
|
}
|
|
body, err := json.Marshal(oaiReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal embedding request: %w", err)
|
|
}
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
if cfg.APIKey != "" {
|
|
httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
|
|
}
|
|
for k, v := range cfg.CustomHeaders {
|
|
httpReq.Header.Set(k, v)
|
|
}
|
|
|
|
resp, err := http.DefaultClient.Do(httpReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("embedding request: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
b, _ := io.ReadAll(resp.Body)
|
|
return nil, fmt.Errorf("embedding error: HTTP %d: %s", resp.StatusCode, string(b))
|
|
}
|
|
|
|
var oaiResp openaiEmbeddingResponse
|
|
if err := json.NewDecoder(resp.Body).Decode(&oaiResp); err != nil {
|
|
return nil, fmt.Errorf("decode embedding response: %w", err)
|
|
}
|
|
|
|
result := &EmbeddingResponse{
|
|
Model: oaiResp.Model,
|
|
InputTokens: oaiResp.Usage.PromptTokens,
|
|
Embeddings: make([][]float64, len(oaiResp.Data)),
|
|
}
|
|
for i, d := range oaiResp.Data {
|
|
result.Embeddings[i] = d.Embedding
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// ── HTTP Layer ──────────────────────────────
|
|
|
|
func (p *OpenAIProvider) doRequest(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (io.ReadCloser, error) {
|
|
endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/chat/completions"
|
|
|
|
// Build OpenAI-format request
|
|
oaiReq := openaiChatRequest{
|
|
Model: req.Model,
|
|
Stream: req.Stream,
|
|
Messages: make([]openaiMessage, 0, len(req.Messages)),
|
|
}
|
|
if req.Stream {
|
|
oaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true}
|
|
}
|
|
if req.MaxTokens > 0 {
|
|
oaiReq.MaxTokens = req.MaxTokens
|
|
}
|
|
if req.Temperature != nil {
|
|
oaiReq.Temperature = req.Temperature
|
|
}
|
|
if req.TopP != nil {
|
|
oaiReq.TopP = req.TopP
|
|
}
|
|
|
|
// Convert tools
|
|
for _, t := range req.Tools {
|
|
oaiReq.Tools = append(oaiReq.Tools, openaiToolDef{
|
|
Type: t.Type,
|
|
Function: openaiToolDefFunction{
|
|
Name: t.Function.Name,
|
|
Description: t.Function.Description,
|
|
Parameters: t.Function.Parameters,
|
|
},
|
|
})
|
|
}
|
|
|
|
// Convert messages — handle all roles including tool
|
|
for _, m := range req.Messages {
|
|
oaiMsg := openaiMessage{
|
|
Role: m.Role,
|
|
Content: m.Content,
|
|
}
|
|
|
|
// Assistant messages with tool calls
|
|
if len(m.ToolCalls) > 0 {
|
|
for _, tc := range m.ToolCalls {
|
|
oaiMsg.ToolCalls = append(oaiMsg.ToolCalls, openaiToolCall{
|
|
ID: tc.ID,
|
|
Type: tc.Type,
|
|
Function: openaiToolCallFunction{
|
|
Name: tc.Function.Name,
|
|
Arguments: tc.Function.Arguments,
|
|
},
|
|
})
|
|
}
|
|
}
|
|
|
|
// Tool result messages
|
|
if m.Role == "tool" {
|
|
oaiMsg.ToolCallID = m.ToolCallID
|
|
oaiMsg.Name = m.Name
|
|
}
|
|
|
|
oaiReq.Messages = append(oaiReq.Messages, oaiMsg)
|
|
}
|
|
|
|
body, err := json.Marshal(oaiReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal request: %w", err)
|
|
}
|
|
|
|
httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
httpReq.Header.Set("Content-Type", "application/json")
|
|
if cfg.APIKey != "" {
|
|
httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
|
|
}
|
|
// Inject provider-level custom headers (OpenRouter, Venice, etc.)
|
|
for k, v := range cfg.CustomHeaders {
|
|
httpReq.Header.Set(k, v)
|
|
}
|
|
|
|
resp, err := http.DefaultClient.Do(httpReq)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("provider request: %w", err)
|
|
}
|
|
|
|
if resp.StatusCode != http.StatusOK {
|
|
b, _ := io.ReadAll(resp.Body)
|
|
resp.Body.Close()
|
|
return nil, fmt.Errorf("provider error: HTTP %d: %s", resp.StatusCode, string(b))
|
|
}
|
|
|
|
return resp.Body, nil
|
|
}
|
|
|
|
// ── OpenAI Wire Types ───────────────────────
|
|
|
|
type openaiMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content,omitempty"`
|
|
ToolCalls []openaiToolCall `json:"tool_calls,omitempty"` // assistant requesting tools
|
|
ToolCallID string `json:"tool_call_id,omitempty"` // role=tool result
|
|
Name string `json:"name,omitempty"` // tool name for role=tool
|
|
}
|
|
|
|
type openaiToolCallFunction struct {
|
|
Name string `json:"name,omitempty"`
|
|
Arguments string `json:"arguments,omitempty"`
|
|
}
|
|
|
|
type openaiToolCall struct {
|
|
Index *int `json:"index,omitempty"` // present in streaming deltas
|
|
ID string `json:"id,omitempty"`
|
|
Type string `json:"type,omitempty"` // "function"
|
|
Function openaiToolCallFunction `json:"function"`
|
|
}
|
|
|
|
type openaiToolDefFunction struct {
|
|
Name string `json:"name"`
|
|
Description string `json:"description"`
|
|
Parameters json.RawMessage `json:"parameters"`
|
|
}
|
|
|
|
type openaiToolDef struct {
|
|
Type string `json:"type"` // "function"
|
|
Function openaiToolDefFunction `json:"function"`
|
|
}
|
|
|
|
type openaiChatRequest struct {
|
|
Model string `json:"model"`
|
|
Messages []openaiMessage `json:"messages"`
|
|
MaxTokens int `json:"max_tokens,omitempty"`
|
|
Temperature *float64 `json:"temperature,omitempty"`
|
|
TopP *float64 `json:"top_p,omitempty"`
|
|
Stream bool `json:"stream"`
|
|
StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
|
|
Tools []openaiToolDef `json:"tools,omitempty"`
|
|
}
|
|
|
|
type openaiStreamOptions struct {
|
|
IncludeUsage bool `json:"include_usage"`
|
|
}
|
|
|
|
type openaiChatResponse struct {
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Message openaiMessage `json:"message"`
|
|
FinishReason string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
Usage openaiUsage `json:"usage"`
|
|
}
|
|
|
|
type openaiUsage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
CompletionTokens int `json:"completion_tokens"`
|
|
PromptTokensDetails *struct {
|
|
CachedTokens int `json:"cached_tokens"`
|
|
} `json:"prompt_tokens_details,omitempty"`
|
|
CompletionTokensDetails *struct {
|
|
ReasoningTokens int `json:"reasoning_tokens"`
|
|
} `json:"completion_tokens_details,omitempty"`
|
|
}
|
|
|
|
type openaiStreamChunk struct {
|
|
Model string `json:"model"`
|
|
Choices []struct {
|
|
Delta struct {
|
|
Content string `json:"content"`
|
|
ReasoningContent string `json:"reasoning_content"`
|
|
ToolCalls []openaiToolCall `json:"tool_calls,omitempty"`
|
|
} `json:"delta"`
|
|
FinishReason string `json:"finish_reason"`
|
|
} `json:"choices"`
|
|
Usage *openaiUsage `json:"usage,omitempty"`
|
|
}
|
|
|
|
type openaiModelsResponse struct {
|
|
Data []struct {
|
|
ID string `json:"id"`
|
|
OwnedBy string `json:"owned_by"`
|
|
ContextLength int `json:"context_length,omitempty"` // Ollama, OpenRouter
|
|
} `json:"data"`
|
|
}
|
|
|
|
// ── Embedding Wire Types ──────────────────
|
|
|
|
type openaiEmbeddingRequest struct {
|
|
Model string `json:"model"`
|
|
Input []string `json:"input"`
|
|
}
|
|
|
|
type openaiEmbeddingResponse struct {
|
|
Data []struct {
|
|
Embedding []float64 `json:"embedding"`
|
|
Index int `json:"index"`
|
|
} `json:"data"`
|
|
Model string `json:"model"`
|
|
Usage struct {
|
|
PromptTokens int `json:"prompt_tokens"`
|
|
} `json:"usage"`
|
|
}
|