From bbcc62afbf704f52515518cf9ebda31dcd7ed4a7 Mon Sep 17 00:00:00 2001 From: Adarsh Prashar Date: Sun, 31 May 2026 03:28:27 +0530 Subject: [PATCH] feat: native OpenAI provider MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements the OpenAI Chat Completions API in internal/provider/openai.go (previously a stub), mirroring the Anthropic provider: maps provider.Request → OpenAI body (system prompt as a leading system message), parses the response into provider.Response with token usage, honors ctx for cancellation, and reads OPENAI_API_KEY. The proxy already routes gpt-*/o1/o3 models here, so those calls now work end to end. Removes the OpenAI stub and refreshes the now-stale "OpenAI is a stub" comments. Tests: success (auth header + system-prepend + usage), API error, missing key, context cancel. Pricing entries for common OpenAI models already exist. Fixes #15 --- internal/app/bootstrap.go | 16 ++-- internal/gateway/gateway.go | 4 +- internal/provider/openai.go | 147 +++++++++++++++++++++++++++++++ internal/provider/openai_test.go | 91 +++++++++++++++++++ internal/provider/stubs.go | 10 +-- 5 files changed, 249 insertions(+), 19 deletions(-) create mode 100644 internal/provider/openai.go create mode 100644 internal/provider/openai_test.go diff --git a/internal/app/bootstrap.go b/internal/app/bootstrap.go index 87df20e..4807f26 100644 --- a/internal/app/bootstrap.go +++ b/internal/app/bootstrap.go @@ -134,15 +134,15 @@ func OpenStore(cfg *config.Config, log *slog.Logger) (storage.Store, error) { return store, nil } -// BuildRegistry constructs the provider registry from config. Anthropic is wired -// natively when a key is present; OpenAI/Bedrock/Ollama are registered as stubs -// so config and routing can reference them. The default provider must be usable. +// BuildRegistry constructs the provider registry from config. Anthropic and +// OpenAI are implemented natively; Bedrock/Ollama are stubs config can reference. +// The default provider must be usable. func BuildRegistry(cfg *config.Config) (*provider.Registry, error) { - // Anthropic is always registered (native v0.1 provider). When the key is - // absent the daemon still boots — health and routing work — and only an actual - // Chat call returns a clear "missing API key" error. OpenAI is registered when - // a key is present; Bedrock/Ollama are stubs config can name before they're - // built out. + // Anthropic is always registered (native provider). When the key is absent the + // daemon still boots — health and routing work — and only an actual Chat call + // returns a clear "missing API key" error. OpenAI is registered (native) when a + // key is present; Bedrock/Ollama are stubs config can name before they're built + // out. ps := []provider.Provider{ provider.NewAnthropic(cfg.AnthropicAPIKey), provider.NewBedrock(), diff --git a/internal/gateway/gateway.go b/internal/gateway/gateway.go index 78b6cff..ca94c53 100644 --- a/internal/gateway/gateway.go +++ b/internal/gateway/gateway.go @@ -217,8 +217,8 @@ func (g *Gateway) resolveRun(httpReq *http.Request) *runs.Run { } // routeModel maps a model id to a provider name by prefix. An empty result routes -// to the registry's default provider. v0.1 implements Anthropic natively; OpenAI -// is a stub (front the long tail with LiteLLM as an upstream forwarder). +// to the registry's default provider. Anthropic and OpenAI are implemented +// natively; front the long tail with LiteLLM as an upstream forwarder. func routeModel(model string) string { m := strings.ToLower(model) switch { diff --git a/internal/provider/openai.go b/internal/provider/openai.go new file mode 100644 index 0000000..8370400 --- /dev/null +++ b/internal/provider/openai.go @@ -0,0 +1,147 @@ +package provider + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" + "time" +) + +// defaultOpenAIBaseURL is the OpenAI API base. Overridable for tests and for +// OpenAI-compatible gateways the user may front. +const defaultOpenAIBaseURL = "https://api.openai.com" + +// OpenAI implements Provider against the OpenAI Chat Completions API. +type OpenAI struct { + apiKey string + baseURL string + http *http.Client +} + +// NewOpenAI constructs an OpenAI provider. +func NewOpenAI(apiKey string) *OpenAI { + return &OpenAI{ + apiKey: apiKey, + baseURL: defaultOpenAIBaseURL, + http: &http.Client{Timeout: 120 * time.Second}, + } +} + +// Name returns the stable provider identifier. +func (o *OpenAI) Name() string { return "openai" } + +// --- wire types (OpenAI Chat Completions API) --- + +type openAIReq struct { + Model string `json:"model"` + Messages []openAIMessage `json:"messages"` + MaxTokens int `json:"max_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` +} + +type openAIMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type openAIResp struct { + ID string `json:"id"` + Model string `json:"model"` + Choices []struct { + Message openAIMessage `json:"message"` + FinishReason string `json:"finish_reason"` + } `json:"choices"` + Usage struct { + PromptTokens int64 `json:"prompt_tokens"` + CompletionTokens int64 `json:"completion_tokens"` + } `json:"usage"` +} + +type openAIError struct { + Error struct { + Message string `json:"message"` + Type string `json:"type"` + Code string `json:"code"` + } `json:"error"` +} + +// Chat performs one chat completion against the OpenAI Chat Completions API. +func (o *OpenAI) Chat(ctx context.Context, req Request) (*Response, error) { + if o.apiKey == "" { + return nil, fmt.Errorf("openai: missing API key") + } + + // OpenAI takes the system prompt as a leading system message. + msgs := make([]openAIMessage, 0, len(req.Messages)+1) + if req.System != "" { + msgs = append(msgs, openAIMessage{Role: string(RoleSystem), Content: req.System}) + } + for _, m := range req.Messages { + msgs = append(msgs, openAIMessage{Role: string(m.Role), Content: m.Content}) + } + + body, err := json.Marshal(openAIReq{ + Model: req.Model, + Messages: msgs, + MaxTokens: req.MaxTokens, // omitted when zero; OpenAI doesn't require it + Temperature: req.Temperature, + }) + if err != nil { + return nil, fmt.Errorf("openai: marshaling request: %w", err) + } + + httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, o.baseURL+"/v1/chat/completions", bytes.NewReader(body)) + if err != nil { + return nil, fmt.Errorf("openai: building request: %w", err) + } + httpReq.Header.Set("content-type", "application/json") + httpReq.Header.Set("authorization", "Bearer "+o.apiKey) + + resp, err := o.http.Do(httpReq) + if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return nil, ctxErr + } + return nil, fmt.Errorf("openai: request failed: %w", err) + } + defer resp.Body.Close() + + raw, err := io.ReadAll(resp.Body) + if err != nil { + return nil, fmt.Errorf("openai: reading response: %w", err) + } + + if resp.StatusCode != http.StatusOK { + var apiErr openAIError + if json.Unmarshal(raw, &apiErr) == nil && apiErr.Error.Message != "" { + return nil, fmt.Errorf("openai: %s (%s, http %d)", apiErr.Error.Message, apiErr.Error.Type, resp.StatusCode) + } + return nil, fmt.Errorf("openai: http %d: %s", resp.StatusCode, strings.TrimSpace(string(raw))) + } + + var out openAIResp + if err := json.Unmarshal(raw, &out); err != nil { + return nil, fmt.Errorf("openai: decoding response: %w", err) + } + + var content, finish string + if len(out.Choices) > 0 { + content = out.Choices[0].Message.Content + finish = out.Choices[0].FinishReason + } + + return &Response{ + ID: out.ID, + Model: out.Model, + Content: content, + FinishReason: finish, + Usage: Usage{ + PromptTokens: out.Usage.PromptTokens, + CompletionTokens: out.Usage.CompletionTokens, + }, + }, nil +} diff --git a/internal/provider/openai_test.go b/internal/provider/openai_test.go new file mode 100644 index 0000000..78f732e --- /dev/null +++ b/internal/provider/openai_test.go @@ -0,0 +1,91 @@ +package provider + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestOpenAIChat_Success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get("authorization") != "Bearer test-key" { + t.Errorf("missing/wrong auth header: %q", r.Header.Get("authorization")) + } + var got openAIReq + if err := json.NewDecoder(r.Body).Decode(&got); err != nil { + t.Fatalf("decode request: %v", err) + } + // System prompt should be the leading system message. + if len(got.Messages) != 2 || got.Messages[0].Role != "system" || got.Messages[0].Content != "be terse" { + t.Errorf("system not prepended: %+v", got.Messages) + } + if got.Messages[1].Role != "user" { + t.Errorf("unexpected messages: %+v", got.Messages) + } + w.Header().Set("content-type", "application/json") + _, _ = w.Write([]byte(`{ + "id":"chatcmpl-1","model":"gpt-4o-2024", + "choices":[{"message":{"role":"assistant","content":"hi there"},"finish_reason":"stop"}], + "usage":{"prompt_tokens":12,"completion_tokens":3,"total_tokens":15} + }`)) + })) + defer srv.Close() + + o := NewOpenAI("test-key") + o.baseURL = srv.URL + + resp, err := o.Chat(context.Background(), Request{ + Model: "gpt-4o", + System: "be terse", + Messages: []Message{{Role: RoleUser, Content: "hi"}}, + }) + if err != nil { + t.Fatalf("Chat: %v", err) + } + if resp.Content != "hi there" || resp.FinishReason != "stop" { + t.Errorf("resp = %+v", resp) + } + if resp.Usage.PromptTokens != 12 || resp.Usage.CompletionTokens != 3 || resp.Usage.Total() != 15 { + t.Errorf("usage = %+v", resp.Usage) + } + if resp.Model != "gpt-4o-2024" { + t.Errorf("model = %q", resp.Model) + } +} + +func TestOpenAIChat_APIError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.WriteHeader(http.StatusTooManyRequests) + _, _ = w.Write([]byte(`{"error":{"message":"rate limit","type":"rate_limit_error","code":"429"}}`)) + })) + defer srv.Close() + + o := NewOpenAI("k") + o.baseURL = srv.URL + _, err := o.Chat(context.Background(), Request{Model: "gpt-4o", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "rate limit") { + t.Fatalf("expected API error surfaced, got %v", err) + } +} + +func TestOpenAIChat_MissingKey(t *testing.T) { + o := NewOpenAI("") + _, err := o.Chat(context.Background(), Request{Model: "gpt-4o", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err == nil || !strings.Contains(err.Error(), "missing API key") { + t.Fatalf("expected missing-key error, got %v", err) + } +} + +func TestOpenAIChat_ContextCancel(t *testing.T) { + o := NewOpenAI("k") + o.baseURL = "http://127.0.0.1:0" + ctx, cancel := context.WithCancel(context.Background()) + cancel() + _, err := o.Chat(ctx, Request{Model: "gpt-4o", Messages: []Message{{Role: RoleUser, Content: "hi"}}}) + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } +} diff --git a/internal/provider/stubs.go b/internal/provider/stubs.go index 70f35fa..619cd48 100644 --- a/internal/provider/stubs.go +++ b/internal/provider/stubs.go @@ -6,15 +6,7 @@ import "context" // routing and config can reference them, but Chat returns ErrNotImplemented until // each is built out. For the long tail of providers, the documented path is to // front RiskKernel with LiteLLM rather than reimplement 100+ vendors here. - -// OpenAI is a stub. Native implementation is planned post-v0.1. -type OpenAI struct{ apiKey string } - -func NewOpenAI(apiKey string) *OpenAI { return &OpenAI{apiKey: apiKey} } -func (o *OpenAI) Name() string { return "openai" } -func (o *OpenAI) Chat(context.Context, Request) (*Response, error) { - return nil, ErrNotImplemented -} +// (Anthropic and OpenAI are implemented natively — see anthropic.go, openai.go.) // Bedrock is a stub. Native implementation is planned post-v0.1. type Bedrock struct{}