diff --git a/agent/llm/handlers/stream.go b/agent/llm/handlers/stream.go index 27d0931f..593e3c14 100644 --- a/agent/llm/handlers/stream.go +++ b/agent/llm/handlers/stream.go @@ -62,10 +62,11 @@ func DefaultStreamHandler(ctx *context.Context) context.StreamFunc { // streamState manages the state of the streaming process type streamState struct { - ctx *context.Context - inGroup bool - currentID string - buffer []byte + ctx *context.Context + inGroup bool + currentID string + currentType string // Track the current message type (text, thinking, tool_call) + buffer []byte } // handleStreamStart handles stream start event @@ -96,6 +97,9 @@ func (s *streamState) handleText(data []byte) int { s.currentID = generateMessageID() } + // Track current message type + s.currentType = message.TypeText + // Append to buffer s.buffer = append(s.buffer, data...) @@ -128,6 +132,9 @@ func (s *streamState) handleThinking(data []byte) int { s.currentID = generateMessageID() } + // Track current message type + s.currentType = message.TypeThinking + // Append to buffer s.buffer = append(s.buffer, data...) @@ -190,9 +197,15 @@ func (s *streamState) handleGroupEnd(data []byte) int { // Send done message with complete content if s.currentID != "" && len(s.buffer) > 0 { + // Use the tracked message type (thinking, text, tool_call, etc.) + msgType := s.currentType + if msgType == "" { + msgType = message.TypeText // Fallback to text if type not set + } + msg := &message.Message{ ID: s.currentID, - Type: message.TypeText, // Default to text + Type: msgType, // Use the actual message type from the group Done: true, Props: map[string]interface{}{ "content": string(s.buffer), @@ -204,6 +217,7 @@ func (s *streamState) handleGroupEnd(data []byte) int { // Reset state s.inGroup = false s.currentID = "" + s.currentType = "" s.buffer = []byte{} return 0 // Continue diff --git a/agent/llm/providers/openai/deepseek_r1_test.go b/agent/llm/providers/openai/deepseek_r1_test.go index da5817da..93b12d2c 100644 --- a/agent/llm/providers/openai/deepseek_r1_test.go +++ b/agent/llm/providers/openai/deepseek_r1_test.go @@ -5,6 +5,7 @@ import ( "strings" "testing" + jsoniter "github.com/json-iterator/go" "github.com/yaoapp/gou/connector" "github.com/yaoapp/gou/plan" "github.com/yaoapp/yao/agent/context" @@ -59,20 +60,58 @@ func TestDeepSeekR1StreamBasic(t *testing.T) { // Create context ctx := newDeepSeekTestContext("test-deepseek-r1-basic", "deepseek.r1") - // Track streaming chunks + // Track streaming chunks and group events var reasoningChunks []string var contentChunks []string + var thinkingGroupEnded bool + var textGroupEnded bool + handler := func(chunkType context.StreamChunkType, data []byte) int { dataStr := string(data) t.Logf("Stream chunk [%s]: %s", chunkType, dataStr) // Track different chunk types - if chunkType == context.ChunkThinking { + switch chunkType { + case context.ChunkThinking: reasoningChunks = append(reasoningChunks, dataStr) - } else if chunkType == context.ChunkText { + case context.ChunkText: contentChunks = append(contentChunks, dataStr) } + // Track group_end events to verify type field + if chunkType == context.ChunkGroupEnd { + // Parse the group_end data to check the type field + var groupEndData struct { + GroupID string `json:"group_id"` + Type string `json:"type"` + Timestamp int64 `json:"timestamp"` + DurationMs int64 `json:"duration_ms"` + ChunkCount int `json:"chunk_count"` + Status string `json:"status"` + } + + if err := jsoniter.Unmarshal(data, &groupEndData); err == nil { + t.Logf("✓ group_end received: type=%s, chunks=%d, duration=%dms", + groupEndData.Type, groupEndData.ChunkCount, groupEndData.DurationMs) + + // Verify the type field matches expected group types + switch groupEndData.Type { + case "thinking": + thinkingGroupEnded = true + if groupEndData.ChunkCount == 0 { + t.Error("thinking group_end should have chunk_count > 0") + } + case "text": + textGroupEnded = true + if groupEndData.ChunkCount == 0 { + t.Error("text group_end should have chunk_count > 0") + } + } + } else { + t.Errorf("Failed to parse group_end data: %v", err) + } + } + return 0 // Continue } @@ -139,6 +178,19 @@ func TestDeepSeekR1StreamBasic(t *testing.T) { t.Logf("Received %d content chunks", len(contentChunks)) } + // Verify group_end events were received with correct types + if !thinkingGroupEnded { + t.Error("❌ Expected thinking group_end event but didn't receive it") + } else { + t.Log("✅ Thinking group_end event received with type='thinking'") + } + + if !textGroupEnded { + t.Error("❌ Expected text group_end event but didn't receive it") + } else { + t.Log("✅ Text group_end event received with type='text'") + } + t.Logf("Final response: %+v", response) }