You've already forked agentic-coding-workflow
feat(infra/llm): neutral types + LLMClient interface + estimate tokenizer + fake
This commit is contained in:
@@ -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)
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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
|
||||||
|
}
|
||||||
@@ -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)
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user