You've already forked agentic-coding-workflow
67 lines
1.8 KiB
Go
67 lines
1.8 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
|
|
}
|