Files
warmbly/internal/pkg/generation/openai_provider.go
T

661 lines
21 KiB
Go

package generation
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"strings"
"sync/atomic"
"time"
)
// openAIProvider is the primary provider: it drives the tool-use agent loop and
// the writing assistant over the OpenAI (or any OpenAI-compatible) chat
// completions API via a lean HTTP client. Using plain HTTP (rather than the SDK
// used for the warmup batch path) keeps the tool-call JSON explicit and lets a
// self-hoster retarget the whole loop with an AI_PROVIDER preset or AI_BASE_URL.
type openAIProvider struct {
apiKey string
baseURL string
modelTrial string
modelPaid string
search SearchClient
local bool
http *http.Client
// Sticky parameter-compatibility flags. Newer OpenAI models (gpt-5.x,
// o-series) reject the legacy max_tokens param (they want
// max_completion_tokens) and any non-default temperature, while most
// OpenAI-compatible backends (Ollama, Groq, OpenRouter) only know the
// legacy shape. Instead of a model-family table that goes stale, the first
// such 400 flips the flag and the call retries adapted; every later call
// uses the adapted shape directly.
useMaxCompletionTokens atomic.Bool
omitTemperature atomic.Bool
// stream_options is OpenAI's usage-in-stream opt-in; some compatible
// backends reject the field entirely.
omitStreamOptions atomic.Bool
}
// defaultOpenAIBaseURL is the public OpenAI API. Overridable for
// OpenAI-compatible self-hosted endpoints (Ollama, vLLM, LocalAI, OpenRouter).
const defaultOpenAIBaseURL = "https://api.openai.com/v1"
// defaultLocalModel is the fall-back model tag for a free/local backend
// (AI_FREE) when no explicit model id is set, so a dev who forgets
// AI_MODEL does not 404 on "gpt-4o-mini". llama3.1 has reliable
// tool-calling in Ollama; override with AI_MODEL (or AI_MODEL_TRIAL / _PAID).
const defaultLocalModel = "llama3.1"
func newOpenAIProvider(cfg ProviderConfig) *openAIProvider {
base := strings.TrimRight(strings.TrimSpace(cfg.OpenAIBaseURL), "/")
if base == "" {
base = defaultOpenAIBaseURL
}
// Empty model ids default to the hosted OpenAI models, or a local-friendly
// tag when this is an explicit free/local backend.
fallbackTrial, fallbackPaid := ModelWritingFreeOpenAI, ModelWritingPaidOpenAI
if cfg.Local {
fallbackTrial, fallbackPaid = defaultLocalModel, defaultLocalModel
}
trial := cfg.OpenAIModelTrial
if trial == "" {
trial = fallbackTrial
}
paid := cfg.OpenAIModelPaid
if paid == "" {
paid = fallbackPaid
}
return &openAIProvider{
apiKey: cfg.OpenAIAPIKey,
baseURL: base,
modelTrial: trial,
modelPaid: paid,
search: cfg.Search,
local: cfg.Local,
http: &http.Client{Timeout: 90 * time.Second},
}
}
func (p *openAIProvider) Name() string { return "openai" }
func (p *openAIProvider) IsLocal() bool { return p.local }
func (p *openAIProvider) ModelForTier(paid bool) string {
if paid {
return p.modelPaid
}
return p.modelTrial
}
// SearchClient exposes the configured web-search backend so the search_web tool
// (M2) can reach it through the provider. Nil when search is not configured.
func (p *openAIProvider) SearchClient() SearchClient { return p.search }
// --- chat-completions wire types (OpenAI-compatible) ---
type oaiTool struct {
Type string `json:"type"`
Function oaiToolFunc `json:"function"`
}
type oaiToolFunc struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Parameters map[string]any `json:"parameters,omitempty"`
}
type oaiToolCall struct {
ID string `json:"id"`
Type string `json:"type"`
Function struct {
Name string `json:"name"`
Arguments string `json:"arguments"`
} `json:"function"`
}
type oaiMessage struct {
Role string `json:"role"`
Content string `json:"content,omitempty"`
ToolCalls []oaiToolCall `json:"tool_calls,omitempty"`
ToolCallID string `json:"tool_call_id,omitempty"`
}
type oaiRequest struct {
Model string `json:"model"`
// Exactly one of MaxTokens / MaxCompletionTokens is set per call, driven
// by the provider's useMaxCompletionTokens compatibility flag.
MaxTokens int `json:"max_tokens,omitempty"`
MaxCompletionTokens int `json:"max_completion_tokens,omitempty"`
Messages []oaiMessage `json:"messages"`
Tools []oaiTool `json:"tools,omitempty"`
ToolChoice string `json:"tool_choice,omitempty"`
Temperature *float64 `json:"temperature,omitempty"`
Stream bool `json:"stream,omitempty"`
StreamOptions *oaiStreamOptions `json:"stream_options,omitempty"`
}
type oaiStreamOptions struct {
IncludeUsage bool `json:"include_usage"`
}
// oaiStreamChunk is one `data:` frame of a streamed completion. Tool-call
// fragments arrive indexed and must be accumulated across chunks.
type oaiStreamChunk struct {
Choices []struct {
Delta struct {
Content string `json:"content"`
ToolCalls []struct {
Index int `json:"index"`
ID string `json:"id"`
Type string `json:"type"`
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 *struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
Error *oaiError `json:"error"`
}
type oaiChoice struct {
Message struct {
Content string `json:"content"`
ToolCalls []oaiToolCall `json:"tool_calls"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
}
type oaiResponse struct {
Choices []oaiChoice `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
Error *oaiError `json:"error"`
}
type oaiError struct {
Message string `json:"message"`
Type string `json:"type"`
Param string `json:"param"`
Code string `json:"code"`
}
// toolDefsToWire converts the provider-agnostic tool defs to OpenAI tools.
func toolDefsToWire(tools []ToolDef) []oaiTool {
if len(tools) == 0 {
return nil
}
out := make([]oaiTool, 0, len(tools))
for _, t := range tools {
out = append(out, oaiTool{
Type: "function",
Function: oaiToolFunc{
Name: t.Name,
Description: t.Description,
Parameters: t.InputSchema,
},
})
}
return out
}
// transcriptToWire renders the running transcript into OpenAI messages, with
// the system prompt prepended.
func transcriptToWire(system string, msgs []AgentMessage) []oaiMessage {
out := make([]oaiMessage, 0, len(msgs)+1)
if strings.TrimSpace(system) != "" {
out = append(out, oaiMessage{Role: "system", Content: system})
}
for _, m := range msgs {
wm := oaiMessage{Role: m.Role, Content: m.Content, ToolCallID: m.ToolCallID}
for _, tc := range m.ToolCalls {
var call oaiToolCall
call.ID = tc.ID
call.Type = "function"
call.Function.Name = tc.Name
call.Function.Arguments = string(tc.Args)
wm.ToolCalls = append(wm.ToolCalls, call)
}
out = append(out, wm)
}
return out
}
// complete performs one chat-completion call, retrying once per parameter when
// the backend rejects the legacy request shape (see the compatibility flags on
// openAIProvider).
func (p *openAIProvider) complete(ctx context.Context, model string, maxTokens int, msgs []oaiMessage, tools []oaiTool, temperature *float64) (*oaiResponse, error) {
for attempt := 0; ; attempt++ {
reqBody := oaiRequest{Model: model, Messages: msgs}
if p.useMaxCompletionTokens.Load() {
reqBody.MaxCompletionTokens = maxTokens
} else {
reqBody.MaxTokens = maxTokens
}
if !p.omitTemperature.Load() {
reqBody.Temperature = temperature
}
if len(tools) > 0 {
reqBody.Tools = tools
reqBody.ToolChoice = "auto"
}
body, err := json.Marshal(reqBody)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", "Bearer "+p.apiKey)
resp, err := p.http.Do(req)
if err != nil {
return nil, err
}
raw, err := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
resp.Body.Close()
if err != nil {
return nil, err
}
var parsed oaiResponse
if err := json.Unmarshal(raw, &parsed); err != nil {
return nil, fmt.Errorf("openai: decode response: %w", err)
}
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
if resp.StatusCode == http.StatusBadRequest && attempt < 2 && p.adaptParams(parsed.Error) {
continue
}
if parsed.Error != nil {
return nil, fmt.Errorf("openai: %s: %s", parsed.Error.Type, parsed.Error.Message)
}
return nil, fmt.Errorf("openai: unexpected status %d", resp.StatusCode)
}
if len(parsed.Choices) == 0 {
return nil, errors.New("openai: empty completion")
}
return &parsed, nil
}
}
// completeStream performs one chat-completion call with stream=true, invoking
// onDelta for each content fragment as it arrives and returning the fully
// accumulated response in the same shape complete() produces (so RunAgent's
// loop is stream-agnostic). Backends that ignore stream=true and answer with
// plain JSON are handled transparently.
func (p *openAIProvider) completeStream(ctx context.Context, model string, maxTokens int, msgs []oaiMessage, tools []oaiTool, temperature *float64, onDelta func(string)) (*oaiResponse, error) {
for attempt := 0; ; attempt++ {
reqBody := oaiRequest{Model: model, Messages: msgs, Stream: true}
if !p.omitStreamOptions.Load() {
reqBody.StreamOptions = &oaiStreamOptions{IncludeUsage: true}
}
if p.useMaxCompletionTokens.Load() {
reqBody.MaxCompletionTokens = maxTokens
} else {
reqBody.MaxTokens = maxTokens
}
if !p.omitTemperature.Load() {
reqBody.Temperature = temperature
}
if len(tools) > 0 {
reqBody.Tools = tools
reqBody.ToolChoice = "auto"
}
body, err := json.Marshal(reqBody)
if err != nil {
return nil, err
}
req, err := http.NewRequestWithContext(ctx, http.MethodPost, p.baseURL+"/chat/completions", bytes.NewReader(body))
if err != nil {
return nil, err
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Accept", "text/event-stream")
req.Header.Set("Authorization", "Bearer "+p.apiKey)
resp, err := p.http.Do(req)
if err != nil {
return nil, err
}
// Errors come back as a plain JSON body regardless of streaming.
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
raw, rerr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
resp.Body.Close()
if rerr != nil {
return nil, rerr
}
var parsed oaiResponse
_ = json.Unmarshal(raw, &parsed)
if resp.StatusCode == http.StatusBadRequest && attempt < 3 && p.adaptParams(parsed.Error) {
continue
}
if parsed.Error != nil {
return nil, fmt.Errorf("openai: %s: %s", parsed.Error.Type, parsed.Error.Message)
}
return nil, fmt.Errorf("openai: unexpected status %d", resp.StatusCode)
}
// A compatible backend may ignore stream=true and answer JSON.
if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/event-stream") {
raw, rerr := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
resp.Body.Close()
if rerr != nil {
return nil, rerr
}
var parsed oaiResponse
if err := json.Unmarshal(raw, &parsed); err != nil {
return nil, fmt.Errorf("openai: decode response: %w", err)
}
if len(parsed.Choices) == 0 {
return nil, errors.New("openai: empty completion")
}
if onDelta != nil && parsed.Choices[0].Message.Content != "" {
onDelta(parsed.Choices[0].Message.Content)
}
return &parsed, nil
}
out, serr := p.readStream(resp.Body, onDelta)
resp.Body.Close()
return out, serr
}
}
// readStream folds the SSE chunk frames into a complete oaiResponse:
// content deltas are forwarded to onDelta and concatenated, tool-call
// fragments are stitched together by index, and the trailing usage frame
// (stream_options.include_usage) fills Usage.
func (p *openAIProvider) readStream(body io.Reader, onDelta func(string)) (*oaiResponse, error) {
var out oaiResponse
out.Choices = make([]oaiChoice, 1)
var content strings.Builder
var calls []oaiToolCall
sawChunk := false
sc := bufio.NewScanner(body)
sc.Buffer(make([]byte, 0, 64<<10), 1<<20)
for sc.Scan() {
line := strings.TrimSpace(sc.Text())
if !strings.HasPrefix(line, "data:") {
continue
}
data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
if data == "" || data == "[DONE]" {
continue
}
var chunk oaiStreamChunk
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
continue
}
if chunk.Error != nil {
return nil, fmt.Errorf("openai: %s: %s", chunk.Error.Type, chunk.Error.Message)
}
sawChunk = true
if chunk.Usage != nil {
out.Usage.PromptTokens = chunk.Usage.PromptTokens
out.Usage.CompletionTokens = chunk.Usage.CompletionTokens
out.Usage.TotalTokens = chunk.Usage.TotalTokens
}
if len(chunk.Choices) == 0 {
continue
}
ch := chunk.Choices[0]
if ch.FinishReason != "" {
out.Choices[0].FinishReason = ch.FinishReason
}
if ch.Delta.Content != "" {
content.WriteString(ch.Delta.Content)
if onDelta != nil {
onDelta(ch.Delta.Content)
}
}
for _, tc := range ch.Delta.ToolCalls {
for len(calls) <= tc.Index {
calls = append(calls, oaiToolCall{Type: "function"})
}
if tc.ID != "" {
calls[tc.Index].ID = tc.ID
}
if tc.Function.Name != "" {
calls[tc.Index].Function.Name = tc.Function.Name
}
calls[tc.Index].Function.Arguments += tc.Function.Arguments
}
}
if err := sc.Err(); err != nil {
return nil, fmt.Errorf("openai: read stream: %w", err)
}
if !sawChunk {
return nil, errors.New("openai: empty stream")
}
out.Choices[0].Message.Content = content.String()
out.Choices[0].Message.ToolCalls = calls
return &out, nil
}
// adaptParams flips the sticky compatibility flag matching a 400 that names a
// parameter this backend rejects. Returns true when a retry makes sense.
func (p *openAIProvider) adaptParams(e *oaiError) bool {
if e == nil {
return false
}
switch {
case e.Param == "max_tokens" || strings.Contains(e.Message, "max_completion_tokens"):
if p.useMaxCompletionTokens.Load() {
return false
}
p.useMaxCompletionTokens.Store(true)
return true
case e.Param == "temperature":
if p.omitTemperature.Load() {
return false
}
p.omitTemperature.Store(true)
return true
case e.Param == "stream_options" || strings.Contains(e.Message, "stream_options"):
if p.omitStreamOptions.Load() {
return false
}
p.omitStreamOptions.Store(true)
return true
}
return false
}
// RunAgent implements the provider-agnostic tool-use loop. See provider.go for
// the approval/resume contract.
func (p *openAIProvider) RunAgent(ctx context.Context, req AgentRequest) (*AgentResult, error) {
model := req.Model
if model == "" {
model = p.modelTrial
}
maxIter := req.MaxIterations
if maxIter <= 0 {
maxIter = defaultMaxIterations
}
maxTokens := req.MaxTokens
if maxTokens <= 0 {
maxTokens = defaultAgentTokens
}
byName := make(map[string]ToolDef, len(req.Tools))
for _, t := range req.Tools {
byName[t.Name] = t
}
wireTools := toolDefsToWire(req.Tools)
// messages is the durable transcript; we mutate a copy so the caller's
// slice is not aliased.
messages := append([]AgentMessage(nil), req.Messages...)
result := &AgentResult{Model: model}
for iter := 0; iter < maxIter; iter++ {
if req.PreIteration != nil {
if err := req.PreIteration(ctx, iter+1); err != nil {
result.Messages = messages
result.StopReason = "stopped"
return result, nil
}
}
result.Iterations++
if req.OnEvent != nil {
req.OnEvent(AgentEvent{Type: EventIteration, Iteration: result.Iterations})
}
// With a listener attached, stream the completion so text reaches the
// client token by token; the closing EventText below carries the full
// block, so deltas stay purely cosmetic.
var resp *oaiResponse
var err error
if req.OnEvent != nil {
resp, err = p.completeStream(ctx, model, maxTokens, transcriptToWire(req.System, messages), wireTools, nil, func(delta string) {
req.OnEvent(AgentEvent{Type: EventTextDelta, Text: delta})
})
} else {
resp, err = p.complete(ctx, model, maxTokens, transcriptToWire(req.System, messages), wireTools, nil)
}
if err != nil {
return nil, err
}
result.TokensUsed += resp.Usage.TotalTokens
choice := resp.Choices[0]
// No tool calls: final answer.
if len(choice.Message.ToolCalls) == 0 {
text := strings.TrimSpace(choice.Message.Content)
messages = append(messages, AgentMessage{Role: "assistant", Content: text})
if req.OnEvent != nil && text != "" {
req.OnEvent(AgentEvent{Type: EventText, Text: text})
}
result.Text = text
result.Messages = messages
result.StopReason = "stop"
return result, nil
}
// Commentary alongside tool calls streams too; close its block so the
// client can finalize it before the tool chips appear.
if req.OnEvent != nil {
if text := strings.TrimSpace(choice.Message.Content); text != "" {
req.OnEvent(AgentEvent{Type: EventText, Text: text})
}
}
// Approval gate: check every requested tool BEFORE executing any, so we
// never leave a partial assistant turn (which the API rejects). On a
// pause, return the transcript as it was at loop entry (without this
// assistant turn) so a resumed run re-issues the call cleanly.
calls := make([]ToolCall, 0, len(choice.Message.ToolCalls))
for _, tc := range choice.Message.ToolCalls {
calls = append(calls, ToolCall{ID: tc.ID, Name: tc.Function.Name, Args: json.RawMessage(tc.Function.Arguments)})
}
if req.Approve != nil {
for _, call := range calls {
tool, ok := byName[call.Name]
if !ok || tool.Risk == RiskRead {
continue
}
if err := req.Approve(ctx, tool, call); err != nil {
if errors.Is(err, ErrApprovalRequired) {
result.Messages = messages
result.StopReason = "approval_required"
result.Pending = &PendingToolCall{Call: call, Risk: tool.Risk}
return result, nil
}
return nil, err
}
}
}
// Execute all tool calls and append the assistant turn + results.
assistant := AgentMessage{Role: "assistant", Content: choice.Message.Content, ToolCalls: calls}
messages = append(messages, assistant)
for _, call := range calls {
if req.OnEvent != nil {
req.OnEvent(AgentEvent{Type: EventToolStart, ToolName: call.Name, ToolArgs: call.Args})
}
out := execToolShared(ctx, byName, call)
if req.OnEvent != nil {
req.OnEvent(AgentEvent{Type: EventToolResult, ToolName: call.Name, ToolResult: out})
}
messages = append(messages, AgentMessage{Role: "tool", ToolCallID: call.ID, Content: out})
}
}
// Exhausted iterations without a final answer.
result.Messages = messages
result.StopReason = "max_iterations"
return result, nil
}
// Complete runs a single tool-less completion with an explicit system prompt.
func (p *openAIProvider) Complete(ctx context.Context, req CompletionRequest) (*WritingResult, error) {
if p == nil {
return nil, ErrNotConfigured
}
model := req.Model
if model == "" {
model = p.modelTrial
}
maxTokens := req.MaxTokens
if maxTokens <= 0 {
maxTokens = defaultAgentTokens
}
resp, err := p.complete(ctx, model, maxTokens, []oaiMessage{
{Role: "system", Content: req.System},
{Role: "user", Content: req.Prompt},
}, nil, req.Temperature)
if err != nil {
return nil, err
}
text := strings.TrimSpace(resp.Choices[0].Message.Content)
if text == "" {
return nil, errors.New("openai: empty completion")
}
return &WritingResult{Text: text, Model: model, TokensUsed: resp.Usage.TotalTokens}, nil
}
// --- WritingGenerator port ---
// GenerateWriting implements WritingGenerator over the chat-completions API,
// behaviorally identical to the prior OpenAI writing path (voice system prompt,
// single completion, writingMaxTokens cap).
func (p *openAIProvider) GenerateWriting(ctx context.Context, model, prompt string, voice VoiceContext) (*WritingResult, error) {
if p == nil {
return nil, ErrNotConfigured
}
if model == "" {
model = p.modelTrial
}
resp, err := p.complete(ctx, model, writingMaxTokens, []oaiMessage{
{Role: "system", Content: BuildVoiceRules(voice)},
{Role: "user", Content: prompt},
}, nil, nil)
if err != nil {
return nil, err
}
text := strings.TrimSpace(resp.Choices[0].Message.Content)
if text == "" {
return nil, errors.New("openai: empty completion")
}
return &WritingResult{Text: text, Model: model, TokensUsed: resp.Usage.TotalTokens}, nil
}