feat(model): 双协议 LLM Provider(AnthropicChNativeAdapter /v1/messages + tool_use.input→Args 直接映射 + 流式重组)+ Provider 感知工厂/路由 + LLM_PROVIDER/anthropic 配置(OpenAI 兼容保留)
This commit is contained in:
1 parent
f9650185f1
commit
3632e04733
11 files changed
+1452
-18
No files matched your search
+9
-1
@@ -30,12 +30,20 @@ DB_SQLITE_PATH=./data/cyrene.db
|
||||
# PostgreSQL + pgvector 为可选增强;未配置时走 SQLite。相关向量能力在 M4 接入。
|
||||
# DB_PG_URL=postgres://cyrene:cyrene_pass@localhost:5432/cyrene_ai?sslmode=disable
|
||||
|
||||
# ---- LLM 多供应商 ----
|
||||
# ---- LLM 多供应商(LLM_PROVIDER 显式声明协议:openai / anthropic,默认 openai) ----
|
||||
# Provider = openai:OpenAI 兼容 /chat/completions(含 OpenAI 兼容网关,也可经它承载 Anthropic 模型,但失去原生 tool input 对象优势)
|
||||
LLM_PROVIDER=openai
|
||||
LLM_BASE_URL=https://api.openai.com/v1
|
||||
LLM_API_KEY=sk-xxxxx
|
||||
LLM_MODEL=gpt-4o
|
||||
LLM_FALLBACK_MODEL=gpt-4o-mini
|
||||
|
||||
# Provider = anthropic:Anthropic Messages /v1/messages(tool_use.input 为结构化对象,工具循环更稳)
|
||||
# ANTHROPIC_BASE_URL=https://api.anthropic.com
|
||||
# ANTHROPIC_API_KEY=sk-ant-xxxxx
|
||||
# ANTHROPIC_MODEL=claude-sonnet-4-20250514
|
||||
# ANTHROPIC_FALLBACK_MODEL=claude-haiku-3-5
|
||||
|
||||
# ---- 记忆嵌入(可选) ----
|
||||
# 不配置时使用内置 n-gram / FTS 向量。配置后必须真实走语义路径(M4 接入)。
|
||||
# EMBEDDING_API_URL=http://localhost:11434/v1
|
||||
|
||||
+2
-2
@@ -246,12 +246,12 @@ func (a *app) wireMemory(ctx context.Context) error {
|
||||
// wireRuntime builds the model adapter, tool dispatcher, guardrails, the agent
|
||||
// and its loop, tracing, persistence and the orchestration scheduler.
|
||||
func (a *app) wireRuntime() error {
|
||||
adapter, aerr := model.NewOpenAIChatChain(a.cfg.LLM)
|
||||
adapter, aerr := model.NewChatChain(a.cfg.LLM)
|
||||
if aerr != nil {
|
||||
a.log.Warn("llm adapter unavailable", "error", aerr)
|
||||
adapter = nil
|
||||
}
|
||||
router := model.NewRouter(a.cfg.LLM.Model)
|
||||
router := model.NewRouter(model.DefaultModelForProvider(a.cfg.LLM))
|
||||
|
||||
dispatcher, derr := tool.New(a.caps, tool.WithDefaultTimeout(defaultToolTimeout))
|
||||
if derr != nil {
|
||||
|
||||
@@ -54,12 +54,27 @@ type DBConfig struct {
|
||||
PGURL string `env:"DB_PG_URL"`
|
||||
}
|
||||
|
||||
// LLMConfig holds the multi-supplier LLM adapter settings.
|
||||
// LLMConfig holds the multi-supplier LLM adapter settings. Provider selects the
|
||||
// wire protocol; every provider has its own independent base/key/model/fallback
|
||||
// so multiple vendors can coexist.
|
||||
type LLMConfig struct {
|
||||
// Provider declares which protocol to use: "openai" (OpenAI-compatible
|
||||
// /chat/completions) or "anthropic" (Anthropic Messages /v1/messages).
|
||||
Provider string `env:"LLM_PROVIDER" default:"openai"`
|
||||
|
||||
// OpenAI-compatible provider (default). Kept flat for backward compatibility
|
||||
// with the historic LLM_BASE_URL/API_KEY/MODEL/FALLBACK_MODEL env names.
|
||||
BaseURL string `env:"LLM_BASE_URL"`
|
||||
APIKey string `env:"LLM_API_KEY"`
|
||||
Model string `env:"LLM_MODEL"`
|
||||
Fallback string `env:"LLM_FALLBACK_MODEL"`
|
||||
|
||||
// Anthropic native provider, used when Provider=anthropic. Independent
|
||||
// config block so each vendor is a self-contained config.
|
||||
AnthropicBaseURL string `env:"ANTHROPIC_BASE_URL"`
|
||||
AnthropicAPIKey string `env:"ANTHROPIC_API_KEY"`
|
||||
AnthropicModel string `env:"ANTHROPIC_MODEL"`
|
||||
AnthropicFallback string `env:"ANTHROPIC_FALLBACK_MODEL"`
|
||||
}
|
||||
|
||||
// EmbeddingConfig holds the optional semantic embedding API settings. When
|
||||
|
||||
@@ -0,0 +1,445 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// Anthropic Messages (/v1/messages) wire protocol constants.
|
||||
const (
|
||||
// anthropicAPIVersion is the Messages API version sent on every request.
|
||||
anthropicAPIVersion = "2023-06-01"
|
||||
// anthropicDefaultMaxTokens is the default max_tokens sent when
|
||||
// Request.MaxTokens is zero. Anthropic requires max_tokens to be present and
|
||||
// positive, so a sensible default is applied rather than omitting it.
|
||||
anthropicDefaultMaxTokens = 1024
|
||||
// anthropicDataPrefix marks a payload embedded as a data: URI.
|
||||
anthropicDataPrefix = "data:"
|
||||
)
|
||||
|
||||
// anthropicRequest is the body sent to /v1/messages. Messages carry Anthropic's
|
||||
// content-block model and the system prompt lives at the request root.
|
||||
type anthropicRequest struct {
|
||||
Model string `json:"model"`
|
||||
System string `json:"system,omitempty"`
|
||||
Messages []anthropicMessage `json:"messages"`
|
||||
Tools []anthropicTool `json:"tools,omitempty"`
|
||||
ToolChoice any `json:"tool_choice,omitempty"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicMessage is one conversation turn. Content is an array of content
|
||||
// blocks (text / image for now; tool_result blocks are added by M2).
|
||||
type anthropicMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content any `json:"content"`
|
||||
}
|
||||
|
||||
// anthropicTool is the wire form of a tool schema.
|
||||
type anthropicTool struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
InputSchema any `json:"input_schema"`
|
||||
}
|
||||
|
||||
// anthropicResponse is the JSON body returned by /v1/messages.
|
||||
type anthropicResponse struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Content []anthropicBlock `json:"content"`
|
||||
Usage *anthropicUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicBlock is a single content block in a response. For tool_use blocks the
|
||||
// Input field is a fully decoded object (never a JSON string).
|
||||
type anthropicBlock struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text,omitempty"`
|
||||
ID string `json:"id,omitempty"`
|
||||
Name string `json:"name,omitempty"`
|
||||
Input map[string]any `json:"input,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicUsage reports token consumption. Anthropic reports input_tokens and
|
||||
// output_tokens rather than a single total.
|
||||
type anthropicUsage struct {
|
||||
InputTokens int `json:"input_tokens"`
|
||||
OutputTokens int `json:"output_tokens"`
|
||||
}
|
||||
|
||||
// AnthropicChatAdapter calls an Anthropic Messages (/v1/messages) endpoint. It is
|
||||
// constructed from a config.LLMConfig and implements Adapter for both
|
||||
// non-streaming and token-level SSE streaming, speaking the provider's native
|
||||
// content-block protocol so tool_use.input arrives as a structured object.
|
||||
type AnthropicChatAdapter struct {
|
||||
name string
|
||||
baseURL string
|
||||
apiKey string
|
||||
model string
|
||||
client *http.Client
|
||||
recorder Recorder
|
||||
}
|
||||
|
||||
// AnthropicOption configures an AnthropicChatAdapter.
|
||||
type AnthropicOption func(*AnthropicChatAdapter)
|
||||
|
||||
// WithAnthropicHTTPClient overrides the default HTTP client (used to inject mocks
|
||||
// or a tuned transport in tests and production).
|
||||
func WithAnthropicHTTPClient(c *http.Client) AnthropicOption {
|
||||
return func(a *AnthropicChatAdapter) { a.client = c }
|
||||
}
|
||||
|
||||
// WithAnthropicRecorder wires a per-call Recorder for CallRecord accounting.
|
||||
func WithAnthropicRecorder(r Recorder) AnthropicOption {
|
||||
return func(a *AnthropicChatAdapter) { a.recorder = r }
|
||||
}
|
||||
|
||||
// WithAnthropicModel overrides the adapter's default model name.
|
||||
func WithAnthropicModel(m string) AnthropicOption {
|
||||
return func(a *AnthropicChatAdapter) { a.model = m }
|
||||
}
|
||||
|
||||
// NewAnthropicChat builds an Anthropic adapter from a supplier config. baseURL is
|
||||
// trimmed of any trailing slash before the /v1/messages path is appended.
|
||||
func NewAnthropicChat(cfg config.LLMConfig, opts ...AnthropicOption) *AnthropicChatAdapter {
|
||||
a := &AnthropicChatAdapter{
|
||||
name: "anthropic",
|
||||
baseURL: strings.TrimRight(cfg.AnthropicBaseURL, "/"),
|
||||
apiKey: cfg.AnthropicAPIKey,
|
||||
model: cfg.AnthropicModel,
|
||||
client: &http.Client{
|
||||
Transport: http.DefaultTransport,
|
||||
Timeout: defaultHTTPTimeout,
|
||||
},
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(a)
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// Name returns the adapter identifier.
|
||||
func (a *AnthropicChatAdapter) Name() string { return a.name }
|
||||
|
||||
// Chat performs a non-streaming completion.
|
||||
func (a *AnthropicChatAdapter) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
|
||||
start := time.Now()
|
||||
if a.baseURL == "" {
|
||||
return Response{}, ErrUpstream
|
||||
}
|
||||
ctx, cancel := withCallTimeout(ctx, opts)
|
||||
defer cancel()
|
||||
model, err := a.resolveModel(req)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Model = model
|
||||
body, err := a.buildRequest(req, false)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
|
||||
}
|
||||
resp, err := a.doRequest(httpReq, false)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return Response{}, a.upstreamError(resp)
|
||||
}
|
||||
var wire anthropicResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&wire); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: decode response: %v", ErrUpstream, err)
|
||||
}
|
||||
out := responseFromAnthropicWire(&wire)
|
||||
a.record(req, start, out.TotalTokens)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ChatStream streams content-block deltas via onChunk and returns the aggregated
|
||||
// response with usage when the provider reports it.
|
||||
func (a *AnthropicChatAdapter) ChatStream(ctx context.Context, req Request, onChunk func(delta string)) (Response, error) {
|
||||
start := time.Now()
|
||||
if a.baseURL == "" {
|
||||
return Response{}, ErrUpstream
|
||||
}
|
||||
model, err := a.resolveModel(req)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
req.Model = model
|
||||
body, err := a.buildRequest(req, true)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
|
||||
if err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
|
||||
}
|
||||
resp, err := a.doRequest(httpReq, true)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return Response{}, a.upstreamError(resp)
|
||||
}
|
||||
out, err := scanAnthropicSSE(resp.Body, onChunk)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
out.Model = model
|
||||
a.record(req, start, out.TotalTokens)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// resolveModel returns the effective model name.
|
||||
func (a *AnthropicChatAdapter) resolveModel(req Request) (string, error) {
|
||||
if req.Model != "" {
|
||||
return req.Model, nil
|
||||
}
|
||||
if a.model != "" {
|
||||
return a.model, nil
|
||||
}
|
||||
return "", ErrNoModel
|
||||
}
|
||||
|
||||
// endpoint returns the full /v1/messages URL.
|
||||
func (a *AnthropicChatAdapter) endpoint() string {
|
||||
return a.baseURL + "/v1/messages"
|
||||
}
|
||||
|
||||
func (a *AnthropicChatAdapter) buildRequest(req Request, stream bool) ([]byte, error) {
|
||||
system, msgs, err := buildAnthropicMessages(req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
tokenLimit := req.MaxTokens
|
||||
if tokenLimit == 0 {
|
||||
tokenLimit = anthropicDefaultMaxTokens
|
||||
}
|
||||
wire := anthropicRequest{
|
||||
Model: req.Model,
|
||||
System: system,
|
||||
Messages: msgs,
|
||||
Tools: buildAnthropicTools(req.Tools),
|
||||
ToolChoice: req.ToolChoice,
|
||||
Temperature: req.Temperature,
|
||||
MaxTokens: tokenLimit,
|
||||
Stream: stream,
|
||||
}
|
||||
buf, err := json.Marshal(wire)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("model: %w: marshal request: %v", ErrBadRequest, err)
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
func (a *AnthropicChatAdapter) doRequest(httpReq *http.Request, stream bool) (*http.Response, error) {
|
||||
if a.apiKey != "" {
|
||||
httpReq.Header.Set("x-api-key", a.apiKey)
|
||||
}
|
||||
httpReq.Header.Set("anthropic-version", anthropicAPIVersion)
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
if stream {
|
||||
httpReq.Header.Set("Accept", "text/event-stream")
|
||||
} else {
|
||||
httpReq.Header.Set("Accept", "application/json")
|
||||
}
|
||||
resp, err := a.client.Do(httpReq)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("model: %w: %v", ErrUpstream, err)
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (a *AnthropicChatAdapter) upstreamError(resp *http.Response) error {
|
||||
snippet, _ := io.ReadAll(io.LimitReader(resp.Body, errBodyCap))
|
||||
return fmt.Errorf("model: %w: status %d: %s", ErrUpstream, resp.StatusCode, strings.TrimSpace(string(snippet)))
|
||||
}
|
||||
|
||||
func (a *AnthropicChatAdapter) record(req Request, start time.Time, tokens int) {
|
||||
if a.recorder == nil {
|
||||
return
|
||||
}
|
||||
a.recorder(CallRecord{
|
||||
Model: req.Model,
|
||||
Purpose: req.Purpose,
|
||||
Tokens: tokens,
|
||||
TookMs: time.Since(start).Milliseconds(),
|
||||
})
|
||||
}
|
||||
|
||||
// buildAnthropicMessages splits a normalized message list into a top-level system
|
||||
// string and the remaining messages in Anthropic's content-block form. All
|
||||
// RoleSystem messages are concatenated into a single system prompt (Anthropic
|
||||
// carries it at the request root, not in messages); messages with no content
|
||||
// parts are dropped because Anthropic rejects empty content.
|
||||
func buildAnthropicMessages(msgs []Message) (string, []anthropicMessage, error) {
|
||||
var sys strings.Builder
|
||||
out := make([]anthropicMessage, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
if m.Role == RoleSystem {
|
||||
if t := messageText(m); t != "" {
|
||||
if sys.Len() > 0 {
|
||||
sys.WriteString("\n")
|
||||
}
|
||||
sys.WriteString(t)
|
||||
}
|
||||
continue
|
||||
}
|
||||
blocks, err := encodeAnthropicContent(m.Content)
|
||||
if err != nil {
|
||||
return "", nil, err
|
||||
}
|
||||
if len(blocks) == 0 {
|
||||
continue
|
||||
}
|
||||
out = append(out, anthropicMessage{Role: anthropicRole(m.Role), Content: blocks})
|
||||
}
|
||||
return sys.String(), out, nil
|
||||
}
|
||||
|
||||
// messageText concatenates the text parts of a message (used for system prompts).
|
||||
func messageText(m Message) string {
|
||||
var b strings.Builder
|
||||
for _, part := range m.Content {
|
||||
if part.Type == ContentTypeText {
|
||||
b.WriteString(part.Text)
|
||||
}
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// anthropicRole maps a normalized Role to the Anthropic role vocabulary. system
|
||||
// messages are extracted to the top-level system field and never reach here.
|
||||
func anthropicRole(r Role) string {
|
||||
switch r {
|
||||
case RoleAssistant:
|
||||
return "assistant"
|
||||
default:
|
||||
// user, tool and any future roles all land on the user turn; tool_result
|
||||
// blocks (M2) belong to the user role in Anthropic.
|
||||
return "user"
|
||||
}
|
||||
}
|
||||
|
||||
// encodeAnthropicContent turns normalized content parts into Anthropic content
|
||||
// blocks.
|
||||
func encodeAnthropicContent(parts []Content) ([]any, error) {
|
||||
blocks := make([]any, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case ContentTypeText:
|
||||
blocks = append(blocks, map[string]any{"type": "text", "text": p.Text})
|
||||
case ContentTypeImage:
|
||||
mediaType, data, err := parseAnthropicImageData(p.ImageURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
blocks = append(blocks, map[string]any{
|
||||
"type": "image",
|
||||
"source": map[string]any{
|
||||
"type": "base64",
|
||||
"media_type": mediaType,
|
||||
"data": data,
|
||||
},
|
||||
})
|
||||
default:
|
||||
return nil, fmt.Errorf("model: %w: unsupported content type %q", ErrBadRequest, p.Type)
|
||||
}
|
||||
}
|
||||
return blocks, nil
|
||||
}
|
||||
|
||||
// parseAnthropicImageData extracts the media type and base64 payload from a data:
|
||||
// URI. Anthropic requires a base64 image source and cannot reference a remote
|
||||
// URL, so any other URI form is rejected as a bad request.
|
||||
func parseAnthropicImageData(uri string) (string, string, error) {
|
||||
if !strings.HasPrefix(uri, anthropicDataPrefix) {
|
||||
return "", "", fmt.Errorf("model: %w: anthropic image requires a data: URI, got %q", ErrBadRequest, uri)
|
||||
}
|
||||
rest := strings.TrimPrefix(uri, anthropicDataPrefix)
|
||||
comma := strings.Index(rest, ",")
|
||||
if comma < 0 {
|
||||
return "", "", fmt.Errorf("model: %w: malformed data: URI %q", ErrBadRequest, uri)
|
||||
}
|
||||
meta := rest[:comma]
|
||||
data := rest[comma+1:]
|
||||
mediaType := meta
|
||||
if idx := strings.Index(meta, ";"); idx >= 0 {
|
||||
mediaType = meta[:idx]
|
||||
}
|
||||
if mediaType == "" {
|
||||
return "", "", fmt.Errorf("model: %w: data: URI missing media type in %q", ErrBadRequest, uri)
|
||||
}
|
||||
return mediaType, data, nil
|
||||
}
|
||||
|
||||
// buildAnthropicTools converts ToolSpec schemas into Anthropic tools.
|
||||
func buildAnthropicTools(tools []ToolSpec) []anthropicTool {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]anthropicTool, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
out = append(out, anthropicTool{
|
||||
Name: t.Name,
|
||||
Description: t.Description,
|
||||
InputSchema: t.Parameters,
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// responseFromAnthropicWire converts a non-streaming /v1/messages response into a
|
||||
// normalized Response. Each text block contributes to the assistant content; each
|
||||
// tool_use block becomes a ToolCall whose Args is taken directly from the decoded
|
||||
// input object.
|
||||
func responseFromAnthropicWire(wire *anthropicResponse) Response {
|
||||
out := Response{Model: wire.Model}
|
||||
var b strings.Builder
|
||||
var toolCalls []ToolCall
|
||||
for _, blk := range wire.Content {
|
||||
switch blk.Type {
|
||||
case "text":
|
||||
b.WriteString(blk.Text)
|
||||
case "tool_use":
|
||||
toolCalls = append(toolCalls, ToolCall{Name: blk.Name, Args: blk.Input})
|
||||
}
|
||||
}
|
||||
out.Content = b.String()
|
||||
out.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: out.Content}}}
|
||||
out.ToolCalls = toolCalls
|
||||
if u := wire.Usage; u != nil {
|
||||
out.PromptTokens = u.InputTokens
|
||||
out.CompletionTokens = u.OutputTokens
|
||||
out.TotalTokens = u.InputTokens + u.OutputTokens
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// AnthropicToolResult builds the content block used to feed a tool result back to
|
||||
// Anthropic (role: user, first content element). toolUseID must equal the id of a
|
||||
// previous tool_use block; isError marks an execution failure.
|
||||
func AnthropicToolResult(toolUseID, content string, isError bool) map[string]any {
|
||||
return map[string]any{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": toolUseID,
|
||||
"content": content,
|
||||
"is_error": isError,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// anthropicStreamEvent is one SSE payload from /v1/messages with stream: true.
|
||||
// Anthropic carries the event kind both in the SSE `event:` field and inside the
|
||||
// JSON `type` field, so the parser keys off the JSON payload type.
|
||||
type anthropicStreamEvent struct {
|
||||
Type string `json:"type"`
|
||||
Message *anthropicStartBlock `json:"message,omitempty"`
|
||||
Index *int `json:"index,omitempty"`
|
||||
Content *anthropicBlock `json:"content_block,omitempty"`
|
||||
Delta *anthropicDelta `json:"delta,omitempty"`
|
||||
Usage *anthropicUsage `json:"usage,omitempty"`
|
||||
Error *anthropicStreamError `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicStartBlock is the payload of a message_start event.
|
||||
type anthropicStartBlock struct {
|
||||
Model string `json:"model"`
|
||||
Usage *anthropicUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicDelta is the incremental payload inside a content_block_delta event.
|
||||
type anthropicDelta struct {
|
||||
Type string `json:"type"` // text_delta | input_json_delta
|
||||
Text string `json:"text,omitempty"`
|
||||
PartialJSON string `json:"partial_json,omitempty"`
|
||||
}
|
||||
|
||||
// anthropicStreamError is the payload of an error event.
|
||||
type anthropicStreamError struct {
|
||||
Type string `json:"type"`
|
||||
Message string `json:"message"`
|
||||
}
|
||||
|
||||
// streamBlock tracks the partial state of one content block during streaming.
|
||||
type streamBlock struct {
|
||||
index int
|
||||
typ string
|
||||
id string
|
||||
name string
|
||||
text strings.Builder
|
||||
input strings.Builder
|
||||
initialArgs map[string]any
|
||||
}
|
||||
|
||||
// scanAnthropicSSE incrementally decodes Anthropic's content-block stream. Text
|
||||
// deltas are forwarded to onChunk; tool_use input_json_delta fragments are
|
||||
// accumulated per block and re-assembled into a structured Args map when the block
|
||||
// stops. The stream must terminate with a message_stop event, otherwise ErrStream
|
||||
// is returned.
|
||||
func scanAnthropicSSE(r io.Reader, onChunk func(string)) (Response, error) {
|
||||
scanner := bufio.NewScanner(r)
|
||||
scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
|
||||
agg := Response{}
|
||||
blocks := map[int]*streamBlock{}
|
||||
var order []int
|
||||
var text strings.Builder
|
||||
sawMessageStop := false
|
||||
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
|
||||
if data == "" {
|
||||
continue
|
||||
}
|
||||
var evt anthropicStreamEvent
|
||||
if err := json.Unmarshal([]byte(data), &evt); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
|
||||
}
|
||||
if err := applyAnthropicEvent(&agg, blocks, &order, &text, evt, onChunk); err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
if evt.Type == "message_stop" {
|
||||
sawMessageStop = true
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
|
||||
}
|
||||
if !sawMessageStop {
|
||||
return Response{}, ErrStream
|
||||
}
|
||||
finalizeAnthropicStream(&agg, blocks, order, text.String())
|
||||
return agg, nil
|
||||
}
|
||||
|
||||
// applyAnthropicEvent folds a single stream event into the running aggregation.
|
||||
func applyAnthropicEvent(
|
||||
agg *Response,
|
||||
blocks map[int]*streamBlock,
|
||||
order *[]int,
|
||||
text *strings.Builder,
|
||||
evt anthropicStreamEvent,
|
||||
onChunk func(string),
|
||||
) error {
|
||||
switch evt.Type {
|
||||
case "message_start":
|
||||
if evt.Message != nil {
|
||||
agg.Model = evt.Message.Model
|
||||
if evt.Message.Usage != nil {
|
||||
agg.PromptTokens = evt.Message.Usage.InputTokens
|
||||
}
|
||||
}
|
||||
case "content_block_start":
|
||||
if evt.Index == nil || evt.Content == nil {
|
||||
return nil
|
||||
}
|
||||
b := &streamBlock{
|
||||
index: *evt.Index,
|
||||
typ: evt.Content.Type,
|
||||
id: evt.Content.ID,
|
||||
name: evt.Content.Name,
|
||||
initialArgs: evt.Content.Input,
|
||||
}
|
||||
if evt.Content.Type == "text" && evt.Content.Text != "" {
|
||||
b.text.WriteString(evt.Content.Text)
|
||||
text.WriteString(evt.Content.Text)
|
||||
if onChunk != nil {
|
||||
onChunk(evt.Content.Text)
|
||||
}
|
||||
}
|
||||
if _, ok := blocks[*evt.Index]; !ok {
|
||||
*order = append(*order, *evt.Index)
|
||||
}
|
||||
blocks[*evt.Index] = b
|
||||
case "content_block_delta":
|
||||
if evt.Index == nil || evt.Delta == nil {
|
||||
return nil
|
||||
}
|
||||
b := blocks[*evt.Index]
|
||||
if b == nil {
|
||||
return nil
|
||||
}
|
||||
switch evt.Delta.Type {
|
||||
case "text_delta":
|
||||
b.text.WriteString(evt.Delta.Text)
|
||||
text.WriteString(evt.Delta.Text)
|
||||
if onChunk != nil {
|
||||
onChunk(evt.Delta.Text)
|
||||
}
|
||||
case "input_json_delta":
|
||||
b.input.WriteString(evt.Delta.PartialJSON)
|
||||
}
|
||||
case "message_delta":
|
||||
if evt.Usage != nil {
|
||||
agg.CompletionTokens = evt.Usage.OutputTokens
|
||||
}
|
||||
case "error":
|
||||
if evt.Error != nil {
|
||||
return fmt.Errorf("model: %w: anthropic stream error: %s: %s", ErrUpstream, evt.Error.Type, evt.Error.Message)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// finalizeAnthropicStream assembles the aggregated Response from the accumulated
|
||||
// blocks: concatenated text plus any tool_use ToolCalls derived per block.
|
||||
func finalizeAnthropicStream(agg *Response, blocks map[int]*streamBlock, order []int, content string) {
|
||||
agg.Content = content
|
||||
agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: content}}}
|
||||
agg.TotalTokens = agg.PromptTokens + agg.CompletionTokens
|
||||
for _, idx := range order {
|
||||
b := blocks[idx]
|
||||
if b == nil || b.typ != "tool_use" {
|
||||
continue
|
||||
}
|
||||
args := b.initialArgs
|
||||
if b.input.Len() > 0 {
|
||||
decoded := map[string]any{}
|
||||
if err := json.Unmarshal([]byte(b.input.String()), &decoded); err == nil {
|
||||
args = decoded
|
||||
} else {
|
||||
args = map[string]any{}
|
||||
}
|
||||
}
|
||||
agg.ToolCalls = append(agg.ToolCalls, ToolCall{Name: b.name, Args: args})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,442 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// anthropicCaptured records what an AnthropicChatAdapter sent to a stub server.
|
||||
type anthropicCaptured struct {
|
||||
method string
|
||||
path string
|
||||
apiKey string
|
||||
anthropicVersion string
|
||||
contentType string
|
||||
accept string
|
||||
body map[string]any
|
||||
}
|
||||
|
||||
func newAnthropicServer(t *testing.T, status int, contentType, body string, ch chan<- anthropicCaptured) *httptest.Server {
|
||||
t.Helper()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
var reqBody map[string]any
|
||||
if ch != nil {
|
||||
_ = json.NewDecoder(r.Body).Decode(&reqBody)
|
||||
}
|
||||
if contentType != "" {
|
||||
w.Header().Set("Content-Type", contentType)
|
||||
}
|
||||
w.WriteHeader(status)
|
||||
if body != "" {
|
||||
_, _ = io.WriteString(w, body)
|
||||
}
|
||||
if ch != nil {
|
||||
ch <- anthropicCaptured{
|
||||
method: r.Method,
|
||||
path: r.URL.Path,
|
||||
apiKey: r.Header.Get("x-api-key"),
|
||||
anthropicVersion: r.Header.Get("anthropic-version"),
|
||||
contentType: r.Header.Get("Content-Type"),
|
||||
accept: r.Header.Get("Accept"),
|
||||
body: reqBody,
|
||||
}
|
||||
}
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
const anthropicTextResponse = `{
|
||||
"id":"msg_1","model":"claude-test","role":"assistant",
|
||||
"content":[{"type":"text","text":"Hello!"}],
|
||||
"usage":{"input_tokens":5,"output_tokens":2}
|
||||
}`
|
||||
|
||||
func TestAnthropicChat_Chat_NonStreamingText(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan anthropicCaptured, 1)
|
||||
srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)
|
||||
|
||||
recs := make([]CallRecord, 0, 1)
|
||||
temp := 1.0
|
||||
ad := NewAnthropicChat(config.LLMConfig{
|
||||
AnthropicBaseURL: srv.URL, AnthropicAPIKey: "ak-test", AnthropicModel: "claude-test",
|
||||
}, WithAnthropicRecorder(func(cr CallRecord) { recs = append(recs, cr) }))
|
||||
|
||||
resp, err := ad.Chat(context.Background(), Request{
|
||||
Messages: []Message{TextMessage(RoleUser, "hi")},
|
||||
Model: "claude-test",
|
||||
Temperature: &temp,
|
||||
MaxTokens: 64,
|
||||
Purpose: "chat",
|
||||
}, Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Chat error: %v", err)
|
||||
}
|
||||
|
||||
got := <-ch
|
||||
if got.method != http.MethodPost {
|
||||
t.Errorf("method = %s, want POST", got.method)
|
||||
}
|
||||
if got.path != "/v1/messages" {
|
||||
t.Errorf("path = %s, want /v1/messages", got.path)
|
||||
}
|
||||
if got.apiKey != "ak-test" {
|
||||
t.Errorf("x-api-key = %q, want ak-test", got.apiKey)
|
||||
}
|
||||
if got.anthropicVersion != "2023-06-01" {
|
||||
t.Errorf("anthropic-version = %q, want 2023-06-01", got.anthropicVersion)
|
||||
}
|
||||
if got.contentType != "application/json" {
|
||||
t.Errorf("content-type = %s, want application/json", got.contentType)
|
||||
}
|
||||
if got.accept != "application/json" {
|
||||
t.Errorf("accept = %s, want application/json", got.accept)
|
||||
}
|
||||
if got.body["model"] != "claude-test" {
|
||||
t.Errorf("body.model = %v, want claude-test", got.body["model"])
|
||||
}
|
||||
if got.body["max_tokens"] != float64(64) {
|
||||
t.Errorf("body.max_tokens = %v, want 64", got.body["max_tokens"])
|
||||
}
|
||||
if got.body["temperature"] != 1.0 {
|
||||
t.Errorf("body.temperature = %v, want 1.0", got.body["temperature"])
|
||||
}
|
||||
if got.body["stream"] != nil && got.body["stream"] != false {
|
||||
t.Errorf("body.stream = %v, want omitted(false)", got.body["stream"])
|
||||
}
|
||||
msgs, _ := got.body["messages"].([]any)
|
||||
if len(msgs) != 1 {
|
||||
t.Fatalf("body.messages = %v, want 1 message", got.body["messages"])
|
||||
}
|
||||
|
||||
if resp.Model != "claude-test" {
|
||||
t.Errorf("resp.Model = %q, want claude-test", resp.Model)
|
||||
}
|
||||
if resp.Content != "Hello!" {
|
||||
t.Errorf("resp.Content = %q, want Hello!", resp.Content)
|
||||
}
|
||||
if resp.TotalTokens != 7 || resp.PromptTokens != 5 || resp.CompletionTokens != 2 {
|
||||
t.Errorf("resp usage = (%d,%d,%d), want (5,2,7)", resp.PromptTokens, resp.CompletionTokens, resp.TotalTokens)
|
||||
}
|
||||
if resp.Message.Role != RoleAssistant || len(resp.Message.Content) != 1 || resp.Message.Content[0].Text != "Hello!" {
|
||||
t.Errorf("resp.Message = %+v", resp.Message)
|
||||
}
|
||||
|
||||
if len(recs) != 1 {
|
||||
t.Fatalf("records = %d, want 1", len(recs))
|
||||
}
|
||||
if recs[0].Model != "claude-test" || recs[0].Purpose != "chat" || recs[0].Tokens != 7 || recs[0].TookMs < 0 {
|
||||
t.Errorf("record = %+v, want model=claude-test purpose=chat tokens=7 took_ms>=0", recs[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_ToolUseInputObject(t *testing.T) {
|
||||
t.Parallel()
|
||||
const resp = `{
|
||||
"id":"msg_2","model":"claude-test",
|
||||
"content":[
|
||||
{"type":"text","text":"Let me check."},
|
||||
{"type":"tool_use","id":"toolu_1","name":"lookup","input":{"location":"Paris","limit":3}}
|
||||
],
|
||||
"usage":{"input_tokens":10,"output_tokens":12}
|
||||
}`
|
||||
ch := make(chan anthropicCaptured, 1)
|
||||
srv := newAnthropicServer(t, http.StatusOK, "application/json", resp, ch)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
|
||||
|
||||
req := Request{
|
||||
Messages: []Message{TextMessage(RoleUser, "hi")},
|
||||
Tools: []ToolSpec{{Name: "lookup", Description: "look up", Parameters: map[string]any{"type": "object"}}},
|
||||
}
|
||||
out, err := ad.Chat(context.Background(), req, Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Chat error: %v", err)
|
||||
}
|
||||
|
||||
got := <-ch
|
||||
tools, ok := got.body["tools"].([]any)
|
||||
if !ok || len(tools) != 1 {
|
||||
t.Fatalf("body.tools = %v, want 1 entry", got.body["tools"])
|
||||
}
|
||||
tool0, _ := tools[0].(map[string]any)
|
||||
if tool0["name"] != "lookup" {
|
||||
t.Errorf("tool name = %v, want lookup", tool0["name"])
|
||||
}
|
||||
if _, ok = tool0["input_schema"]; !ok {
|
||||
t.Errorf("tool input_schema missing from %v", tool0)
|
||||
}
|
||||
|
||||
if out.Content != "Let me check." {
|
||||
t.Errorf("out.Content = %q, want Let me check.", out.Content)
|
||||
}
|
||||
if len(out.ToolCalls) != 1 {
|
||||
t.Fatalf("out.ToolCalls = %+v, want 1", out.ToolCalls)
|
||||
}
|
||||
tc := out.ToolCalls[0]
|
||||
if tc.Name != "lookup" {
|
||||
t.Errorf("tool call name = %q, want lookup", tc.Name)
|
||||
}
|
||||
if tc.Args == nil {
|
||||
t.Fatalf("tool call args = nil, want structured object")
|
||||
}
|
||||
if tc.Args["location"] != "Paris" {
|
||||
t.Errorf("tool call args.location = %v, want Paris (structured, not JSON string)", tc.Args["location"])
|
||||
}
|
||||
if tc.Args["limit"] != float64(3) {
|
||||
t.Errorf("tool call args.limit = %v, want 3", tc.Args["limit"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_SystemAndToolChoice(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan anthropicCaptured, 1)
|
||||
srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
|
||||
|
||||
req := Request{
|
||||
Messages: []Message{
|
||||
TextMessage(RoleSystem, "Be concise."),
|
||||
TextMessage(RoleUser, "hi"),
|
||||
},
|
||||
ToolChoice: map[string]any{"type": "auto"},
|
||||
}
|
||||
if _, err := ad.Chat(context.Background(), req, Options{}); err != nil {
|
||||
t.Fatalf("Chat error: %v", err)
|
||||
}
|
||||
|
||||
got := <-ch
|
||||
if got.body["system"] != "Be concise." {
|
||||
t.Errorf("body.system = %v, want Be concise.", got.body["system"])
|
||||
}
|
||||
msgs, _ := got.body["messages"].([]any)
|
||||
if len(msgs) != 1 {
|
||||
t.Fatalf("body.messages = %v, want 1 message (system extracted)", got.body["messages"])
|
||||
}
|
||||
if _, ok := got.body["tool_choice"].(map[string]any); !ok {
|
||||
t.Errorf("body.tool_choice = %v, want an object", got.body["tool_choice"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_ImageContent(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan anthropicCaptured, 1)
|
||||
srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, ch)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
|
||||
|
||||
req := Request{Messages: []Message{ImageMessage(RoleUser, "what is this", "data:image/png;base64,abc")}}
|
||||
if _, err := ad.Chat(context.Background(), req, Options{}); err != nil {
|
||||
t.Fatalf("Chat error: %v", err)
|
||||
}
|
||||
|
||||
got := <-ch
|
||||
msgs, _ := got.body["messages"].([]any)
|
||||
part0, _ := msgs[0].(map[string]any)
|
||||
content, _ := part0["content"].([]any)
|
||||
if len(content) != 2 {
|
||||
t.Fatalf("content parts = %d, want 2 (text+image)", len(content))
|
||||
}
|
||||
found := false
|
||||
for _, c := range content {
|
||||
cp, _ := c.(map[string]any)
|
||||
if cp["type"] != "image" {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
src, _ := cp["source"].(map[string]any)
|
||||
if src["type"] != "base64" {
|
||||
t.Errorf("image source type = %v, want base64", src["type"])
|
||||
}
|
||||
if src["media_type"] != "image/png" {
|
||||
t.Errorf("image media_type = %v, want image/png", src["media_type"])
|
||||
}
|
||||
if src["data"] != "abc" {
|
||||
t.Errorf("image data = %v, want abc", src["data"])
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("no image block found in %v", content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_UpstreamError(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := newAnthropicServer(t, http.StatusInternalServerError, "application/json", `{"error":"boom"}`, nil)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "m"})
|
||||
_, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{})
|
||||
if err == nil {
|
||||
t.Fatal("want error, got nil")
|
||||
}
|
||||
if !errors.Is(err, ErrUpstream) {
|
||||
t.Fatalf("err = %v, want ErrUpstream", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_NoBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicModel: "m"})
|
||||
if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) {
|
||||
t.Fatalf("err = %v, want ErrUpstream", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_Chat_NoModel(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := newAnthropicServer(t, http.StatusOK, "application/json", anthropicTextResponse, nil)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL})
|
||||
if _, err := ad.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrNoModel) {
|
||||
t.Fatalf("err = %v, want ErrNoModel", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_ChatStream_Text(t *testing.T) {
|
||||
t.Parallel()
|
||||
sse := strings.Join([]string{
|
||||
`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hel"}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"lo"}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_stop","index":0}`,
|
||||
``,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":2}}`,
|
||||
``,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n")
|
||||
ch := make(chan anthropicCaptured, 1)
|
||||
srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, ch)
|
||||
|
||||
recs := make([]CallRecord, 0, 1)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"},
|
||||
WithAnthropicRecorder(func(cr CallRecord) { recs = append(recs, cr) }))
|
||||
|
||||
var deltas []string
|
||||
resp, err := ad.ChatStream(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test", Purpose: "chat"},
|
||||
func(d string) { deltas = append(deltas, d) })
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream error: %v", err)
|
||||
}
|
||||
|
||||
got := <-ch
|
||||
if got.accept != "text/event-stream" {
|
||||
t.Errorf("accept = %s, want text/event-stream", got.accept)
|
||||
}
|
||||
if got.body["stream"] != true {
|
||||
t.Errorf("body.stream = %v, want true", got.body["stream"])
|
||||
}
|
||||
|
||||
if resp.Content != "Hello" {
|
||||
t.Errorf("resp.Content = %q, want Hello", resp.Content)
|
||||
}
|
||||
if resp.Model != "claude-test" {
|
||||
t.Errorf("resp.Model = %q, want claude-test", resp.Model)
|
||||
}
|
||||
if resp.TotalTokens != 7 {
|
||||
t.Errorf("resp.TotalTokens = %d, want 7", resp.TotalTokens)
|
||||
}
|
||||
if !reflect.DeepEqual(deltas, []string{"Hel", "lo"}) {
|
||||
t.Errorf("deltas = %v, want [Hel lo]", deltas)
|
||||
}
|
||||
if len(recs) != 1 || recs[0].Tokens != 7 || recs[0].Purpose != "chat" {
|
||||
t.Errorf("records = %+v, want 1 record tokens=7 purpose=chat", recs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_ChatStream_ToolUseReassembly(t *testing.T) {
|
||||
t.Parallel()
|
||||
sse := strings.Join([]string{
|
||||
`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":10}}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"tool_use","id":"toolu_1","name":"lookup","input":{}}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"{\"location\":"}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"input_json_delta","partial_json":"\"Paris\"}"}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_stop","index":0}`,
|
||||
``,
|
||||
`data: {"type":"message_delta","delta":{"stop_reason":"tool_use"},"usage":{"output_tokens":8}}`,
|
||||
``,
|
||||
`data: {"type":"message_stop"}`,
|
||||
``,
|
||||
}, "\n")
|
||||
srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, nil)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
|
||||
|
||||
var deltas []string
|
||||
resp, err := ad.ChatStream(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test"},
|
||||
func(d string) { deltas = append(deltas, d) })
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream error: %v", err)
|
||||
}
|
||||
|
||||
if len(resp.ToolCalls) != 1 {
|
||||
t.Fatalf("resp.ToolCalls = %+v, want 1", resp.ToolCalls)
|
||||
}
|
||||
tc := resp.ToolCalls[0]
|
||||
if tc.Name != "lookup" {
|
||||
t.Errorf("tool call name = %q, want lookup", tc.Name)
|
||||
}
|
||||
if tc.Args["location"] != "Paris" {
|
||||
t.Errorf("tool call args.location = %v, want Paris (reassembled from partial_json)", tc.Args["location"])
|
||||
}
|
||||
if resp.Content != "" {
|
||||
t.Errorf("resp.Content = %q, want empty (tool-only response)", resp.Content)
|
||||
}
|
||||
if resp.TotalTokens != 18 {
|
||||
t.Errorf("resp.TotalTokens = %d, want 18", resp.TotalTokens)
|
||||
}
|
||||
if len(deltas) != 0 {
|
||||
t.Errorf("deltas = %v, want none (only tool_use)", deltas)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_ChatStream_NoMessageStop(t *testing.T) {
|
||||
t.Parallel()
|
||||
sse := strings.Join([]string{
|
||||
`data: {"type":"message_start","message":{"model":"claude-test","usage":{"input_tokens":5}}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}`,
|
||||
``,
|
||||
`data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hi"}}`,
|
||||
``,
|
||||
}, "\n")
|
||||
srv := newAnthropicServer(t, http.StatusOK, "text/event-stream", sse, nil)
|
||||
ad := NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: srv.URL, AnthropicModel: "claude-test"})
|
||||
if _, err := ad.ChatStream(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "claude-test"}, nil); !errors.Is(err, ErrStream) {
|
||||
t.Fatalf("err = %v, want ErrStream", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicChat_ToolResult(t *testing.T) {
|
||||
t.Parallel()
|
||||
blk := AnthropicToolResult("toolu_1", `{"ok":true}`, true)
|
||||
if blk["type"] != "tool_result" {
|
||||
t.Errorf("type = %v, want tool_result", blk["type"])
|
||||
}
|
||||
if blk["tool_use_id"] != "toolu_1" {
|
||||
t.Errorf("tool_use_id = %v, want toolu_1", blk["tool_use_id"])
|
||||
}
|
||||
if blk["content"] != `{"ok":true}` {
|
||||
t.Errorf("content = %v, want the raw result", blk["content"])
|
||||
}
|
||||
if blk["is_error"] != true {
|
||||
t.Errorf("is_error = %v, want true", blk["is_error"])
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,87 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// Provider names accepted by the provider-aware factory. These are the only two
|
||||
// protocols the factory understands today; adding a vendor means adding a case
|
||||
// here plus a matching adapter, not touching the rest of the package.
|
||||
const (
|
||||
// ProviderOpenAI selects the OpenAI-compatible /chat/completions protocol.
|
||||
ProviderOpenAI = "openai"
|
||||
// ProviderAnthropic selects the Anthropic Messages /v1/messages protocol.
|
||||
ProviderAnthropic = "anthropic"
|
||||
)
|
||||
|
||||
// ErrUnknownProvider is returned by the provider-aware factory when the
|
||||
// configured provider name is not recognized.
|
||||
var ErrUnknownProvider = errors.New("model: unknown provider")
|
||||
|
||||
// NewChatAdapter builds a single provider adapter selected by cfg.Provider. The
|
||||
// empty provider defaults to openai for backward compatibility. An unrecognized
|
||||
// provider returns ErrUnknownProvider.
|
||||
func NewChatAdapter(cfg config.LLMConfig) (Adapter, error) {
|
||||
switch cfg.Provider {
|
||||
case ProviderOpenAI, "":
|
||||
return NewOpenAIChat(cfg), nil
|
||||
case ProviderAnthropic:
|
||||
return NewAnthropicChat(cfg), nil
|
||||
default:
|
||||
return nil, fmt.Errorf("model: %w: provider %q", ErrUnknownProvider, cfg.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
// NewChatChain builds a fallback Chain for the configured provider: the primary
|
||||
// model first, then the provider's fallback model when set and distinct. Each
|
||||
// step is served by an adapter of the selected protocol. The empty provider
|
||||
// defaults to openai. An unrecognized provider returns ErrUnknownProvider.
|
||||
func NewChatChain(cfg config.LLMConfig) (*Chain, error) {
|
||||
switch cfg.Provider {
|
||||
case ProviderOpenAI, "":
|
||||
return newOpenAIChain(cfg)
|
||||
case ProviderAnthropic:
|
||||
return newAnthropicChain(cfg)
|
||||
default:
|
||||
return nil, fmt.Errorf("model: %w: provider %q", ErrUnknownProvider, cfg.Provider)
|
||||
}
|
||||
}
|
||||
|
||||
// newOpenAIChain builds a fallback chain from the OpenAI-compatible config block.
|
||||
// Each step is an OpenAIChatAdapter against the configured endpoint.
|
||||
func newOpenAIChain(cfg config.LLMConfig, opts ...OpenAIOption) (*Chain, error) {
|
||||
if cfg.BaseURL == "" || cfg.Model == "" {
|
||||
return nil, ErrNoModel
|
||||
}
|
||||
steps := []Step{{Model: cfg.Model, Adapter: NewOpenAIChat(cfg, opts...)}}
|
||||
if cfg.Fallback != "" && cfg.Fallback != cfg.Model {
|
||||
steps = append(steps, Step{Model: cfg.Fallback, Adapter: NewOpenAIChat(cfg, opts...)})
|
||||
}
|
||||
return NewChain(steps...), nil
|
||||
}
|
||||
|
||||
// newAnthropicChain builds a fallback chain from the Anthropic config block. Each
|
||||
// step is an AnthropicChatAdapter against the configured /v1/messages endpoint.
|
||||
func newAnthropicChain(cfg config.LLMConfig) (*Chain, error) {
|
||||
if cfg.AnthropicBaseURL == "" || cfg.AnthropicModel == "" {
|
||||
return nil, ErrNoModel
|
||||
}
|
||||
steps := []Step{{Model: cfg.AnthropicModel, Adapter: NewAnthropicChat(cfg)}}
|
||||
if cfg.AnthropicFallback != "" && cfg.AnthropicFallback != cfg.AnthropicModel {
|
||||
steps = append(steps, Step{Model: cfg.AnthropicFallback, Adapter: NewAnthropicChat(cfg)})
|
||||
}
|
||||
return NewChain(steps...), nil
|
||||
}
|
||||
|
||||
// DefaultModelForProvider returns the primary model name for the configured
|
||||
// provider: the flat OpenAI Model for "openai", the Anthropic model for
|
||||
// "anthropic". Used by callers (e.g. cmd/cyrene) to seed a Router's default.
|
||||
func DefaultModelForProvider(cfg config.LLMConfig) string {
|
||||
if cfg.Provider == ProviderAnthropic {
|
||||
return cfg.AnthropicModel
|
||||
}
|
||||
return cfg.Model
|
||||
}
|
||||
@@ -0,0 +1,199 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
func TestNewChatAdapter_Provider(t *testing.T) {
|
||||
t.Parallel()
|
||||
cases := []struct {
|
||||
name string
|
||||
provider string
|
||||
wantType any
|
||||
wantErr bool
|
||||
}{
|
||||
{"openai", ProviderOpenAI, (*OpenAIChatAdapter)(nil), false},
|
||||
{"empty defaults to openai", "", (*OpenAIChatAdapter)(nil), false},
|
||||
{"anthropic", ProviderAnthropic, (*AnthropicChatAdapter)(nil), false},
|
||||
{"unknown", "gemini", nil, true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
c := c
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
ad, err := NewChatAdapter(config.LLMConfig{Provider: c.provider, Model: "m", AnthropicModel: "m"})
|
||||
if c.wantErr {
|
||||
if !errors.Is(err, ErrUnknownProvider) {
|
||||
t.Fatalf("err = %v, want ErrUnknownProvider", err)
|
||||
}
|
||||
if ad != nil {
|
||||
t.Fatalf("adapter = %v, want nil on error", ad)
|
||||
}
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
t.Fatalf("NewChatAdapter error: %v", err)
|
||||
}
|
||||
switch c.wantType.(type) {
|
||||
case *OpenAIChatAdapter:
|
||||
if _, ok := ad.(*OpenAIChatAdapter); !ok {
|
||||
t.Fatalf("adapter = %T, want *OpenAIChatAdapter", ad)
|
||||
}
|
||||
case *AnthropicChatAdapter:
|
||||
if _, ok := ad.(*AnthropicChatAdapter); !ok {
|
||||
t.Fatalf("adapter = %T, want *AnthropicChatAdapter", ad)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewChatChain_Provider(t *testing.T) {
|
||||
t.Parallel()
|
||||
t.Run("openai primary and fallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
chain, err := NewChatChain(config.LLMConfig{Provider: ProviderOpenAI, BaseURL: "http://x", Model: "m", Fallback: "fb"})
|
||||
if err != nil {
|
||||
t.Fatalf("NewChatChain error: %v", err)
|
||||
}
|
||||
if len(chain.steps) != 2 {
|
||||
t.Fatalf("steps = %d, want 2", len(chain.steps))
|
||||
}
|
||||
if chain.steps[0].Model != "m" || chain.steps[1].Model != "fb" {
|
||||
t.Errorf("step models = [%s %s], want [m fb]", chain.steps[0].Model, chain.steps[1].Model)
|
||||
}
|
||||
if _, ok := chain.steps[0].Adapter.(*OpenAIChatAdapter); !ok {
|
||||
t.Errorf("step[0] adapter = %T, want *OpenAIChatAdapter", chain.steps[0].Adapter)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("anthropic primary and fallback", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
chain, err := NewChatChain(config.LLMConfig{
|
||||
Provider: ProviderAnthropic,
|
||||
AnthropicBaseURL: "http://x",
|
||||
AnthropicModel: "claude-main",
|
||||
AnthropicFallback: "claude-fb",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("NewChatChain error: %v", err)
|
||||
}
|
||||
if len(chain.steps) != 2 {
|
||||
t.Fatalf("steps = %d, want 2", len(chain.steps))
|
||||
}
|
||||
if chain.steps[0].Model != "claude-main" || chain.steps[1].Model != "claude-fb" {
|
||||
t.Errorf("step models = [%s %s], want [claude-main claude-fb]", chain.steps[0].Model, chain.steps[1].Model)
|
||||
}
|
||||
if _, ok := chain.steps[0].Adapter.(*AnthropicChatAdapter); !ok {
|
||||
t.Errorf("step[0] adapter = %T, want *AnthropicChatAdapter", chain.steps[0].Adapter)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("missing model returns ErrNoModel", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := NewChatChain(config.LLMConfig{Provider: ProviderOpenAI, BaseURL: "http://x"}); !errors.Is(err, ErrNoModel) {
|
||||
t.Fatalf("openai err = %v, want ErrNoModel", err)
|
||||
}
|
||||
if _, err := NewChatChain(config.LLMConfig{Provider: ProviderAnthropic, AnthropicBaseURL: "http://x"}); !errors.Is(err, ErrNoModel) {
|
||||
t.Fatalf("anthropic err = %v, want ErrNoModel", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unknown provider returns ErrUnknownProvider", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if _, err := NewChatChain(config.LLMConfig{Provider: "gemini"}); !errors.Is(err, ErrUnknownProvider) {
|
||||
t.Fatalf("err = %v, want ErrUnknownProvider", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestNewOpenAIChatChain_StillWorks(t *testing.T) {
|
||||
t.Parallel()
|
||||
chain, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x", Model: "m", Fallback: "fb"})
|
||||
if err != nil {
|
||||
t.Fatalf("NewOpenAIChatChain error: %v", err)
|
||||
}
|
||||
if len(chain.steps) != 2 {
|
||||
t.Fatalf("steps = %d, want 2", len(chain.steps))
|
||||
}
|
||||
if _, ok := chain.steps[0].Adapter.(*OpenAIChatAdapter); !ok {
|
||||
t.Errorf("step[0] adapter = %T, want *OpenAIChatAdapter", chain.steps[0].Adapter)
|
||||
}
|
||||
}
|
||||
|
||||
func TestDefaultModelForProvider(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := DefaultModelForProvider(config.LLMConfig{Provider: ProviderOpenAI, Model: "gpt"}); got != "gpt" {
|
||||
t.Errorf("openai = %q, want gpt", got)
|
||||
}
|
||||
if got := DefaultModelForProvider(config.LLMConfig{Provider: ProviderAnthropic, AnthropicModel: "claude"}); got != "claude" {
|
||||
t.Errorf("anthropic = %q, want claude", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouter_RouteProvider(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default",
|
||||
WithDefaultProvider(ProviderAnthropic),
|
||||
WithRoleProvider(Role("engineer"), ProviderOpenAI),
|
||||
)
|
||||
cases := []struct {
|
||||
role Role
|
||||
want string
|
||||
}{
|
||||
{Role("engineer"), ProviderOpenAI},
|
||||
{Role("user"), ProviderAnthropic},
|
||||
{Role("assistant"), ProviderAnthropic},
|
||||
}
|
||||
for _, c := range cases {
|
||||
c := c
|
||||
t.Run(string(c.role), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := r.RouteProvider(c.role); got != c.want {
|
||||
t.Fatalf("RouteProvider(%q) = %q, want %q", c.role, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestChain_AnthropicFallback exercises the fallback chain against two Anthropic
|
||||
// adapters: the primary step fails and the backup step succeeds, proving the
|
||||
// adapter composes cleanly inside a Chain.
|
||||
func TestChain_AnthropicFallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = io.WriteString(w, `{"error":"down"}`)
|
||||
}))
|
||||
t.Cleanup(primary.Close)
|
||||
backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, `{"id":"msg_b","model":"backup","content":[{"type":"text","text":"Hello!"}],"usage":{"input_tokens":5,"output_tokens":2}}`)
|
||||
}))
|
||||
t.Cleanup(backup.Close)
|
||||
|
||||
chain := NewChain(
|
||||
Step{Model: "main", Adapter: NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: primary.URL, AnthropicModel: "main"})},
|
||||
Step{Model: "backup", Adapter: NewAnthropicChat(config.LLMConfig{AnthropicBaseURL: backup.URL, AnthropicModel: "backup"})},
|
||||
)
|
||||
|
||||
resp, err := chain.Chat(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Purpose: "chat"}, Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Chain.Chat error: %v", err)
|
||||
}
|
||||
if resp.Content != "Hello!" {
|
||||
t.Errorf("resp.Content = %q, want Hello!", resp.Content)
|
||||
}
|
||||
if resp.Model != "backup" {
|
||||
t.Errorf("resp.Model = %q, want backup", resp.Model)
|
||||
}
|
||||
}
|
||||
@@ -29,16 +29,10 @@ func NewChain(steps ...Step) *Chain {
|
||||
// NewOpenAIChatChain builds a fallback chain from a single supplier config: the
|
||||
// primary model first, then cfg.Fallback (when set and distinct), both served
|
||||
// by an OpenAIChatAdapter against the same OpenAI-compatible endpoint. It
|
||||
// returns ErrNoModel when no base URL or primary model is configured.
|
||||
// delegates to the provider-aware builder so the OpenAI path and the factory stay
|
||||
// in sync. It returns ErrNoModel when no base URL or primary model is configured.
|
||||
func NewOpenAIChatChain(cfg config.LLMConfig, opts ...OpenAIOption) (*Chain, error) {
|
||||
if cfg.BaseURL == "" || cfg.Model == "" {
|
||||
return nil, ErrNoModel
|
||||
}
|
||||
steps := []Step{{Model: cfg.Model, Adapter: NewOpenAIChat(cfg, opts...)}}
|
||||
if cfg.Fallback != "" && cfg.Fallback != cfg.Model {
|
||||
steps = append(steps, Step{Model: cfg.Fallback, Adapter: NewOpenAIChat(cfg, opts...)})
|
||||
}
|
||||
return NewChain(steps...), nil
|
||||
return newOpenAIChain(cfg, opts...)
|
||||
}
|
||||
|
||||
// Name returns the chain identifier.
|
||||
|
||||
@@ -87,6 +87,18 @@ type ToolSpec struct {
|
||||
Parameters any
|
||||
}
|
||||
|
||||
// ToolCall is a single assistant-originated tool invocation parsed from a
|
||||
// provider's tool_use content block. Args is the structured argument object
|
||||
// (never a serialized JSON string), so a tool loop can bind it directly without
|
||||
// re-parsing.
|
||||
type ToolCall struct {
|
||||
// Name is the tool being invoked.
|
||||
Name string
|
||||
// Args are the tool arguments, decoded from the provider's structured
|
||||
// tool_use.input object. It is never a JSON-encoded string.
|
||||
Args map[string]any
|
||||
}
|
||||
|
||||
// Request is a completion request, normalized across providers.
|
||||
type Request struct {
|
||||
// Model overrides the adapter's configured default model. Empty means the
|
||||
@@ -101,6 +113,11 @@ type Request struct {
|
||||
Purpose string
|
||||
// Tools is the interface-level placeholder for M2 function calling.
|
||||
Tools []ToolSpec
|
||||
// ToolChoice optionally directs how the provider selects a tool. It is an
|
||||
// opaque provider-agnostic directive (e.g. {"type":"auto"} or {"type":"tool",
|
||||
// "name":"x"}); empty means the provider default. Passed through verbatim to
|
||||
// protocols that support it (Anthropic's tool_choice).
|
||||
ToolChoice any
|
||||
}
|
||||
|
||||
// Response is a normalized completion result.
|
||||
@@ -110,6 +127,9 @@ type Response struct {
|
||||
Content string
|
||||
// Message is the normalized assistant message.
|
||||
Message Message
|
||||
// ToolCalls carries any tool_use blocks the assistant emitted. Empty when the
|
||||
// response is pure text.
|
||||
ToolCalls []ToolCall
|
||||
// PromptTokens, CompletionTokens and TotalTokens report provider usage when
|
||||
// available; zero when the provider does not return it.
|
||||
PromptTokens int
|
||||
|
||||
@@ -2,11 +2,16 @@ package model
|
||||
|
||||
// Router maps roles to model names. An unknown role resolves to the configured
|
||||
// default model. When a request carries image content and a vision model is
|
||||
// configured, the request routes to the vision model.
|
||||
// configured, the request routes to the vision model. A Router also carries a
|
||||
// provider per role (defaulting to the global provider) so a multi-vendor
|
||||
// deployment can pair a role's model with the correct protocol; the provider
|
||||
// selection is surfaced via RouteProvider.
|
||||
type Router struct {
|
||||
defaultModel string
|
||||
byRole map[Role]string
|
||||
visionModel string
|
||||
defaultModel string
|
||||
byRole map[Role]string
|
||||
visionModel string
|
||||
defaultProvider string
|
||||
providerByRole map[Role]string
|
||||
}
|
||||
|
||||
// RouterOption configures a Router.
|
||||
@@ -22,10 +27,28 @@ func WithVisionModel(model string) RouterOption {
|
||||
return func(r *Router) { r.visionModel = model }
|
||||
}
|
||||
|
||||
// WithDefaultProvider sets the provider returned by RouteProvider for roles with
|
||||
// no explicit binding. It defaults to openai.
|
||||
func WithDefaultProvider(provider string) RouterOption {
|
||||
return func(r *Router) { r.defaultProvider = provider }
|
||||
}
|
||||
|
||||
// WithRoleProvider binds a provider to a role, so a multi-vendor deployment can
|
||||
// route a role to a different vendor's protocol (e.g. one role on anthropic, one
|
||||
// on openai). Pair it with WithRoleModel for the role's model name.
|
||||
func WithRoleProvider(role Role, provider string) RouterOption {
|
||||
return func(r *Router) { r.providerByRole[role] = provider }
|
||||
}
|
||||
|
||||
// NewRouter builds a Router with a default model applied to any role that has
|
||||
// no explicit binding.
|
||||
func NewRouter(defaultModel string, opts ...RouterOption) *Router {
|
||||
r := &Router{defaultModel: defaultModel, byRole: map[Role]string{}}
|
||||
r := &Router{
|
||||
defaultModel: defaultModel,
|
||||
byRole: map[Role]string{},
|
||||
defaultProvider: ProviderOpenAI,
|
||||
providerByRole: map[Role]string{},
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(r)
|
||||
}
|
||||
@@ -40,6 +63,18 @@ func (r *Router) Route(role Role) string {
|
||||
return r.defaultModel
|
||||
}
|
||||
|
||||
// RouteProvider returns the provider for a role, falling back to the router's
|
||||
// default provider. It is the routing hook for a multi-vendor deployment: each
|
||||
// role can carry its own provider+model pairing, and roles without a binding fall
|
||||
// back to the global provider. The runtime/loop is expected to build a per-role
|
||||
// adapter or chain from this decision (the provider-aware factory is the builder).
|
||||
func (r *Router) RouteProvider(role Role) string {
|
||||
if p, ok := r.providerByRole[role]; ok && p != "" {
|
||||
return p
|
||||
}
|
||||
return r.defaultProvider
|
||||
}
|
||||
|
||||
// RouteRequest picks the model for a request and reports whether the request
|
||||
// needs a vision model. When the request carries image content and a vision
|
||||
// model is configured, the vision model is returned with needsVision set to
|
||||
|
||||
Reference in new issue
Block a user