feat(model): LLM 多供应商适配(Adapter/OpenAI 非流式+SSE 流式/per-role 路由/主备降级/多模态检测/调用记录)
This commit is contained in:
1 parent
51b2bdfba2
commit
d09ccc7570
9 files changed
+1285
No files matched your search
@@ -0,0 +1,41 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
)
|
||||
|
||||
// Errors returned by the model package. Each is wrapped with %w on the way out
|
||||
// so callers can test for a category while keeping the underlying cause.
|
||||
var (
|
||||
// ErrNoModel is returned when no model name can be resolved.
|
||||
ErrNoModel = errors.New("model: no model configured")
|
||||
// ErrNoSteps is returned when a fallback chain has no steps.
|
||||
ErrNoSteps = errors.New("model: empty fallback chain")
|
||||
// ErrUpstream is returned when the provider request fails (network, status,
|
||||
// decode) or every fallback step fails.
|
||||
ErrUpstream = errors.New("model: upstream request failed")
|
||||
// ErrBadRequest is returned when the request cannot be encoded for the wire.
|
||||
ErrBadRequest = errors.New("model: invalid request")
|
||||
// ErrStream is returned when a stream ends before completion.
|
||||
ErrStream = errors.New("model: stream closed before completion")
|
||||
)
|
||||
|
||||
// Recorder receives a CallRecord after a successfully completed LLM call. It is
|
||||
// invoked by adapters so the runtime (M9 monitor/frontend) can collect usage
|
||||
// without coupling to a storage backend.
|
||||
type Recorder func(CallRecord)
|
||||
|
||||
// Adapter abstracts a single LLM provider. Implementations are vendor-specific
|
||||
// but speak the normalized Request/Response vocabulary; a fallback Chain or a
|
||||
// Router can compose them without knowing which provider is behind each one.
|
||||
type Adapter interface {
|
||||
// Name returns a stable identifier for the provider/adapter.
|
||||
Name() string
|
||||
// Chat performs a non-streaming completion and returns the full response.
|
||||
Chat(ctx context.Context, req Request, opts Options) (Response, error)
|
||||
// ChatStream streams token-level deltas via onChunk and returns the
|
||||
// aggregated response (with usage when the provider reports it). onChunk may
|
||||
// be nil to discard deltas while still aggregating.
|
||||
ChatStream(ctx context.Context, req Request, onChunk func(delta string)) (Response, error)
|
||||
}
|
||||
@@ -0,0 +1,99 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// Step binds a model name to an Adapter for one entry in a fallback chain. An
|
||||
// empty Model keeps the incoming Request.Model.
|
||||
type Step struct {
|
||||
Model string
|
||||
Adapter Adapter
|
||||
}
|
||||
|
||||
// Chain is an Adapter that tries its steps in order and returns the first
|
||||
// successful response. It implements Adapter, so it can be used anywhere a
|
||||
// single provider is expected.
|
||||
type Chain struct {
|
||||
steps []Step
|
||||
}
|
||||
|
||||
// NewChain builds a fallback chain from the given steps.
|
||||
func NewChain(steps ...Step) *Chain {
|
||||
return &Chain{steps: steps}
|
||||
}
|
||||
|
||||
// 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.
|
||||
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
|
||||
}
|
||||
|
||||
// Name returns the chain identifier.
|
||||
func (c *Chain) Name() string { return "fallback" }
|
||||
|
||||
// Chat tries each step in order and returns the first success. If every step
|
||||
// fails, the last error is returned wrapped as ErrUpstream.
|
||||
func (c *Chain) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
|
||||
if len(c.steps) == 0 {
|
||||
return Response{}, ErrNoSteps
|
||||
}
|
||||
var lastErr error
|
||||
for _, s := range c.steps {
|
||||
stepReq := req
|
||||
if s.Model != "" {
|
||||
stepReq.Model = s.Model
|
||||
}
|
||||
resp, err := s.Adapter.Chat(ctx, stepReq, opts)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return Response{}, fmt.Errorf("model: %w: %v", ErrUpstream, lastErr)
|
||||
}
|
||||
|
||||
// ChatStream tries each step in order. If a step fails before delivering any
|
||||
// delta, the chain moves to the next step; once a step has delivered output it
|
||||
// is committed, so a mid-stream failure surfaces as an error rather than mixing
|
||||
// content across providers.
|
||||
func (c *Chain) ChatStream(ctx context.Context, req Request, onChunk func(delta string)) (Response, error) {
|
||||
if len(c.steps) == 0 {
|
||||
return Response{}, ErrNoSteps
|
||||
}
|
||||
var lastErr error
|
||||
for _, s := range c.steps {
|
||||
stepReq := req
|
||||
if s.Model != "" {
|
||||
stepReq.Model = s.Model
|
||||
}
|
||||
delivered := false
|
||||
wrapped := func(delta string) {
|
||||
delivered = true
|
||||
if onChunk != nil {
|
||||
onChunk(delta)
|
||||
}
|
||||
}
|
||||
resp, err := s.Adapter.ChatStream(ctx, stepReq, wrapped)
|
||||
if err == nil {
|
||||
return resp, nil
|
||||
}
|
||||
if delivered {
|
||||
return Response{}, fmt.Errorf("model: %w: stream step %q failed after partial output: %v", ErrUpstream, s.Model, err)
|
||||
}
|
||||
lastErr = err
|
||||
}
|
||||
return Response{}, fmt.Errorf("model: %w: %v", ErrUpstream, lastErr)
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
func TestChain_Chat_Fallback(t *testing.T) {
|
||||
t.Parallel()
|
||||
var primaryCalls, backupCalls atomic.Int64
|
||||
primary := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
primaryCalls.Add(1)
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = io.WriteString(w, `{"error":"down"}`)
|
||||
}))
|
||||
t.Cleanup(primary.Close)
|
||||
backup := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
backupCalls.Add(1)
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_, _ = io.WriteString(w, okCompletion)
|
||||
}))
|
||||
t.Cleanup(backup.Close)
|
||||
|
||||
records := make([]CallRecord, 0, 1)
|
||||
chain := NewChain(
|
||||
Step{Model: "main", Adapter: NewOpenAIChat(
|
||||
config.LLMConfig{BaseURL: primary.URL, Model: "main"},
|
||||
WithRecorder(func(cr CallRecord) { records = append(records, cr) }),
|
||||
)},
|
||||
Step{Model: "backup", Adapter: NewOpenAIChat(
|
||||
config.LLMConfig{BaseURL: backup.URL, Model: "backup"},
|
||||
WithRecorder(func(cr CallRecord) { records = append(records, cr) }),
|
||||
)},
|
||||
)
|
||||
|
||||
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 primaryCalls.Load() != 1 {
|
||||
t.Errorf("primary calls = %d, want 1", primaryCalls.Load())
|
||||
}
|
||||
if backupCalls.Load() != 1 {
|
||||
t.Errorf("backup calls = %d, want 1", backupCalls.Load())
|
||||
}
|
||||
if len(records) != 1 {
|
||||
t.Fatalf("records = %d, want just the successful backup record", len(records))
|
||||
}
|
||||
r := records[0]
|
||||
if r.Model != "backup" || r.Tokens != 7 || r.Purpose != "chat" || r.TookMs < 0 {
|
||||
t.Errorf("record = %+v, want model=backup tokens=7 purpose=chat took_ms>=0", r)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChain_Chat_AllStepsFail(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
w.WriteHeader(http.StatusInternalServerError)
|
||||
_, _ = io.WriteString(w, `{"error":"down"}`)
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
chain := NewChain(Step{Model: "m", Adapter: NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "m"})})
|
||||
|
||||
if _, err := chain.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrUpstream) {
|
||||
t.Fatalf("err = %v, want ErrUpstream", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChain_Chat_EmptyChain(t *testing.T) {
|
||||
t.Parallel()
|
||||
chain := NewChain()
|
||||
if _, err := chain.Chat(context.Background(), Request{Messages: []Message{TextMessage(RoleUser, "hi")}}, Options{}); !errors.Is(err, ErrNoSteps) {
|
||||
t.Fatalf("err = %v, want ErrNoSteps", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewOpenAIChatChain_Steps(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 (primary+fallback)", len(chain.steps))
|
||||
}
|
||||
if chain.steps[0].Model != "m" || chain.steps[1].Model != "fb" {
|
||||
t.Errorf("steps = %+v, want models m then fb", chain.steps)
|
||||
}
|
||||
|
||||
noFB, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x", Model: "m"})
|
||||
if err != nil {
|
||||
t.Fatalf("NewOpenAIChatChain error: %v", err)
|
||||
}
|
||||
if len(noFB.steps) != 1 {
|
||||
t.Fatalf("steps = %d, want 1 (no fallback)", len(noFB.steps))
|
||||
}
|
||||
|
||||
if _, err := NewOpenAIChatChain(config.LLMConfig{BaseURL: "http://x"}); !errors.Is(err, ErrNoModel) {
|
||||
t.Fatalf("err = %v, want ErrNoModel", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,137 @@
|
||||
// Package model provides a vendor-agnostic LLM adapter layer for Cyrene v2.
|
||||
//
|
||||
// It defines a normalized Message/Request/Response vocabulary, a single
|
||||
// Adapter interface implemented by OpenAIChatAdapter (which speaks the
|
||||
// OpenAI-compatible /chat/completions protocol used by OpenAI, DashScope/百炼,
|
||||
// local gateways and proxies), a per-role Router, an automatic fallback Chain,
|
||||
// multimodal (vision) detection, and per-call CallRecord accounting.
|
||||
//
|
||||
// The full tool loop (schema registration, argument binding, result re-feed) is
|
||||
// deferred to M2; this package only carries interface-level placeholders
|
||||
// (ToolSpec) so the wire contract is ready without executing any loop.
|
||||
package model
|
||||
|
||||
import "time"
|
||||
|
||||
// Role identifies who authored a message. The values are vendor-agnostic and
|
||||
// translated to each provider's own role vocabulary by the adapter.
|
||||
type Role string
|
||||
|
||||
const (
|
||||
// RoleSystem is a system/instruction message.
|
||||
RoleSystem Role = "system"
|
||||
// RoleUser is an end-user message.
|
||||
RoleUser Role = "user"
|
||||
// RoleAssistant is an assistant (model) message.
|
||||
RoleAssistant Role = "assistant"
|
||||
// RoleTool is a tool-result message (reserved for the M2 tool loop).
|
||||
RoleTool Role = "tool"
|
||||
)
|
||||
|
||||
// ContentType describes the kind of payload inside a Content part.
|
||||
type ContentType string
|
||||
|
||||
const (
|
||||
// ContentTypeText is a plain-text content part.
|
||||
ContentTypeText ContentType = "text"
|
||||
// ContentTypeImage is an image content part.
|
||||
ContentTypeImage ContentType = "image_url"
|
||||
)
|
||||
|
||||
// Content is a single, possibly multimodal, part of a message.
|
||||
type Content struct {
|
||||
// Type distinguishes a text part from an image part.
|
||||
Type ContentType
|
||||
// Text is the text payload when Type == ContentTypeText.
|
||||
Text string
|
||||
// ImageURL is the image source when Type == ContentTypeImage. It may be a
|
||||
// data: URI or an http(s):// URL.
|
||||
ImageURL string
|
||||
}
|
||||
|
||||
// Message is a single conversation turn, normalized across providers.
|
||||
type Message struct {
|
||||
Role Role
|
||||
// Content carries one or more parts. A text-only message uses a single
|
||||
// ContentTypeText part.
|
||||
Content []Content
|
||||
// ToolCallID and ToolCallName are reserved as interface-level placeholders
|
||||
// for the M2 tool loop; they are ignored by current adapters.
|
||||
ToolCallID string
|
||||
ToolCallName string
|
||||
}
|
||||
|
||||
// TextMessage builds a text-only message.
|
||||
func TextMessage(role Role, text string) Message {
|
||||
return Message{Role: role, Content: []Content{{Type: ContentTypeText, Text: text}}}
|
||||
}
|
||||
|
||||
// ImageMessage builds a message carrying optional text plus an image. imageURL
|
||||
// may be a data: URI or an http(s):// URL. An empty text omits the text part.
|
||||
func ImageMessage(role Role, text, imageURL string) Message {
|
||||
parts := make([]Content, 0, 2)
|
||||
if text != "" {
|
||||
parts = append(parts, Content{Type: ContentTypeText, Text: text})
|
||||
}
|
||||
parts = append(parts, Content{Type: ContentTypeImage, ImageURL: imageURL})
|
||||
return Message{Role: role, Content: parts}
|
||||
}
|
||||
|
||||
// ToolSpec is an interface-level placeholder for a tool/function schema. The
|
||||
// full tool loop (registration, argument binding, result re-feed) is M2; here
|
||||
// it only lets a request advertise tools to the provider in the wire contract.
|
||||
type ToolSpec struct {
|
||||
Name string
|
||||
Description string
|
||||
// Parameters is the JSON Schema for the tool arguments. Populated from M2.
|
||||
Parameters any
|
||||
}
|
||||
|
||||
// Request is a completion request, normalized across providers.
|
||||
type Request struct {
|
||||
// Model overrides the adapter's configured default model. Empty means the
|
||||
// adapter default (or the chain step model) applies.
|
||||
Model string
|
||||
Messages []Message
|
||||
// Temperature is the sampling temperature. nil means the provider default.
|
||||
Temperature *float64
|
||||
// MaxTokens caps the generated tokens. Zero means the provider default.
|
||||
MaxTokens int
|
||||
// Purpose labels the call for CallRecord accounting (e.g. "chat").
|
||||
Purpose string
|
||||
// Tools is the interface-level placeholder for M2 function calling.
|
||||
Tools []ToolSpec
|
||||
}
|
||||
|
||||
// Response is a normalized completion result.
|
||||
type Response struct {
|
||||
// Model is the model that produced the response.
|
||||
Model string
|
||||
Content string
|
||||
// Message is the normalized assistant message.
|
||||
Message Message
|
||||
// PromptTokens, CompletionTokens and TotalTokens report provider usage when
|
||||
// available; zero when the provider does not return it.
|
||||
PromptTokens int
|
||||
CompletionTokens int
|
||||
TotalTokens int
|
||||
}
|
||||
|
||||
// CallRecord describes a single successfully completed LLM call, for the M9
|
||||
// monitor/frontend. A failed attempt emits no record.
|
||||
type CallRecord struct {
|
||||
// Model is the model that served the request.
|
||||
Model string
|
||||
// Purpose labels what the call was for.
|
||||
Purpose string
|
||||
// Tokens is the total token consumption reported by the provider.
|
||||
Tokens int
|
||||
// TookMs is the end-to-end call latency in milliseconds.
|
||||
TookMs int64
|
||||
}
|
||||
|
||||
// Options carries per-call knobs layered on top of the Request.
|
||||
type Options struct {
|
||||
// Timeout bounds a single HTTP call. Zero uses the adapter client timeout.
|
||||
Timeout time.Duration
|
||||
}
|
||||
@@ -0,0 +1,56 @@
|
||||
package model
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestTextMessage(t *testing.T) {
|
||||
t.Parallel()
|
||||
msg := TextMessage(RoleUser, "hi")
|
||||
if msg.Role != RoleUser {
|
||||
t.Fatalf("role = %q, want %q", msg.Role, RoleUser)
|
||||
}
|
||||
if len(msg.Content) != 1 || msg.Content[0].Type != ContentTypeText || msg.Content[0].Text != "hi" {
|
||||
t.Fatalf("content = %+v, want a single text part", msg.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageMessage(t *testing.T) {
|
||||
t.Parallel()
|
||||
msg := ImageMessage(RoleUser, "look", "data:image/png;base64,abc")
|
||||
if len(msg.Content) != 2 {
|
||||
t.Fatalf("content len = %d, want 2 (text + image)", len(msg.Content))
|
||||
}
|
||||
if msg.Content[1].Type != ContentTypeImage || msg.Content[1].ImageURL != "data:image/png;base64,abc" {
|
||||
t.Fatalf("image part = %+v", msg.Content[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestImageMessage_NoText(t *testing.T) {
|
||||
t.Parallel()
|
||||
msg := ImageMessage(RoleUser, "", "data:image/png;base64,abc")
|
||||
if len(msg.Content) != 1 || msg.Content[0].Type != ContentTypeImage {
|
||||
t.Fatalf("content = %+v, want a single image part", msg.Content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestHasImage(t *testing.T) {
|
||||
t.Parallel()
|
||||
cases := []struct {
|
||||
name string
|
||||
msgs []Message
|
||||
want bool
|
||||
}{
|
||||
{"empty", nil, false},
|
||||
{"text only", []Message{TextMessage(RoleUser, "hi")}, false},
|
||||
{"image", []Message{ImageMessage(RoleUser, "", "data:image/png;base64,x")}, true},
|
||||
{"text and image", []Message{ImageMessage(RoleUser, "desc", "data:image/png;base64,x")}, true},
|
||||
}
|
||||
for _, c := range cases {
|
||||
c := c
|
||||
t.Run(c.name, func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := HasImage(c.msgs); got != c.want {
|
||||
t.Fatalf("HasImage = %v, want %v", got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,425 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
// Wire constraints for the OpenAI-compatible /chat/completions protocol.
|
||||
const (
|
||||
// sseDataPrefix marks an SSE event line carrying a JSON payload.
|
||||
sseDataPrefix = "data: "
|
||||
// sseDone is the terminal SSE sentinel.
|
||||
sseDone = "[DONE]"
|
||||
// sseBufInit is the initial bufio.Scanner read buffer size.
|
||||
sseBufInit = 64 * 1024
|
||||
// maxSSELineBytes caps a single SSE event line.
|
||||
maxSSELineBytes = 1024 * 1024
|
||||
// errBodyCap bounds how much of an error body is read for the message.
|
||||
errBodyCap = 512
|
||||
// defaultHTTPTimeout bounds a single provider call when Options.Timeout is 0.
|
||||
defaultHTTPTimeout = 180 * time.Second
|
||||
)
|
||||
|
||||
// openAIChatRequest is the JSON body sent to /chat/completions. Field names
|
||||
// match the provider protocol so the struct marshals directly.
|
||||
type openAIChatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []openAIMessage `json:"messages"`
|
||||
Temperature *float64 `json:"temperature,omitempty"`
|
||||
MaxTokens int `json:"max_tokens,omitempty"`
|
||||
Stream bool `json:"stream,omitempty"`
|
||||
StreamOptions *openAIStreamOpts `json:"stream_options,omitempty"`
|
||||
Tools []openAITool `json:"tools,omitempty"`
|
||||
}
|
||||
|
||||
// openAIStreamOpts asks the provider to include usage in the final stream chunk.
|
||||
type openAIStreamOpts struct {
|
||||
IncludeUsage bool `json:"include_usage"`
|
||||
}
|
||||
|
||||
// openAIMessage is one request message in the wire body. Content is either a
|
||||
// plain string (text-only) or an array of {type,text|image_url} parts.
|
||||
type openAIMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content any `json:"content"`
|
||||
}
|
||||
|
||||
// openAITool is the wire form of a function tool schema (M2 placeholder).
|
||||
type openAITool struct {
|
||||
Type string `json:"type"`
|
||||
Function openAIToolFn `json:"function"`
|
||||
}
|
||||
|
||||
// openAIToolFn carries the function definition.
|
||||
type openAIToolFn struct {
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description,omitempty"`
|
||||
Parameters any `json:"parameters,omitempty"`
|
||||
}
|
||||
|
||||
// openAIChatResponse is the JSON body returned by /chat/completions, for both
|
||||
// the non-streaming payload and each SSE stream chunk.
|
||||
type openAIChatResponse struct {
|
||||
ID string `json:"id"`
|
||||
Model string `json:"model"`
|
||||
Choices []openAIChoice `json:"choices"`
|
||||
Usage *openAIUsage `json:"usage,omitempty"`
|
||||
}
|
||||
|
||||
// openAIChoice is one candidate completion.
|
||||
type openAIChoice struct {
|
||||
Index int `json:"index"`
|
||||
Message openAIResponseMessage `json:"message"`
|
||||
Delta *openAIDelta `json:"delta,omitempty"`
|
||||
FinishReason string `json:"finish_reason,omitempty"`
|
||||
}
|
||||
|
||||
// openAIResponseMessage is the assistant message inside a non-streaming choice.
|
||||
type openAIResponseMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// openAIDelta is the incremental payload inside a stream chunk.
|
||||
type openAIDelta struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
}
|
||||
|
||||
// openAIUsage reports token consumption.
|
||||
type openAIUsage struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
}
|
||||
|
||||
// OpenAIChatAdapter calls an OpenAI-compatible /chat/completions endpoint. It is
|
||||
// constructed from a config.LLMConfig and implements Adapter for both
|
||||
// non-streaming and token-level SSE streaming.
|
||||
type OpenAIChatAdapter struct {
|
||||
name string
|
||||
baseURL string
|
||||
apiKey string
|
||||
model string
|
||||
client *http.Client
|
||||
recorder Recorder
|
||||
}
|
||||
|
||||
// OpenAIOption configures an OpenAIChatAdapter.
|
||||
type OpenAIOption func(*OpenAIChatAdapter)
|
||||
|
||||
// WithHTTPClient overrides the default HTTP client (used to inject mocks or a
|
||||
// tuned transport in tests and production).
|
||||
func WithHTTPClient(c *http.Client) OpenAIOption {
|
||||
return func(a *OpenAIChatAdapter) { a.client = c }
|
||||
}
|
||||
|
||||
// WithRecorder wires a per-call Recorder for CallRecord accounting.
|
||||
func WithRecorder(r Recorder) OpenAIOption {
|
||||
return func(a *OpenAIChatAdapter) { a.recorder = r }
|
||||
}
|
||||
|
||||
// WithModel overrides the adapter's default model name.
|
||||
func WithModel(m string) OpenAIOption {
|
||||
return func(a *OpenAIChatAdapter) { a.model = m }
|
||||
}
|
||||
|
||||
// NewOpenAIChat builds an adapter from a supplier config. baseURL is trimmed of
|
||||
// any trailing slash before the /chat/completions path is appended.
|
||||
func NewOpenAIChat(cfg config.LLMConfig, opts ...OpenAIOption) *OpenAIChatAdapter {
|
||||
a := &OpenAIChatAdapter{
|
||||
name: "openai",
|
||||
baseURL: strings.TrimRight(cfg.BaseURL, "/"),
|
||||
apiKey: cfg.APIKey,
|
||||
model: cfg.Model,
|
||||
client: &http.Client{
|
||||
Transport: http.DefaultTransport,
|
||||
Timeout: defaultHTTPTimeout,
|
||||
},
|
||||
}
|
||||
for _, o := range opts {
|
||||
o(a)
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
// Name returns the adapter identifier.
|
||||
func (a *OpenAIChatAdapter) Name() string { return a.name }
|
||||
|
||||
// Chat performs a non-streaming completion.
|
||||
func (a *OpenAIChatAdapter) 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 openAIChatResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&wire); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: decode response: %v", ErrUpstream, err)
|
||||
}
|
||||
out := responseFromWire(&wire)
|
||||
a.record(req, start, out.TotalTokens)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// ChatStream streams token-level deltas via onChunk and returns the aggregated
|
||||
// response with usage when the provider reports it in the final chunk.
|
||||
func (a *OpenAIChatAdapter) 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 := scanSSE(resp.Body, onChunk)
|
||||
if err != nil {
|
||||
return Response{}, err
|
||||
}
|
||||
out.Model = model
|
||||
a.record(req, start, out.TotalTokens)
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// withCallTimeout wraps ctx with opts.Timeout when it is positive.
|
||||
func withCallTimeout(parent context.Context, opts Options) (context.Context, context.CancelFunc) {
|
||||
if opts.Timeout > 0 {
|
||||
return context.WithTimeout(parent, opts.Timeout)
|
||||
}
|
||||
return parent, func() {}
|
||||
}
|
||||
|
||||
// resolveModel returns the effective model name.
|
||||
func (a *OpenAIChatAdapter) 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 /chat/completions URL.
|
||||
func (a *OpenAIChatAdapter) endpoint() string {
|
||||
return a.baseURL + "/chat/completions"
|
||||
}
|
||||
|
||||
func (a *OpenAIChatAdapter) buildRequest(req Request, stream bool) ([]byte, error) {
|
||||
msgs, err := buildMessages(req.Messages)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
wire := openAIChatRequest{
|
||||
Model: req.Model,
|
||||
Messages: msgs,
|
||||
Temperature: req.Temperature,
|
||||
MaxTokens: req.MaxTokens,
|
||||
Stream: stream,
|
||||
Tools: buildTools(req.Tools),
|
||||
}
|
||||
if stream {
|
||||
wire.StreamOptions = &openAIStreamOpts{IncludeUsage: true}
|
||||
}
|
||||
buf, err := json.Marshal(wire)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("model: %w: marshal request: %v", ErrBadRequest, err)
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
func buildMessages(msgs []Message) ([]openAIMessage, error) {
|
||||
out := make([]openAIMessage, 0, len(msgs))
|
||||
for _, m := range msgs {
|
||||
content, err := encodeContent(m.Content)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out = append(out, openAIMessage{Role: string(m.Role), Content: content})
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func encodeContent(parts []Content) (any, error) {
|
||||
if len(parts) == 0 {
|
||||
return "", nil
|
||||
}
|
||||
if len(parts) == 1 && parts[0].Type == ContentTypeText {
|
||||
return parts[0].Text, nil
|
||||
}
|
||||
arr := make([]any, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case ContentTypeText:
|
||||
arr = append(arr, map[string]any{"type": "text", "text": p.Text})
|
||||
case ContentTypeImage:
|
||||
arr = append(arr, map[string]any{
|
||||
"type": "image_url",
|
||||
"image_url": map[string]any{"url": p.ImageURL},
|
||||
})
|
||||
default:
|
||||
return nil, fmt.Errorf("model: %w: unsupported content type %q", ErrBadRequest, p.Type)
|
||||
}
|
||||
}
|
||||
return arr, nil
|
||||
}
|
||||
|
||||
func buildTools(tools []ToolSpec) []openAITool {
|
||||
if len(tools) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]openAITool, 0, len(tools))
|
||||
for _, t := range tools {
|
||||
out = append(out, openAITool{Type: "function", Function: openAIToolFn{
|
||||
Name: t.Name,
|
||||
Description: t.Description,
|
||||
Parameters: t.Parameters,
|
||||
}})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func (a *OpenAIChatAdapter) doRequest(httpReq *http.Request, stream bool) (*http.Response, error) {
|
||||
if a.apiKey != "" {
|
||||
httpReq.Header.Set("Authorization", "Bearer "+a.apiKey)
|
||||
}
|
||||
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 *OpenAIChatAdapter) 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 *OpenAIChatAdapter) 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(),
|
||||
})
|
||||
}
|
||||
|
||||
func responseFromWire(wire *openAIChatResponse) Response {
|
||||
out := Response{Model: wire.Model}
|
||||
if len(wire.Choices) > 0 {
|
||||
text := wire.Choices[0].Message.Content
|
||||
out.Content = text
|
||||
out.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: text}}}
|
||||
}
|
||||
if u := wire.Usage; u != nil {
|
||||
out.PromptTokens = u.PromptTokens
|
||||
out.CompletionTokens = u.CompletionTokens
|
||||
out.TotalTokens = u.TotalTokens
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// scanSSE incrementally decodes stream chunks, accumulates the full content
|
||||
// with a strings.Builder (never reading the whole body at once) and captures
|
||||
// usage from the final chunk.
|
||||
func scanSSE(r io.Reader, onChunk func(string)) (Response, error) {
|
||||
scanner := bufio.NewScanner(r)
|
||||
scanner.Buffer(make([]byte, 0, sseBufInit), maxSSELineBytes)
|
||||
agg := Response{}
|
||||
var b strings.Builder
|
||||
for scanner.Scan() {
|
||||
line := strings.TrimSpace(scanner.Text())
|
||||
if line == "" || !strings.HasPrefix(line, sseDataPrefix) {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimSpace(strings.TrimPrefix(line, sseDataPrefix))
|
||||
if data == sseDone {
|
||||
break
|
||||
}
|
||||
var chunk openAIChatResponse
|
||||
if err := json.Unmarshal([]byte(data), &chunk); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: decode stream chunk: %v", ErrUpstream, err)
|
||||
}
|
||||
if u := chunk.Usage; u != nil {
|
||||
agg.PromptTokens = u.PromptTokens
|
||||
agg.CompletionTokens = u.CompletionTokens
|
||||
agg.TotalTokens = u.TotalTokens
|
||||
}
|
||||
if len(chunk.Choices) == 0 || chunk.Choices[0].Delta == nil {
|
||||
continue
|
||||
}
|
||||
delta := chunk.Choices[0].Delta.Content
|
||||
if delta == "" {
|
||||
continue
|
||||
}
|
||||
b.WriteString(delta)
|
||||
if onChunk != nil {
|
||||
onChunk(delta)
|
||||
}
|
||||
}
|
||||
if err := scanner.Err(); err != nil {
|
||||
return Response{}, fmt.Errorf("model: %w: read stream: %v", ErrUpstream, err)
|
||||
}
|
||||
agg.Content = b.String()
|
||||
agg.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: agg.Content}}}
|
||||
return agg, nil
|
||||
}
|
||||
@@ -0,0 +1,256 @@
|
||||
package model
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
|
||||
)
|
||||
|
||||
const okCompletion = `{
|
||||
"id": "cmpl-1",
|
||||
"model": "gpt-test",
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "Hello!"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 2, "total_tokens": 7}
|
||||
}`
|
||||
|
||||
// capturedRequest records what an OpenAIChatAdapter sent to a stub server.
|
||||
type capturedRequest struct {
|
||||
method string
|
||||
path string
|
||||
contentType string
|
||||
body map[string]any
|
||||
}
|
||||
|
||||
func newCompletionServer(t *testing.T, status int, contentType, body string, ch chan<- capturedRequest) *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 <- capturedRequest{
|
||||
method: r.Method,
|
||||
path: r.URL.Path,
|
||||
contentType: r.Header.Get("Content-Type"),
|
||||
body: reqBody,
|
||||
}
|
||||
}
|
||||
}))
|
||||
t.Cleanup(srv.Close)
|
||||
return srv
|
||||
}
|
||||
|
||||
func TestOpenAIChat_Chat_NonStreaming(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
|
||||
recs := make([]CallRecord, 0, 1)
|
||||
temp := 1.0
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, APIKey: "sk-test", Model: "gpt-test"},
|
||||
WithRecorder(func(cr CallRecord) { recs = append(recs, cr) }))
|
||||
|
||||
ctx := context.Background()
|
||||
resp, err := ad.Chat(ctx, Request{
|
||||
Messages: []Message{TextMessage(RoleUser, "hi")},
|
||||
Model: "gpt-test",
|
||||
Temperature: &temp,
|
||||
MaxTokens: 64,
|
||||
Purpose: "chat",
|
||||
}, Options{})
|
||||
if err != nil {
|
||||
t.Fatalf("Chat error: %v", err)
|
||||
}
|
||||
|
||||
req := <-ch
|
||||
if req.method != http.MethodPost {
|
||||
t.Errorf("method = %s, want POST", req.method)
|
||||
}
|
||||
if req.path != "/chat/completions" {
|
||||
t.Errorf("path = %s, want /chat/completions", req.path)
|
||||
}
|
||||
if req.contentType != "application/json" {
|
||||
t.Errorf("content-type = %s, want application/json", req.contentType)
|
||||
}
|
||||
if req.body["model"] != "gpt-test" {
|
||||
t.Errorf("body.model = %v, want gpt-test", req.body["model"])
|
||||
}
|
||||
if req.body["stream"] != nil && req.body["stream"] != false {
|
||||
t.Errorf("body.stream = %v, want omitted(false)", req.body["stream"])
|
||||
}
|
||||
if req.body["temperature"] != 1.0 {
|
||||
t.Errorf("body.temperature = %v, want 1.0", req.body["temperature"])
|
||||
}
|
||||
if req.body["max_tokens"] != float64(64) {
|
||||
t.Errorf("body.max_tokens = %v, want 64", req.body["max_tokens"])
|
||||
}
|
||||
msgs, _ := req.body["messages"].([]any)
|
||||
if len(msgs) != 1 {
|
||||
t.Fatalf("body.messages = %v, want 1 message", req.body["messages"])
|
||||
}
|
||||
|
||||
if resp.Model != "gpt-test" {
|
||||
t.Errorf("resp.Model = %q, want gpt-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 != "gpt-test" || recs[0].Purpose != "chat" || recs[0].Tokens != 7 || recs[0].TookMs < 0 {
|
||||
t.Errorf("record = %+v, want model=gpt-test purpose=chat tokens=7 took_ms>=0", recs[0])
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIChat_Chat_ImageContent(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-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_url" {
|
||||
continue
|
||||
}
|
||||
found = true
|
||||
iu, _ := cp["image_url"].(map[string]any)
|
||||
if iu["url"] != "data:image/png;base64,abc" {
|
||||
t.Fatalf("image url = %v, want data:image/png;base64,abc", iu["url"])
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
t.Fatalf("no image_url part found in %v", content)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIChat_Chat_ToolsPlaceholder(t *testing.T) {
|
||||
t.Parallel()
|
||||
ch := make(chan capturedRequest, 1)
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, ch)
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"})
|
||||
|
||||
req := Request{Messages: []Message{TextMessage(RoleUser, "hi")},
|
||||
Tools: []ToolSpec{{Name: "lookup", Description: "look it up", Parameters: map[string]any{"type": "object"}}}}
|
||||
if _, err := ad.Chat(context.Background(), req, Options{}); 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"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIChat_ChatStream_SSE(t *testing.T) {
|
||||
t.Parallel()
|
||||
sse := strings.Join([]string{
|
||||
`data: {"model":"gpt-test","choices":[{"index":0,"delta":{"role":"assistant","content":"Hel"}}]}`,
|
||||
``,
|
||||
`data: {"choices":[{"delta":{"content":"lo"}}]}`,
|
||||
``,
|
||||
`data: {"model":"gpt-test","choices":[],"usage":{"prompt_tokens":5,"completion_tokens":2,"total_tokens":7}}`,
|
||||
``,
|
||||
`data: [DONE]`,
|
||||
``,
|
||||
}, "\n")
|
||||
srv := newCompletionServer(t, http.StatusOK, "text/event-stream", sse, nil)
|
||||
|
||||
recs := make([]CallRecord, 0, 1)
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "gpt-test"},
|
||||
WithRecorder(func(cr CallRecord) { recs = append(recs, cr) }))
|
||||
|
||||
var deltas []string
|
||||
resp, err := ad.ChatStream(context.Background(),
|
||||
Request{Messages: []Message{TextMessage(RoleUser, "hi")}, Model: "gpt-test", Purpose: "chat"},
|
||||
func(d string) { deltas = append(deltas, d) })
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream error: %v", err)
|
||||
}
|
||||
if resp.Content != "Hello" {
|
||||
t.Errorf("resp.Content = %q, want Hello", resp.Content)
|
||||
}
|
||||
if resp.Model != "gpt-test" {
|
||||
t.Errorf("resp.Model = %q, want gpt-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 TestOpenAIChat_Chat_UpstreamError(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := newCompletionServer(t, http.StatusInternalServerError, "application/json", `{"error":"boom"}`, nil)
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: srv.URL, Model: "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 TestOpenAIChat_Chat_NoBaseURL(t *testing.T) {
|
||||
t.Parallel()
|
||||
ad := NewOpenAIChat(config.LLMConfig{})
|
||||
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 TestOpenAIChat_Chat_NoModel(t *testing.T) {
|
||||
t.Parallel()
|
||||
srv := newCompletionServer(t, http.StatusOK, "application/json", okCompletion, nil)
|
||||
ad := NewOpenAIChat(config.LLMConfig{BaseURL: 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)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
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.
|
||||
type Router struct {
|
||||
defaultModel string
|
||||
byRole map[Role]string
|
||||
visionModel string
|
||||
}
|
||||
|
||||
// RouterOption configures a Router.
|
||||
type RouterOption func(*Router)
|
||||
|
||||
// WithRoleModel binds a model to a role.
|
||||
func WithRoleModel(role Role, model string) RouterOption {
|
||||
return func(r *Router) { r.byRole[role] = model }
|
||||
}
|
||||
|
||||
// WithVisionModel configures the dedicated vision model.
|
||||
func WithVisionModel(model string) RouterOption {
|
||||
return func(r *Router) { r.visionModel = model }
|
||||
}
|
||||
|
||||
// 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{}}
|
||||
for _, o := range opts {
|
||||
o(r)
|
||||
}
|
||||
return r
|
||||
}
|
||||
|
||||
// Route returns the model for a role, falling back to the default model.
|
||||
func (r *Router) Route(role Role) string {
|
||||
if m, ok := r.byRole[role]; ok && m != "" {
|
||||
return m
|
||||
}
|
||||
return r.defaultModel
|
||||
}
|
||||
|
||||
// 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
|
||||
// true. Otherwise the role of the last message drives routing (user by
|
||||
// default).
|
||||
func (r *Router) RouteRequest(req Request) (model string, needsVision bool) {
|
||||
role := RoleUser
|
||||
if n := len(req.Messages); n > 0 {
|
||||
role = req.Messages[n-1].Role
|
||||
}
|
||||
needsVision = HasImage(req.Messages)
|
||||
if needsVision && r.visionModel != "" {
|
||||
return r.visionModel, true
|
||||
}
|
||||
return r.Route(role), needsVision
|
||||
}
|
||||
|
||||
// HasImage reports whether any message carries an image content part.
|
||||
func HasImage(msgs []Message) bool {
|
||||
for _, m := range msgs {
|
||||
for _, part := range m.Content {
|
||||
if part.Type == ContentTypeImage {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
@@ -0,0 +1,88 @@
|
||||
package model
|
||||
|
||||
import "testing"
|
||||
|
||||
func TestRouter_Route_PerRole(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default",
|
||||
WithRoleModel(RoleUser, "user-model"),
|
||||
WithRoleModel(RoleAssistant, "assistant-model"),
|
||||
WithVisionModel("vision-model"),
|
||||
)
|
||||
cases := []struct {
|
||||
role Role
|
||||
want string
|
||||
}{
|
||||
{RoleUser, "user-model"},
|
||||
{RoleAssistant, "assistant-model"},
|
||||
{RoleSystem, "default"},
|
||||
{Role("unknown"), "default"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
c := c
|
||||
t.Run(string(c.role), func(t *testing.T) {
|
||||
t.Parallel()
|
||||
if got := r.Route(c.role); got != c.want {
|
||||
t.Fatalf("Route(%q) = %q, want %q", c.role, got, c.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouter_RouteRequest_Text(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default",
|
||||
WithRoleModel(RoleUser, "user-model"),
|
||||
WithRoleModel(RoleAssistant, "assistant-model"),
|
||||
WithVisionModel("vision-model"),
|
||||
)
|
||||
req := Request{Messages: []Message{TextMessage(RoleAssistant, "hi")}}
|
||||
model, needsVision := r.RouteRequest(req)
|
||||
if needsVision {
|
||||
t.Errorf("needsVision = true, want false for text-only")
|
||||
}
|
||||
if model != "assistant-model" {
|
||||
t.Errorf("model = %q, want assistant-model (last-message role)", model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouter_RouteRequest_Vision(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default",
|
||||
WithRoleModel(RoleUser, "user-model"),
|
||||
WithVisionModel("vision-model"),
|
||||
)
|
||||
req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
|
||||
model, needsVision := r.RouteRequest(req)
|
||||
if !needsVision {
|
||||
t.Errorf("needsVision = false, want true for image content")
|
||||
}
|
||||
if model != "vision-model" {
|
||||
t.Errorf("model = %q, want vision-model", model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouter_RouteRequest_VisionUnconfigured(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
|
||||
req := Request{Messages: []Message{ImageMessage(RoleUser, "", "data:image/png;base64,abc")}}
|
||||
model, needsVision := r.RouteRequest(req)
|
||||
if !needsVision {
|
||||
t.Errorf("needsVision = false, want true when image present even without a vision model")
|
||||
}
|
||||
if model != "user-model" {
|
||||
t.Errorf("model = %q, want user-model (role fallback)", model)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRouter_RouteRequest_EmptyMessages(t *testing.T) {
|
||||
t.Parallel()
|
||||
r := NewRouter("default", WithRoleModel(RoleUser, "user-model"))
|
||||
model, needsVision := r.RouteRequest(Request{})
|
||||
if needsVision {
|
||||
t.Errorf("needsVision = true, want false")
|
||||
}
|
||||
if model != "user-model" {
|
||||
t.Errorf("model = %q, want user-model (default role for empty messages)", model)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user