diff --git a/internal/infra/llm/client.go b/internal/infra/llm/client.go new file mode 100644 index 0000000..690546a --- /dev/null +++ b/internal/infra/llm/client.go @@ -0,0 +1,81 @@ +// Package llm 提供与具体 provider 解耦的 LLM 抽象层。 +// 所有 provider 实现都把 SDK 的 message/event 形态映射到本包定义的中立类型。 +package llm + +import "context" + +// Role 标识消息发送方。 +type Role string + +const ( + RoleUser Role = "user" + RoleAssistant Role = "assistant" +) + +// ContentBlock 是消息内容的最小单元。MVP 仅 text/image/document;tool_use/tool_result 为 mcp 预留。 +type ContentBlock struct { + Type string // "text" | "image" | "document" | "tool_use" | "tool_result" + Text string // type=text + MimeType string // type=image|document + Data []byte // type=image|document(原始字节,未编码;各 provider 实现内部 base64) +} + +// Message 是一轮对话的中立表示。 +type Message struct { + Role Role + Blocks []ContentBlock +} + +// ToolSpec 为 mcp 模块预留;MVP 始终空切片。 +type ToolSpec struct { + Name string + Description string + InputSchema map[string]any +} + +// Request 是 Stream/CountTokens 共用的输入。 +type Request struct { + SystemPrompt string + Messages []Message + MaxOutputTokens int + Reasoning bool + Tools []ToolSpec // mcp 预留 +} + +// StreamEventType 枚举所有流事件。 +type StreamEventType string + +const ( + EventTextDelta StreamEventType = "text_delta" + EventThinkingDelta StreamEventType = "thinking_delta" + EventToolCall StreamEventType = "tool_call" // 预留 + EventUsage StreamEventType = "usage" + EventDone StreamEventType = "done" + EventError StreamEventType = "error" +) + +// Usage 携带本次调用的 token 计数。 +type Usage struct { + PromptTokens int + CompletionTokens int + ThinkingTokens int +} + +// StreamEvent 是流通道里的统一事件结构。 +type StreamEvent struct { + Type StreamEventType + Text string // text_delta / thinking_delta + Usage *Usage // usage + Err error // error + Reason string // done.finish_reason +} + +// LLMClient 是 Anthropic / OpenAI / Gemini 各自实现的统一接口。 +type LLMClient interface { + Provider() string + // Stream 启动一次流式请求,返回事件 channel。channel 关闭即表示流结束。 + // ctx 取消时实现层须立即终止上游 stream。 + Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) + // CountTokens 返回精确 token 数(远程接口);仅 admin 离线场景使用。 + CountTokens(ctx context.Context, modelID string, req Request) (int, error) +} diff --git a/internal/infra/llm/fake.go b/internal/infra/llm/fake.go new file mode 100644 index 0000000..b793adc --- /dev/null +++ b/internal/infra/llm/fake.go @@ -0,0 +1,53 @@ +package llm + +import ( + "context" + "errors" + "sync" +) + +// FakeClient 是测试用 LLMClient,可编程产出固定 stream 事件。 +// 注意:本文件 NOT _test.go 后缀,因下游 integration test 需作为依赖导入。 +type FakeClient struct { + mu sync.Mutex + ProviderID string + Events []StreamEvent // 流事件序列(按顺序 emit) + StreamError error // Stream() 立即返回的错误(如果非 nil) + TokenCount int // CountTokens 返回值 + StreamCalls int +} + +// NewFake 构造一个 FakeClient,默认 provider="fake"。 +func NewFake() *FakeClient { return &FakeClient{ProviderID: "fake"} } + +func (f *FakeClient) Provider() string { return f.ProviderID } + +func (f *FakeClient) Stream(ctx context.Context, modelID string, req Request) (<-chan StreamEvent, error) { + f.mu.Lock() + f.StreamCalls++ + if f.StreamError != nil { + err := f.StreamError + f.mu.Unlock() + return nil, err + } + events := append([]StreamEvent(nil), f.Events...) + f.mu.Unlock() + + ch := make(chan StreamEvent, len(events)+1) + go func() { + defer close(ch) + for _, e := range events { + select { + case <-ctx.Done(): + ch <- StreamEvent{Type: EventError, Err: errors.New("context canceled")} + return + case ch <- e: + } + } + }() + return ch, nil +} + +func (f *FakeClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) { + return f.TokenCount, nil +} diff --git a/internal/infra/llm/tokenizer.go b/internal/infra/llm/tokenizer.go new file mode 100644 index 0000000..6f7f341 --- /dev/null +++ b/internal/infra/llm/tokenizer.go @@ -0,0 +1,30 @@ +package llm + +import "unicode/utf8" + +// EstimateTokens 是兜底估算:按 rune 数 / 4 向上取整。 +// 真正的 tiktoken/远程精确计数由各 provider 的 CountTokens 实现。 +func EstimateTokens(s string) int { + if s == "" { + return 0 + } + runes := utf8.RuneCountInString(s) + if runes == 0 { + return 0 + } + return (runes + 3) / 4 +} + +// EstimateRequestTokens 累加 system prompt + 所有 messages 的 text block。 +// image/document block 暂按 0 计(多模态 token 估算各家差异大,热路径不依赖)。 +func EstimateRequestTokens(req Request) int { + total := EstimateTokens(req.SystemPrompt) + for _, m := range req.Messages { + for _, b := range m.Blocks { + if b.Type == "text" { + total += EstimateTokens(b.Text) + } + } + } + return total +} diff --git a/internal/infra/llm/tokenizer_test.go b/internal/infra/llm/tokenizer_test.go new file mode 100644 index 0000000..a06cc0d --- /dev/null +++ b/internal/infra/llm/tokenizer_test.go @@ -0,0 +1,27 @@ +package llm + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestEstimateTokens_AsciiAndCJK(t *testing.T) { + require.Equal(t, 0, EstimateTokens("")) + require.InDelta(t, 1, EstimateTokens("hi"), 1) // 2 chars / 4 ≈ 0~1 + require.InDelta(t, 25, EstimateTokens(string(make([]byte, 100))), 1) // 100 chars / 4 = 25 + // CJK 按 rune 数计算 + require.InDelta(t, 2, EstimateTokens("你好"), 1) // 2 runes / 4 ≈ 0~1,向上取整 1 +} + +func TestEstimateRequestTokens_SumsAllBlocks(t *testing.T) { + req := Request{ + SystemPrompt: "You are helpful.", // 16 chars / 4 = 4 + Messages: []Message{ + {Role: RoleUser, Blocks: []ContentBlock{{Type: "text", Text: "hello world"}}}, // 11 chars / 4 = 3 + {Role: RoleAssistant, Blocks: []ContentBlock{{Type: "text", Text: "hi"}}}, // 2 chars / 4 = 1 + }, + } + got := EstimateRequestTokens(req) + require.InDelta(t, 8, got, 2) +}