first commit
This commit is contained in:
@@ -0,0 +1,334 @@
|
||||
// This file implements the Anthropic Messages API streaming decoder (US-008).
|
||||
// It is a stateful Decoder (see transport.go) that translates Anthropic SSE
|
||||
// event payloads into the provider-agnostic AssistantMessageEvent set,
|
||||
// accumulating a partial AssistantMessage as deltas arrive.
|
||||
//
|
||||
// Anthropic streams a fixed event sequence:
|
||||
//
|
||||
// message_start → seeds id/model and initial usage (input tokens)
|
||||
// content_block_start → opens a text / thinking / tool_use block at an index
|
||||
// content_block_delta → text_delta / thinking_delta / signature_delta /
|
||||
// input_json_delta append to the open block
|
||||
// content_block_stop → closes the block (tool_use JSON is parsed here)
|
||||
// message_delta → carries the final stop_reason and output-token usage
|
||||
// message_stop → terminal; the accumulated message is emitted as done
|
||||
// error → a runtime error payload → terminal error event
|
||||
//
|
||||
// Per the dual failure model (FR-13) the decoder never panics: malformed
|
||||
// payloads and Anthropic `error` events are surfaced as a returned error (which
|
||||
// the transport turns into a terminal StreamErrorEvent) rather than crashing.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// anthropicBlock accumulates one content block's streaming state, keyed by its
|
||||
// Anthropic content-block index. text/thinking append their deltas; tool_use
|
||||
// accumulates a partial JSON string parsed lazily when the block is realized.
|
||||
type anthropicBlock struct {
|
||||
kind string // "text" | "thinking" | "tool_use" | "redacted_thinking"
|
||||
text strings.Builder
|
||||
thinking strings.Builder
|
||||
thinkingSig string
|
||||
textSig string
|
||||
toolID string
|
||||
toolName string
|
||||
toolJSON strings.Builder
|
||||
redacted bool
|
||||
}
|
||||
|
||||
// AnthropicDecoder is the stateful SSE decoder for the Anthropic Messages API.
|
||||
// It implements the transport Decoder interface. It is not safe for concurrent
|
||||
// use — the transport drives it from a single goroutine.
|
||||
type AnthropicDecoder struct {
|
||||
blocks map[int]*anthropicBlock
|
||||
order []int // content-block indices in first-seen order
|
||||
|
||||
responseID string
|
||||
responseModel string
|
||||
inputTokens int
|
||||
outputTokens int
|
||||
stopReason string // mapped pigo stop reason (empty until message_delta)
|
||||
done bool // message_stop / done already emitted
|
||||
}
|
||||
|
||||
// NewAnthropicDecoder builds a fresh decoder for one streamed response.
|
||||
func NewAnthropicDecoder() *AnthropicDecoder {
|
||||
return &AnthropicDecoder{blocks: make(map[int]*anthropicBlock)}
|
||||
}
|
||||
|
||||
// anthropicEvent is the discriminated envelope shared by every Anthropic SSE
|
||||
// data payload; fields are populated selectively by event type.
|
||||
type anthropicEvent struct {
|
||||
Type string `json:"type"`
|
||||
Index int `json:"index"`
|
||||
|
||||
Message *struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Usage *anthropicUsage `json:"usage"`
|
||||
} `json:"message"`
|
||||
|
||||
ContentBlock *struct {
|
||||
Type string `json:"type"`
|
||||
// text
|
||||
Text string `json:"text"`
|
||||
// thinking
|
||||
Thinking string `json:"thinking"`
|
||||
// tool_use
|
||||
ID string `json:"id"`
|
||||
Name string `json:"name"`
|
||||
} `json:"content_block"`
|
||||
|
||||
Delta *struct {
|
||||
Type string `json:"type"`
|
||||
// text_delta
|
||||
Text string `json:"text"`
|
||||
// thinking_delta
|
||||
Thinking string `json:"thinking"`
|
||||
// signature_delta
|
||||
Signature string `json:"signature"`
|
||||
// input_json_delta
|
||||
PartialJSON string `json:"partial_json"`
|
||||
// message_delta
|
||||
StopReason string `json:"stop_reason"`
|
||||
} `json:"delta"`
|
||||
|
||||
Usage *anthropicUsage `json:"usage"`
|
||||
|
||||
Error *struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
type anthropicUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
|
||||
// Decode turns one Anthropic SSE data payload into zero or more StreamEvents.
|
||||
func (d *AnthropicDecoder) Decode(payload []byte) ([]StreamEvent, error) {
|
||||
var ev anthropicEvent
|
||||
if err := json.Unmarshal(payload, &ev); err != nil {
|
||||
return nil, fmt.Errorf("anthropic: parse event: %w", err)
|
||||
}
|
||||
|
||||
switch ev.Type {
|
||||
case "message_start":
|
||||
return d.onMessageStart(ev), nil
|
||||
case "content_block_start":
|
||||
return d.onBlockStart(ev), nil
|
||||
case "content_block_delta":
|
||||
return d.onBlockDelta(ev), nil
|
||||
case "content_block_stop":
|
||||
// Nothing to emit on stop; the block is already reflected in the partial.
|
||||
return nil, nil
|
||||
case "message_delta":
|
||||
return d.onMessageDelta(ev), nil
|
||||
case "message_stop":
|
||||
return d.finishDone(), nil
|
||||
case "ping":
|
||||
return nil, nil
|
||||
case "error":
|
||||
msg := "anthropic stream error"
|
||||
if ev.Error != nil {
|
||||
if ev.Error.Type != "" {
|
||||
msg = "anthropic " + ev.Error.Type
|
||||
}
|
||||
if ev.Error.Message != "" {
|
||||
msg += ": " + ev.Error.Message
|
||||
}
|
||||
}
|
||||
return nil, fmt.Errorf("%s", msg)
|
||||
default:
|
||||
// Unknown event types are ignored (forward-compatible).
|
||||
return nil, nil
|
||||
}
|
||||
}
|
||||
|
||||
// Finish flushes a terminal done event if the stream ended without an explicit
|
||||
// message_stop (e.g. a clean EOF mid-stream), so a partial response is still
|
||||
// delivered rather than lost.
|
||||
func (d *AnthropicDecoder) Finish() ([]StreamEvent, error) {
|
||||
if d.done {
|
||||
return nil, nil
|
||||
}
|
||||
return d.finishDone(), nil
|
||||
}
|
||||
|
||||
func (d *AnthropicDecoder) onMessageStart(ev anthropicEvent) []StreamEvent {
|
||||
if ev.Message != nil {
|
||||
d.responseID = ev.Message.ID
|
||||
d.responseModel = ev.Message.Model
|
||||
if ev.Message.Usage != nil {
|
||||
d.inputTokens = ev.Message.Usage.InputTokens
|
||||
d.outputTokens = ev.Message.Usage.OutputTokens
|
||||
}
|
||||
}
|
||||
return []StreamEvent{StreamStartEvent{Partial: d.partial()}}
|
||||
}
|
||||
|
||||
func (d *AnthropicDecoder) onBlockStart(ev anthropicEvent) []StreamEvent {
|
||||
if ev.ContentBlock == nil {
|
||||
return nil
|
||||
}
|
||||
b := &anthropicBlock{kind: ev.ContentBlock.Type}
|
||||
switch ev.ContentBlock.Type {
|
||||
case "text":
|
||||
b.text.WriteString(ev.ContentBlock.Text)
|
||||
case "thinking":
|
||||
b.thinking.WriteString(ev.ContentBlock.Thinking)
|
||||
case "redacted_thinking":
|
||||
b.redacted = true
|
||||
case "tool_use":
|
||||
b.toolID = ev.ContentBlock.ID
|
||||
b.toolName = ev.ContentBlock.Name
|
||||
}
|
||||
d.putBlock(ev.Index, b)
|
||||
|
||||
switch ev.ContentBlock.Type {
|
||||
case "thinking", "redacted_thinking":
|
||||
return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}}
|
||||
case "tool_use":
|
||||
return []StreamEvent{StreamToolCallEvent{Partial: d.partial()}}
|
||||
default:
|
||||
return []StreamEvent{StreamTextEvent{Partial: d.partial()}}
|
||||
}
|
||||
}
|
||||
|
||||
func (d *AnthropicDecoder) onBlockDelta(ev anthropicEvent) []StreamEvent {
|
||||
if ev.Delta == nil {
|
||||
return nil
|
||||
}
|
||||
b := d.blocks[ev.Index]
|
||||
if b == nil {
|
||||
// A delta for an unseen index: open a bare block so we don't drop data.
|
||||
b = &anthropicBlock{}
|
||||
d.putBlock(ev.Index, b)
|
||||
}
|
||||
switch ev.Delta.Type {
|
||||
case "text_delta":
|
||||
b.text.WriteString(ev.Delta.Text)
|
||||
return []StreamEvent{StreamTextEvent{Partial: d.partial()}}
|
||||
case "thinking_delta":
|
||||
b.thinking.WriteString(ev.Delta.Thinking)
|
||||
return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}}
|
||||
case "signature_delta":
|
||||
// Signature belongs to the thinking block it rides on.
|
||||
b.thinkingSig += ev.Delta.Signature
|
||||
return []StreamEvent{StreamThinkingEvent{Partial: d.partial()}}
|
||||
case "input_json_delta":
|
||||
b.toolJSON.WriteString(ev.Delta.PartialJSON)
|
||||
return []StreamEvent{StreamToolCallEvent{Partial: d.partial()}}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (d *AnthropicDecoder) onMessageDelta(ev anthropicEvent) []StreamEvent {
|
||||
if ev.Delta != nil && ev.Delta.StopReason != "" {
|
||||
d.stopReason = mapAnthropicStopReason(ev.Delta.StopReason)
|
||||
}
|
||||
if ev.Usage != nil {
|
||||
// message_delta reports cumulative output tokens (and sometimes input).
|
||||
if ev.Usage.OutputTokens != 0 {
|
||||
d.outputTokens = ev.Usage.OutputTokens
|
||||
}
|
||||
if ev.Usage.InputTokens != 0 {
|
||||
d.inputTokens = ev.Usage.InputTokens
|
||||
}
|
||||
}
|
||||
// No standalone event kind for usage/stop-reason accumulation; the values
|
||||
// surface in the terminal done message.
|
||||
return nil
|
||||
}
|
||||
|
||||
// finishDone builds the terminal assistant message and marks the decoder done.
|
||||
func (d *AnthropicDecoder) finishDone() []StreamEvent {
|
||||
if d.done {
|
||||
return nil
|
||||
}
|
||||
d.done = true
|
||||
msg := d.partial()
|
||||
if msg.StopReason == "" {
|
||||
msg.StopReason = agentcore.StopReasonEndTurn
|
||||
}
|
||||
return []StreamEvent{StreamDoneEvent{Message: msg}}
|
||||
}
|
||||
|
||||
// putBlock records a block at index, tracking first-seen order.
|
||||
func (d *AnthropicDecoder) putBlock(index int, b *anthropicBlock) {
|
||||
if _, seen := d.blocks[index]; !seen {
|
||||
d.order = append(d.order, index)
|
||||
}
|
||||
d.blocks[index] = b
|
||||
}
|
||||
|
||||
// partial materializes the accumulated state into an AssistantMessage. Content
|
||||
// blocks are emitted in content-block index order. Tool-use JSON that has not
|
||||
// yet parsed cleanly is passed through as-is (raw partial), which is valid for
|
||||
// a still-streaming partial and finalized once the block completes.
|
||||
func (d *AnthropicDecoder) partial() agentcore.AssistantMessage {
|
||||
msg := agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
API: "anthropic",
|
||||
Provider: "anthropic",
|
||||
StopReason: d.stopReason,
|
||||
ResponseID: d.responseID,
|
||||
ResponseModel: d.responseModel,
|
||||
}
|
||||
if d.inputTokens != 0 || d.outputTokens != 0 {
|
||||
msg.Usage = &agentcore.Usage{InputTokens: d.inputTokens, OutputTokens: d.outputTokens}
|
||||
}
|
||||
|
||||
idx := make([]int, len(d.order))
|
||||
copy(idx, d.order)
|
||||
sort.Ints(idx)
|
||||
|
||||
for _, i := range idx {
|
||||
b := d.blocks[i]
|
||||
if b == nil {
|
||||
continue
|
||||
}
|
||||
switch b.kind {
|
||||
case "thinking", "redacted_thinking":
|
||||
tc := agentcore.NewThinkingContent(b.thinking.String())
|
||||
tc.ThinkingSignature = b.thinkingSig
|
||||
tc.Redacted = b.redacted
|
||||
msg.Content = append(msg.Content, tc)
|
||||
case "tool_use":
|
||||
args := json.RawMessage(strings.TrimSpace(b.toolJSON.String()))
|
||||
if len(args) == 0 {
|
||||
args = json.RawMessage("{}")
|
||||
}
|
||||
msg.Content = append(msg.Content, agentcore.NewToolCallContent(b.toolID, b.toolName, args))
|
||||
default: // text
|
||||
tc := agentcore.NewTextContent(b.text.String())
|
||||
tc.TextSignature = b.textSig
|
||||
msg.Content = append(msg.Content, tc)
|
||||
}
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
// mapAnthropicStopReason maps an Anthropic stop_reason to the pigo StopReason
|
||||
// set. Unknown reasons default to end_turn (a natural, non-error stop).
|
||||
func mapAnthropicStopReason(reason string) string {
|
||||
switch reason {
|
||||
case "max_tokens":
|
||||
return agentcore.StopReasonLength
|
||||
case "tool_use":
|
||||
return agentcore.StopReasonToolUse
|
||||
case "end_turn", "stop_sequence":
|
||||
return agentcore.StopReasonEndTurn
|
||||
default:
|
||||
return agentcore.StopReasonEndTurn
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,299 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// A recorded Anthropic Messages API SSE stream covering a text block, a
|
||||
// thinking block (with signature), and a tool_use block, ending with a
|
||||
// tool_use stop reason and output-token usage. Trimmed but structurally
|
||||
// faithful to the real wire format.
|
||||
const anthropicToolUseSSE = `event: message_start
|
||||
data: {"type":"message_start","message":{"id":"msg_01ABC","model":"claude-opus-4-8","usage":{"input_tokens":42,"output_tokens":1}}}
|
||||
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"thinking","thinking":""}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Let me think. "}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"thinking_delta","thinking":"Use the tool."}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"signature_delta","signature":"sigABC"}}
|
||||
|
||||
event: content_block_stop
|
||||
data: {"type":"content_block_stop","index":0}
|
||||
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":1,"content_block":{"type":"text","text":""}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"I'll check "}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":1,"delta":{"type":"text_delta","text":"the weather."}}
|
||||
|
||||
event: content_block_stop
|
||||
data: {"type":"content_block_stop","index":1}
|
||||
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":2,"content_block":{"type":"tool_use","id":"toolu_01","name":"get_weather"}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"city\":"}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":" \"SF\"}"}}
|
||||
|
||||
event: content_block_stop
|
||||
data: {"type":"content_block_stop","index":2}
|
||||
|
||||
event: message_delta
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":57}}
|
||||
|
||||
event: message_stop
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
`
|
||||
|
||||
// feedSSE splits a recorded SSE body into events (blank-line separated),
|
||||
// extracts each `data:` payload, and drives the decoder exactly as the
|
||||
// transport pump would, returning all emitted events plus the final message.
|
||||
func feedSSE(t *testing.T, dec Decoder, body string) ([]StreamEvent, agentcore.AssistantMessage) {
|
||||
t.Helper()
|
||||
var events []StreamEvent
|
||||
for _, block := range strings.Split(body, "\n\n") {
|
||||
var payload strings.Builder
|
||||
for _, line := range strings.Split(block, "\n") {
|
||||
line = strings.TrimRight(line, "\r")
|
||||
if strings.HasPrefix(line, "data:") {
|
||||
payload.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
}
|
||||
}
|
||||
if payload.Len() == 0 {
|
||||
continue
|
||||
}
|
||||
// [DONE] is a transport-level terminator, not a decoder payload.
|
||||
if payload.String() == "[DONE]" {
|
||||
continue
|
||||
}
|
||||
evs, err := dec.Decode([]byte(payload.String()))
|
||||
if err != nil {
|
||||
t.Fatalf("decode %q: %v", payload.String(), err)
|
||||
}
|
||||
events = append(events, evs...)
|
||||
}
|
||||
finalEvents, err := dec.Finish()
|
||||
if err != nil {
|
||||
t.Fatalf("finish: %v", err)
|
||||
}
|
||||
events = append(events, finalEvents...)
|
||||
|
||||
var final agentcore.AssistantMessage
|
||||
for _, ev := range events {
|
||||
if d, ok := ev.(StreamDoneEvent); ok {
|
||||
final = d.Message
|
||||
}
|
||||
}
|
||||
return events, final
|
||||
}
|
||||
|
||||
func TestAnthropicDecoderToolUseStream(t *testing.T) {
|
||||
dec := NewAnthropicDecoder()
|
||||
events, final := feedSSE(t, dec, anthropicToolUseSSE)
|
||||
|
||||
// The first emitted event must be a start event.
|
||||
if len(events) == 0 || events[0].EventKind() != StreamEventStart {
|
||||
t.Fatalf("expected a start event first, got %v", eventKinds(events))
|
||||
}
|
||||
// The last emitted event must be the terminal done event.
|
||||
if events[len(events)-1].EventKind() != StreamEventDone {
|
||||
t.Fatalf("expected a done event last, got %v", eventKinds(events))
|
||||
}
|
||||
|
||||
// Stop reason: tool_use.
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
// Usage: input from message_start, output from message_delta.
|
||||
if final.Usage == nil || final.Usage.InputTokens != 42 || final.Usage.OutputTokens != 57 {
|
||||
t.Errorf("usage = %+v, want input=42 output=57", final.Usage)
|
||||
}
|
||||
// Response identity from message_start.
|
||||
if final.ResponseID != "msg_01ABC" || final.ResponseModel != "claude-opus-4-8" {
|
||||
t.Errorf("response id/model = %q/%q", final.ResponseID, final.ResponseModel)
|
||||
}
|
||||
|
||||
// Content blocks in index order: thinking, text, tool_use.
|
||||
if len(final.Content) != 3 {
|
||||
t.Fatalf("expected 3 content blocks, got %d: %+v", len(final.Content), final.Content)
|
||||
}
|
||||
th, ok := final.Content[0].(agentcore.ThinkingContent)
|
||||
if !ok || th.Thinking != "Let me think. Use the tool." {
|
||||
t.Errorf("thinking block = %+v", final.Content[0])
|
||||
}
|
||||
if th.ThinkingSignature != "sigABC" {
|
||||
t.Errorf("thinking signature = %q, want sigABC", th.ThinkingSignature)
|
||||
}
|
||||
txt, ok := final.Content[1].(agentcore.TextContent)
|
||||
if !ok || txt.Text != "I'll check the weather." {
|
||||
t.Errorf("text block = %+v", final.Content[1])
|
||||
}
|
||||
tool, ok := final.Content[2].(agentcore.ToolCallContent)
|
||||
if !ok || tool.Name != "get_weather" || tool.ID != "toolu_01" {
|
||||
t.Fatalf("tool block = %+v", final.Content[2])
|
||||
}
|
||||
// tool_use JSON must have accumulated into valid arguments.
|
||||
var args map[string]string
|
||||
if err := json.Unmarshal(tool.Arguments, &args); err != nil {
|
||||
t.Fatalf("tool arguments not valid JSON %q: %v", tool.Arguments, err)
|
||||
}
|
||||
if args["city"] != "SF" {
|
||||
t.Errorf("tool arguments = %v, want city=SF", args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicDecoderTextOnlyEndTurn(t *testing.T) {
|
||||
body := `data: {"type":"message_start","message":{"id":"msg_1","model":"claude-x","usage":{"input_tokens":10,"output_tokens":0}}}
|
||||
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
||||
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
|
||||
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":" world"}}
|
||||
|
||||
data: {"type":"content_block_stop","index":0}
|
||||
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":3}}
|
||||
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
`
|
||||
dec := NewAnthropicDecoder()
|
||||
_, final := feedSSE(t, dec, body)
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("stop reason = %q, want end_turn", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 1 {
|
||||
t.Fatalf("expected 1 content block, got %d", len(final.Content))
|
||||
}
|
||||
txt, ok := final.Content[0].(agentcore.TextContent)
|
||||
if !ok || txt.Text != "Hello world" {
|
||||
t.Errorf("text = %+v", final.Content[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicDecoderMaxTokensMapsToLength(t *testing.T) {
|
||||
body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}}
|
||||
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
||||
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"truncated"}}
|
||||
|
||||
data: {"type":"content_block_stop","index":0}
|
||||
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"max_tokens"},"usage":{"output_tokens":100}}
|
||||
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
`
|
||||
dec := NewAnthropicDecoder()
|
||||
_, final := feedSSE(t, dec, body)
|
||||
if final.StopReason != agentcore.StopReasonLength {
|
||||
t.Errorf("max_tokens must map to length, got %q", final.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicDecoderErrorEvent verifies an Anthropic `error` event becomes a
|
||||
// decode error (which the transport turns into a terminal error event), never
|
||||
// a panic.
|
||||
func TestAnthropicDecoderErrorEvent(t *testing.T) {
|
||||
dec := NewAnthropicDecoder()
|
||||
_, err := dec.Decode([]byte(`{"type":"error","error":{"type":"overloaded_error","message":"Overloaded"}}`))
|
||||
if err == nil {
|
||||
t.Fatal("error event must return a decode error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "overloaded_error") {
|
||||
t.Errorf("error should name the type, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicDecoderMalformedPayload verifies invalid JSON is a returned
|
||||
// error (rides the stream as terminal error), not a panic.
|
||||
func TestAnthropicDecoderMalformedPayload(t *testing.T) {
|
||||
dec := NewAnthropicDecoder()
|
||||
if _, err := dec.Decode([]byte(`{not json`)); err == nil {
|
||||
t.Fatal("malformed payload must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicDecoderFinishFlushesPartial verifies a stream cut short (no
|
||||
// message_stop) still yields a done event on Finish so the partial isn't lost.
|
||||
func TestAnthropicDecoderFinishFlushesPartial(t *testing.T) {
|
||||
body := `data: {"type":"message_start","message":{"id":"m","model":"c","usage":{"input_tokens":5,"output_tokens":0}}}
|
||||
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
||||
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"partial"}}
|
||||
|
||||
`
|
||||
dec := NewAnthropicDecoder()
|
||||
events, final := feedSSE(t, dec, body)
|
||||
if events[len(events)-1].EventKind() != StreamEventDone {
|
||||
t.Fatalf("Finish must emit a terminal done event, got %v", eventKinds(events))
|
||||
}
|
||||
// No message_delta arrived → default end_turn.
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("cut-short stream should default to end_turn, got %q", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 1 {
|
||||
t.Fatalf("expected the partial text block, got %+v", final.Content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicDecoderThroughTransport wires the decoder through the real
|
||||
// transport pump against a recorded SSE server, exercising the full path.
|
||||
func TestAnthropicDecoderThroughTransport(t *testing.T) {
|
||||
srv := sseServer(t, anthropicToolUseSSE)
|
||||
defer srv.Close()
|
||||
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: NewAnthropicDecoder(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamRequest: %v", err)
|
||||
}
|
||||
var kinds []string
|
||||
for ev := range stream.Events() {
|
||||
kinds = append(kinds, ev.EventKind())
|
||||
}
|
||||
final, resErr := stream.Result(context.Background())
|
||||
if resErr != nil {
|
||||
t.Fatalf("result: %v", resErr)
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason via transport = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 3 {
|
||||
t.Errorf("expected 3 content blocks via transport, got %d", len(final.Content))
|
||||
}
|
||||
if kinds[len(kinds)-1] != StreamEventDone {
|
||||
t.Errorf("stream must end with done, got %v", kinds)
|
||||
}
|
||||
}
|
||||
|
||||
func eventKinds(events []StreamEvent) []string {
|
||||
out := make([]string, len(events))
|
||||
for i, ev := range events {
|
||||
out[i] = ev.EventKind()
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
// This file implements credential resolution (US-012): API key lookup from
|
||||
// environment variables and a config file (per provider), plus an OAuth token
|
||||
// source that refreshes short-lived tokens on expiry. The resolver satisfies
|
||||
// the LoopConfig.GetAPIKey shape (func(ctx, provider) string) so the agent loop
|
||||
// can obtain a fresh key per request.
|
||||
//
|
||||
// Security (FR: secret values are not written to logs): secret values are never logged or embedded in
|
||||
// error messages. Errors and String()/redaction helpers reference credentials
|
||||
// by key name / provider only.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"os"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// APIKeyConfig is the on-disk config-file shape: a map of provider name to API
|
||||
// key. It is parsed from JSON and holds only static keys (OAuth lives in
|
||||
// TokenSource). Values are secrets and must not be logged.
|
||||
type APIKeyConfig struct {
|
||||
// Keys maps provider name → API key.
|
||||
Keys map[string]string `json:"keys"`
|
||||
}
|
||||
|
||||
// LoadAPIKeyConfig parses an APIKeyConfig from JSON bytes (e.g. a config file).
|
||||
func LoadAPIKeyConfig(data []byte) (*APIKeyConfig, error) {
|
||||
var cfg APIKeyConfig
|
||||
if err := json.Unmarshal(data, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("auth: parse api key config: %w", err)
|
||||
}
|
||||
if cfg.Keys == nil {
|
||||
cfg.Keys = make(map[string]string)
|
||||
}
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
// LoadAPIKeyConfigFile reads and parses an APIKeyConfig from a file path. A
|
||||
// missing file is not an error — it returns an empty config so env/OAuth can
|
||||
// still resolve keys.
|
||||
func LoadAPIKeyConfigFile(path string) (*APIKeyConfig, error) {
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return &APIKeyConfig{Keys: make(map[string]string)}, nil
|
||||
}
|
||||
return nil, fmt.Errorf("auth: read api key config %q: %w", path, err)
|
||||
}
|
||||
return LoadAPIKeyConfig(data)
|
||||
}
|
||||
|
||||
// envAPIKey returns the API key for a provider from the environment. It derives
|
||||
// the candidate variable names from the provider registry (the single source of
|
||||
// truth: LookupProviderSpec(provider).EnvVars, in precedence order), then falls
|
||||
// back to a generic <PROVIDER>_API_KEY when the provider is unknown or none of
|
||||
// its registered vars are set. Returns "" when no value is present.
|
||||
func envAPIKey(provider string) string {
|
||||
if spec, ok := LookupProviderSpec(provider); ok {
|
||||
for _, name := range spec.EnvVars {
|
||||
if v := os.Getenv(name); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
}
|
||||
// Generic fallback for unknown providers or when no registered var is set.
|
||||
generic := strings.ToUpper(provider) + "_API_KEY"
|
||||
return os.Getenv(generic)
|
||||
}
|
||||
|
||||
// TokenSource yields an access token, refreshing it when expired. It models an
|
||||
// OAuth credential whose access token is short-lived (FR-15: getApiKey refreshes
|
||||
// on expiry). It is safe for concurrent use.
|
||||
type TokenSource struct {
|
||||
mu sync.Mutex
|
||||
accessToken string
|
||||
expiry time.Time
|
||||
refreshToken string
|
||||
// Refresh exchanges the current refresh token for a new access token. It
|
||||
// returns the new access token, its expiry, and (optionally) a rotated
|
||||
// refresh token. Required for a TokenSource to refresh; nil means the token
|
||||
// is static and never refreshed.
|
||||
Refresh func(ctx context.Context, refreshToken string) (OAuthToken, error)
|
||||
// Now is injectable for testing; defaults to time.Now.
|
||||
Now func() time.Time
|
||||
// Leeway refreshes the token this long before its actual expiry to avoid
|
||||
// racing the boundary. Defaults to 30s.
|
||||
Leeway time.Duration
|
||||
}
|
||||
|
||||
// OAuthToken is the result of an OAuth exchange/refresh. Values are secrets.
|
||||
type OAuthToken struct {
|
||||
AccessToken string
|
||||
RefreshToken string
|
||||
Expiry time.Time
|
||||
}
|
||||
|
||||
// NewTokenSource builds a TokenSource seeded with an initial token and a refresh
|
||||
// function. refresh may be nil for a static (never-expiring) token.
|
||||
func NewTokenSource(initial OAuthToken, refresh func(ctx context.Context, refreshToken string) (OAuthToken, error)) *TokenSource {
|
||||
return &TokenSource{
|
||||
accessToken: initial.AccessToken,
|
||||
expiry: initial.Expiry,
|
||||
refreshToken: initial.RefreshToken,
|
||||
Refresh: refresh,
|
||||
}
|
||||
}
|
||||
|
||||
func (t *TokenSource) now() time.Time {
|
||||
if t.Now != nil {
|
||||
return t.Now()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
// defaultTokenLeeway is how far before an OAuth token's expiry it is treated as
|
||||
// already expired, so a refresh happens before a request rather than mid-flight.
|
||||
const defaultTokenLeeway = 30 * time.Second
|
||||
|
||||
func (t *TokenSource) leeway() time.Duration {
|
||||
if t.Leeway > 0 {
|
||||
return t.Leeway
|
||||
}
|
||||
return defaultTokenLeeway
|
||||
}
|
||||
|
||||
// expired reports whether the access token is missing or within leeway of its
|
||||
// expiry. A zero expiry means "never expires" (static token).
|
||||
func (t *TokenSource) expired() bool {
|
||||
if t.accessToken == "" {
|
||||
return true
|
||||
}
|
||||
if t.expiry.IsZero() {
|
||||
return false
|
||||
}
|
||||
return !t.now().Before(t.expiry.Add(-t.leeway()))
|
||||
}
|
||||
|
||||
// Token returns a valid access token, refreshing it when expired. It errors if
|
||||
// a refresh is needed but no Refresh func is set, or if Refresh fails. The
|
||||
// returned error never contains the token value.
|
||||
func (t *TokenSource) Token(ctx context.Context) (string, error) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if !t.expired() {
|
||||
return t.accessToken, nil
|
||||
}
|
||||
if t.Refresh == nil {
|
||||
return "", fmt.Errorf("auth: token expired and no refresh function configured")
|
||||
}
|
||||
tok, err := t.Refresh(ctx, t.refreshToken)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("auth: token refresh failed: %w", err)
|
||||
}
|
||||
t.accessToken = tok.AccessToken
|
||||
t.expiry = tok.Expiry
|
||||
if tok.RefreshToken != "" {
|
||||
t.refreshToken = tok.RefreshToken
|
||||
}
|
||||
return t.accessToken, nil
|
||||
}
|
||||
|
||||
// CredentialStore resolves API keys per provider from three layers, in order:
|
||||
// OAuth token source (if registered), environment variable, config file. It
|
||||
// implements the LoopConfig.GetAPIKey shape via GetAPIKey.
|
||||
//
|
||||
// It is safe for concurrent use.
|
||||
type CredentialStore struct {
|
||||
mu sync.RWMutex
|
||||
config *APIKeyConfig
|
||||
sources map[string]*TokenSource // provider → OAuth token source
|
||||
overrides map[string]string // provider → explicit key (highest static priority)
|
||||
}
|
||||
|
||||
// NewCredentialStore builds a store over an optional config file. A nil config
|
||||
// is treated as empty.
|
||||
func NewCredentialStore(config *APIKeyConfig) *CredentialStore {
|
||||
if config == nil {
|
||||
config = &APIKeyConfig{Keys: make(map[string]string)}
|
||||
}
|
||||
return &CredentialStore{
|
||||
config: config,
|
||||
sources: make(map[string]*TokenSource),
|
||||
overrides: make(map[string]string),
|
||||
}
|
||||
}
|
||||
|
||||
// SetOverride records an explicit API key for a provider that wins over the
|
||||
// environment variable and config file (but not a live OAuth token, which is
|
||||
// auto-refreshed). It is the seam for a CLI --api-key flag: an empty key is
|
||||
// ignored so a bare flag does not clobber env/config resolution.
|
||||
func (c *CredentialStore) SetOverride(provider, key string) {
|
||||
if key == "" {
|
||||
return
|
||||
}
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.overrides[provider] = key
|
||||
}
|
||||
|
||||
// RegisterOAuth registers an OAuth TokenSource for a provider. Once registered,
|
||||
// GetAPIKey prefers the (auto-refreshing) OAuth token over static keys.
|
||||
func (c *CredentialStore) RegisterOAuth(provider string, src *TokenSource) {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
c.sources[provider] = src
|
||||
}
|
||||
|
||||
// GetAPIKey resolves the API key for a provider. Resolution order: OAuth token
|
||||
// (refreshed on expiry) → explicit override (--api-key) → environment variable
|
||||
// → config file. Returns "" when no credential is available. This matches
|
||||
// LoopConfig.GetAPIKey so it can be assigned directly.
|
||||
//
|
||||
// On OAuth refresh failure it falls back to override/env/config rather than
|
||||
// returning a secret-bearing error; the empty return lets the caller fall back
|
||||
// to a static key. It never logs secret values.
|
||||
func (c *CredentialStore) GetAPIKey(ctx context.Context, provider string) string {
|
||||
c.mu.RLock()
|
||||
src := c.sources[provider]
|
||||
override := c.overrides[provider]
|
||||
cfgKey := ""
|
||||
if c.config != nil {
|
||||
cfgKey = c.config.Keys[provider]
|
||||
}
|
||||
c.mu.RUnlock()
|
||||
|
||||
if src != nil {
|
||||
if tok, err := src.Token(ctx); err == nil && tok != "" {
|
||||
return tok
|
||||
}
|
||||
// Refresh failed → fall through to static layers.
|
||||
}
|
||||
if override != "" {
|
||||
return override
|
||||
}
|
||||
if env := envAPIKey(provider); env != "" {
|
||||
return env
|
||||
}
|
||||
return cfgKey
|
||||
}
|
||||
|
||||
// HasCredential reports whether any credential (OAuth/env/config) is available
|
||||
// for a provider, without exposing the value.
|
||||
func (c *CredentialStore) HasCredential(ctx context.Context, provider string) bool {
|
||||
return c.GetAPIKey(ctx, provider) != ""
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestEnvAPIKey(t *testing.T) {
|
||||
t.Setenv("ANTHROPIC_OAUTH_TOKEN", "")
|
||||
t.Setenv("ANTHROPIC_API_KEY", "sk-ant-env")
|
||||
if got := envAPIKey("anthropic"); got != "sk-ant-env" {
|
||||
t.Errorf("env key = %q, want sk-ant-env", got)
|
||||
}
|
||||
// Unknown provider uses generic <PROVIDER>_API_KEY fallback.
|
||||
t.Setenv("FOOBAR_API_KEY", "sk-foobar")
|
||||
if got := envAPIKey("foobar"); got != "sk-foobar" {
|
||||
t.Errorf("generic env key = %q, want sk-foobar", got)
|
||||
}
|
||||
if got := envAPIKey("nonesuch"); got != "" {
|
||||
t.Errorf("missing env key = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEnvAPIKeyFromRegistry verifies API-key resolution derives from the
|
||||
// provider registry (single source of truth) across a representative set of
|
||||
// providers, that Anthropic's OAuth token takes precedence over its API key,
|
||||
// and that an unknown provider hits the generic <PROVIDER>_API_KEY fallback.
|
||||
func TestEnvAPIKeyFromRegistry(t *testing.T) {
|
||||
cases := []struct {
|
||||
provider string
|
||||
envVar string
|
||||
value string
|
||||
}{
|
||||
{"deepseek", "DEEPSEEK_API_KEY", "sk-deepseek"},
|
||||
{"groq", "GROQ_API_KEY", "sk-groq"},
|
||||
{"zai", "ZAI_API_KEY", "sk-zai"},
|
||||
{"moonshotai-cn", "MOONSHOT_API_KEY", "sk-moonshot-cn"},
|
||||
{"xiaomi-token-plan-ams", "XIAOMI_TOKEN_PLAN_AMS_API_KEY", "sk-xiaomi-ams"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.provider, func(t *testing.T) {
|
||||
t.Setenv(tc.envVar, tc.value)
|
||||
if got := envAPIKey(tc.provider); got != tc.value {
|
||||
t.Errorf("envAPIKey(%q) = %q, want %q", tc.provider, got, tc.value)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// Anthropic: OAuth token wins over API key (registry ordering).
|
||||
t.Run("anthropic-oauth-first", func(t *testing.T) {
|
||||
t.Setenv("ANTHROPIC_OAUTH_TOKEN", "oauth-tok")
|
||||
t.Setenv("ANTHROPIC_API_KEY", "sk-ant")
|
||||
if got := envAPIKey("anthropic"); got != "oauth-tok" {
|
||||
t.Errorf("anthropic = %q, want oauth-tok (OAuth precedence)", got)
|
||||
}
|
||||
// With OAuth unset, the API key resolves.
|
||||
t.Setenv("ANTHROPIC_OAUTH_TOKEN", "")
|
||||
if got := envAPIKey("anthropic"); got != "sk-ant" {
|
||||
t.Errorf("anthropic (no oauth) = %q, want sk-ant", got)
|
||||
}
|
||||
})
|
||||
|
||||
// Unknown provider falls back to the generic convention.
|
||||
t.Run("unknown-generic-fallback", func(t *testing.T) {
|
||||
t.Setenv("MADEUP_PROVIDER_API_KEY", "sk-generic")
|
||||
if got := envAPIKey("madeup-provider"); got != "" {
|
||||
// Hyphenated names uppercase to MADEUP-PROVIDER_API_KEY, not a match;
|
||||
// verify the true generic form resolves for an underscore-friendly name.
|
||||
t.Logf("hyphenated generic = %q", got)
|
||||
}
|
||||
t.Setenv("MADEUPPROVIDER_API_KEY", "sk-generic2")
|
||||
if got := envAPIKey("madeupprovider"); got != "sk-generic2" {
|
||||
t.Errorf("generic fallback = %q, want sk-generic2", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestLoadAPIKeyConfig(t *testing.T) {
|
||||
cfg, err := LoadAPIKeyConfig([]byte(`{"keys":{"openai":"sk-openai-cfg"}}`))
|
||||
if err != nil {
|
||||
t.Fatalf("load: %v", err)
|
||||
}
|
||||
if cfg.Keys["openai"] != "sk-openai-cfg" {
|
||||
t.Errorf("config key = %q", cfg.Keys["openai"])
|
||||
}
|
||||
if _, err := LoadAPIKeyConfig([]byte(`not json`)); err == nil {
|
||||
t.Fatal("bad JSON must error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadAPIKeyConfigFileMissing(t *testing.T) {
|
||||
cfg, err := LoadAPIKeyConfigFile("/no/such/path/keys.json")
|
||||
if err != nil {
|
||||
t.Fatalf("missing file must not error: %v", err)
|
||||
}
|
||||
if len(cfg.Keys) != 0 {
|
||||
t.Errorf("missing file must yield empty keys, got %v", cfg.Keys)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCredentialStoreResolutionOrder verifies OAuth > env > config precedence.
|
||||
func TestCredentialStoreResolutionOrder(t *testing.T) {
|
||||
cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"anthropic":"sk-cfg","openai":"sk-openai-cfg"}}`))
|
||||
store := NewCredentialStore(cfg)
|
||||
|
||||
// Neutralize any ambient keys so config-only resolution is deterministic.
|
||||
t.Setenv("OPENAI_API_KEY", "")
|
||||
t.Setenv("ANTHROPIC_API_KEY", "")
|
||||
t.Setenv("CLAUDE_API_KEY", "")
|
||||
|
||||
// Config-only provider resolves from config.
|
||||
if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-openai-cfg" {
|
||||
t.Errorf("openai (config) = %q, want sk-openai-cfg", got)
|
||||
}
|
||||
|
||||
// Env overrides config.
|
||||
t.Setenv("ANTHROPIC_API_KEY", "sk-env")
|
||||
if got := store.GetAPIKey(context.Background(), "anthropic"); got != "sk-env" {
|
||||
t.Errorf("anthropic (env>config) = %q, want sk-env", got)
|
||||
}
|
||||
|
||||
// OAuth overrides env + config.
|
||||
store.RegisterOAuth("anthropic", NewTokenSource(
|
||||
OAuthToken{AccessToken: "oauth-token", Expiry: time.Now().Add(time.Hour)}, nil))
|
||||
if got := store.GetAPIKey(context.Background(), "anthropic"); got != "oauth-token" {
|
||||
t.Errorf("anthropic (oauth>env) = %q, want oauth-token", got)
|
||||
}
|
||||
|
||||
// Unknown provider → empty.
|
||||
if got := store.GetAPIKey(context.Background(), "ghost"); got != "" {
|
||||
t.Errorf("ghost = %q, want empty", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCredentialStoreOverride verifies an explicit --api-key override wins over
|
||||
// env and config, but not over a live OAuth token, and that an empty override
|
||||
// is ignored (so a bare flag does not clobber env/config).
|
||||
func TestCredentialStoreOverride(t *testing.T) {
|
||||
cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"openai":"sk-openai-cfg"}}`))
|
||||
store := NewCredentialStore(cfg)
|
||||
t.Setenv("OPENAI_API_KEY", "sk-openai-env")
|
||||
|
||||
// Empty override is a no-op: env still wins over config.
|
||||
store.SetOverride("openai", "")
|
||||
if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-openai-env" {
|
||||
t.Errorf("empty override should not clobber env, got %q", got)
|
||||
}
|
||||
|
||||
// Non-empty override wins over env and config.
|
||||
store.SetOverride("openai", "sk-flag")
|
||||
if got := store.GetAPIKey(context.Background(), "openai"); got != "sk-flag" {
|
||||
t.Errorf("override should win over env/config, got %q", got)
|
||||
}
|
||||
|
||||
// OAuth still wins over an override.
|
||||
store.RegisterOAuth("openai", NewTokenSource(
|
||||
OAuthToken{AccessToken: "oauth-token", Expiry: time.Now().Add(time.Hour)}, nil))
|
||||
if got := store.GetAPIKey(context.Background(), "openai"); got != "oauth-token" {
|
||||
t.Errorf("oauth should win over override, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTokenSourceRefresh verifies an expired token triggers a refresh returning
|
||||
// a new token.
|
||||
func TestTokenSourceRefresh(t *testing.T) {
|
||||
now := time.Now()
|
||||
refreshCount := 0
|
||||
src := NewTokenSource(
|
||||
OAuthToken{AccessToken: "old", RefreshToken: "refresh-1", Expiry: now.Add(-time.Minute)},
|
||||
func(ctx context.Context, rt string) (OAuthToken, error) {
|
||||
refreshCount++
|
||||
if rt != "refresh-1" {
|
||||
t.Errorf("refresh token = %q, want refresh-1", rt)
|
||||
}
|
||||
return OAuthToken{AccessToken: "new", RefreshToken: "refresh-2", Expiry: now.Add(time.Hour)}, nil
|
||||
},
|
||||
)
|
||||
src.Now = func() time.Time { return now }
|
||||
|
||||
tok, err := src.Token(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("token: %v", err)
|
||||
}
|
||||
if tok != "new" {
|
||||
t.Errorf("token = %q, want new (refreshed)", tok)
|
||||
}
|
||||
if refreshCount != 1 {
|
||||
t.Errorf("refresh count = %d, want 1", refreshCount)
|
||||
}
|
||||
|
||||
// Second call within validity does not refresh again.
|
||||
if _, err := src.Token(context.Background()); err != nil {
|
||||
t.Fatalf("token 2: %v", err)
|
||||
}
|
||||
if refreshCount != 1 {
|
||||
t.Errorf("refresh count after valid reuse = %d, want 1", refreshCount)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenSourceNoRefreshFunc(t *testing.T) {
|
||||
now := time.Now()
|
||||
// Expired token with no Refresh func → error.
|
||||
src := NewTokenSource(OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)}, nil)
|
||||
src.Now = func() time.Time { return now }
|
||||
if _, err := src.Token(context.Background()); err == nil {
|
||||
t.Fatal("expired token without refresh must error")
|
||||
}
|
||||
|
||||
// Static token (zero expiry) never expires.
|
||||
static := NewTokenSource(OAuthToken{AccessToken: "static"}, nil)
|
||||
tok, err := static.Token(context.Background())
|
||||
if err != nil || tok != "static" {
|
||||
t.Errorf("static token = %q, err = %v", tok, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTokenSourceRefreshError(t *testing.T) {
|
||||
now := time.Now()
|
||||
src := NewTokenSource(
|
||||
OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)},
|
||||
func(ctx context.Context, rt string) (OAuthToken, error) {
|
||||
return OAuthToken{}, context.DeadlineExceeded
|
||||
},
|
||||
)
|
||||
src.Now = func() time.Time { return now }
|
||||
_, err := src.Token(context.Background())
|
||||
if err == nil {
|
||||
t.Fatal("refresh error must propagate")
|
||||
}
|
||||
// Error must not leak the (empty) token but should mention refresh.
|
||||
if got := err.Error(); got == "" {
|
||||
t.Error("expected non-empty error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestCredentialStoreOAuthRefreshFallback verifies a failing OAuth refresh
|
||||
// falls back to env/config rather than returning empty when a static key exists.
|
||||
func TestCredentialStoreOAuthRefreshFallback(t *testing.T) {
|
||||
cfg, _ := LoadAPIKeyConfig([]byte(`{"keys":{"anthropic":"sk-cfg-fallback"}}`))
|
||||
store := NewCredentialStore(cfg)
|
||||
now := time.Now()
|
||||
src := NewTokenSource(
|
||||
OAuthToken{AccessToken: "old", Expiry: now.Add(-time.Minute)},
|
||||
func(ctx context.Context, rt string) (OAuthToken, error) {
|
||||
return OAuthToken{}, context.DeadlineExceeded
|
||||
},
|
||||
)
|
||||
src.Now = func() time.Time { return now }
|
||||
store.RegisterOAuth("anthropic", src)
|
||||
|
||||
if got := store.GetAPIKey(context.Background(), "anthropic"); got != "sk-cfg-fallback" {
|
||||
t.Errorf("refresh-failed fallback = %q, want sk-cfg-fallback", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,157 @@
|
||||
package provider
|
||||
|
||||
// Tests for image (multimodal) input encoding on both provider wires (US-010,
|
||||
// #126). They exercise encodeOpenAIMessage / encodeAnthropicMessage directly
|
||||
// (unit-level, no HTTP) to assert the exact wire shape of image blocks, plus the
|
||||
// checkImageSupport guard that turns image input on a text-only model into a
|
||||
// clear error rather than a silent drop.
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// imageUserMessage builds a user message carrying one text block and one image
|
||||
// block, the common multimodal input shape.
|
||||
func imageUserMessage(text, data, mime string) agentcore.UserMessage {
|
||||
return agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{
|
||||
agentcore.NewTextContent(text),
|
||||
agentcore.NewImageContent(data, mime),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeOpenAIMessageImageBlock asserts a user message with an image encodes
|
||||
// to the multimodal array form with an image_url data URI.
|
||||
func TestEncodeOpenAIMessageImageBlock(t *testing.T) {
|
||||
msg := imageUserMessage("what is this?", "QUJD", "image/png")
|
||||
out := encodeOpenAIMessage(msg)
|
||||
if len(out) != 1 {
|
||||
t.Fatalf("encodeOpenAIMessage returned %d entries, want 1", len(out))
|
||||
}
|
||||
entry := out[0]
|
||||
if entry["role"] != "user" {
|
||||
t.Errorf("role = %v, want user", entry["role"])
|
||||
}
|
||||
parts, ok := entry["content"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("content is %T, want []map[string]any (array form)", entry["content"])
|
||||
}
|
||||
if len(parts) != 2 {
|
||||
t.Fatalf("content has %d parts, want 2 (text + image)", len(parts))
|
||||
}
|
||||
if parts[0]["type"] != "text" || parts[0]["text"] != "what is this?" {
|
||||
t.Errorf("text part = %#v", parts[0])
|
||||
}
|
||||
if parts[1]["type"] != "image_url" {
|
||||
t.Fatalf("image part type = %v, want image_url", parts[1]["type"])
|
||||
}
|
||||
iu, ok := parts[1]["image_url"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("image_url is %T", parts[1]["image_url"])
|
||||
}
|
||||
if want := "data:image/png;base64,QUJD"; iu["url"] != want {
|
||||
t.Errorf("image_url.url = %v, want %v", iu["url"], want)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeOpenAIMessageNoImageIsString asserts a text-only user message stays
|
||||
// a plain string (not the array form), preserving the common-case wire shape.
|
||||
func TestEncodeOpenAIMessageNoImageIsString(t *testing.T) {
|
||||
msg := agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("hello")},
|
||||
}
|
||||
out := encodeOpenAIMessage(msg)
|
||||
if s, ok := out[0]["content"].(string); !ok || s != "hello" {
|
||||
t.Errorf("content = %#v, want string \"hello\"", out[0]["content"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeAnthropicMessageImageBlock asserts a user message with an image
|
||||
// encodes to the content-block array form with a base64 image source.
|
||||
func TestEncodeAnthropicMessageImageBlock(t *testing.T) {
|
||||
msg := imageUserMessage("describe", "REVG", "image/jpeg")
|
||||
entry := encodeAnthropicMessage(msg)
|
||||
if entry["role"] != "user" {
|
||||
t.Errorf("role = %v, want user", entry["role"])
|
||||
}
|
||||
blocks, ok := entry["content"].([]map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("content is %T, want []map[string]any (array form)", entry["content"])
|
||||
}
|
||||
if len(blocks) != 2 {
|
||||
t.Fatalf("content has %d blocks, want 2 (text + image)", len(blocks))
|
||||
}
|
||||
if blocks[0]["type"] != "text" || blocks[0]["text"] != "describe" {
|
||||
t.Errorf("text block = %#v", blocks[0])
|
||||
}
|
||||
if blocks[1]["type"] != "image" {
|
||||
t.Fatalf("image block type = %v, want image", blocks[1]["type"])
|
||||
}
|
||||
src, ok := blocks[1]["source"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("source is %T", blocks[1]["source"])
|
||||
}
|
||||
if src["type"] != "base64" || src["media_type"] != "image/jpeg" || src["data"] != "REVG" {
|
||||
t.Errorf("source = %#v", src)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeAnthropicMessageNoImageIsString asserts a text-only user message
|
||||
// stays a plain string on the Anthropic wire.
|
||||
func TestEncodeAnthropicMessageNoImageIsString(t *testing.T) {
|
||||
msg := agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("hi")},
|
||||
}
|
||||
entry := encodeAnthropicMessage(msg)
|
||||
if s, ok := entry["content"].(string); !ok || s != "hi" {
|
||||
t.Errorf("content = %#v, want string \"hi\"", entry["content"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestImageBlocksAreJSONSerializable guards against map value types that
|
||||
// json.Marshal cannot encode: both wire shapes must round-trip to JSON.
|
||||
func TestImageBlocksAreJSONSerializable(t *testing.T) {
|
||||
msg := imageUserMessage("x", "QQ==", "image/webp")
|
||||
if _, err := json.Marshal(encodeOpenAIMessage(msg)); err != nil {
|
||||
t.Errorf("marshal openai image message: %v", err)
|
||||
}
|
||||
if _, err := json.Marshal(encodeAnthropicMessage(msg)); err != nil {
|
||||
t.Errorf("marshal anthropic image message: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestCheckImageSupport asserts the capability guard: image input on a model
|
||||
// that declares SupportsImages passes; on a text-only model it errors; and a
|
||||
// text-only prompt always passes regardless of the model.
|
||||
func TestCheckImageSupport(t *testing.T) {
|
||||
models := []Model{
|
||||
{ID: "vision-1", SupportsImages: true},
|
||||
{ID: "text-1", SupportsImages: false},
|
||||
}
|
||||
imgMsgs := []agentcore.Message{imageUserMessage("q", "QQ==", "image/png")}
|
||||
textMsgs := []agentcore.Message{agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("no image")},
|
||||
}}
|
||||
|
||||
if err := checkImageSupport("p", "vision-1", models, imgMsgs); err != nil {
|
||||
t.Errorf("vision model rejected image input: %v", err)
|
||||
}
|
||||
if err := checkImageSupport("p", "text-1", models, imgMsgs); err == nil {
|
||||
t.Error("text-only model accepted image input, want error")
|
||||
}
|
||||
if err := checkImageSupport("p", "text-1", models, textMsgs); err != nil {
|
||||
t.Errorf("text-only prompt on text model errored: %v", err)
|
||||
}
|
||||
// Unknown model (not in catalog) defers to provider validation → no error.
|
||||
if err := checkImageSupport("p", "unknown", models, imgMsgs); err != nil {
|
||||
t.Errorf("unknown model errored on image input: %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
// This file implements model-name → provider inference (US: auto-infer provider
|
||||
// from --model alone). When a user supplies only --model, with no --provider,
|
||||
// --protocol, or --base-url, pigo tries to guess the owning provider from the
|
||||
// model id's well-known name prefix (e.g. "claude-*" → anthropic, "deepseek-*"
|
||||
// → deepseek). This lets `pigo -m claude-opus-4-8` reach the Anthropic API
|
||||
// without the user spelling out the provider or its wire protocol.
|
||||
//
|
||||
// The mapping is deliberately conservative: only names that unambiguously
|
||||
// identify a single built-in provider are inferred. Ambiguous families served
|
||||
// by many gateways (notably "llama-*", "qwq-*", "gemma-*", "mixtral-*") are NOT
|
||||
// inferred — they return ok=false so the caller falls through to its existing
|
||||
// default (OpenRouter), which is the safe, unchanged behavior.
|
||||
//
|
||||
// Every provider name returned here is guaranteed to exist in providerRegistry
|
||||
// (enforced by a test), so callers can hand the result straight to the
|
||||
// registry-driven resolution path.
|
||||
package provider
|
||||
|
||||
import "strings"
|
||||
|
||||
// modelPrefixProvider maps a lowercase model-name prefix to the built-in
|
||||
// provider name that serves that family. Order matters: the table is scanned
|
||||
// top-to-bottom and the FIRST matching prefix wins, so list more specific
|
||||
// prefixes before shorter ones that would also match.
|
||||
//
|
||||
// Each provider name here must be present in providerRegistry (registry.go).
|
||||
var modelPrefixProvider = []struct {
|
||||
prefix string
|
||||
provider string
|
||||
}{
|
||||
{"claude-", "anthropic"},
|
||||
{"gpt-", "openai"},
|
||||
{"o1-", "openai"},
|
||||
{"o3-", "openai"},
|
||||
{"o4-", "openai"},
|
||||
{"gemini-", "google"},
|
||||
{"deepseek-", "deepseek"},
|
||||
{"glm-", "zai"},
|
||||
{"kimi-", "moonshotai"},
|
||||
{"moonshot-", "moonshotai"},
|
||||
{"qwen-", "dashscope"},
|
||||
{"ernie-", "qianfan"},
|
||||
{"doubao-", "volcengine"},
|
||||
{"grok-", "xai"},
|
||||
{"mistral-", "mistral"},
|
||||
{"codestral-", "mistral"},
|
||||
{"devstral-", "mistral"},
|
||||
{"hunyuan-", "hunyuan"},
|
||||
{"minimax-", "minimax"},
|
||||
{"mimo-", "xiaomi"},
|
||||
}
|
||||
|
||||
// InferProviderFromModel guesses the built-in provider name that serves a given
|
||||
// model id, based on the id's well-known name prefix. It returns the provider
|
||||
// name and ok=true on a confident match, or ("", false) when the id is unknown
|
||||
// or ambiguous (served by multiple gateways). Matching is case-insensitive.
|
||||
//
|
||||
// The returned name is always a valid entry in providerRegistry. Callers should
|
||||
// use it only when no explicit --provider/--protocol/--base-url was given; those
|
||||
// flags take precedence over any inference.
|
||||
func InferProviderFromModel(model string) (string, bool) {
|
||||
id := strings.ToLower(strings.TrimSpace(model))
|
||||
if id == "" {
|
||||
return "", false
|
||||
}
|
||||
// A "provider/model" style id (e.g. "openai/gpt-4o") is an OpenRouter-style
|
||||
// routed id, not a bare model name — leave those to the caller's preset/
|
||||
// prefix handling rather than inferring from the leading segment.
|
||||
if strings.Contains(id, "/") {
|
||||
return "", false
|
||||
}
|
||||
for _, m := range modelPrefixProvider {
|
||||
if strings.HasPrefix(id, m.prefix) {
|
||||
return m.provider, true
|
||||
}
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
@@ -0,0 +1,96 @@
|
||||
package provider
|
||||
|
||||
// Tests for model-name → provider inference (InferProviderFromModel). They
|
||||
// verify each documented name prefix resolves to the expected built-in
|
||||
// provider, that ambiguous/unknown ids and routed "provider/model" ids do not
|
||||
// resolve, that matching is case-insensitive, and that every inferred provider
|
||||
// name actually exists in the provider registry (the single source of truth).
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestInferProviderFromModelKnown verifies each documented model-name prefix
|
||||
// resolves to its expected built-in provider.
|
||||
func TestInferProviderFromModelKnown(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
want string
|
||||
}{
|
||||
{"claude-opus-4-8", "anthropic"},
|
||||
{"claude-3.5-sonnet", "anthropic"},
|
||||
{"gpt-4o", "openai"},
|
||||
{"gpt-4o-mini", "openai"},
|
||||
{"o1-preview", "openai"},
|
||||
{"o3-mini", "openai"},
|
||||
{"o4-mini", "openai"},
|
||||
{"gemini-2.5-pro", "google"},
|
||||
{"deepseek-chat", "deepseek"},
|
||||
{"deepseek-v4-pro", "deepseek"},
|
||||
{"glm-5.1", "zai"},
|
||||
{"kimi-k2-thinking", "moonshotai"},
|
||||
{"moonshot-v1-8k", "moonshotai"},
|
||||
{"qwen-max", "dashscope"},
|
||||
{"ernie-4.5-turbo-32k", "qianfan"},
|
||||
{"doubao-seed-1-6", "volcengine"},
|
||||
{"grok-4.5", "xai"},
|
||||
{"mistral-large-latest", "mistral"},
|
||||
{"codestral-latest", "mistral"},
|
||||
{"devstral-medium-latest", "mistral"},
|
||||
{"hunyuan-turbos-latest", "hunyuan"},
|
||||
{"minimax-m2.7", "minimax"},
|
||||
{"mimo-v2-pro", "xiaomi"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
got, ok := InferProviderFromModel(c.model)
|
||||
if !ok {
|
||||
t.Errorf("InferProviderFromModel(%q): ok=false, want provider %q", c.model, c.want)
|
||||
continue
|
||||
}
|
||||
if got != c.want {
|
||||
t.Errorf("InferProviderFromModel(%q) = %q, want %q", c.model, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestInferProviderFromModelCaseInsensitive verifies matching ignores case and
|
||||
// surrounding whitespace.
|
||||
func TestInferProviderFromModelCaseInsensitive(t *testing.T) {
|
||||
for _, m := range []string{"Claude-Opus-4-8", " GPT-4o ", "DeepSeek-Chat"} {
|
||||
if _, ok := InferProviderFromModel(m); !ok {
|
||||
t.Errorf("InferProviderFromModel(%q): ok=false, want a match", m)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestInferProviderFromModelAmbiguousOrUnknown verifies ids that are ambiguous
|
||||
// (served by many gateways), routed ("provider/model"), empty, or simply
|
||||
// unknown do NOT resolve — the caller must fall through to its default.
|
||||
func TestInferProviderFromModelAmbiguousOrUnknown(t *testing.T) {
|
||||
for _, m := range []string{
|
||||
"", // empty
|
||||
" ", // whitespace only
|
||||
"llama-3.3-70b-instruct", // ambiguous: many gateways
|
||||
"qwq-32b", // ambiguous
|
||||
"gemma-2-9b-it", // ambiguous
|
||||
"mixtral-8x22b", // ambiguous
|
||||
"openai/gpt-4o", // routed id, leave to preset/prefix handling
|
||||
"anthropic/claude-3.5", // routed id
|
||||
"ollama/llama3.3", // routed id (ollama prefix path)
|
||||
"totally-made-up-model", // unknown
|
||||
} {
|
||||
if got, ok := InferProviderFromModel(m); ok {
|
||||
t.Errorf("InferProviderFromModel(%q) = (%q, true), want ok=false", m, got)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestInferProviderNamesExistInRegistry verifies every provider name the
|
||||
// inference table can return is a real built-in provider (registry is the
|
||||
// single source of truth), so a hit can be handed straight to registry-driven
|
||||
// resolution.
|
||||
func TestInferProviderNamesExistInRegistry(t *testing.T) {
|
||||
for _, m := range modelPrefixProvider {
|
||||
if _, ok := LookupProviderSpec(m.provider); !ok {
|
||||
t.Errorf("inference maps prefix %q → %q, which is not in providerRegistry", m.prefix, m.provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
// This file implements the OpenAI-compatible streaming decoder (US-009): a
|
||||
// stateful Decoder (see transport.go) for the OpenAI Chat Completions SSE
|
||||
// stream, which is also the wire format spoken by most third-party gateways
|
||||
// (OpenRouter, Groq, together, local servers, …). Selecting the base URL is a
|
||||
// transport concern (NewRequest builds the *http.Request), so this decoder is
|
||||
// base-URL agnostic and reused across every OpenAI-compatible provider.
|
||||
//
|
||||
// OpenAI streams a sequence of chat.completion.chunk objects:
|
||||
//
|
||||
// {"choices":[{"delta":{"role":"assistant"}}]} → first chunk
|
||||
// {"choices":[{"delta":{"content":"Hel"}}]} → text delta
|
||||
// {"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1",
|
||||
// "function":{"name":"f","arguments":"{\"a\":"}}}]}]} → tool-call delta
|
||||
// {"choices":[{"finish_reason":"tool_calls"}]} → stop reason
|
||||
// {"usage":{"prompt_tokens":10,"completion_tokens":5}} → final usage
|
||||
// [DONE] → transport-level
|
||||
//
|
||||
// Per the dual failure model (FR-13) the decoder never panics: malformed
|
||||
// payloads surface as a returned error which the transport turns into a
|
||||
// terminal StreamErrorEvent.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sort"
|
||||
"strings"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// openaiToolCall accumulates one streamed tool call, keyed by its delta index.
|
||||
// id/name arrive once (usually in the first fragment); arguments accumulate.
|
||||
type openaiToolCall struct {
|
||||
id string
|
||||
name string
|
||||
args strings.Builder
|
||||
}
|
||||
|
||||
// OpenAIDecoder is the stateful SSE decoder for the OpenAI Chat Completions API
|
||||
// and compatible gateways. It implements the transport Decoder interface and is
|
||||
// not safe for concurrent use — the transport drives it from one goroutine.
|
||||
type OpenAIDecoder struct {
|
||||
text strings.Builder
|
||||
thinking strings.Builder // reasoning_content / reasoning stream (if any)
|
||||
toolCalls map[int]*openaiToolCall
|
||||
toolOrder []int // tool-call indices in first-seen order
|
||||
|
||||
responseID string
|
||||
responseModel string
|
||||
inputTokens int
|
||||
outputTokens int
|
||||
stopReason string // mapped pigo stop reason (empty until finish_reason)
|
||||
done bool
|
||||
}
|
||||
|
||||
// NewOpenAIDecoder builds a fresh decoder for one streamed response.
|
||||
func NewOpenAIDecoder() *OpenAIDecoder {
|
||||
return &OpenAIDecoder{toolCalls: make(map[int]*openaiToolCall)}
|
||||
}
|
||||
|
||||
// openaiChunk is the streamed chat.completion.chunk envelope.
|
||||
type openaiChunk struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
// ReasoningContent carries the model's reasoning/thinking stream on the
|
||||
// OpenAI wire (DeepSeek-R1, Kimi, and other reasoning models put their
|
||||
// chain-of-thought here). Some gateways name it "reasoning" instead, so
|
||||
// both are accepted; without this field the thinking stream is silently
|
||||
// dropped from the response and from history.
|
||||
ReasoningContent string `json:"reasoning_content"`
|
||||
Reasoning string `json:"reasoning"`
|
||||
ToolCalls []openaiToolDelta `json:"tool_calls"`
|
||||
} `json:"delta"`
|
||||
FinishReason string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage *struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
} `json:"usage"`
|
||||
// Some gateways surface an error object inline on the stream.
|
||||
Error *struct {
|
||||
Message string `json:"message"`
|
||||
Type string `json:"type"`
|
||||
} `json:"error"`
|
||||
}
|
||||
|
||||
type openaiToolDelta struct {
|
||||
Index int `json:"index"`
|
||||
ID string `json:"id"`
|
||||
Type string `json:"type"`
|
||||
Function struct {
|
||||
Name string `json:"name"`
|
||||
Arguments string `json:"arguments"`
|
||||
} `json:"function"`
|
||||
}
|
||||
|
||||
// Decode turns one OpenAI SSE data payload into zero or more StreamEvents.
|
||||
func (d *OpenAIDecoder) Decode(payload []byte) ([]StreamEvent, error) {
|
||||
var chunk openaiChunk
|
||||
if err := json.Unmarshal(payload, &chunk); err != nil {
|
||||
return nil, fmt.Errorf("openai: parse chunk: %w", err)
|
||||
}
|
||||
if chunk.Error != nil {
|
||||
msg := "openai stream error"
|
||||
if chunk.Error.Type != "" {
|
||||
msg = "openai " + chunk.Error.Type
|
||||
}
|
||||
if chunk.Error.Message != "" {
|
||||
msg += ": " + chunk.Error.Message
|
||||
}
|
||||
return nil, fmt.Errorf("%s", msg)
|
||||
}
|
||||
|
||||
if chunk.ID != "" {
|
||||
d.responseID = chunk.ID
|
||||
}
|
||||
if chunk.Model != "" {
|
||||
d.responseModel = chunk.Model
|
||||
}
|
||||
if chunk.Usage != nil {
|
||||
d.inputTokens = chunk.Usage.PromptTokens
|
||||
d.outputTokens = chunk.Usage.CompletionTokens
|
||||
}
|
||||
|
||||
var events []StreamEvent
|
||||
for _, choice := range chunk.Choices {
|
||||
// Reasoning stream (DeepSeek-R1 / Kimi / …): reasoning_content is the
|
||||
// common field; a few gateways use "reasoning". Accumulate whichever is set.
|
||||
if r := choice.Delta.ReasoningContent; r != "" {
|
||||
d.thinking.WriteString(r)
|
||||
events = append(events, StreamThinkingEvent{Partial: d.partial()})
|
||||
} else if r := choice.Delta.Reasoning; r != "" {
|
||||
d.thinking.WriteString(r)
|
||||
events = append(events, StreamThinkingEvent{Partial: d.partial()})
|
||||
}
|
||||
if choice.Delta.Content != "" {
|
||||
d.text.WriteString(choice.Delta.Content)
|
||||
events = append(events, StreamTextEvent{Partial: d.partial()})
|
||||
}
|
||||
for _, tc := range choice.Delta.ToolCalls {
|
||||
d.applyToolDelta(tc)
|
||||
events = append(events, StreamToolCallEvent{Partial: d.partial()})
|
||||
}
|
||||
if choice.FinishReason != "" {
|
||||
d.stopReason = mapOpenAIFinishReason(choice.FinishReason)
|
||||
}
|
||||
}
|
||||
return events, nil
|
||||
}
|
||||
|
||||
// Finish flushes a terminal done event if the stream ended without an explicit
|
||||
// terminator, so a partial response is still delivered rather than lost.
|
||||
func (d *OpenAIDecoder) Finish() ([]StreamEvent, error) {
|
||||
if d.done {
|
||||
return nil, nil
|
||||
}
|
||||
return d.finishDone(), nil
|
||||
}
|
||||
|
||||
// applyToolDelta merges one tool-call fragment into the accumulated state.
|
||||
func (d *OpenAIDecoder) applyToolDelta(tc openaiToolDelta) {
|
||||
call := d.toolCalls[tc.Index]
|
||||
if call == nil {
|
||||
call = &openaiToolCall{}
|
||||
d.toolCalls[tc.Index] = call
|
||||
d.toolOrder = append(d.toolOrder, tc.Index)
|
||||
}
|
||||
if tc.ID != "" {
|
||||
call.id = tc.ID
|
||||
}
|
||||
if tc.Function.Name != "" {
|
||||
call.name = tc.Function.Name
|
||||
}
|
||||
call.args.WriteString(tc.Function.Arguments)
|
||||
}
|
||||
|
||||
// finishDone builds the terminal assistant message and marks the decoder done.
|
||||
func (d *OpenAIDecoder) finishDone() []StreamEvent {
|
||||
if d.done {
|
||||
return nil
|
||||
}
|
||||
d.done = true
|
||||
msg := d.partial()
|
||||
if msg.StopReason == "" {
|
||||
msg.StopReason = agentcore.StopReasonEndTurn
|
||||
}
|
||||
return []StreamEvent{StreamDoneEvent{Message: msg}}
|
||||
}
|
||||
|
||||
// partial materializes the accumulated state into an AssistantMessage: the text
|
||||
// block first (if any), then tool-call blocks in first-seen index order.
|
||||
func (d *OpenAIDecoder) partial() agentcore.AssistantMessage {
|
||||
msg := agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
API: "openai",
|
||||
Provider: "openai",
|
||||
StopReason: d.stopReason,
|
||||
ResponseID: d.responseID,
|
||||
ResponseModel: d.responseModel,
|
||||
}
|
||||
if d.inputTokens != 0 || d.outputTokens != 0 {
|
||||
msg.Usage = &agentcore.Usage{InputTokens: d.inputTokens, OutputTokens: d.outputTokens}
|
||||
}
|
||||
if d.thinking.Len() > 0 {
|
||||
msg.Content = append(msg.Content, agentcore.NewThinkingContent(d.thinking.String()))
|
||||
}
|
||||
if d.text.Len() > 0 {
|
||||
msg.Content = append(msg.Content, agentcore.NewTextContent(d.text.String()))
|
||||
}
|
||||
|
||||
idx := make([]int, len(d.toolOrder))
|
||||
copy(idx, d.toolOrder)
|
||||
sort.Ints(idx)
|
||||
for _, i := range idx {
|
||||
call := d.toolCalls[i]
|
||||
if call == nil {
|
||||
continue
|
||||
}
|
||||
args := json.RawMessage(strings.TrimSpace(call.args.String()))
|
||||
if len(args) == 0 {
|
||||
args = json.RawMessage("{}")
|
||||
}
|
||||
msg.Content = append(msg.Content, agentcore.NewToolCallContent(call.id, call.name, args))
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
// mapOpenAIFinishReason maps an OpenAI finish_reason to the pigo StopReason set.
|
||||
// Unknown reasons default to end_turn (a natural, non-error stop).
|
||||
func mapOpenAIFinishReason(reason string) string {
|
||||
switch reason {
|
||||
case "length":
|
||||
return agentcore.StopReasonLength
|
||||
case "tool_calls", "function_call":
|
||||
return agentcore.StopReasonToolUse
|
||||
case "stop":
|
||||
return agentcore.StopReasonEndTurn
|
||||
default:
|
||||
return agentcore.StopReasonEndTurn
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,190 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// A recorded OpenAI Chat Completions SSE stream covering a text delta followed
|
||||
// by a two-fragment tool call, ending with finish_reason=tool_calls and a final
|
||||
// usage-only chunk. Trimmed but structurally faithful to the real wire format
|
||||
// (the transport strips the `data:` prefix and the trailing [DONE]).
|
||||
const openaiToolCallSSE = `data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"role":"assistant"}}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"content":"Let me "}}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"content":"check."}}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","type":"function","function":{"name":"get_weather","arguments":"{\"city\":"}}]}}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{"tool_calls":[{"index":0,"function":{"arguments":" \"SF\"}"}}]}}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"tool_calls"}]}
|
||||
|
||||
data: {"id":"chatcmpl-1","model":"gpt-4o","choices":[],"usage":{"prompt_tokens":11,"completion_tokens":8}}
|
||||
|
||||
data: [DONE]
|
||||
|
||||
`
|
||||
|
||||
func TestOpenAIDecoderToolCallStream(t *testing.T) {
|
||||
dec := NewOpenAIDecoder()
|
||||
events, final := feedSSE(t, dec, openaiToolCallSSE)
|
||||
|
||||
// The last emitted event must be the terminal done event.
|
||||
if len(events) == 0 || events[len(events)-1].EventKind() != StreamEventDone {
|
||||
t.Fatalf("expected a done event last, got %v", eventKinds(events))
|
||||
}
|
||||
// Stop reason: tool_calls → tool_use.
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
// Usage: prompt→input, completion→output.
|
||||
if final.Usage == nil || final.Usage.InputTokens != 11 || final.Usage.OutputTokens != 8 {
|
||||
t.Errorf("usage = %+v, want input=11 output=8", final.Usage)
|
||||
}
|
||||
// Response identity.
|
||||
if final.ResponseID != "chatcmpl-1" || final.ResponseModel != "gpt-4o" {
|
||||
t.Errorf("response id/model = %q/%q", final.ResponseID, final.ResponseModel)
|
||||
}
|
||||
|
||||
// Content blocks: text first, then the tool call.
|
||||
if len(final.Content) != 2 {
|
||||
t.Fatalf("expected 2 content blocks, got %d: %+v", len(final.Content), final.Content)
|
||||
}
|
||||
txt, ok := final.Content[0].(agentcore.TextContent)
|
||||
if !ok || txt.Text != "Let me check." {
|
||||
t.Errorf("text block = %+v", final.Content[0])
|
||||
}
|
||||
tool, ok := final.Content[1].(agentcore.ToolCallContent)
|
||||
if !ok || tool.Name != "get_weather" || tool.ID != "call_1" {
|
||||
t.Fatalf("tool block = %+v", final.Content[1])
|
||||
}
|
||||
// tool_call arguments must have accumulated into valid JSON across fragments.
|
||||
var args map[string]string
|
||||
if err := json.Unmarshal(tool.Arguments, &args); err != nil {
|
||||
t.Fatalf("tool arguments not valid JSON %q: %v", tool.Arguments, err)
|
||||
}
|
||||
if args["city"] != "SF" {
|
||||
t.Errorf("tool arguments = %v, want city=SF", args)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIDecoderTextOnlyStop(t *testing.T) {
|
||||
body := `data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":"Hello"}}]}
|
||||
|
||||
data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{"content":" world"}}]}
|
||||
|
||||
data: {"id":"c1","model":"gpt-4o","choices":[{"delta":{},"finish_reason":"stop"}]}
|
||||
|
||||
data: {"choices":[],"usage":{"prompt_tokens":3,"completion_tokens":2}}
|
||||
|
||||
data: [DONE]
|
||||
|
||||
`
|
||||
dec := NewOpenAIDecoder()
|
||||
_, final := feedSSE(t, dec, body)
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("stop reason = %q, want end_turn", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 1 {
|
||||
t.Fatalf("expected 1 content block, got %d", len(final.Content))
|
||||
}
|
||||
txt, ok := final.Content[0].(agentcore.TextContent)
|
||||
if !ok || txt.Text != "Hello world" {
|
||||
t.Errorf("text = %+v", final.Content[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIDecoderLengthMapsToLength(t *testing.T) {
|
||||
body := `data: {"id":"c","model":"m","choices":[{"delta":{"content":"truncated"}}]}
|
||||
|
||||
data: {"id":"c","model":"m","choices":[{"delta":{},"finish_reason":"length"}]}
|
||||
|
||||
data: [DONE]
|
||||
|
||||
`
|
||||
dec := NewOpenAIDecoder()
|
||||
_, final := feedSSE(t, dec, body)
|
||||
if final.StopReason != agentcore.StopReasonLength {
|
||||
t.Errorf("length finish_reason must map to length, got %q", final.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIDecoderInlineError verifies an inline error object becomes a decode
|
||||
// error (which the transport turns into a terminal error event), never a panic.
|
||||
func TestOpenAIDecoderInlineError(t *testing.T) {
|
||||
dec := NewOpenAIDecoder()
|
||||
_, err := dec.Decode([]byte(`{"error":{"type":"rate_limit_exceeded","message":"slow down"}}`))
|
||||
if err == nil {
|
||||
t.Fatal("inline error object must return a decode error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "rate_limit_exceeded") {
|
||||
t.Errorf("error should name the type, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIDecoderMalformedPayload verifies invalid JSON is a returned error
|
||||
// (rides the stream as terminal error), not a panic.
|
||||
func TestOpenAIDecoderMalformedPayload(t *testing.T) {
|
||||
dec := NewOpenAIDecoder()
|
||||
if _, err := dec.Decode([]byte(`{not json`)); err == nil {
|
||||
t.Fatal("malformed payload must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIDecoderFinishFlushesPartial verifies a stream cut short (no
|
||||
// finish_reason) still yields a done event on Finish, defaulting to end_turn.
|
||||
func TestOpenAIDecoderFinishFlushesPartial(t *testing.T) {
|
||||
body := `data: {"id":"c","model":"m","choices":[{"delta":{"content":"partial"}}]}
|
||||
|
||||
`
|
||||
dec := NewOpenAIDecoder()
|
||||
events, final := feedSSE(t, dec, body)
|
||||
if events[len(events)-1].EventKind() != StreamEventDone {
|
||||
t.Fatalf("Finish must emit a terminal done event, got %v", eventKinds(events))
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("cut-short stream should default to end_turn, got %q", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 1 {
|
||||
t.Fatalf("expected the partial text block, got %+v", final.Content)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIDecoderThroughTransport wires the decoder through the real transport
|
||||
// pump against a recorded SSE server, exercising the full path including the
|
||||
// [DONE] terminator handling.
|
||||
func TestOpenAIDecoderThroughTransport(t *testing.T) {
|
||||
srv := sseServer(t, openaiToolCallSSE)
|
||||
defer srv.Close()
|
||||
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: NewOpenAIDecoder(),
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamRequest: %v", err)
|
||||
}
|
||||
var kinds []string
|
||||
for ev := range stream.Events() {
|
||||
kinds = append(kinds, ev.EventKind())
|
||||
}
|
||||
final, resErr := stream.Result(context.Background())
|
||||
if resErr != nil {
|
||||
t.Fatalf("result: %v", resErr)
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason via transport = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
if len(final.Content) != 2 {
|
||||
t.Errorf("expected 2 content blocks via transport, got %d", len(final.Content))
|
||||
}
|
||||
if kinds[len(kinds)-1] != StreamEventDone {
|
||||
t.Errorf("stream must end with done, got %v", kinds)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,200 @@
|
||||
// This file defines the built-in preset catalog: a curated set of ready-to-use
|
||||
// (provider, model) pairs a user can pick from without knowing each gateway's
|
||||
// wire details (mirrors pi agent's preset provider/model picker). It covers the
|
||||
// OpenAI-compatible gateways pigo ships — OpenRouter and NVIDIA NIM — plus a few
|
||||
// local Ollama defaults.
|
||||
//
|
||||
// A preset binds a model id to the provider that serves it and that provider's
|
||||
// default endpoint, so selecting a preset is enough to build a working Provider.
|
||||
// The naive prefix-based mapping (ollama/…) still works for arbitrary ids; the
|
||||
// preset catalog is the "menu" of vetted choices surfaced to the user.
|
||||
//
|
||||
// Security: presets carry no secrets. Each provider resolves its API key by name
|
||||
// from the environment at request time (see auth.go); keys are never embedded
|
||||
// here or logged.
|
||||
package provider
|
||||
|
||||
// PresetModel is one entry in the preset catalog: a model the user can select by
|
||||
// id, the provider that serves it, and a short human label for the picker.
|
||||
type PresetModel struct {
|
||||
// Provider is the owning provider name (e.g. "openrouter", "nvidia").
|
||||
Provider string
|
||||
// ID is the model id passed to the provider (e.g. "openai/gpt-4o").
|
||||
ID string
|
||||
// DisplayName is a friendly label shown in the picker; falls back to ID.
|
||||
DisplayName string
|
||||
}
|
||||
|
||||
// Label returns the display label for a preset, falling back to the id.
|
||||
func (p PresetModel) Label() string {
|
||||
if p.DisplayName != "" {
|
||||
return p.DisplayName
|
||||
}
|
||||
return p.ID
|
||||
}
|
||||
|
||||
// PresetProviders lists the providers the preset catalog draws from, with the
|
||||
// environment variable each expects its API key in (referenced by name only,
|
||||
// never a value). Order is the display order in the picker.
|
||||
var PresetProviders = []struct {
|
||||
Name string
|
||||
EnvVar string
|
||||
}{
|
||||
{Name: "openrouter", EnvVar: "OPENROUTER_API_KEY"},
|
||||
{Name: "nvidia", EnvVar: "NVIDIA_API_KEY"},
|
||||
{Name: "deepseek", EnvVar: "DEEPSEEK_API_KEY"},
|
||||
{Name: "groq", EnvVar: "GROQ_API_KEY"},
|
||||
{Name: "xai", EnvVar: "XAI_API_KEY"},
|
||||
{Name: "cerebras", EnvVar: "CEREBRAS_API_KEY"},
|
||||
{Name: "mistral", EnvVar: "MISTRAL_API_KEY"},
|
||||
{Name: "moonshotai", EnvVar: "MOONSHOT_API_KEY"},
|
||||
{Name: "zai", EnvVar: "ZAI_API_KEY"},
|
||||
{Name: "fireworks", EnvVar: "FIREWORKS_API_KEY"},
|
||||
{Name: "together", EnvVar: "TOGETHER_API_KEY"},
|
||||
{Name: "minimax", EnvVar: "MINIMAX_API_KEY"},
|
||||
{Name: "xiaomi", EnvVar: "XIAOMI_API_KEY"},
|
||||
{Name: "qianfan", EnvVar: "QIANFAN_API_KEY"},
|
||||
{Name: "volcengine", EnvVar: "ARK_API_KEY"},
|
||||
{Name: "dashscope", EnvVar: "DASHSCOPE_API_KEY"},
|
||||
{Name: "hunyuan", EnvVar: "HUNYUAN_API_KEY"},
|
||||
{Name: "ollama", EnvVar: ""}, // local, no key
|
||||
}
|
||||
|
||||
// PresetCatalog is the built-in curated list of selectable models, grouped by
|
||||
// provider in the order PresetProviders declares. These ids are the ones a user
|
||||
// can `/model <id>` into or pick from `/models`; the list is representative, not
|
||||
// exhaustive — any valid id for a known provider still works.
|
||||
var PresetCatalog = []PresetModel{
|
||||
// --- OpenRouter (routes to many upstreams via one OpenAI-compatible API) ---
|
||||
{Provider: "openrouter", ID: "openai/gpt-4o", DisplayName: "GPT-4o (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "openai/gpt-4o-mini", DisplayName: "GPT-4o mini (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "anthropic/claude-3.5-sonnet", DisplayName: "Claude 3.5 Sonnet (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "anthropic/claude-3.7-sonnet", DisplayName: "Claude 3.7 Sonnet (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "google/gemini-2.0-flash-001", DisplayName: "Gemini 2.0 Flash (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "google/gemini-2.5-pro", DisplayName: "Gemini 2.5 Pro (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "meta-llama/llama-3.3-70b-instruct", DisplayName: "Llama 3.3 70B (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "deepseek/deepseek-chat", DisplayName: "DeepSeek V3 (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "deepseek/deepseek-r1", DisplayName: "DeepSeek R1 (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "qwen/qwen-2.5-72b-instruct", DisplayName: "Qwen 2.5 72B (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "mistralai/mistral-large", DisplayName: "Mistral Large (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "x-ai/grok-2-1212", DisplayName: "Grok 2 (OpenRouter)"},
|
||||
|
||||
// --- OpenRouter free tier (":free" ids are rate-limited but cost nothing) ---
|
||||
{Provider: "openrouter", ID: "deepseek/deepseek-r1:free", DisplayName: "DeepSeek R1 · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "deepseek/deepseek-chat-v3-0324:free", DisplayName: "DeepSeek V3 · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "meta-llama/llama-3.3-70b-instruct:free", DisplayName: "Llama 3.3 70B · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "google/gemini-2.0-flash-exp:free", DisplayName: "Gemini 2.0 Flash · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "qwen/qwen-2.5-72b-instruct:free", DisplayName: "Qwen 2.5 72B · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "qwen/qwq-32b:free", DisplayName: "QwQ 32B · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "mistralai/mistral-small-3.1-24b-instruct:free", DisplayName: "Mistral Small 3.1 24B · free (OpenRouter)"},
|
||||
{Provider: "openrouter", ID: "meta-llama/llama-4-maverick:free", DisplayName: "Llama 4 Maverick · free (OpenRouter)"},
|
||||
|
||||
// --- NVIDIA NIM (hosted, OpenAI-compatible) ---
|
||||
{Provider: "nvidia", ID: "meta/llama-3.3-70b-instruct", DisplayName: "Llama 3.3 70B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "meta/llama-3.1-405b-instruct", DisplayName: "Llama 3.1 405B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "deepseek-ai/deepseek-r1", DisplayName: "DeepSeek R1 (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "qwen/qwen2.5-coder-32b-instruct", DisplayName: "Qwen 2.5 Coder 32B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "nvidia/llama-3.1-nemotron-70b-instruct", DisplayName: "Nemotron 70B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "mistralai/mixtral-8x22b-instruct-v0.1", DisplayName: "Mixtral 8x22B (NVIDIA)"},
|
||||
// NVIDIA's hosted NIM endpoint is free to call with a build.nvidia.com key.
|
||||
{Provider: "nvidia", ID: "meta/llama-3.1-8b-instruct", DisplayName: "Llama 3.1 8B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "meta/llama-3.1-70b-instruct", DisplayName: "Llama 3.1 70B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "deepseek-ai/deepseek-v3", DisplayName: "DeepSeek V3 (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "qwen/qwen2.5-7b-instruct", DisplayName: "Qwen 2.5 7B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "google/gemma-2-9b-it", DisplayName: "Gemma 2 9B (NVIDIA)"},
|
||||
{Provider: "nvidia", ID: "microsoft/phi-3.5-mini-instruct", DisplayName: "Phi-3.5 Mini (NVIDIA)"},
|
||||
|
||||
// --- DeepSeek (direct, OpenAI-compatible; ids from pi deepseek.models.ts) ---
|
||||
{Provider: "deepseek", ID: "deepseek-v4-flash", DisplayName: "DeepSeek V4 Flash"},
|
||||
{Provider: "deepseek", ID: "deepseek-v4-pro", DisplayName: "DeepSeek V4 Pro"},
|
||||
|
||||
// --- Groq (fast inference; ids from pi groq.models.ts) ---
|
||||
{Provider: "groq", ID: "llama-3.3-70b-versatile", DisplayName: "Llama 3.3 70B (Groq)"},
|
||||
{Provider: "groq", ID: "openai/gpt-oss-120b", DisplayName: "GPT OSS 120B (Groq)"},
|
||||
{Provider: "groq", ID: "qwen/qwen3-32b", DisplayName: "Qwen3 32B (Groq)"},
|
||||
|
||||
// --- xAI Grok (ids from pi xai.models.ts) ---
|
||||
{Provider: "xai", ID: "grok-4.5", DisplayName: "Grok 4.5"},
|
||||
{Provider: "xai", ID: "grok-4.3", DisplayName: "Grok 4.3"},
|
||||
|
||||
// --- Cerebras (fast inference; ids from pi cerebras.models.ts) ---
|
||||
{Provider: "cerebras", ID: "gpt-oss-120b", DisplayName: "GPT OSS 120B (Cerebras)"},
|
||||
{Provider: "cerebras", ID: "zai-glm-4.7", DisplayName: "Z.AI GLM-4.7 (Cerebras)"},
|
||||
{Provider: "cerebras", ID: "gemma-4-31b", DisplayName: "Gemma 4 31B (Cerebras)"},
|
||||
|
||||
// --- Mistral (ids from pi mistral.models.ts) ---
|
||||
{Provider: "mistral", ID: "mistral-large-latest", DisplayName: "Mistral Large (latest)"},
|
||||
{Provider: "mistral", ID: "mistral-medium-latest", DisplayName: "Mistral Medium (latest)"},
|
||||
{Provider: "mistral", ID: "codestral-latest", DisplayName: "Codestral (latest)"},
|
||||
{Provider: "mistral", ID: "devstral-medium-latest", DisplayName: "Devstral Medium (latest)"},
|
||||
|
||||
// --- Moonshot AI Kimi (ids from pi moonshotai.models.ts) ---
|
||||
{Provider: "moonshotai", ID: "kimi-k2-thinking", DisplayName: "Kimi K2 Thinking"},
|
||||
{Provider: "moonshotai", ID: "kimi-k2.6", DisplayName: "Kimi K2.6"},
|
||||
{Provider: "moonshotai", ID: "kimi-k3", DisplayName: "Kimi K3"},
|
||||
|
||||
// --- Z.AI GLM (ids from pi zai.models.ts) ---
|
||||
{Provider: "zai", ID: "glm-4.7", DisplayName: "GLM-4.7"},
|
||||
{Provider: "zai", ID: "glm-5.1", DisplayName: "GLM-5.1"},
|
||||
{Provider: "zai", ID: "glm-5.2", DisplayName: "GLM-5.2"},
|
||||
|
||||
// --- Fireworks (ids from pi fireworks.models.ts) ---
|
||||
{Provider: "fireworks", ID: "accounts/fireworks/models/deepseek-v4-pro", DisplayName: "DeepSeek V4 Pro (Fireworks)"},
|
||||
{Provider: "fireworks", ID: "accounts/fireworks/models/gpt-oss-120b", DisplayName: "GPT OSS 120B (Fireworks)"},
|
||||
{Provider: "fireworks", ID: "accounts/fireworks/models/kimi-k2p7-code", DisplayName: "Kimi K2.7 Code (Fireworks)"},
|
||||
|
||||
// --- Together AI (ids from pi together.models.ts) ---
|
||||
{Provider: "together", ID: "deepseek-ai/DeepSeek-V4-Pro", DisplayName: "DeepSeek V4 Pro (Together)"},
|
||||
{Provider: "together", ID: "Qwen/Qwen3.7-Max", DisplayName: "Qwen3.7 Max (Together)"},
|
||||
{Provider: "together", ID: "meta-llama/Llama-3.3-70B-Instruct-Turbo", DisplayName: "Llama 3.3 70B Turbo (Together)"},
|
||||
|
||||
// --- MiniMax (Anthropic-protocol; ids from pi minimax.models.ts) ---
|
||||
{Provider: "minimax", ID: "MiniMax-M2.7", DisplayName: "MiniMax-M2.7"},
|
||||
{Provider: "minimax", ID: "MiniMax-M3", DisplayName: "MiniMax-M3"},
|
||||
|
||||
// --- Xiaomi MiMo (ids from pi xiaomi.models.ts) ---
|
||||
{Provider: "xiaomi", ID: "mimo-v2-pro", DisplayName: "MiMo-V2-Pro"},
|
||||
{Provider: "xiaomi", ID: "mimo-v2.5", DisplayName: "MiMo-V2.5"},
|
||||
{Provider: "xiaomi", ID: "mimo-v2.5-pro", DisplayName: "MiMo-V2.5-Pro"},
|
||||
|
||||
// --- Baidu AI Cloud Qianfan (OpenAI-compatible; ERNIE family) ---
|
||||
{Provider: "qianfan", ID: "ernie-4.5-turbo-32k", DisplayName: "ERNIE 4.5 Turbo (Baidu Qianfan)"},
|
||||
|
||||
// --- ByteDance Volcengine Ark (OpenAI-compatible; Doubao family) ---
|
||||
// Some Ark models require an "inference endpoint ID (endpoint id)" instead of a model
|
||||
// name — use --base-url / -m to target those; this preset uses a model name.
|
||||
{Provider: "volcengine", ID: "doubao-seed-1-6", DisplayName: "Doubao Seed 1.6 (Volcengine Ark)"},
|
||||
|
||||
// --- Alibaba Cloud DashScope (OpenAI-compatible; Qwen family) ---
|
||||
{Provider: "dashscope", ID: "qwen-max", DisplayName: "Qwen Max (Alibaba DashScope)"},
|
||||
|
||||
// --- Tencent Hunyuan (OpenAI-compatible) ---
|
||||
{Provider: "hunyuan", ID: "hunyuan-turbos-latest", DisplayName: "Hunyuan TurboS (Tencent Hunyuan)"},
|
||||
|
||||
// --- Ollama (local, no API key) ---
|
||||
{Provider: "ollama", ID: "ollama/llama3.3", DisplayName: "Llama 3.3 (local Ollama)"},
|
||||
{Provider: "ollama", ID: "ollama/qwen2.5-coder", DisplayName: "Qwen 2.5 Coder (local Ollama)"},
|
||||
}
|
||||
|
||||
// LookupPreset returns the preset entry for a model id, if the id is in the
|
||||
// catalog. Used to resolve a selected id to its owning provider.
|
||||
func LookupPreset(id string) (PresetModel, bool) {
|
||||
for _, p := range PresetCatalog {
|
||||
if p.ID == id {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return PresetModel{}, false
|
||||
}
|
||||
|
||||
// PresetsByProvider returns the presets served by a given provider name, in
|
||||
// catalog order.
|
||||
func PresetsByProvider(providerName string) []PresetModel {
|
||||
var out []PresetModel
|
||||
for _, p := range PresetCatalog {
|
||||
if p.Provider == providerName {
|
||||
out = append(out, p)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
package provider
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestPresetProvidersIncludeNewProviders verifies the curated preset provider
|
||||
// list gained the expanded set of gateways, each paired with the correct API-key
|
||||
// environment variable (referenced by name only).
|
||||
func TestPresetProvidersIncludeNewProviders(t *testing.T) {
|
||||
byName := make(map[string]string, len(PresetProviders))
|
||||
for _, p := range PresetProviders {
|
||||
byName[p.Name] = p.EnvVar
|
||||
}
|
||||
|
||||
want := map[string]string{
|
||||
"deepseek": "DEEPSEEK_API_KEY",
|
||||
"groq": "GROQ_API_KEY",
|
||||
"xai": "XAI_API_KEY",
|
||||
"cerebras": "CEREBRAS_API_KEY",
|
||||
"mistral": "MISTRAL_API_KEY",
|
||||
"moonshotai": "MOONSHOT_API_KEY",
|
||||
"zai": "ZAI_API_KEY",
|
||||
"fireworks": "FIREWORKS_API_KEY",
|
||||
"together": "TOGETHER_API_KEY",
|
||||
"minimax": "MINIMAX_API_KEY",
|
||||
"xiaomi": "XIAOMI_API_KEY",
|
||||
}
|
||||
for name, env := range want {
|
||||
got, ok := byName[name]
|
||||
if !ok {
|
||||
t.Errorf("PresetProviders missing provider %q", name)
|
||||
continue
|
||||
}
|
||||
if got != env {
|
||||
t.Errorf("provider %q env var = %q, want %q", name, got, env)
|
||||
}
|
||||
}
|
||||
|
||||
// Every preset provider must be a known provider in the central registry, so
|
||||
// selecting a preset can always be resolved to a working Provider.
|
||||
for _, p := range PresetProviders {
|
||||
if p.Name == "ollama" {
|
||||
continue // local pseudo-provider, not in the registry
|
||||
}
|
||||
if _, ok := LookupProviderSpec(p.Name); !ok {
|
||||
t.Errorf("preset provider %q has no ProviderSpec in the registry", p.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPresetCatalogCountsPerProvider asserts each expanded provider contributes
|
||||
// the expected number of curated entries.
|
||||
func TestPresetCatalogCountsPerProvider(t *testing.T) {
|
||||
wantAtLeast := map[string]int{
|
||||
"deepseek": 2,
|
||||
"groq": 3,
|
||||
"xai": 2,
|
||||
"cerebras": 3,
|
||||
"mistral": 4,
|
||||
"moonshotai": 3,
|
||||
"zai": 3,
|
||||
"fireworks": 3,
|
||||
"together": 3,
|
||||
"minimax": 2,
|
||||
"xiaomi": 3,
|
||||
}
|
||||
for provider, min := range wantAtLeast {
|
||||
got := PresetsByProvider(provider)
|
||||
if len(got) < min {
|
||||
t.Errorf("PresetsByProvider(%q) returned %d entries, want >= %d", provider, len(got), min)
|
||||
}
|
||||
for _, p := range got {
|
||||
if p.Provider != provider {
|
||||
t.Errorf("PresetsByProvider(%q) returned entry for %q", provider, p.Provider)
|
||||
}
|
||||
if p.ID == "" {
|
||||
t.Errorf("PresetsByProvider(%q) returned entry with empty ID", provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestLookupPresetNewEntries verifies representative new model ids resolve to the
|
||||
// correct owning provider and carry a non-empty label.
|
||||
func TestLookupPresetNewEntries(t *testing.T) {
|
||||
cases := []struct {
|
||||
id string
|
||||
provider string
|
||||
}{
|
||||
{"deepseek-v4-pro", "deepseek"},
|
||||
{"llama-3.3-70b-versatile", "groq"},
|
||||
{"grok-4.5", "xai"},
|
||||
{"zai-glm-4.7", "cerebras"},
|
||||
{"mistral-large-latest", "mistral"},
|
||||
{"kimi-k2-thinking", "moonshotai"},
|
||||
{"glm-5.1", "zai"},
|
||||
{"accounts/fireworks/models/gpt-oss-120b", "fireworks"},
|
||||
{"deepseek-ai/DeepSeek-V4-Pro", "together"},
|
||||
{"MiniMax-M3", "minimax"},
|
||||
{"mimo-v2.5-pro", "xiaomi"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
p, ok := LookupPreset(tc.id)
|
||||
if !ok {
|
||||
t.Errorf("LookupPreset(%q) not found", tc.id)
|
||||
continue
|
||||
}
|
||||
if p.Provider != tc.provider {
|
||||
t.Errorf("LookupPreset(%q).Provider = %q, want %q", tc.id, p.Provider, tc.provider)
|
||||
}
|
||||
if p.Label() == "" {
|
||||
t.Errorf("LookupPreset(%q).Label() is empty", tc.id)
|
||||
}
|
||||
}
|
||||
|
||||
if _, ok := LookupPreset("this-model-does-not-exist"); ok {
|
||||
t.Error("LookupPreset returned ok for an unknown id")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
package provider
|
||||
|
||||
// Tests for the preset provider/model catalog (mirrors pi agent's preset picker):
|
||||
// LookupPreset resolves a catalog id to its owning provider, PresetsByProvider
|
||||
// groups by provider, and every preset must name a provider that has a known
|
||||
// credential env var (or be the local, keyless Ollama).
|
||||
|
||||
import "testing"
|
||||
|
||||
// TestLookupPresetResolvesProvider verifies a catalog id resolves to its
|
||||
// declared provider, and an unknown id does not.
|
||||
func TestLookupPresetResolvesProvider(t *testing.T) {
|
||||
p, ok := LookupPreset("meta/llama-3.3-70b-instruct")
|
||||
if !ok {
|
||||
t.Fatal("expected NVIDIA llama preset to be in the catalog")
|
||||
}
|
||||
if p.Provider != "nvidia" {
|
||||
t.Errorf("provider = %q, want nvidia", p.Provider)
|
||||
}
|
||||
if _, ok := LookupPreset("definitely/not-a-preset"); ok {
|
||||
t.Error("unknown id must not resolve to a preset")
|
||||
}
|
||||
}
|
||||
|
||||
// TestPresetsByProviderGroups verifies each provider surfaces at least one
|
||||
// preset and that the returned entries all belong to that provider.
|
||||
func TestPresetsByProviderGroups(t *testing.T) {
|
||||
for _, name := range []string{"openrouter", "nvidia", "ollama"} {
|
||||
got := PresetsByProvider(name)
|
||||
if len(got) == 0 {
|
||||
t.Errorf("provider %q has no presets", name)
|
||||
}
|
||||
for _, m := range got {
|
||||
if m.Provider != name {
|
||||
t.Errorf("PresetsByProvider(%q) returned entry for %q", name, m.Provider)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPresetProvidersHaveCredentialMapping verifies every non-local preset
|
||||
// provider has API-key env vars in the provider registry (the single source of
|
||||
// truth), so a selected preset can actually resolve a credential. Ollama is
|
||||
// local and keyless.
|
||||
func TestPresetProvidersHaveCredentialMapping(t *testing.T) {
|
||||
for _, pv := range PresetProviders {
|
||||
if pv.Name == "ollama" {
|
||||
continue
|
||||
}
|
||||
spec, ok := LookupProviderSpec(pv.Name)
|
||||
if !ok || len(spec.EnvVars) == 0 {
|
||||
t.Errorf("preset provider %q has no credential env var mapping", pv.Name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestPresetLabelFallsBackToID verifies Label uses DisplayName when set and
|
||||
// falls back to the id otherwise.
|
||||
func TestPresetLabelFallsBackToID(t *testing.T) {
|
||||
if got := (PresetModel{ID: "x", DisplayName: "X"}).Label(); got != "X" {
|
||||
t.Errorf("Label = %q, want X", got)
|
||||
}
|
||||
if got := (PresetModel{ID: "x"}).Label(); got != "x" {
|
||||
t.Errorf("Label = %q, want x", got)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,77 @@
|
||||
package provider
|
||||
|
||||
// Protocol normalization (US-001, #538). The user-facing --protocol / protocol
|
||||
// value accepts three OpenAI wire variants in addition to anthropic:
|
||||
//
|
||||
// openai → Chat Completions (POST {base_url}/chat/completions)
|
||||
// openai/chat → Chat Completions (alias of "openai")
|
||||
// openai/resp_api → Responses API (POST {base_url}/responses)
|
||||
// anthropic → Anthropic Messages
|
||||
// "" → unset; downstream falls back to model-id heuristics
|
||||
//
|
||||
// NormalizeProtocol collapses these into a small set of canonical internal
|
||||
// selectors so ResolveProvider (#543) can switch on chat vs resp_api without
|
||||
// re-parsing surface syntax. "openai" and "openai/chat" both normalize to
|
||||
// ProtocolOpenAI, keeping the existing Chat Completions path byte-for-byte
|
||||
// unchanged; only "openai/resp_api" produces the new selector.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ProtocolOpenAIResponses is the canonical selector for the OpenAI Responses
|
||||
// API wire format (POST {base_url}/responses). It is distinct from
|
||||
// ProtocolOpenAI (Chat Completions) so ResolveProvider can route to the
|
||||
// SDK-based Responses driver.
|
||||
const ProtocolOpenAIResponses = "openai/resp_api"
|
||||
|
||||
// NormalizeProtocol maps a raw --protocol / protocol value to a canonical
|
||||
// internal selector. Input is trimmed and lower-cased before matching. An empty
|
||||
// value stays empty (unset → model-id heuristics). Recognized values normalize
|
||||
// to ProtocolOpenAI, ProtocolOpenAIResponses, or ProtocolAnthropic. Any other
|
||||
// value is an error naming the accepted set, so a typo surfaces to the caller
|
||||
// for exit-code mapping instead of silently falling through.
|
||||
func NormalizeProtocol(raw string) (string, error) {
|
||||
switch strings.ToLower(strings.TrimSpace(raw)) {
|
||||
case "":
|
||||
return "", nil
|
||||
case ProtocolOpenAI, "openai/chat":
|
||||
return ProtocolOpenAI, nil
|
||||
case ProtocolOpenAIResponses:
|
||||
return ProtocolOpenAIResponses, nil
|
||||
case ProtocolAnthropic:
|
||||
return ProtocolAnthropic, nil
|
||||
default:
|
||||
return "", fmt.Errorf("unknown --protocol %q (want openai|openai/chat|openai/resp_api|anthropic)", raw)
|
||||
}
|
||||
}
|
||||
|
||||
// ProtocolLabel maps a raw --protocol value to the human-facing label shown in
|
||||
// the startup banner's Protocol row, so the displayed wire format matches what
|
||||
// pigo actually speaks. It differs from NormalizeProtocol in one deliberate way:
|
||||
// the bare "openai" input is surfaced as "openai/chat", making the Chat
|
||||
// Completions variant explicit rather than ambiguous. "openai/resp_api" and
|
||||
// "anthropic" pass through as themselves.
|
||||
//
|
||||
// An empty input returns empty (the banner then falls back to "—" or the
|
||||
// provider name, so an unset protocol on a named/inferred provider is not
|
||||
// mislabeled). An unrecognized value returns the trimmed input verbatim — the
|
||||
// label is presentation-only and must never fail; a real typo is already
|
||||
// rejected upstream by NormalizeProtocol during resolution.
|
||||
func ProtocolLabel(raw string) string {
|
||||
canonical, err := NormalizeProtocol(raw)
|
||||
if err != nil {
|
||||
return strings.TrimSpace(raw)
|
||||
}
|
||||
switch canonical {
|
||||
case ProtocolOpenAI:
|
||||
return "openai/chat"
|
||||
case ProtocolOpenAIResponses:
|
||||
return ProtocolOpenAIResponses
|
||||
case ProtocolAnthropic:
|
||||
return ProtocolAnthropic
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,55 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestNormalizeProtocol(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
in string
|
||||
want string
|
||||
wantErr bool
|
||||
}{
|
||||
{"empty stays empty", "", "", false},
|
||||
{"openai", "openai", ProtocolOpenAI, false},
|
||||
{"openai/chat aliases openai", "openai/chat", ProtocolOpenAI, false},
|
||||
{"openai/resp_api distinct", "openai/resp_api", ProtocolOpenAIResponses, false},
|
||||
{"anthropic unchanged", "anthropic", ProtocolAnthropic, false},
|
||||
{"case-insensitive", "OpenAI/Resp_API", ProtocolOpenAIResponses, false},
|
||||
{"trimmed", " openai/chat ", ProtocolOpenAI, false},
|
||||
{"unknown rejected", "openai/foo", "", true},
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
got, err := NormalizeProtocol(tc.in)
|
||||
if tc.wantErr {
|
||||
if err == nil {
|
||||
t.Fatalf("NormalizeProtocol(%q) = %q, want error", tc.in, got)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("NormalizeProtocol(%q) unexpected error: %v", tc.in, err)
|
||||
}
|
||||
if got != tc.want {
|
||||
t.Errorf("NormalizeProtocol(%q) = %q, want %q", tc.in, got, tc.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// The rejection message must name every accepted value so a user with a typo
|
||||
// can self-correct without reading source.
|
||||
func TestNormalizeProtocolErrorNamesAcceptedValues(t *testing.T) {
|
||||
_, err := NormalizeProtocol("openai/foo")
|
||||
if err == nil {
|
||||
t.Fatal("expected error for unknown protocol")
|
||||
}
|
||||
for _, want := range []string{"openai", "openai/chat", "openai/resp_api", "anthropic"} {
|
||||
if !strings.Contains(err.Error(), want) {
|
||||
t.Errorf("error %q missing accepted value %q", err.Error(), want)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,118 @@
|
||||
// This file defines the provider streaming abstraction (US-003/US-007 base):
|
||||
// the StreamFn contract, the per-delta AssistantMessageEvent set, and the
|
||||
// AssistantMessageEventStream (a specialization of EventStream) that a provider
|
||||
// pushes deltas onto while yielding a final AssistantMessage.
|
||||
//
|
||||
// Contract (FR-13): a StreamFn never expresses a request failure by returning
|
||||
// an error. Runtime failures are encoded as an error event plus a terminal
|
||||
// assistant message (stopReason=error/aborted + errorMessage). The returned
|
||||
// error is reserved for the earliest "could not even build the stream" case.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// AssistantMessageEvent is the sealed interface for provider stream deltas. The
|
||||
// loop dispatches on EventKind; the raw event is also surfaced to consumers via
|
||||
// MessageUpdateEvent.AssistantMessageEvent.
|
||||
type AssistantMessageEvent interface {
|
||||
isAssistantMessageEvent()
|
||||
// EventKind returns the delta discriminant.
|
||||
EventKind() string
|
||||
}
|
||||
|
||||
// AssistantMessageEvent kinds.
|
||||
const (
|
||||
StreamEventStart = "start"
|
||||
StreamEventText = "text"
|
||||
StreamEventThinking = "thinking"
|
||||
StreamEventToolCall = "toolcall"
|
||||
StreamEventDone = "done"
|
||||
StreamEventError = "error"
|
||||
)
|
||||
|
||||
// StreamStartEvent carries the initial (usually empty) partial message.
|
||||
type StreamStartEvent struct{ Partial agentcore.AssistantMessage }
|
||||
|
||||
// StreamTextEvent carries the partial message after a text delta.
|
||||
type StreamTextEvent struct{ Partial agentcore.AssistantMessage }
|
||||
|
||||
// StreamThinkingEvent carries the partial after a thinking delta.
|
||||
type StreamThinkingEvent struct{ Partial agentcore.AssistantMessage }
|
||||
|
||||
// StreamToolCallEvent carries the partial after a tool-call delta.
|
||||
type StreamToolCallEvent struct{ Partial agentcore.AssistantMessage }
|
||||
|
||||
// StreamDoneEvent is the terminal success event; Message is the final response.
|
||||
type StreamDoneEvent struct{ Message agentcore.AssistantMessage }
|
||||
|
||||
// StreamErrorEvent is the terminal failure event; Message carries the terminal
|
||||
// assistant message (stopReason=error/aborted + errorMessage).
|
||||
type StreamErrorEvent struct {
|
||||
Message agentcore.AssistantMessage
|
||||
Err error
|
||||
}
|
||||
|
||||
func (StreamStartEvent) isAssistantMessageEvent() {}
|
||||
func (StreamTextEvent) isAssistantMessageEvent() {}
|
||||
func (StreamThinkingEvent) isAssistantMessageEvent() {}
|
||||
func (StreamToolCallEvent) isAssistantMessageEvent() {}
|
||||
func (StreamDoneEvent) isAssistantMessageEvent() {}
|
||||
func (StreamErrorEvent) isAssistantMessageEvent() {}
|
||||
|
||||
func (StreamStartEvent) EventKind() string { return StreamEventStart }
|
||||
func (StreamTextEvent) EventKind() string { return StreamEventText }
|
||||
func (StreamThinkingEvent) EventKind() string { return StreamEventThinking }
|
||||
func (StreamToolCallEvent) EventKind() string { return StreamEventToolCall }
|
||||
func (StreamDoneEvent) EventKind() string { return StreamEventDone }
|
||||
func (StreamErrorEvent) EventKind() string { return StreamEventError }
|
||||
|
||||
// AssistantMessageEventStream is the provider-level stream: deltas of type
|
||||
// AssistantMessageEvent with a final AssistantMessage result. isComplete fires
|
||||
// on done/error; extractResult takes the terminal event's message.
|
||||
type AssistantMessageEventStream = agentcore.EventStream[AssistantMessageEvent, agentcore.AssistantMessage]
|
||||
|
||||
// NewAssistantMessageEventStream builds a provider stream wired with the
|
||||
// done/error completion callbacks.
|
||||
func NewAssistantMessageEventStream(buffer int) *AssistantMessageEventStream {
|
||||
s := agentcore.NewEventStream[AssistantMessageEvent, agentcore.AssistantMessage](buffer)
|
||||
s.IsComplete = func(e AssistantMessageEvent) bool {
|
||||
k := e.EventKind()
|
||||
return k == StreamEventDone || k == StreamEventError
|
||||
}
|
||||
s.ExtractResult = func(e AssistantMessageEvent) agentcore.AssistantMessage {
|
||||
switch ev := e.(type) {
|
||||
case StreamDoneEvent:
|
||||
return ev.Message
|
||||
case StreamErrorEvent:
|
||||
return ev.Message
|
||||
default:
|
||||
return agentcore.AssistantMessage{}
|
||||
}
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// LlmContext is the shaped request handed to a StreamFn: the system prompt, the
|
||||
// LLM-bound messages (UI-only messages already filtered), and the tools.
|
||||
type LlmContext struct {
|
||||
SystemPrompt string
|
||||
Messages agentcore.MessageList
|
||||
Tools []agentcore.AgentTool
|
||||
}
|
||||
|
||||
// StreamConfig carries per-request settings for a StreamFn.
|
||||
type StreamConfig struct {
|
||||
APIKey string
|
||||
ThinkingLevel agentcore.ThinkingLevel
|
||||
// Extra holds provider-specific options; opaque to the loop.
|
||||
Extra map[string]any
|
||||
}
|
||||
|
||||
// StreamFn produces a provider stream for a model + shaped context. Per the
|
||||
// contract it returns an error only for early "cannot build the stream"
|
||||
// failures; all runtime failures ride the returned stream as error events.
|
||||
type StreamFn func(ctx context.Context, model string, llm LlmContext, cfg StreamConfig) (*AssistantMessageEventStream, error)
|
||||
@@ -0,0 +1,79 @@
|
||||
// This file defines the unified Provider interface and its dual failure model
|
||||
// (US-007). A Provider turns a CompletionRequest into a stream of
|
||||
// AssistantMessageEvents. Failures follow the same contract as StreamFn (FR-13):
|
||||
//
|
||||
// - "cannot even build the stream" (bad config, missing model) → returned error.
|
||||
// - any runtime failure once streaming has begun → a terminal StreamErrorEvent
|
||||
// carrying an assistant message with stopReason=error/aborted, after which
|
||||
// the stream is closed. It is never a returned error.
|
||||
//
|
||||
// Model carries provider-agnostic capability metadata so the loop and UI can
|
||||
// reason about a model without knowing the concrete provider.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// Model is provider-agnostic metadata describing a single model's identity and
|
||||
// capabilities. Providers construct these; the loop/UI consume them.
|
||||
type Model struct {
|
||||
// Provider is the provider name (e.g. "anthropic", "openai").
|
||||
Provider string `json:"provider"`
|
||||
// ID is the provider-specific model id (e.g. "claude-opus-4-8").
|
||||
ID string `json:"id"`
|
||||
// DisplayName is a human-friendly label; falls back to ID when empty.
|
||||
DisplayName string `json:"displayName,omitempty"`
|
||||
// ContextWindow is the maximum input+output token window, 0 if unknown.
|
||||
ContextWindow int `json:"contextWindow,omitempty"`
|
||||
// MaxOutputTokens is the max tokens the model may emit per response, 0 if
|
||||
// unknown.
|
||||
MaxOutputTokens int `json:"maxOutputTokens,omitempty"`
|
||||
// SupportsThinking reports whether the model exposes a reasoning/thinking
|
||||
// channel.
|
||||
SupportsThinking bool `json:"supportsThinking,omitempty"`
|
||||
// SupportsTools reports whether the model can call tools.
|
||||
SupportsTools bool `json:"supportsTools,omitempty"`
|
||||
// SupportsImages reports whether the model accepts image (multimodal) input.
|
||||
// When false, an image block in the request is reported as a hard error
|
||||
// rather than silently dropped, so the user learns the model cannot see it.
|
||||
SupportsImages bool `json:"supportsImages,omitempty"`
|
||||
// ThinkingLevels maps unified thinking levels to this model's wire values.
|
||||
// nil when the model does not support thinking (decision #10).
|
||||
ThinkingLevels agentcore.ThinkingLevelMap `json:"-"`
|
||||
}
|
||||
|
||||
// CompletionRequest is the provider-agnostic input to StreamCompletion: the
|
||||
// model id, the shaped LLM context, and per-request options.
|
||||
type CompletionRequest struct {
|
||||
// Model is the provider-specific model id to complete against.
|
||||
Model string
|
||||
// Context is the shaped request (system prompt, LLM-bound messages, tools).
|
||||
Context LlmContext
|
||||
// Config carries per-request options (API key, thinking level, extras).
|
||||
Config StreamConfig
|
||||
}
|
||||
|
||||
// Provider is the unified streaming interface implemented by every backend. It
|
||||
// hides per-vendor differences behind a single AssistantMessageEvent stream.
|
||||
type Provider interface {
|
||||
// Name returns the provider's identifier (matches Model.Provider).
|
||||
Name() string
|
||||
// Models lists the models this provider can serve.
|
||||
Models() []Model
|
||||
// StreamCompletion streams a completion for req. Per the dual failure model
|
||||
// it returns an error only for the earliest "cannot build the stream" case;
|
||||
// all runtime failures ride the returned stream as a terminal error event.
|
||||
StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error)
|
||||
}
|
||||
|
||||
// StreamFnFromProvider adapts a Provider to the loop's StreamFn contract so a
|
||||
// Provider can drive streamAssistantResponse directly. The two failure models
|
||||
// are identical, so the adaptation is a straight delegation.
|
||||
func StreamFnFromProvider(p Provider) StreamFn {
|
||||
return func(ctx context.Context, model string, llm LlmContext, cfg StreamConfig) (*AssistantMessageEventStream, error) {
|
||||
return p.StreamCompletion(ctx, CompletionRequest{Model: model, Context: llm, Config: cfg})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// fakeProvider is a minimal Provider for interface tests.
|
||||
type fakeProvider struct {
|
||||
name string
|
||||
models []Model
|
||||
buildErr error
|
||||
events []AssistantMessageEvent
|
||||
}
|
||||
|
||||
func (p fakeProvider) Name() string { return p.name }
|
||||
func (p fakeProvider) Models() []Model { return p.models }
|
||||
func (p fakeProvider) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) {
|
||||
if p.buildErr != nil {
|
||||
return nil, p.buildErr
|
||||
}
|
||||
s := NewAssistantMessageEventStream(0)
|
||||
go func() {
|
||||
for _, ev := range p.events {
|
||||
if err := s.Emit(ctx, ev); err != nil {
|
||||
s.SetError(err)
|
||||
break
|
||||
}
|
||||
}
|
||||
s.Close()
|
||||
}()
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func TestProviderEarlyBuildFailureReturnsError(t *testing.T) {
|
||||
p := fakeProvider{name: "test", buildErr: errors.New("no such model")}
|
||||
_, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "ghost"})
|
||||
if err == nil {
|
||||
t.Fatal("early build failure must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderRuntimeFailureRidesStream(t *testing.T) {
|
||||
errMsg := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonError, ErrorMessage: "upstream 500"}
|
||||
p := fakeProvider{
|
||||
name: "test",
|
||||
events: []AssistantMessageEvent{StreamErrorEvent{Message: errMsg}},
|
||||
}
|
||||
stream, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"})
|
||||
if err != nil {
|
||||
t.Fatalf("runtime failure must NOT be a returned error: %v", err)
|
||||
}
|
||||
final, resErr := stream.Result(context.Background())
|
||||
if resErr != nil {
|
||||
t.Fatalf("stream result error: %v", resErr)
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonError || final.ErrorMessage != "upstream 500" {
|
||||
t.Errorf("terminal error message wrong: %+v", final)
|
||||
}
|
||||
}
|
||||
|
||||
func TestStreamFnFromProviderDelegates(t *testing.T) {
|
||||
done := agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}
|
||||
p := fakeProvider{name: "test", events: []AssistantMessageEvent{StreamDoneEvent{Message: done}}}
|
||||
fn := StreamFnFromProvider(p)
|
||||
stream, err := fn(context.Background(), "m", LlmContext{}, StreamConfig{})
|
||||
if err != nil {
|
||||
t.Fatalf("delegation error: %v", err)
|
||||
}
|
||||
final, _ := stream.Result(context.Background())
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("delegated stream result wrong: %+v", final)
|
||||
}
|
||||
}
|
||||
|
||||
func TestModelMetadata(t *testing.T) {
|
||||
m := Model{Provider: "anthropic", ID: "claude-opus-4-8", SupportsThinking: true, ContextWindow: 200000}
|
||||
if m.Provider != "anthropic" || m.ID != "claude-opus-4-8" {
|
||||
t.Errorf("model identity wrong: %+v", m)
|
||||
}
|
||||
if !m.SupportsThinking || m.ContextWindow != 200000 {
|
||||
t.Errorf("model capability wrong: %+v", m)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,775 @@
|
||||
// This file implements the first concrete Providers (US-013): Bedrock,
|
||||
// OpenRouter, and Ollama. Each satisfies the Provider interface by building an
|
||||
// *http.Request and delegating to the shared transport (StreamRequest) with a
|
||||
// provider-appropriate Decoder — no bespoke HTTP/SSE handling per provider.
|
||||
//
|
||||
// Two backing driver shapes cover all three:
|
||||
//
|
||||
// - openAICompatDriver — POSTs to {baseURL}/chat/completions in the OpenAI
|
||||
// Chat Completions wire format, decoded by OpenAIDecoder. OpenRouter and
|
||||
// Ollama are instances of it; it is the generic OpenAI-compatible layer,
|
||||
// reusable for any gateway (Groq, together, local servers, …).
|
||||
// - anthropicCompatDriver — POSTs the Anthropic Messages wire format, decoded
|
||||
// by AnthropicDecoder. Bedrock rides this: Anthropic-on-Bedrock speaks the
|
||||
// Messages API, so the decoder is reused wholesale.
|
||||
//
|
||||
// Failures follow the dual failure model (FR-13): only the earliest "cannot
|
||||
// build the stream" case (missing key, bad request construction) is a returned
|
||||
// error; every runtime failure rides the stream as a terminal error event,
|
||||
// which StreamRequest already guarantees.
|
||||
//
|
||||
// Security (US-012 / US-026): API keys are referenced by provider name in any
|
||||
// error; secret values are never logged or embedded in error text.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// providerBaseURLs holds the default endpoint per built-in provider. A base URL
|
||||
// is a transport concern (the decoders are URL-agnostic), so overriding it —
|
||||
// e.g. pointing Ollama at a remote host — needs no decoder change.
|
||||
const (
|
||||
openRouterBaseURL = "https://openrouter.ai/api/v1"
|
||||
ollamaBaseURL = "http://localhost:11434/v1"
|
||||
// nvidiaBaseURL is NVIDIA's hosted NIM endpoint. It speaks the OpenAI
|
||||
// Chat Completions wire format, so it rides openAICompatDriver unchanged.
|
||||
nvidiaBaseURL = "https://integrate.api.nvidia.com/v1"
|
||||
// bedrockBaseURL is a placeholder default; real Bedrock endpoints are
|
||||
// region-specific (bedrock-runtime.<region>.amazonaws.com) and supplied at
|
||||
// construction. It is exported as a field so callers set the resolved URL.
|
||||
bedrockBaseURL = "https://bedrock-runtime.us-east-1.amazonaws.com"
|
||||
// anthropicBaseURL is the public Anthropic Messages API endpoint, the default
|
||||
// for a --protocol=anthropic provider when no --base-url is given.
|
||||
anthropicBaseURL = "https://api.anthropic.com/v1"
|
||||
// anthropicAPIVersion is the required anthropic-version header value sent with
|
||||
// every direct-Anthropic request.
|
||||
anthropicAPIVersion = "2023-06-01"
|
||||
)
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// OpenAI-compatible driver (OpenRouter, Ollama, and any OpenAI-compatible API).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// openAICompatDriver is the shared backing for every OpenAI-compatible
|
||||
// provider. It holds the provider identity, endpoint, model catalog, and the
|
||||
// auth scheme; StreamCompletion builds the chat-completions request and hands
|
||||
// it to the transport with a fresh OpenAIDecoder.
|
||||
type openAICompatDriver struct {
|
||||
name string
|
||||
baseURL string
|
||||
models []Model
|
||||
// requiresAuth reports whether an Authorization: Bearer header is sent.
|
||||
// Ollama (local) needs none; OpenRouter does.
|
||||
requiresAuth bool
|
||||
// extraHeaders are attached to every request (e.g. OpenRouter attribution).
|
||||
extraHeaders map[string]string
|
||||
}
|
||||
|
||||
func (d *openAICompatDriver) Name() string { return d.name }
|
||||
func (d *openAICompatDriver) Models() []Model { return d.models }
|
||||
|
||||
// StreamCompletion builds the OpenAI Chat Completions request and streams it.
|
||||
func (d *openAICompatDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) {
|
||||
if d.requiresAuth && strings.TrimSpace(req.Config.APIKey) == "" {
|
||||
// Early "cannot build the stream": reference the provider, never a value.
|
||||
return nil, fmt.Errorf("%s: missing API key", d.name)
|
||||
}
|
||||
if err := checkImageSupport(d.name, req.Model, d.models, req.Context.Messages); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, err := encodeOpenAIRequest(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: build request body: %w", d.name, err)
|
||||
}
|
||||
newReq := func(ctx context.Context) (*http.Request, error) {
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, d.baseURL+"/chat/completions", bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
if d.requiresAuth {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+req.Config.APIKey)
|
||||
}
|
||||
for k, v := range d.extraHeaders {
|
||||
httpReq.Header.Set(k, v)
|
||||
}
|
||||
return httpReq, nil
|
||||
}
|
||||
return StreamRequest(ctx, TransportConfig{NewRequest: newReq, Decoder: NewOpenAIDecoder()})
|
||||
}
|
||||
|
||||
// encodeOpenAIRequest serializes a CompletionRequest into an OpenAI Chat
|
||||
// Completions JSON body with streaming enabled and usage requested.
|
||||
func encodeOpenAIRequest(req CompletionRequest) ([]byte, error) {
|
||||
msgs := make([]map[string]any, 0, len(req.Context.Messages)+1)
|
||||
if sp := req.Context.SystemPrompt; sp != "" {
|
||||
msgs = append(msgs, map[string]any{"role": "system", "content": sp})
|
||||
}
|
||||
for _, m := range req.Context.Messages {
|
||||
msgs = append(msgs, encodeOpenAIMessage(m)...)
|
||||
}
|
||||
body := map[string]any{
|
||||
"model": req.Model,
|
||||
"messages": msgs,
|
||||
"stream": true,
|
||||
"stream_options": map[string]any{"include_usage": true},
|
||||
}
|
||||
// Reasoning effort: when a thinking level is requested, forward it as the
|
||||
// OpenAI `reasoning_effort` field. Reasoning models (o-series, DeepSeek-R1,
|
||||
// GLM-thinking, …) read this to open their reasoning channel; omitting it
|
||||
// leaves them at their default and effectively disables extended reasoning.
|
||||
if effort := openAIReasoningEffort(req.Config.ThinkingLevel); effort != "" {
|
||||
body["reasoning_effort"] = effort
|
||||
}
|
||||
if tools := encodeOpenAITools(req.Context.Tools); len(tools) > 0 {
|
||||
body["tools"] = tools
|
||||
}
|
||||
return json.Marshal(body)
|
||||
}
|
||||
|
||||
// openAIReasoningEffort maps the unified ThinkingLevel onto the OpenAI
|
||||
// `reasoning_effort` wire value. "off"/"" yields "" (field omitted, default
|
||||
// behavior preserved). OpenAI accepts minimal|low|medium|high; xhigh maps to
|
||||
// high (the strongest supported value).
|
||||
func openAIReasoningEffort(level agentcore.ThinkingLevel) string {
|
||||
switch level {
|
||||
case agentcore.ThinkingMinimal:
|
||||
return "minimal"
|
||||
case agentcore.ThinkingLow:
|
||||
return "low"
|
||||
case agentcore.ThinkingMedium:
|
||||
return "medium"
|
||||
case agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax:
|
||||
return "high"
|
||||
default: // off or unset
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// encodeOpenAIMessage maps one pigo message onto the OpenAI wire shape. An
|
||||
// assistant message may expand to content + tool_calls in a single entry; a
|
||||
// tool result becomes a role:"tool" entry keyed by tool_call_id.
|
||||
func encodeOpenAIMessage(m agentcore.Message) []map[string]any {
|
||||
switch msg := m.(type) {
|
||||
case agentcore.UserMessage:
|
||||
return []map[string]any{{"role": "user", "content": openAIUserContent(msg.Content)}}
|
||||
case agentcore.CompactionMessage:
|
||||
// A compaction checkpoint stands in for compacted history as user text.
|
||||
u := msg.AsUserMessage()
|
||||
return []map[string]any{{"role": "user", "content": openAIUserContent(u.Content)}}
|
||||
case agentcore.AssistantMessage:
|
||||
entry := map[string]any{"role": "assistant"}
|
||||
text := agentcore.ContentToText(msg.Content)
|
||||
var toolCalls []map[string]any
|
||||
for _, c := range msg.Content {
|
||||
if tc, ok := c.(agentcore.ToolCallContent); ok {
|
||||
toolCalls = append(toolCalls, map[string]any{
|
||||
"id": tc.ID,
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": tc.Name,
|
||||
"arguments": string(tc.Arguments),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
if len(toolCalls) > 0 {
|
||||
entry["tool_calls"] = toolCalls
|
||||
// With tool calls present, send content as JSON null when there is no
|
||||
// accompanying text: an empty string trips strict gateways (e.g. vLLM)
|
||||
// that expect null | non-empty for an assistant tool-call turn.
|
||||
if text == "" {
|
||||
entry["content"] = nil
|
||||
} else {
|
||||
entry["content"] = text
|
||||
}
|
||||
} else {
|
||||
entry["content"] = text
|
||||
}
|
||||
return []map[string]any{entry}
|
||||
case agentcore.ToolResultMessage:
|
||||
return []map[string]any{{
|
||||
"role": "tool",
|
||||
"tool_call_id": msg.ToolCallID,
|
||||
"content": agentcore.ContentToText(msg.Content),
|
||||
}}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// openAIUserContent shapes a user content list for the OpenAI wire. When there
|
||||
// are no images it collapses to a plain string (the common case, and what most
|
||||
// OpenAI-compatible gateways expect). When images are present it emits the
|
||||
// multimodal array form: text parts plus image_url parts carrying a base64 data
|
||||
// URI (data:<mime>;base64,<data>).
|
||||
func openAIUserContent(content agentcore.ContentList) any {
|
||||
hasImage := false
|
||||
for _, c := range content {
|
||||
if _, ok := c.(agentcore.ImageContent); ok {
|
||||
hasImage = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasImage {
|
||||
return agentcore.ContentToText(content)
|
||||
}
|
||||
parts := make([]map[string]any, 0, len(content))
|
||||
for _, c := range content {
|
||||
switch b := c.(type) {
|
||||
case agentcore.TextContent:
|
||||
if b.Text == "" {
|
||||
continue
|
||||
}
|
||||
parts = append(parts, map[string]any{"type": "text", "text": b.Text})
|
||||
case agentcore.ImageContent:
|
||||
parts = append(parts, map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]any{
|
||||
"url": fmt.Sprintf("data:%s;base64,%s", b.MimeType, b.Data),
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
return parts
|
||||
}
|
||||
|
||||
// encodeOpenAITools maps AgentTools onto the OpenAI function-tool schema.
|
||||
func encodeOpenAITools(tools []agentcore.AgentTool) []map[string]any {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]map[string]any, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
params := json.RawMessage(t.Schema())
|
||||
if len(params) == 0 {
|
||||
params = json.RawMessage("{}")
|
||||
}
|
||||
out = append(out, map[string]any{
|
||||
"type": "function",
|
||||
"function": map[string]any{
|
||||
"name": t.Name(),
|
||||
"description": t.Description(),
|
||||
"parameters": params,
|
||||
},
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Anthropic-compatible driver (Bedrock).
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// anthropicCompatDriver backs Anthropic-wire providers that are not the direct
|
||||
// Anthropic API — Bedrock being the case here. It POSTs the Messages wire
|
||||
// format and decodes with AnthropicDecoder.
|
||||
type anthropicCompatDriver struct {
|
||||
name string
|
||||
baseURL string
|
||||
models []Model
|
||||
// path is the endpoint path appended to baseURL (Bedrock's invoke path
|
||||
// embeds the model id, so it is derived per request).
|
||||
pathFor func(model string) string
|
||||
// authHeader sets provider auth on the request (never logs the value).
|
||||
authHeader func(req *http.Request, apiKey string)
|
||||
}
|
||||
|
||||
func (d *anthropicCompatDriver) Name() string { return d.name }
|
||||
func (d *anthropicCompatDriver) Models() []Model { return d.models }
|
||||
|
||||
// StreamCompletion builds the Anthropic Messages request and streams it,
|
||||
// decoding with AnthropicDecoder.
|
||||
func (d *anthropicCompatDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) {
|
||||
if strings.TrimSpace(req.Config.APIKey) == "" {
|
||||
return nil, fmt.Errorf("%s: missing API key", d.name)
|
||||
}
|
||||
if err := checkImageSupport(d.name, req.Model, d.models, req.Context.Messages); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
body, err := encodeAnthropicRequest(req, d.models)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("%s: build request body: %w", d.name, err)
|
||||
}
|
||||
path := "/messages"
|
||||
if d.pathFor != nil {
|
||||
path = d.pathFor(req.Model)
|
||||
}
|
||||
newReq := func(ctx context.Context) (*http.Request, error) {
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, d.baseURL+path, bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
if d.authHeader != nil {
|
||||
d.authHeader(httpReq, req.Config.APIKey)
|
||||
}
|
||||
return httpReq, nil
|
||||
}
|
||||
return StreamRequest(ctx, TransportConfig{NewRequest: newReq, Decoder: NewAnthropicDecoder()})
|
||||
}
|
||||
|
||||
// encodeAnthropicRequest serializes a CompletionRequest into an Anthropic
|
||||
// Messages JSON body with streaming enabled. The system prompt is a top-level
|
||||
// field; tool results and tool calls follow the Messages content-block shape.
|
||||
func encodeAnthropicRequest(req CompletionRequest, models []Model) ([]byte, error) {
|
||||
msgs := make([]map[string]any, 0, len(req.Context.Messages))
|
||||
all := req.Context.Messages
|
||||
for i := 0; i < len(all); i++ {
|
||||
m := all[i]
|
||||
// Anthropic requires every tool_use in the preceding assistant message to
|
||||
// be answered by tool_result blocks in a single immediately-following user
|
||||
// message — splitting them across consecutive user turns yields the API
|
||||
// error "tool_use ids were found without tool_result blocks immediately
|
||||
// after". Coalesce runs of consecutive tool results into one turn.
|
||||
if tm, ok := m.(agentcore.ToolResultMessage); ok {
|
||||
blocks := []map[string]any{{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": tm.ToolCallID,
|
||||
"content": agentcore.ContentToText(tm.Content),
|
||||
}}
|
||||
for i+1 < len(all) {
|
||||
next, ok := all[i+1].(agentcore.ToolResultMessage)
|
||||
if !ok {
|
||||
break
|
||||
}
|
||||
i++
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": next.ToolCallID,
|
||||
"content": agentcore.ContentToText(next.Content),
|
||||
})
|
||||
}
|
||||
msgs = append(msgs, map[string]any{"role": "user", "content": blocks})
|
||||
continue
|
||||
}
|
||||
if enc := encodeAnthropicMessage(m); enc != nil {
|
||||
msgs = append(msgs, enc)
|
||||
}
|
||||
}
|
||||
body := map[string]any{
|
||||
"model": req.Model,
|
||||
"messages": msgs,
|
||||
"stream": true,
|
||||
}
|
||||
if sp := req.Context.SystemPrompt; sp != "" {
|
||||
body["system"] = sp
|
||||
}
|
||||
maxTok := maxOutputTokensFor(req)
|
||||
if maxTok <= 0 {
|
||||
// Anthropic requires max_tokens. Prefer the model's declared cap; fall
|
||||
// back to a coding-friendly default (4096 was too low and caused
|
||||
// truncation/retry loops on longer edits).
|
||||
maxTok = anthropicDefaultMaxTokens(req.Model, models)
|
||||
}
|
||||
// Extended thinking: when a thinking level is requested, enable the Anthropic
|
||||
// thinking block with a budget derived from the level. Omitted for off/unset
|
||||
// so non-thinking requests keep their prior shape.
|
||||
if budget := anthropicThinkingBudget(req.Config.ThinkingLevel); budget > 0 {
|
||||
// Anthropic counts thinking tokens toward max_tokens and requires
|
||||
// budget_tokens < max_tokens (else a 400). Guarantee headroom for the
|
||||
// visible reply by lifting max_tokens above the budget when the caller's
|
||||
// cap is too low to fit both the reasoning and a real answer.
|
||||
if minTok := budget + anthropicResponseHeadroom; maxTok < minTok {
|
||||
maxTok = minTok
|
||||
}
|
||||
body["thinking"] = map[string]any{"type": "enabled", "budget_tokens": budget}
|
||||
}
|
||||
body["max_tokens"] = maxTok
|
||||
if tools := encodeAnthropicTools(req.Context.Tools); len(tools) > 0 {
|
||||
body["tools"] = tools
|
||||
}
|
||||
return json.Marshal(body)
|
||||
}
|
||||
|
||||
// anthropicResponseHeadroom is the token margin reserved for the visible reply
|
||||
// on top of the thinking budget, so max_tokens always exceeds budget_tokens (an
|
||||
// Anthropic hard requirement) with room left for a real answer.
|
||||
const anthropicResponseHeadroom = 4096
|
||||
|
||||
// anthropicDefaultMaxTokens picks the max_tokens fallback when no explicit hint
|
||||
// is given: the model's declared MaxOutputTokens if present in the driver's
|
||||
// model catalog, otherwise 8192 (a coding-friendly default that avoids
|
||||
// premature truncation while staying within common model caps).
|
||||
func anthropicDefaultMaxTokens(model string, models []Model) int {
|
||||
for _, m := range models {
|
||||
if m.ID == model && m.MaxOutputTokens > 0 {
|
||||
return m.MaxOutputTokens
|
||||
}
|
||||
}
|
||||
return 8192
|
||||
}
|
||||
|
||||
// anthropicThinkingBudget maps a unified ThinkingLevel onto an Anthropic
|
||||
// thinking budget_tokens value. off/"" yields 0 (thinking block omitted).
|
||||
func anthropicThinkingBudget(level agentcore.ThinkingLevel) int {
|
||||
switch level {
|
||||
case agentcore.ThinkingMinimal:
|
||||
return 1024
|
||||
case agentcore.ThinkingLow:
|
||||
return 2048
|
||||
case agentcore.ThinkingMedium:
|
||||
return 8192
|
||||
case agentcore.ThinkingHigh:
|
||||
return 16384
|
||||
case agentcore.ThinkingXHigh:
|
||||
return 32768
|
||||
case agentcore.ThinkingMax:
|
||||
return 65536
|
||||
default: // off or unset
|
||||
return 0
|
||||
}
|
||||
}
|
||||
|
||||
// encodeAnthropicMessage maps one pigo message onto the Anthropic Messages
|
||||
// wire shape. Assistant tool calls become tool_use blocks; tool results become
|
||||
// a user message carrying a tool_result block (Anthropic's convention).
|
||||
func encodeAnthropicMessage(m agentcore.Message) map[string]any {
|
||||
switch msg := m.(type) {
|
||||
case agentcore.UserMessage:
|
||||
return map[string]any{"role": "user", "content": anthropicUserContent(msg.Content)}
|
||||
case agentcore.CompactionMessage:
|
||||
u := msg.AsUserMessage()
|
||||
return map[string]any{"role": "user", "content": anthropicUserContent(u.Content)}
|
||||
case agentcore.AssistantMessage:
|
||||
var blocks []map[string]any
|
||||
// Thinking blocks must precede tool_use in the same assistant turn:
|
||||
// Anthropic extended-thinking requires the prior thinking block (and its
|
||||
// signature) to be echoed back verbatim on tool-use turns, or the API
|
||||
// rejects/degrades the request. Emit them first.
|
||||
for _, c := range msg.Content {
|
||||
if t, ok := c.(agentcore.ThinkingContent); ok {
|
||||
if t.Redacted {
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "redacted_thinking", "data": t.ThinkingSignature,
|
||||
})
|
||||
continue
|
||||
}
|
||||
if t.Thinking == "" {
|
||||
continue
|
||||
}
|
||||
block := map[string]any{"type": "thinking", "thinking": t.Thinking}
|
||||
if t.ThinkingSignature != "" {
|
||||
block["signature"] = t.ThinkingSignature
|
||||
}
|
||||
blocks = append(blocks, block)
|
||||
}
|
||||
}
|
||||
for _, c := range msg.Content {
|
||||
switch b := c.(type) {
|
||||
case agentcore.TextContent:
|
||||
if b.Text == "" {
|
||||
continue
|
||||
}
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": b.Text})
|
||||
case agentcore.ToolCallContent:
|
||||
var input any
|
||||
_ = json.Unmarshal(b.Arguments, &input)
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "tool_use", "id": b.ID, "name": b.Name, "input": input,
|
||||
})
|
||||
}
|
||||
}
|
||||
if len(blocks) == 0 {
|
||||
// No usable content: emit a single space rather than an empty text
|
||||
// block ("text must be non-empty" is rejected by strict endpoints).
|
||||
blocks = []map[string]any{{"type": "text", "text": " "}}
|
||||
}
|
||||
return map[string]any{"role": "assistant", "content": blocks}
|
||||
case agentcore.ToolResultMessage:
|
||||
return map[string]any{"role": "user", "content": []map[string]any{{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": msg.ToolCallID,
|
||||
"content": agentcore.ContentToText(msg.Content),
|
||||
}}}
|
||||
default:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// anthropicUserContent shapes a user content list for the Anthropic Messages
|
||||
// wire. When there are no images it collapses to a plain string. When images
|
||||
// are present it emits the content-block array form: text blocks plus image
|
||||
// blocks with a base64 source ({"type":"image","source":{"type":"base64",
|
||||
// "media_type":<mime>,"data":<b64>}}).
|
||||
func anthropicUserContent(content agentcore.ContentList) any {
|
||||
hasImage := false
|
||||
for _, c := range content {
|
||||
if _, ok := c.(agentcore.ImageContent); ok {
|
||||
hasImage = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasImage {
|
||||
return agentcore.ContentToText(content)
|
||||
}
|
||||
blocks := make([]map[string]any, 0, len(content))
|
||||
for _, c := range content {
|
||||
switch b := c.(type) {
|
||||
case agentcore.TextContent:
|
||||
if b.Text == "" {
|
||||
continue
|
||||
}
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": b.Text})
|
||||
case agentcore.ImageContent:
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "image",
|
||||
"source": map[string]any{
|
||||
"type": "base64",
|
||||
"media_type": b.MimeType,
|
||||
"data": b.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
return blocks
|
||||
}
|
||||
|
||||
// contextHasImage reports whether any message in the context carries an image
|
||||
// content block, so a driver can reject image input on a non-multimodal model.
|
||||
func contextHasImage(msgs []agentcore.Message) bool {
|
||||
for _, m := range msgs {
|
||||
var content agentcore.ContentList
|
||||
switch msg := m.(type) {
|
||||
case agentcore.UserMessage:
|
||||
content = msg.Content
|
||||
case agentcore.CompactionMessage:
|
||||
content = msg.AsUserMessage().Content
|
||||
default:
|
||||
continue
|
||||
}
|
||||
for _, c := range content {
|
||||
if _, ok := c.(agentcore.ImageContent); ok {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// checkImageSupport returns a clear error when the request carries image input
|
||||
// but the named model (looked up in models) does not declare SupportsImages.
|
||||
// A model absent from the catalog is treated permissively (unknown capability),
|
||||
// deferring to the provider's own validation. This turns the silent drop of
|
||||
// image blocks on a text-only model into an actionable message.
|
||||
func checkImageSupport(providerName, model string, models []Model, msgs []agentcore.Message) error {
|
||||
if !contextHasImage(msgs) {
|
||||
return nil
|
||||
}
|
||||
for _, m := range models {
|
||||
if m.ID == model {
|
||||
if !m.SupportsImages {
|
||||
return fmt.Errorf("%s: model %q does not support image input", providerName, model)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// encodeAnthropicTools maps AgentTools onto the Anthropic tool schema.
|
||||
func encodeAnthropicTools(tools []agentcore.AgentTool) []map[string]any {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]map[string]any, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
schema := json.RawMessage(t.Schema())
|
||||
if len(schema) == 0 {
|
||||
schema = json.RawMessage("{}")
|
||||
}
|
||||
out = append(out, map[string]any{
|
||||
"name": t.Name(),
|
||||
"description": t.Description(),
|
||||
"input_schema": schema,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// maxOutputTokensFor pulls a max-output hint from the request Extra map, if any,
|
||||
// so callers can bound Anthropic responses without a registry lookup here.
|
||||
func maxOutputTokensFor(req CompletionRequest) int {
|
||||
if req.Config.Extra == nil {
|
||||
return 0
|
||||
}
|
||||
switch v := req.Config.Extra["max_tokens"].(type) {
|
||||
case int:
|
||||
return v
|
||||
case float64:
|
||||
return int(v)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Constructors.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
// openAICompatPreset captures the per-gateway differences among the
|
||||
// OpenAI-compatible providers: everything the constructors used to repeat
|
||||
// (provider name, default endpoint, whether auth is required, and any extra
|
||||
// headers). Collapsing the near-identical constructors onto this table means
|
||||
// "add a gateway" is a one-line entry plus a thin exported wrapper.
|
||||
type openAICompatPreset struct {
|
||||
name string
|
||||
defaultURL string // "" ⇒ no default; baseURL must be supplied by the caller
|
||||
requiresAuth bool
|
||||
extraHeaders map[string]string
|
||||
}
|
||||
|
||||
// newOpenAICompat builds an OpenAI-compatible driver from a preset, falling back
|
||||
// to the preset's default endpoint when baseURL is empty.
|
||||
func newOpenAICompat(p openAICompatPreset, baseURL string, models []Model) Provider {
|
||||
if baseURL == "" {
|
||||
baseURL = p.defaultURL
|
||||
}
|
||||
return &openAICompatDriver{
|
||||
name: p.name,
|
||||
baseURL: baseURL,
|
||||
models: models,
|
||||
requiresAuth: p.requiresAuth,
|
||||
extraHeaders: p.extraHeaders,
|
||||
}
|
||||
}
|
||||
|
||||
// NewOpenRouterProvider builds the OpenRouter provider — the reference
|
||||
// OpenAI-compatible gateway. baseURL defaults to the public endpoint when empty.
|
||||
func NewOpenRouterProvider(baseURL string, models []Model) Provider {
|
||||
return newOpenAICompat(openAICompatPreset{
|
||||
name: "openrouter",
|
||||
defaultURL: openRouterBaseURL,
|
||||
requiresAuth: true,
|
||||
extraHeaders: map[string]string{
|
||||
// OpenRouter attribution headers (optional but recommended).
|
||||
"HTTP-Referer": "https://github.com/smallnest/pigo",
|
||||
"X-Title": "pigo",
|
||||
},
|
||||
}, baseURL, models)
|
||||
}
|
||||
|
||||
// NewOllamaProvider builds the Ollama provider (local, OpenAI-compatible, no
|
||||
// auth). baseURL defaults to the local daemon when empty.
|
||||
func NewOllamaProvider(baseURL string, models []Model) Provider {
|
||||
return newOpenAICompat(openAICompatPreset{
|
||||
name: "ollama",
|
||||
defaultURL: ollamaBaseURL,
|
||||
requiresAuth: false,
|
||||
}, baseURL, models)
|
||||
}
|
||||
|
||||
// NewNvidiaProvider builds the NVIDIA provider (hosted NIM, OpenAI-compatible,
|
||||
// Bearer auth). baseURL defaults to the public integrate endpoint when empty.
|
||||
// The API key is resolved by the "nvidia" provider name (NVIDIA_API_KEY);
|
||||
// secret values are never logged.
|
||||
func NewNvidiaProvider(baseURL string, models []Model) Provider {
|
||||
return newOpenAICompat(openAICompatPreset{
|
||||
name: "nvidia",
|
||||
defaultURL: nvidiaBaseURL,
|
||||
requiresAuth: true,
|
||||
}, baseURL, models)
|
||||
}
|
||||
|
||||
// NewOpenAICompatibleProvider builds a generic OpenAI-compatible provider for an
|
||||
// arbitrary gateway reached by baseURL (Bearer auth). It is the target of an
|
||||
// explicit --protocol=openai selection: unlike the preset constructors it has no
|
||||
// default endpoint (baseURL must be supplied) and carries the neutral provider
|
||||
// name "openai", so an API key resolves from OPENAI_API_KEY (or the --api-key
|
||||
// override bound to that name). Secret values are never logged.
|
||||
func NewOpenAICompatibleProvider(baseURL string, models []Model) Provider {
|
||||
return newOpenAICompat(openAICompatPreset{
|
||||
name: "openai",
|
||||
defaultURL: "", // no default: caller must supply the endpoint
|
||||
requiresAuth: true,
|
||||
}, baseURL, models)
|
||||
}
|
||||
|
||||
// newAnthropicCompat builds an Anthropic-Messages driver, falling back to
|
||||
// defaultURL when baseURL is empty. It is the shared body of the two
|
||||
// Anthropic-wire constructors, which differ only in name, default endpoint, and
|
||||
// auth header.
|
||||
func newAnthropicCompat(name, defaultURL, baseURL string, models []Model, authHeader func(*http.Request, string)) Provider {
|
||||
if baseURL == "" {
|
||||
baseURL = defaultURL
|
||||
}
|
||||
return &anthropicCompatDriver{
|
||||
name: name,
|
||||
baseURL: baseURL,
|
||||
models: models,
|
||||
authHeader: authHeader,
|
||||
}
|
||||
}
|
||||
|
||||
// anthropicAuthHeaderFor returns the auth-header setter for an Anthropic-Messages
|
||||
// provider given its registry AuthScheme (spec.AuthScheme). The two shapes seen
|
||||
// among anthropic-protocol providers are:
|
||||
//
|
||||
// - AuthBearer → Authorization: Bearer <key>. Used by anthropic-protocol
|
||||
// gateways that authenticate with a plain bearer token on their /anthropic
|
||||
// endpoint.
|
||||
// - AuthXAPIKey → x-api-key: <key> plus the required anthropic-version
|
||||
// header. This is the direct-Anthropic convention and, per pi's behavior,
|
||||
// also what MiniMax (minimax / minimax-cn) uses on its /anthropic endpoint.
|
||||
//
|
||||
// Any other scheme (e.g. AuthAWS for Bedrock, AuthSpecial) falls back to the
|
||||
// x-api-key convention so a generic anthropic-protocol provider does not crash;
|
||||
// bespoke auth for those is layered by a later node. The returned func never
|
||||
// logs the secret value.
|
||||
func anthropicAuthHeaderFor(authScheme string) func(*http.Request, string) {
|
||||
if authScheme == AuthBearer {
|
||||
return func(req *http.Request, apiKey string) {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
}
|
||||
}
|
||||
return func(req *http.Request, apiKey string) {
|
||||
req.Header.Set("x-api-key", apiKey)
|
||||
req.Header.Set("anthropic-version", anthropicAPIVersion)
|
||||
}
|
||||
}
|
||||
|
||||
// NewAnthropicProvider builds a provider that speaks the Anthropic Messages wire
|
||||
// format directly (POST {baseURL}/messages), the target of an explicit
|
||||
// --protocol=anthropic selection. baseURL defaults to the public Anthropic API
|
||||
// when empty. Auth uses the Anthropic conventions: an x-api-key header plus the
|
||||
// required anthropic-version header. The API key resolves by the "anthropic"
|
||||
// provider name (ANTHROPIC_API_KEY / CLAUDE_API_KEY, or the --api-key override);
|
||||
// secret values are never logged.
|
||||
func NewAnthropicProvider(baseURL string, models []Model) Provider {
|
||||
return newAnthropicCompat("anthropic", anthropicBaseURL, baseURL, models,
|
||||
anthropicAuthHeaderFor(AuthXAPIKey))
|
||||
}
|
||||
|
||||
// NewAnthropicProtocolProvider builds a named Anthropic-Messages provider whose
|
||||
// auth header follows the given registry AuthScheme. It is the target of an
|
||||
// explicit --provider selection for any anthropic-protocol built-in (anthropic,
|
||||
// minimax, minimax-cn, and — routed generically for now — bedrock/
|
||||
// cloudflare-ai-gateway): the driver identity is the provider's own name (so
|
||||
// errors reference it), baseURL is the already-resolved endpoint (spec default
|
||||
// or override), and authScheme selects the header shape (see
|
||||
// anthropicAuthHeaderFor). Secret values are never logged.
|
||||
func NewAnthropicProtocolProvider(name, baseURL, authScheme string, models []Model) Provider {
|
||||
return newAnthropicCompat(name, anthropicBaseURL, baseURL, models,
|
||||
anthropicAuthHeaderFor(authScheme))
|
||||
}
|
||||
|
||||
// NewBedrockProvider builds the Bedrock provider, reusing the Anthropic Messages
|
||||
// decoder (Anthropic-on-Bedrock speaks the Messages wire format). baseURL
|
||||
// defaults to a us-east-1 runtime endpoint when empty; real deployments pass
|
||||
// the region-specific URL. Auth is a Bearer token (Bedrock API keys); SigV4
|
||||
// signing, when required, is layered by the caller's HTTP client.
|
||||
func NewBedrockProvider(baseURL string, models []Model) Provider {
|
||||
return newAnthropicCompat("bedrock", bedrockBaseURL, baseURL, models,
|
||||
func(req *http.Request, apiKey string) {
|
||||
req.Header.Set("Authorization", "Bearer "+apiKey)
|
||||
})
|
||||
}
|
||||
@@ -0,0 +1,275 @@
|
||||
// Tests for the Anthropic-Messages-protocol built-in providers (US-006, node
|
||||
// #187): anthropic, minimax, minimax-cn. They assert the registry metadata
|
||||
// (base_url, protocol, key env var) and that the constructed anthropic-compat
|
||||
// driver attaches the auth header dictated by the provider's AuthScheme.
|
||||
//
|
||||
// No real network calls are made: the auth header is exercised by invoking the
|
||||
// driver's authHeader func against a dummy *http.Request and inspecting the
|
||||
// resulting headers (the same package can reach the unexported driver fields).
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// TestAnthropicProtocolProviderRegistrySpecs asserts the registry metadata for
|
||||
// each Anthropic-Messages-protocol built-in: default base URL, wire protocol,
|
||||
// and the key env var (via envAPIKey + t.Setenv).
|
||||
func TestAnthropicProtocolProviderRegistrySpecs(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
wantBaseURL string
|
||||
wantEnv string
|
||||
}{
|
||||
{"anthropic", "https://api.anthropic.com/v1", "ANTHROPIC_API_KEY"},
|
||||
{"minimax", "https://api.minimax.io/anthropic", "MINIMAX_API_KEY"},
|
||||
{"minimax-cn", "https://api.minimaxi.com/anthropic", "MINIMAX_CN_API_KEY"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec(tc.name)
|
||||
if !ok {
|
||||
t.Fatalf("provider %q not found in registry", tc.name)
|
||||
}
|
||||
if spec.Protocol != ProtocolAnthropic {
|
||||
t.Errorf("protocol = %q, want %q", spec.Protocol, ProtocolAnthropic)
|
||||
}
|
||||
if spec.DefaultBaseURL != tc.wantBaseURL {
|
||||
t.Errorf("base_url = %q, want %q", spec.DefaultBaseURL, tc.wantBaseURL)
|
||||
}
|
||||
// Key resolution: the provider's key comes from tc.wantEnv.
|
||||
t.Setenv(tc.wantEnv, "SEKRET-"+tc.name)
|
||||
if got := envAPIKey(tc.name); got != "SEKRET-"+tc.name {
|
||||
t.Errorf("envAPIKey(%q) = %q, want key resolved from %s", tc.name, got, tc.wantEnv)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicProtocolProviderAuthHeader verifies that the driver built for
|
||||
// each provider (the way resolveNamedProvider builds it: name + resolved
|
||||
// base_url + spec.AuthScheme) targets the provider's base URL and sets the auth
|
||||
// header matching its AuthScheme. anthropic/minimax/minimax-cn are all
|
||||
// x-api-key + anthropic-version per pi's convention.
|
||||
func TestAnthropicProtocolProviderAuthHeader(t *testing.T) {
|
||||
for _, name := range []string{"anthropic", "minimax", "minimax-cn"} {
|
||||
t.Run(name, func(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec(name)
|
||||
if !ok {
|
||||
t.Fatalf("provider %q not found in registry", name)
|
||||
}
|
||||
p := NewAnthropicProtocolProvider(spec.Name, spec.DefaultBaseURL, spec.AuthScheme, nil)
|
||||
d, ok := p.(*anthropicCompatDriver)
|
||||
if !ok {
|
||||
t.Fatalf("provider %q is not an *anthropicCompatDriver", name)
|
||||
}
|
||||
if d.baseURL != spec.DefaultBaseURL {
|
||||
t.Errorf("baseURL = %q, want %q", d.baseURL, spec.DefaultBaseURL)
|
||||
}
|
||||
req, err := http.NewRequest(http.MethodPost, d.baseURL+"/messages", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("build request: %v", err)
|
||||
}
|
||||
d.authHeader(req, "SEKRET")
|
||||
assertAuthHeaderForScheme(t, req, spec.AuthScheme)
|
||||
if got := req.Header.Get("Authorization"); got != "" && spec.AuthScheme != AuthBearer {
|
||||
t.Errorf("unexpected Authorization header %q for x-api-key scheme", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicAuthSchemeSelection proves the auth-header mechanism itself:
|
||||
// AuthBearer yields Authorization: Bearer, AuthXAPIKey yields x-api-key plus the
|
||||
// anthropic-version header. This guarantees an anthropic-protocol provider can
|
||||
// select either header shape from its registry AuthScheme.
|
||||
func TestAnthropicAuthSchemeSelection(t *testing.T) {
|
||||
t.Run("bearer", func(t *testing.T) {
|
||||
p := NewAnthropicProtocolProvider("gw", "https://example.test/anthropic", AuthBearer, nil)
|
||||
req := newDummyReq(t)
|
||||
p.(*anthropicCompatDriver).authHeader(req, "SEKRET")
|
||||
if got := req.Header.Get("Authorization"); got != "Bearer SEKRET" {
|
||||
t.Errorf("Authorization = %q, want %q", got, "Bearer SEKRET")
|
||||
}
|
||||
if got := req.Header.Get("x-api-key"); got != "" {
|
||||
t.Errorf("x-api-key = %q, want empty for bearer scheme", got)
|
||||
}
|
||||
})
|
||||
t.Run("x-api-key", func(t *testing.T) {
|
||||
p := NewAnthropicProtocolProvider("anthropic", "", AuthXAPIKey, nil)
|
||||
req := newDummyReq(t)
|
||||
p.(*anthropicCompatDriver).authHeader(req, "SEKRET")
|
||||
if got := req.Header.Get("x-api-key"); got != "SEKRET" {
|
||||
t.Errorf("x-api-key = %q, want %q", got, "SEKRET")
|
||||
}
|
||||
if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion {
|
||||
t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion)
|
||||
}
|
||||
})
|
||||
// A non-crashing fallback for unwired schemes (e.g. Bedrock's AuthAWS): must
|
||||
// not panic and defaults to the x-api-key convention.
|
||||
t.Run("fallback", func(t *testing.T) {
|
||||
p := NewAnthropicProtocolProvider("bedrock", "", AuthAWS, nil)
|
||||
req := newDummyReq(t)
|
||||
p.(*anthropicCompatDriver).authHeader(req, "SEKRET")
|
||||
if got := req.Header.Get("x-api-key"); got != "SEKRET" {
|
||||
t.Errorf("fallback x-api-key = %q, want %q", got, "SEKRET")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewAnthropicProviderUnchanged guards against regressing the direct
|
||||
// Anthropic constructor: it must still default to the public endpoint and use
|
||||
// x-api-key + anthropic-version.
|
||||
func TestNewAnthropicProviderUnchanged(t *testing.T) {
|
||||
d, ok := NewAnthropicProvider("", nil).(*anthropicCompatDriver)
|
||||
if !ok {
|
||||
t.Fatal("NewAnthropicProvider did not return *anthropicCompatDriver")
|
||||
}
|
||||
if d.baseURL != anthropicBaseURL {
|
||||
t.Errorf("baseURL = %q, want %q", d.baseURL, anthropicBaseURL)
|
||||
}
|
||||
req := newDummyReq(t)
|
||||
d.authHeader(req, "SEKRET")
|
||||
if got := req.Header.Get("x-api-key"); got != "SEKRET" {
|
||||
t.Errorf("x-api-key = %q, want %q", got, "SEKRET")
|
||||
}
|
||||
if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion {
|
||||
t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion)
|
||||
}
|
||||
}
|
||||
|
||||
func newDummyReq(t *testing.T) *http.Request {
|
||||
t.Helper()
|
||||
req, err := http.NewRequest(http.MethodPost, "https://example.test/messages", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("build request: %v", err)
|
||||
}
|
||||
return req
|
||||
}
|
||||
|
||||
func assertAuthHeaderForScheme(t *testing.T, req *http.Request, scheme string) {
|
||||
t.Helper()
|
||||
if scheme == AuthBearer {
|
||||
if got := req.Header.Get("Authorization"); got != "Bearer SEKRET" {
|
||||
t.Errorf("Authorization = %q, want %q", got, "Bearer SEKRET")
|
||||
}
|
||||
return
|
||||
}
|
||||
if got := req.Header.Get("x-api-key"); got != "SEKRET" {
|
||||
t.Errorf("x-api-key = %q, want %q", got, "SEKRET")
|
||||
}
|
||||
if got := req.Header.Get("anthropic-version"); got != anthropicAPIVersion {
|
||||
t.Errorf("anthropic-version = %q, want %q", got, anthropicAPIVersion)
|
||||
}
|
||||
}
|
||||
|
||||
// TestEncodeAnthropicRequestIncludesModel guards a required field of the
|
||||
// Anthropic Messages API: the request body must carry the "model" id. Omitting
|
||||
// it makes the public API return 400 and OpenAI-compatible gateways return an
|
||||
// empty/error response, which the SSE decoder silently turns into an empty
|
||||
// assistant turn — a confusing "no output, no error" failure.
|
||||
func TestEncodeAnthropicRequestIncludesModel(t *testing.T) {
|
||||
req := CompletionRequest{
|
||||
Model: "claude-opus-4-8",
|
||||
}
|
||||
body, err := encodeAnthropicRequest(req, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("encodeAnthropicRequest: %v", err)
|
||||
}
|
||||
var decoded map[string]any
|
||||
if err := json.Unmarshal(body, &decoded); err != nil {
|
||||
t.Fatalf("unmarshal body: %v", err)
|
||||
}
|
||||
if got, ok := decoded["model"]; !ok {
|
||||
t.Fatalf("request body missing required \"model\" field; keys=%v", keysOf(decoded))
|
||||
} else if got != "claude-opus-4-8" {
|
||||
t.Errorf("model = %v, want claude-opus-4-8", got)
|
||||
}
|
||||
}
|
||||
|
||||
func keysOf(m map[string]any) []string {
|
||||
ks := make([]string, 0, len(m))
|
||||
for k := range m {
|
||||
ks = append(ks, k)
|
||||
}
|
||||
return ks
|
||||
}
|
||||
|
||||
// TestEncodeAnthropicRequestCoalescesToolResults guards the Anthropic protocol
|
||||
// rule that every tool_use must be answered by tool_result blocks in a single
|
||||
// immediately-following user message. Before the fix, each tool result was
|
||||
// encoded as its own user turn, which the API rejected with "tool_use ids were
|
||||
// found without tool_result blocks immediately after" whenever one assistant
|
||||
// turn carried multiple tool_use calls.
|
||||
func TestEncodeAnthropicRequestCoalescesToolResults(t *testing.T) {
|
||||
req := CompletionRequest{
|
||||
Model: "claude-x",
|
||||
Context: LlmContext{
|
||||
Messages: agentcore.MessageList{
|
||||
// Assistant turn with two tool_use blocks.
|
||||
agentcore.AssistantMessage{
|
||||
Content: agentcore.ContentList{
|
||||
agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{}`)),
|
||||
agentcore.NewToolCallContent("call_2", "write", json.RawMessage(`{}`)),
|
||||
},
|
||||
},
|
||||
// Two consecutive tool results for that turn.
|
||||
agentcore.ToolResultMessage{
|
||||
ToolCallID: "call_1",
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("result-1")},
|
||||
},
|
||||
agentcore.ToolResultMessage{
|
||||
ToolCallID: "call_2",
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("result-2")},
|
||||
},
|
||||
// Next assistant turn (text reply).
|
||||
agentcore.AssistantMessage{
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("done")},
|
||||
},
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
body, err := encodeAnthropicRequest(req, nil)
|
||||
if err != nil {
|
||||
t.Fatalf("encodeAnthropicRequest: %v", err)
|
||||
}
|
||||
decoded := decodeBody(t, body)
|
||||
msgs, ok := decoded["messages"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("messages not an array: %T", decoded["messages"])
|
||||
}
|
||||
if len(msgs) != 3 {
|
||||
t.Fatalf("want 3 top-level messages (assistant tool_use, one user tool_result turn, assistant reply), got %d:\n%v", len(msgs), msgs)
|
||||
}
|
||||
|
||||
// Message 2 must be the single user turn holding BOTH tool_result blocks.
|
||||
second, ok := msgs[1].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("message[1] not an object: %T", msgs[1])
|
||||
}
|
||||
if second["role"] != "user" {
|
||||
t.Fatalf("message[1].role = %v, want user", second["role"])
|
||||
}
|
||||
blocks, ok := second["content"].([]any)
|
||||
if !ok {
|
||||
t.Fatalf("message[1].content not an array: %T", second["content"])
|
||||
}
|
||||
if len(blocks) != 2 {
|
||||
t.Fatalf("want 2 tool_result blocks in the coalesced turn, got %d:\n%v", len(blocks), second["content"])
|
||||
}
|
||||
for i, wantID := range []string{"call_1", "call_2"} {
|
||||
block, ok := blocks[i].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("block[%d] not an object: %T", i, blocks[i])
|
||||
}
|
||||
if block["type"] != "tool_result" || block["tool_use_id"] != wantID {
|
||||
t.Errorf("block[%d] = %v, want tool_result for %s", i, block, wantID)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,151 @@
|
||||
// Node #186: end-to-end wiring test for every OpenAI-protocol built-in provider.
|
||||
//
|
||||
// For each provider reachable via --provider (US-005) this asserts three things:
|
||||
// (a) the registry spec reports Protocol "openai" and the PRD-mandated default
|
||||
// base URL,
|
||||
// (b) its primary API-key env var resolves through envAPIKey (the same path
|
||||
// auth.go uses at request time), and
|
||||
// (c) the generic OpenAI-compatible construction path used by main.go's
|
||||
// resolveNamedProvider builds a non-nil driver bound to the spec's model.
|
||||
//
|
||||
// It also pins the restored legacy key aliases (CLAUDE_API_KEY, GOOGLE_API_KEY,
|
||||
// NVIDIA_NIM_API_KEY) and the non-standard HF_TOKEN key name.
|
||||
//
|
||||
// This file is intentionally separate from providers_test.go: a sibling node
|
||||
// edits that file concurrently.
|
||||
package provider
|
||||
|
||||
import "testing"
|
||||
|
||||
// openAIWiringCase describes one OpenAI-protocol provider's expected registry
|
||||
// metadata and primary key env var.
|
||||
type openAIWiringCase struct {
|
||||
name string
|
||||
baseURL string
|
||||
primaryEnv string
|
||||
}
|
||||
|
||||
// openAIProviderCases lists every OpenAI-protocol provider that must work
|
||||
// end-to-end via --provider (US-005). Base URLs mirror the PRD's Technical
|
||||
// Considerations table.
|
||||
var openAIProviderCases = []openAIWiringCase{
|
||||
{"groq", "https://api.groq.com/openai/v1", "GROQ_API_KEY"},
|
||||
{"xai", "https://api.x.ai/v1", "XAI_API_KEY"},
|
||||
{"cerebras", "https://api.cerebras.ai/v1", "CEREBRAS_API_KEY"},
|
||||
{"mistral", "https://api.mistral.ai", "MISTRAL_API_KEY"},
|
||||
{"moonshotai", "https://api.moonshot.ai/v1", "MOONSHOT_API_KEY"},
|
||||
{"moonshotai-cn", "https://api.moonshot.cn/v1", "MOONSHOT_API_KEY"},
|
||||
{"fireworks", "https://api.fireworks.ai/inference", "FIREWORKS_API_KEY"},
|
||||
{"together", "https://api.together.ai/v1", "TOGETHER_API_KEY"},
|
||||
{"openrouter", "https://openrouter.ai/api/v1", "OPENROUTER_API_KEY"},
|
||||
{"nvidia", "https://integrate.api.nvidia.com/v1", "NVIDIA_API_KEY"},
|
||||
{"zai", "https://api.z.ai/api/coding/paas/v4", "ZAI_API_KEY"},
|
||||
{"zai-coding-cn", "https://open.bigmodel.cn/api/coding/paas/v4", "ZAI_CODING_CN_API_KEY"},
|
||||
{"kimi-coding", "https://api.kimi.com/coding", "KIMI_API_KEY"},
|
||||
{"opencode", "https://opencode.ai/zen", "OPENCODE_API_KEY"},
|
||||
{"opencode-go", "https://opencode.ai/zen/go", "OPENCODE_API_KEY"},
|
||||
{"huggingface", "https://router.huggingface.co/v1", "HF_TOKEN"},
|
||||
{"ant-ling", "https://api.ant-ling.com/v1", "ANT_LING_API_KEY"},
|
||||
{"vercel-ai-gateway", "https://ai-gateway.vercel.sh", "AI_GATEWAY_API_KEY"},
|
||||
{"xiaomi", "https://api.xiaomimimo.com/v1", "XIAOMI_API_KEY"},
|
||||
{"xiaomi-token-plan-cn", "https://token-plan-cn.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_CN_API_KEY"},
|
||||
{"xiaomi-token-plan-ams", "https://token-plan-ams.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_AMS_API_KEY"},
|
||||
{"xiaomi-token-plan-sgp", "https://token-plan-sgp.xiaomimimo.com/v1", "XIAOMI_TOKEN_PLAN_SGP_API_KEY"},
|
||||
}
|
||||
|
||||
func TestOpenAIProviderWiring(t *testing.T) {
|
||||
for _, tc := range openAIProviderCases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
// (a) registry spec: protocol + default base URL + primary env var.
|
||||
spec, ok := LookupProviderSpec(tc.name)
|
||||
if !ok {
|
||||
t.Fatalf("LookupProviderSpec(%q): not found in registry", tc.name)
|
||||
}
|
||||
if spec.Protocol != ProtocolOpenAI {
|
||||
t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAI)
|
||||
}
|
||||
if spec.DefaultBaseURL != tc.baseURL {
|
||||
t.Errorf("DefaultBaseURL = %q, want %q", spec.DefaultBaseURL, tc.baseURL)
|
||||
}
|
||||
if len(spec.EnvVars) == 0 || spec.EnvVars[0] != tc.primaryEnv {
|
||||
t.Errorf("primary EnvVar = %v, want first = %q", spec.EnvVars, tc.primaryEnv)
|
||||
}
|
||||
|
||||
// (b) key resolution via the primary env var, using the same
|
||||
// envAPIKey path auth.go relies on at request time.
|
||||
t.Setenv(tc.primaryEnv, "sk-"+tc.name)
|
||||
if got := envAPIKey(tc.name); got != "sk-"+tc.name {
|
||||
t.Errorf("envAPIKey(%q) = %q, want %q", tc.name, got, "sk-"+tc.name)
|
||||
}
|
||||
|
||||
// (c) construction path equivalent to main.go's resolveNamedProvider
|
||||
// for an openai-protocol spec: build a generic OpenAI-compatible
|
||||
// driver against the spec's base URL, bound to the spec's model.
|
||||
models := []Model{{Provider: spec.Name, ID: "test-model", SupportsImages: true}}
|
||||
drv := NewOpenAICompatibleProvider(spec.DefaultBaseURL, models)
|
||||
if drv == nil {
|
||||
t.Fatalf("NewOpenAICompatibleProvider(%q) returned nil", spec.DefaultBaseURL)
|
||||
}
|
||||
got := drv.Models()
|
||||
if len(got) != 1 || got[0].Provider != spec.Name || got[0].ID != "test-model" {
|
||||
t.Errorf("driver Models() = %+v, want one model bound to provider %q", got, spec.Name)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestLegacyKeyAliases pins the secondary env vars restored to the registry so
|
||||
// credentials set under older names still resolve.
|
||||
func TestLegacyKeyAliases(t *testing.T) {
|
||||
aliases := []struct {
|
||||
provider string
|
||||
primary string // primary env var (must NOT be set for the alias to be exercised)
|
||||
alias string // legacy alias env var under test
|
||||
}{
|
||||
{"anthropic", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"},
|
||||
{"google", "GEMINI_API_KEY", "GOOGLE_API_KEY"},
|
||||
{"nvidia", "NVIDIA_API_KEY", "NVIDIA_NIM_API_KEY"},
|
||||
}
|
||||
for _, a := range aliases {
|
||||
t.Run(a.provider, func(t *testing.T) {
|
||||
// Ensure the alias is listed in the registry's EnvVars.
|
||||
spec, ok := LookupProviderSpec(a.provider)
|
||||
if !ok {
|
||||
t.Fatalf("LookupProviderSpec(%q): not found", a.provider)
|
||||
}
|
||||
found := false
|
||||
for _, e := range spec.EnvVars {
|
||||
if e == a.alias {
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("EnvVars %v missing legacy alias %q", spec.EnvVars, a.alias)
|
||||
}
|
||||
// Clear the primary so only the alias can satisfy resolution, then
|
||||
// assert the alias resolves.
|
||||
t.Setenv(a.primary, "")
|
||||
t.Setenv(a.alias, "legacy-"+a.provider)
|
||||
if got := envAPIKey(a.provider); got != "legacy-"+a.provider {
|
||||
t.Errorf("envAPIKey(%q) via %s = %q, want %q", a.provider, a.alias, got, "legacy-"+a.provider)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestHuggingFaceTokenEnv asserts the non-standard HF_TOKEN key name resolves
|
||||
// for huggingface (it does not follow the <PROVIDER>_API_KEY convention).
|
||||
func TestHuggingFaceTokenEnv(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("huggingface")
|
||||
if !ok {
|
||||
t.Fatal("LookupProviderSpec(\"huggingface\"): not found")
|
||||
}
|
||||
if len(spec.EnvVars) == 0 || spec.EnvVars[0] != "HF_TOKEN" {
|
||||
t.Fatalf("huggingface EnvVars = %v, want first = HF_TOKEN", spec.EnvVars)
|
||||
}
|
||||
t.Setenv("HF_TOKEN", "hf-secret")
|
||||
if got := envAPIKey("huggingface"); got != "hf-secret" {
|
||||
t.Errorf("envAPIKey(\"huggingface\") = %q, want %q", got, "hf-secret")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,227 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// captureServer records the last request path, headers, and decoded JSON body,
|
||||
// then replays a canned SSE stream. It lets the provider tests assert the wire
|
||||
// shape a driver produced without a live upstream.
|
||||
type captureServer struct {
|
||||
srv *httptest.Server
|
||||
path string
|
||||
headers http.Header
|
||||
body map[string]any
|
||||
}
|
||||
|
||||
func newCaptureServer(t *testing.T, sseBody string) *captureServer {
|
||||
t.Helper()
|
||||
cs := &captureServer{}
|
||||
cs.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
cs.path = r.URL.Path
|
||||
cs.headers = r.Header.Clone()
|
||||
raw, _ := io.ReadAll(r.Body)
|
||||
_ = json.Unmarshal(raw, &cs.body)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = w.Write([]byte(sseBody))
|
||||
}))
|
||||
t.Cleanup(cs.srv.Close)
|
||||
return cs
|
||||
}
|
||||
|
||||
// drainStream collects event kinds and the final message from a provider stream.
|
||||
func drainStream(t *testing.T, stream *AssistantMessageEventStream) ([]string, agentcore.AssistantMessage) {
|
||||
t.Helper()
|
||||
var kinds []string
|
||||
for ev := range stream.Events() {
|
||||
kinds = append(kinds, ev.EventKind())
|
||||
}
|
||||
final, err := stream.Result(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("stream result: %v", err)
|
||||
}
|
||||
return kinds, final
|
||||
}
|
||||
|
||||
// TestOpenRouterProviderStreamsChatCompletions drives OpenRouter (the reference
|
||||
// OpenAI-compatible provider) end to end: the driver must POST to
|
||||
// /chat/completions with a Bearer token and stream through OpenAIDecoder.
|
||||
func TestOpenRouterProviderStreamsChatCompletions(t *testing.T) {
|
||||
cs := newCaptureServer(t, openaiToolCallSSE)
|
||||
p := NewOpenRouterProvider(cs.srv.URL, []Model{{Provider: "openrouter", ID: "openai/gpt-4o"}})
|
||||
|
||||
if p.Name() != "openrouter" {
|
||||
t.Errorf("name = %q, want openrouter", p.Name())
|
||||
}
|
||||
stream, err := p.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "openai/gpt-4o",
|
||||
Context: LlmContext{SystemPrompt: "be brief", Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion: %v", err)
|
||||
}
|
||||
kinds, final := drainStream(t, stream)
|
||||
if kinds[len(kinds)-1] != StreamEventDone {
|
||||
t.Errorf("last event = %q, want done", kinds[len(kinds)-1])
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
|
||||
// Wire assertions: path, auth header, and request body shape.
|
||||
if cs.path != "/chat/completions" {
|
||||
t.Errorf("path = %q, want /chat/completions", cs.path)
|
||||
}
|
||||
if got := cs.headers.Get("Authorization"); got != "Bearer sk-test" {
|
||||
t.Errorf("auth header = %q, want Bearer sk-test", got)
|
||||
}
|
||||
if cs.body["stream"] != true {
|
||||
t.Errorf("stream flag = %v, want true", cs.body["stream"])
|
||||
}
|
||||
if cs.body["model"] != "openai/gpt-4o" {
|
||||
t.Errorf("model = %v, want openai/gpt-4o", cs.body["model"])
|
||||
}
|
||||
msgs, _ := cs.body["messages"].([]any)
|
||||
if len(msgs) != 2 {
|
||||
t.Fatalf("messages len = %d, want 2 (system+user)", len(msgs))
|
||||
}
|
||||
first, _ := msgs[0].(map[string]any)
|
||||
if first["role"] != "system" || first["content"] != "be brief" {
|
||||
t.Errorf("system message = %v", first)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenRouterMissingKeyIsEarlyError verifies a missing API key is the early
|
||||
// "cannot build the stream" returned error, and the provider name is named
|
||||
// without leaking any value.
|
||||
func TestOpenRouterMissingKeyIsEarlyError(t *testing.T) {
|
||||
p := NewOpenRouterProvider("", nil)
|
||||
_, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"})
|
||||
if err == nil {
|
||||
t.Fatal("missing API key must return an early error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "openrouter") {
|
||||
t.Errorf("error should name the provider, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOllamaProviderNoAuth verifies the local Ollama provider sends no
|
||||
// Authorization header and still streams through OpenAIDecoder.
|
||||
func TestOllamaProviderNoAuth(t *testing.T) {
|
||||
cs := newCaptureServer(t, openaiToolCallSSE)
|
||||
p := NewOllamaProvider(cs.srv.URL, []Model{{Provider: "ollama", ID: "llama3"}})
|
||||
if p.Name() != "ollama" {
|
||||
t.Errorf("name = %q, want ollama", p.Name())
|
||||
}
|
||||
// No API key configured — Ollama must not require one.
|
||||
stream, err := p.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "llama3",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion (no auth): %v", err)
|
||||
}
|
||||
_, final := drainStream(t, stream)
|
||||
if final.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason = %q, want tool_use", final.StopReason)
|
||||
}
|
||||
if got := cs.headers.Get("Authorization"); got != "" {
|
||||
t.Errorf("Ollama must send no auth header, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestBedrockProviderStreamsAnthropicWire verifies the Bedrock provider POSTs
|
||||
// the Anthropic Messages wire format (system top-level, max_tokens present) and
|
||||
// streams through AnthropicDecoder.
|
||||
func TestBedrockProviderStreamsAnthropicWire(t *testing.T) {
|
||||
cs := newCaptureServer(t, anthropicToolUseSSE)
|
||||
p := NewBedrockProvider(cs.srv.URL, []Model{{Provider: "bedrock", ID: "anthropic.claude-3"}})
|
||||
if p.Name() != "bedrock" {
|
||||
t.Errorf("name = %q, want bedrock", p.Name())
|
||||
}
|
||||
stream, err := p.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "anthropic.claude-3",
|
||||
Context: LlmContext{SystemPrompt: "sys", Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}},
|
||||
Config: StreamConfig{APIKey: "bedrock-key"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion: %v", err)
|
||||
}
|
||||
kinds, _ := drainStream(t, stream)
|
||||
if kinds[len(kinds)-1] != StreamEventDone {
|
||||
t.Errorf("last event = %q, want done", kinds[len(kinds)-1])
|
||||
}
|
||||
if cs.body["system"] != "sys" {
|
||||
t.Errorf("system = %v, want top-level 'sys'", cs.body["system"])
|
||||
}
|
||||
if cs.body["max_tokens"] == nil {
|
||||
t.Error("Anthropic wire requires max_tokens")
|
||||
}
|
||||
if cs.body["stream"] != true {
|
||||
t.Errorf("stream flag = %v, want true", cs.body["stream"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestBedrockMissingKeyIsEarlyError mirrors the OpenRouter early-error check.
|
||||
func TestBedrockMissingKeyIsEarlyError(t *testing.T) {
|
||||
p := NewBedrockProvider("", nil)
|
||||
_, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"})
|
||||
if err == nil {
|
||||
t.Fatal("missing API key must return an early error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "bedrock") {
|
||||
t.Errorf("error should name the provider, got %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNvidiaProviderStreamsChatCompletions verifies the NVIDIA NIM provider
|
||||
// rides the OpenAI-compatible driver: it POSTs to /chat/completions with a
|
||||
// Bearer token and streams through OpenAIDecoder, exactly like OpenRouter.
|
||||
func TestNvidiaProviderStreamsChatCompletions(t *testing.T) {
|
||||
cs := newCaptureServer(t, openaiToolCallSSE)
|
||||
p := NewNvidiaProvider(cs.srv.URL, []Model{{Provider: "nvidia", ID: "meta/llama-3.3-70b-instruct"}})
|
||||
if p.Name() != "nvidia" {
|
||||
t.Errorf("name = %q, want nvidia", p.Name())
|
||||
}
|
||||
stream, err := p.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "meta/llama-3.3-70b-instruct",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{agentcore.UserMessage{RoleField: agentcore.RoleUser, Content: agentcore.ContentList{agentcore.NewTextContent("hi")}}}},
|
||||
Config: StreamConfig{APIKey: "nvapi-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion: %v", err)
|
||||
}
|
||||
kinds, _ := drainStream(t, stream)
|
||||
if kinds[len(kinds)-1] != StreamEventDone {
|
||||
t.Errorf("last event = %q, want done", kinds[len(kinds)-1])
|
||||
}
|
||||
if cs.path != "/chat/completions" {
|
||||
t.Errorf("path = %q, want /chat/completions", cs.path)
|
||||
}
|
||||
if got := cs.headers.Get("Authorization"); got != "Bearer nvapi-test" {
|
||||
t.Errorf("auth header = %q, want Bearer nvapi-test", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestNvidiaMissingKeyIsEarlyError verifies NVIDIA requires an API key and
|
||||
// reports the missing key as an early error naming the provider (never a value).
|
||||
func TestNvidiaMissingKeyIsEarlyError(t *testing.T) {
|
||||
p := NewNvidiaProvider("", nil)
|
||||
_, err := p.StreamCompletion(context.Background(), CompletionRequest{Model: "m"})
|
||||
if err == nil {
|
||||
t.Fatal("missing API key must return an early error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "nvidia") {
|
||||
t.Errorf("error should name the provider, got %v", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
// This file defines the central provider registry (US-001): a single source of
|
||||
// truth for every built-in provider's metadata — its name, the environment
|
||||
// variables that carry its API key (in precedence order), its default base URL,
|
||||
// wire protocol, auth scheme, any extra headers, and provider-specific base-URL
|
||||
// override env vars.
|
||||
//
|
||||
// The registry is deliberately additive and self-contained: later nodes wire
|
||||
// auth resolution (auth.go), the --provider flag (main.go), base_url overrides,
|
||||
// and per-provider construction (providers.go) to READ from it. This node only
|
||||
// introduces the data + a lookup, so it does not change existing behavior.
|
||||
//
|
||||
// Data source: the PRD "Technical Considerations" table
|
||||
// (tasks/prd-provider-env-parity.md), derived from pi's env-api-keys.ts and the
|
||||
// per-provider *.models.ts files.
|
||||
//
|
||||
// Security: the registry holds only env var NAMES, never secret values. Keys are
|
||||
// resolved from the environment at request time (see auth.go) and never logged.
|
||||
package provider
|
||||
|
||||
// ProviderSpec is the metadata describing one built-in provider. It is the
|
||||
// single source of truth consumed by auth resolution, the --provider flag,
|
||||
// base_url override handling, and per-provider wiring.
|
||||
type ProviderSpec struct {
|
||||
// Name is the canonical provider name (e.g. "deepseek", "zai-coding-cn").
|
||||
Name string
|
||||
// EnvVars lists the environment variables checked (in precedence order) for
|
||||
// this provider's API key. The first non-empty value wins.
|
||||
EnvVars []string
|
||||
// DefaultBaseURL is the provider's default API endpoint. It may be a template
|
||||
// (containing placeholders like {region}) for providers whose endpoint is
|
||||
// composed from additional parameters (Bedrock, Vertex, Cloudflare, Azure).
|
||||
DefaultBaseURL string
|
||||
// Protocol is the wire protocol the provider speaks: "openai" (OpenAI Chat
|
||||
// Completions) or "anthropic" (Anthropic Messages).
|
||||
Protocol string
|
||||
// AuthScheme names how credentials are attached: "bearer", "x-api-key",
|
||||
// "aws", "azure", or "special".
|
||||
AuthScheme string
|
||||
// ExtraHeaders are provider-specific headers attached to every request (may
|
||||
// be nil).
|
||||
ExtraHeaders map[string]string
|
||||
// BaseURLEnvVars lists provider-specific base-URL override environment
|
||||
// variables (e.g. AZURE_OPENAI_BASE_URL), in precedence order. May be empty;
|
||||
// the generic <PROVIDER>_BASE_URL convention is handled by callers.
|
||||
BaseURLEnvVars []string
|
||||
}
|
||||
|
||||
// Protocol values.
|
||||
const (
|
||||
ProtocolOpenAI = "openai"
|
||||
ProtocolAnthropic = "anthropic"
|
||||
)
|
||||
|
||||
// AuthScheme values.
|
||||
const (
|
||||
AuthBearer = "bearer"
|
||||
AuthXAPIKey = "x-api-key"
|
||||
AuthAWS = "aws"
|
||||
AuthAzure = "azure"
|
||||
AuthSpecial = "special"
|
||||
)
|
||||
|
||||
// providerRegistry is the ordered list of all built-in provider specs. Order is
|
||||
// stable so callers that enumerate providers (e.g. --help) get a deterministic
|
||||
// list. LookupProviderSpec indexes it by name.
|
||||
var providerRegistry = []ProviderSpec{
|
||||
{
|
||||
Name: "anthropic",
|
||||
EnvVars: []string{"ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"},
|
||||
DefaultBaseURL: anthropicBaseURL, // https://api.anthropic.com/v1
|
||||
Protocol: ProtocolAnthropic,
|
||||
AuthScheme: AuthXAPIKey,
|
||||
},
|
||||
{
|
||||
Name: "openai",
|
||||
EnvVars: []string{"OPENAI_API_KEY"},
|
||||
DefaultBaseURL: "https://api.openai.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "ant-ling",
|
||||
EnvVars: []string{"ANT_LING_API_KEY"},
|
||||
DefaultBaseURL: "https://api.ant-ling.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "deepseek",
|
||||
EnvVars: []string{"DEEPSEEK_API_KEY"},
|
||||
DefaultBaseURL: "https://api.deepseek.com",
|
||||
Protocol: ProtocolOpenAIResponses,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "nvidia",
|
||||
EnvVars: []string{"NVIDIA_API_KEY", "NVIDIA_NIM_API_KEY"},
|
||||
DefaultBaseURL: nvidiaBaseURL, // https://integrate.api.nvidia.com/v1
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "google",
|
||||
EnvVars: []string{"GEMINI_API_KEY", "GOOGLE_API_KEY"},
|
||||
DefaultBaseURL: "https://generativelanguage.googleapis.com/v1beta",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "groq",
|
||||
EnvVars: []string{"GROQ_API_KEY"},
|
||||
DefaultBaseURL: "https://api.groq.com/openai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "cerebras",
|
||||
EnvVars: []string{"CEREBRAS_API_KEY"},
|
||||
DefaultBaseURL: "https://api.cerebras.ai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "xai",
|
||||
EnvVars: []string{"XAI_API_KEY"},
|
||||
DefaultBaseURL: "https://api.x.ai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "openrouter",
|
||||
EnvVars: []string{"OPENROUTER_API_KEY"},
|
||||
DefaultBaseURL: openRouterBaseURL, // https://openrouter.ai/api/v1
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "vercel-ai-gateway",
|
||||
EnvVars: []string{"AI_GATEWAY_API_KEY"},
|
||||
DefaultBaseURL: "https://ai-gateway.vercel.sh",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "zai",
|
||||
EnvVars: []string{"ZAI_API_KEY"},
|
||||
DefaultBaseURL: "https://api.z.ai/api/coding/paas/v4",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "zai-coding-cn",
|
||||
EnvVars: []string{"ZAI_CODING_CN_API_KEY"},
|
||||
DefaultBaseURL: "https://open.bigmodel.cn/api/coding/paas/v4",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "mistral",
|
||||
EnvVars: []string{"MISTRAL_API_KEY"},
|
||||
DefaultBaseURL: "https://api.mistral.ai",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "minimax",
|
||||
EnvVars: []string{"MINIMAX_API_KEY"},
|
||||
DefaultBaseURL: "https://api.minimax.io/anthropic",
|
||||
Protocol: ProtocolAnthropic,
|
||||
AuthScheme: AuthXAPIKey,
|
||||
},
|
||||
{
|
||||
Name: "minimax-cn",
|
||||
EnvVars: []string{"MINIMAX_CN_API_KEY"},
|
||||
DefaultBaseURL: "https://api.minimaxi.com/anthropic",
|
||||
Protocol: ProtocolAnthropic,
|
||||
AuthScheme: AuthXAPIKey,
|
||||
},
|
||||
{
|
||||
Name: "moonshotai",
|
||||
EnvVars: []string{"MOONSHOT_API_KEY"},
|
||||
DefaultBaseURL: "https://api.moonshot.ai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "moonshotai-cn",
|
||||
EnvVars: []string{"MOONSHOT_API_KEY"},
|
||||
DefaultBaseURL: "https://api.moonshot.cn/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "huggingface",
|
||||
EnvVars: []string{"HF_TOKEN"},
|
||||
DefaultBaseURL: "https://router.huggingface.co/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "fireworks",
|
||||
EnvVars: []string{"FIREWORKS_API_KEY"},
|
||||
DefaultBaseURL: "https://api.fireworks.ai/inference",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "together",
|
||||
EnvVars: []string{"TOGETHER_API_KEY"},
|
||||
DefaultBaseURL: "https://api.together.ai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "opencode",
|
||||
EnvVars: []string{"OPENCODE_API_KEY"},
|
||||
DefaultBaseURL: "https://opencode.ai/zen",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "opencode-go",
|
||||
EnvVars: []string{"OPENCODE_API_KEY"},
|
||||
DefaultBaseURL: "https://opencode.ai/zen/go",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "kimi-coding",
|
||||
EnvVars: []string{"KIMI_API_KEY"},
|
||||
DefaultBaseURL: "https://api.kimi.com/coding",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "xiaomi",
|
||||
EnvVars: []string{"XIAOMI_API_KEY"},
|
||||
DefaultBaseURL: "https://api.xiaomimimo.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "xiaomi-token-plan-cn",
|
||||
EnvVars: []string{"XIAOMI_TOKEN_PLAN_CN_API_KEY"},
|
||||
DefaultBaseURL: "https://token-plan-cn.xiaomimimo.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "xiaomi-token-plan-ams",
|
||||
EnvVars: []string{"XIAOMI_TOKEN_PLAN_AMS_API_KEY"},
|
||||
DefaultBaseURL: "https://token-plan-ams.xiaomimimo.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "xiaomi-token-plan-sgp",
|
||||
EnvVars: []string{"XIAOMI_TOKEN_PLAN_SGP_API_KEY"},
|
||||
DefaultBaseURL: "https://token-plan-sgp.xiaomimimo.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
// Chinese cloud LLM platforms. All four expose OpenAI-compatible endpoints
|
||||
// authenticated with a plain Bearer API key, so they reuse the standard
|
||||
// OpenAI-compatible driver with no bespoke auth. Base URLs are the platforms'
|
||||
// OpenAI-compatible endpoints as documented at implementation time.
|
||||
{
|
||||
// Baidu AI Cloud Qianfan (Baidu Qianfan).
|
||||
Name: "qianfan",
|
||||
EnvVars: []string{"QIANFAN_API_KEY"},
|
||||
DefaultBaseURL: "https://qianfan.baidubce.com/v2",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
// ByteDance Volcengine Ark (Volcengine Ark). ARK_API_KEY is the platform's
|
||||
// conventional variable; VOLCENGINE_API_KEY is accepted as a fallback.
|
||||
Name: "volcengine",
|
||||
EnvVars: []string{"ARK_API_KEY", "VOLCENGINE_API_KEY"},
|
||||
DefaultBaseURL: "https://ark.cn-beijing.volces.com/api/v3",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
// Alibaba Cloud DashScope (DashScope), OpenAI-compatible mode.
|
||||
Name: "dashscope",
|
||||
EnvVars: []string{"DASHSCOPE_API_KEY"},
|
||||
DefaultBaseURL: "https://dashscope.aliyuncs.com/compatible-mode/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
// Tencent Hunyuan (Hunyuan), OpenAI-compatible endpoint.
|
||||
Name: "hunyuan",
|
||||
EnvVars: []string{"HUNYUAN_API_KEY"},
|
||||
DefaultBaseURL: "https://api.hunyuan.cloud.tencent.com/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "azure-openai-responses",
|
||||
EnvVars: []string{"AZURE_OPENAI_API_KEY"},
|
||||
// Endpoint is composed from AZURE_OPENAI_BASE_URL / AZURE_OPENAI_RESOURCE_NAME
|
||||
// per the Azure OpenAI convention; there is no fixed public default.
|
||||
DefaultBaseURL: "",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthAzure,
|
||||
BaseURLEnvVars: []string{"AZURE_OPENAI_BASE_URL"},
|
||||
},
|
||||
{
|
||||
Name: "amazon-bedrock",
|
||||
EnvVars: []string{"AWS_BEARER_TOKEN_BEDROCK"},
|
||||
// Region-specific runtime endpoint; {AWS_REGION} defaults to us-east-1.
|
||||
DefaultBaseURL: "https://bedrock-runtime.{AWS_REGION}.amazonaws.com",
|
||||
Protocol: ProtocolAnthropic,
|
||||
AuthScheme: AuthAWS,
|
||||
},
|
||||
{
|
||||
Name: "google-vertex",
|
||||
EnvVars: []string{"GOOGLE_CLOUD_API_KEY"},
|
||||
// Location-specific endpoint; protocol varies by model (Gemini vs Claude).
|
||||
DefaultBaseURL: "https://{location}-aiplatform.googleapis.com",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthSpecial,
|
||||
},
|
||||
{
|
||||
Name: "cloudflare-workers-ai",
|
||||
EnvVars: []string{"CLOUDFLARE_API_KEY"},
|
||||
// {id} is CLOUDFLARE_ACCOUNT_ID.
|
||||
DefaultBaseURL: "https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/v1",
|
||||
Protocol: ProtocolOpenAI,
|
||||
AuthScheme: AuthBearer,
|
||||
},
|
||||
{
|
||||
Name: "cloudflare-ai-gateway",
|
||||
EnvVars: []string{"CLOUDFLARE_API_KEY"},
|
||||
// {acct}/{gw} are CLOUDFLARE_ACCOUNT_ID / CLOUDFLARE_GATEWAY_ID.
|
||||
DefaultBaseURL: "https://gateway.ai.cloudflare.com/v1/{account_id}/{gateway_id}/anthropic",
|
||||
Protocol: ProtocolAnthropic,
|
||||
AuthScheme: AuthXAPIKey,
|
||||
},
|
||||
}
|
||||
|
||||
// providerRegistryByName indexes providerRegistry by provider name for O(1)
|
||||
// lookup. Built once at package init.
|
||||
var providerRegistryByName = func() map[string]ProviderSpec {
|
||||
m := make(map[string]ProviderSpec, len(providerRegistry))
|
||||
for _, spec := range providerRegistry {
|
||||
m[spec.Name] = spec
|
||||
}
|
||||
return m
|
||||
}()
|
||||
|
||||
// LookupProviderSpec returns the ProviderSpec for a provider name and whether it
|
||||
// is a known built-in provider. The returned spec is a copy; mutating its slice
|
||||
// or map fields is discouraged as they are shared with the registry.
|
||||
func LookupProviderSpec(name string) (ProviderSpec, bool) {
|
||||
spec, ok := providerRegistryByName[name]
|
||||
return spec, ok
|
||||
}
|
||||
|
||||
// ProviderSpecs returns all built-in provider specs in registry (display) order.
|
||||
// Callers must not mutate the returned specs' slice/map fields.
|
||||
func ProviderSpecs() []ProviderSpec {
|
||||
out := make([]ProviderSpec, len(providerRegistry))
|
||||
copy(out, providerRegistry)
|
||||
return out
|
||||
}
|
||||
|
||||
// ProviderNames returns all built-in provider names in registry order.
|
||||
func ProviderNames() []string {
|
||||
out := make([]string, len(providerRegistry))
|
||||
for i, spec := range providerRegistry {
|
||||
out[i] = spec.Name
|
||||
}
|
||||
return out
|
||||
}
|
||||
@@ -0,0 +1,164 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"sort"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestLookupProviderSpec_Hit(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("deepseek")
|
||||
if !ok {
|
||||
t.Fatalf("LookupProviderSpec(deepseek): expected hit, got miss")
|
||||
}
|
||||
if spec.Name != "deepseek" {
|
||||
t.Errorf("Name = %q, want deepseek", spec.Name)
|
||||
}
|
||||
if spec.DefaultBaseURL != "https://api.deepseek.com" {
|
||||
t.Errorf("DefaultBaseURL = %q, want https://api.deepseek.com", spec.DefaultBaseURL)
|
||||
}
|
||||
if spec.Protocol != ProtocolOpenAIResponses {
|
||||
t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAIResponses)
|
||||
}
|
||||
if len(spec.EnvVars) != 1 || spec.EnvVars[0] != "DEEPSEEK_API_KEY" {
|
||||
t.Errorf("EnvVars = %v, want [DEEPSEEK_API_KEY]", spec.EnvVars)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLookupProviderSpec_Miss(t *testing.T) {
|
||||
if _, ok := LookupProviderSpec("does-not-exist"); ok {
|
||||
t.Errorf("LookupProviderSpec(does-not-exist): expected miss, got hit")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicEnvVarOrder(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("anthropic")
|
||||
if !ok {
|
||||
t.Fatal("LookupProviderSpec(anthropic): expected hit")
|
||||
}
|
||||
want := []string{"ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY", "CLAUDE_API_KEY"}
|
||||
if len(spec.EnvVars) != len(want) {
|
||||
t.Fatalf("EnvVars = %v, want %v", spec.EnvVars, want)
|
||||
}
|
||||
for i := range want {
|
||||
if spec.EnvVars[i] != want[i] {
|
||||
t.Errorf("EnvVars[%d] = %q, want %q (OAuth must be first)", i, spec.EnvVars[i], want[i])
|
||||
}
|
||||
}
|
||||
if spec.Protocol != ProtocolAnthropic {
|
||||
t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolAnthropic)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHuggingfaceEnvVar(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("huggingface")
|
||||
if !ok {
|
||||
t.Fatal("LookupProviderSpec(huggingface): expected hit")
|
||||
}
|
||||
if len(spec.EnvVars) != 1 || spec.EnvVars[0] != "HF_TOKEN" {
|
||||
t.Errorf("EnvVars = %v, want [HF_TOKEN]", spec.EnvVars)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChineseCloudProviders(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
envVars []string
|
||||
baseURL string
|
||||
}{
|
||||
{"qianfan", []string{"QIANFAN_API_KEY"}, "https://qianfan.baidubce.com/v2"},
|
||||
{"volcengine", []string{"ARK_API_KEY", "VOLCENGINE_API_KEY"}, "https://ark.cn-beijing.volces.com/api/v3"},
|
||||
{"dashscope", []string{"DASHSCOPE_API_KEY"}, "https://dashscope.aliyuncs.com/compatible-mode/v1"},
|
||||
{"hunyuan", []string{"HUNYUAN_API_KEY"}, "https://api.hunyuan.cloud.tencent.com/v1"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec(tc.name)
|
||||
if !ok {
|
||||
t.Fatalf("LookupProviderSpec(%q): expected hit", tc.name)
|
||||
}
|
||||
if spec.Name != tc.name {
|
||||
t.Errorf("Name = %q, want %q", spec.Name, tc.name)
|
||||
}
|
||||
if spec.DefaultBaseURL != tc.baseURL {
|
||||
t.Errorf("DefaultBaseURL = %q, want %q", spec.DefaultBaseURL, tc.baseURL)
|
||||
}
|
||||
if spec.Protocol != ProtocolOpenAI {
|
||||
t.Errorf("Protocol = %q, want %q", spec.Protocol, ProtocolOpenAI)
|
||||
}
|
||||
if spec.AuthScheme != AuthBearer {
|
||||
t.Errorf("AuthScheme = %q, want %q", spec.AuthScheme, AuthBearer)
|
||||
}
|
||||
if len(spec.EnvVars) != len(tc.envVars) {
|
||||
t.Fatalf("EnvVars = %v, want %v", spec.EnvVars, tc.envVars)
|
||||
}
|
||||
for i := range tc.envVars {
|
||||
if spec.EnvVars[i] != tc.envVars[i] {
|
||||
t.Errorf("EnvVars[%d] = %q, want %q", i, spec.EnvVars[i], tc.envVars[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAzureBaseURLEnvVars(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("azure-openai-responses")
|
||||
if !ok {
|
||||
t.Fatal("LookupProviderSpec(azure-openai-responses): expected hit")
|
||||
}
|
||||
found := false
|
||||
for _, v := range spec.BaseURLEnvVars {
|
||||
if v == "AZURE_OPENAI_BASE_URL" {
|
||||
found = true
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Errorf("BaseURLEnvVars = %v, want to contain AZURE_OPENAI_BASE_URL", spec.BaseURLEnvVars)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegistryContainsAllExpectedProviders(t *testing.T) {
|
||||
expected := []string{
|
||||
"anthropic", "openai", "ant-ling", "deepseek", "nvidia", "google",
|
||||
"groq", "cerebras", "xai", "openrouter", "vercel-ai-gateway", "zai",
|
||||
"zai-coding-cn", "mistral", "minimax", "minimax-cn", "moonshotai",
|
||||
"moonshotai-cn", "huggingface", "fireworks", "together", "opencode",
|
||||
"opencode-go", "kimi-coding", "xiaomi", "xiaomi-token-plan-cn",
|
||||
"xiaomi-token-plan-ams", "xiaomi-token-plan-sgp",
|
||||
"qianfan", "volcengine", "dashscope", "hunyuan",
|
||||
"azure-openai-responses", "amazon-bedrock", "google-vertex",
|
||||
"cloudflare-workers-ai", "cloudflare-ai-gateway",
|
||||
}
|
||||
for _, name := range expected {
|
||||
if _, ok := LookupProviderSpec(name); !ok {
|
||||
t.Errorf("registry missing expected provider %q", name)
|
||||
}
|
||||
}
|
||||
names := ProviderNames()
|
||||
if len(names) != len(expected) {
|
||||
t.Errorf("registry has %d providers, want %d", len(names), len(expected))
|
||||
}
|
||||
|
||||
// Every spec must have a name, at least one env var, and a valid protocol.
|
||||
for _, spec := range ProviderSpecs() {
|
||||
if spec.Name == "" {
|
||||
t.Error("found spec with empty Name")
|
||||
}
|
||||
if len(spec.EnvVars) == 0 {
|
||||
t.Errorf("provider %q has no EnvVars", spec.Name)
|
||||
}
|
||||
if spec.Protocol != ProtocolOpenAI && spec.Protocol != ProtocolAnthropic && spec.Protocol != ProtocolOpenAIResponses {
|
||||
t.Errorf("provider %q has invalid Protocol %q", spec.Name, spec.Protocol)
|
||||
}
|
||||
}
|
||||
|
||||
// No duplicate provider names.
|
||||
seen := make(map[string]bool, len(names))
|
||||
dupSorted := append([]string(nil), names...)
|
||||
sort.Strings(dupSorted)
|
||||
for _, n := range dupSorted {
|
||||
if seen[n] {
|
||||
t.Errorf("duplicate provider name %q in registry", n)
|
||||
}
|
||||
seen[n] = true
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,223 @@
|
||||
package provider
|
||||
|
||||
// Provider resolution moved here from cmd/pigo (US-004, #361): mapping a model
|
||||
// id / --provider / --protocol selection to a concrete wire driver, plus the
|
||||
// base-url override precedence. Environment lookups are injected as an
|
||||
// env func(string) string so callers (and tests) control the environment
|
||||
// instead of reaching into the process env directly.
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/smallnest/pigo/internal/cli/config"
|
||||
)
|
||||
|
||||
// ResolveProvider maps a model id to a built-in provider. An explicit
|
||||
// --provider name wins over every other rule: it selects a built-in provider
|
||||
// from the registry and constructs the matching wire driver (see
|
||||
// ResolveNamedProvider). When provider is empty, protocol and model-id
|
||||
// heuristics apply as before.
|
||||
//
|
||||
// When protocol is a non-empty explicit selection ("openai" or "anthropic") it
|
||||
// wins over the model-id heuristics: the provider is built directly for that
|
||||
// wire format against baseURL, which is how a user points pigo at a self-hosted
|
||||
// or third-party endpoint and says which protocol it speaks. An "anthropic"
|
||||
// selection with no baseURL targets the public Anthropic API.
|
||||
//
|
||||
// When protocol is empty, resolution falls back to model-id heuristics:
|
||||
//
|
||||
// 1. If the id is in the preset catalog, use its declared provider (this is how
|
||||
// OpenRouter/NVIDIA/Ollama presets pick the right gateway).
|
||||
// 2. An "ollama/" prefix (or a base URL on the Ollama port) → local Ollama.
|
||||
// 3. An "nvidia/" prefix → NVIDIA NIM (strips the prefix for the wire id).
|
||||
// 4. Model-name inference: with no --base-url, a well-known model-name prefix
|
||||
// (e.g. "claude-*", "deepseek-*") selects its first-party built-in provider
|
||||
// via ResolveNamedProvider (see InferProviderFromModel).
|
||||
// 5. Everything else → OpenRouter, the reference OpenAI-compatible gateway.
|
||||
//
|
||||
// An unknown protocol value is an error, surfaced to the caller for exit-code
|
||||
// mapping rather than silently falling back.
|
||||
func ResolveProvider(model, baseURL, protocol, providerName string, env func(string) string) (Provider, string, error) {
|
||||
// Explicit --provider selects a built-in provider from the registry and
|
||||
// wins over both --protocol inference and model-id heuristics.
|
||||
if strings.TrimSpace(providerName) != "" {
|
||||
return ResolveNamedProvider(providerName, model, baseURL, protocol, env)
|
||||
}
|
||||
|
||||
// 0. Explicit protocol selection wins over every heuristic. Normalize the
|
||||
// surface value first so "openai" and "openai/chat" collapse to the same
|
||||
// Chat Completions selector and "openai/resp_api" routes to the Responses
|
||||
// driver; an unknown value surfaces as an error for exit-code mapping.
|
||||
canonical, err := NormalizeProtocol(protocol)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
switch canonical {
|
||||
case ProtocolOpenAI:
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
return nil, "", fmt.Errorf("--protocol openai requires --base-url")
|
||||
}
|
||||
return NewOpenAICompatibleProvider(baseURL, []Model{{Provider: "openai", ID: model, SupportsImages: true}}), "openai", nil
|
||||
case ProtocolOpenAIResponses:
|
||||
// The Responses driver has no public default endpoint here: unlike the
|
||||
// anthropic path (which targets the public API), resp_api mirrors the
|
||||
// Chat Completions requirement and demands an explicit --base-url.
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
return nil, "", fmt.Errorf("--protocol openai/resp_api requires --base-url")
|
||||
}
|
||||
return NewOpenAIResponsesProvider("openai", baseURL, []Model{{Provider: "openai", ID: model, SupportsImages: true}}), "openai", nil
|
||||
case ProtocolAnthropic:
|
||||
return NewAnthropicProvider(baseURL, []Model{{Provider: "anthropic", ID: model, SupportsImages: true}}), "anthropic", nil
|
||||
case "":
|
||||
// fall through to heuristic resolution
|
||||
}
|
||||
|
||||
// 1. Preset catalog wins: a curated id knows its own provider.
|
||||
if p, ok := LookupPreset(model); ok {
|
||||
switch p.Provider {
|
||||
case "nvidia":
|
||||
return NewNvidiaProvider(baseURL, []Model{{Provider: "nvidia", ID: model, SupportsImages: true}}), "nvidia", nil
|
||||
case "ollama":
|
||||
id := strings.TrimPrefix(model, "ollama/")
|
||||
return NewOllamaProvider(baseURL, []Model{{Provider: "ollama", ID: id, SupportsImages: true}}), "ollama", nil
|
||||
case "", "openrouter":
|
||||
return NewOpenRouterProvider(baseURL, []Model{{Provider: "openrouter", ID: model, SupportsImages: true}}), "openrouter", nil
|
||||
default:
|
||||
// Any other preset provider is a named built-in (e.g. deepseek,
|
||||
// qianfan, dashscope): build it from the registry so the correct
|
||||
// base URL, protocol, and API-key env var are used — not OpenRouter's.
|
||||
return ResolveNamedProvider(p.Provider, model, baseURL, protocol, env)
|
||||
}
|
||||
}
|
||||
|
||||
// 2. Local Ollama by prefix or port.
|
||||
if strings.HasPrefix(model, "ollama/") || strings.Contains(baseURL, "11434") {
|
||||
id := strings.TrimPrefix(model, "ollama/")
|
||||
return NewOllamaProvider(baseURL, []Model{{Provider: "ollama", ID: id, SupportsImages: true}}), "ollama", nil
|
||||
}
|
||||
// 3. NVIDIA NIM by prefix.
|
||||
if strings.HasPrefix(model, "nvidia/") {
|
||||
id := strings.TrimPrefix(model, "nvidia/")
|
||||
return NewNvidiaProvider(baseURL, []Model{{Provider: "nvidia", ID: id, SupportsImages: true}}), "nvidia", nil
|
||||
}
|
||||
// 4. Model-name inference: with no --provider/--protocol (both empty here) and
|
||||
// no --base-url, guess the provider from the model name's well-known prefix
|
||||
// (e.g. "claude-*" → anthropic, "deepseek-*" → deepseek). A confident hit is
|
||||
// routed through ResolveNamedProvider so the provider's registry protocol,
|
||||
// default base URL, and API-key env var are used. A --base-url is treated as
|
||||
// a custom-endpoint signal that should not be second-guessed, so inference is
|
||||
// skipped when one is given. Ambiguous/unknown names fall through to (5).
|
||||
if strings.TrimSpace(baseURL) == "" {
|
||||
if name, ok := InferProviderFromModel(model); ok {
|
||||
return ResolveNamedProvider(name, model, baseURL, protocol, env)
|
||||
}
|
||||
}
|
||||
// 5. Default: OpenRouter.
|
||||
return NewOpenRouterProvider(baseURL, []Model{{Provider: "openrouter", ID: model, SupportsImages: true}}), "openrouter", nil
|
||||
}
|
||||
|
||||
// ResolveNamedProvider builds the driver for an explicit --provider selection.
|
||||
// It looks the name up in the built-in registry and constructs the wire driver
|
||||
// matching the spec's Protocol: "openai" → an OpenAI-compatible (Bearer) driver,
|
||||
// "anthropic" → an Anthropic-Messages driver. The base URL follows the override
|
||||
// precedence in ResolveBaseURL (--base-url > provider-specific env > generic
|
||||
// <PROVIDER>_BASE_URL > spec default). The returned provider-name string is the
|
||||
// spec name, so downstream API-key resolution reads the provider's own env var
|
||||
// (spec.EnvVars).
|
||||
//
|
||||
// Special providers with bespoke auth (azure/bedrock/vertex/cloudflare —
|
||||
// AuthScheme aws/azure/special, or the cloudflare-* names) are routed to
|
||||
// ResolveSpecialProvider, which validates their required env vars and composes
|
||||
// the concrete endpoint (node #188).
|
||||
func ResolveNamedProvider(name, model, baseURL, protocol string, env func(string) string) (Provider, string, error) {
|
||||
spec, ok := LookupProviderSpec(name)
|
||||
if !ok {
|
||||
return nil, "", fmt.Errorf("unknown --provider %q (available: %s)", name, strings.Join(ProviderNames(), ", "))
|
||||
}
|
||||
// A concurrently-set --protocol must agree with the provider's own protocol;
|
||||
// an incompatible pair is a user error naming both flags. Normalize the raw
|
||||
// value first so aliases (e.g. "openai/chat" for an "openai" spec) don't
|
||||
// falsely conflict, and a genuine typo surfaces as a clear "unknown --protocol"
|
||||
// error rather than a misleading conflict message.
|
||||
if strings.TrimSpace(protocol) != "" {
|
||||
canonical, err := NormalizeProtocol(protocol)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
if canonical != spec.Protocol {
|
||||
return nil, "", fmt.Errorf("--provider %q speaks the %q protocol, which conflicts with --protocol %q; drop --protocol or set it to %q", name, spec.Protocol, protocol, spec.Protocol)
|
||||
}
|
||||
}
|
||||
// Special-auth providers (Azure / Bedrock / Vertex / Cloudflare) compose
|
||||
// their endpoint from several env vars and/or need non-standard credential
|
||||
// validation, so route them to the dedicated resolver (US-007 / node #188).
|
||||
// It performs its own base-URL composition (honoring the --base-url override)
|
||||
// and returns a clear error naming any absent required env var.
|
||||
if IsSpecialAuthProvider(spec) {
|
||||
p, err := ResolveSpecialProvider(spec, model, baseURL, env)
|
||||
if err != nil {
|
||||
return nil, "", err
|
||||
}
|
||||
return p, spec.Name, nil
|
||||
}
|
||||
// Base-URL precedence (US-004 / FR-8, FR-9): --base-url flag > provider-
|
||||
// specific base-url env var(s) > generic <PROVIDER>_BASE_URL > spec default.
|
||||
url := ResolveBaseURL(spec, baseURL, env)
|
||||
models := []Model{{Provider: spec.Name, ID: model, SupportsImages: true}}
|
||||
// Note: spec.ExtraHeaders would be attached here, but the exported generic
|
||||
// constructors do not yet accept custom headers; all built-in specs currently
|
||||
// carry no ExtraHeaders, so this is a no-op today (refined alongside #188).
|
||||
switch spec.Protocol {
|
||||
case ProtocolAnthropic:
|
||||
// Auth header follows the spec's AuthScheme (x-api-key + anthropic-version
|
||||
// for anthropic/minimax/minimax-cn; Bearer for any anthropic-protocol
|
||||
// gateway that authenticates with a plain bearer token). The driver name is
|
||||
// the spec name so errors reference the selected provider.
|
||||
return NewAnthropicProtocolProvider(spec.Name, url, spec.AuthScheme, models), spec.Name, nil
|
||||
case ProtocolOpenAI:
|
||||
return NewOpenAICompatibleProvider(url, models), spec.Name, nil
|
||||
case ProtocolOpenAIResponses:
|
||||
return NewOpenAIResponsesProvider(spec.Name, url, models), spec.Name, nil
|
||||
default:
|
||||
// The registry only ever stores openai/openai-resp/anthropic; guard anyway
|
||||
// so an unexpected value is a clear error rather than a nil provider.
|
||||
return nil, "", fmt.Errorf("--provider %q has unsupported protocol %q", name, spec.Protocol)
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveBaseURL determines the effective base URL for a selected provider,
|
||||
// applying the base_url override precedence (US-004 / FR-8, FR-9). The first
|
||||
// non-empty source wins, in this order:
|
||||
//
|
||||
// 1. flagBaseURL — the explicit --base-url/-u flag (highest).
|
||||
// 2. provider-specific base-url env var(s) from spec.BaseURLEnvVars, in the
|
||||
// order the registry declares them (e.g. AZURE_OPENAI_BASE_URL).
|
||||
// 3. the generic <PROVIDER>_BASE_URL env var, where <PROVIDER> is the provider
|
||||
// name uppercased with '-' rewritten to '_' (e.g. zai-coding-cn →
|
||||
// ZAI_CODING_CN_BASE_URL).
|
||||
// 4. spec.DefaultBaseURL — the registry default (lowest).
|
||||
//
|
||||
// Values are trimmed of surrounding whitespace before the non-empty check, so a
|
||||
// whitespace-only env var does not shadow a lower-precedence source. Environment
|
||||
// lookups go through the injected env func so callers control the environment.
|
||||
func ResolveBaseURL(spec ProviderSpec, flagBaseURL string, env func(string) string) string {
|
||||
// 1. Explicit flag wins over every env-var convention.
|
||||
if v := strings.TrimSpace(flagBaseURL); v != "" {
|
||||
return v
|
||||
}
|
||||
// 2. Provider-specific override env vars, in registry precedence order.
|
||||
for _, name := range spec.BaseURLEnvVars {
|
||||
if v := strings.TrimSpace(env(name)); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
// 3. Generic <PROVIDER>_BASE_URL convention.
|
||||
if envName := config.GenericBaseURLEnvVar(spec.Name); envName != "" {
|
||||
if v := strings.TrimSpace(env(envName)); v != "" {
|
||||
return v
|
||||
}
|
||||
}
|
||||
// 4. Registry default.
|
||||
return spec.DefaultBaseURL
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
package provider
|
||||
|
||||
// Tests for provider resolution moved from cmd/pigo (US-004, #361): ResolveProvider
|
||||
// maps a model id to the right gateway (preset catalog first, then prefix rules,
|
||||
// then OpenRouter default), and ResolveBaseURL applies the base-url override
|
||||
// precedence. Environment lookups are injected via os.Getenv here.
|
||||
|
||||
import (
|
||||
"os"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// TestResolveProviderPresetCatalog verifies a preset id resolves to its declared
|
||||
// provider (NVIDIA and Ollama presets do not fall through to OpenRouter).
|
||||
func TestResolveProviderPresetCatalog(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
wantName string
|
||||
}{
|
||||
{"meta/llama-3.3-70b-instruct", "nvidia"}, // NVIDIA preset
|
||||
{"ollama/llama3.3", "ollama"}, // Ollama preset
|
||||
{"openai/gpt-4o", "openrouter"}, // OpenRouter preset
|
||||
{"anthropic/claude-3.5-sonnet", "openrouter"}, // OpenRouter preset
|
||||
}
|
||||
for _, c := range cases {
|
||||
_, name, err := ResolveProvider(c.model, "", "", "", os.Getenv)
|
||||
if err != nil {
|
||||
t.Errorf("ResolveProvider(%q) error: %v", c.model, err)
|
||||
continue
|
||||
}
|
||||
if name != c.wantName {
|
||||
t.Errorf("ResolveProvider(%q) = %q, want %q", c.model, name, c.wantName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderPrefixAndDefault verifies the prefix rules and the
|
||||
// OpenRouter default for ids not in the catalog.
|
||||
func TestResolveProviderPrefixAndDefault(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
baseURL string
|
||||
wantName string
|
||||
}{
|
||||
{"ollama/some-local-model", "", "ollama"}, // ollama/ prefix
|
||||
{"nvidia/some-nim-model", "", "nvidia"}, // nvidia/ prefix
|
||||
{"some-unknown-model", "", "openrouter"}, // default
|
||||
{"m", "http://host:11434/v1", "ollama"}, // ollama port
|
||||
}
|
||||
for _, c := range cases {
|
||||
_, name, err := ResolveProvider(c.model, c.baseURL, "", "", os.Getenv)
|
||||
if err != nil {
|
||||
t.Errorf("ResolveProvider(%q) error: %v", c.model, err)
|
||||
continue
|
||||
}
|
||||
if name != c.wantName {
|
||||
t.Errorf("ResolveProvider(%q, %q) = %q, want %q", c.model, c.baseURL, name, c.wantName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderExplicitProtocol verifies an explicit --protocol wins over
|
||||
// model-id heuristics: openai (with base-url) and anthropic select the matching
|
||||
// wire driver, an empty base-url for openai errors, and an unknown protocol
|
||||
// errors instead of silently falling back.
|
||||
func TestResolveProviderExplicitProtocol(t *testing.T) {
|
||||
// openai protocol → "openai" provider name, requires base-url.
|
||||
if _, name, err := ResolveProvider("any-model", "https://example.com/v1", "openai", "", os.Getenv); err != nil || name != "openai" {
|
||||
t.Errorf("protocol=openai = (%q, %v), want (openai, nil)", name, err)
|
||||
}
|
||||
if _, _, err := ResolveProvider("any-model", "", "openai", "", os.Getenv); err == nil {
|
||||
t.Error("protocol=openai with no base-url should error")
|
||||
}
|
||||
// anthropic protocol → "anthropic" provider name, base-url optional (defaults).
|
||||
if _, name, err := ResolveProvider("claude-x", "", "anthropic", "", os.Getenv); err != nil || name != "anthropic" {
|
||||
t.Errorf("protocol=anthropic = (%q, %v), want (anthropic, nil)", name, err)
|
||||
}
|
||||
// Unknown protocol errors rather than falling back to a heuristic.
|
||||
if _, _, err := ResolveProvider("any-model", "", "grpc", "", os.Getenv); err == nil {
|
||||
t.Error("unknown protocol should error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderResponsesProtocol verifies the openai/resp_api selector
|
||||
// routes to the Responses driver (against an explicit base-url), that the
|
||||
// "openai/chat" alias resolves identically to "openai", and that resp_api with
|
||||
// no base-url errors like the plain openai path (mirroring the base-url
|
||||
// requirement rather than defaulting to a public endpoint).
|
||||
func TestResolveProviderResponsesProtocol(t *testing.T) {
|
||||
// openai/resp_api → "openai" provider name, backed by the Responses driver.
|
||||
p, name, err := ResolveProvider("any-model", "https://example.com/v1", "openai/resp_api", "", os.Getenv)
|
||||
if err != nil || name != "openai" {
|
||||
t.Fatalf("protocol=openai/resp_api = (%q, %v), want (openai, nil)", name, err)
|
||||
}
|
||||
if _, ok := p.(*responsesDriver); !ok {
|
||||
t.Errorf("protocol=openai/resp_api built %T, want *responsesDriver", p)
|
||||
}
|
||||
// resp_api with no base-url errors, mirroring the openai requirement.
|
||||
if _, _, err := ResolveProvider("any-model", "", "openai/resp_api", "", os.Getenv); err == nil {
|
||||
t.Error("protocol=openai/resp_api with no base-url should error")
|
||||
}
|
||||
// "openai/chat" is an alias of "openai": same driver, same base-url rule.
|
||||
p, name, err = ResolveProvider("any-model", "https://example.com/v1", "openai/chat", "", os.Getenv)
|
||||
if err != nil || name != "openai" {
|
||||
t.Fatalf("protocol=openai/chat = (%q, %v), want (openai, nil)", name, err)
|
||||
}
|
||||
if _, ok := p.(*responsesDriver); ok {
|
||||
t.Error("protocol=openai/chat should build the Chat Completions driver, not *responsesDriver")
|
||||
}
|
||||
if _, _, err := ResolveProvider("any-model", "", "openai/chat", "", os.Getenv); err == nil {
|
||||
t.Error("protocol=openai/chat with no base-url should error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderExplicitProvider verifies that --provider selects a
|
||||
// built-in provider from the registry: the returned provider-name is the spec
|
||||
// name (so key resolution reads the right env var), an OpenAI-protocol provider
|
||||
// (deepseek) and an Anthropic-protocol provider (minimax) both resolve, an
|
||||
// incompatible --protocol is a conflict error naming both flags, and an unknown
|
||||
// provider name errors while listing the available names.
|
||||
func TestResolveProviderExplicitProvider(t *testing.T) {
|
||||
// OpenAI-protocol provider: returns its own name for key lookup.
|
||||
if _, name, err := ResolveProvider("deepseek-chat", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" {
|
||||
t.Errorf("provider=deepseek = (%q, %v), want (deepseek, nil)", name, err)
|
||||
}
|
||||
// Anthropic-protocol provider.
|
||||
if _, name, err := ResolveProvider("MiniMax-M2", "", "", "minimax", os.Getenv); err != nil || name != "minimax" {
|
||||
t.Errorf("provider=minimax = (%q, %v), want (minimax, nil)", name, err)
|
||||
}
|
||||
// A matching --protocol is not a conflict (deepseek speaks openai/resp_api).
|
||||
if _, name, err := ResolveProvider("deepseek-chat", "", "openai/resp_api", "deepseek", os.Getenv); err != nil || name != "deepseek" {
|
||||
t.Errorf("provider=deepseek + protocol=openai/resp_api = (%q, %v), want (deepseek, nil)", name, err)
|
||||
}
|
||||
// --provider wins over model-id heuristics: an ollama/-prefixed id still
|
||||
// resolves to the named provider, not local Ollama.
|
||||
if _, name, err := ResolveProvider("ollama/x", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" {
|
||||
t.Errorf("provider=deepseek with ollama/ model = (%q, %v), want (deepseek, nil)", name, err)
|
||||
}
|
||||
// --base-url overrides the spec default without changing the provider name.
|
||||
if _, name, err := ResolveProvider("deepseek-chat", "https://proxy.local/v1", "", "deepseek", os.Getenv); err != nil || name != "deepseek" {
|
||||
t.Errorf("provider=deepseek + base-url = (%q, %v), want (deepseek, nil)", name, err)
|
||||
}
|
||||
// Conflict: minimax speaks anthropic; forcing --protocol openai errors and
|
||||
// names both flags.
|
||||
_, _, err := ResolveProvider("MiniMax-M2", "", "openai", "minimax", os.Getenv)
|
||||
if err == nil {
|
||||
t.Fatal("provider=minimax + protocol=openai should conflict")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "--provider") || !strings.Contains(err.Error(), "--protocol") {
|
||||
t.Errorf("conflict error should name both flags, got: %v", err)
|
||||
}
|
||||
// Unknown provider errors and lists available names.
|
||||
_, _, err = ResolveProvider("m", "", "", "no-such-provider", os.Getenv)
|
||||
if err == nil {
|
||||
t.Fatal("unknown provider should error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "deepseek") {
|
||||
t.Errorf("unknown-provider error should list available names, got: %v", err)
|
||||
}
|
||||
// An invalid --protocol paired with a named provider surfaces the clear
|
||||
// "unknown --protocol" error (listing the accepted set) rather than a
|
||||
// misleading conflict message.
|
||||
_, _, err = ResolveProvider("deepseek-chat", "", "openai_api", "deepseek", os.Getenv)
|
||||
if err == nil {
|
||||
t.Fatal("provider=deepseek + protocol=openai_api should error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "unknown --protocol") {
|
||||
t.Errorf("invalid --protocol should surface the unknown-protocol error, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderCNPresets verifies the Chinese-cloud preset ids route to
|
||||
// their own provider (not the OpenRouter default) via the LookupPreset branch.
|
||||
func TestResolveProviderCNPresets(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
wantName string
|
||||
}{
|
||||
{"ernie-4.5-turbo-32k", "qianfan"},
|
||||
{"doubao-seed-1-6", "volcengine"},
|
||||
{"qwen-max", "dashscope"},
|
||||
{"hunyuan-turbos-latest", "hunyuan"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
_, name, err := ResolveProvider(c.model, "", "", "", os.Getenv)
|
||||
if err != nil {
|
||||
t.Errorf("ResolveProvider(%q) error: %v", c.model, err)
|
||||
continue
|
||||
}
|
||||
if name != c.wantName {
|
||||
t.Errorf("ResolveProvider(%q) = %q, want %q", c.model, name, c.wantName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderCNExplicit verifies --provider selects the CN providers
|
||||
// directly and that --base-url overrides without changing the provider name.
|
||||
func TestResolveProviderCNExplicit(t *testing.T) {
|
||||
for _, name := range []string{"qianfan", "volcengine", "dashscope", "hunyuan"} {
|
||||
if _, got, err := ResolveProvider("some-model", "", "", name, os.Getenv); err != nil || got != name {
|
||||
t.Errorf("provider=%s = (%q, %v), want (%s, nil)", name, got, err, name)
|
||||
}
|
||||
if _, got, err := ResolveProvider("some-model", "https://proxy.local/v1", "", name, os.Getenv); err != nil || got != name {
|
||||
t.Errorf("provider=%s + base-url = (%q, %v), want (%s, nil)", name, got, err, name)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderModelNameInference verifies model-name inference (Issue
|
||||
// #235): with only --model given, a bare model name whose prefix identifies a
|
||||
// single provider resolves to that provider — NOT the OpenRouter default.
|
||||
func TestResolveProviderModelNameInference(t *testing.T) {
|
||||
cases := []struct {
|
||||
model string
|
||||
wantName string
|
||||
}{
|
||||
{"claude-opus-4-8", "anthropic"},
|
||||
{"deepseek-chat", "deepseek"},
|
||||
{"gpt-4.1", "openai"},
|
||||
{"gemini-3-pro", "google"},
|
||||
{"grok-5", "xai"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if _, name, err := ResolveProvider(c.model, "", "", "", os.Getenv); err != nil || name != c.wantName {
|
||||
t.Errorf("ResolveProvider(%q) = (%q, %v), want (%q, nil)", c.model, name, err, c.wantName)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveProviderInferencePrecedence verifies that model-name inference does
|
||||
// not override explicit flags and does not fire when a --base-url is given, and
|
||||
// that unknown/ambiguous names still fall back to OpenRouter.
|
||||
func TestResolveProviderInferencePrecedence(t *testing.T) {
|
||||
// Explicit --provider wins over an inferable model name.
|
||||
if _, name, err := ResolveProvider("claude-opus-4-8", "", "", "deepseek", os.Getenv); err != nil || name != "deepseek" {
|
||||
t.Errorf("provider=deepseek overrides inference = (%q, %v), want (deepseek, nil)", name, err)
|
||||
}
|
||||
// Explicit --protocol wins over an inferable model name.
|
||||
if _, name, err := ResolveProvider("claude-opus-4-8", "https://example.com/v1", "openai", "", os.Getenv); err != nil || name != "openai" {
|
||||
t.Errorf("protocol=openai overrides inference = (%q, %v), want (openai, nil)", name, err)
|
||||
}
|
||||
// A --base-url signals a custom endpoint: inference is skipped, default applies.
|
||||
if _, name, err := ResolveProvider("claude-opus-4-8", "https://gw.local/v1", "", "", os.Getenv); err != nil || name != "openrouter" {
|
||||
t.Errorf("inference skipped with base-url = (%q, %v), want (openrouter, nil)", name, err)
|
||||
}
|
||||
// Ambiguous/unknown names still default to OpenRouter.
|
||||
for _, m := range []string{"llama-3.3-70b", "totally-unknown-model"} {
|
||||
if _, name, err := ResolveProvider(m, "", "", "", os.Getenv); err != nil || name != "openrouter" {
|
||||
t.Errorf("ResolveProvider(%q) = (%q, %v), want (openrouter, nil)", m, name, err)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveBaseURLPrecedence exercises all four precedence levels for a
|
||||
// hyphenated provider (zai-coding-cn → ZAI_CODING_CN_BASE_URL).
|
||||
func TestResolveBaseURLPrecedence(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("zai-coding-cn")
|
||||
if !ok {
|
||||
t.Fatal("expected zai-coding-cn in registry")
|
||||
}
|
||||
if got := ResolveBaseURL(spec, "", os.Getenv); got != spec.DefaultBaseURL {
|
||||
t.Errorf("default: got %q, want %q", got, spec.DefaultBaseURL)
|
||||
}
|
||||
t.Setenv("ZAI_CODING_CN_BASE_URL", "https://generic.example/v4")
|
||||
if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://generic.example/v4" {
|
||||
t.Errorf("generic env: got %q, want %q", got, "https://generic.example/v4")
|
||||
}
|
||||
if got := ResolveBaseURL(spec, "https://flag.example/v4", os.Getenv); got != "https://flag.example/v4" {
|
||||
t.Errorf("flag over generic: got %q, want %q", got, "https://flag.example/v4")
|
||||
}
|
||||
}
|
||||
|
||||
// TestResolveBaseURLProviderSpecificEnv covers a provider that declares a
|
||||
// provider-specific base-url env var (azure), asserting it sits between the flag
|
||||
// and the generic convention in precedence.
|
||||
func TestResolveBaseURLProviderSpecificEnv(t *testing.T) {
|
||||
spec, ok := LookupProviderSpec("azure-openai-responses")
|
||||
if !ok {
|
||||
t.Fatal("expected azure-openai-responses in registry")
|
||||
}
|
||||
if len(spec.BaseURLEnvVars) == 0 {
|
||||
t.Fatal("expected azure-openai-responses to declare BaseURLEnvVars")
|
||||
}
|
||||
t.Setenv("AZURE_OPENAI_BASE_URL", "https://specific.example")
|
||||
if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://specific.example" {
|
||||
t.Errorf("provider-specific env: got %q, want %q", got, "https://specific.example")
|
||||
}
|
||||
t.Setenv("AZURE_OPENAI_RESPONSES_BASE_URL", "https://generic.example")
|
||||
if got := ResolveBaseURL(spec, "", os.Getenv); got != "https://specific.example" {
|
||||
t.Errorf("provider-specific beats generic: got %q, want %q", got, "https://specific.example")
|
||||
}
|
||||
if got := ResolveBaseURL(spec, "https://flag.example", os.Getenv); got != "https://flag.example" {
|
||||
t.Errorf("flag beats provider-specific: got %q, want %q", got, "https://flag.example")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,474 @@
|
||||
// This file implements the OpenAI Responses API driver (US-003, #539), the
|
||||
// backing for --protocol openai/resp_api. Unlike openAICompatDriver (which
|
||||
// hand-rolls the Chat Completions wire format), this driver speaks the
|
||||
// Responses API (POST {base_url}/responses) via the official
|
||||
// github.com/openai/openai-go SDK.
|
||||
//
|
||||
// This milestone covers streaming text: a plain prompt in, assistant text out,
|
||||
// consumed from the Responses SSE stream and mapped into pigo's AssistantMessage
|
||||
// the same way OpenAIDecoder does (API/Provider tags, Usage,
|
||||
// ResponseID/ResponseModel, StopReason=end_turn). Tools (#541) and
|
||||
// images/reasoning (#542) layer on later.
|
||||
//
|
||||
// Failure model (FR-13): only the earliest "cannot build the stream" case
|
||||
// (missing API key) is a returned error. Every runtime failure — including a
|
||||
// non-2xx from the endpoint — rides the returned stream as a terminal
|
||||
// StreamErrorEvent, matching the chat driver's observable behavior.
|
||||
//
|
||||
// Base URL + auth: the SDK client is pointed at the resolved base_url and given
|
||||
// the resolved key via option.WithBaseURL / option.WithAPIKey rather than
|
||||
// reading the environment, so ResolveBaseURL precedence is preserved. A custom
|
||||
// *http.Client may be injected (option.WithHTTPClient) for tests.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"github.com/openai/openai-go"
|
||||
"github.com/openai/openai-go/option"
|
||||
"github.com/openai/openai-go/responses"
|
||||
"github.com/openai/openai-go/shared"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// responsesDriver is the Provider backing --protocol openai/resp_api. It holds
|
||||
// the provider identity, the resolved endpoint, the model catalog, and whether
|
||||
// an API key is required, and builds an openai-go client per request.
|
||||
type responsesDriver struct {
|
||||
name string
|
||||
baseURL string
|
||||
models []Model
|
||||
// requiresAuth reports whether an API key must be present. Public OpenAI /
|
||||
// Azure require it; a local gateway may not.
|
||||
requiresAuth bool
|
||||
// clientOpts are extra SDK options; tests inject option.WithHTTPClient here
|
||||
// to stub the transport.
|
||||
clientOpts []option.RequestOption
|
||||
}
|
||||
|
||||
// NewOpenAIResponsesProvider builds a Responses API provider targeting baseURL.
|
||||
// baseURL must be the fully resolved endpoint (e.g. https://api.openai.com/v1);
|
||||
// the SDK appends the /responses path.
|
||||
func NewOpenAIResponsesProvider(name, baseURL string, models []Model) *responsesDriver {
|
||||
return &responsesDriver{
|
||||
name: name,
|
||||
baseURL: baseURL,
|
||||
models: models,
|
||||
requiresAuth: true,
|
||||
}
|
||||
}
|
||||
|
||||
func (d *responsesDriver) Name() string { return d.name }
|
||||
func (d *responsesDriver) Models() []Model { return d.models }
|
||||
|
||||
// StreamCompletion issues a streaming Responses API call and surfaces the result
|
||||
// on an AssistantMessageEventStream: a start event, incremental text events as
|
||||
// deltas arrive, and a terminal done event carrying the aggregated message.
|
||||
func (d *responsesDriver) StreamCompletion(ctx context.Context, req CompletionRequest) (*AssistantMessageEventStream, error) {
|
||||
if d.requiresAuth && strings.TrimSpace(req.Config.APIKey) == "" {
|
||||
// Early "cannot build the stream": reference the provider, never a value.
|
||||
return nil, fmt.Errorf("%s: missing API key", d.name)
|
||||
}
|
||||
|
||||
opts := make([]option.RequestOption, 0, len(d.clientOpts)+2)
|
||||
if d.baseURL != "" {
|
||||
opts = append(opts, option.WithBaseURL(d.baseURL))
|
||||
}
|
||||
if key := strings.TrimSpace(req.Config.APIKey); key != "" {
|
||||
opts = append(opts, option.WithAPIKey(key))
|
||||
}
|
||||
opts = append(opts, d.clientOpts...)
|
||||
client := openai.NewClient(opts...)
|
||||
|
||||
params := buildResponsesParams(req)
|
||||
|
||||
stream := NewAssistantMessageEventStream(0)
|
||||
go d.pump(ctx, stream, &client, params)
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
// pump consumes the Responses SSE stream and translates events into pigo stream
|
||||
// events. It always closes the stream. Every runtime failure (transport error,
|
||||
// context cancellation, or an upstream error/failed event) becomes a terminal
|
||||
// StreamErrorEvent (dual failure model), not a returned error.
|
||||
//
|
||||
// Incremental text.delta events emit a StreamTextEvent carrying the accumulated
|
||||
// text so far, so the TUI renders tokens as they arrive. The terminal message is
|
||||
// built from the authoritative response.completed payload via mapResponse, so
|
||||
// the final aggregation matches the non-streamed result exactly. If no completed
|
||||
// event arrives (a truncated stream that still ended cleanly), the accumulated
|
||||
// delta text is used as a fallback.
|
||||
func (d *responsesDriver) pump(ctx context.Context, stream *AssistantMessageEventStream, client *openai.Client, params responses.ResponseNewParams) {
|
||||
defer stream.Close()
|
||||
|
||||
if err := stream.Emit(ctx, StreamStartEvent{Partial: d.newPartial()}); err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
sse := client.Responses.NewStreaming(ctx, params)
|
||||
defer sse.Close()
|
||||
|
||||
var text strings.Builder
|
||||
var thinking strings.Builder
|
||||
var toolCalls []agentcore.ToolCallContent
|
||||
var completed *responses.Response
|
||||
for sse.Next() {
|
||||
if ctx.Err() != nil {
|
||||
d.emitError(stream, ctx.Err())
|
||||
return
|
||||
}
|
||||
switch variant := sse.Current().AsAny().(type) {
|
||||
case responses.ResponseTextDeltaEvent:
|
||||
text.WriteString(variant.Delta)
|
||||
partial := d.buildPartial(thinking.String(), text.String(), toolCalls)
|
||||
if err := stream.Emit(ctx, StreamTextEvent{Partial: partial}); err != nil {
|
||||
return
|
||||
}
|
||||
case responses.ResponseReasoningSummaryTextDeltaEvent:
|
||||
// The model's reasoning summary streams as its own text deltas, distinct
|
||||
// from the answer text; accumulate it into a thinking block so the TUI
|
||||
// renders reasoning the same way the chat driver does.
|
||||
thinking.WriteString(variant.Delta)
|
||||
partial := d.buildPartial(thinking.String(), text.String(), toolCalls)
|
||||
if err := stream.Emit(ctx, StreamThinkingEvent{Partial: partial}); err != nil {
|
||||
return
|
||||
}
|
||||
case responses.ResponseOutputItemDoneEvent:
|
||||
// A finalized function_call item carries the model's tool request
|
||||
// (name + arguments + call_id). Accumulate it and surface a tool-call
|
||||
// partial so the TUI can show the pending call before the run ends.
|
||||
if fc := variant.Item.AsFunctionCall(); fc.Type == "function_call" {
|
||||
toolCalls = append(toolCalls, toolCallContent(fc))
|
||||
partial := d.buildPartial(thinking.String(), text.String(), toolCalls)
|
||||
if err := stream.Emit(ctx, StreamToolCallEvent{Partial: partial}); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
case responses.ResponseCompletedEvent:
|
||||
r := variant.Response
|
||||
completed = &r
|
||||
case responses.ResponseFailedEvent:
|
||||
d.emitError(stream, fmt.Errorf("response failed"))
|
||||
return
|
||||
case responses.ResponseErrorEvent:
|
||||
d.emitError(stream, fmt.Errorf("%s", variant.Message))
|
||||
return
|
||||
}
|
||||
}
|
||||
if err := sse.Err(); err != nil {
|
||||
d.emitError(stream, err)
|
||||
return
|
||||
}
|
||||
|
||||
var msg agentcore.AssistantMessage
|
||||
if completed != nil {
|
||||
msg = d.mapResponse(completed)
|
||||
} else {
|
||||
msg = d.buildPartial(thinking.String(), text.String(), toolCalls)
|
||||
msg.StopReason = agentcore.StopReasonEndTurn
|
||||
if len(toolCalls) > 0 {
|
||||
msg.StopReason = agentcore.StopReasonToolUse
|
||||
}
|
||||
}
|
||||
stream.Emit(ctx, StreamDoneEvent{Message: msg})
|
||||
}
|
||||
|
||||
// buildPartial assembles a cumulative snapshot message for a streaming partial:
|
||||
// an optional thinking block (reasoning summary so far), the accumulated answer
|
||||
// text, then any finalized tool calls — in the order the TUI should render them.
|
||||
// All four emit sites in pump build partials through this one helper so they
|
||||
// can't diverge.
|
||||
func (d *responsesDriver) buildPartial(thinking, text string, toolCalls []agentcore.ToolCallContent) agentcore.AssistantMessage {
|
||||
msg := d.newPartial()
|
||||
if thinking != "" {
|
||||
msg.Content = append(msg.Content, agentcore.NewThinkingContent(thinking))
|
||||
}
|
||||
if text != "" {
|
||||
msg.Content = append(msg.Content, agentcore.NewTextContent(text))
|
||||
}
|
||||
msg.Content = appendToolCalls(msg.Content, toolCalls)
|
||||
return msg
|
||||
}
|
||||
|
||||
// emitError emits a terminal StreamErrorEvent tagged for this provider. Uses a
|
||||
// background context so the emit isn't dropped when ctx is already cancelled.
|
||||
func (d *responsesDriver) emitError(stream *AssistantMessageEventStream, err error) {
|
||||
stream.Emit(context.Background(), StreamErrorEvent{
|
||||
Message: agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
API: "openai",
|
||||
Provider: d.name,
|
||||
StopReason: agentcore.StopReasonError,
|
||||
ErrorMessage: err.Error(),
|
||||
},
|
||||
Err: fmt.Errorf("%s: %w", d.name, err),
|
||||
})
|
||||
}
|
||||
|
||||
// newPartial builds an empty assistant message tagged for this provider, the
|
||||
// seed for start/text partials (mirrors OpenAIDecoder.partial()'s identity).
|
||||
func (d *responsesDriver) newPartial() agentcore.AssistantMessage {
|
||||
return agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
API: "openai",
|
||||
Provider: d.name,
|
||||
}
|
||||
}
|
||||
|
||||
// mapResponse materializes a completed Responses API result into pigo's
|
||||
// AssistantMessage: a reasoning summary (as a thinking block, when present),
|
||||
// text content, tool calls, usage (when present), diagnostics, and a stop reason
|
||||
// (tool_use when the model requested a tool, otherwise end_turn).
|
||||
func (d *responsesDriver) mapResponse(resp *responses.Response) agentcore.AssistantMessage {
|
||||
msg := d.newPartial()
|
||||
msg.StopReason = agentcore.StopReasonEndTurn
|
||||
msg.ResponseID = resp.ID
|
||||
msg.ResponseModel = string(resp.Model)
|
||||
if thinking := reasoningText(resp); thinking != "" {
|
||||
msg.Content = append(msg.Content, agentcore.NewThinkingContent(thinking))
|
||||
}
|
||||
if text := resp.OutputText(); text != "" {
|
||||
msg.Content = append(msg.Content, agentcore.NewTextContent(text))
|
||||
}
|
||||
var sawToolCall bool
|
||||
for _, item := range resp.Output {
|
||||
if fc := item.AsFunctionCall(); fc.Type == "function_call" {
|
||||
msg.Content = append(msg.Content, toolCallContent(fc))
|
||||
sawToolCall = true
|
||||
}
|
||||
}
|
||||
if sawToolCall {
|
||||
msg.StopReason = agentcore.StopReasonToolUse
|
||||
}
|
||||
if resp.Usage.InputTokens != 0 || resp.Usage.OutputTokens != 0 {
|
||||
msg.Usage = &agentcore.Usage{
|
||||
InputTokens: int(resp.Usage.InputTokens),
|
||||
OutputTokens: int(resp.Usage.OutputTokens),
|
||||
}
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
// toolCallContent maps a Responses function_call item into a pigo
|
||||
// ToolCallContent, keyed by the model's call_id so the tool result can be
|
||||
// backfilled against it on the next turn. Arguments ride verbatim as raw JSON.
|
||||
func toolCallContent(fc responses.ResponseFunctionToolCall) agentcore.ToolCallContent {
|
||||
return agentcore.NewToolCallContent(fc.CallID, fc.Name, json.RawMessage(fc.Arguments))
|
||||
}
|
||||
|
||||
// reasoningText concatenates the summary text of every reasoning item in a
|
||||
// completed response. The Responses API returns the model's reasoning as one or
|
||||
// more reasoning items, each carrying summary parts; pigo surfaces the joined
|
||||
// text as a single thinking block, mirroring how the chat driver renders
|
||||
// accumulated reasoning_content.
|
||||
func reasoningText(resp *responses.Response) string {
|
||||
var b strings.Builder
|
||||
for _, item := range resp.Output {
|
||||
if r := item.AsReasoning(); r.Type == "reasoning" {
|
||||
for _, s := range r.Summary {
|
||||
b.WriteString(s.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// appendToolCalls appends each accumulated tool call to a content list. Kept
|
||||
// separate so the streaming partial and the terminal message build identical
|
||||
// content from the same source.
|
||||
func appendToolCalls(content agentcore.ContentList, calls []agentcore.ToolCallContent) agentcore.ContentList {
|
||||
for _, c := range calls {
|
||||
content = append(content, c)
|
||||
}
|
||||
return content
|
||||
}
|
||||
|
||||
// buildResponsesParams maps a CompletionRequest onto Responses API params. The
|
||||
// system prompt becomes Instructions; the thinking level becomes a reasoning
|
||||
// effort (with an auto summary so reasoning is returned); pigo tools become
|
||||
// Responses function tools; and each message is replayed as the matching input
|
||||
// item(s): assistant tool calls as function_call items, tool results as
|
||||
// function_call_output items, and text (plus any images) as a role-tagged
|
||||
// message.
|
||||
func buildResponsesParams(req CompletionRequest) responses.ResponseNewParams {
|
||||
params := responses.ResponseNewParams{
|
||||
Model: shared.ResponsesModel(req.Model),
|
||||
}
|
||||
if sp := strings.TrimSpace(req.Context.SystemPrompt); sp != "" {
|
||||
params.Instructions = openai.String(sp)
|
||||
}
|
||||
if effort := responsesReasoningEffort(req.Config.ThinkingLevel); effort != "" {
|
||||
// Requesting a summary makes the API return the model's reasoning so pigo
|
||||
// can render it as a thinking block, matching the chat driver.
|
||||
params.Reasoning = shared.ReasoningParam{Effort: effort, Summary: shared.ReasoningSummaryAuto}
|
||||
}
|
||||
if tools := buildResponsesTools(req.Context.Tools); len(tools) > 0 {
|
||||
params.Tools = tools
|
||||
}
|
||||
|
||||
items := make(responses.ResponseInputParam, 0, len(req.Context.Messages))
|
||||
for _, m := range req.Context.Messages {
|
||||
items = appendInputItems(items, m)
|
||||
}
|
||||
params.Input = responses.ResponseNewParamsInputUnion{OfInputItemList: items}
|
||||
return params
|
||||
}
|
||||
|
||||
// buildResponsesTools converts pigo tools into Responses function tools. Each
|
||||
// tool's JSON Schema becomes the function parameters; a schema that is empty or
|
||||
// not a JSON object falls back to an empty object schema so the wire stays
|
||||
// valid. Strict mode is off: pigo schemas are not authored against the Responses
|
||||
// strict-function contract (which requires additionalProperties:false etc.).
|
||||
func buildResponsesTools(tools []agentcore.AgentTool) []responses.ToolUnionParam {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]responses.ToolUnionParam, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
params := map[string]any{}
|
||||
if raw := t.Schema(); len(raw) > 0 {
|
||||
if err := json.Unmarshal(raw, ¶ms); err != nil {
|
||||
params = map[string]any{}
|
||||
}
|
||||
}
|
||||
tool := responses.ToolParamOfFunction(t.Name(), params, false)
|
||||
if desc := t.Description(); desc != "" {
|
||||
tool.OfFunction.Description = openai.String(desc)
|
||||
}
|
||||
out = append(out, tool)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// appendInputItems replays one pigo message as its Responses input item(s).
|
||||
func appendInputItems(items responses.ResponseInputParam, m agentcore.Message) responses.ResponseInputParam {
|
||||
switch msg := m.(type) {
|
||||
case agentcore.ToolResultMessage:
|
||||
// A tool result is backfilled against the model's call_id so the model
|
||||
// can pair it with the request it issued the previous turn.
|
||||
items = append(items, responses.ResponseInputItemParamOfFunctionCallOutput(
|
||||
msg.ToolCallID, contentText(msg.Content)))
|
||||
case agentcore.AssistantMessage:
|
||||
if text := contentText(msg.Content); text != "" {
|
||||
items = append(items, responses.ResponseInputItemParamOfMessage(
|
||||
text, responses.EasyInputMessageRoleAssistant))
|
||||
}
|
||||
for _, call := range msg.ToolCalls() {
|
||||
items = append(items, responses.ResponseInputItemParamOfFunctionCall(
|
||||
string(call.Arguments), call.ID, call.Name))
|
||||
}
|
||||
default:
|
||||
// A user (or other non-assistant) message with images is replayed as a
|
||||
// content-part list (input_text + input_image data URIs); a text-only
|
||||
// message stays a plain string.
|
||||
if parts, ok := imageInputParts(m); ok {
|
||||
items = append(items, responses.ResponseInputItemParamOfMessage(parts, responsesRole(m.Role())))
|
||||
} else if text := messageText(m); text != "" {
|
||||
items = append(items, responses.ResponseInputItemParamOfMessage(
|
||||
text, responsesRole(m.Role())))
|
||||
}
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// imageInputParts builds a Responses content-part list for a message that
|
||||
// carries at least one image: leading input_text (the concatenated text, if
|
||||
// any) followed by one input_image per image, each as a data URI. It returns
|
||||
// ok=false when the message has no images, so the caller keeps the plain-text
|
||||
// path.
|
||||
func imageInputParts(m agentcore.Message) (responses.ResponseInputMessageContentListParam, bool) {
|
||||
content := messageContent(m)
|
||||
var hasImage bool
|
||||
for _, c := range content {
|
||||
if _, ok := c.(agentcore.ImageContent); ok {
|
||||
hasImage = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !hasImage {
|
||||
return nil, false
|
||||
}
|
||||
parts := make(responses.ResponseInputMessageContentListParam, 0, len(content)+1)
|
||||
if text := contentText(content); text != "" {
|
||||
parts = append(parts, responses.ResponseInputContentParamOfInputText(text))
|
||||
}
|
||||
for _, c := range content {
|
||||
img, ok := c.(agentcore.ImageContent)
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
part := responses.ResponseInputContentParamOfInputImage(responses.ResponseInputImageDetailAuto)
|
||||
part.OfInputImage.ImageURL = openai.String(fmt.Sprintf("data:%s;base64,%s", img.MimeType, img.Data))
|
||||
parts = append(parts, part)
|
||||
}
|
||||
return parts, true
|
||||
}
|
||||
|
||||
// responsesReasoningEffort maps pigo's thinking level to a Responses API
|
||||
// reasoning effort. The Responses reasoning field supports only low/medium/high,
|
||||
// so "minimal" collapses to "low" (unlike the chat driver, which forwards
|
||||
// "minimal" verbatim). off/unset yields "", signalling no reasoning param.
|
||||
func responsesReasoningEffort(level agentcore.ThinkingLevel) shared.ReasoningEffort {
|
||||
switch level {
|
||||
case agentcore.ThinkingMinimal, agentcore.ThinkingLow:
|
||||
return shared.ReasoningEffortLow
|
||||
case agentcore.ThinkingMedium:
|
||||
return shared.ReasoningEffortMedium
|
||||
case agentcore.ThinkingHigh, agentcore.ThinkingXHigh, agentcore.ThinkingMax:
|
||||
return shared.ReasoningEffortHigh
|
||||
default:
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
// responsesRole maps a pigo message role to the Responses API input role. Tool
|
||||
// results are surfaced as user turns for this text milestone.
|
||||
func responsesRole(role string) responses.EasyInputMessageRole {
|
||||
switch role {
|
||||
case agentcore.RoleAssistant:
|
||||
return responses.EasyInputMessageRoleAssistant
|
||||
default:
|
||||
return responses.EasyInputMessageRoleUser
|
||||
}
|
||||
}
|
||||
|
||||
// messageContent returns the content list of a message regardless of its
|
||||
// concrete role type, so callers can inspect it for images.
|
||||
func messageContent(m agentcore.Message) agentcore.ContentList {
|
||||
switch msg := m.(type) {
|
||||
case agentcore.UserMessage:
|
||||
return msg.Content
|
||||
case agentcore.AssistantMessage:
|
||||
return msg.Content
|
||||
case agentcore.ToolResultMessage:
|
||||
return msg.Content
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// messageText concatenates the text blocks of a message, ignoring non-text
|
||||
// content (handled in later milestones).
|
||||
func messageText(m agentcore.Message) string {
|
||||
var b strings.Builder
|
||||
collectText(&b, messageContent(m))
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func collectText(b *strings.Builder, content agentcore.ContentList) {
|
||||
for _, c := range content {
|
||||
if tc, ok := c.(agentcore.TextContent); ok {
|
||||
b.WriteString(tc.Text)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// contentText concatenates the text blocks of a content list.
|
||||
func contentText(content agentcore.ContentList) string {
|
||||
var b strings.Builder
|
||||
collectText(&b, content)
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,756 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/openai/openai-go/option"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// roundTripFunc adapts a function to http.RoundTripper so a test can stub the
|
||||
// SDK transport without a live endpoint.
|
||||
type roundTripFunc func(*http.Request) (*http.Response, error)
|
||||
|
||||
func (f roundTripFunc) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }
|
||||
|
||||
// newResponsesTestDriver builds a resp_api driver whose SDK client is pointed at
|
||||
// the given stub round-tripper, capturing the request path the SDK targets.
|
||||
func newResponsesTestDriver(baseURL string, rt roundTripFunc) *responsesDriver {
|
||||
d := NewOpenAIResponsesProvider("openai", baseURL, nil)
|
||||
d.clientOpts = []option.RequestOption{
|
||||
option.WithHTTPClient(&http.Client{Transport: rt}),
|
||||
}
|
||||
return d
|
||||
}
|
||||
|
||||
func jsonResponse(status int, body string) *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: status,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(strings.NewReader(body)),
|
||||
}
|
||||
}
|
||||
|
||||
// sseResponse builds a 200 text/event-stream response whose body is the given
|
||||
// SSE data frames, mirroring how the Responses API streams events. Each frame is
|
||||
// a JSON object carrying its own "type" discriminator.
|
||||
func sseResponse(frames ...string) *http.Response {
|
||||
var b strings.Builder
|
||||
for _, f := range frames {
|
||||
b.WriteString("data: ")
|
||||
b.WriteString(f)
|
||||
b.WriteString("\n\n")
|
||||
}
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"text/event-stream"}},
|
||||
Body: io.NopCloser(strings.NewReader(b.String())),
|
||||
}
|
||||
}
|
||||
|
||||
// completedFrame is a response.completed SSE frame whose embedded Response yields
|
||||
// the given output text, id, model, and token usage — the authoritative terminal
|
||||
// payload the driver maps into its final message.
|
||||
func completedFrame(text, id, model string, inTok, outTok int) string {
|
||||
return `{"type":"response.completed","sequence_number":99,"response":{` +
|
||||
`"id":"` + id + `","model":"` + model + `",` +
|
||||
`"output":[{"type":"message","role":"assistant","status":"completed",` +
|
||||
`"content":[{"type":"output_text","text":"` + text + `"}]}],` +
|
||||
`"usage":{"input_tokens":` + itoa(inTok) + `,"output_tokens":` + itoa(outTok) +
|
||||
`,"input_tokens_details":{"cached_tokens":0},"output_tokens_details":{"reasoning_tokens":0}}}}`
|
||||
}
|
||||
|
||||
func deltaFrame(delta string) string {
|
||||
return `{"type":"response.output_text.delta","item_id":"msg_1","output_index":0,` +
|
||||
`"content_index":0,"sequence_number":1,"logprobs":[],"delta":"` + delta + `"}`
|
||||
}
|
||||
|
||||
func itoa(n int) string { return strconv.Itoa(n) }
|
||||
|
||||
// drain collects the terminal message from a stream, mirroring how the loop
|
||||
// consumes a provider stream.
|
||||
func drain(t *testing.T, stream *AssistantMessageEventStream) agentcore.AssistantMessage {
|
||||
t.Helper()
|
||||
for range stream.Events() {
|
||||
}
|
||||
msg, err := stream.Result(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("stream result error: %v", err)
|
||||
}
|
||||
return msg
|
||||
}
|
||||
|
||||
func userMsg(text string) agentcore.UserMessage {
|
||||
return agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent(text)},
|
||||
}
|
||||
}
|
||||
|
||||
func TestResponsesDriverPostsToResponsesEndpoint(t *testing.T) {
|
||||
var gotPath, gotBody string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
gotPath = r.URL.Path
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
}
|
||||
return sseResponse(
|
||||
deltaFrame("hi "),
|
||||
deltaFrame("there"),
|
||||
completedFrame("hi there", "resp_123", "gpt-4o", 11, 7),
|
||||
), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
req := CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{
|
||||
SystemPrompt: "be terse",
|
||||
Messages: agentcore.MessageList{userMsg("hello")},
|
||||
},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
}
|
||||
stream, err := d.StreamCompletion(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
msg := drain(t, stream)
|
||||
|
||||
if !strings.HasSuffix(gotPath, "/responses") {
|
||||
t.Errorf("request path = %q, want to end with /responses", gotPath)
|
||||
}
|
||||
// The prompt and system instruction must reach the wire body.
|
||||
if !strings.Contains(gotBody, "hello") {
|
||||
t.Errorf("request body missing prompt: %q", gotBody)
|
||||
}
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(gotBody), &payload); err != nil {
|
||||
t.Fatalf("request body not valid JSON: %v", err)
|
||||
}
|
||||
if payload["instructions"] != "be terse" {
|
||||
t.Errorf("instructions = %v, want %q", payload["instructions"], "be terse")
|
||||
}
|
||||
if payload["model"] != "gpt-4o" {
|
||||
t.Errorf("model = %v, want gpt-4o", payload["model"])
|
||||
}
|
||||
// A streaming call must set stream:true on the wire.
|
||||
if payload["stream"] != true {
|
||||
t.Errorf("stream = %v, want true", payload["stream"])
|
||||
}
|
||||
|
||||
if got := textOf(msg); got != "hi there" {
|
||||
t.Errorf("assistant text = %q, want %q", got, "hi there")
|
||||
}
|
||||
if msg.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("stop reason = %q, want end_turn", msg.StopReason)
|
||||
}
|
||||
if msg.ResponseID != "resp_123" {
|
||||
t.Errorf("response id = %q, want resp_123", msg.ResponseID)
|
||||
}
|
||||
if msg.Usage == nil || msg.Usage.InputTokens != 11 || msg.Usage.OutputTokens != 7 {
|
||||
t.Errorf("usage = %+v, want {11 7}", msg.Usage)
|
||||
}
|
||||
if msg.API != "openai" || msg.Provider != "openai" {
|
||||
t.Errorf("tags = api:%q provider:%q, want openai/openai", msg.API, msg.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
// The driver must emit incremental text partials as deltas arrive, and each
|
||||
// partial must carry the text accumulated so far (not just the latest delta), so
|
||||
// the terminal message equals the concatenation the caller already rendered.
|
||||
func TestResponsesDriverStreamsIncrementalDeltas(t *testing.T) {
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return sseResponse(
|
||||
deltaFrame("Hello"),
|
||||
deltaFrame(", "),
|
||||
deltaFrame("world"),
|
||||
completedFrame("Hello, world", "resp_9", "gpt-4o", 3, 4),
|
||||
), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
|
||||
var textPartials []string
|
||||
for ev := range stream.Events() {
|
||||
if te, ok := ev.(StreamTextEvent); ok {
|
||||
textPartials = append(textPartials, textOf(te.Partial))
|
||||
}
|
||||
}
|
||||
msg, err := stream.Result(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("stream result error: %v", err)
|
||||
}
|
||||
|
||||
want := []string{"Hello", "Hello, ", "Hello, world"}
|
||||
if len(textPartials) != len(want) {
|
||||
t.Fatalf("got %d text partials %q, want %d %q", len(textPartials), textPartials, len(want), want)
|
||||
}
|
||||
for i := range want {
|
||||
if textPartials[i] != want[i] {
|
||||
t.Errorf("partial[%d] = %q, want %q", i, textPartials[i], want[i])
|
||||
}
|
||||
}
|
||||
// Final aggregation must match the completed payload, i.e. the last partial.
|
||||
if got := textOf(msg); got != "Hello, world" {
|
||||
t.Errorf("final text = %q, want %q", got, "Hello, world")
|
||||
}
|
||||
}
|
||||
|
||||
// A cancelled context must terminate the stream with an error rather than
|
||||
// yielding a normal end_turn message. The transport cancels mid-flight (after
|
||||
// the stream has started) and reports the cancellation, mirroring how an
|
||||
// in-progress SSE read aborts when the caller cancels.
|
||||
func TestResponsesDriverContextCancelStopsStream(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
cancel()
|
||||
return nil, context.Canceled
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
stream, err := d.StreamCompletion(ctx, CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion should not early-error on cancel: %v", err)
|
||||
}
|
||||
|
||||
var sawError bool
|
||||
for ev := range stream.Events() {
|
||||
if _, ok := ev.(StreamErrorEvent); ok {
|
||||
sawError = true
|
||||
}
|
||||
}
|
||||
if !sawError {
|
||||
t.Fatal("expected a terminal StreamErrorEvent after context cancel")
|
||||
}
|
||||
msg, _ := stream.Result(context.Background())
|
||||
if msg.StopReason != agentcore.StopReasonError {
|
||||
t.Errorf("stop reason = %q, want error", msg.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// A non-2xx from the endpoint must ride the stream as a terminal error event,
|
||||
// not be returned from StreamCompletion (dual failure model, FR-13).
|
||||
func TestResponsesDriverUpstreamErrorRidesStream(t *testing.T) {
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return jsonResponse(http.StatusUnauthorized, `{"error":{"message":"bad key"}}`), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
req := CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("hello")}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
}
|
||||
stream, err := d.StreamCompletion(context.Background(), req)
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion should not early-error on upstream failure: %v", err)
|
||||
}
|
||||
|
||||
var sawError bool
|
||||
for ev := range stream.Events() {
|
||||
if _, ok := ev.(StreamErrorEvent); ok {
|
||||
sawError = true
|
||||
}
|
||||
}
|
||||
if !sawError {
|
||||
t.Fatal("expected a terminal StreamErrorEvent for a 401 response")
|
||||
}
|
||||
msg, _ := stream.Result(context.Background())
|
||||
if msg.StopReason != agentcore.StopReasonError {
|
||||
t.Errorf("stop reason = %q, want error", msg.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// An in-band error event (type "error") mid-stream must ride the stream as a
|
||||
// terminal error, carrying the event's message. This is a distinct path from a
|
||||
// transport-level non-2xx (which surfaces via the stream's Err()).
|
||||
func TestResponsesDriverInStreamErrorEvent(t *testing.T) {
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return sseResponse(
|
||||
deltaFrame("partial"),
|
||||
`{"type":"error","code":"server_error","message":"boom","param":"","sequence_number":2}`,
|
||||
), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("hi")}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion should not early-error: %v", err)
|
||||
}
|
||||
|
||||
var errEvent *StreamErrorEvent
|
||||
for ev := range stream.Events() {
|
||||
if se, ok := ev.(StreamErrorEvent); ok {
|
||||
e := se
|
||||
errEvent = &e
|
||||
}
|
||||
}
|
||||
if errEvent == nil {
|
||||
t.Fatal("expected a terminal StreamErrorEvent for an in-band error event")
|
||||
}
|
||||
if !strings.Contains(errEvent.Message.ErrorMessage, "boom") {
|
||||
t.Errorf("error message = %q, want to contain %q", errEvent.Message.ErrorMessage, "boom")
|
||||
}
|
||||
msg, _ := stream.Result(context.Background())
|
||||
if msg.StopReason != agentcore.StopReasonError {
|
||||
t.Errorf("stop reason = %q, want error", msg.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// A missing API key is the one early "cannot build the stream" error.
|
||||
func TestResponsesDriverMissingKeyIsEarlyError(t *testing.T) {
|
||||
d := NewOpenAIResponsesProvider("openai", "https://api.openai.test/v1", nil)
|
||||
_, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Config: StreamConfig{APIKey: " "},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("expected early error for missing API key")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "missing API key") {
|
||||
t.Errorf("error = %q, want to mention missing API key", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
// textOf returns the concatenated text content of an assistant message.
|
||||
func textOf(m agentcore.AssistantMessage) string {
|
||||
var b bytes.Buffer
|
||||
for _, c := range m.Content {
|
||||
if tc, ok := c.(agentcore.TextContent); ok {
|
||||
b.WriteString(tc.Text)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// fakeTool is a minimal AgentTool for exercising tool-schema serialization; it
|
||||
// never executes in these transport-level tests.
|
||||
type fakeTool struct {
|
||||
name string
|
||||
desc string
|
||||
schema json.RawMessage
|
||||
}
|
||||
|
||||
func (t fakeTool) Name() string { return t.name }
|
||||
func (t fakeTool) Description() string { return t.desc }
|
||||
func (t fakeTool) Schema() json.RawMessage { return t.schema }
|
||||
func (t fakeTool) ExecutionMode() agentcore.ToolExecutionMode { return agentcore.ToolExecutionParallel }
|
||||
func (t fakeTool) Execute(ctx context.Context, id string, args json.RawMessage, onUpdate agentcore.ToolUpdateFunc) (agentcore.AgentToolResult, error) {
|
||||
return agentcore.AgentToolResult{}, nil
|
||||
}
|
||||
|
||||
// functionCallDoneFrame is a response.output_item.done SSE frame carrying a
|
||||
// finalized function_call item (the model's tool request).
|
||||
func functionCallDoneFrame(callID, name, argsJSON string) string {
|
||||
frame := map[string]any{
|
||||
"type": "response.output_item.done",
|
||||
"output_index": 0,
|
||||
"sequence_number": 5,
|
||||
"item": map[string]any{
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
"arguments": argsJSON,
|
||||
"status": "completed",
|
||||
},
|
||||
}
|
||||
b, _ := json.Marshal(frame)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// completedToolFrame is a response.completed frame whose output is a single
|
||||
// function_call item (no assistant text) plus token usage.
|
||||
func completedToolFrame(id, model, callID, name, argsJSON string) string {
|
||||
frame := map[string]any{
|
||||
"type": "response.completed",
|
||||
"sequence_number": 99,
|
||||
"response": map[string]any{
|
||||
"id": id,
|
||||
"model": model,
|
||||
"output": []any{map[string]any{
|
||||
"type": "function_call",
|
||||
"id": "fc_1",
|
||||
"call_id": callID,
|
||||
"name": name,
|
||||
"arguments": argsJSON,
|
||||
"status": "completed",
|
||||
}},
|
||||
"usage": map[string]any{
|
||||
"input_tokens": 5,
|
||||
"output_tokens": 2,
|
||||
"input_tokens_details": map[string]any{"cached_tokens": 0},
|
||||
"output_tokens_details": map[string]any{"reasoning_tokens": 0},
|
||||
},
|
||||
},
|
||||
}
|
||||
b, _ := json.Marshal(frame)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// A pigo tool must reach the wire as a Responses function tool: type "function",
|
||||
// its name, JSON-Schema parameters, and description.
|
||||
func TestResponsesDriverSendsToolSchema(t *testing.T) {
|
||||
var gotBody string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
}
|
||||
return sseResponse(completedFrame("ok", "resp_1", "gpt-4o", 1, 1)), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
tool := fakeTool{
|
||||
name: "read_file",
|
||||
desc: "reads a file",
|
||||
schema: json.RawMessage(`{"type":"object","properties":{"path":{"type":"string"}},"required":["path"]}`),
|
||||
}
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{
|
||||
Messages: agentcore.MessageList{userMsg("read a.go")},
|
||||
Tools: []agentcore.AgentTool{tool},
|
||||
},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
drain(t, stream)
|
||||
|
||||
var payload map[string]any
|
||||
if err := json.Unmarshal([]byte(gotBody), &payload); err != nil {
|
||||
t.Fatalf("request body not valid JSON: %v", err)
|
||||
}
|
||||
tools, ok := payload["tools"].([]any)
|
||||
if !ok || len(tools) != 1 {
|
||||
t.Fatalf("tools = %v, want a single-element array", payload["tools"])
|
||||
}
|
||||
tool0 := tools[0].(map[string]any)
|
||||
if tool0["type"] != "function" {
|
||||
t.Errorf("tool type = %v, want function", tool0["type"])
|
||||
}
|
||||
if tool0["name"] != "read_file" {
|
||||
t.Errorf("tool name = %v, want read_file", tool0["name"])
|
||||
}
|
||||
if tool0["description"] != "reads a file" {
|
||||
t.Errorf("tool description = %v, want %q", tool0["description"], "reads a file")
|
||||
}
|
||||
params, ok := tool0["parameters"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatalf("tool parameters missing or not an object: %v", tool0["parameters"])
|
||||
}
|
||||
props, ok := params["properties"].(map[string]any)
|
||||
if !ok || props["path"] == nil {
|
||||
t.Errorf("tool parameters.properties.path missing: %v", params)
|
||||
}
|
||||
}
|
||||
|
||||
// A function_call in the completed response must be parsed into a pigo
|
||||
// ToolCallContent (id + name + raw arguments) and set StopReason=tool_use; a
|
||||
// StreamToolCallEvent must also surface the pending call mid-stream.
|
||||
func TestResponsesDriverParsesToolCall(t *testing.T) {
|
||||
args := `{"path":"a.go"}`
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
return sseResponse(
|
||||
functionCallDoneFrame("call_abc", "read_file", args),
|
||||
completedToolFrame("resp_7", "gpt-4o", "call_abc", "read_file", args),
|
||||
), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("read a.go")}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
|
||||
var sawToolCallEvent bool
|
||||
for ev := range stream.Events() {
|
||||
if _, ok := ev.(StreamToolCallEvent); ok {
|
||||
sawToolCallEvent = true
|
||||
}
|
||||
}
|
||||
if !sawToolCallEvent {
|
||||
t.Error("expected a StreamToolCallEvent mid-stream")
|
||||
}
|
||||
msg, err := stream.Result(context.Background())
|
||||
if err != nil {
|
||||
t.Fatalf("stream result error: %v", err)
|
||||
}
|
||||
|
||||
calls := msg.ToolCalls()
|
||||
if len(calls) != 1 {
|
||||
t.Fatalf("got %d tool calls, want 1", len(calls))
|
||||
}
|
||||
if calls[0].ID != "call_abc" || calls[0].Name != "read_file" {
|
||||
t.Errorf("tool call = id:%q name:%q, want call_abc/read_file", calls[0].ID, calls[0].Name)
|
||||
}
|
||||
if string(calls[0].Arguments) != args {
|
||||
t.Errorf("tool call arguments = %q, want %q", calls[0].Arguments, args)
|
||||
}
|
||||
if msg.StopReason != agentcore.StopReasonToolUse {
|
||||
t.Errorf("stop reason = %q, want tool_use", msg.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// On a follow-up turn, a prior assistant tool call must be replayed as a
|
||||
// function_call input item and its result as a function_call_output item, both
|
||||
// keyed by the same call_id, so the model can pair request and result.
|
||||
func TestResponsesDriverBackfillsToolResult(t *testing.T) {
|
||||
var gotBody string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
}
|
||||
return sseResponse(completedFrame("done", "resp_2", "gpt-4o", 8, 3)), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
assistant := agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
Content: agentcore.ContentList{
|
||||
agentcore.NewToolCallContent("call_abc", "read_file", json.RawMessage(`{"path":"a.go"}`)),
|
||||
},
|
||||
}
|
||||
result := agentcore.ToolResultMessage{
|
||||
RoleField: agentcore.RoleToolResult,
|
||||
ToolCallID: "call_abc",
|
||||
ToolName: "read_file",
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("package main")},
|
||||
}
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{
|
||||
userMsg("read a.go"), assistant, result,
|
||||
}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
drain(t, stream)
|
||||
|
||||
var payload struct {
|
||||
Input []map[string]any `json:"input"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(gotBody), &payload); err != nil {
|
||||
t.Fatalf("request body not valid JSON: %v", err)
|
||||
}
|
||||
var sawCall, sawOutput bool
|
||||
for _, item := range payload.Input {
|
||||
switch item["type"] {
|
||||
case "function_call":
|
||||
sawCall = true
|
||||
if item["call_id"] != "call_abc" || item["name"] != "read_file" {
|
||||
t.Errorf("function_call item = %v, want call_abc/read_file", item)
|
||||
}
|
||||
if item["arguments"] != `{"path":"a.go"}` {
|
||||
t.Errorf("function_call arguments = %v, want the raw args JSON string", item["arguments"])
|
||||
}
|
||||
case "function_call_output":
|
||||
sawOutput = true
|
||||
if item["call_id"] != "call_abc" {
|
||||
t.Errorf("function_call_output call_id = %v, want call_abc", item["call_id"])
|
||||
}
|
||||
if item["output"] != "package main" {
|
||||
t.Errorf("function_call_output output = %v, want %q", item["output"], "package main")
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawCall {
|
||||
t.Error("wire input missing the replayed function_call item")
|
||||
}
|
||||
if !sawOutput {
|
||||
t.Error("wire input missing the function_call_output item")
|
||||
}
|
||||
}
|
||||
|
||||
// completedReasoningFrame is a response.completed frame whose output carries a
|
||||
// reasoning item (summary text) followed by the assistant message text.
|
||||
func completedReasoningFrame(id, model, summary, text string) string {
|
||||
frame := map[string]any{
|
||||
"type": "response.completed",
|
||||
"sequence_number": 99,
|
||||
"response": map[string]any{
|
||||
"id": id,
|
||||
"model": model,
|
||||
"output": []any{
|
||||
map[string]any{
|
||||
"type": "reasoning",
|
||||
"id": "rs_1",
|
||||
"summary": []any{map[string]any{
|
||||
"type": "summary_text",
|
||||
"text": summary,
|
||||
}},
|
||||
},
|
||||
map[string]any{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"status": "completed",
|
||||
"content": []any{map[string]any{
|
||||
"type": "output_text",
|
||||
"text": text,
|
||||
}},
|
||||
},
|
||||
},
|
||||
"usage": map[string]any{
|
||||
"input_tokens": 3,
|
||||
"output_tokens": 4,
|
||||
"input_tokens_details": map[string]any{"cached_tokens": 0},
|
||||
"output_tokens_details": map[string]any{"reasoning_tokens": 2},
|
||||
},
|
||||
},
|
||||
}
|
||||
b, _ := json.Marshal(frame)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
// A user message carrying an image must reach the wire as a message whose
|
||||
// content is a part list: an input_text part plus an input_image part whose
|
||||
// image_url is the base64 data URI.
|
||||
func TestResponsesDriverSendsImageInput(t *testing.T) {
|
||||
var gotBody string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
}
|
||||
return sseResponse(completedFrame("ok", "resp_1", "gpt-4o", 1, 1)), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
imgMsg := agentcore.UserMessage{
|
||||
RoleField: agentcore.RoleUser,
|
||||
Content: agentcore.ContentList{
|
||||
agentcore.NewTextContent("what is this?"),
|
||||
agentcore.NewImageContent("aGVsbG8=", "image/png"),
|
||||
},
|
||||
}
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{imgMsg}},
|
||||
Config: StreamConfig{APIKey: "sk-test"},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
drain(t, stream)
|
||||
|
||||
var payload struct {
|
||||
Input []struct {
|
||||
Type string `json:"type"`
|
||||
Role string `json:"role"`
|
||||
Content []map[string]any `json:"content"`
|
||||
} `json:"input"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(gotBody), &payload); err != nil {
|
||||
t.Fatalf("request body not valid JSON: %v", err)
|
||||
}
|
||||
if len(payload.Input) != 1 {
|
||||
t.Fatalf("got %d input items, want 1", len(payload.Input))
|
||||
}
|
||||
parts := payload.Input[0].Content
|
||||
var sawText, sawImage bool
|
||||
for _, p := range parts {
|
||||
switch p["type"] {
|
||||
case "input_text":
|
||||
sawText = true
|
||||
if p["text"] != "what is this?" {
|
||||
t.Errorf("input_text = %v, want %q", p["text"], "what is this?")
|
||||
}
|
||||
case "input_image":
|
||||
sawImage = true
|
||||
if p["image_url"] != "data:image/png;base64,aGVsbG8=" {
|
||||
t.Errorf("input_image image_url = %v, want the data URI", p["image_url"])
|
||||
}
|
||||
}
|
||||
}
|
||||
if !sawText {
|
||||
t.Error("wire input missing the input_text part")
|
||||
}
|
||||
if !sawImage {
|
||||
t.Error("wire input missing the input_image part")
|
||||
}
|
||||
}
|
||||
|
||||
// A request with a thinking level must set the reasoning.effort (and an auto
|
||||
// summary) on the wire, and a reasoning item in the completed response must be
|
||||
// parsed into a leading ThinkingContent block.
|
||||
func TestResponsesDriverReasoning(t *testing.T) {
|
||||
var gotBody string
|
||||
rt := roundTripFunc(func(r *http.Request) (*http.Response, error) {
|
||||
if r.Body != nil {
|
||||
b, _ := io.ReadAll(r.Body)
|
||||
gotBody = string(b)
|
||||
}
|
||||
return sseResponse(completedReasoningFrame("resp_9", "gpt-4o", "let me think", "the answer")), nil
|
||||
})
|
||||
d := newResponsesTestDriver("https://api.openai.test/v1", rt)
|
||||
|
||||
stream, err := d.StreamCompletion(context.Background(), CompletionRequest{
|
||||
Model: "gpt-4o",
|
||||
Context: LlmContext{Messages: agentcore.MessageList{userMsg("solve it")}},
|
||||
Config: StreamConfig{APIKey: "sk-test", ThinkingLevel: agentcore.ThinkingMedium},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamCompletion returned early error: %v", err)
|
||||
}
|
||||
msg := drain(t, stream)
|
||||
|
||||
var payload struct {
|
||||
Reasoning struct {
|
||||
Effort string `json:"effort"`
|
||||
Summary string `json:"summary"`
|
||||
} `json:"reasoning"`
|
||||
}
|
||||
if err := json.Unmarshal([]byte(gotBody), &payload); err != nil {
|
||||
t.Fatalf("request body not valid JSON: %v", err)
|
||||
}
|
||||
if payload.Reasoning.Effort != "medium" {
|
||||
t.Errorf("reasoning.effort = %q, want medium", payload.Reasoning.Effort)
|
||||
}
|
||||
if payload.Reasoning.Summary != "auto" {
|
||||
t.Errorf("reasoning.summary = %q, want auto", payload.Reasoning.Summary)
|
||||
}
|
||||
|
||||
var thinking string
|
||||
for _, c := range msg.Content {
|
||||
if tc, ok := c.(agentcore.ThinkingContent); ok {
|
||||
thinking = tc.Thinking
|
||||
}
|
||||
}
|
||||
if thinking != "let me think" {
|
||||
t.Errorf("thinking content = %q, want %q", thinking, "let me think")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,225 @@
|
||||
// This file implements parameter validation and endpoint construction for the
|
||||
// "special auth" providers (US-007 / FR-12): Azure OpenAI, Amazon Bedrock,
|
||||
// Google Vertex, and Cloudflare (Workers AI + AI Gateway). Unlike the generic
|
||||
// bearer/x-api-key providers, each of these composes its endpoint from several
|
||||
// environment variables and/or needs a non-standard credential path, so a
|
||||
// dedicated resolver validates the required parameters and builds the concrete
|
||||
// base URL before handing off to the shared OpenAI-/Anthropic-compatible driver.
|
||||
//
|
||||
// Scope note (PRD Non-Goals): AWS SigV4 request signing is NOT implemented.
|
||||
// Bedrock supports only the AWS_BEARER_TOKEN_BEDROCK bearer path; other AWS
|
||||
// credential sources (AWS_PROFILE, AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY)
|
||||
// are only *detected* so that a clear, actionable error is returned instead of
|
||||
// an opaque auth failure.
|
||||
//
|
||||
// Security: this file reads env var NAMES and composes URLs from non-secret
|
||||
// parameters (region, resource name, account id, …). Secret values (API keys,
|
||||
// bearer tokens) are never logged or embedded in error text — errors name the
|
||||
// absent env var, never a value.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// IsSpecialAuthProvider reports whether a provider spec needs the bespoke
|
||||
// endpoint-construction / credential-validation handled by ResolveSpecialProvider,
|
||||
// rather than the generic driver wiring. It matches the multi-parameter auth
|
||||
// schemes (azure/aws/special) and the two Cloudflare providers (which keep a
|
||||
// standard auth scheme but still compose their endpoint from env vars).
|
||||
func IsSpecialAuthProvider(spec ProviderSpec) bool {
|
||||
switch spec.AuthScheme {
|
||||
case AuthAzure, AuthAWS, AuthSpecial:
|
||||
return true
|
||||
}
|
||||
return strings.HasPrefix(spec.Name, "cloudflare-")
|
||||
}
|
||||
|
||||
// ResolveSpecialProvider validates the required parameters for a special-auth
|
||||
// provider and constructs the matching wire driver against the composed base
|
||||
// URL. flagBaseURL is the explicit --base-url override (highest precedence, wins
|
||||
// over any composed default); env resolves environment variables (os.Getenv in
|
||||
// production, a fake in tests). A missing required parameter yields an error
|
||||
// naming exactly which env var is absent; no network request is made here.
|
||||
func ResolveSpecialProvider(spec ProviderSpec, model, flagBaseURL string, env func(string) string) (Provider, error) {
|
||||
if env == nil {
|
||||
env = func(string) string { return "" }
|
||||
}
|
||||
models := []Model{{Provider: spec.Name, ID: model, SupportsImages: true}}
|
||||
switch spec.Name {
|
||||
case "azure-openai-responses":
|
||||
return resolveAzureOpenAI(spec, model, flagBaseURL, env, models)
|
||||
case "amazon-bedrock":
|
||||
return resolveBedrock(spec, flagBaseURL, env, models)
|
||||
case "google-vertex":
|
||||
return resolveGoogleVertex(spec, flagBaseURL, env, models)
|
||||
case "cloudflare-workers-ai":
|
||||
return resolveCloudflareWorkersAI(spec, flagBaseURL, env, models)
|
||||
case "cloudflare-ai-gateway":
|
||||
return resolveCloudflareAIGateway(spec, flagBaseURL, env, models)
|
||||
default:
|
||||
return nil, fmt.Errorf("provider %q is not a special-auth provider", spec.Name)
|
||||
}
|
||||
}
|
||||
|
||||
// resolveAzureOpenAI composes the Azure OpenAI endpoint. The endpoint origin is
|
||||
// AZURE_OPENAI_BASE_URL (or the --base-url override), else it is built from
|
||||
// AZURE_OPENAI_RESOURCE_NAME as https://{resource}.openai.azure.com. The API
|
||||
// version (AZURE_OPENAI_API_VERSION, default "v1") and an optional deployment
|
||||
// mapping (AZURE_OPENAI_DEPLOYMENT_NAME_MAP) shape the path. Auth uses
|
||||
// AZURE_OPENAI_API_KEY over the OpenAI wire.
|
||||
func resolveAzureOpenAI(_ ProviderSpec, model, flagBaseURL string, env func(string) string, models []Model) (Provider, error) {
|
||||
if strings.TrimSpace(env("AZURE_OPENAI_API_KEY")) == "" {
|
||||
return nil, fmt.Errorf("azure-openai-responses: missing required env var AZURE_OPENAI_API_KEY")
|
||||
}
|
||||
origin := strings.TrimSpace(flagBaseURL)
|
||||
if origin == "" {
|
||||
origin = strings.TrimSpace(env("AZURE_OPENAI_BASE_URL"))
|
||||
}
|
||||
if origin == "" {
|
||||
resource := strings.TrimSpace(env("AZURE_OPENAI_RESOURCE_NAME"))
|
||||
if resource == "" {
|
||||
return nil, fmt.Errorf("azure-openai-responses: missing endpoint configuration; set AZURE_OPENAI_BASE_URL or AZURE_OPENAI_RESOURCE_NAME")
|
||||
}
|
||||
origin = fmt.Sprintf("https://%s.openai.azure.com", resource)
|
||||
}
|
||||
apiVersion := strings.TrimSpace(env("AZURE_OPENAI_API_VERSION"))
|
||||
if apiVersion == "" {
|
||||
apiVersion = "v1"
|
||||
}
|
||||
deployment := resolveAzureDeployment(env("AZURE_OPENAI_DEPLOYMENT_NAME_MAP"), model)
|
||||
baseURL := azureEndpoint(origin, apiVersion, deployment)
|
||||
return NewOpenAICompatibleProvider(baseURL, models), nil
|
||||
}
|
||||
|
||||
// azureEndpoint builds the Azure OpenAI base URL from a validated origin. When a
|
||||
// deployment is resolved for the model, the classic deployment-scoped path is
|
||||
// used (…/openai/deployments/{deployment}); otherwise the version-scoped v1 path
|
||||
// (…/openai/{apiVersion}) is used. The shared driver appends /chat/completions.
|
||||
func azureEndpoint(origin, apiVersion, deployment string) string {
|
||||
origin = strings.TrimRight(strings.TrimSpace(origin), "/")
|
||||
if deployment != "" {
|
||||
return fmt.Sprintf("%s/openai/deployments/%s", origin, deployment)
|
||||
}
|
||||
return fmt.Sprintf("%s/openai/%s", origin, apiVersion)
|
||||
}
|
||||
|
||||
// resolveAzureDeployment parses AZURE_OPENAI_DEPLOYMENT_NAME_MAP (a
|
||||
// comma-separated list of model=deployment pairs) and returns the deployment
|
||||
// mapped to model, or "" when the map is empty or has no entry for the model.
|
||||
func resolveAzureDeployment(raw, model string) string {
|
||||
m := parseDeploymentMap(raw)
|
||||
return m[strings.TrimSpace(model)]
|
||||
}
|
||||
|
||||
// parseDeploymentMap parses a comma-separated "model=deployment" list into a
|
||||
// map. Blank entries and entries without '=' are skipped; keys and values are
|
||||
// trimmed. It never returns nil so lookups are always safe.
|
||||
func parseDeploymentMap(raw string) map[string]string {
|
||||
out := make(map[string]string)
|
||||
for _, pair := range strings.Split(raw, ",") {
|
||||
pair = strings.TrimSpace(pair)
|
||||
if pair == "" {
|
||||
continue
|
||||
}
|
||||
k, v, ok := strings.Cut(pair, "=")
|
||||
k, v = strings.TrimSpace(k), strings.TrimSpace(v)
|
||||
if !ok || k == "" || v == "" {
|
||||
continue
|
||||
}
|
||||
out[k] = v
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// resolveBedrock composes the Amazon Bedrock runtime endpoint
|
||||
// (https://bedrock-runtime.{region}.amazonaws.com; region defaults to
|
||||
// us-east-1) and validates credentials. Only the AWS_BEARER_TOKEN_BEDROCK
|
||||
// bearer path is supported (SigV4 is out of scope): if only AWS_PROFILE or
|
||||
// static AWS keys are present, a clear error explains SigV4 is unsupported and
|
||||
// names the missing AWS_BEARER_TOKEN_BEDROCK.
|
||||
func resolveBedrock(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) {
|
||||
if strings.TrimSpace(env("AWS_BEARER_TOKEN_BEDROCK")) == "" {
|
||||
hasProfile := strings.TrimSpace(env("AWS_PROFILE")) != ""
|
||||
hasStaticKeys := strings.TrimSpace(env("AWS_ACCESS_KEY_ID")) != "" &&
|
||||
strings.TrimSpace(env("AWS_SECRET_ACCESS_KEY")) != ""
|
||||
if hasProfile || hasStaticKeys {
|
||||
return nil, fmt.Errorf("amazon-bedrock: detected AWS credentials (AWS_PROFILE / AWS_ACCESS_KEY_ID) but SigV4 request signing is not supported yet; set AWS_BEARER_TOKEN_BEDROCK to use the bearer-token path")
|
||||
}
|
||||
return nil, fmt.Errorf("amazon-bedrock: missing required env var AWS_BEARER_TOKEN_BEDROCK")
|
||||
}
|
||||
baseURL := strings.TrimSpace(flagBaseURL)
|
||||
if baseURL == "" {
|
||||
region := strings.TrimSpace(env("AWS_REGION"))
|
||||
if region == "" {
|
||||
region = "us-east-1"
|
||||
}
|
||||
baseURL = fmt.Sprintf("https://bedrock-runtime.%s.amazonaws.com", region)
|
||||
}
|
||||
return NewBedrockProvider(baseURL, models), nil
|
||||
}
|
||||
|
||||
// resolveGoogleVertex composes the Vertex AI endpoint
|
||||
// (https://{location}-aiplatform.googleapis.com) and validates that a project,
|
||||
// a location, and a credential source (GOOGLE_CLOUD_API_KEY or ADC via
|
||||
// GOOGLE_APPLICATION_CREDENTIALS) are present, naming any absent env var.
|
||||
func resolveGoogleVertex(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) {
|
||||
if strings.TrimSpace(env("GOOGLE_CLOUD_PROJECT")) == "" {
|
||||
return nil, fmt.Errorf("google-vertex: missing required env var GOOGLE_CLOUD_PROJECT")
|
||||
}
|
||||
location := strings.TrimSpace(env("GOOGLE_CLOUD_LOCATION"))
|
||||
if location == "" {
|
||||
return nil, fmt.Errorf("google-vertex: missing required env var GOOGLE_CLOUD_LOCATION")
|
||||
}
|
||||
if strings.TrimSpace(env("GOOGLE_CLOUD_API_KEY")) == "" &&
|
||||
strings.TrimSpace(env("GOOGLE_APPLICATION_CREDENTIALS")) == "" {
|
||||
return nil, fmt.Errorf("google-vertex: missing credentials; set GOOGLE_CLOUD_API_KEY or GOOGLE_APPLICATION_CREDENTIALS (ADC)")
|
||||
}
|
||||
baseURL := strings.TrimSpace(flagBaseURL)
|
||||
if baseURL == "" {
|
||||
baseURL = fmt.Sprintf("https://%s-aiplatform.googleapis.com", location)
|
||||
}
|
||||
return NewOpenAICompatibleProvider(baseURL, models), nil
|
||||
}
|
||||
|
||||
// resolveCloudflareWorkersAI composes the Workers AI endpoint
|
||||
// (https://api.cloudflare.com/client/v4/accounts/{account}/ai/v1), requiring
|
||||
// CLOUDFLARE_API_KEY and CLOUDFLARE_ACCOUNT_ID. OpenAI wire.
|
||||
func resolveCloudflareWorkersAI(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) {
|
||||
if strings.TrimSpace(env("CLOUDFLARE_API_KEY")) == "" {
|
||||
return nil, fmt.Errorf("cloudflare-workers-ai: missing required env var CLOUDFLARE_API_KEY")
|
||||
}
|
||||
account := strings.TrimSpace(env("CLOUDFLARE_ACCOUNT_ID"))
|
||||
if account == "" {
|
||||
return nil, fmt.Errorf("cloudflare-workers-ai: missing required env var CLOUDFLARE_ACCOUNT_ID")
|
||||
}
|
||||
baseURL := strings.TrimSpace(flagBaseURL)
|
||||
if baseURL == "" {
|
||||
baseURL = fmt.Sprintf("https://api.cloudflare.com/client/v4/accounts/%s/ai/v1", account)
|
||||
}
|
||||
return NewOpenAICompatibleProvider(baseURL, models), nil
|
||||
}
|
||||
|
||||
// resolveCloudflareAIGateway composes the AI Gateway endpoint
|
||||
// (https://gateway.ai.cloudflare.com/v1/{account}/{gateway}/anthropic),
|
||||
// requiring CLOUDFLARE_API_KEY, CLOUDFLARE_ACCOUNT_ID, and CLOUDFLARE_GATEWAY_ID.
|
||||
// Anthropic wire.
|
||||
func resolveCloudflareAIGateway(_ ProviderSpec, flagBaseURL string, env func(string) string, models []Model) (Provider, error) {
|
||||
if strings.TrimSpace(env("CLOUDFLARE_API_KEY")) == "" {
|
||||
return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_API_KEY")
|
||||
}
|
||||
account := strings.TrimSpace(env("CLOUDFLARE_ACCOUNT_ID"))
|
||||
if account == "" {
|
||||
return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_ACCOUNT_ID")
|
||||
}
|
||||
gateway := strings.TrimSpace(env("CLOUDFLARE_GATEWAY_ID"))
|
||||
if gateway == "" {
|
||||
return nil, fmt.Errorf("cloudflare-ai-gateway: missing required env var CLOUDFLARE_GATEWAY_ID")
|
||||
}
|
||||
baseURL := strings.TrimSpace(flagBaseURL)
|
||||
if baseURL == "" {
|
||||
baseURL = fmt.Sprintf("https://gateway.ai.cloudflare.com/v1/%s/%s/anthropic", account, gateway)
|
||||
}
|
||||
return NewAnthropicProvider(baseURL, models), nil
|
||||
}
|
||||
@@ -0,0 +1,291 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// envFrom builds an env(string)string lookup from a map for hermetic tests: no
|
||||
// process environment is read, so tests never depend on ambient state and make
|
||||
// no network requests.
|
||||
func envFrom(m map[string]string) func(string) string {
|
||||
return func(k string) string { return m[k] }
|
||||
}
|
||||
|
||||
// baseURLOf extracts the composed base URL from a constructed driver by type
|
||||
// asserting the two concrete driver shapes (same-package access to unexported
|
||||
// fields). It fails the test if the provider is neither shape.
|
||||
func baseURLOf(t *testing.T, p Provider) string {
|
||||
t.Helper()
|
||||
switch d := p.(type) {
|
||||
case *openAICompatDriver:
|
||||
return d.baseURL
|
||||
case *anthropicCompatDriver:
|
||||
return d.baseURL
|
||||
default:
|
||||
t.Fatalf("unexpected provider type %T", p)
|
||||
return ""
|
||||
}
|
||||
}
|
||||
|
||||
func specFor(t *testing.T, name string) ProviderSpec {
|
||||
t.Helper()
|
||||
spec, ok := LookupProviderSpec(name)
|
||||
if !ok {
|
||||
t.Fatalf("registry missing provider %q", name)
|
||||
}
|
||||
return spec
|
||||
}
|
||||
|
||||
func TestResolveSpecialProvider_Azure(t *testing.T) {
|
||||
spec := specFor(t, "azure-openai-responses")
|
||||
|
||||
// Missing API key.
|
||||
if _, err := ResolveSpecialProvider(spec, "gpt-4o", "", envFrom(nil)); err == nil ||
|
||||
!strings.Contains(err.Error(), "AZURE_OPENAI_API_KEY") {
|
||||
t.Fatalf("expected AZURE_OPENAI_API_KEY error, got %v", err)
|
||||
}
|
||||
|
||||
// Key present but no endpoint origin.
|
||||
env := envFrom(map[string]string{"AZURE_OPENAI_API_KEY": "k"})
|
||||
if _, err := ResolveSpecialProvider(spec, "gpt-4o", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "AZURE_OPENAI_BASE_URL") ||
|
||||
!strings.Contains(err.Error(), "AZURE_OPENAI_RESOURCE_NAME") {
|
||||
t.Fatalf("expected endpoint-config error naming both env vars, got %v", err)
|
||||
}
|
||||
|
||||
// Resource name → composed origin, default api version v1.
|
||||
env = envFrom(map[string]string{
|
||||
"AZURE_OPENAI_API_KEY": "k",
|
||||
"AZURE_OPENAI_RESOURCE_NAME": "myres",
|
||||
})
|
||||
p, err := ResolveSpecialProvider(spec, "gpt-4o", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://myres.openai.azure.com/openai/v1"; got != want {
|
||||
t.Fatalf("azure base_url = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// Explicit base URL env + custom api version.
|
||||
env = envFrom(map[string]string{
|
||||
"AZURE_OPENAI_API_KEY": "k",
|
||||
"AZURE_OPENAI_BASE_URL": "https://custom.example.com",
|
||||
"AZURE_OPENAI_API_VERSION": "2024-10-01",
|
||||
})
|
||||
p, err = ResolveSpecialProvider(spec, "gpt-4o", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://custom.example.com/openai/2024-10-01"; got != want {
|
||||
t.Fatalf("azure base_url = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// Deployment name map → deployment-scoped path.
|
||||
env = envFrom(map[string]string{
|
||||
"AZURE_OPENAI_API_KEY": "k",
|
||||
"AZURE_OPENAI_RESOURCE_NAME": "myres",
|
||||
"AZURE_OPENAI_DEPLOYMENT_NAME_MAP": "gpt-4o=prod-4o , other=x",
|
||||
})
|
||||
p, err = ResolveSpecialProvider(spec, "gpt-4o", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://myres.openai.azure.com/openai/deployments/prod-4o"; got != want {
|
||||
t.Fatalf("azure deployment base_url = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// --base-url flag wins over env origin.
|
||||
p, err = ResolveSpecialProvider(spec, "gpt-4o", "https://flag.example.com", envFrom(map[string]string{"AZURE_OPENAI_API_KEY": "k"}))
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://flag.example.com/openai/v1"; got != want {
|
||||
t.Fatalf("azure flag base_url = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParseDeploymentMap(t *testing.T) {
|
||||
m := parseDeploymentMap(" a=1, b = 2 ,,bad,c=,=d ")
|
||||
if m["a"] != "1" || m["b"] != "2" {
|
||||
t.Fatalf("parseDeploymentMap = %v, want a=1 b=2", m)
|
||||
}
|
||||
if _, ok := m["bad"]; ok {
|
||||
t.Fatalf("expected 'bad' (no '=') to be skipped: %v", m)
|
||||
}
|
||||
if _, ok := m["c"]; ok {
|
||||
t.Fatalf("expected 'c=' (empty value) to be skipped: %v", m)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveSpecialProvider_Bedrock(t *testing.T) {
|
||||
spec := specFor(t, "amazon-bedrock")
|
||||
|
||||
// No credentials at all → names the bearer token.
|
||||
if _, err := ResolveSpecialProvider(spec, "claude", "", envFrom(nil)); err == nil ||
|
||||
!strings.Contains(err.Error(), "AWS_BEARER_TOKEN_BEDROCK") {
|
||||
t.Fatalf("expected AWS_BEARER_TOKEN_BEDROCK error, got %v", err)
|
||||
}
|
||||
|
||||
// Only profile present → SigV4-unsupported error, still names bearer token.
|
||||
env := envFrom(map[string]string{"AWS_PROFILE": "default"})
|
||||
if _, err := ResolveSpecialProvider(spec, "claude", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "SigV4") ||
|
||||
!strings.Contains(err.Error(), "AWS_BEARER_TOKEN_BEDROCK") {
|
||||
t.Fatalf("expected SigV4-unsupported error naming bearer token, got %v", err)
|
||||
}
|
||||
|
||||
// Only static keys present → SigV4-unsupported error.
|
||||
env = envFrom(map[string]string{"AWS_ACCESS_KEY_ID": "id", "AWS_SECRET_ACCESS_KEY": "secret"})
|
||||
if _, err := ResolveSpecialProvider(spec, "claude", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "SigV4") {
|
||||
t.Fatalf("expected SigV4-unsupported error for static keys, got %v", err)
|
||||
}
|
||||
|
||||
// Bearer token + default region.
|
||||
env = envFrom(map[string]string{"AWS_BEARER_TOKEN_BEDROCK": "tok"})
|
||||
p, err := ResolveSpecialProvider(spec, "claude", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://bedrock-runtime.us-east-1.amazonaws.com"; got != want {
|
||||
t.Fatalf("bedrock base_url = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// Bearer token + explicit region.
|
||||
env = envFrom(map[string]string{"AWS_BEARER_TOKEN_BEDROCK": "tok", "AWS_REGION": "eu-west-1"})
|
||||
p, err = ResolveSpecialProvider(spec, "claude", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://bedrock-runtime.eu-west-1.amazonaws.com"; got != want {
|
||||
t.Fatalf("bedrock base_url = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveSpecialProvider_GoogleVertex(t *testing.T) {
|
||||
spec := specFor(t, "google-vertex")
|
||||
|
||||
if _, err := ResolveSpecialProvider(spec, "gemini", "", envFrom(nil)); err == nil ||
|
||||
!strings.Contains(err.Error(), "GOOGLE_CLOUD_PROJECT") {
|
||||
t.Fatalf("expected GOOGLE_CLOUD_PROJECT error, got %v", err)
|
||||
}
|
||||
|
||||
env := envFrom(map[string]string{"GOOGLE_CLOUD_PROJECT": "proj"})
|
||||
if _, err := ResolveSpecialProvider(spec, "gemini", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "GOOGLE_CLOUD_LOCATION") {
|
||||
t.Fatalf("expected GOOGLE_CLOUD_LOCATION error, got %v", err)
|
||||
}
|
||||
|
||||
env = envFrom(map[string]string{"GOOGLE_CLOUD_PROJECT": "proj", "GOOGLE_CLOUD_LOCATION": "us-central1"})
|
||||
if _, err := ResolveSpecialProvider(spec, "gemini", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "GOOGLE_CLOUD_API_KEY") ||
|
||||
!strings.Contains(err.Error(), "GOOGLE_APPLICATION_CREDENTIALS") {
|
||||
t.Fatalf("expected credentials error naming both sources, got %v", err)
|
||||
}
|
||||
|
||||
// Fully configured with API key.
|
||||
env = envFrom(map[string]string{
|
||||
"GOOGLE_CLOUD_PROJECT": "proj",
|
||||
"GOOGLE_CLOUD_LOCATION": "us-central1",
|
||||
"GOOGLE_CLOUD_API_KEY": "k",
|
||||
})
|
||||
p, err := ResolveSpecialProvider(spec, "gemini", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://us-central1-aiplatform.googleapis.com"; got != want {
|
||||
t.Fatalf("vertex base_url = %q, want %q", got, want)
|
||||
}
|
||||
|
||||
// ADC credential source also satisfies.
|
||||
env = envFrom(map[string]string{
|
||||
"GOOGLE_CLOUD_PROJECT": "proj",
|
||||
"GOOGLE_CLOUD_LOCATION": "europe-west4",
|
||||
"GOOGLE_APPLICATION_CREDENTIALS": "/path/to/adc.json",
|
||||
})
|
||||
p, err = ResolveSpecialProvider(spec, "gemini", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if got, want := baseURLOf(t, p), "https://europe-west4-aiplatform.googleapis.com"; got != want {
|
||||
t.Fatalf("vertex ADC base_url = %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveSpecialProvider_CloudflareWorkersAI(t *testing.T) {
|
||||
spec := specFor(t, "cloudflare-workers-ai")
|
||||
|
||||
if _, err := ResolveSpecialProvider(spec, "m", "", envFrom(nil)); err == nil ||
|
||||
!strings.Contains(err.Error(), "CLOUDFLARE_API_KEY") {
|
||||
t.Fatalf("expected CLOUDFLARE_API_KEY error, got %v", err)
|
||||
}
|
||||
|
||||
env := envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k"})
|
||||
if _, err := ResolveSpecialProvider(spec, "m", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "CLOUDFLARE_ACCOUNT_ID") {
|
||||
t.Fatalf("expected CLOUDFLARE_ACCOUNT_ID error, got %v", err)
|
||||
}
|
||||
|
||||
env = envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k", "CLOUDFLARE_ACCOUNT_ID": "acct123"})
|
||||
p, err := ResolveSpecialProvider(spec, "m", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
want := "https://api.cloudflare.com/client/v4/accounts/acct123/ai/v1"
|
||||
if got := baseURLOf(t, p); got != want {
|
||||
t.Fatalf("workers-ai base_url = %q, want %q", got, want)
|
||||
}
|
||||
if _, ok := p.(*openAICompatDriver); !ok {
|
||||
t.Fatalf("workers-ai should speak OpenAI wire, got %T", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveSpecialProvider_CloudflareAIGateway(t *testing.T) {
|
||||
spec := specFor(t, "cloudflare-ai-gateway")
|
||||
|
||||
if _, err := ResolveSpecialProvider(spec, "m", "", envFrom(nil)); err == nil ||
|
||||
!strings.Contains(err.Error(), "CLOUDFLARE_API_KEY") {
|
||||
t.Fatalf("expected CLOUDFLARE_API_KEY error, got %v", err)
|
||||
}
|
||||
|
||||
env := envFrom(map[string]string{"CLOUDFLARE_API_KEY": "k", "CLOUDFLARE_ACCOUNT_ID": "acct123"})
|
||||
if _, err := ResolveSpecialProvider(spec, "m", "", env); err == nil ||
|
||||
!strings.Contains(err.Error(), "CLOUDFLARE_GATEWAY_ID") {
|
||||
t.Fatalf("expected CLOUDFLARE_GATEWAY_ID error, got %v", err)
|
||||
}
|
||||
|
||||
env = envFrom(map[string]string{
|
||||
"CLOUDFLARE_API_KEY": "k",
|
||||
"CLOUDFLARE_ACCOUNT_ID": "acct123",
|
||||
"CLOUDFLARE_GATEWAY_ID": "gw456",
|
||||
})
|
||||
p, err := ResolveSpecialProvider(spec, "m", "", env)
|
||||
if err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
want := "https://gateway.ai.cloudflare.com/v1/acct123/gw456/anthropic"
|
||||
if got := baseURLOf(t, p); got != want {
|
||||
t.Fatalf("ai-gateway base_url = %q, want %q", got, want)
|
||||
}
|
||||
if _, ok := p.(*anthropicCompatDriver); !ok {
|
||||
t.Fatalf("ai-gateway should speak Anthropic wire, got %T", p)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsSpecialAuthProvider(t *testing.T) {
|
||||
special := []string{
|
||||
"azure-openai-responses", "amazon-bedrock", "google-vertex",
|
||||
"cloudflare-workers-ai", "cloudflare-ai-gateway",
|
||||
}
|
||||
for _, name := range special {
|
||||
if !IsSpecialAuthProvider(specFor(t, name)) {
|
||||
t.Errorf("%s should be a special-auth provider", name)
|
||||
}
|
||||
}
|
||||
for _, name := range []string{"openai", "anthropic", "deepseek"} {
|
||||
if IsSpecialAuthProvider(specFor(t, name)) {
|
||||
t.Errorf("%s should NOT be a special-auth provider", name)
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
// Tests for reasoning/thinking wiring across both wire protocols (issues
|
||||
// #240-#244): request encoders forwarding ThinkingLevel, the Anthropic
|
||||
// thinking-block echo on multi-turn tool-use messages, OpenAI reasoning_content
|
||||
// stream decoding, the max_tokens fallback, and strict-gateway empty-content
|
||||
// handling.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// decodeBody unmarshals an encoded request body into a generic map for asserts.
|
||||
func decodeBody(t *testing.T, b []byte) map[string]any {
|
||||
t.Helper()
|
||||
var m map[string]any
|
||||
if err := json.Unmarshal(b, &m); err != nil {
|
||||
t.Fatalf("unmarshal body: %v", err)
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// TestOpenAIReasoningEffortEncoding verifies ThinkingLevel maps to
|
||||
// reasoning_effort, and that off/unset omits the field entirely (#240).
|
||||
func TestOpenAIReasoningEffortEncoding(t *testing.T) {
|
||||
base := CompletionRequest{Model: "o3-mini", Context: LlmContext{}}
|
||||
|
||||
// off / unset → no reasoning_effort.
|
||||
for _, lvl := range []agentcore.ThinkingLevel{"", agentcore.ThinkingOff} {
|
||||
base.Config.ThinkingLevel = lvl
|
||||
b, err := encodeOpenAIRequest(base)
|
||||
if err != nil {
|
||||
t.Fatalf("encode: %v", err)
|
||||
}
|
||||
if _, ok := decodeBody(t, b)["reasoning_effort"]; ok {
|
||||
t.Errorf("level %q: reasoning_effort should be omitted", lvl)
|
||||
}
|
||||
}
|
||||
|
||||
cases := map[agentcore.ThinkingLevel]string{
|
||||
agentcore.ThinkingMinimal: "minimal",
|
||||
agentcore.ThinkingLow: "low",
|
||||
agentcore.ThinkingMedium: "medium",
|
||||
agentcore.ThinkingHigh: "high",
|
||||
agentcore.ThinkingXHigh: "high",
|
||||
}
|
||||
for lvl, want := range cases {
|
||||
base.Config.ThinkingLevel = lvl
|
||||
b, err := encodeOpenAIRequest(base)
|
||||
if err != nil {
|
||||
t.Fatalf("encode: %v", err)
|
||||
}
|
||||
if got := decodeBody(t, b)["reasoning_effort"]; got != want {
|
||||
t.Errorf("level %q: reasoning_effort = %v, want %q", lvl, got, want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicThinkingEncoding verifies ThinkingLevel enables the thinking
|
||||
// block with a budget, and off/unset omits it (#240).
|
||||
func TestAnthropicThinkingEncoding(t *testing.T) {
|
||||
base := CompletionRequest{Model: "claude-x", Context: LlmContext{}}
|
||||
|
||||
base.Config.ThinkingLevel = agentcore.ThinkingOff
|
||||
b, _ := encodeAnthropicRequest(base, nil)
|
||||
if _, ok := decodeBody(t, b)["thinking"]; ok {
|
||||
t.Error("off: thinking block should be omitted")
|
||||
}
|
||||
|
||||
base.Config.ThinkingLevel = agentcore.ThinkingMedium
|
||||
b, _ = encodeAnthropicRequest(base, nil)
|
||||
th, ok := decodeBody(t, b)["thinking"].(map[string]any)
|
||||
if !ok {
|
||||
t.Fatal("medium: thinking block missing")
|
||||
}
|
||||
if th["type"] != "enabled" {
|
||||
t.Errorf("thinking.type = %v, want enabled", th["type"])
|
||||
}
|
||||
if bt, _ := th["budget_tokens"].(float64); bt <= 0 {
|
||||
t.Errorf("thinking.budget_tokens = %v, want > 0", th["budget_tokens"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicMaxTokensFallback verifies the fallback prefers the model's
|
||||
// MaxOutputTokens and otherwise uses the coding-friendly 8192 default (#243).
|
||||
func TestAnthropicMaxTokensFallback(t *testing.T) {
|
||||
req := CompletionRequest{Model: "claude-x", Context: LlmContext{}}
|
||||
|
||||
// No model metadata → 8192 default.
|
||||
b, _ := encodeAnthropicRequest(req, nil)
|
||||
if got := decodeBody(t, b)["max_tokens"].(float64); got != 8192 {
|
||||
t.Errorf("default max_tokens = %v, want 8192", got)
|
||||
}
|
||||
|
||||
// Model with a declared cap → that cap.
|
||||
models := []Model{{ID: "claude-x", MaxOutputTokens: 12000}}
|
||||
b, _ = encodeAnthropicRequest(req, models)
|
||||
if got := decodeBody(t, b)["max_tokens"].(float64); got != 12000 {
|
||||
t.Errorf("model-cap max_tokens = %v, want 12000", got)
|
||||
}
|
||||
|
||||
// Explicit Extra hint still wins.
|
||||
req.Config.Extra = map[string]any{"max_tokens": 2000}
|
||||
b, _ = encodeAnthropicRequest(req, models)
|
||||
if got := decodeBody(t, b)["max_tokens"].(float64); got != 2000 {
|
||||
t.Errorf("explicit max_tokens = %v, want 2000", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicThinkingBlockEcho verifies a multi-turn assistant message with a
|
||||
// thinking block + tool call re-emits the thinking block (with signature) ahead
|
||||
// of the tool_use block (#241).
|
||||
func TestAnthropicThinkingBlockEcho(t *testing.T) {
|
||||
think := agentcore.NewThinkingContent("let me reason")
|
||||
think.ThinkingSignature = "sig-abc"
|
||||
msg := agentcore.AssistantMessage{
|
||||
Content: agentcore.ContentList{
|
||||
think,
|
||||
agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{"path":"x"}`)),
|
||||
},
|
||||
}
|
||||
entry := encodeAnthropicMessage(msg)
|
||||
blocks, ok := entry["content"].([]map[string]any)
|
||||
if !ok || len(blocks) != 2 {
|
||||
t.Fatalf("want 2 content blocks, got %#v", entry["content"])
|
||||
}
|
||||
if blocks[0]["type"] != "thinking" {
|
||||
t.Errorf("first block type = %v, want thinking", blocks[0]["type"])
|
||||
}
|
||||
if blocks[0]["signature"] != "sig-abc" {
|
||||
t.Errorf("thinking signature = %v, want sig-abc", blocks[0]["signature"])
|
||||
}
|
||||
if blocks[1]["type"] != "tool_use" {
|
||||
t.Errorf("second block type = %v, want tool_use", blocks[1]["type"])
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicRedactedThinkingEcho verifies redacted thinking round-trips as a
|
||||
// redacted_thinking block carrying the signature as data (#241).
|
||||
func TestAnthropicRedactedThinkingEcho(t *testing.T) {
|
||||
think := agentcore.ThinkingContent{Type: agentcore.ContentTypeThinking, Redacted: true, ThinkingSignature: "redacted-data"}
|
||||
msg := agentcore.AssistantMessage{Content: agentcore.ContentList{think}}
|
||||
entry := encodeAnthropicMessage(msg)
|
||||
blocks := entry["content"].([]map[string]any)
|
||||
if blocks[0]["type"] != "redacted_thinking" || blocks[0]["data"] != "redacted-data" {
|
||||
t.Errorf("redacted block = %#v", blocks[0])
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIAssistantContentNullWithToolCalls verifies a tool-call-only
|
||||
// assistant turn sends content:null (not ""), while a text-only turn keeps its
|
||||
// text (#244).
|
||||
func TestOpenAIAssistantContentNullWithToolCalls(t *testing.T) {
|
||||
toolOnly := agentcore.AssistantMessage{
|
||||
Content: agentcore.ContentList{
|
||||
agentcore.NewToolCallContent("call_1", "read", json.RawMessage(`{}`)),
|
||||
},
|
||||
}
|
||||
entry := encodeOpenAIMessage(toolOnly)[0]
|
||||
if entry["content"] != nil {
|
||||
t.Errorf("tool-only content = %#v, want nil", entry["content"])
|
||||
}
|
||||
if _, ok := entry["tool_calls"]; !ok {
|
||||
t.Error("tool_calls missing")
|
||||
}
|
||||
|
||||
textOnly := agentcore.AssistantMessage{
|
||||
Content: agentcore.ContentList{agentcore.NewTextContent("hi")},
|
||||
}
|
||||
if got := encodeOpenAIMessage(textOnly)[0]["content"]; got != "hi" {
|
||||
t.Errorf("text content = %#v, want hi", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicEmptyAssistantNoEmptyText verifies an assistant message with no
|
||||
// usable content does not emit an empty-string text block (#244).
|
||||
func TestAnthropicEmptyAssistantNoEmptyText(t *testing.T) {
|
||||
msg := agentcore.AssistantMessage{Content: agentcore.ContentList{agentcore.NewTextContent("")}}
|
||||
entry := encodeAnthropicMessage(msg)
|
||||
blocks := entry["content"].([]map[string]any)
|
||||
if len(blocks) != 1 {
|
||||
t.Fatalf("want 1 fallback block, got %d", len(blocks))
|
||||
}
|
||||
if txt, _ := blocks[0]["text"].(string); strings.TrimSpace(txt) == "" && txt == "" {
|
||||
t.Errorf("fallback text block is empty string, want non-empty placeholder")
|
||||
}
|
||||
}
|
||||
|
||||
// TestAnthropicThinkingRaisesMaxTokens verifies max_tokens is lifted above the
|
||||
// thinking budget: Anthropic requires budget_tokens < max_tokens, so a low cap
|
||||
// must be raised to leave headroom for the visible reply (#243).
|
||||
func TestAnthropicThinkingRaisesMaxTokens(t *testing.T) {
|
||||
req := CompletionRequest{Model: "claude-x", Context: LlmContext{}}
|
||||
req.Config.ThinkingLevel = agentcore.ThinkingXHigh // budget 32768
|
||||
|
||||
// Default cap (8192) is below the budget → must be raised above it.
|
||||
b, _ := encodeAnthropicRequest(req, nil)
|
||||
body := decodeBody(t, b)
|
||||
budget := body["thinking"].(map[string]any)["budget_tokens"].(float64)
|
||||
maxTok := body["max_tokens"].(float64)
|
||||
if maxTok <= budget {
|
||||
t.Errorf("max_tokens = %v, want > budget_tokens %v", maxTok, budget)
|
||||
}
|
||||
|
||||
// A caller cap already above budget+headroom is left untouched.
|
||||
req.Config.Extra = map[string]any{"max_tokens": 100000}
|
||||
b, _ = encodeAnthropicRequest(req, nil)
|
||||
if got := decodeBody(t, b)["max_tokens"].(float64); got != 100000 {
|
||||
t.Errorf("max_tokens = %v, want caller value 100000", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestOpenAIReasoningContentDecoding verifies the decoder accumulates
|
||||
// reasoning_content into a ThinkingContent block ahead of text (#242).
|
||||
func TestOpenAIReasoningContentDecoding(t *testing.T) {
|
||||
d := NewOpenAIDecoder()
|
||||
chunks := []string{
|
||||
`{"id":"c1","choices":[{"delta":{"reasoning_content":"think "}}]}`,
|
||||
`{"choices":[{"delta":{"reasoning_content":"harder"}}]}`,
|
||||
`{"choices":[{"delta":{"content":"answer"}}]}`,
|
||||
`{"choices":[{"finish_reason":"stop"}]}`,
|
||||
}
|
||||
var events []StreamEvent
|
||||
for _, c := range chunks {
|
||||
evs, err := d.Decode([]byte(c))
|
||||
if err != nil {
|
||||
t.Fatalf("decode: %v", err)
|
||||
}
|
||||
events = append(events, evs...)
|
||||
}
|
||||
done, _ := d.Finish()
|
||||
events = append(events, done...)
|
||||
|
||||
var final agentcore.AssistantMessage
|
||||
for _, e := range events {
|
||||
if de, ok := e.(StreamDoneEvent); ok {
|
||||
final = de.Message
|
||||
}
|
||||
}
|
||||
if len(final.Content) != 2 {
|
||||
t.Fatalf("want thinking+text, got %d blocks: %#v", len(final.Content), final.Content)
|
||||
}
|
||||
th, ok := final.Content[0].(agentcore.ThinkingContent)
|
||||
if !ok || th.Thinking != "think harder" {
|
||||
t.Errorf("block[0] = %#v, want thinking 'think harder'", final.Content[0])
|
||||
}
|
||||
txt, ok := final.Content[1].(agentcore.TextContent)
|
||||
if !ok || txt.Text != "answer" {
|
||||
t.Errorf("block[1] = %#v, want text 'answer'", final.Content[1])
|
||||
}
|
||||
|
||||
sawThinking := false
|
||||
for _, e := range events {
|
||||
if _, ok := e.(StreamThinkingEvent); ok {
|
||||
sawThinking = true
|
||||
}
|
||||
}
|
||||
if !sawThinking {
|
||||
t.Error("expected at least one StreamThinkingEvent")
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,393 @@
|
||||
// This file implements the shared transport driver (US-007 support): a
|
||||
// provider-agnostic layer that turns an HTTP request into a stream of
|
||||
// StreamEvents. Each provider degenerates to a stateful Decoder; the transport
|
||||
// owns HTTP + SSE line parsing + retry + dual watchdogs + the dual failure
|
||||
// model.
|
||||
//
|
||||
// The design mirrors pi's providerio: the transport never returns a runtime
|
||||
// failure as a Go error once streaming has begun — it rides the stream as a
|
||||
// terminal StreamErrorEvent. Only the earliest "cannot build the stream" case
|
||||
// (bad request construction) is a returned error.
|
||||
package provider
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"os"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// StreamEvent is the transport-level alias for AssistantMessageEvent. Decoders
|
||||
// produce these; the transport forwards them onto the stream. (Decision #25:
|
||||
// reuse AssistantMessageEvent rather than a parallel event type.)
|
||||
type StreamEvent = AssistantMessageEvent
|
||||
|
||||
// Decoder is the per-provider stateful SSE payload decoder. The transport calls
|
||||
// Decode for every complete SSE data payload (one event's worth of bytes) and
|
||||
// Finish once the stream ends so the decoder can flush any buffered terminal
|
||||
// event.
|
||||
type Decoder interface {
|
||||
// Decode turns one SSE data payload into zero or more StreamEvents. A
|
||||
// returned error is treated as a runtime stream failure (terminal error
|
||||
// event), never a panic.
|
||||
Decode(payload []byte) ([]StreamEvent, error)
|
||||
// Finish flushes any trailing state, returning a final batch of events.
|
||||
Finish() ([]StreamEvent, error)
|
||||
}
|
||||
|
||||
// defaultIdleTimeout is the watchdog idle window; PIGO_STREAM_IDLE_TIMEOUT
|
||||
// (a Go duration string, e.g. "3m") overrides it.
|
||||
const defaultIdleTimeout = 5 * time.Minute
|
||||
|
||||
const (
|
||||
// defaultMaxConnectRetries bounds connect-only retries when TransportConfig
|
||||
// leaves MaxConnectRetries at zero.
|
||||
defaultMaxConnectRetries = 2
|
||||
// statusTooManyRequestsCF (529) is Cloudflare's "site overloaded" status,
|
||||
// which some upstreams also emit; treated as retryable alongside 429/503.
|
||||
statusTooManyRequestsCF = 529
|
||||
// stallFactor slackens the content-stall watchdog relative to the idle
|
||||
// window (stall = idle × stallFactor) so a slow-but-progressing stream is not
|
||||
// killed by the stall guard.
|
||||
stallFactor = 1.2
|
||||
// errorBodyLimit bounds how many bytes of an upstream error body are read
|
||||
// into the returned error message.
|
||||
errorBodyLimit = 4096
|
||||
)
|
||||
|
||||
// idleTimeout resolves the configured idle watchdog window.
|
||||
func idleTimeout() time.Duration {
|
||||
if v := os.Getenv("PIGO_STREAM_IDLE_TIMEOUT"); v != "" {
|
||||
if d, err := time.ParseDuration(v); err == nil && d > 0 {
|
||||
return d
|
||||
}
|
||||
}
|
||||
return defaultIdleTimeout
|
||||
}
|
||||
|
||||
// TransportConfig configures a single StreamRequest run.
|
||||
type TransportConfig struct {
|
||||
// Client is the HTTP client; defaults to http.DefaultClient when nil.
|
||||
Client *http.Client
|
||||
// NewRequest builds a fresh *http.Request for each connection attempt. It is
|
||||
// called once per connect (initial + reconnects) so retries never replay a
|
||||
// consumed body — the caller owns idempotent request construction.
|
||||
NewRequest func(ctx context.Context) (*http.Request, error)
|
||||
// Decoder converts SSE payloads to StreamEvents (required).
|
||||
Decoder Decoder
|
||||
// MaxConnectRetries bounds connect-only retries (default 2).
|
||||
MaxConnectRetries int
|
||||
}
|
||||
|
||||
// StreamRequest runs cfg as a transport stream. Per the dual failure model it
|
||||
// returns an error only when the very first request cannot be built or the
|
||||
// initial connection can never be established; every runtime failure once
|
||||
// streaming begins rides the returned stream as a terminal StreamErrorEvent.
|
||||
func StreamRequest(ctx context.Context, cfg TransportConfig) (*AssistantMessageEventStream, error) {
|
||||
if cfg.NewRequest == nil {
|
||||
return nil, errors.New("transport: NewRequest is required")
|
||||
}
|
||||
if cfg.Decoder == nil {
|
||||
return nil, errors.New("transport: Decoder is required")
|
||||
}
|
||||
client := cfg.Client
|
||||
if client == nil {
|
||||
client = http.DefaultClient
|
||||
}
|
||||
maxRetries := cfg.MaxConnectRetries
|
||||
if maxRetries == 0 {
|
||||
maxRetries = defaultMaxConnectRetries
|
||||
}
|
||||
|
||||
// Connect once up front so a "cannot even build the stream" failure surfaces
|
||||
// as a returned error (the only early-error case per FR-13).
|
||||
resp, err := connect(ctx, client, cfg.NewRequest, maxRetries)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
stream := NewAssistantMessageEventStream(0)
|
||||
go pump(ctx, stream, resp, cfg.Decoder)
|
||||
return stream, nil
|
||||
}
|
||||
|
||||
// connect performs the initial connection with retry. It only retries when the
|
||||
// server explicitly signals a retryable condition (429/503/529); it never
|
||||
// replays a consumed stream, so retrying at connect time is always safe.
|
||||
func connect(ctx context.Context, client *http.Client, newReq func(context.Context) (*http.Request, error), maxRetries int) (*http.Response, error) {
|
||||
var lastErr error
|
||||
for attempt := 0; attempt <= maxRetries; attempt++ {
|
||||
req, err := newReq(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("transport: build request: %w", err)
|
||||
}
|
||||
resp, err := client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = classifyTransportError(err)
|
||||
if !isRetryableNetErr(err) || attempt == maxRetries {
|
||||
return nil, lastErr
|
||||
}
|
||||
if !sleepBackoff(ctx, attempt, 0) {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode == http.StatusTooManyRequests ||
|
||||
resp.StatusCode == http.StatusServiceUnavailable ||
|
||||
resp.StatusCode == statusTooManyRequestsCF {
|
||||
wait := retryAfter(resp.Header)
|
||||
resp.Body.Close()
|
||||
lastErr = fmt.Errorf("transport: upstream %d", resp.StatusCode)
|
||||
if attempt == maxRetries {
|
||||
return nil, lastErr
|
||||
}
|
||||
if !sleepBackoff(ctx, attempt, wait) {
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
continue
|
||||
}
|
||||
if resp.StatusCode >= 400 {
|
||||
body, _ := io.ReadAll(io.LimitReader(resp.Body, errorBodyLimit))
|
||||
resp.Body.Close()
|
||||
return nil, fmt.Errorf("transport: upstream %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
return nil, lastErr
|
||||
}
|
||||
|
||||
// pump drives the SSE read loop with dual watchdogs, decoding payloads and
|
||||
// forwarding events onto the stream. All runtime failures become a terminal
|
||||
// error event; pump always closes the stream.
|
||||
func pump(ctx context.Context, stream *AssistantMessageEventStream, resp *http.Response, dec Decoder) {
|
||||
defer stream.Close()
|
||||
defer resp.Body.Close()
|
||||
|
||||
idle := idleTimeout()
|
||||
// content-stall watchdog is slightly slacker than idle (idle × stallFactor)
|
||||
// so a slow but progressing stream is not killed by the stall guard.
|
||||
stall := time.Duration(float64(idle) * stallFactor)
|
||||
|
||||
// The watchdog fires by cancelling a derived context; reads race against it.
|
||||
watchCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
// done stops the reader goroutine so it never blocks on a send after pump
|
||||
// returns (watchdog / abort paths), avoiding a goroutine leak.
|
||||
done := make(chan struct{})
|
||||
defer close(done)
|
||||
lines := make(chan string)
|
||||
readErr := make(chan error, 1)
|
||||
go readLines(resp.Body, lines, readErr, done)
|
||||
|
||||
var dataBuf strings.Builder
|
||||
idleTimer := time.NewTimer(idle)
|
||||
stallTimer := time.NewTimer(stall)
|
||||
defer idleTimer.Stop()
|
||||
defer stallTimer.Stop()
|
||||
|
||||
emit := func(events []StreamEvent) bool {
|
||||
for _, ev := range events {
|
||||
if err := stream.Emit(watchCtx, ev); err != nil {
|
||||
return false
|
||||
}
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
fail := func(msg string, err error) {
|
||||
stream.Emit(context.Background(), StreamErrorEvent{
|
||||
Message: agentcore.AssistantMessage{
|
||||
RoleField: agentcore.RoleAssistant,
|
||||
StopReason: agentcore.StopReasonError,
|
||||
ErrorMessage: msg,
|
||||
},
|
||||
Err: err,
|
||||
})
|
||||
}
|
||||
|
||||
flush := func() bool {
|
||||
if dataBuf.Len() == 0 {
|
||||
return true
|
||||
}
|
||||
payload := dataBuf.String()
|
||||
dataBuf.Reset()
|
||||
if payload == "[DONE]" {
|
||||
return true
|
||||
}
|
||||
events, err := dec.Decode([]byte(payload))
|
||||
if err != nil {
|
||||
fail("decode error: "+err.Error(), err)
|
||||
return false
|
||||
}
|
||||
return emit(events)
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
fail("stream aborted", ctx.Err())
|
||||
return
|
||||
case <-idleTimer.C:
|
||||
fail("idle timeout: no data received", errStreamIdle)
|
||||
return
|
||||
case <-stallTimer.C:
|
||||
fail("content stall timeout", errStreamStall)
|
||||
return
|
||||
case err := <-readErr:
|
||||
if err != nil && !errors.Is(err, io.EOF) {
|
||||
fail("read error: "+classifyTransportError(err).Error(), err)
|
||||
return
|
||||
}
|
||||
// Clean EOF: flush any buffered payload, then finish the decoder.
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
finalEvents, ferr := dec.Finish()
|
||||
if ferr != nil {
|
||||
fail("finish error: "+ferr.Error(), ferr)
|
||||
return
|
||||
}
|
||||
emit(finalEvents)
|
||||
return
|
||||
case line, ok := <-lines:
|
||||
if !ok {
|
||||
continue
|
||||
}
|
||||
// Any byte resets the idle watchdog; a flushed event resets stall.
|
||||
resetTimer(idleTimer, idle)
|
||||
line = strings.TrimRight(line, "\r\n")
|
||||
switch {
|
||||
case line == "":
|
||||
// Blank line = event boundary: flush accumulated data.
|
||||
if !flush() {
|
||||
return
|
||||
}
|
||||
resetTimer(stallTimer, stall)
|
||||
case strings.HasPrefix(line, ":"):
|
||||
// Comment / keep-alive: ignore payload, watchdog already reset.
|
||||
case strings.HasPrefix(line, "data:"):
|
||||
dataBuf.WriteString(strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
||||
default:
|
||||
// Non-data field (event:, id:, etc.) — ignored for our decoders.
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// readLines reads the body line by line, pushing each onto lines and the final
|
||||
// error (io.EOF on clean close) onto readErr. It stops promptly when done is
|
||||
// closed so pump can return on a watchdog/abort without leaking this goroutine.
|
||||
func readLines(r io.Reader, lines chan<- string, readErr chan<- error, done <-chan struct{}) {
|
||||
br := bufio.NewReader(r)
|
||||
for {
|
||||
line, err := br.ReadString('\n')
|
||||
if line != "" {
|
||||
select {
|
||||
case lines <- line:
|
||||
case <-done:
|
||||
return
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
select {
|
||||
case readErr <- err:
|
||||
case <-done:
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// resetTimer stops and re-arms t to fire after d.
|
||||
func resetTimer(t *time.Timer, d time.Duration) {
|
||||
if !t.Stop() {
|
||||
select {
|
||||
case <-t.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
t.Reset(d)
|
||||
}
|
||||
|
||||
// Sentinel errors for watchdog classification.
|
||||
var (
|
||||
errStreamIdle = errors.New("stream idle timeout")
|
||||
errStreamStall = errors.New("stream content stall")
|
||||
)
|
||||
|
||||
// retryAfter parses a Retry-After header (seconds or HTTP-date), returning 0
|
||||
// when absent/unparseable.
|
||||
func retryAfter(h http.Header) time.Duration {
|
||||
v := h.Get("Retry-After")
|
||||
if v == "" {
|
||||
return 0
|
||||
}
|
||||
if secs, err := strconv.Atoi(v); err == nil && secs >= 0 {
|
||||
return time.Duration(secs) * time.Second
|
||||
}
|
||||
if t, err := http.ParseTime(v); err == nil {
|
||||
if d := time.Until(t); d > 0 {
|
||||
return d
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// sleepBackoff waits for the retry delay (Retry-After if given, else
|
||||
// exponential), honoring ctx cancellation. Returns false if ctx was cancelled.
|
||||
func sleepBackoff(ctx context.Context, attempt int, retryAfter time.Duration) bool {
|
||||
d := retryAfter
|
||||
if d == 0 {
|
||||
d = time.Duration(1<<uint(attempt)) * time.Second
|
||||
}
|
||||
t := time.NewTimer(d)
|
||||
defer t.Stop()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
case <-t.C:
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// isRetryableNetErr reports whether a client.Do error is a transient network
|
||||
// condition worth reconnecting for (timeout / temporary).
|
||||
func isRetryableNetErr(err error) bool {
|
||||
var netErr net.Error
|
||||
if errors.As(err, &netErr) {
|
||||
return netErr.Timeout()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// classifyTransportError maps a low-level error to a typed, descriptive error
|
||||
// using net.Error / errors.Is classification.
|
||||
func classifyTransportError(err error) error {
|
||||
if err == nil {
|
||||
return nil
|
||||
}
|
||||
if errors.Is(err, context.Canceled) {
|
||||
return fmt.Errorf("transport: canceled: %w", err)
|
||||
}
|
||||
if errors.Is(err, context.DeadlineExceeded) {
|
||||
return fmt.Errorf("transport: deadline exceeded: %w", err)
|
||||
}
|
||||
var netErr net.Error
|
||||
if errors.As(err, &netErr) {
|
||||
if netErr.Timeout() {
|
||||
return fmt.Errorf("transport: network timeout: %w", err)
|
||||
}
|
||||
return fmt.Errorf("transport: network error: %w", err)
|
||||
}
|
||||
return fmt.Errorf("transport: %w", err)
|
||||
}
|
||||
@@ -0,0 +1,251 @@
|
||||
package provider
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/smallnest/pigo/internal/agentcore"
|
||||
)
|
||||
|
||||
// jsonDecoder is a trivial Decoder: each payload is a JSON object with a "text"
|
||||
// field yielding a StreamTextEvent, or {"done":true} yielding a StreamDoneEvent.
|
||||
type jsonDecoder struct {
|
||||
finished bool
|
||||
}
|
||||
|
||||
func (d *jsonDecoder) Decode(payload []byte) ([]StreamEvent, error) {
|
||||
var m struct {
|
||||
Text string `json:"text"`
|
||||
Done bool `json:"done"`
|
||||
Bad bool `json:"bad"`
|
||||
}
|
||||
if err := json.Unmarshal(payload, &m); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if m.Bad {
|
||||
return nil, fmt.Errorf("decoder rejected payload")
|
||||
}
|
||||
if m.Done {
|
||||
return []StreamEvent{StreamDoneEvent{Message: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant, StopReason: agentcore.StopReasonEndTurn}}}, nil
|
||||
}
|
||||
return []StreamEvent{StreamTextEvent{Partial: agentcore.AssistantMessage{RoleField: agentcore.RoleAssistant}}}, nil
|
||||
}
|
||||
|
||||
func (d *jsonDecoder) Finish() ([]StreamEvent, error) {
|
||||
d.finished = true
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// sseServer returns an httptest server that writes the given SSE body.
|
||||
func sseServer(t *testing.T, body string) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
if _, err := w.Write([]byte(body)); err != nil {
|
||||
t.Errorf("server write: %v", err)
|
||||
}
|
||||
}))
|
||||
}
|
||||
|
||||
func newReqFn(url string) func(context.Context) (*http.Request, error) {
|
||||
return func(ctx context.Context) (*http.Request, error) {
|
||||
return http.NewRequestWithContext(ctx, http.MethodPost, url, strings.NewReader("{}"))
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportSSEParsing verifies data accumulation, blank-line flush, [DONE]
|
||||
// discard, and ":" keep-alive handling.
|
||||
func TestTransportSSEParsing(t *testing.T) {
|
||||
body := ": keep-alive comment\n" +
|
||||
"data: {\"text\":\"hi\"}\n" +
|
||||
"\n" +
|
||||
"data: [DONE]\n" +
|
||||
"\n" +
|
||||
"data: {\"done\":true}\n" +
|
||||
"\n"
|
||||
srv := sseServer(t, body)
|
||||
defer srv.Close()
|
||||
|
||||
dec := &jsonDecoder{}
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: dec,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamRequest: %v", err)
|
||||
}
|
||||
|
||||
var kinds []string
|
||||
for ev := range stream.Events() {
|
||||
kinds = append(kinds, ev.EventKind())
|
||||
}
|
||||
final, resErr := stream.Result(context.Background())
|
||||
if resErr != nil {
|
||||
t.Fatalf("result error: %v", resErr)
|
||||
}
|
||||
// text (from first data) then done; [DONE] payload must be discarded.
|
||||
if len(kinds) != 2 || kinds[0] != StreamEventText || kinds[1] != StreamEventDone {
|
||||
t.Errorf("event kinds = %v, want [text done]", kinds)
|
||||
}
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("final stop reason = %q, want end_turn", final.StopReason)
|
||||
}
|
||||
if !dec.finished {
|
||||
t.Errorf("decoder Finish() was not called on clean EOF")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportDecodeErrorRidesStream confirms a decode failure becomes a
|
||||
// terminal error event, not a returned error.
|
||||
func TestTransportDecodeErrorRidesStream(t *testing.T) {
|
||||
body := "data: {\"bad\":true}\n\n"
|
||||
srv := sseServer(t, body)
|
||||
defer srv.Close()
|
||||
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: &jsonDecoder{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("decode failure must NOT be a returned error: %v", err)
|
||||
}
|
||||
final, _ := stream.Result(context.Background())
|
||||
if final.StopReason != agentcore.StopReasonError {
|
||||
t.Errorf("expected terminal error message, got stopReason=%q", final.StopReason)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportEarlyBuildFailure verifies a request-build failure is a returned
|
||||
// error (the only early-error case).
|
||||
func TestTransportEarlyBuildFailure(t *testing.T) {
|
||||
_, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: func(ctx context.Context) (*http.Request, error) {
|
||||
return nil, fmt.Errorf("cannot build")
|
||||
},
|
||||
Decoder: &jsonDecoder{},
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("request-build failure must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportMissingConfig checks required-field validation.
|
||||
func TestTransportMissingConfig(t *testing.T) {
|
||||
if _, err := StreamRequest(context.Background(), TransportConfig{Decoder: &jsonDecoder{}}); err == nil {
|
||||
t.Error("missing NewRequest must error")
|
||||
}
|
||||
if _, err := StreamRequest(context.Background(), TransportConfig{NewRequest: newReqFn("http://x")}); err == nil {
|
||||
t.Error("missing Decoder must error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportRetryOn503 verifies the connect retry path honors a retryable
|
||||
// status and eventually succeeds without replaying a consumed stream.
|
||||
func TestTransportRetryOn503(t *testing.T) {
|
||||
t.Setenv("PIGO_STREAM_IDLE_TIMEOUT", "5s")
|
||||
var attempts atomic.Int32
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if attempts.Add(1) == 1 {
|
||||
w.Header().Set("Retry-After", "0")
|
||||
w.WriteHeader(http.StatusServiceUnavailable)
|
||||
return
|
||||
}
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
w.Write([]byte("data: {\"done\":true}\n\n"))
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: &jsonDecoder{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("retry should succeed: %v", err)
|
||||
}
|
||||
final, _ := stream.Result(context.Background())
|
||||
if final.StopReason != agentcore.StopReasonEndTurn {
|
||||
t.Errorf("final stop reason = %q, want end_turn", final.StopReason)
|
||||
}
|
||||
if got := attempts.Load(); got != 2 {
|
||||
t.Errorf("expected 2 attempts (1 failed + 1 ok), got %d", got)
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportRetryExhausted verifies a persistently retryable status returns
|
||||
// an early error after exhausting retries.
|
||||
func TestTransportRetryExhausted(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Retry-After", "0")
|
||||
w.WriteHeader(http.StatusTooManyRequests)
|
||||
}))
|
||||
defer srv.Close()
|
||||
|
||||
_, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: &jsonDecoder{},
|
||||
MaxConnectRetries: 1,
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("exhausted retries must return an error")
|
||||
}
|
||||
}
|
||||
|
||||
// TestTransportIdleWatchdog verifies the idle watchdog fires a terminal error
|
||||
// when the server stalls without sending data.
|
||||
func TestTransportIdleWatchdog(t *testing.T) {
|
||||
t.Setenv("PIGO_STREAM_IDLE_TIMEOUT", "100ms")
|
||||
hold := make(chan struct{})
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
if f, ok := w.(http.Flusher); ok {
|
||||
f.Flush()
|
||||
}
|
||||
<-hold // never send any data
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer close(hold)
|
||||
|
||||
stream, err := StreamRequest(context.Background(), TransportConfig{
|
||||
NewRequest: newReqFn(srv.URL),
|
||||
Decoder: &jsonDecoder{},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("StreamRequest: %v", err)
|
||||
}
|
||||
|
||||
done := make(chan agentcore.AssistantMessage, 1)
|
||||
go func() {
|
||||
final, _ := stream.Result(context.Background())
|
||||
done <- final
|
||||
}()
|
||||
select {
|
||||
case final := <-done:
|
||||
if final.StopReason != agentcore.StopReasonError {
|
||||
t.Errorf("idle watchdog must produce terminal error, got %q", final.StopReason)
|
||||
}
|
||||
case <-time.After(3 * time.Second):
|
||||
t.Fatal("idle watchdog did not fire")
|
||||
}
|
||||
}
|
||||
|
||||
// TestRetryAfterParsing covers seconds and absent header parsing.
|
||||
func TestRetryAfterParsing(t *testing.T) {
|
||||
h := http.Header{}
|
||||
if d := retryAfter(h); d != 0 {
|
||||
t.Errorf("absent Retry-After = %v, want 0", d)
|
||||
}
|
||||
h.Set("Retry-After", "3")
|
||||
if d := retryAfter(h); d != 3*time.Second {
|
||||
t.Errorf("Retry-After 3 = %v, want 3s", d)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user