diff --git a/README.md b/README.md index 45e7c1b..b83b2f6 100644 --- a/README.md +++ b/README.md @@ -633,26 +633,31 @@ carry every caveat into any recommendation. The AI text is printed below a labeled rule (`πŸ€– generated by … β€” verify before acting`); if the model errors or the key is unset, the deterministic report still stands. -This is the **only** command that sends data off the machine β€” the same PII-free -Context you can see with `inspect --json`. It works with **OpenAI or Google -Gemini**, and the key is always read from the environment (never a flag). pgbot -picks the provider automatically: `OPENAI_API_KEY` β†’ OpenAI, `GEMINI_API_KEY` (or -`GOOGLE_API_KEY`) β†’ Gemini. Set `PGBOT_AI_PROVIDER=openai|gemini` to force one when -both are present. +With a remote model, this sends the same PII-free Context shown by +`inspect --json`. Before sending it, pgbot identifies the provider, host, and +model and asks for confirmation. Local endpoints are identified as local and do +not require confirmation. -``` -# OpenAI -export OPENAI_API_KEY=sk-… -pgbot explain "$DATABASE_URL" # gpt-4o-mini by default +| Provider | Key | Default model | API | +|---|---|---|---| +| Gemini | `GEMINI_API_KEY` / `GOOGLE_API_KEY` | `gemini-flash-latest` | `generateContent` | +| Anthropic | `ANTHROPIC_API_KEY` | `claude-opus-5` | `/v1/messages` | +| OpenAI | `OPENAI_API_KEY` | `gpt-5.6-terra` | `/chat/completions` | +| xAI | `XAI_API_KEY` / `GROK_API_KEY` | `grok-4.6` | `/responses` | + +The OpenAI provider also supports compatible services such as OpenRouter, +Groq, Together, DeepSeek, Mistral, Ollama, vLLM, and LM Studio. -# …or Google Gemini -export GEMINI_API_KEY=… # from Google AI Studio +``` +export OPENAI_API_KEY=… pgbot explain "$DATABASE_URL" ``` -Override the model or endpoint per provider: `PGBOT_OPENAI_MODEL` / -`PGBOT_OPENAI_URL` (any OpenAI-compatible endpoint works β€” Azure OpenAI, -OpenRouter, a local server) and `PGBOT_GEMINI_MODEL` / `PGBOT_GEMINI_URL`. +Use `PGBOT_AI_PROVIDER` to select a provider explicitly. `PGBOT_AI_MODEL`, +`PGBOT_AI_BASE_URL`, `PGBOT_AI_API_KEY`, and `PGBOT_AI_REASONING_EFFORT` +override its defaults. Existing `PGBOT_GEMINI_MODEL` and `PGBOT_GEMINI_URL` +and `PGBOT_OPENAI_MODEL` and `PGBOT_OPENAI_URL` settings remain supported. Keys +are read only from environment variables. **Exit codes** (a stable contract for CI): `0` clean Β· `1` warnings Β· `2` critical findings Β· `3` connection/execution failure Β· `64` usage error (bad flags/args). diff --git a/cmd/pgbot/ask.go b/cmd/pgbot/ask.go index 2ea30ec..8a014c0 100644 --- a/cmd/pgbot/ask.go +++ b/cmd/pgbot/ask.go @@ -25,8 +25,8 @@ func newAskCmd() *cobra.Command { Short: "Ask an AI about your database, grounded on pgbot's findings", Long: "Runs the same read-only inspection, then answers your question using ONLY the\n" + "deterministic findings (the model can't reach into the database). Connection\n" + - "comes from --url or $DATABASE_URL. Sends the PII-free findings to an AI provider β€”\n" + - "set $OPENAI_API_KEY (OpenAI) or $GEMINI_API_KEY (Google Gemini).", + "comes from --url or $DATABASE_URL. Sends the PII-free findings to the model you\n" + + "configured β€” Gemini, Anthropic, OpenAI, xAI, or an OpenAI-compatible endpoint.", Args: cobra.MinimumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { return runAsk(cmd, strings.Join(args, " "), url, f, yes) @@ -38,17 +38,16 @@ func newAskCmd() *cobra.Command { fl.IntVar(&f.ashHz, "ash-hz", 10, "active-session sampling rate in Hz (0 disables)") fl.BoolVar(&f.noStore, "no-store", false, "do not read or write the local baseline store") fl.BoolVar(&f.strictPooler, "strict-pooler", false, "refuse (exit 3) behind a transaction pooler") - fl.BoolVar(&yes, "yes", false, "skip the 'this sends data to the AI provider' confirmation prompt") + fl.BoolVar(&yes, "yes", false, "skip the data-disclosure confirmation prompt") return cmd } func runAsk(cmd *cobra.Command, question, url string, f inspectFlags, yes bool) error { - client, err := ai.NewFromEnv() + llm, err := ai.Resolve() if err != nil { return err } - fmt.Fprintf(os.Stderr, "pgbot ask: this sends the PII-free findings to %s (model %s).\n", client.Vendor(), client.ModelName()) - if !yes && isInteractive() && !confirm() { + if !confirmDisclosure("pgbot ask", llm, yes) { return fmt.Errorf("aborted") } @@ -66,8 +65,12 @@ func runAsk(cmd *cobra.Command, question, url string, f inspectFlags, yes bool) return err } - answer, aiErr := ai.Ask(ctx, client, c, question) - printAnswer(useColor(false), client.ModelName(), answer, aiErr) + // Give the model its own deadline instead of the remainder of collection's + // budget. Local models may need substantially longer than hosted providers. + aiCtx, aiCancel := context.WithTimeout(cmd.Context(), 3*time.Minute) + defer aiCancel() + answer, aiErr := ai.Ask(aiCtx, llm, c, question) + printAnswer(useColor(false), llm.Model(), answer, aiErr) if aiErr == nil { // `ask` prints only the model's prose, so a destructive-action guard that // the model may have reworded away must be reasserted here, verbatim from @@ -126,12 +129,3 @@ func printAnswer(color bool, modelName, text string, aiErr error) { fmt.Println() fmt.Println(st.Dim("β€” " + modelName + " Β· a reading of pgbot's findings; verify before acting")) } - -// confirm reads a y/N from stdin. -func confirm() bool { - fmt.Fprint(os.Stderr, "Continue? [y/N] ") - var resp string - fmt.Fscanln(os.Stdin, &resp) - r := strings.ToLower(strings.TrimSpace(resp)) - return r == "y" || r == "yes" -} diff --git a/cmd/pgbot/explain.go b/cmd/pgbot/explain.go index bbe1322..5ed32fe 100644 --- a/cmd/pgbot/explain.go +++ b/cmd/pgbot/explain.go @@ -26,12 +26,12 @@ func newExplainCmd() *cobra.Command { Use: "explain ", Short: "Inspect, then have an AI explain the findings in plain language", Long: "Runs the same read-only inspection as `pgbot inspect`, prints the deterministic\n" + - "report, then sends the PII-free findings to an AI provider for a plain-language\n" + + "report, then sends the PII-free findings to a model for a plain-language\n" + "explanation. The findings are still computed locally in Go β€” the model only\n" + "explains them, never invents them.\n\n" + - "The key is read from $OPENAI_API_KEY (OpenAI) or $GEMINI_API_KEY (Google Gemini),\n" + - "never a flag; set PGBOT_AI_PROVIDER to force one. This is the only pgbot command\n" + - "that sends data off the machine; the payload is the same PII-free Context you can\n" + + "Configure Gemini, Anthropic, OpenAI, xAI, or an OpenAI-compatible endpoint. Keys\n" + + "are read from the environment, never a flag. With a remote model, the payload is\n" + + "the same PII-free Context you can\n" + "inspect with `pgbot inspect --json`.", Args: cobra.MaximumNArgs(1), RunE: func(cmd *cobra.Command, args []string) error { @@ -46,7 +46,7 @@ func newExplainCmd() *cobra.Command { fl.BoolVar(&f.strictPooler, "strict-pooler", false, "refuse (exit 3) if connected through a transaction pooler") fl.IntVar(&f.ashHz, "ash-hz", 10, "active-session sampling rate in Hz (0 disables the wait-event profile)") fl.DurationVar(&f.window, "window", 5*time.Second, "active-session sampling window") - fl.BoolVar(&yes, "yes", false, "skip the 'this sends data to the AI provider' confirmation prompt") + fl.BoolVar(&yes, "yes", false, "skip the data-disclosure confirmation prompt") fl.StringVar(&f.config, "config", "", "path to .pgbot.toml (default: discover from cwd upward)") fl.StringArrayVar(&f.ignore, "ignore", nil, "suppress a finding for this run: finding[:object] (repeatable)") fl.StringVar(&f.failOn, "fail-on", "warn", "exit non-zero on findings at/above this severity: critical|warn|info|none") @@ -59,21 +59,12 @@ func runExplain(cmd *cobra.Command, args []string, f inspectFlags, yes bool) err } // Build the model client first β€” fail fast before we connect if the key is // missing, so the user isn't surprised after a full inspection. - client, err := ai.NewFromEnv() + llm, err := ai.Resolve() if err != nil { return err } - - // This is the one command that sends data off the box. Say so, loudly, and - // require an explicit go-ahead unless --yes (or non-interactive). - fmt.Fprintf(os.Stderr, "pgbot explain: this sends the PII-free findings (same as `inspect --json`) to %s (model %s).\n", client.Vendor(), client.ModelName()) - if !yes && isInteractive() { - fmt.Fprint(os.Stderr, "Continue? [y/N] ") - var resp string - fmt.Fscanln(os.Stdin, &resp) - if r := strings.ToLower(strings.TrimSpace(resp)); r != "y" && r != "yes" { - return fmt.Errorf("aborted") - } + if !confirmDisclosure("pgbot explain", llm, yes) { + return fmt.Errorf("aborted") } connString := firstNonEmpty(argAt(args, 0), os.Getenv("DATABASE_URL"), os.Getenv("PGBOT_DATABASE_URL")) @@ -126,8 +117,12 @@ func runExplain(cmd *cobra.Command, args []string, f inspectFlags, yes bool) err // 2. The AI explanation β€” clearly labeled as model-generated. If it fails, the // deterministic report above still stands; we just note the explanation is // unavailable and exit on the findings' code. - explanation, aiErr := ai.Explain(ctx, client, c) - printAISection(color, client.ModelName(), explanation, aiErr) + // Give the model its own deadline instead of the remainder of collection's + // budget. Local models may need substantially longer than hosted providers. + aiCtx, aiCancel := context.WithTimeout(cmd.Context(), 3*time.Minute) + defer aiCancel() + explanation, aiErr := ai.Explain(aiCtx, llm, c) + printAISection(color, llm.Model(), explanation, aiErr) // The destructive-action guards, reasserted by code AFTER the model text β€” so a // reworded or truncated explanation can never be the thing that drops them. if aiErr == nil { @@ -138,6 +133,26 @@ func runExplain(cmd *cobra.Command, args []string, f inspectFlags, yes bool) err return nil } +// confirmDisclosure identifies the remote destination before data is sent. +// Local endpoints do not send findings off the machine and need no confirmation. +func confirmDisclosure(cmdName string, llm ai.LanguageModel, yes bool) bool { + if ai.Local(llm.Endpoint()) { + fmt.Fprintf(os.Stderr, "%s: using a local model at %s (%s) β€” the findings do not leave this machine.\n", + cmdName, ai.Host(llm.Endpoint()), llm.Model()) + return true + } + fmt.Fprintf(os.Stderr, "%s: this sends the PII-free findings (same as `inspect --json`) to %s at %s (model %s).\n", + cmdName, llm.Provider(), ai.Host(llm.Endpoint()), llm.Model()) + if yes || !isInteractive() { + return true + } + fmt.Fprint(os.Stderr, "Continue? [y/N] ") + var resp string + fmt.Fscanln(os.Stdin, &resp) + r := strings.ToLower(strings.TrimSpace(resp)) + return r == "y" || r == "yes" +} + // printAISection renders the labeled AI block. The banner makes it unmistakable // that this text is model-generated and must be verified before acting. func printAISection(color bool, modelName, text string, aiErr error) { diff --git a/internal/ai/anthropic.go b/internal/ai/anthropic.go new file mode 100644 index 0000000..9ca7792 --- /dev/null +++ b/internal/ai/anthropic.go @@ -0,0 +1,144 @@ +package ai + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +const ( + // The most capable current model. `pgbot explain` is a one-shot call on a + // small payload, so this is cheap in absolute terms β€” but set + // $PGBOT_AI_MODEL=claude-haiku-4-5 if you run it in a tight loop. + defaultAnthropicModel = "claude-opus-5" + defaultAnthropicURL = "https://api.anthropic.com" + anthropicVersion = "2023-06-01" +) + +// AnthropicProvider talks to the Messages API. +type AnthropicProvider struct { + APIKey string + BaseURL string + HTTP *http.Client +} + +func (p *AnthropicProvider) Name() string { return "anthropic" } + +func (p *AnthropicProvider) LanguageModel(_ context.Context, modelID string) (LanguageModel, error) { + if modelID == "" { + modelID = defaultAnthropicModel + } + return &anthropicModel{provider: p, model: modelID}, nil +} + +type anthropicModel struct { + provider *AnthropicProvider + model string +} + +func (m *anthropicModel) Provider() string { return "anthropic" } +func (m *anthropicModel) Model() string { return m.model } +func (m *anthropicModel) Endpoint() string { return m.provider.BaseURL } + +// ---- wire types (only the fields we use) ---- + +type messagesRequest struct { + Model string `json:"model"` + System string `json:"system,omitempty"` + // Required by the API, and it caps thinking + visible text together: current + // models think before they answer, so a tight value truncates the explanation + // mid-sentence. Same headroom, same reason, as the Gemini path. + MaxTokens int `json:"max_tokens"` + Messages []anthropicMessage `json:"messages"` + // Deliberately no `temperature`: it is REMOVED on current models (Opus 5, + // Sonnet 5, Opus 4.7+) and sending it returns a 400. Call.Temperature is a + // hint, and this provider ignores it. +} + +type anthropicMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type messagesResponse struct { + Content []struct { + Type string `json:"type"` + Text string `json:"text"` + } `json:"content"` + StopReason string `json:"stop_reason"` + StopDetails *struct { + Category string `json:"category"` + } `json:"stop_details"` + Error *struct { + Type string `json:"type"` + Message string `json:"message"` + } `json:"error"` +} + +// Generate sends one system + user turn and returns the model's text. No retries +// β€” a failed explanation must not hang the CLI. +func (m *anthropicModel) Generate(ctx context.Context, c Call) (*Response, error) { + maxTokens := 8192 + if c.MaxOutputTokens != nil { + maxTokens = int(*c.MaxOutputTokens) + } + buf, err := json.Marshal(messagesRequest{ + Model: m.model, + System: c.System, + MaxTokens: maxTokens, + Messages: []anthropicMessage{{Role: "user", Content: c.Prompt}}, + }) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, m.provider.BaseURL+"/v1/messages", bytes.NewReader(buf)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("anthropic-version", anthropicVersion) + req.Header.Set("x-api-key", m.provider.APIKey) // header, never a query param + + resp, err := m.provider.HTTP.Do(req) + if err != nil { + return nil, fmt.Errorf("calling Anthropic: %w", err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + + var mr messagesResponse + if err := json.Unmarshal(body, &mr); err != nil { + return nil, fmt.Errorf("anthropic returned unparseable response (HTTP %d)", resp.StatusCode) + } + if mr.Error != nil { + return nil, fmt.Errorf("anthropic error (%s): %s", mr.Error.Type, mr.Error.Message) + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("anthropic HTTP %d", resp.StatusCode) + } + // A refusal is a successful HTTP 200 with empty content β€” check it before + // reading the blocks, or it looks like an unexplained empty answer. + if mr.StopReason == "refusal" { + reason := "safety" + if mr.StopDetails != nil && mr.StopDetails.Category != "" { + reason = mr.StopDetails.Category + } + return nil, fmt.Errorf("anthropic declined the prompt (%s)", reason) + } + var sb strings.Builder + for _, blk := range mr.Content { + if blk.Type == "text" { + sb.WriteString(blk.Text) + } + } + out := strings.TrimSpace(sb.String()) + if out == "" { + return nil, fmt.Errorf("anthropic returned an empty explanation (finish: %s)", mr.StopReason) + } + return &Response{Text: out, FinishReason: mr.StopReason}, nil +} diff --git a/internal/ai/anthropic_test.go b/internal/ai/anthropic_test.go new file mode 100644 index 0000000..3db319c --- /dev/null +++ b/internal/ai/anthropic_test.go @@ -0,0 +1,114 @@ +package ai + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func anthropicFor(url, key, model string) LanguageModel { + p := &AnthropicProvider{APIKey: key, BaseURL: url, HTTP: http.DefaultClient} + m, _ := p.LanguageModel(context.Background(), model) + return m +} + +func TestAnthropic_success(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("x-api-key"); got != "k123" { + t.Errorf("key must travel in the x-api-key header, got %q", got) + } + if got := r.Header.Get("anthropic-version"); got != anthropicVersion { + t.Errorf("anthropic-version header missing or wrong: %q", got) + } + if strings.Contains(r.URL.String(), "k123") { + t.Error("API key leaked into the URL") + } + if r.URL.Path != "/v1/messages" { + t.Errorf("unexpected path %s", r.URL.Path) + } + body, _ := io.ReadAll(r.Body) + + // temperature is REMOVED on current Anthropic models β€” sending it is a 400, + // so it must never appear even though Call carries one. + var raw map[string]any + if err := json.Unmarshal(body, &raw); err != nil { + t.Fatalf("bad request body: %v", err) + } + if _, ok := raw["temperature"]; ok { + t.Error("temperature must not be sent to Anthropic β€” current models reject it") + } + if raw["max_tokens"] == nil { + t.Error("max_tokens is required by the Messages API") + } + if raw["system"] != "be terse" { + t.Errorf("system prompt not sent as a top-level field: %v", raw["system"]) + } + + json.NewEncoder(w).Encode(map[string]any{ + "content": []map[string]string{{"type": "text", "text": "Looks healthy."}}, + "stop_reason": "end_turn", + }) + })) + defer srv.Close() + + out, err := anthropicFor(srv.URL, "k123", "m1").Generate(context.Background(), Call{ + System: "be terse", Prompt: "the report", Temperature: f64(0.2), + }) + if err != nil { + t.Fatal(err) + } + if out.Text != "Looks healthy." { + t.Errorf("unexpected output %q", out.Text) + } +} + +func TestAnthropic_apiError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + json.NewEncoder(w).Encode(map[string]any{ + "type": "error", + "error": map[string]string{"type": "authentication_error", "message": "invalid x-api-key"}, + }) + })) + defer srv.Close() + + _, err := anthropicFor(srv.URL, "bad", "m1").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "invalid x-api-key") { + t.Errorf("expected a surfaced API error, got %v", err) + } +} + +// A refusal is a successful HTTP 200 with empty content β€” it must not read as an +// unexplained blank answer. +func TestAnthropic_refusal(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "content": []any{}, + "stop_reason": "refusal", + "stop_details": map[string]string{"type": "refusal", "category": "cyber"}, + }) + })) + defer srv.Close() + + _, err := anthropicFor(srv.URL, "k", "m").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "declined") { + t.Errorf("a refusal must surface as an error naming the reason, got %v", err) + } + if err != nil && !strings.Contains(err.Error(), "cyber") { + t.Errorf("refusal category should be reported, got %v", err) + } +} + +func TestAnthropic_emptyContent(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"content": []any{}, "stop_reason": "max_tokens"}) + })) + defer srv.Close() + if _, err := anthropicFor(srv.URL, "k", "m").Generate(context.Background(), Call{Prompt: "x"}); err == nil { + t.Error("empty content should be an error, not an empty success") + } +} diff --git a/internal/ai/explain.go b/internal/ai/explain.go index db085f4..5ab86a1 100644 --- a/internal/ai/explain.go +++ b/internal/ai/explain.go @@ -117,22 +117,36 @@ func BuildAskPrompt(c *model.Context, question string) (system, user string) { // Explain builds the prompt and calls the model, returning the labeled-elsewhere // explanation text. A nil/empty findings set still gets an explanation (the model -// is told to confirm health briefly). The Provider may be OpenAI or Gemini. -func Explain(ctx context.Context, p Provider, mc *model.Context) (string, error) { - if p == nil { - return "", fmt.Errorf("no AI client") - } +// is told to confirm health briefly). +func Explain(ctx context.Context, m LanguageModel, mc *model.Context) (string, error) { system, user := BuildExplainPrompt(mc) - return p.Generate(ctx, system, user) + return generate(ctx, m, system, user) } // Ask answers a specific question grounded on the report. -func Ask(ctx context.Context, p Provider, mc *model.Context, question string) (string, error) { - if p == nil { - return "", fmt.Errorf("no AI client") - } +func Ask(ctx context.Context, m LanguageModel, mc *model.Context, question string) (string, error) { system, user := BuildAskPrompt(mc, question) - return p.Generate(ctx, system, user) + return generate(ctx, m, system, user) +} + +// generate is the one place the explanation layer meets a provider. Low +// temperature because we want the model reading the findings, not riffing on +// them; 8192 output tokens because current models spend part of that budget +// thinking before the visible text. +func generate(ctx context.Context, m LanguageModel, system, user string) (string, error) { + if m == nil { + return "", fmt.Errorf("no model configured") + } + resp, err := m.Generate(ctx, Call{ + System: system, + Prompt: user, + Temperature: f64(0.2), + MaxOutputTokens: i64(8192), + }) + if err != nil { + return "", err + } + return resp.Text, nil } func topBuckets(b []model.WaitBucket, n int) []model.WaitBucket { diff --git a/internal/ai/gemini.go b/internal/ai/gemini.go index eaa981f..4928f56 100644 --- a/internal/ai/gemini.go +++ b/internal/ai/gemini.go @@ -1,10 +1,3 @@ -// Package ai is pgbot's OPTIONAL explanation layer. It never produces findings β€” -// those are computed deterministically in package findings. It only takes an -// already-computed, PII-free Context and asks a model to explain and prioritize -// it in plain language. Everything it emits is labeled as model-generated. -// -// It talks to Google's Generative Language REST API directly over net/http so -// pgbot keeps its single-static-binary, minimal-dependency promise (no SDK). package ai import ( @@ -14,60 +7,44 @@ import ( "fmt" "io" "net/http" - "os" "strings" - "time" ) const ( // A moving alias, on purpose: pinned versions get retired (e.g. gemini-2.5-flash // now 404s for new projects), and an explanation feature doesn't need a frozen - // model. Override with $PGBOT_GEMINI_MODEL to pin one. - defaultModel = "gemini-flash-latest" - defaultBaseURL = "https://generativelanguage.googleapis.com/v1beta" + // model. Override with $PGBOT_AI_MODEL to pin one. + defaultGeminiModel = "gemini-flash-latest" + defaultGeminiURL = "https://generativelanguage.googleapis.com/v1beta" ) -// Client is a minimal Gemini generateContent client. -type Client struct { +// GeminiProvider talks to Google's Generative Language REST API. Both AI-Studio +// "auth" keys (AQ.…) and legacy standard keys (AIza…) work β€” they travel in the +// same header. +type GeminiProvider struct { APIKey string - Model string BaseURL string HTTP *http.Client } -// NewGeminiFromEnv builds a client from the environment. The key comes ONLY from -// GEMINI_API_KEY (or GOOGLE_API_KEY, the SDK's other convention) β€” never a flag, -// so it can't leak into shell history or the process list. PGBOT_GEMINI_MODEL and -// PGBOT_GEMINI_URL override the defaults. Both AI-Studio "auth" keys (AQ.…) and -// legacy standard keys (AIza…) work β€” they travel in the same header. -func NewGeminiFromEnv() (*Client, error) { - key := strings.TrimSpace(os.Getenv("GEMINI_API_KEY")) - if key == "" { - key = strings.TrimSpace(os.Getenv("GOOGLE_API_KEY")) - } - if key == "" { - return nil, fmt.Errorf("GEMINI_API_KEY (or GOOGLE_API_KEY) is not set β€” export it to use `pgbot explain` (the key is never read from a flag)") +func (p *GeminiProvider) Name() string { return "gemini" } + +func (p *GeminiProvider) LanguageModel(_ context.Context, modelID string) (LanguageModel, error) { + if modelID == "" { + modelID = defaultGeminiModel } - model := envOr("PGBOT_GEMINI_MODEL", defaultModel) - base := envOr("PGBOT_GEMINI_URL", defaultBaseURL) - return &Client{ - APIKey: key, - Model: model, - BaseURL: strings.TrimRight(base, "/"), - HTTP: &http.Client{Timeout: 60 * time.Second}, - }, nil + return &geminiModel{provider: p, model: modelID}, nil } -func (c *Client) ModelName() string { return c.Model } -func (c *Client) Vendor() string { return "Google Gemini" } - -func envOr(name, def string) string { - if v := strings.TrimSpace(os.Getenv(name)); v != "" { - return v - } - return def +type geminiModel struct { + provider *GeminiProvider + model string } +func (m *geminiModel) Provider() string { return "gemini" } +func (m *geminiModel) Model() string { return m.model } +func (m *geminiModel) Endpoint() string { return m.provider.BaseURL } + // ---- wire types (only the fields we use) ---- type genRequest struct { @@ -110,52 +87,59 @@ type apiError struct { // Generate sends one system + user turn and returns the model's text. It never // retries with backoff loops β€” a failed explanation must not hang the CLI; the // caller degrades to the deterministic report alone. -func (c *Client) Generate(ctx context.Context, system, user string) (string, error) { +func (m *geminiModel) Generate(ctx context.Context, c Call) (*Response, error) { + cfg := generationConfig{Temperature: 0.2, MaxOutputTokens: 8192} + if c.Temperature != nil { + cfg.Temperature = *c.Temperature + } + if c.MaxOutputTokens != nil { + cfg.MaxOutputTokens = int(*c.MaxOutputTokens) + } reqBody := genRequest{ - Contents: []content{{Role: "user", Parts: []part{{Text: user}}}}, + Contents: []content{{Role: "user", Parts: []part{{Text: c.Prompt}}}}, // Headroom for the answer AND for the internal "thinking" that current // Gemini models spend before the visible text β€” a tight cap truncates the // explanation mid-sentence. - GenerationConfig: generationConfig{Temperature: 0.2, MaxOutputTokens: 8192}, + GenerationConfig: cfg, } - if system != "" { - reqBody.SystemInstruction = &content{Parts: []part{{Text: system}}} + if c.System != "" { + reqBody.SystemInstruction = &content{Parts: []part{{Text: c.System}}} } buf, err := json.Marshal(reqBody) if err != nil { - return "", err + return nil, err } - url := fmt.Sprintf("%s/models/%s:generateContent", c.BaseURL, c.Model) + url := fmt.Sprintf("%s/models/%s:generateContent", m.provider.BaseURL, m.model) req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(buf)) if err != nil { - return "", err + return nil, err } req.Header.Set("Content-Type", "application/json") - req.Header.Set("x-goog-api-key", c.APIKey) // header, not ?key= β€” keeps the key out of any URL log + req.Header.Set("x-goog-api-key", m.provider.APIKey) // header, not ?key= β€” keeps the key out of any URL log - resp, err := c.HTTP.Do(req) + resp, err := m.provider.HTTP.Do(req) if err != nil { - return "", fmt.Errorf("calling Gemini: %w", err) + return nil, fmt.Errorf("calling Gemini: %w", err) } defer resp.Body.Close() body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) var gr genResponse if err := json.Unmarshal(body, &gr); err != nil { - return "", fmt.Errorf("gemini returned unparseable response (HTTP %d)", resp.StatusCode) + return nil, fmt.Errorf("gemini returned unparseable response (HTTP %d)", resp.StatusCode) } if gr.Error != nil { - return "", fmt.Errorf("gemini error (%d %s): %s", gr.Error.Code, gr.Error.Status, gr.Error.Message) + return nil, fmt.Errorf("gemini error (%d %s): %s", gr.Error.Code, gr.Error.Status, gr.Error.Message) } if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("gemini HTTP %d", resp.StatusCode) + return nil, fmt.Errorf("gemini HTTP %d", resp.StatusCode) } if gr.PromptFeedback != nil && gr.PromptFeedback.BlockReason != "" { - return "", fmt.Errorf("gemini blocked the prompt: %s", gr.PromptFeedback.BlockReason) + return nil, fmt.Errorf("gemini blocked the prompt: %s", gr.PromptFeedback.BlockReason) } if len(gr.Candidates) == 0 { - return "", fmt.Errorf("gemini returned no candidates") + return nil, fmt.Errorf("gemini returned no candidates") } var sb strings.Builder for _, p := range gr.Candidates[0].Content.Parts { @@ -163,7 +147,7 @@ func (c *Client) Generate(ctx context.Context, system, user string) (string, err } out := strings.TrimSpace(sb.String()) if out == "" { - return "", fmt.Errorf("gemini returned an empty explanation (finish: %s)", gr.Candidates[0].FinishReason) + return nil, fmt.Errorf("gemini returned an empty explanation (finish: %s)", gr.Candidates[0].FinishReason) } - return out, nil + return &Response{Text: out, FinishReason: gr.Candidates[0].FinishReason}, nil } diff --git a/internal/ai/gemini_test.go b/internal/ai/gemini_test.go index 97ad85f..0215ccd 100644 --- a/internal/ai/gemini_test.go +++ b/internal/ai/gemini_test.go @@ -10,33 +10,14 @@ import ( "testing" ) -func TestNewGeminiFromEnv(t *testing.T) { - t.Setenv("GEMINI_API_KEY", "") - t.Setenv("GOOGLE_API_KEY", "") - if _, err := NewGeminiFromEnv(); err == nil { - t.Error("missing key must be an error") - } - // GOOGLE_API_KEY is accepted as a fallback (the SDK's other convention). - t.Setenv("GOOGLE_API_KEY", "fromgoogle") - if c, err := NewGeminiFromEnv(); err != nil || c.APIKey != "fromgoogle" { - t.Errorf("GOOGLE_API_KEY fallback not honored: %v / %+v", err, c) - } - t.Setenv("GOOGLE_API_KEY", "") - t.Setenv("GEMINI_API_KEY", "secret") - t.Setenv("PGBOT_GEMINI_MODEL", "gemini-3-pro") - c, err := NewGeminiFromEnv() - if err != nil { - t.Fatal(err) - } - if c.APIKey != "secret" || c.Model != "gemini-3-pro" { - t.Errorf("env not honored: %+v", c) - } - if c.BaseURL != defaultBaseURL { - t.Errorf("default base URL wrong: %s", c.BaseURL) - } +// geminiFor builds a model pointed at a test server, bypassing the environment. +func geminiFor(url, key, model string) LanguageModel { + p := &GeminiProvider{APIKey: key, BaseURL: url, HTTP: http.DefaultClient} + m, _ := p.LanguageModel(context.Background(), model) + return m } -func TestGenerate_success(t *testing.T) { +func TestGemini_success(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if got := r.Header.Get("x-goog-api-key"); got != "k123" { t.Errorf("API key must travel in the x-goog-api-key header, got %q", got) @@ -63,37 +44,44 @@ func TestGenerate_success(t *testing.T) { })) defer srv.Close() - c := &Client{APIKey: "k123", Model: "m1", BaseURL: srv.URL, HTTP: srv.Client()} - out, err := c.Generate(context.Background(), "be terse", "the report") + out, err := geminiFor(srv.URL, "k123", "m1").Generate(context.Background(), Call{System: "be terse", Prompt: "the report"}) if err != nil { t.Fatal(err) } - if out != "Looks healthy." { - t.Errorf("unexpected output %q", out) + if out.Text != "Looks healthy." { + t.Errorf("unexpected output %q", out.Text) } } -func TestGenerate_apiError(t *testing.T) { +func TestGemini_apiError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusBadRequest) json.NewEncoder(w).Encode(genResponse{Error: &apiError{Code: 400, Status: "INVALID_ARGUMENT", Message: "API key not valid"}}) })) defer srv.Close() - c := &Client{APIKey: "bad", Model: "m1", BaseURL: srv.URL, HTTP: srv.Client()} - _, err := c.Generate(context.Background(), "", "x") + _, err := geminiFor(srv.URL, "bad", "m1").Generate(context.Background(), Call{Prompt: "x"}) if err == nil || !strings.Contains(err.Error(), "API key not valid") { t.Errorf("expected a surfaced API error, got %v", err) } } -func TestGenerate_noCandidates(t *testing.T) { +func TestGemini_noCandidates(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { json.NewEncoder(w).Encode(genResponse{}) // 200, no candidates })) defer srv.Close() - c := &Client{APIKey: "k", Model: "m", BaseURL: srv.URL, HTTP: srv.Client()} - if _, err := c.Generate(context.Background(), "", "x"); err == nil { + if _, err := geminiFor(srv.URL, "k", "m").Generate(context.Background(), Call{Prompt: "x"}); err == nil { t.Error("no candidates should be an error, not an empty success") } } + +func TestGemini_defaultModel(t *testing.T) { + m := geminiFor("https://example.test", "k", "") + if m.Model() != defaultGeminiModel { + t.Errorf("empty model id should fall back to the default, got %q", m.Model()) + } + if m.Provider() != "gemini" { + t.Errorf("provider name wrong: %q", m.Provider()) + } +} diff --git a/internal/ai/openai.go b/internal/ai/openai.go index 5b8818b..0bd34d1 100644 --- a/internal/ai/openai.go +++ b/internal/ai/openai.go @@ -7,54 +7,86 @@ import ( "fmt" "io" "net/http" - "os" "strings" - "time" ) const ( - // A small, cheap chat model is plenty for an explanation task. Override with - // $PGBOT_OPENAI_MODEL to pin a bigger or newer one. - defaultOpenAIModel = "gpt-4o-mini" + defaultOpenAIModel = "gpt-5.6-terra" defaultOpenAIURL = "https://api.openai.com/v1" + + // Applied only to reasoning models. Override with $PGBOT_AI_REASONING_EFFORT + // (none, low, medium, high, xhigh, max). + defaultReasoningEffort = "xhigh" + + // A reasoning model's token cap covers hidden reasoning AND the visible + // answer. At xhigh the reasoning alone can exceed the 8192 that suffices for + // a non-reasoning model, leaving an empty response with finish_reason + // "length" β€” so give reasoning models a floor well above it. + reasoningTokenFloor = 32000 ) -// OpenAIClient is a minimal Chat Completions client. Because it speaks the plain -// OpenAI wire format, it also works against any OpenAI-compatible endpoint via -// $PGBOT_OPENAI_URL (Azure OpenAI, OpenRouter, a local server, …). -type OpenAIClient struct { +// OpenAIProvider speaks /chat/completions. That one wire format covers far more +// than OpenAI: OpenRouter, Groq, Together, Fireworks, DeepSeek, xAI, Mistral, and +// every local server (Ollama, vLLM, LM Studio) accept it, so pointing +// $PGBOT_AI_BASE_URL at any of them is the whole integration. +type OpenAIProvider struct { APIKey string - Model string BaseURL string HTTP *http.Client + + // ReasoningEffort is sent only to models that accept it; empty means the + // provider's own default. + ReasoningEffort string } -// NewOpenAIFromEnv builds a client from $OPENAI_API_KEY. $PGBOT_OPENAI_MODEL and -// $PGBOT_OPENAI_URL override the defaults. The key is read from the environment -// only β€” never a flag β€” so it stays out of shell history and the process list. -func NewOpenAIFromEnv() (*OpenAIClient, error) { - key := strings.TrimSpace(os.Getenv("OPENAI_API_KEY")) - if key == "" { - return nil, fmt.Errorf("OPENAI_API_KEY is not set β€” export it to use `pgbot explain`/`ask` with OpenAI (the key is never read from a flag)") +func (p *OpenAIProvider) Name() string { return "openai" } + +func (p *OpenAIProvider) LanguageModel(_ context.Context, modelID string) (LanguageModel, error) { + if modelID == "" { + modelID = defaultOpenAIModel } - return &OpenAIClient{ - APIKey: key, - Model: envOr("PGBOT_OPENAI_MODEL", defaultOpenAIModel), - BaseURL: strings.TrimRight(envOr("PGBOT_OPENAI_URL", defaultOpenAIURL), "/"), - HTTP: &http.Client{Timeout: 60 * time.Second}, - }, nil + return &openaiModel{provider: p, model: modelID}, nil } -func (c *OpenAIClient) ModelName() string { return c.Model } -func (c *OpenAIClient) Vendor() string { return "OpenAI" } +type openaiModel struct { + provider *OpenAIProvider + model string +} + +func (m *openaiModel) Provider() string { return "openai" } +func (m *openaiModel) Model() string { return m.model } +func (m *openaiModel) Endpoint() string { return m.provider.BaseURL } // ---- wire types (only the fields we use) ---- type chatRequest struct { - Model string `json:"model"` - Messages []chatMessage `json:"messages"` - Temperature float64 `json:"temperature,omitempty"` - MaxCompletionTokens int `json:"max_completion_tokens,omitempty"` + Model string `json:"model"` + Messages []chatMessage `json:"messages"` + Temperature *float64 `json:"temperature,omitempty"` + // Exactly one of these is set. OpenAI's reasoning models reject max_tokens + // outright ("use max_completion_tokens"), while every other compatible + // server (Ollama, Groq, vLLM, LM Studio) understands only max_tokens β€” and + // those are the reason this provider exists. reasoningModel picks. + MaxTokens *int64 `json:"max_tokens,omitempty"` + MaxCompletionTokens *int64 `json:"max_completion_tokens,omitempty"` + ReasoningEffort string `json:"reasoning_effort,omitempty"` +} + +// reasoningModel reports whether a model id names an OpenAI-family reasoning +// model, which takes a different request shape from a plain chat model. Keyed on +// the id rather than the endpoint so it also fires for the same models reached +// through OpenRouter ("openai/gpt-5.6-terra") or Azure. +func reasoningModel(id string) bool { + id = strings.ToLower(id) + if i := strings.LastIndex(id, "/"); i >= 0 { // strip an "openai/" vendor prefix + id = id[i+1:] + } + for _, p := range []string{"gpt-5", "o1", "o3", "o4"} { + if strings.HasPrefix(id, p) { + return true + } + } + return false } type chatMessage struct { @@ -67,66 +99,86 @@ type chatResponse struct { Message chatMessage `json:"message"` FinishReason string `json:"finish_reason"` } `json:"choices"` - Error *struct { - Message string `json:"message"` - Type string `json:"type"` - Code string `json:"code"` - } `json:"error"` } -// Generate sends one system + user turn and returns the model's text. Like the -// Gemini client it does not retry β€” a failed explanation must not hang the CLI; -// the caller degrades to the deterministic report alone. -func (c *OpenAIClient) Generate(ctx context.Context, system, user string) (string, error) { +// Generate sends one system + user turn and returns the model's text. Like every +// provider here it never retries β€” a failed explanation must not hang the CLI. +func (m *openaiModel) Generate(ctx context.Context, c Call) (*Response, error) { msgs := make([]chatMessage, 0, 2) - if system != "" { - msgs = append(msgs, chatMessage{Role: "system", Content: system}) + if c.System != "" { + msgs = append(msgs, chatMessage{Role: "system", Content: c.System}) + } + msgs = append(msgs, chatMessage{Role: "user", Content: c.Prompt}) + + reqBody := chatRequest{Model: m.model, Messages: msgs} + if reasoningModel(m.model) { + // No temperature: the reasoning families reject sampling parameters, and + // the low temperature we ask for elsewhere buys nothing on a model that + // deliberates before answering. + limit := int64(reasoningTokenFloor) + if c.MaxOutputTokens != nil && *c.MaxOutputTokens > limit { + limit = *c.MaxOutputTokens + } + reqBody.MaxCompletionTokens = &limit + reqBody.ReasoningEffort = m.provider.ReasoningEffort + if reqBody.ReasoningEffort == "" { + reqBody.ReasoningEffort = defaultReasoningEffort + } + } else { + reqBody.Temperature = c.Temperature + reqBody.MaxTokens = c.MaxOutputTokens } - msgs = append(msgs, chatMessage{Role: "user", Content: user}) - - buf, err := json.Marshal(chatRequest{ - Model: c.Model, - Messages: msgs, - // Low temperature for a deterministic-ish explanation; headroom for output. - // Note: some reasoning models only accept the default temperature β€” set - // PGBOT_OPENAI_MODEL accordingly if you pin one. - Temperature: 0.2, - MaxCompletionTokens: 8192, - }) + + buf, err := json.Marshal(reqBody) if err != nil { - return "", err + return nil, err } - req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+"/chat/completions", bytes.NewReader(buf)) + req, err := http.NewRequestWithContext(ctx, http.MethodPost, m.provider.BaseURL+"/chat/completions", bytes.NewReader(buf)) if err != nil { - return "", err + return nil, err } req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+c.APIKey) // header, never a URL param + if m.provider.APIKey != "" { // local servers often want no auth at all + req.Header.Set("Authorization", "Bearer "+m.provider.APIKey) // header, never a query param + } - resp, err := c.HTTP.Do(req) + resp, err := m.provider.HTTP.Do(req) if err != nil { - return "", fmt.Errorf("calling OpenAI: %w", err) + return nil, fmt.Errorf("calling %s: %w", Host(m.provider.BaseURL), err) } defer resp.Body.Close() body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + // Errors first, decoded leniently: OpenAI nests an object under "error", + // xAI returns a bare string there. + var werr wireError + if err := json.Unmarshal(body, &werr); err == nil { + if msg := werr.message(); msg != "" { + return nil, fmt.Errorf("%s error (HTTP %d): %s", Host(m.provider.BaseURL), resp.StatusCode, msg) + } + } + var cr chatResponse if err := json.Unmarshal(body, &cr); err != nil { - return "", fmt.Errorf("openai returned unparseable response (HTTP %d)", resp.StatusCode) - } - if cr.Error != nil { - return "", fmt.Errorf("openai error: %s", cr.Error.Message) + return nil, fmt.Errorf("%s returned unparseable response (HTTP %d)", Host(m.provider.BaseURL), resp.StatusCode) } if resp.StatusCode != http.StatusOK { - return "", fmt.Errorf("openai HTTP %d", resp.StatusCode) + return nil, fmt.Errorf("%s HTTP %d", Host(m.provider.BaseURL), resp.StatusCode) } if len(cr.Choices) == 0 { - return "", fmt.Errorf("openai returned no choices") + return nil, fmt.Errorf("%s returned no choices", Host(m.provider.BaseURL)) } - out := strings.TrimSpace(cr.Choices[0].Message.Content) + choice := cr.Choices[0] + out := strings.TrimSpace(choice.Message.Content) if out == "" { - return "", fmt.Errorf("openai returned an empty explanation (finish: %s)", cr.Choices[0].FinishReason) + // A reasoning model can burn the whole cap on hidden thinking and return + // no visible text at all β€” common on small local models. Say what to do + // about it rather than just reporting emptiness. + if choice.FinishReason == "length" { + return nil, fmt.Errorf("%s hit the output cap before writing an answer β€” the model spent it all on internal reasoning; lower $PGBOT_AI_REASONING_EFFORT, raise the cap, or pick a non-reasoning model", Host(m.provider.BaseURL)) + } + return nil, fmt.Errorf("%s returned an empty explanation (finish: %s)", Host(m.provider.BaseURL), choice.FinishReason) } - return out, nil + return &Response{Text: out, FinishReason: choice.FinishReason}, nil } diff --git a/internal/ai/openai_test.go b/internal/ai/openai_test.go index 9a307bd..c391998 100644 --- a/internal/ai/openai_test.go +++ b/internal/ai/openai_test.go @@ -10,79 +10,217 @@ import ( "testing" ) -func TestOpenAIGenerate_success(t *testing.T) { +func openaiFor(url, key, model string) LanguageModel { + p := &OpenAIProvider{APIKey: key, BaseURL: url, HTTP: http.DefaultClient} + m, _ := p.LanguageModel(context.Background(), model) + return m +} + +func TestOpenAI_success(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if got := r.Header.Get("Authorization"); got != "Bearer sk-test" { - t.Errorf("API key must travel in the Authorization: Bearer header, got %q", got) + if got := r.Header.Get("Authorization"); got != "Bearer k123" { + t.Errorf("key must travel in the Authorization header, got %q", got) } - if !strings.HasSuffix(r.URL.Path, "/chat/completions") { - t.Errorf("unexpected path %s", r.URL.Path) - } - if strings.Contains(r.URL.String(), "sk-test") { + if strings.Contains(r.URL.String(), "k123") { t.Error("API key leaked into the URL") } + if r.URL.Path != "/chat/completions" { + t.Errorf("unexpected path %s", r.URL.Path) + } body, _ := io.ReadAll(r.Body) var req chatRequest if err := json.Unmarshal(body, &req); err != nil { t.Fatalf("bad request body: %v", err) } - if req.Model != "gpt-test" || len(req.Messages) != 2 { - t.Errorf("request shape wrong: %+v", req) + if len(req.Messages) != 2 || req.Messages[0].Role != "system" || req.Messages[1].Role != "user" { + t.Errorf("request shape wrong: %+v", req.Messages) } - if req.Messages[0].Role != "system" || req.Messages[1].Role != "user" { - t.Errorf("expected system then user message, got %+v", req.Messages) + if req.Model != "m1" { + t.Errorf("model not sent: %q", req.Model) } - var resp chatResponse - resp.Choices = append(resp.Choices, struct { - Message chatMessage `json:"message"` - FinishReason string `json:"finish_reason"` - }{Message: chatMessage{Role: "assistant", Content: "Looks healthy."}, FinishReason: "stop"}) - json.NewEncoder(w).Encode(resp) + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{ + "message": map[string]string{"role": "assistant", "content": "Looks healthy."}, + "finish_reason": "stop", + }}, + }) })) defer srv.Close() - c := &OpenAIClient{APIKey: "sk-test", Model: "gpt-test", BaseURL: srv.URL, HTTP: srv.Client()} - out, err := c.Generate(context.Background(), "be terse", "the report") + out, err := openaiFor(srv.URL, "k123", "m1").Generate(context.Background(), Call{System: "be terse", Prompt: "the report"}) if err != nil { t.Fatal(err) } - if out != "Looks healthy." { - t.Errorf("unexpected output %q", out) - } - if c.Vendor() != "OpenAI" || c.ModelName() != "gpt-test" { - t.Errorf("Vendor/ModelName wrong: %q %q", c.Vendor(), c.ModelName()) + if out.Text != "Looks healthy." { + t.Errorf("unexpected output %q", out.Text) } } -func TestOpenAIGenerate_apiError(t *testing.T) { +func TestOpenAI_apiError(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.WriteHeader(http.StatusUnauthorized) - w.Write([]byte(`{"error":{"message":"Incorrect API key provided","type":"invalid_request_error"}}`)) + json.NewEncoder(w).Encode(map[string]any{ + "error": map[string]any{"message": "Incorrect API key provided", "type": "invalid_request_error"}, + }) })) defer srv.Close() - c := &OpenAIClient{APIKey: "bad", Model: "m", BaseURL: srv.URL, HTTP: srv.Client()} - _, err := c.Generate(context.Background(), "", "x") - if err == nil || !strings.Contains(err.Error(), "Incorrect API key") { + + _, err := openaiFor(srv.URL, "bad", "m1").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "Incorrect API key provided") { t.Errorf("expected a surfaced API error, got %v", err) } } -func TestNewOpenAIFromEnv(t *testing.T) { - t.Setenv("OPENAI_API_KEY", "") - if _, err := NewOpenAIFromEnv(); err == nil { - t.Error("missing key should error") +func TestOpenAI_noChoices(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{"choices": []any{}}) // 200, nothing to say + })) + defer srv.Close() + if _, err := openaiFor(srv.URL, "k", "m").Generate(context.Background(), Call{Prompt: "x"}); err == nil { + t.Error("no choices should be an error, not an empty success") + } +} + +// OpenAI's reasoning models reject max_tokens outright and take a reasoning +// effort. Getting this wrong 400s the *default* configuration, so pin the shape. +func TestOpenAI_reasoningRequestShape(t *testing.T) { + var raw map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + if err := json.Unmarshal(b, &raw); err != nil { + t.Fatalf("bad request body: %v", err) + } + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{"message": map[string]string{"content": "ok"}}}, + }) + })) + defer srv.Close() + + p := &OpenAIProvider{APIKey: "k", BaseURL: srv.URL, HTTP: http.DefaultClient} + m, _ := p.LanguageModel(context.Background(), defaultOpenAIModel) + if _, err := m.Generate(context.Background(), Call{ + Prompt: "x", Temperature: f64(0.2), MaxOutputTokens: i64(8192), + }); err != nil { + t.Fatal(err) + } + + if _, ok := raw["max_tokens"]; ok { + t.Error("max_tokens must not be sent to a reasoning model β€” it is rejected") + } + if _, ok := raw["temperature"]; ok { + t.Error("temperature must not be sent to a reasoning model") + } + if raw["reasoning_effort"] != defaultReasoningEffort { + t.Errorf("reasoning_effort = %v, want %q", raw["reasoning_effort"], defaultReasoningEffort) + } + // The cap covers hidden reasoning too, so the 8192 the caller asked for must + // be raised rather than passed through. + if got, ok := raw["max_completion_tokens"].(float64); !ok || int(got) != reasoningTokenFloor { + t.Errorf("max_completion_tokens = %v, want the reasoning floor %d", raw["max_completion_tokens"], reasoningTokenFloor) + } +} + +// A non-reasoning model β€” including everything served by a local runtime β€” must +// keep the plain shape, or Ollama/vLLM stop working. +func TestOpenAI_plainModelKeepsLegacyShape(t *testing.T) { + var raw map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + json.Unmarshal(b, &raw) + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{"message": map[string]string{"content": "ok"}}}, + }) + })) + defer srv.Close() + + _, err := openaiFor(srv.URL, "", "llama3.1").Generate(context.Background(), Call{ + Prompt: "x", Temperature: f64(0.2), MaxOutputTokens: i64(8192), + }) + if err != nil { + t.Fatal(err) + } + if _, ok := raw["max_completion_tokens"]; ok { + t.Error("max_completion_tokens must not be sent to a plain chat model") + } + if _, ok := raw["reasoning_effort"]; ok { + t.Error("reasoning_effort must not be sent to a plain chat model") + } + if got, ok := raw["max_tokens"].(float64); !ok || int(got) != 8192 { + t.Errorf("max_tokens = %v, want 8192", raw["max_tokens"]) + } + if raw["temperature"] == nil { + t.Error("temperature should still be sent to a plain chat model") + } +} + +func TestReasoningModelDetection(t *testing.T) { + reasoning := []string{"gpt-5.6-terra", "gpt-5.6-sol", "GPT-5.6-Luna", "openai/gpt-5.6-terra", "o3-mini", "o4-mini"} + plain := []string{"gpt-4o-mini", "llama3.1", "qwen2.5:7b-instruct", "mistral-large", "deepseek-chat"} + for _, id := range reasoning { + if !reasoningModel(id) { + t.Errorf("%q should be detected as a reasoning model", id) + } } - t.Setenv("OPENAI_API_KEY", "sk-abc") - c, err := NewOpenAIFromEnv() - if err != nil || c.APIKey != "sk-abc" { - t.Fatalf("unexpected: %v %+v", err, c) + for _, id := range plain { + if reasoningModel(id) { + t.Errorf("%q should NOT be detected as a reasoning model", id) + } } - if c.Model != defaultOpenAIModel { - t.Errorf("default model = %q, want %q", c.Model, defaultOpenAIModel) +} + +func TestOpenAI_reasoningEffortOverride(t *testing.T) { + var raw map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + json.Unmarshal(b, &raw) + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{"message": map[string]string{"content": "ok"}}}, + }) + })) + defer srv.Close() + + p := &OpenAIProvider{APIKey: "k", BaseURL: srv.URL, HTTP: http.DefaultClient, ReasoningEffort: "low"} + m, _ := p.LanguageModel(context.Background(), "gpt-5.6-terra") + if _, err := m.Generate(context.Background(), Call{Prompt: "x"}); err != nil { + t.Fatal(err) + } + if raw["reasoning_effort"] != "low" { + t.Errorf("reasoning_effort = %v, want the configured override", raw["reasoning_effort"]) + } +} + +// Small local reasoning models routinely spend the whole cap on hidden thinking +// and return empty content β€” the error has to say what to do about it. +func TestOpenAI_reasoningAteTheBudget(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{ + "message": map[string]string{"content": "", "reasoning": "Thinking Process: …"}, + "finish_reason": "length", + }}, + }) + })) + defer srv.Close() + + _, err := openaiFor(srv.URL, "k", "m").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "output cap") { + t.Errorf("empty content with finish_reason=length should explain the cause, got %v", err) } - t.Setenv("PGBOT_OPENAI_MODEL", "gpt-4.1") - c, _ = NewOpenAIFromEnv() - if c.Model != "gpt-4.1" { - t.Errorf("model override not honored: %q", c.Model) +} + +// A local server (Ollama, LM Studio) usually wants no auth at all β€” sending an +// empty Bearer header makes some of them 401. +func TestOpenAI_noKeySendsNoAuthHeader(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if _, ok := r.Header["Authorization"]; ok { + t.Error("no key configured, but an Authorization header was sent") + } + json.NewEncoder(w).Encode(map[string]any{ + "choices": []map[string]any{{"message": map[string]string{"content": "ok"}}}, + }) + })) + defer srv.Close() + if _, err := openaiFor(srv.URL, "", "llama3.1").Generate(context.Background(), Call{Prompt: "x"}); err != nil { + t.Fatal(err) } } diff --git a/internal/ai/provider.go b/internal/ai/provider.go index 0ef4d70..bbbeac5 100644 --- a/internal/ai/provider.go +++ b/internal/ai/provider.go @@ -1,52 +1,109 @@ +// Package ai is pgbot's OPTIONAL explanation layer. It never produces findings β€” +// those are computed deterministically in package findings. It only takes an +// already-computed, PII-free Context and asks a model to explain and prioritize +// it in plain language. Everything it emits is labeled as model-generated. +// +// The model is yours to choose: Gemini, Anthropic, OpenAI, or any OpenAI-compatible +// endpoint (OpenRouter, Groq, Together, DeepSeek, xAI, Mistral, Ollama, vLLM, +// LM Studio). Each provider is a few hundred lines of net/http so pgbot keeps its +// single-static-binary, minimal-dependency promise β€” no vendor SDKs. package ai import ( "context" - "fmt" - "os" + "encoding/json" + "net/url" "strings" ) -// Provider is an LLM backend for pgbot's OPTIONAL explanation layer. Backends are -// interchangeable and each is called over plain net/http (no SDK), preserving the -// single-static-binary promise. A Provider NEVER produces findings β€” those are -// computed deterministically in package findings; a Provider only explains them. +// Provider is a named source of language models. +// +// This mirrors charmbracelet/fantasy's Provider/LanguageModel pair on purpose, so +// a fantasy-backed implementation could drop in later β€” but we implement it over +// net/http instead of depending on fantasy, which pulls the real vendor SDKs +// (anthropic-sdk-go, openai-go, google.golang.org/genai, aws-sdk-go-v2) and takes +// the binary from 23 MB to ~65 MB for one non-streaming POST. The interface is +// narrowed to the single call pgbot makes: one system turn, one user turn, no +// tools, no streaming. type Provider interface { - // Generate sends one system + user turn and returns the model's text. - Generate(ctx context.Context, system, user string) (string, error) - // ModelName is the concrete model id, for the "model X" line. - ModelName() string - // Vendor is the human destination the data is sent to, for the consent line. - Vendor() string + Name() string + LanguageModel(ctx context.Context, modelID string) (LanguageModel, error) } -// NewFromEnv picks a provider from the environment. Selection order: -// -// 1. PGBOT_AI_PROVIDER=openai|gemini forces one explicitly. -// 2. otherwise auto-detect: OPENAI_API_KEY β†’ OpenAI; GEMINI_API_KEY / -// GOOGLE_API_KEY β†’ Gemini. If both keys are present, OpenAI wins (set -// PGBOT_AI_PROVIDER=gemini to override). -// -// The key is ALWAYS read from the environment, never a flag, so it can't leak -// into shell history or the process list. -func NewFromEnv() (Provider, error) { - switch strings.ToLower(strings.TrimSpace(os.Getenv("PGBOT_AI_PROVIDER"))) { - case "openai": - return NewOpenAIFromEnv() - case "gemini", "google": - return NewGeminiFromEnv() - case "": - // fall through to auto-detect by which key is present - default: - return nil, fmt.Errorf("PGBOT_AI_PROVIDER=%q is not recognized (use \"openai\" or \"gemini\")", - strings.TrimSpace(os.Getenv("PGBOT_AI_PROVIDER"))) +// LanguageModel is one model at one endpoint, ready to answer a single turn. +type LanguageModel interface { + Generate(ctx context.Context, c Call) (*Response, error) + Provider() string // "gemini" | "openai" | "anthropic" + Model() string // resolved model id β€” shown in the AI banner + Endpoint() string // base URL we POST to β€” powers the consent prompt +} + +// Call is one system+user turn. Temperature and MaxOutputTokens are hints: a +// provider omits either when its API rejects it (Anthropic's current models +// return 400 for `temperature`), which is why both are pointers. +type Call struct { + System string + Prompt string + Temperature *float64 + MaxOutputTokens *int64 +} + +// Response is the model's text plus why it stopped. +type Response struct { + Text string + FinishReason string +} + +func f64(v float64) *float64 { return &v } +func i64(v int64) *int64 { return &v } + +// wireError decodes the incompatible error payloads these APIs actually return: +// OpenAI and Anthropic nest an object under "error", xAI puts a bare string +// there. Modelling it as a struct silently fails on xAI and loses the message β€” +// which is the one thing the user needs when a key or model id is wrong. +type wireError struct { + Error json.RawMessage `json:"error"` +} + +func (w wireError) message() string { + if len(w.Error) == 0 || string(w.Error) == "null" { + return "" + } + var s string + if json.Unmarshal(w.Error, &s) == nil { // {"error": "Incorrect API key provided."} + return s + } + var o struct { + Message string `json:"message"` + } + if json.Unmarshal(w.Error, &o) == nil && o.Message != "" { // {"error": {"message": …}} + return o.Message } - if strings.TrimSpace(os.Getenv("OPENAI_API_KEY")) != "" { - return NewOpenAIFromEnv() + return strings.TrimSpace(string(w.Error)) +} + +// Local reports whether a base URL points at this machine. `pgbot explain` is the +// only command that can send data off the box β€” when it can't (Ollama, vLLM, LM +// Studio on localhost), the consent prompt should say so rather than warn about a +// disclosure that isn't happening. +func Local(endpoint string) bool { + u, err := url.Parse(endpoint) + if err != nil { + return false + } + switch strings.ToLower(u.Hostname()) { + case "localhost", "127.0.0.1", "::1", "0.0.0.0": + return true } - if strings.TrimSpace(os.Getenv("GEMINI_API_KEY")) != "" || strings.TrimSpace(os.Getenv("GOOGLE_API_KEY")) != "" { - return NewGeminiFromEnv() + return false +} + +// Host is the display form of an endpoint for the consent prompt β€” host[:port], +// never the full path, and never anything that could carry a key. +func Host(endpoint string) string { + u, err := url.Parse(endpoint) + if err != nil || u.Host == "" { + return endpoint } - return nil, fmt.Errorf("no AI key found β€” set OPENAI_API_KEY (OpenAI) or GEMINI_API_KEY (Google Gemini) " + - "to use `pgbot explain`/`ask` (keys are read from the environment, never a flag)") + return u.Host } diff --git a/internal/ai/provider_test.go b/internal/ai/provider_test.go deleted file mode 100644 index dc8df38..0000000 --- a/internal/ai/provider_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package ai - -import "testing" - -// clearAIEnv blanks every key/override NewFromEnv reads, so each case starts clean. -func clearAIEnv(t *testing.T) { - t.Helper() - for _, k := range []string{"PGBOT_AI_PROVIDER", "OPENAI_API_KEY", "GEMINI_API_KEY", "GOOGLE_API_KEY"} { - t.Setenv(k, "") - } -} - -func TestNewFromEnv_selection(t *testing.T) { - t.Run("no keys errors", func(t *testing.T) { - clearAIEnv(t) - if _, err := NewFromEnv(); err == nil { - t.Error("no key should error") - } - }) - t.Run("gemini key picks Gemini", func(t *testing.T) { - clearAIEnv(t) - t.Setenv("GEMINI_API_KEY", "g") - p, err := NewFromEnv() - if err != nil || p.Vendor() != "Google Gemini" { - t.Fatalf("expected Gemini, got %v / %v", p, err) - } - }) - t.Run("openai key picks OpenAI", func(t *testing.T) { - clearAIEnv(t) - t.Setenv("OPENAI_API_KEY", "sk") - p, err := NewFromEnv() - if err != nil || p.Vendor() != "OpenAI" { - t.Fatalf("expected OpenAI, got %v / %v", p, err) - } - }) - t.Run("both present: OpenAI wins", func(t *testing.T) { - clearAIEnv(t) - t.Setenv("OPENAI_API_KEY", "sk") - t.Setenv("GEMINI_API_KEY", "g") - p, _ := NewFromEnv() - if p.Vendor() != "OpenAI" { - t.Errorf("with both keys OpenAI should win, got %q", p.Vendor()) - } - }) - t.Run("PGBOT_AI_PROVIDER=gemini forces Gemini even with an OpenAI key", func(t *testing.T) { - clearAIEnv(t) - t.Setenv("OPENAI_API_KEY", "sk") - t.Setenv("GEMINI_API_KEY", "g") - t.Setenv("PGBOT_AI_PROVIDER", "gemini") - p, err := NewFromEnv() - if err != nil || p.Vendor() != "Google Gemini" { - t.Fatalf("forced gemini not honored: %v / %v", p, err) - } - }) - t.Run("unknown provider errors", func(t *testing.T) { - clearAIEnv(t) - t.Setenv("PGBOT_AI_PROVIDER", "llama") - if _, err := NewFromEnv(); err == nil { - t.Error("unknown provider should error") - } - }) -} diff --git a/internal/ai/resolve.go b/internal/ai/resolve.go new file mode 100644 index 0000000..d940d51 --- /dev/null +++ b/internal/ai/resolve.go @@ -0,0 +1,183 @@ +package ai + +import ( + "context" + "fmt" + "net/http" + "os" + "strings" + "time" +) + +// Resolve builds the model to use from the environment. Keys come ONLY from the +// environment β€” never a flag β€” so they can't leak into shell history or the +// process list. That invariant is enforced here, once, for every provider. +// +// Precedence: +// +// PGBOT_AI_PROVIDER explicit: gemini | openai | anthropic | xai +// otherwise auto-detected from whichever key is set, +// OpenAI first to preserve existing behavior +// PGBOT_AI_MODEL model id (else the provider's default) +// PGBOT_AI_BASE_URL endpoint (else the provider's default) +// PGBOT_AI_API_KEY key (else the provider's conventional variable) +// PGBOT_AI_REASONING_EFFORT none|low|medium|high|xhigh|max, for reasoning models +// +// Existing PGBOT_OPENAI_* and PGBOT_GEMINI_* settings remain provider-scoped +// aliases for the general model and URL overrides. +func Resolve() (LanguageModel, error) { + name := strings.ToLower(envOr("PGBOT_AI_PROVIDER", "")) + if name == "" { + var err error + if name, err = detectProvider(); err != nil { + return nil, err + } + } + + key := envOr("PGBOT_AI_API_KEY", "") + base := envOr("PGBOT_AI_BASE_URL", "") + model := envOr("PGBOT_AI_MODEL", "") + // Generous because a local model on CPU can take minutes on a full report; + // hosted providers answer in seconds and never come near this. + httpc := &http.Client{Timeout: 3 * time.Minute} + + var p Provider + switch name { + case "gemini", "google": + if key == "" { + key = firstEnv("GEMINI_API_KEY", "GOOGLE_API_KEY") + } + if base == "" { + base = envOr("PGBOT_GEMINI_URL", defaultGeminiURL) + } + if model == "" { + model = envOr("PGBOT_GEMINI_MODEL", "") + } + p = &GeminiProvider{APIKey: key, BaseURL: trimURL(base), HTTP: httpc} + + case "anthropic", "claude": + if key == "" { + key = firstEnv("ANTHROPIC_API_KEY") + } + if base == "" { + base = defaultAnthropicURL + } + p = &AnthropicProvider{APIKey: key, BaseURL: trimURL(base), HTTP: httpc} + + case "xai", "grok", "responses": + // The Responses API, which xAI documents as its primary interface. Same + // endpoint shape at OpenAI, so PGBOT_AI_PROVIDER=responses + an OpenAI key + // and base URL works too. + label := "xai" + if key == "" { + key = firstEnv("XAI_API_KEY", "GROK_API_KEY") + } + if key == "" && firstEnv("OPENAI_API_KEY") != "" { + key, label = firstEnv("OPENAI_API_KEY"), "openai" + if base == "" { + base = envOr("PGBOT_OPENAI_URL", defaultOpenAIURL) + } + if model == "" { + model = envOr("PGBOT_OPENAI_MODEL", defaultOpenAIModel) + } + } + if base == "" { + base = defaultXAIURL + } + p = &ResponsesProvider{ + APIKey: key, + BaseURL: trimURL(base), + HTTP: httpc, + Label: label, + ReasoningEffort: envOr("PGBOT_AI_REASONING_EFFORT", ""), + } + + case "openai", "openrouter", "ollama", "openai-compatible": + // One /chat/completions client serves them all; only the endpoint and the + // conventional key variable differ. + if key == "" { + key = firstEnv("OPENAI_API_KEY", "OPENROUTER_API_KEY") + } + if base == "" { + base = envOr("PGBOT_OPENAI_URL", "") + if base == "" { + base = defaultOpenAIURL + if os.Getenv("OPENAI_API_KEY") == "" && os.Getenv("OPENROUTER_API_KEY") != "" { + base = "https://openrouter.ai/api/v1" + } + } + } + if model == "" { + model = envOr("PGBOT_OPENAI_MODEL", "") + } + p = &OpenAIProvider{ + APIKey: key, + BaseURL: trimURL(base), + HTTP: httpc, + ReasoningEffort: envOr("PGBOT_AI_REASONING_EFFORT", ""), + } + + default: + return nil, fmt.Errorf("unknown PGBOT_AI_PROVIDER %q (want gemini, openai, anthropic, or xai)", name) + } + + // A local endpoint (Ollama, vLLM, LM Studio) usually has no key at all, and + // nothing leaves the machine β€” don't demand one there. + if key == "" && !Local(base) { + return nil, fmt.Errorf("no API key for provider %q β€” set PGBOT_AI_API_KEY (or %s). The key is never read from a flag", name, keyVarsFor(name)) + } + return p.LanguageModel(context.Background(), model) +} + +// detectProvider picks a provider from whichever key is present. OpenAI remains +// first because that was the established precedence before more providers were +// added. PGBOT_AI_PROVIDER removes any ambiguity when several keys are present. +func detectProvider() (string, error) { + switch { + case firstEnv("OPENAI_API_KEY", "OPENROUTER_API_KEY") != "": + return "openai", nil + case firstEnv("GEMINI_API_KEY", "GOOGLE_API_KEY") != "": + return "gemini", nil + case firstEnv("ANTHROPIC_API_KEY") != "": + return "anthropic", nil + case firstEnv("XAI_API_KEY", "GROK_API_KEY") != "": + return "xai", nil + case envOr("PGBOT_AI_API_KEY", "") != "" || Local(envOr("PGBOT_AI_BASE_URL", "")): + // A generic key or a local endpoint with no provider named: /chat/completions + // is the safe assumption β€” it's what every local server speaks. + return "openai", nil + } + return "", fmt.Errorf("no model configured for `pgbot explain` β€” set one of GEMINI_API_KEY, ANTHROPIC_API_KEY, OPENAI_API_KEY, OPENROUTER_API_KEY, or XAI_API_KEY " + + "(or PGBOT_AI_PROVIDER + PGBOT_AI_API_KEY, or PGBOT_AI_BASE_URL for a local model). Keys are never read from a flag") +} + +func keyVarsFor(name string) string { + switch name { + case "gemini", "google": + return "GEMINI_API_KEY / GOOGLE_API_KEY" + case "anthropic", "claude": + return "ANTHROPIC_API_KEY" + case "xai", "grok", "responses": + return "XAI_API_KEY / GROK_API_KEY" + default: + return "OPENAI_API_KEY / OPENROUTER_API_KEY" + } +} + +func envOr(name, def string) string { + if v := strings.TrimSpace(os.Getenv(name)); v != "" { + return v + } + return def +} + +func firstEnv(names ...string) string { + for _, n := range names { + if v := strings.TrimSpace(os.Getenv(n)); v != "" { + return v + } + } + return "" +} + +func trimURL(s string) string { return strings.TrimRight(s, "/") } diff --git a/internal/ai/resolve_test.go b/internal/ai/resolve_test.go new file mode 100644 index 0000000..7fa2389 --- /dev/null +++ b/internal/ai/resolve_test.go @@ -0,0 +1,203 @@ +package ai + +import ( + "strings" + "testing" +) + +// clearEnv unsets every variable Resolve reads, so each case starts from nothing +// and can't be perturbed by the developer's own shell. +func clearEnv(t *testing.T) { + t.Helper() + for _, k := range []string{ + "PGBOT_AI_PROVIDER", "PGBOT_AI_MODEL", "PGBOT_AI_BASE_URL", "PGBOT_AI_API_KEY", + "PGBOT_AI_REASONING_EFFORT", "GEMINI_API_KEY", "GOOGLE_API_KEY", "ANTHROPIC_API_KEY", + "OPENAI_API_KEY", "OPENROUTER_API_KEY", "XAI_API_KEY", "GROK_API_KEY", + "PGBOT_GEMINI_MODEL", "PGBOT_GEMINI_URL", "PGBOT_OPENAI_MODEL", "PGBOT_OPENAI_URL", + } { + t.Setenv(k, "") + } +} + +func TestResolve_noConfiguration(t *testing.T) { + clearEnv(t) + _, err := Resolve() + if err == nil { + t.Fatal("no key anywhere must be an error") + } + // The error has to name every accepted variable, not just Gemini's. + for _, want := range []string{"GEMINI_API_KEY", "ANTHROPIC_API_KEY", "OPENAI_API_KEY", "never read from a flag"} { + if !strings.Contains(err.Error(), want) { + t.Errorf("error should mention %q, got: %v", want, err) + } + } +} + +func TestResolve_autodetect(t *testing.T) { + cases := []struct { + name string + env map[string]string + wantProvider string + wantEndpoint string + }{ + {"gemini", map[string]string{"GEMINI_API_KEY": "k"}, "gemini", defaultGeminiURL}, + {"google fallback", map[string]string{"GOOGLE_API_KEY": "k"}, "gemini", defaultGeminiURL}, + {"anthropic", map[string]string{"ANTHROPIC_API_KEY": "k"}, "anthropic", defaultAnthropicURL}, + {"openai", map[string]string{"OPENAI_API_KEY": "k"}, "openai", defaultOpenAIURL}, + {"openrouter", map[string]string{"OPENROUTER_API_KEY": "k"}, "openai", "https://openrouter.ai/api/v1"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + clearEnv(t) + for k, v := range tc.env { + t.Setenv(k, v) + } + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != tc.wantProvider { + t.Errorf("provider = %q, want %q", m.Provider(), tc.wantProvider) + } + if m.Endpoint() != tc.wantEndpoint { + t.Errorf("endpoint = %q, want %q", m.Endpoint(), tc.wantEndpoint) + } + }) + } +} + +// OpenAI already won auto-detection before Anthropic and xAI were added. +func TestResolve_openAIWinsAutodetect(t *testing.T) { + clearEnv(t) + t.Setenv("GEMINI_API_KEY", "g") + t.Setenv("ANTHROPIC_API_KEY", "a") + t.Setenv("OPENAI_API_KEY", "o") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != "openai" { + t.Errorf("openai must win auto-detection, got %q", m.Provider()) + } +} + +func TestResolve_explicitProviderBeatsAutodetect(t *testing.T) { + clearEnv(t) + t.Setenv("GEMINI_API_KEY", "g") + t.Setenv("ANTHROPIC_API_KEY", "a") + t.Setenv("PGBOT_AI_PROVIDER", "anthropic") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != "anthropic" { + t.Errorf("PGBOT_AI_PROVIDER must win, got %q", m.Provider()) + } + if m.Model() != defaultAnthropicModel { + t.Errorf("model = %q, want the anthropic default", m.Model()) + } +} + +func TestResolve_modelAndURLOverrides(t *testing.T) { + clearEnv(t) + t.Setenv("GEMINI_API_KEY", "k") + t.Setenv("PGBOT_AI_MODEL", "gemini-3-pro") + t.Setenv("PGBOT_AI_BASE_URL", "https://proxy.example/v1/") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Model() != "gemini-3-pro" { + t.Errorf("model = %q", m.Model()) + } + if m.Endpoint() != "https://proxy.example/v1" { + t.Errorf("trailing slash should be trimmed, got %q", m.Endpoint()) + } +} + +// The pre-BYOK variables keep working. +func TestResolve_legacyGeminiVars(t *testing.T) { + clearEnv(t) + t.Setenv("GEMINI_API_KEY", "k") + t.Setenv("PGBOT_GEMINI_MODEL", "gemini-3-pro") + t.Setenv("PGBOT_GEMINI_URL", "https://legacy.example/v1beta") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Model() != "gemini-3-pro" { + t.Errorf("PGBOT_GEMINI_MODEL not honored: %q", m.Model()) + } + if m.Endpoint() != "https://legacy.example/v1beta" { + t.Errorf("PGBOT_GEMINI_URL not honored: %q", m.Endpoint()) + } +} + +func TestResolve_legacyOpenAIVars(t *testing.T) { + clearEnv(t) + t.Setenv("OPENAI_API_KEY", "k") + t.Setenv("PGBOT_OPENAI_MODEL", "gpt-4.1") + t.Setenv("PGBOT_OPENAI_URL", "https://legacy-openai.example/v1/") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Model() != "gpt-4.1" { + t.Errorf("PGBOT_OPENAI_MODEL not honored: %q", m.Model()) + } + if m.Endpoint() != "https://legacy-openai.example/v1" { + t.Errorf("PGBOT_OPENAI_URL not honored: %q", m.Endpoint()) + } +} + +// A local model needs no key and discloses nothing; a remote one still must have +// a key before we'd send findings to it. +func TestResolve_localEndpointNeedsNoKey(t *testing.T) { + clearEnv(t) + t.Setenv("PGBOT_AI_BASE_URL", "http://localhost:11434/v1") + t.Setenv("PGBOT_AI_MODEL", "llama3.1") + m, err := Resolve() + if err != nil { + t.Fatalf("a local endpoint should not require a key: %v", err) + } + if m.Provider() != "openai" { + t.Errorf("a bare local endpoint should default to the OpenAI-compatible client, got %q", m.Provider()) + } + if !Local(m.Endpoint()) { + t.Errorf("%q should be recognized as local", m.Endpoint()) + } +} + +func TestResolve_remoteEndpointRequiresKey(t *testing.T) { + clearEnv(t) + t.Setenv("PGBOT_AI_PROVIDER", "openai") + t.Setenv("PGBOT_AI_BASE_URL", "https://api.example.com/v1") + if _, err := Resolve(); err == nil { + t.Error("a remote endpoint with no key must be an error") + } +} + +func TestResolve_unknownProvider(t *testing.T) { + clearEnv(t) + t.Setenv("PGBOT_AI_PROVIDER", "bedrock") + t.Setenv("PGBOT_AI_API_KEY", "k") + if _, err := Resolve(); err == nil || !strings.Contains(err.Error(), "unknown") { + t.Errorf("unknown provider should be rejected clearly, got %v", err) + } +} + +func TestLocalAndHost(t *testing.T) { + for _, u := range []string{"http://localhost:11434/v1", "http://127.0.0.1:8000/v1", "http://[::1]:1234/v1"} { + if !Local(u) { + t.Errorf("%q should be local", u) + } + } + for _, u := range []string{defaultGeminiURL, defaultAnthropicURL, "https://openrouter.ai/api/v1"} { + if Local(u) { + t.Errorf("%q should not be local", u) + } + } + if got := Host("https://api.anthropic.com/v1/messages"); got != "api.anthropic.com" { + t.Errorf("Host should be host[:port] only, got %q", got) + } +} diff --git a/internal/ai/responses.go b/internal/ai/responses.go new file mode 100644 index 0000000..b18ce37 --- /dev/null +++ b/internal/ai/responses.go @@ -0,0 +1,170 @@ +package ai + +import ( + "bytes" + "context" + "encoding/json" + "fmt" + "io" + "net/http" + "strings" +) + +const ( + defaultXAIModel = "grok-4.6" + defaultXAIURL = "https://api.x.ai/v1" +) + +// ResponsesProvider speaks the Responses API (POST /responses) β€” the newer +// surface both xAI and OpenAI prefer over /chat/completions. pgbot uses it for +// xAI, where it is the documented primary interface. +// +// It is deliberately NOT the default for the OpenAI-compatible world: only +// OpenAI and xAI implement /responses, while Ollama, vLLM, LM Studio, Groq, +// Together, DeepSeek and Mistral implement only /chat/completions. This provider +// is additive β€” OpenAIProvider stays the compatibility path. +type ResponsesProvider struct { + APIKey string + BaseURL string + HTTP *http.Client + + // Label is the provider name shown in the consent prompt and AI banner + // ("xai", "openai") β€” the endpoint is shared, the vendor is not. + Label string + + // ReasoningEffort is sent as reasoning.effort when set. Left empty by default + // because the accepted vocabulary differs per vendor; the model's own default + // is the safe choice. + ReasoningEffort string +} + +func (p *ResponsesProvider) Name() string { + if p.Label != "" { + return p.Label + } + return "responses" +} + +func (p *ResponsesProvider) LanguageModel(_ context.Context, modelID string) (LanguageModel, error) { + if modelID == "" { + modelID = defaultXAIModel + } + return &responsesModel{provider: p, model: modelID}, nil +} + +type responsesModel struct { + provider *ResponsesProvider + model string +} + +func (m *responsesModel) Provider() string { return m.provider.Name() } +func (m *responsesModel) Model() string { return m.model } +func (m *responsesModel) Endpoint() string { return m.provider.BaseURL } + +// ---- wire types (only the fields we use) ---- + +type responsesRequest struct { + Model string `json:"model"` + Instructions string `json:"instructions,omitempty"` // the system turn + Input string `json:"input"` // the user turn + // Store is explicitly false. The Responses API defaults it to TRUE, which + // retains the request server-side for later retrieval β€” a quiet downgrade of + // the disclosure `pgbot explain` asks the user to consent to. We send the + // findings once and keep nothing on the vendor's side. + Store bool `json:"store"` + MaxOutputTokens *int64 `json:"max_output_tokens,omitempty"` + Temperature *float64 `json:"temperature,omitempty"` + Reasoning *reasoningCfg `json:"reasoning,omitempty"` +} + +type reasoningCfg struct { + Effort string `json:"effort,omitempty"` +} + +type responsesResponse struct { + Status string `json:"status"` // "completed" | "incomplete" | … + Output []struct { + Type string `json:"type"` // "reasoning" | "message" | … + Role string `json:"role"` + Content []struct { + Type string `json:"type"` // "output_text" | … + Text string `json:"text"` + } `json:"content"` + } `json:"output"` + IncompleteDetails *struct { + Reason string `json:"reason"` + } `json:"incomplete_details"` +} + +// Generate sends one instructions + input turn and returns the model's text. No +// retries β€” a failed explanation must not hang the CLI. +func (m *responsesModel) Generate(ctx context.Context, c Call) (*Response, error) { + reqBody := responsesRequest{ + Model: m.model, + Instructions: c.System, + Input: c.Prompt, + Store: false, + MaxOutputTokens: c.MaxOutputTokens, + Temperature: c.Temperature, + } + if e := m.provider.ReasoningEffort; e != "" { + reqBody.Reasoning = &reasoningCfg{Effort: e} + } + buf, err := json.Marshal(reqBody) + if err != nil { + return nil, err + } + + req, err := http.NewRequestWithContext(ctx, http.MethodPost, m.provider.BaseURL+"/responses", bytes.NewReader(buf)) + if err != nil { + return nil, err + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Authorization", "Bearer "+m.provider.APIKey) // header, never a query param + + resp, err := m.provider.HTTP.Do(req) + if err != nil { + return nil, fmt.Errorf("calling %s: %w", m.provider.Name(), err) + } + defer resp.Body.Close() + body, _ := io.ReadAll(io.LimitReader(resp.Body, 1<<20)) + + // Errors first: this API returns them as a top-level field, and xAI's are a + // bare string where OpenAI's are an object. + var werr wireError + if err := json.Unmarshal(body, &werr); err == nil { + if msg := werr.message(); msg != "" { + return nil, fmt.Errorf("%s error (HTTP %d): %s", m.provider.Name(), resp.StatusCode, msg) + } + } + if resp.StatusCode != http.StatusOK { + return nil, fmt.Errorf("%s HTTP %d", m.provider.Name(), resp.StatusCode) + } + + var rr responsesResponse + if err := json.Unmarshal(body, &rr); err != nil { + return nil, fmt.Errorf("%s returned unparseable response (HTTP %d)", m.provider.Name(), resp.StatusCode) + } + + // The answer is one entry in `output` β€” the rest is reasoning, which we never + // surface. There is no `output_text` on the wire; that is an SDK convenience. + var sb strings.Builder + for _, o := range rr.Output { + if o.Type != "message" { + continue + } + for _, ct := range o.Content { + if ct.Type == "output_text" { + sb.WriteString(ct.Text) + } + } + } + out := strings.TrimSpace(sb.String()) + if out == "" { + if rr.IncompleteDetails != nil && rr.IncompleteDetails.Reason == "max_output_tokens" { + return nil, fmt.Errorf("%s hit the output cap before writing an answer β€” the model spent it all on internal reasoning; lower $PGBOT_AI_REASONING_EFFORT or raise the cap", m.provider.Name()) + } + return nil, fmt.Errorf("%s returned an empty explanation (status: %s)", m.provider.Name(), rr.Status) + } + return &Response{Text: out, FinishReason: rr.Status}, nil +} diff --git a/internal/ai/responses_test.go b/internal/ai/responses_test.go new file mode 100644 index 0000000..472d448 --- /dev/null +++ b/internal/ai/responses_test.go @@ -0,0 +1,192 @@ +package ai + +import ( + "context" + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func responsesFor(url, key, model string) LanguageModel { + p := &ResponsesProvider{APIKey: key, BaseURL: url, HTTP: http.DefaultClient, Label: "xai"} + m, _ := p.LanguageModel(context.Background(), model) + return m +} + +// The shape below mirrors a real grok-4.6 response: a reasoning entry the user +// must never see, followed by the message. +func TestResponses_success(t *testing.T) { + var raw map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if got := r.Header.Get("Authorization"); got != "Bearer k123" { + t.Errorf("key must travel in the Authorization header, got %q", got) + } + if strings.Contains(r.URL.String(), "k123") { + t.Error("API key leaked into the URL") + } + if r.URL.Path != "/responses" { + t.Errorf("unexpected path %s", r.URL.Path) + } + b, _ := io.ReadAll(r.Body) + if err := json.Unmarshal(b, &raw); err != nil { + t.Fatalf("bad request body: %v", err) + } + json.NewEncoder(w).Encode(map[string]any{ + "status": "completed", + "output": []map[string]any{ + {"type": "reasoning", "summary": []map[string]string{{"type": "summary_text", "text": "internal"}}}, + {"type": "message", "role": "assistant", "content": []map[string]string{ + {"type": "output_text", "text": "Looks healthy."}, + }}, + }, + }) + })) + defer srv.Close() + + out, err := responsesFor(srv.URL, "k123", "grok-4.6").Generate(context.Background(), Call{ + System: "be terse", Prompt: "the report", MaxOutputTokens: i64(8192), + }) + if err != nil { + t.Fatal(err) + } + if out.Text != "Looks healthy." { + t.Errorf("unexpected output %q β€” reasoning entries must not leak in", out.Text) + } + if raw["instructions"] != "be terse" || raw["input"] != "the report" { + t.Errorf("system/user must map to instructions/input, got %v / %v", raw["instructions"], raw["input"]) + } + // The API defaults store to true, which would retain the findings server-side. + if raw["store"] != false { + t.Errorf("store must be explicitly false, got %v", raw["store"]) + } +} + +// xAI returns a bare string under "error" where OpenAI returns an object. +// Decoding it as a struct loses the message the user needs. +func TestResponses_xaiStringError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusBadRequest) + w.Write([]byte(`{"code":"invalid-argument","error":"Incorrect API key provided. You can obtain an API key from https://console.x.ai."}`)) + })) + defer srv.Close() + + _, err := responsesFor(srv.URL, "bad", "grok-4.6").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "Incorrect API key provided") { + t.Errorf("xAI's string error must be surfaced verbatim, got %v", err) + } +} + +func TestResponses_openAIObjectError(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.WriteHeader(http.StatusUnauthorized) + w.Write([]byte(`{"error":{"message":"Invalid Authentication","type":"invalid_request_error"}}`)) + })) + defer srv.Close() + + _, err := responsesFor(srv.URL, "bad", "gpt-5.6-terra").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "Invalid Authentication") { + t.Errorf("OpenAI's object error must be surfaced too, got %v", err) + } +} + +// Reasoning can consume the whole budget, leaving a message entry with no text. +func TestResponses_truncated(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + json.NewEncoder(w).Encode(map[string]any{ + "status": "incomplete", + "incomplete_details": map[string]string{"reason": "max_output_tokens"}, + "output": []map[string]any{{"type": "reasoning"}}, + }) + })) + defer srv.Close() + + _, err := responsesFor(srv.URL, "k", "grok-4.6").Generate(context.Background(), Call{Prompt: "x"}) + if err == nil || !strings.Contains(err.Error(), "output cap") { + t.Errorf("a truncated response should explain the cause, got %v", err) + } +} + +func TestResponses_reasoningEffortOmittedUnlessSet(t *testing.T) { + var raw map[string]any + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + b, _ := io.ReadAll(r.Body) + json.Unmarshal(b, &raw) + json.NewEncoder(w).Encode(map[string]any{ + "status": "completed", + "output": []map[string]any{{"type": "message", "content": []map[string]string{{"type": "output_text", "text": "ok"}}}}, + }) + })) + defer srv.Close() + + if _, err := responsesFor(srv.URL, "k", "grok-4.6").Generate(context.Background(), Call{Prompt: "x"}); err != nil { + t.Fatal(err) + } + if _, ok := raw["reasoning"]; ok { + t.Error("reasoning must be omitted unless an effort is configured β€” vocabularies differ per vendor") + } + + p := &ResponsesProvider{APIKey: "k", BaseURL: srv.URL, HTTP: http.DefaultClient, Label: "xai", ReasoningEffort: "low"} + m, _ := p.LanguageModel(context.Background(), "grok-4.6") + if _, err := m.Generate(context.Background(), Call{Prompt: "x"}); err != nil { + t.Fatal(err) + } + got, _ := raw["reasoning"].(map[string]any) + if got == nil || got["effort"] != "low" { + t.Errorf("configured effort should be sent, got %v", raw["reasoning"]) + } +} + +func TestResolve_xai(t *testing.T) { + clearEnv(t) + t.Setenv("XAI_API_KEY", "k") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != "xai" { + t.Errorf("provider = %q, want xai", m.Provider()) + } + if m.Model() != defaultXAIModel { + t.Errorf("model = %q, want %q", m.Model(), defaultXAIModel) + } + if m.Endpoint() != defaultXAIURL { + t.Errorf("endpoint = %q, want %q", m.Endpoint(), defaultXAIURL) + } +} + +// PGBOT_AI_PROVIDER=responses with an OpenAI key must not inherit Grok's model. +func TestResolve_responsesWithOpenAIKey(t *testing.T) { + clearEnv(t) + t.Setenv("PGBOT_AI_PROVIDER", "responses") + t.Setenv("OPENAI_API_KEY", "k") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if m.Provider() != "openai" { + t.Errorf("provider = %q, want openai", m.Provider()) + } + if m.Model() != defaultOpenAIModel { + t.Errorf("model = %q, want %q", m.Model(), defaultOpenAIModel) + } + if m.Endpoint() != defaultOpenAIURL { + t.Errorf("endpoint = %q, want %q", m.Endpoint(), defaultOpenAIURL) + } +} + +// An existing OpenAI/Ollama setup must keep the chat/completions provider β€” +// /responses is not implemented by the local runtimes. +func TestResolve_openAIKeyStillUsesChatCompletions(t *testing.T) { + clearEnv(t) + t.Setenv("OPENAI_API_KEY", "k") + m, err := Resolve() + if err != nil { + t.Fatal(err) + } + if _, ok := m.(*openaiModel); !ok { + t.Errorf("auto-detected OpenAI must use the chat/completions provider, got %T", m) + } +}