66"""
77
88import logging
9- from typing import Any
9+
10+ from claude_agent_sdk import (
11+ AssistantMessage ,
12+ SystemMessage ,
13+ TextBlock ,
14+ ThinkingBlock ,
15+ ToolUseBlock ,
16+ )
17+ from claude_agent_sdk .types import StreamEvent
1018
1119from providers .base import AIStreamParser , ParsedStreamMessage
1220
1624class ClaudeStreamParser (AIStreamParser ):
1725 """Parser for Claude SDK streaming messages.
1826
19- This wraps the existing parsing logic from sdk/stream_parser.py
20- and adapts it to the unified ParsedStreamMessage format.
21-
2227 SDK Message Types:
28+ - StreamEvent: Partial message updates with raw Anthropic API events
2329 - AssistantMessage: content=[TextBlock, ThinkingBlock, ToolUseBlock, ...]
2430 - SystemMessage: subtype='sessionStarted', data={'session_id': ...}
2531 - ResultMessage: Final result with session_id, duration_ms, is_error
2632 """
2733
2834 @staticmethod
2935 def parse_message (
30- message : Any ,
36+ message : AssistantMessage | SystemMessage | StreamEvent | object ,
3137 current_response : str ,
3238 current_thinking : str ,
3339 ) -> ParsedStreamMessage :
3440 """Parse a streaming message from Claude SDK.
3541
3642 Args:
37- message: SDK message object (AssistantMessage, SystemMessage, etc.)
43+ message: SDK message object (StreamEvent, AssistantMessage, SystemMessage, etc.)
3844 current_response: Accumulated response text so far
3945 current_thinking: Accumulated thinking text so far
4046
@@ -43,70 +49,77 @@ def parse_message(
4349 """
4450 content_delta = ""
4551 thinking_delta = ""
46- new_session_id = None
47- skip_tool_called = False
48- memory_entries = []
49- anthropic_calls = []
52+ new_session_id : str | None = None
53+ memory_entries : list [str ] = []
54+
55+ # Handle StreamEvent for partial message updates (real-time streaming)
56+ if isinstance (message , StreamEvent ):
57+ event = message .event
58+ if isinstance (event , dict ):
59+ event_type = event .get ("type" , "" )
60+
61+ # Handle content_block_delta events (text streaming)
62+ if event_type == "content_block_delta" :
63+ delta = event .get ("delta" , {})
64+ delta_type = delta .get ("type" , "" )
65+
66+ if delta_type == "text_delta" :
67+ content_delta = delta .get ("text" , "" )
68+ elif delta_type == "thinking_delta" :
69+ thinking_delta = delta .get ("thinking" , "" )
70+
71+ # Extract session_id from StreamEvent
72+ if message .session_id :
73+ new_session_id = message .session_id
74+
75+ return ParsedStreamMessage (
76+ response_text = current_response + content_delta ,
77+ thinking_text = current_thinking + thinking_delta ,
78+ session_id = new_session_id ,
79+ skip_used = False ,
80+ memory_entries = memory_entries ,
81+ anthropic_calls = [],
82+ )
5083
5184 # Extract session_id from SystemMessage
52- if hasattr (message , "__class__" ) and message .__class__ .__name__ == "SystemMessage" :
53- if hasattr (message , "data" ) and isinstance (message .data , dict ):
54- if "session_id" in message .data :
55- new_session_id = message .data ["session_id" ]
56- logger .debug (f"Extracted session_id: { new_session_id } )
57-
58- # Handle content
59- if hasattr (message , "text" ):
60- content_delta = message .text
61- elif hasattr (message , "content" ):
62- if isinstance (message .content , str ):
63- content_delta = message .content
64- elif isinstance (message .content , list ):
65- for block in message .content :
66- block_type = getattr (block , "type" , None ) or (
67- block .get ("type" ) if isinstance (block , dict ) else None
68- )
69-
70- # Check for tool calls
71- # Note: MCP tools (skip, anthropic) are detected via PostToolUse hooks
72- # Only memorize is detected here as fallback for memory_entries
73- if block_type == "tool_use" :
74- tool_name = getattr (block , "name" , None ) or (
75- block .get ("name" ) if isinstance (block , dict ) else None
76- )
77-
78- if tool_name and tool_name .endswith ("__memorize" ):
79- tool_input = getattr (block , "input" , None ) or (
80- block .get ("input" ) if isinstance (block , dict ) else None
81- )
82- if tool_input and isinstance (tool_input , dict ):
83- memory_entry = tool_input .get ("memory_entry" , "" )
84- if memory_entry :
85- memory_entries .append (memory_entry )
86- logger .info (f"Agent recorded memory: { memory_entry } )
87-
88- # Handle thinking blocks
89- block_class_name = block .__class__ .__name__ if hasattr (block , "__class__" ) else ""
90- if block_class_name == "ThinkingBlock" or (hasattr (block , "type" ) and block .type == "thinking" ):
91- if hasattr (block , "thinking" ):
92- thinking_delta = block .thinking
93- elif hasattr (block , "text" ):
94- thinking_delta = block .text
95- elif isinstance (block , dict ) and block .get ("type" ) == "thinking" :
96- thinking_delta = block .get ("thinking" , block .get ("text" , "" ))
97- else :
98- # Handle text content blocks
99- if hasattr (block , "text" ):
100- content_delta += block .text
101- elif isinstance (block , dict ) and block .get ("type" ) == "text" :
102- content_delta += block .get ("text" , "" )
85+ if isinstance (message , SystemMessage ):
86+ if isinstance (message .data , dict ) and "session_id" in message .data :
87+ new_session_id = message .data ["session_id" ]
88+ logger .debug (f"Extracted session_id: { new_session_id } )
89+
90+ # Handle AssistantMessage content
91+ if isinstance (message , AssistantMessage ):
92+ # Track if we've already streamed content via StreamEvent
93+ # If so, skip adding text to avoid duplication
94+ skip_content = bool (current_response )
95+ skip_thinking = bool (current_thinking )
96+
97+ for block in message .content :
98+ # Check for memorize tool calls
99+ if isinstance (block , ToolUseBlock ):
100+ if block .name .endswith ("__memorize" ):
101+ if isinstance (block .input , dict ):
102+ memory_entry = block .input .get ("memory_entry" , "" )
103+ if memory_entry :
104+ memory_entries .append (memory_entry )
105+ logger .info (f"Agent recorded memory: { memory_entry } )
106+
107+ # Handle thinking blocks (skip if already streamed)
108+ elif isinstance (block , ThinkingBlock ):
109+ if not skip_thinking :
110+ thinking_delta = block .thinking
111+
112+ # Handle text blocks (skip if already streamed)
113+ elif isinstance (block , TextBlock ):
114+ if not skip_content :
115+ content_delta += block .text
103116
104117 # Return accumulated text with deltas applied
105118 return ParsedStreamMessage (
106119 response_text = current_response + content_delta ,
107120 thinking_text = current_thinking + thinking_delta ,
108121 session_id = new_session_id ,
109- skip_used = skip_tool_called ,
122+ skip_used = False ,
110123 memory_entries = memory_entries ,
111- anthropic_calls = anthropic_calls ,
124+ anthropic_calls = [] ,
112125 )
0 commit comments