diff --git a/agent/context/interfaces.go b/agent/context/interfaces.go index ac7c1a36..ab6318cb 100644 --- a/agent/context/interfaces.go +++ b/agent/context/interfaces.go @@ -7,6 +7,7 @@ type StreamChunkType string // Stream chunk type constants - indicates what type of content is in the current chunk const ( + // Content chunk types - actual data from the LLM ChunkText StreamChunkType = "text" // Regular text content ChunkThinking StreamChunkType = "thinking" // Reasoning/thinking content (o1, DeepSeek R1) ChunkToolCall StreamChunkType = "tool_call" // Tool/function call @@ -14,6 +15,12 @@ const ( ChunkMetadata StreamChunkType = "metadata" // Metadata (usage, finish_reason, etc.) ChunkError StreamChunkType = "error" // Error chunk ChunkUnknown StreamChunkType = "unknown" // Unknown/unrecognized chunk type + + // Lifecycle event types - stream and group boundaries + ChunkStreamStart StreamChunkType = "stream_start" // Stream begins (entire request starts) + ChunkStreamEnd StreamChunkType = "stream_end" // Stream ends (entire request completes) + ChunkGroupStart StreamChunkType = "group_start" // Message group begins (text/tool_call/thinking group starts) + ChunkGroupEnd StreamChunkType = "group_end" // Message group ends (text/tool_call/thinking group completes) ) // StreamFunc the streaming function callback diff --git a/agent/context/types_llm.go b/agent/context/types_llm.go index c78dc69b..94a37847 100644 --- a/agent/context/types_llm.go +++ b/agent/context/types_llm.go @@ -162,3 +162,62 @@ type JSONSchema struct { Schema interface{} `json:"schema"` // Required: JSON schema (*jsonschema.Schema or map[string]interface{}) Strict *bool `json:"strict,omitempty"` // Optional: whether to enforce strict schema validation (default: true) } + +// ============================================================================ +// Stream Lifecycle Event Data Structures +// ============================================================================ +// These structures define the data format for stream lifecycle events. +// They provide a standardized way to communicate stream boundaries and metadata +// to the frontend, enabling better UI/UX (progress indicators, timing, etc.). + +// StreamStartData represents the data for stream_start event +// Sent when a streaming request begins +type StreamStartData struct { + RequestID string `json:"request_id"` // Unique identifier for this request + Timestamp int64 `json:"timestamp"` // Unix timestamp when stream started + Model string `json:"model,omitempty"` // Model being used (e.g., "gpt-4o") + Capabilities map[string]interface{} `json:"capabilities,omitempty"` // Model capabilities for this request +} + +// StreamEndData represents the data for stream_end event +// Sent when a streaming request completes (successfully or with error) +type StreamEndData struct { + RequestID string `json:"request_id"` // Corresponding request ID + Timestamp int64 `json:"timestamp"` // Unix timestamp when stream ended + DurationMs int64 `json:"duration_ms"` // Total duration in milliseconds + Status string `json:"status"` // "completed" | "error" | "cancelled" + Error string `json:"error,omitempty"` // Error message if status is "error" + Usage *UsageInfo `json:"usage,omitempty"` // Token usage statistics +} + +// GroupStartData represents the data for group_start event +// Sent when a logical message group begins (text, tool_call, thinking, etc.) +type GroupStartData struct { + GroupID string `json:"group_id"` // Unique identifier for this group + Type string `json:"type"` // Group type: "text" | "thinking" | "tool_call" | "refusal" + Timestamp int64 `json:"timestamp"` // Unix timestamp when group started + ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Tool call metadata (if type is "tool_call") + Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata (for custom providers or future extensions) +} + +// GroupEndData represents the data for group_end event +// Sent when a logical message group completes +type GroupEndData struct { + GroupID string `json:"group_id"` // Corresponding group ID + Type string `json:"type"` // Group type (same as in group_start) + Timestamp int64 `json:"timestamp"` // Unix timestamp when group ended + DurationMs int64 `json:"duration_ms"` // Duration of this group in milliseconds + ChunkCount int `json:"chunk_count"` // Number of data chunks in this group + Status string `json:"status"` // "completed" | "partial" | "error" + ToolCall *GroupToolCallInfo `json:"tool_call,omitempty"` // Complete tool call info (if type is "tool_call") + Extra map[string]interface{} `json:"extra,omitempty"` // Additional metadata +} + +// GroupToolCallInfo contains tool call information for group events +// Used in both group_start (partial info) and group_end (complete info) +type GroupToolCallInfo struct { + ID string `json:"id"` // Tool call ID (e.g., "call_abc123") + Name string `json:"name"` // Function name (may be partial in group_start) + Arguments string `json:"arguments,omitempty"` // Complete arguments (only in group_end) + Index int `json:"index"` // Index in the tool calls array +} diff --git a/agent/llm/providers/openai/openai.go b/agent/llm/providers/openai/openai.go index acfc4bfb..62d14b52 100644 --- a/agent/llm/providers/openai/openai.go +++ b/agent/llm/providers/openai/openai.go @@ -15,6 +15,93 @@ import ( "github.com/yaoapp/yao/utils/jsonschema" ) +// startGroup starts a new group and sends group_start event +func (gt *groupTracker) startGroup(groupType context.StreamChunkType, handler context.StreamFunc) { + if gt.active { + // End previous group first + gt.endGroup(handler) + } + + gt.active = true + gt.groupID = fmt.Sprintf("grp_%d", time.Now().UnixNano()) + gt.groupType = groupType + gt.startTime = time.Now().UnixMilli() + gt.chunkCount = 0 + gt.toolCallInfo = nil + + if handler != nil { + startData := &context.GroupStartData{ + GroupID: gt.groupID, + Type: string(groupType), + Timestamp: gt.startTime, + } + if startJSON, err := jsoniter.Marshal(startData); err == nil { + handler(context.ChunkGroupStart, startJSON) + } + } +} + +// startToolCallGroup starts a new tool call group with tool call info +func (gt *groupTracker) startToolCallGroup(toolCallInfo *context.GroupToolCallInfo, handler context.StreamFunc) { + if gt.active { + gt.endGroup(handler) + } + + gt.active = true + gt.groupID = fmt.Sprintf("grp_tool_%d", time.Now().UnixNano()) + gt.groupType = context.ChunkToolCall + gt.startTime = time.Now().UnixMilli() + gt.chunkCount = 0 + gt.toolCallInfo = toolCallInfo + + if handler != nil { + startData := &context.GroupStartData{ + GroupID: gt.groupID, + Type: string(context.ChunkToolCall), + Timestamp: gt.startTime, + ToolCall: toolCallInfo, + } + if startJSON, err := jsoniter.Marshal(startData); err == nil { + handler(context.ChunkGroupStart, startJSON) + } + } +} + +// incrementChunk increments the chunk count for the current group +func (gt *groupTracker) incrementChunk() { + if gt.active { + gt.chunkCount++ + } +} + +// endGroup ends the current group and sends group_end event +func (gt *groupTracker) endGroup(handler context.StreamFunc) { + if !gt.active { + return + } + + if handler != nil { + endData := &context.GroupEndData{ + GroupID: gt.groupID, + Type: string(gt.groupType), + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Now().UnixMilli() - gt.startTime, + ChunkCount: gt.chunkCount, + Status: "completed", + } + if gt.toolCallInfo != nil { + endData.ToolCall = gt.toolCallInfo + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkGroupEnd, endJSON) + } + } + + gt.active = false + gt.groupID = "" + gt.toolCallInfo = nil +} + // Provider OpenAI-compatible provider // Supports: vision, tool calls, streaming, JSON mode type Provider struct { @@ -98,9 +185,38 @@ func (p *Provider) Stream(ctx *context.Context, messages []context.Message, opti // streamWithRetry performs a single streaming request attempt func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Message, options *context.CompletionOptions, handler context.StreamFunc) (*context.CompletionResponse, error) { + streamStartTime := time.Now() + requestID := fmt.Sprintf("req_%d", streamStartTime.UnixNano()) + + // Send stream_start event + if handler != nil { + model, _ := p.GetModel() + startData := &context.StreamStartData{ + RequestID: requestID, + Timestamp: streamStartTime.UnixMilli(), + Model: model, + } + if startJSON, err := jsoniter.Marshal(startData); err == nil { + handler(context.ChunkStreamStart, startJSON) + } + } + // Build request body requestBody, err := p.buildRequestBody(messages, options, true) if err != nil { + // Send stream_end with error + if handler != nil { + endData := &context.StreamEndData{ + RequestID: requestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(streamStartTime).Milliseconds(), + Status: "error", + Error: err.Error(), + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkStreamEnd, endJSON) + } + } return nil, fmt.Errorf("failed to build request body: %w", err) } @@ -135,6 +251,9 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess toolCalls: make(map[int]*accumulatedToolCall), } + // Group tracker for lifecycle events + groupTracker := &groupTracker{} + // Stream handler streamHandler := func(data []byte) int { if len(data) == 0 { @@ -181,17 +300,29 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // Handle content if delta.Content != "" { + // Start text group if not active + if !groupTracker.active || groupTracker.groupType != context.ChunkText { + groupTracker.startGroup(context.ChunkText, handler) + } + accumulator.content += delta.Content if handler != nil { handler(context.ChunkText, []byte(delta.Content)) + groupTracker.incrementChunk() } } // Handle refusal if delta.Refusal != "" { + // Start refusal group if not active + if !groupTracker.active || groupTracker.groupType != context.ChunkRefusal { + groupTracker.startGroup(context.ChunkRefusal, handler) + } + accumulator.refusal += delta.Refusal if handler != nil { handler(context.ChunkRefusal, []byte(delta.Refusal)) + groupTracker.incrementChunk() } } @@ -200,6 +331,16 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess for _, tc := range delta.ToolCalls { if _, exists := accumulator.toolCalls[tc.Index]; !exists { accumulator.toolCalls[tc.Index] = &accumulatedToolCall{} + + // Start new tool call group when we first see this tool call + if tc.ID != "" { + toolCallInfo := &context.GroupToolCallInfo{ + ID: tc.ID, + Name: tc.Function.Name, // May be partial or empty initially + Index: tc.Index, + } + groupTracker.startToolCallGroup(toolCallInfo, handler) + } } accTC := accumulator.toolCalls[tc.Index] @@ -211,9 +352,17 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess } if tc.Function.Name != "" { accTC.functionName = tc.Function.Name + // Update tool call info in tracker + if groupTracker.active && groupTracker.toolCallInfo != nil { + groupTracker.toolCallInfo.Name = tc.Function.Name + } } if tc.Function.Arguments != "" { accTC.functionArgs += tc.Function.Arguments + // Update tool call info in tracker + if groupTracker.active && groupTracker.toolCallInfo != nil { + groupTracker.toolCallInfo.Arguments = accTC.functionArgs + } } } @@ -221,6 +370,7 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if handler != nil { toolCallData, _ := jsoniter.Marshal(delta.ToolCalls) handler(context.ChunkToolCall, toolCallData) + groupTracker.incrementChunk() } } @@ -259,10 +409,25 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess err = req.Stream(goCtx, "POST", requestBody, streamHandler) if err != nil { + // End current group if active + groupTracker.endGroup(handler) + // Notify handler of error if provided if handler != nil { errData := []byte(err.Error()) handler(context.ChunkError, errData) + + // Send stream_end with error + endData := &context.StreamEndData{ + RequestID: requestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(streamStartTime).Milliseconds(), + Status: "error", + Error: err.Error(), + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkStreamEnd, endJSON) + } } return nil, fmt.Errorf("streaming request failed: %w", err) } @@ -271,10 +436,26 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess if accumulator.id == "" { log.Warn("OpenAI stream completed but no data was received (accumulator.id is empty)") err := fmt.Errorf("no data received from OpenAI API") + + // End current group if active + groupTracker.endGroup(handler) + // Notify handler of error if provided if handler != nil { errData := []byte(err.Error()) handler(context.ChunkError, errData) + + // Send stream_end with error + endData := &context.StreamEndData{ + RequestID: requestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(streamStartTime).Milliseconds(), + Status: "error", + Error: err.Error(), + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkStreamEnd, endJSON) + } } return nil, err } @@ -311,11 +492,45 @@ func (p *Provider) streamWithRetry(ctx *context.Context, messages []context.Mess // Validate tool call results if schema is provided if err := p.validateToolCallResults(options, toolCalls); err != nil { + // End current group + groupTracker.endGroup(handler) + + // Send stream_end with validation error + if handler != nil { + endData := &context.StreamEndData{ + RequestID: requestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(streamStartTime).Milliseconds(), + Status: "error", + Error: "tool call validation failed", + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkStreamEnd, endJSON) + } + } + // Tool call validation failed, need to retry with error feedback return nil, fmt.Errorf("tool call validation failed: %w", err) } } + // End final group if still active + groupTracker.endGroup(handler) + + // Send stream_end event (success) + if handler != nil { + endData := &context.StreamEndData{ + RequestID: requestID, + Timestamp: time.Now().UnixMilli(), + DurationMs: time.Since(streamStartTime).Milliseconds(), + Status: "completed", + Usage: response.Usage, + } + if endJSON, err := jsoniter.Marshal(endData); err == nil { + handler(context.ChunkStreamEnd, endJSON) + } + } + return response, nil } diff --git a/agent/llm/providers/openai/openai_test.go b/agent/llm/providers/openai/openai_test.go index 955788a0..a3ac56cf 100644 --- a/agent/llm/providers/openai/openai_test.go +++ b/agent/llm/providers/openai/openai_test.go @@ -1249,6 +1249,145 @@ func TestOpenAIProxySupport(t *testing.T) { t.Log("HTTP proxy support is implemented via http.GetTransport using environment variables") } +// TestOpenAIStreamLifecycleEvents tests that lifecycle events are sent correctly +func TestOpenAIStreamLifecycleEvents(t *testing.T) { + test.Prepare(t, config.Conf) + defer test.Clean() + + conn, err := connector.Select("openai.gpt-4o") + if err != nil { + t.Fatalf("Failed to select connector: %v", err) + } + + trueVal := true + options := &context.CompletionOptions{ + Capabilities: &context.ModelCapabilities{ + Streaming: &trueVal, + ToolCalls: &trueVal, + }, + } + + llmInstance, err := llm.New(conn, options) + if err != nil { + t.Fatalf("Failed to create LLM instance: %v", err) + } + + messages := []context.Message{ + { + Role: context.RoleUser, + Content: "Say 'hello' in one word", + }, + } + + ctx := newTestContext("test-lifecycle", "openai.gpt-4o") + + // Track lifecycle events + var events []string + var streamStartReceived, streamEndReceived bool + var groupStartReceived, groupEndReceived bool + + handler := func(chunkType context.StreamChunkType, data []byte) int { + events = append(events, string(chunkType)) + + switch chunkType { + case context.ChunkStreamStart: + streamStartReceived = true + var startData context.StreamStartData + if err := json.Unmarshal(data, &startData); err == nil { + t.Logf("✓ stream_start: request_id=%s, model=%s", startData.RequestID, startData.Model) + if startData.RequestID == "" { + t.Error("stream_start missing request_id") + } + } else { + t.Errorf("Failed to parse stream_start data: %v", err) + } + + case context.ChunkStreamEnd: + streamEndReceived = true + var endData context.StreamEndData + if err := json.Unmarshal(data, &endData); err == nil { + t.Logf("✓ stream_end: status=%s, duration=%dms", endData.Status, endData.DurationMs) + if endData.Status != "completed" { + t.Errorf("stream_end status should be 'completed', got '%s'", endData.Status) + } + if endData.DurationMs <= 0 { + t.Error("stream_end duration should be > 0") + } + } else { + t.Errorf("Failed to parse stream_end data: %v", err) + } + + case context.ChunkGroupStart: + groupStartReceived = true + var startData context.GroupStartData + if err := json.Unmarshal(data, &startData); err == nil { + t.Logf("✓ group_start: type=%s, group_id=%s", startData.Type, startData.GroupID) + if startData.GroupID == "" { + t.Error("group_start missing group_id") + } + } else { + t.Errorf("Failed to parse group_start data: %v", err) + } + + case context.ChunkGroupEnd: + groupEndReceived = true + var endData context.GroupEndData + if err := json.Unmarshal(data, &endData); err == nil { + t.Logf("✓ group_end: type=%s, chunks=%d, duration=%dms", + endData.Type, endData.ChunkCount, endData.DurationMs) + if endData.ChunkCount <= 0 { + t.Error("group_end should have chunk_count > 0") + } + } else { + t.Errorf("Failed to parse group_end data: %v", err) + } + + case context.ChunkText: + t.Logf(" text chunk: %s", string(data)) + } + + return 0 + } + + response, err := llmInstance.Stream(ctx, messages, options, handler) + if err != nil { + t.Fatalf("Stream failed: %v", err) + } + + if response == nil { + t.Fatal("Response is nil") + } + + // Validate that all lifecycle events were received + if !streamStartReceived { + t.Error("stream_start event was not received") + } + if !streamEndReceived { + t.Error("stream_end event was not received") + } + if !groupStartReceived { + t.Error("group_start event was not received") + } + if !groupEndReceived { + t.Error("group_end event was not received") + } + + // Validate event order: stream_start should be first, stream_end should be last + if len(events) < 4 { + t.Errorf("Expected at least 4 events, got %d", len(events)) + } else { + if events[0] != "stream_start" { + t.Errorf("First event should be stream_start, got %s", events[0]) + } + if events[len(events)-1] != "stream_end" { + t.Errorf("Last event should be stream_end, got %s", events[len(events)-1]) + } + } + + t.Logf("Total events received: %d", len(events)) + t.Log("Lifecycle events test completed successfully") +} + // TestOpenAIStreamWithTemperature tests different temperature settings func TestOpenAIStreamWithTemperature(t *testing.T) { test.Prepare(t, config.Conf) diff --git a/agent/llm/providers/openai/types.go b/agent/llm/providers/openai/types.go index dbcc78eb..0cd475a9 100644 --- a/agent/llm/providers/openai/types.go +++ b/agent/llm/providers/openai/types.go @@ -80,3 +80,13 @@ type accumulatedToolCall struct { functionName string functionArgs string } + +// groupTracker tracks the current group state for lifecycle events +type groupTracker struct { + active bool // Whether a group is currently active + groupID string // Current group ID + groupType context.StreamChunkType // Current group type + startTime int64 // Group start timestamp + chunkCount int // Number of chunks in this group + toolCallInfo *context.GroupToolCallInfo // Tool call info if group is tool_call type +}