You've already forked agentic-coding-workflow
126 lines
3.6 KiB
Go
126 lines
3.6 KiB
Go
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
|
|
// BlockCh, if non-nil, causes Stream goroutine to block after emitting all
|
|
// Events until ctx is cancelled or BlockCh is closed. Useful for cancel tests.
|
|
BlockCh chan struct{}
|
|
}
|
|
|
|
// 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()
|
|
|
|
blockCh := f.BlockCh
|
|
|
|
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:
|
|
}
|
|
}
|
|
// If a block channel is set, pause here until ctx is cancelled or BlockCh
|
|
// is closed. This lets cancel tests confirm mid-stream cancellation.
|
|
if blockCh != nil {
|
|
select {
|
|
case <-ctx.Done():
|
|
case <-blockCh:
|
|
}
|
|
}
|
|
}()
|
|
return ch, nil
|
|
}
|
|
|
|
func (f *FakeClient) CountTokens(ctx context.Context, modelID string, req Request) (int, error) {
|
|
return f.TokenCount, nil
|
|
}
|
|
|
|
// FakeEmbedder is a deterministic Embedder for tests. Each input maps to a
|
|
// fixed-length vector derived from its bytes, so identical text always yields the
|
|
// same vector (lets tests assert idempotency / ranking without a real endpoint).
|
|
type FakeEmbedder struct {
|
|
Dims int // output dimension; defaults to PlatformEmbeddingDims when 0
|
|
Err error // if non-nil, Embed returns it
|
|
mu sync.Mutex
|
|
Calls int
|
|
Inputs []string // accumulated inputs across calls (for assertions)
|
|
}
|
|
|
|
// NewFakeEmbedder constructs a FakeEmbedder with the platform dimension.
|
|
func NewFakeEmbedder() *FakeEmbedder { return &FakeEmbedder{Dims: PlatformEmbeddingDims} }
|
|
|
|
var _ Embedder = (*FakeEmbedder)(nil)
|
|
|
|
func (e *FakeEmbedder) dims() int {
|
|
if e.Dims > 0 {
|
|
return e.Dims
|
|
}
|
|
return PlatformEmbeddingDims
|
|
}
|
|
|
|
// EmbedDims returns the configured dimension.
|
|
func (e *FakeEmbedder) EmbedDims(string) int { return e.dims() }
|
|
|
|
// Embed returns one deterministic vector per input. The vector is a normalized
|
|
// projection of a small hash of the input across the dimension, so distinct
|
|
// inputs are distinguishable and identical inputs are identical.
|
|
func (e *FakeEmbedder) Embed(_ context.Context, _ string, inputs []string) ([][]float32, error) {
|
|
e.mu.Lock()
|
|
e.Calls++
|
|
e.Inputs = append(e.Inputs, inputs...)
|
|
err := e.Err
|
|
e.mu.Unlock()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
d := e.dims()
|
|
out := make([][]float32, len(inputs))
|
|
for i, s := range inputs {
|
|
v := make([]float32, d)
|
|
// Seed from a stable rolling hash of the bytes.
|
|
var h uint32 = 2166136261
|
|
for _, b := range []byte(s) {
|
|
h ^= uint32(b)
|
|
h *= 16777619
|
|
}
|
|
for j := 0; j < d; j++ {
|
|
h ^= uint32(j) * 2654435761
|
|
h *= 16777619
|
|
// Map to [-1,1).
|
|
v[j] = float32(int32(h)) / float32(1<<31)
|
|
}
|
|
out[i] = v
|
|
}
|
|
return out, nil
|
|
}
|