feat: 重构用eino框架 #128
@@ -1,239 +0,0 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
// OpenAIService 基于 OpenAI Chat Completions API 的 LLM 实现。
|
||||
type OpenAIService struct {
|
||||
apiKey string
|
||||
model string
|
||||
endpoint string
|
||||
timeout time.Duration
|
||||
logger *zap.SugaredLogger
|
||||
client *http.Client
|
||||
}
|
||||
|
||||
// NewOpenAIService 创建 OpenAI LLM 服务。
|
||||
// model、endpoint 由 config 层保证非空。
|
||||
func NewOpenAIService(apiKey, model, endpoint string, timeoutSec, httpClientTimeoutSec int, logger *zap.SugaredLogger) *OpenAIService {
|
||||
timeout := time.Duration(timeoutSec) * time.Second
|
||||
if timeout <= 0 {
|
||||
timeout = 10 * time.Second
|
||||
}
|
||||
httpClientTimeout := time.Duration(httpClientTimeoutSec) * time.Second
|
||||
if httpClientTimeout <= 0 {
|
||||
httpClientTimeout = 60 * time.Second
|
||||
}
|
||||
return &OpenAIService{
|
||||
apiKey: apiKey,
|
||||
model: model,
|
||||
endpoint: endpoint,
|
||||
timeout: timeout,
|
||||
logger: logger,
|
||||
client: &http.Client{Timeout: httpClientTimeout},
|
||||
}
|
||||
}
|
||||
|
||||
// --- OpenAI API 请求/响应结构 ---
|
||||
|
||||
type chatRequest struct {
|
||||
Model string `json:"model"`
|
||||
Messages []chatMessage `json:"messages"`
|
||||
Stream bool `json:"stream"`
|
||||
}
|
||||
|
||||
type chatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content []contentPart `json:"content"`
|
||||
}
|
||||
|
||||
type contentPart struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
ImageURL *imageURL `json:"image_url,omitempty"`
|
||||
}
|
||||
|
||||
type imageURL struct {
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// streamDelta SSE 流式响应的单个 delta。
|
||||
type streamDelta struct {
|
||||
Choices []struct {
|
||||
Delta struct {
|
||||
Content string `json:"content"`
|
||||
} `json:"delta"`
|
||||
FinishReason *string `json:"finish_reason"`
|
||||
} `json:"choices"`
|
||||
Usage *struct {
|
||||
PromptTokens int `json:"prompt_tokens"`
|
||||
CompletionTokens int `json:"completion_tokens"`
|
||||
TotalTokens int `json:"total_tokens"`
|
||||
} `json:"usage"`
|
||||
Model string `json:"model"`
|
||||
}
|
||||
|
||||
// ChatStream 实现 llm.Service。调用 OpenAI Chat Completions API 流式推理。
|
||||
func (o *OpenAIService) ChatStream(ctx context.Context, req Request) (<-chan Chunk, error) {
|
||||
// 构建请求
|
||||
messages := o.buildMessages(req)
|
||||
|
||||
body := chatRequest{
|
||||
Model: o.model,
|
||||
Messages: messages,
|
||||
Stream: true,
|
||||
}
|
||||
|
||||
payload, err := json.Marshal(body)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llm: marshal request: %w", err)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("llm: marshal request: %w", err)
|
||||
}
|
||||
|
||||
// 创建带超时的 context
|
||||
ctx, cancel := context.WithTimeout(ctx, o.timeout)
|
||||
|
||||
httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint+"/chat/completions", bytes.NewReader(payload))
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("llm: create request: %w", err)
|
||||
}
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
httpReq.Header.Set("Authorization", "Bearer "+o.apiKey)
|
||||
|
||||
resp, err := o.client.Do(httpReq)
|
||||
if err != nil {
|
||||
cancel()
|
||||
return nil, fmt.Errorf("llm: send request: %w", err)
|
||||
}
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
cancel()
|
||||
bodyBytes, _ := io.ReadAll(resp.Body)
|
||||
resp.Body.Close()
|
||||
return nil, fmt.Errorf("llm: api error (status %d): %s", resp.StatusCode, string(bodyBytes))
|
||||
}
|
||||
|
||||
// 启动 goroutine 解析 SSE 流
|
||||
ch := make(chan Chunk, 64)
|
||||
go func() {
|
||||
defer close(ch)
|
||||
defer cancel()
|
||||
defer resp.Body.Close()
|
||||
|
||||
o.parseSSEStream(resp.Body, ch)
|
||||
}()
|
||||
|
||||
return ch, nil
|
||||
}
|
||||
|
||||
// parseSSEStream 解析 SSE 流，将 delta 发送到 channel。
|
||||
func (o *OpenAIService) parseSSEStream(body io.Reader, ch chan<- Chunk) {
|
||||
scanner := bufio.NewScanner(body)
|
||||
scanner.Buffer(make([]byte, 0, 64*1024), 256*1024)
|
||||
|
||||
var fullText strings.Builder
|
||||
var lastModel string
|
||||
|
||||
for scanner.Scan() {
|
||||
line := scanner.Text()
|
||||
|
||||
// SSE 格式：data: {...}
|
||||
if !strings.HasPrefix(line, "data: ") {
|
||||
continue
|
||||
}
|
||||
data := strings.TrimPrefix(line, "data: ")
|
||||
if data == "[DONE]" {
|
||||
// 流结束，发送最终 chunk
|
||||
ch <- Chunk{Delta: "", Done: true, Model: lastModel}
|
||||
return
|
||||
}
|
||||
|
||||
var delta streamDelta
|
||||
if err := json.Unmarshal([]byte(data), &delta); err != nil {
|
||||
o.logger.Warnw("llm: unmarshal delta failed", "error", err, "data", data)
|
||||
continue
|
||||
}
|
||||
|
||||
if delta.Model != "" {
|
||||
lastModel = delta.Model
|
||||
}
|
||||
|
||||
// 提取增量文本
|
||||
if len(delta.Choices) > 0 {
|
||||
content := delta.Choices[0].Delta.Content
|
||||
if content != "" {
|
||||
fullText.WriteString(content)
|
||||
ch <- Chunk{Delta: content, Done: false, Model: lastModel}
|
||||
}
|
||||
|
||||
// 某些模型在最后一个 choice 中携带 usage
|
||||
if delta.Choices[0].FinishReason != nil && delta.Usage != nil {
|
||||
ch <- Chunk{
|
||||
Delta: "",
|
||||
Done: true,
|
||||
Model: lastModel,
|
||||
TokensUsed: &TokenUsage{
|
||||
Prompt: delta.Usage.PromptTokens,
|
||||
Completion: delta.Usage.CompletionTokens,
|
||||
Total: delta.Usage.TotalTokens,
|
||||
},
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// scanner 结束但没收到 [DONE]
|
||||
if err := scanner.Err(); err != nil {
|
||||
o.logger.Warnw("llm: scan error", "error", err)
|
||||
}
|
||||
ch <- Chunk{Delta: "", Done: true, Model: lastModel}
|
||||
}
|
||||
|
||||
// buildMessages 构建 OpenAI Chat API 的 messages 数组。
|
||||
func (o *OpenAIService) buildMessages(req Request) []chatMessage {
|
||||
var messages []chatMessage
|
||||
|
||||
// System prompt（情景覆盖优先）
|
||||
messages = append(messages, chatMessage{
|
||||
Role: "system",
|
||||
Content: []contentPart{{Type: "text", Text: BuildSystemPrompt(req.Language, "", req.SystemPrompt)}},
|
||||
})
|
||||
|
||||
// 历史消息
|
||||
for _, msg := range req.History {
|
||||
messages = append(messages, chatMessage{
|
||||
Role: msg.Role,
|
||||
Content: []contentPart{{Type: "text", Text: msg.Content}},
|
||||
})
|
||||
}
|
||||
|
||||
// 当前用户消息（图像 + 文本）
|
||||
var parts []contentPart
|
||||
if len(req.Image) > 0 {
|
||||
b64 := base64.StdEncoding.EncodeToString(req.Image)
|
||||
parts = append(parts, contentPart{
|
||||
Type: "image_url",
|
||||
ImageURL: &imageURL{URL: "data:image/jpeg;base64," + b64},
|
||||
})
|
||||
}
|
||||
parts = append(parts, contentPart{Type: "text", Text: req.Text})
|
||||
messages = append(messages, chatMessage{Role: "user", Content: parts})
|
||||
|
||||
return messages
|
||||
}
|
||||
@@ -1,251 +0,0 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"github.com/hhs/camtalk/internal/models"
|
||||
)
|
||||
|
||||
// mockLLMServer 创建模拟 OpenAI SSE 流式响应的 HTTP 服务器。
|
||||
func mockLLMServer(t *testing.T, handler http.HandlerFunc) *httptest.Server {
|
||||
t.Helper()
|
||||
return httptest.NewServer(handler)
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_Success(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
// 验证请求
|
||||
if r.Method != http.MethodPost {
|
||||
t.Errorf("method = %s, want POST", r.Method)
|
||||
}
|
||||
if !strings.Contains(r.URL.Path, "/chat/completions") {
|
||||
t.Errorf("path = %s, should contain /chat/completions", r.URL.Path)
|
||||
}
|
||||
auth := r.Header.Get("Authorization")
|
||||
if auth != "Bearer test-key" {
|
||||
t.Errorf("Authorization = %q, want %q", auth, "Bearer test-key")
|
||||
}
|
||||
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
flusher, ok := w.(http.Flusher)
|
||||
if !ok {
|
||||
t.Fatal("ResponseWriter does not support Flusher")
|
||||
}
|
||||
|
||||
// 发送几个 delta
|
||||
deltas := []string{"你好", "世界", "！"}
|
||||
for _, d := range deltas {
|
||||
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"%s\"}}],\"model\":\"gpt-4o\"}\n\n", d)
|
||||
flusher.Flush()
|
||||
}
|
||||
|
||||
// 发送 [DONE]
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, 60, zap.NewNop().Sugar())
|
||||
|
||||
ch, err := svc.ChatStream(context.Background(), Request{
|
||||
Text: "这是什么？",
|
||||
Language: "zh-CN",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream() error: %v", err)
|
||||
}
|
||||
|
||||
var chunks []Chunk
|
||||
for c := range ch {
|
||||
chunks = append(chunks, c)
|
||||
}
|
||||
|
||||
// 应该有 3 个文本 chunk + 1 个 Done chunk
|
||||
if len(chunks) != 4 {
|
||||
t.Fatalf("got %d chunks, want 4", len(chunks))
|
||||
}
|
||||
|
||||
// 验证文本内容
|
||||
if chunks[0].Delta != "你好" {
|
||||
t.Errorf("chunk[0].Delta = %q, want %q", chunks[0].Delta, "你好")
|
||||
}
|
||||
if chunks[1].Delta != "世界" {
|
||||
t.Errorf("chunk[1].Delta = %q, want %q", chunks[1].Delta, "世界")
|
||||
}
|
||||
|
||||
// 验证最后一个 chunk 是 Done
|
||||
last := chunks[len(chunks)-1]
|
||||
if !last.Done {
|
||||
t.Error("last chunk should be Done")
|
||||
}
|
||||
if last.Model != "gpt-4o" {
|
||||
t.Errorf("last chunk Model = %q, want %q", last.Model, "gpt-4o")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_WithImage(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}],\"model\":\"gpt-4o\"}\n\n")
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, 60, zap.NewNop().Sugar())
|
||||
|
||||
ch, err := svc.ChatStream(context.Background(), Request{
|
||||
Image: []byte("fake-jpeg-data"),
|
||||
Text: "描述图片",
|
||||
Language: "zh-CN",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream() error: %v", err)
|
||||
}
|
||||
|
||||
// 消费 channel
|
||||
for range ch {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_WithHistory(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"ok\"}}],\"model\":\"gpt-4o\"}\n\n")
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, 60, zap.NewNop().Sugar())
|
||||
|
||||
ch, err := svc.ChatStream(context.Background(), Request{
|
||||
Text: "继续",
|
||||
Language: "zh-CN",
|
||||
History: []models.Message{
|
||||
{Role: "user", Content: "你好"},
|
||||
{Role: "assistant", Content: "你好！有什么可以帮助你的吗？"},
|
||||
},
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream() error: %v", err)
|
||||
}
|
||||
|
||||
for range ch {
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_APIError(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "application/json")
|
||||
w.WriteHeader(http.StatusUnauthorized)
|
||||
fmt.Fprintf(w, `{"error":{"message":"Invalid API key"}}`)
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("bad-key", "gpt-4o", srv.URL, 10, 60, zap.NewNop().Sugar())
|
||||
|
||||
_, err := svc.ChatStream(context.Background(), Request{
|
||||
Text: "test",
|
||||
})
|
||||
if err == nil {
|
||||
t.Fatal("ChatStream() should return error for 401")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "401") {
|
||||
t.Errorf("error should mention 401, got: %v", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_Timeout(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
// 模拟慢响应
|
||||
time.Sleep(5 * time.Second)
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"late\"}}]}\n\n")
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 1, 60, zap.NewNop().Sugar()) // 1s timeout
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
|
||||
ch, err := svc.ChatStream(ctx, Request{Text: "test"})
|
||||
if err != nil {
|
||||
// 超时可能在建立连接时或读取时发生
|
||||
return
|
||||
}
|
||||
|
||||
// 如果连接成功，消费 channel 应该超时
|
||||
var gotContent bool
|
||||
for c := range ch {
|
||||
if c.Delta != "" {
|
||||
gotContent = true
|
||||
}
|
||||
}
|
||||
if gotContent {
|
||||
t.Error("should not receive content before timeout")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIService_ChatStream_UsageInResponse(t *testing.T) {
|
||||
srv := mockLLMServer(t, func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
// 带 usage 的最后一个 chunk
|
||||
fmt.Fprintf(w, "data: {\"choices\":[{\"delta\":{\"content\":\"hi\"},\"finish_reason\":\"stop\"}],\"model\":\"gpt-4o\",\"usage\":{\"prompt_tokens\":10,\"completion_tokens\":5,\"total_tokens\":15}}\n\n")
|
||||
fmt.Fprintf(w, "data: [DONE]\n\n")
|
||||
})
|
||||
defer srv.Close()
|
||||
|
||||
svc := NewOpenAIService("test-key", "gpt-4o", srv.URL, 10, 60, zap.NewNop().Sugar())
|
||||
|
||||
ch, err := svc.ChatStream(context.Background(), Request{Text: "test"})
|
||||
if err != nil {
|
||||
t.Fatalf("ChatStream() error: %v", err)
|
||||
}
|
||||
|
||||
var last Chunk
|
||||
for c := range ch {
|
||||
last = c
|
||||
}
|
||||
|
||||
if !last.Done {
|
||||
t.Error("last chunk should be Done")
|
||||
}
|
||||
if last.TokensUsed == nil {
|
||||
t.Fatal("last chunk should have TokensUsed")
|
||||
}
|
||||
if last.TokensUsed.Total != 15 {
|
||||
t.Errorf("TokensUsed.Total = %d, want 15", last.TokensUsed.Total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildSystemPrompt(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
language string
|
||||
detailLevel string
|
||||
wantContain string
|
||||
}{
|
||||
{"chinese default", "zh-CN", "", "视觉助手"},
|
||||
{"chinese high", "zh-CN", "high", "更详细"},
|
||||
{"english default", "en", "", "visual assistant"},
|
||||
{"english high", "en", "high", "detailed"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := BuildSystemPrompt(tt.language, tt.detailLevel, "")
|
||||
if !strings.Contains(got, tt.wantContain) {
|
||||
t.Errorf("BuildSystemPrompt(%q, %q, \"\") should contain %q", tt.language, tt.detailLevel, tt.wantContain)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
235
backend/internal/eino/graph_test.go
Normal file
235
backend/internal/eino/graph_test.go
Normal file
@@ -0,0 +1,235 @@
|
||||
package eino
|
||||
|
||||
import (
|
||||
"context"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/hhs/camtalk/internal/ai/stt"
|
||||
"github.com/hhs/camtalk/internal/ai/tts"
|
||||
"github.com/hhs/camtalk/internal/models"
|
||||
"github.com/hhs/camtalk/internal/orchestrator"
|
||||
)
|
||||
|
||||
// --- Mock STT Service ---
|
||||
|
||||
type mockSTTService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *mockSTTService) Recognize(ctx context.Context, audio []byte, opts stt.Options) (string, error) {
|
||||
args := m.Called(ctx, audio, opts)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
// --- Mock TTS Service ---
|
||||
|
||||
type mockTTSService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *mockTTSService) SynthesizeStream(ctx context.Context, textStream <-chan string, opts tts.Options) (<-chan tts.Chunk, error) {
|
||||
args := m.Called(ctx, textStream, opts)
|
||||
return args.Get(0).(<-chan tts.Chunk), args.Error(1)
|
||||
}
|
||||
|
||||
// --- Mock Sender ---
|
||||
|
||||
type mockSender struct {
|
||||
mock.Mock
|
||||
STTResults []models.WsSTTResult
|
||||
LLMChunks []models.WsLLMChunk
|
||||
LLMDones []models.WsLLMDone
|
||||
TTSAudios []models.WsTTSAudio
|
||||
Errors []models.WsError
|
||||
}
|
||||
|
||||
func (m *mockSender) SendSTTResult(result models.WsSTTResult) error {
|
||||
m.STTResults = append(m.STTResults, result)
|
||||
return m.Called(result).Error(0)
|
||||
}
|
||||
|
||||
func (m *mockSender) SendLLMChunk(chunk models.WsLLMChunk) error {
|
||||
m.LLMChunks = append(m.LLMChunks, chunk)
|
||||
return m.Called(chunk).Error(0)
|
||||
}
|
||||
|
||||
func (m *mockSender) SendLLMDone(done models.WsLLMDone) error {
|
||||
m.LLMDones = append(m.LLMDones, done)
|
||||
return m.Called(done).Error(0)
|
||||
}
|
||||
|
||||
func (m *mockSender) SendTTSAudio(audio models.WsTTSAudio) error {
|
||||
m.TTSAudios = append(m.TTSAudios, audio)
|
||||
return m.Called(audio).Error(0)
|
||||
}
|
||||
|
||||
func (m *mockSender) SendError(err models.WsError) error {
|
||||
m.Errors = append(m.Errors, err)
|
||||
return m.Called(err).Error(0)
|
||||
}
|
||||
|
||||
// --- Tests ---
|
||||
|
||||
func TestDetectImageMimeType(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
data []byte
|
||||
expected string
|
||||
}{
|
||||
{"JPEG", []byte{0xFF, 0xD8, 0xFF, 0xE0}, "image/jpeg"},
|
||||
{"PNG", []byte{0x89, 0x50, 0x4E, 0x47}, "image/png"},
|
||||
{"GIF", []byte{0x47, 0x49, 0x46, 0x38}, "image/gif"},
|
||||
{"WebP", []byte{0x52, 0x49, 0x46, 0x46}, "image/webp"},
|
||||
{"Unknown", []byte{0x00, 0x00, 0x00}, "image/jpeg"},
|
||||
{"Short", []byte{0xFF}, "image/jpeg"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := detectImageMimeType(tt.data)
|
||||
assert.Equal(t, tt.expected, result)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPipelineInput(t *testing.T) {
|
||||
req := models.WsQuery{
|
||||
Text: "你好",
|
||||
RequestID: "req-1",
|
||||
}
|
||||
sess := &models.Session{
|
||||
Config: models.SessionConfig{
|
||||
Language: "zh-CN",
|
||||
Scenario: "free_chat",
|
||||
TTSEnabled: true,
|
||||
},
|
||||
}
|
||||
|
||||
input := buildPipelineInput(req, "sess-1", sess, nil, nil)
|
||||
require.Equal(t, "你好", input.Text)
|
||||
require.Equal(t, "sess-1", input.SessionID)
|
||||
require.Equal(t, "req-1", input.RequestID)
|
||||
require.Equal(t, "zh-CN", input.Language)
|
||||
require.Equal(t, "free_chat", input.Scenario)
|
||||
require.True(t, input.TTSEnabled)
|
||||
}
|
||||
|
||||
func TestBuildPipelineInput_WithAudioData(t *testing.T) {
|
||||
req := models.WsQuery{
|
||||
Audio: "base64audio",
|
||||
RequestID: "req-2",
|
||||
}
|
||||
sess := &models.Session{
|
||||
Config: models.SessionConfig{
|
||||
Language: "en",
|
||||
Scenario: "free_chat",
|
||||
TTSEnabled: false,
|
||||
},
|
||||
}
|
||||
|
||||
audioData := []byte("fake-audio-bytes")
|
||||
imageData := []byte("fake-image-bytes")
|
||||
|
||||
input := buildPipelineInput(req, "sess-2", sess, audioData, imageData)
|
||||
require.Equal(t, audioData, input.AudioData)
|
||||
require.Equal(t, imageData, input.ImageData)
|
||||
require.False(t, input.TTSEnabled)
|
||||
require.Equal(t, "en", input.Language)
|
||||
}
|
||||
|
||||
func TestPipelineState_AppendAndGet(t *testing.T) {
|
||||
state := genLocalState(context.Background())
|
||||
|
||||
state.AppendText("Hello ")
|
||||
state.AppendText("World")
|
||||
|
||||
require.Equal(t, "Hello World", state.GetFullResponse())
|
||||
}
|
||||
|
||||
func TestPipelineState_ConcurrentAccess(t *testing.T) {
|
||||
state := genLocalState(context.Background())
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
for i := 0; i < 100; i++ {
|
||||
state.AppendText("a")
|
||||
}
|
||||
close(done)
|
||||
}()
|
||||
|
||||
for i := 0; i < 100; i++ {
|
||||
_ = state.GetFullResponse()
|
||||
}
|
||||
|
||||
<-done
|
||||
require.Equal(t, 100, len(state.GetFullResponse()))
|
||||
}
|
||||
|
||||
func TestContextInjection(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
sender := &mockSender{}
|
||||
ctx = WithSender(ctx, sender)
|
||||
ctx = WithRequestID(ctx, "req-123")
|
||||
ctx = WithSessionID(ctx, "sess-456")
|
||||
ctx = WithStartTime(ctx, time.Now())
|
||||
ctx = WithPipelineState(ctx, genLocalState(ctx))
|
||||
|
||||
require.NotNil(t, senderFromCtx(ctx))
|
||||
require.Equal(t, "req-123", requestIDFromCtx(ctx))
|
||||
require.NotNil(t, stateFromCtx(ctx))
|
||||
}
|
||||
|
||||
func TestLatencyFromCtx(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
|
||||
// No start time set
|
||||
require.Equal(t, int64(0), latencyFromCtx(ctx))
|
||||
|
||||
// With start time
|
||||
start := time.Now().Add(-100 * time.Millisecond)
|
||||
ctx = WithStartTime(ctx, start)
|
||||
latency := latencyFromCtx(ctx)
|
||||
require.Greater(t, latency, int64(0))
|
||||
require.Less(t, latency, int64(1000)) // should be < 1 second
|
||||
}
|
||||
|
||||
func TestEinoOrchestrator_ImplementsInterface(t *testing.T) {
|
||||
// Compile-time check that EinoOrchestrator implements orchestrator.Orchestrator
|
||||
var _ orchestrator.Orchestrator = (*EinoOrchestrator)(nil)
|
||||
}
|
||||
|
||||
func TestNewSTTLambda_ReturnsNonNil(t *testing.T) {
|
||||
mockSTT := &mockSTTService{}
|
||||
lambda := NewSTTLambda(mockSTT)
|
||||
require.NotNil(t, lambda)
|
||||
}
|
||||
|
||||
func TestNewHistoryLambda_ReturnsNonNil(t *testing.T) {
|
||||
fetcher := func(ctx context.Context, sessionID string, limit int) ([]models.Message, error) {
|
||||
return nil, nil
|
||||
}
|
||||
lambda := NewHistoryLambda(fetcher, 10)
|
||||
require.NotNil(t, lambda)
|
||||
}
|
||||
|
||||
func TestNewSplitterLambda_ReturnsNonNil(t *testing.T) {
|
||||
lambda := NewSplitterLambda()
|
||||
require.NotNil(t, lambda)
|
||||
}
|
||||
|
||||
func TestNewTTSLambda_ReturnsNonNil(t *testing.T) {
|
||||
mockTTS := &mockTTSService{}
|
||||
lambda := NewTTSLambda(mockTTS, "alloy", 1.0, "mp3", 24000)
|
||||
require.NotNil(t, lambda)
|
||||
}
|
||||
|
||||
func TestNewDoneLambda_ReturnsNonNil(t *testing.T) {
|
||||
lambda := NewDoneLambda("test-model")
|
||||
require.NotNil(t, lambda)
|
||||
}
|
||||
@@ -1,403 +0,0 @@
|
||||
package orchestrator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/hhs/camtalk/internal/ai/llm"
|
||||
"github.com/hhs/camtalk/internal/ai/stt"
|
||||
"github.com/hhs/camtalk/internal/ai/tts"
|
||||
"github.com/hhs/camtalk/internal/config"
|
||||
"github.com/hhs/camtalk/internal/logger"
|
||||
"github.com/hhs/camtalk/internal/models"
|
||||
"github.com/hhs/camtalk/internal/session"
|
||||
)
|
||||
|
||||
// Pipeline 实现 Orchestrator 接口，管理 STT → LLM → TTS 流式管道。
|
||||
type Pipeline struct {
|
||||
sttService stt.Service
|
||||
llmService llm.Service
|
||||
ttsService tts.Service
|
||||
sessionMgr session.Manager
|
||||
model string // LLM 模型名，用于 llm_done 上报
|
||||
ttsVoice string // TTS 音色
|
||||
ttsSpeed float64 // TTS 语速
|
||||
ttsOutputFmt string // TTS 输出格式
|
||||
ttsSampleRate int // TTS 输出采样率
|
||||
}
|
||||
|
||||
// New 创建 Pipeline 实例。
|
||||
func New(
|
||||
sttService stt.Service,
|
||||
llmService llm.Service,
|
||||
ttsService tts.Service,
|
||||
sessionMgr session.Manager,
|
||||
cfg *config.Config,
|
||||
) *Pipeline {
|
||||
return &Pipeline{
|
||||
sttService: sttService,
|
||||
llmService: llmService,
|
||||
ttsService: ttsService,
|
||||
sessionMgr: sessionMgr,
|
||||
model: cfg.AI.LLM.Model,
|
||||
ttsVoice: cfg.AI.TTS.Voice,
|
||||
ttsSpeed: cfg.AI.TTS.Speed,
|
||||
ttsOutputFmt: cfg.AI.TTS.OutputFormat,
|
||||
ttsSampleRate: cfg.AI.TTS.SampleRate,
|
||||
}
|
||||
}
|
||||
|
||||
// ProcessQuery 实现 Orchestrator 接口。
|
||||
func (p *Pipeline) ProcessQuery(
|
||||
ctx context.Context,
|
||||
sessionID string,
|
||||
req models.WsQuery,
|
||||
history []models.Message,
|
||||
sender Sender,
|
||||
) error {
|
||||
log := logger.Log
|
||||
startTime := time.Now()
|
||||
|
||||
// 解码音频数据（文本输入模式可跳过）
|
||||
var audio []byte
|
||||
if req.Text == "" && req.Audio != "" {
|
||||
var err error
|
||||
audio, err = base64.StdEncoding.DecodeString(req.Audio)
|
||||
if err != nil {
|
||||
log.Errorw("音频解码失败", "error", err)
|
||||
sender.SendError(models.WsError{
|
||||
Type: "error",
|
||||
RequestID: req.RequestID,
|
||||
Code: "INVALID_MESSAGE",
|
||||
Message: "音频数据解码失败",
|
||||
})
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 解码图片数据（可选）
|
||||
var image []byte
|
||||
if req.Image != "" {
|
||||
var err error
|
||||
image, err = base64.StdEncoding.DecodeString(req.Image)
|
||||
if err != nil {
|
||||
log.Errorw("图片解码失败", "error", err)
|
||||
sender.SendError(models.WsError{
|
||||
Type: "error",
|
||||
RequestID: req.RequestID,
|
||||
Code: "INVALID_MESSAGE",
|
||||
Message: "图片数据解码失败",
|
||||
})
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 设置活跃请求
|
||||
if err := p.sessionMgr.SetActiveRequest(ctx, sessionID, req.RequestID); err != nil {
|
||||
log.Errorw("设置活跃请求失败", "error", err)
|
||||
}
|
||||
defer p.sessionMgr.ClearActiveRequest(ctx, sessionID)
|
||||
|
||||
// 获取会话配置
|
||||
sess, err := p.sessionMgr.Get(ctx, sessionID)
|
||||
if err != nil {
|
||||
log.Errorw("获取会话失败", "error", err)
|
||||
sender.SendError(models.WsError{
|
||||
Type: "error",
|
||||
RequestID: req.RequestID,
|
||||
Code: "SESSION_NOT_FOUND",
|
||||
Message: "会话不存在",
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// Step 1: 获取用户文本（语音识别或直接使用输入文本）
|
||||
var userText string
|
||||
if req.Text != "" {
|
||||
// 文本输入模式：跳过 STT，直接使用用户输入的文本
|
||||
log.Infow("使用文本输入", "request_id", req.RequestID, "text", req.Text)
|
||||
userText = req.Text
|
||||
|
||||
// 发送 stt_result 以保持前端消息流一致性
|
||||
if err := sender.SendSTTResult(models.WsSTTResult{
|
||||
Type: "stt_result",
|
||||
RequestID: req.RequestID,
|
||||
Text: userText,
|
||||
IsFinal: true,
|
||||
}); err != nil {
|
||||
log.Errorw("发送 STT 结果失败", "error", err)
|
||||
}
|
||||
} else {
|
||||
// 语音模式：执行 STT 语音识别
|
||||
log.Infow("开始语音识别", "request_id", req.RequestID, "audio_bytes", len(audio))
|
||||
sttResult, err := p.sttService.Recognize(ctx, audio, stt.Options{
|
||||
Encoding: "pcm_s16le",
|
||||
SampleRate: 16000,
|
||||
Language: sess.Config.Language,
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorw("语音识别失败", "error", err, "audio_bytes", len(audio))
|
||||
sender.SendError(models.WsError{
|
||||
Type: "error",
|
||||
RequestID: req.RequestID,
|
||||
Code: "STT_ERROR",
|
||||
Message: "语音识别失败: " + err.Error(),
|
||||
})
|
||||
return err
|
||||
}
|
||||
userText = sttResult
|
||||
|
||||
// STT 返回空文本：未识别到语音，发送结果后直接返回（不调 LLM）
|
||||
if strings.TrimSpace(userText) == "" {
|
||||
log.Infow("语音识别结果为空", "request_id", req.RequestID)
|
||||
userText = "（未识别到语音）"
|
||||
if err := sender.SendSTTResult(models.WsSTTResult{
|
||||
Type: "stt_result",
|
||||
RequestID: req.RequestID,
|
||||
Text: userText,
|
||||
IsFinal: true,
|
||||
}); err != nil {
|
||||
log.Errorw("发送 STT 结果失败", "error", err)
|
||||
}
|
||||
// 发送空的 llm_done 以结束本轮处理
|
||||
latency := time.Since(startTime).Milliseconds()
|
||||
_ = sender.SendLLMDone(models.WsLLMDone{
|
||||
Type: "llm_done",
|
||||
RequestID: req.RequestID,
|
||||
FullText: "",
|
||||
Model: p.model,
|
||||
LatencyMs: latency,
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
// 发送 STT 结果
|
||||
if err := sender.SendSTTResult(models.WsSTTResult{
|
||||
Type: "stt_result",
|
||||
RequestID: req.RequestID,
|
||||
Text: userText,
|
||||
IsFinal: true,
|
||||
}); err != nil {
|
||||
log.Errorw("发送 STT 结果失败", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 追加用户消息到历史
|
||||
p.sessionMgr.AppendMessage(ctx, sessionID, models.Message{
|
||||
Role: "user",
|
||||
Content: userText,
|
||||
})
|
||||
|
||||
// Step 2+3: LLM 流式推理 + TTS 并行合成
|
||||
log.Infow("开始 LLM 推理", "request_id", req.RequestID, "scenario", sess.Config.Scenario)
|
||||
llmReq := llm.Request{
|
||||
Image: image,
|
||||
Text: userText,
|
||||
History: history,
|
||||
Language: sess.Config.Language,
|
||||
SystemPrompt: llm.GetScenarioPrompt(sess.Config.Scenario, sess.Config.Language),
|
||||
}
|
||||
|
||||
llmStream, err := p.llmService.ChatStream(ctx, llmReq)
|
||||
if err != nil {
|
||||
log.Errorw("LLM 流式推理启动失败", "error", err)
|
||||
sender.SendError(models.WsError{
|
||||
Type: "error",
|
||||
RequestID: req.RequestID,
|
||||
Code: "LLM_ERROR",
|
||||
Message: "LLM 推理失败",
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
// 创建句子切分器
|
||||
sentenceCh := make(chan string, 4)
|
||||
splitter := NewSplitter(sentenceCh)
|
||||
|
||||
// 并行：LLM 消费 + TTS 合成
|
||||
var wg sync.WaitGroup
|
||||
var fullText string
|
||||
var ttsErr error
|
||||
|
||||
// goroutine 1: 消费 LLM token + 句子切分
|
||||
var tokenUsage *llm.TokenUsage
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer close(sentenceCh)
|
||||
fullText, tokenUsage = p.consumeLLMStream(ctx, llmStream, req.RequestID, sender, splitter)
|
||||
}()
|
||||
|
||||
// goroutine 2: TTS 合成（如果启用）
|
||||
if sess.Config.TTSEnabled {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
log.Infow("开始 TTS 合成", "request_id", req.RequestID)
|
||||
ttsErr = p.synthesizeTTS(ctx, sentenceCh, req.RequestID, sender)
|
||||
}()
|
||||
} else {
|
||||
// 如果 TTS 未启用，需要消费 sentenceCh 防止阻塞
|
||||
go func() {
|
||||
for range sentenceCh {
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
// 等待所有 goroutine 完成
|
||||
wg.Wait()
|
||||
|
||||
// TTS 失败静默跳过
|
||||
if ttsErr != nil {
|
||||
log.Warnw("TTS 合成失败（已跳过）", "error", ttsErr)
|
||||
}
|
||||
|
||||
// 追加助手消息到历史
|
||||
p.sessionMgr.AppendMessage(ctx, sessionID, models.Message{
|
||||
Role: "assistant",
|
||||
Content: fullText,
|
||||
})
|
||||
|
||||
// 发送 llm_done
|
||||
latency := time.Since(startTime).Milliseconds()
|
||||
done := models.WsLLMDone{
|
||||
Type: "llm_done",
|
||||
RequestID: req.RequestID,
|
||||
FullText: fullText,
|
||||
Model: p.model,
|
||||
LatencyMs: latency,
|
||||
}
|
||||
if tokenUsage != nil {
|
||||
done.TokensUsed = struct {
|
||||
Prompt int `json:"prompt"`
|
||||
Completion int `json:"completion"`
|
||||
Total int `json:"total"`
|
||||
}{
|
||||
Prompt: tokenUsage.Prompt,
|
||||
Completion: tokenUsage.Completion,
|
||||
Total: tokenUsage.Total,
|
||||
}
|
||||
}
|
||||
if err := sender.SendLLMDone(done); err != nil {
|
||||
log.Errorw("发送 llm_done 失败", "error", err)
|
||||
}
|
||||
|
||||
log.Infow("查询处理完成",
|
||||
"request_id", req.RequestID,
|
||||
"latency_ms", latency,
|
||||
"text_length", utf8.RuneCountInString(fullText),
|
||||
)
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// consumeLLMStream 消费 LLM 流式输出，发送 llm_chunk 并进行句子切分。
|
||||
// 返回完整文本和 token 用量。
|
||||
func (p *Pipeline) consumeLLMStream(
|
||||
ctx context.Context,
|
||||
stream <-chan llm.Chunk,
|
||||
requestID string,
|
||||
sender Sender,
|
||||
splitter *Splitter,
|
||||
) (string, *llm.TokenUsage) {
|
||||
log := logger.Log
|
||||
var fullText strings.Builder
|
||||
var tokenUsage *llm.TokenUsage
|
||||
|
||||
for chunk := range stream {
|
||||
// 检查上下文是否已取消
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Infow("LLM 流被中断", "request_id", requestID)
|
||||
return fullText.String(), tokenUsage
|
||||
default:
|
||||
}
|
||||
|
||||
if chunk.Done {
|
||||
// 流结束，记录 token 用量
|
||||
if chunk.TokensUsed != nil {
|
||||
tokenUsage = chunk.TokensUsed
|
||||
log.Infow("LLM 用量统计",
|
||||
"request_id", requestID,
|
||||
"prompt_tokens", tokenUsage.Prompt,
|
||||
"completion_tokens", tokenUsage.Completion,
|
||||
"total_tokens", tokenUsage.Total,
|
||||
)
|
||||
}
|
||||
break
|
||||
}
|
||||
|
||||
// 累积全文
|
||||
fullText.WriteString(chunk.Delta)
|
||||
|
||||
// 发送 llm_chunk
|
||||
if err := sender.SendLLMChunk(models.WsLLMChunk{
|
||||
Type: "llm_chunk",
|
||||
RequestID: requestID,
|
||||
Delta: chunk.Delta,
|
||||
Role: "assistant",
|
||||
}); err != nil {
|
||||
log.Errorw("发送 llm_chunk 失败", "error", err)
|
||||
}
|
||||
|
||||
// 句子切分
|
||||
splitter.Feed(chunk.Delta)
|
||||
}
|
||||
|
||||
// 刷新切分器中的剩余文本
|
||||
splitter.Flush()
|
||||
|
||||
return fullText.String(), tokenUsage
|
||||
}
|
||||
|
||||
// synthesizeTTS 从句子 channel 读取文本，进行 TTS 合成并发送音频。
|
||||
func (p *Pipeline) synthesizeTTS(
|
||||
ctx context.Context,
|
||||
sentenceCh <-chan string,
|
||||
requestID string,
|
||||
sender Sender,
|
||||
) error {
|
||||
log := logger.Log
|
||||
|
||||
ttsStream, err := p.ttsService.SynthesizeStream(ctx, sentenceCh, tts.Options{
|
||||
Voice: p.ttsVoice,
|
||||
Speed: p.ttsSpeed,
|
||||
OutputFmt: p.ttsOutputFmt,
|
||||
SampleRate: p.ttsSampleRate,
|
||||
})
|
||||
if err != nil {
|
||||
log.Errorw("TTS 合成启动失败", "error", err)
|
||||
return err
|
||||
}
|
||||
|
||||
// 消费 TTS 音频流
|
||||
for chunk := range ttsStream {
|
||||
// 检查上下文是否已取消
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
log.Infow("TTS 流被中断", "request_id", requestID)
|
||||
return ctx.Err()
|
||||
default:
|
||||
}
|
||||
|
||||
// Base64 编码音频数据
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(chunk.Audio)
|
||||
|
||||
if err := sender.SendTTSAudio(models.WsTTSAudio{
|
||||
Type: "tts_audio",
|
||||
RequestID: requestID,
|
||||
Audio: audioBase64,
|
||||
MimeType: "audio/mp3",
|
||||
IsLast: chunk.IsLast,
|
||||
Final: chunk.Final,
|
||||
}); err != nil {
|
||||
log.Errorw("发送 tts_audio 失败", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
@@ -1,713 +0,0 @@
|
||||
package orchestrator
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/mock"
|
||||
|
||||
"github.com/hhs/camtalk/internal/ai/llm"
|
||||
"github.com/hhs/camtalk/internal/ai/stt"
|
||||
"github.com/hhs/camtalk/internal/ai/tts"
|
||||
"github.com/hhs/camtalk/internal/config"
|
||||
"github.com/hhs/camtalk/internal/logger"
|
||||
"github.com/hhs/camtalk/internal/models"
|
||||
"github.com/hhs/camtalk/internal/session"
|
||||
)
|
||||
|
||||
func init() {
|
||||
logger.Init("debug", "console")
|
||||
}
|
||||
|
||||
// MockSTTService mock STT 服务
|
||||
type MockSTTService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockSTTService) Recognize(ctx context.Context, audio []byte, opts stt.Options) (string, error) {
|
||||
args := m.Called(ctx, audio, opts)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
// MockLLMService mock LLM 服务
|
||||
type MockLLMService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockLLMService) ChatStream(ctx context.Context, req llm.Request) (<-chan llm.Chunk, error) {
|
||||
args := m.Called(ctx, req)
|
||||
if args.Get(0) == nil {
|
||||
return nil, args.Error(1)
|
||||
}
|
||||
return args.Get(0).(<-chan llm.Chunk), args.Error(1)
|
||||
}
|
||||
|
||||
// MockTTSService mock TTS 服务
|
||||
type MockTTSService struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockTTSService) SynthesizeStream(ctx context.Context, textStream <-chan string, opts tts.Options) (<-chan tts.Chunk, error) {
|
||||
args := m.Called(ctx, textStream, opts)
|
||||
if args.Get(0) == nil {
|
||||
return nil, args.Error(1)
|
||||
}
|
||||
return args.Get(0).(<-chan tts.Chunk), args.Error(1)
|
||||
}
|
||||
|
||||
// MockSessionManager mock 会话管理器
|
||||
type MockSessionManager struct {
|
||||
mock.Mock
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) Create(ctx context.Context, userID string, config models.SessionConfig) (string, error) {
|
||||
args := m.Called(ctx, userID, config)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) UpdateTitle(ctx context.Context, sessionID string, title string) error {
|
||||
args := m.Called(ctx, sessionID, title)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) ListByUser(ctx context.Context, userID string, page, size int) ([]session.ConversationSummary, int, error) {
|
||||
args := m.Called(ctx, userID, page, size)
|
||||
return args.Get(0).([]session.ConversationSummary), args.Int(1), args.Error(2)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) Get(ctx context.Context, sessionID string) (*models.Session, error) {
|
||||
args := m.Called(ctx, sessionID)
|
||||
if args.Get(0) == nil {
|
||||
return nil, args.Error(1)
|
||||
}
|
||||
return args.Get(0).(*models.Session), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) UpdateConfig(ctx context.Context, sessionID string, patch models.SessionConfigPatch) error {
|
||||
args := m.Called(ctx, sessionID, patch)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) GetHistory(ctx context.Context, sessionID string, limit int) ([]models.Message, error) {
|
||||
args := m.Called(ctx, sessionID, limit)
|
||||
return args.Get(0).([]models.Message), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) AppendMessage(ctx context.Context, sessionID string, msg models.Message) error {
|
||||
args := m.Called(ctx, sessionID, msg)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) SetActiveRequest(ctx context.Context, sessionID string, requestID string) error {
|
||||
args := m.Called(ctx, sessionID, requestID)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) GetActiveRequestID(ctx context.Context, sessionID string) (string, error) {
|
||||
args := m.Called(ctx, sessionID)
|
||||
return args.String(0), args.Error(1)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) ClearActiveRequest(ctx context.Context, sessionID string) error {
|
||||
args := m.Called(ctx, sessionID)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) Touch(ctx context.Context, sessionID string) error {
|
||||
args := m.Called(ctx, sessionID)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) Destroy(ctx context.Context, sessionID string) error {
|
||||
args := m.Called(ctx, sessionID)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSessionManager) ActiveCount() int {
|
||||
args := m.Called()
|
||||
return args.Int(0)
|
||||
}
|
||||
|
||||
// MockSender mock WebSocket 发送器
|
||||
type MockSender struct {
|
||||
mock.Mock
|
||||
STTResults []models.WsSTTResult
|
||||
LLMChunks []models.WsLLMChunk
|
||||
LLMDones []models.WsLLMDone
|
||||
TTSAudios []models.WsTTSAudio
|
||||
Errors []models.WsError
|
||||
}
|
||||
|
||||
func NewMockSender() *MockSender {
|
||||
return &MockSender{
|
||||
STTResults: make([]models.WsSTTResult, 0),
|
||||
LLMChunks: make([]models.WsLLMChunk, 0),
|
||||
LLMDones: make([]models.WsLLMDone, 0),
|
||||
TTSAudios: make([]models.WsTTSAudio, 0),
|
||||
Errors: make([]models.WsError, 0),
|
||||
}
|
||||
}
|
||||
|
||||
func (m *MockSender) SendSTTResult(result models.WsSTTResult) error {
|
||||
m.STTResults = append(m.STTResults, result)
|
||||
args := m.Called(result)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSender) SendLLMChunk(chunk models.WsLLMChunk) error {
|
||||
m.LLMChunks = append(m.LLMChunks, chunk)
|
||||
args := m.Called(chunk)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSender) SendLLMDone(done models.WsLLMDone) error {
|
||||
m.LLMDones = append(m.LLMDones, done)
|
||||
args := m.Called(done)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSender) SendTTSAudio(audio models.WsTTSAudio) error {
|
||||
m.TTSAudios = append(m.TTSAudios, audio)
|
||||
args := m.Called(audio)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
func (m *MockSender) SendError(err models.WsError) error {
|
||||
m.Errors = append(m.Errors, err)
|
||||
args := m.Called(err)
|
||||
return args.Error(0)
|
||||
}
|
||||
|
||||
// 辅助函数：创建 LLM 流式响应
|
||||
func createLLMStream(chunks []llm.Chunk) <-chan llm.Chunk {
|
||||
ch := make(chan llm.Chunk, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
ch <- chunk
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
|
||||
// 辅助函数：创建 TTS 流式响应
|
||||
func createTTSStream(chunks []tts.Chunk) <-chan tts.Chunk {
|
||||
ch := make(chan tts.Chunk, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
ch <- chunk
|
||||
}
|
||||
close(ch)
|
||||
return ch
|
||||
}
|
||||
|
||||
// TestProcessQuery_Success 测试完整流程
|
||||
func TestProcessQuery_Success(t *testing.T) {
|
||||
// 准备测试数据
|
||||
audioData := []byte("test audio")
|
||||
imageData := []byte("test image")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
imageBase64 := base64.StdEncoding.EncodeToString(imageData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Image: imageBase64,
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
session := &models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
TTSEnabled: true,
|
||||
DetailLevel: "low",
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}
|
||||
|
||||
// 创建 mock
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
// 设置 mock 期望
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(session, nil)
|
||||
mockSession.On("AppendMessage", mock.Anything, "session-123", mock.Anything).Return(nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, stt.Options{
|
||||
Encoding: "pcm_s16le",
|
||||
SampleRate: 16000,
|
||||
Language: "zh-CN",
|
||||
}).Return("你好，世界", nil)
|
||||
|
||||
mockSender.On("SendSTTResult", mock.Anything).Return(nil)
|
||||
|
||||
llmChunks := []llm.Chunk{
|
||||
{Delta: "你好"},
|
||||
{Delta: "，世界！"},
|
||||
{Done: true, TokensUsed: &llm.TokenUsage{Prompt: 10, Completion: 5, Total: 15}},
|
||||
}
|
||||
mockLLM.On("ChatStream", mock.Anything, mock.Anything).Return(createLLMStream(llmChunks), nil)
|
||||
|
||||
mockSender.On("SendLLMChunk", mock.Anything).Return(nil)
|
||||
mockSender.On("SendLLMDone", mock.Anything).Return(nil)
|
||||
|
||||
ttsChunks := []tts.Chunk{
|
||||
{Audio: []byte("audio1"), IsLast: false},
|
||||
{Audio: []byte("audio2"), IsLast: true},
|
||||
}
|
||||
mockTTS.On("SynthesizeStream", mock.Anything, mock.Anything, mock.Anything).Return(createTTSStream(ttsChunks), nil)
|
||||
|
||||
mockSender.On("SendTTSAudio", mock.Anything).Return(nil)
|
||||
|
||||
// 创建 Pipeline
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
// 执行
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
// 验证
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, mockSender.STTResults, 1)
|
||||
assert.Equal(t, "你好，世界", mockSender.STTResults[0].Text)
|
||||
assert.Len(t, mockSender.LLMChunks, 2)
|
||||
assert.Len(t, mockSender.LLMDones, 1)
|
||||
assert.Len(t, mockSender.TTSAudios, 2)
|
||||
|
||||
mockSTT.AssertExpectations(t)
|
||||
mockLLM.AssertExpectations(t)
|
||||
mockTTS.AssertExpectations(t)
|
||||
mockSession.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestProcessQuery_STTError 测试 STT 失败降级
|
||||
func TestProcessQuery_STTError(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(&models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}, nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, mock.Anything).
|
||||
Return("", errors.New("STT service unavailable"))
|
||||
|
||||
mockSender.On("SendError", mock.Anything).Return(nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Len(t, mockSender.Errors, 1)
|
||||
assert.Equal(t, "STT_ERROR", mockSender.Errors[0].Code)
|
||||
|
||||
mockSTT.AssertExpectations(t)
|
||||
mockLLM.AssertNotCalled(t, "ChatStream")
|
||||
mockTTS.AssertNotCalled(t, "SynthesizeStream")
|
||||
}
|
||||
|
||||
// TestProcessQuery_LLMError 测试 LLM 失败降级
|
||||
func TestProcessQuery_LLMError(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
session := &models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
TTSEnabled: true,
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(session, nil)
|
||||
mockSession.On("AppendMessage", mock.Anything, "session-123", mock.Anything).Return(nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, mock.Anything).Return("你好", nil)
|
||||
mockSender.On("SendSTTResult", mock.Anything).Return(nil)
|
||||
|
||||
mockLLM.On("ChatStream", mock.Anything, mock.Anything).
|
||||
Return(nil, errors.New("LLM service unavailable"))
|
||||
|
||||
mockSender.On("SendError", mock.Anything).Return(nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Len(t, mockSender.Errors, 1)
|
||||
assert.Equal(t, "LLM_ERROR", mockSender.Errors[0].Code)
|
||||
|
||||
mockSTT.AssertExpectations(t)
|
||||
mockLLM.AssertExpectations(t)
|
||||
mockTTS.AssertNotCalled(t, "SynthesizeStream")
|
||||
}
|
||||
|
||||
// TestProcessQuery_TTSError 测试 TTS 失败静默跳过
|
||||
func TestProcessQuery_TTSError(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
session := &models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
TTSEnabled: true,
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(session, nil)
|
||||
mockSession.On("AppendMessage", mock.Anything, "session-123", mock.Anything).Return(nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, mock.Anything).Return("你好", nil)
|
||||
mockSender.On("SendSTTResult", mock.Anything).Return(nil)
|
||||
|
||||
llmChunks := []llm.Chunk{
|
||||
{Delta: "你好"},
|
||||
{Done: true},
|
||||
}
|
||||
mockLLM.On("ChatStream", mock.Anything, mock.Anything).Return(createLLMStream(llmChunks), nil)
|
||||
mockSender.On("SendLLMChunk", mock.Anything).Return(nil)
|
||||
mockSender.On("SendLLMDone", mock.Anything).Return(nil)
|
||||
|
||||
mockTTS.On("SynthesizeStream", mock.Anything, mock.Anything, mock.Anything).
|
||||
Return(nil, errors.New("TTS service unavailable"))
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
// TTS 失败应该静默跳过，不返回错误
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, mockSender.LLMDones, 1)
|
||||
assert.Len(t, mockSender.TTSAudios, 0)
|
||||
|
||||
mockSTT.AssertExpectations(t)
|
||||
mockLLM.AssertExpectations(t)
|
||||
mockTTS.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestProcessQuery_ContextCancelled 测试上下文取消（Interrupt）
|
||||
func TestProcessQuery_ContextCancelled(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
session := &models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
TTSEnabled: true,
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(session, nil)
|
||||
mockSession.On("AppendMessage", mock.Anything, "session-123", mock.Anything).Return(nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, mock.Anything).Return("你好", nil)
|
||||
mockSender.On("SendSTTResult", mock.Anything).Return(nil)
|
||||
|
||||
// 创建一个会延迟的 LLM 流，以便我们可以取消上下文
|
||||
llmCh := make(chan llm.Chunk)
|
||||
go func() {
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
llmCh <- llm.Chunk{Delta: "你"}
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
llmCh <- llm.Chunk{Delta: "好"}
|
||||
close(llmCh)
|
||||
}()
|
||||
|
||||
mockLLM.On("ChatStream", mock.Anything, mock.Anything).Return((<-chan llm.Chunk)(llmCh), nil)
|
||||
mockSender.On("SendLLMChunk", mock.Anything).Return(nil)
|
||||
mockSender.On("SendLLMDone", mock.Anything).Return(nil)
|
||||
|
||||
// 创建一个会延迟的 TTS 流
|
||||
ttsCh := make(chan tts.Chunk)
|
||||
go func() {
|
||||
time.Sleep(200 * time.Millisecond)
|
||||
close(ttsCh)
|
||||
}()
|
||||
mockTTS.On("SynthesizeStream", mock.Anything, mock.Anything, mock.Anything).Return((<-chan tts.Chunk)(ttsCh), nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
// 创建可取消的上下文
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
|
||||
// 在 50ms 后取消
|
||||
go func() {
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
cancel()
|
||||
}()
|
||||
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
// 上下文取消后，流程应该正常完成（中断流但不返回错误）
|
||||
assert.NoError(t, err)
|
||||
|
||||
mockSTT.AssertExpectations(t)
|
||||
}
|
||||
|
||||
// TestProcessQuery_DisabledTTS 测试 TTS 未启用的情况
|
||||
func TestProcessQuery_DisabledTTS(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
session := &models.Session{
|
||||
ID: "session-123",
|
||||
Config: models.SessionConfig{
|
||||
TTSEnabled: false, // TTS 未启用
|
||||
Language: "zh-CN",
|
||||
},
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(session, nil)
|
||||
mockSession.On("AppendMessage", mock.Anything, "session-123", mock.Anything).Return(nil)
|
||||
|
||||
mockSTT.On("Recognize", mock.Anything, audioData, mock.Anything).Return("你好", nil)
|
||||
mockSender.On("SendSTTResult", mock.Anything).Return(nil)
|
||||
|
||||
llmChunks := []llm.Chunk{
|
||||
{Delta: "你好"},
|
||||
{Done: true},
|
||||
}
|
||||
mockLLM.On("ChatStream", mock.Anything, mock.Anything).Return(createLLMStream(llmChunks), nil)
|
||||
mockSender.On("SendLLMChunk", mock.Anything).Return(nil)
|
||||
mockSender.On("SendLLMDone", mock.Anything).Return(nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
assert.NoError(t, err)
|
||||
assert.Len(t, mockSender.LLMDones, 1)
|
||||
assert.Len(t, mockSender.TTSAudios, 0)
|
||||
|
||||
// TTS 不应该被调用
|
||||
mockTTS.AssertNotCalled(t, "SynthesizeStream")
|
||||
}
|
||||
|
||||
// TestSplitter 测试句子切分器
|
||||
func TestSplitter(t *testing.T) {
|
||||
ch := make(chan string, 10)
|
||||
splitter := NewSplitter(ch)
|
||||
|
||||
// 输入包含多个句子的文本
|
||||
splitter.Feed("你好。")
|
||||
splitter.Feed("世界！")
|
||||
splitter.Feed("这是")
|
||||
splitter.Feed("一个测试。")
|
||||
splitter.Flush()
|
||||
|
||||
// 应该有 3 个句子
|
||||
assert.Equal(t, 3, len(ch))
|
||||
assert.Equal(t, "你好。", <-ch)
|
||||
assert.Equal(t, "世界！", <-ch)
|
||||
assert.Equal(t, "这是一个测试。", <-ch)
|
||||
}
|
||||
|
||||
// TestSplitter_NoDelimiter 测试没有分隔符的情况
|
||||
func TestSplitter_NoDelimiter(t *testing.T) {
|
||||
ch := make(chan string, 10)
|
||||
splitter := NewSplitter(ch)
|
||||
|
||||
splitter.Feed("没有分隔符的文本")
|
||||
splitter.Flush()
|
||||
|
||||
// 应该有 1 个句子（Flush 会发送剩余内容）
|
||||
assert.Equal(t, 1, len(ch))
|
||||
assert.Equal(t, "没有分隔符的文本", <-ch)
|
||||
}
|
||||
|
||||
// TestSplitter_Empty 测试空输入
|
||||
func TestSplitter_Empty(t *testing.T) {
|
||||
ch := make(chan string, 10)
|
||||
splitter := NewSplitter(ch)
|
||||
|
||||
splitter.Flush()
|
||||
|
||||
// 应该没有句子
|
||||
assert.Equal(t, 0, len(ch))
|
||||
}
|
||||
|
||||
// TestProcessQuery_InvalidAudio 测试无效音频数据
|
||||
func TestProcessQuery_InvalidAudio(t *testing.T) {
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: "invalid-base64!!!",
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSender.On("SendError", mock.Anything).Return(nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Len(t, mockSender.Errors, 1)
|
||||
assert.Equal(t, "INVALID_MESSAGE", mockSender.Errors[0].Code)
|
||||
}
|
||||
|
||||
// TestProcessQuery_SessionNotFound 测试会话不存在
|
||||
func TestProcessQuery_SessionNotFound(t *testing.T) {
|
||||
audioData := []byte("test audio")
|
||||
audioBase64 := base64.StdEncoding.EncodeToString(audioData)
|
||||
|
||||
req := models.WsQuery{
|
||||
Type: "query",
|
||||
RequestID: "req-123",
|
||||
Audio: audioBase64,
|
||||
}
|
||||
|
||||
mockSTT := new(MockSTTService)
|
||||
mockLLM := new(MockLLMService)
|
||||
mockTTS := new(MockTTSService)
|
||||
mockSession := new(MockSessionManager)
|
||||
mockSender := NewMockSender()
|
||||
|
||||
mockSession.On("SetActiveRequest", mock.Anything, "session-123", "req-123").Return(nil)
|
||||
mockSession.On("ClearActiveRequest", mock.Anything, "session-123").Return(nil)
|
||||
mockSession.On("Get", mock.Anything, "session-123").Return(nil, errors.New("session not found"))
|
||||
|
||||
mockSender.On("SendError", mock.Anything).Return(nil)
|
||||
|
||||
pipeline := New(mockSTT, mockLLM, mockTTS, mockSession, &config.Config{
|
||||
AI: config.AIConfig{
|
||||
LLM: config.LLMConfig{Model: "gpt-4o"},
|
||||
TTS: config.TTSConfig{Voice: "alloy", Speed: 1.0, OutputFormat: "mp3", SampleRate: 24000},
|
||||
},
|
||||
})
|
||||
|
||||
ctx := context.Background()
|
||||
err := pipeline.ProcessQuery(ctx, "session-123", req, nil, mockSender)
|
||||
|
||||
assert.Error(t, err)
|
||||
assert.Len(t, mockSender.Errors, 1)
|
||||
assert.Equal(t, "SESSION_NOT_FOUND", mockSender.Errors[0].Code)
|
||||
}
|
||||
@@ -1,55 +0,0 @@
|
||||
package orchestrator
|
||||
|
||||
import "strings"
|
||||
|
||||
// sentenceDelimiters 句子分隔符集合。
|
||||
var sentenceDelimiters = map[rune]bool{
|
||||
'。': true,
|
||||
'！': true,
|
||||
'？': true,
|
||||
'\n': true,
|
||||
'.': true,
|
||||
'!': true,
|
||||
'?': true,
|
||||
}
|
||||
|
||||
// Splitter 句子切分器。
|
||||
// 将流式文本按句子边界切分，发送到 channel 供 TTS 合成。
|
||||
type Splitter struct {
|
||||
ch chan<- string
|
||||
buffer strings.Builder
|
||||
}
|
||||
|
||||
// NewSplitter 创建句子切分器。
|
||||
// ch 用于接收切分后的句子文本。
|
||||
func NewSplitter(ch chan<- string) *Splitter {
|
||||
return &Splitter{
|
||||
ch: ch,
|
||||
}
|
||||
}
|
||||
|
||||
// Feed 输入增量文本，遇到句子分隔符时发送完整句子。
|
||||
func (s *Splitter) Feed(delta string) {
|
||||
for _, r := range delta {
|
||||
s.buffer.WriteRune(r)
|
||||
if sentenceDelimiters[r] {
|
||||
s.flushBuffer()
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Flush 刷新缓冲区中的剩余文本（即使没有句子分隔符）。
|
||||
func (s *Splitter) Flush() {
|
||||
if s.buffer.Len() > 0 {
|
||||
s.flushBuffer()
|
||||
}
|
||||
}
|
||||
|
||||
// flushBuffer 将缓冲区内容发送到 channel 并清空。
|
||||
func (s *Splitter) flushBuffer() {
|
||||
text := strings.TrimSpace(s.buffer.String())
|
||||
if text != "" {
|
||||
s.ch <- text
|
||||
}
|
||||
s.buffer.Reset()
|
||||
}
|
||||
Reference in New Issue
Block a user