Skip to content

A
aisbf
Commit a303acc3 authored by Your Name's avatar Your Name
Browse files
Options

fix: Fix SDK async streaming - use create(stream=True) instead of stream() 

The Anthropic SDK's messages.stream() is a synchronous context manager,
not async. For async streaming, we need to use messages.create(..., stream=True)
which returns an async iterator of ServerSentEvent objects.

Changed from:
    async with client.messages.stream(**request_kwargs) as stream:
To:
    stream = await client.messages.create(**request_kwargs, stream=True)
    async for event in stream:
parent 0ba372d8
Hide whitespace changes
Inline Side-by-side
Showing with 139 additions and 138 deletions
aisbf/providers.py
View file @ a303acc3
...@@ -3932,10 +3932,10 @@ class ClaudeProviderHandler(BaseProviderHandler): ...@@ -3932,10 +3932,10 @@ class ClaudeProviderHandler(BaseProviderHandler):
async def _handle_streaming_request_sdk(self, client, request_kwargs: Dict, model: str): async def _handle_streaming_request_sdk(self, client, request_kwargs: Dict, model: str):
""" """
Handle streaming request using Anthropic SDK. Handle streaming request using Anthropic SDK's async streaming API.
The SDK handles proper streaming event parsing and retries. Uses client.messages.create(..., stream=True) which returns an async iterator
We convert SDK events to OpenAI-compatible SSE chunks. of ServerSentEvent objects. We parse these events and convert to OpenAI SSE chunks.
""" """
import logging import logging
import json import json
...@@ -3961,164 +3961,165 @@ class ClaudeProviderHandler(BaseProviderHandler): ...@@ -3961,164 +3961,165 @@ class ClaudeProviderHandler(BaseProviderHandler):
idle_timeout = self.stream_idle_timeout idle_timeout = self.stream_idle_timeout
try: try:
# Use SDK's streaming API with context manager # Use SDK's async streaming API - create(stream=True) returns async iterator
async with client.messages.stream(**request_kwargs) as stream: stream = await client.messages.create(**request_kwargs, stream=True)
async for event in stream:
# Update idle watchdog async for event in stream:
last_event_time = time.time() # Update idle watchdog
last_event_time = time.time()
# Check for idle timeout
if time.time() - last_event_time > idle_timeout: # Check for idle timeout
logger.error(f"ClaudeProviderHandler: Stream idle timeout ({idle_timeout}s)") if time.time() - last_event_time > idle_timeout:
raise TimeoutError(f"Stream idle for {idle_timeout}s") logger.error(f"ClaudeProviderHandler: Stream idle timeout ({idle_timeout}s)")
raise TimeoutError(f"Stream idle for {idle_timeout}s")
# Handle different event types from SDK
event_type = getattr(event, 'type', None) # Handle different event types from SDK
event_type = getattr(event, 'type', None)
if event_type == 'content_block_start':
content_block = getattr(event, 'content_block', None) if event_type == 'content_block_start':
if content_block: content_block = getattr(event, 'content_block', None)
block_type = getattr(content_block, 'type', '') if content_block:
block_type = getattr(content_block, 'type', '')
if block_type == 'tool_use':
tool_call = { if block_type == 'tool_use':
'index': content_block_index, tool_call = {
'id': getattr(content_block, 'id', ''), 'index': content_block_index,
'type': 'function', 'id': getattr(content_block, 'id', ''),
'function': { 'type': 'function',
'name': getattr(content_block, 'name', ''), 'function': {
'arguments': '' 'name': getattr(content_block, 'name', ''),
} 'arguments': ''
}
current_tool_calls.append(tool_call)
logger.debug(f"ClaudeProviderHandler: Tool use block started: {tool_call['function']['name']}")
elif block_type == 'thinking':
accumulated_thinking = ""
is_redacted_thinking = False
thinking_signature = ""
logger.debug(f"ClaudeProviderHandler: Thinking block started")
elif block_type == 'redacted_thinking':
accumulated_thinking = ""
is_redacted_thinking = True
thinking_signature = ""
logger.debug(f"ClaudeProviderHandler: Redacted thinking block started")
content_block_index += 1
elif event_type == 'content_block_delta':
delta = getattr(event, 'delta', None)
if delta:
delta_type = getattr(delta, 'type', '')
if delta_type == 'text_delta':
text = getattr(delta, 'text', '')
accumulated_content += text
openai_delta = {'content': text}
if first_chunk:
openai_delta['role'] = 'assistant'
first_chunk = False
openai_chunk = {
'id': completion_id,
'object': 'chat.completion.chunk',
'created': created_time,
'model': f'{self.provider_id}/{model}',
'choices': [{
'index': 0,
'delta': openai_delta,
'finish_reason': None
}]
} }
}
yield f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n".encode('utf-8') current_tool_calls.append(tool_call)
logger.debug(f"ClaudeProviderHandler: Tool use block started: {tool_call['function']['name']}")
elif delta_type == 'input_json_delta':
partial_json = getattr(delta, 'partial_json', '') elif block_type == 'thinking':
if current_tool_calls: accumulated_thinking = ""
current_tool_calls[-1]['function']['arguments'] += partial_json is_redacted_thinking = False
thinking_signature = ""
elif delta_type == 'thinking_delta': logger.debug(f"ClaudeProviderHandler: Thinking block started")
thinking_text = getattr(delta, 'thinking', '')
accumulated_thinking += thinking_text elif block_type == 'redacted_thinking':
logger.debug(f"ClaudeProviderHandler: Thinking delta: {len(thinking_text)} chars") accumulated_thinking = ""
is_redacted_thinking = True
thinking_signature = ""
logger.debug(f"ClaudeProviderHandler: Redacted thinking block started")
content_block_index += 1
elif event_type == 'content_block_delta':
delta = getattr(event, 'delta', None)
if delta:
delta_type = getattr(delta, 'type', '')
if delta_type == 'text_delta':
text = getattr(delta, 'text', '')
accumulated_content += text
elif delta_type == 'signature_delta': openai_delta = {'content': text}
signature = getattr(delta, 'signature', '') if first_chunk:
thinking_signature = signature openai_delta['role'] = 'assistant'
logger.debug(f"ClaudeProviderHandler: Thinking signature received") first_chunk = False
elif event_type == 'content_block_stop':
if current_tool_calls:
tool_call = current_tool_calls[-1]
try:
args = json.loads(tool_call['function']['arguments']) if tool_call['function']['arguments'] else {}
tool_call['function']['arguments'] = json.dumps(args)
except json.JSONDecodeError:
logger.warning(f"ClaudeProviderHandler: Invalid tool call arguments JSON")
tool_call['function']['arguments'] = '{}'
tool_call_chunk = { openai_chunk = {
'id': completion_id, 'id': completion_id,
'object': 'chat.completion.chunk', 'object': 'chat.completion.chunk',
'created': created_time, 'created': created_time,
'model': f'{self.provider_id}/{model}', 'model': f'{self.provider_id}/{model}',
'choices': [{ 'choices': [{
'index': 0, 'index': 0,
'delta': { 'delta': openai_delta,
'tool_calls': [{
'index': tool_call['index'],
'id': tool_call['id'],
'type': tool_call['type'],
'function': tool_call['function']
}]
},
'finish_reason': None 'finish_reason': None
}] }]
} }
yield f"data: {json.dumps(tool_call_chunk, ensure_ascii=False)}\n\n".encode('utf-8') yield f"data: {json.dumps(openai_chunk, ensure_ascii=False)}\n\n".encode('utf-8')
logger.debug(f"ClaudeProviderHandler: Emitted tool call: {tool_call['function']['name']}")
elif accumulated_thinking: elif delta_type == 'input_json_delta':
block_type = "redacted_thinking" if is_redacted_thinking else "thinking" partial_json = getattr(delta, 'partial_json', '')
logger.info(f"ClaudeProviderHandler: {block_type} block completed ({len(accumulated_thinking)} chars)") if current_tool_calls:
accumulated_thinking = "" current_tool_calls[-1]['function']['arguments'] += partial_json
is_redacted_thinking = False
thinking_signature = "" elif delta_type == 'thinking_delta':
thinking_text = getattr(delta, 'thinking', '')
elif event_type == 'message_delta': accumulated_thinking += thinking_text
usage = getattr(event, 'usage', None) logger.debug(f"ClaudeProviderHandler: Thinking delta: {len(thinking_text)} chars")
if usage:
logger.debug(f"ClaudeProviderHandler: Streaming usage update: {usage}") elif delta_type == 'signature_delta':
signature = getattr(delta, 'signature', '')
# Track cache tokens for analytics thinking_signature = signature
cache_read = getattr(usage, 'cache_read_input_tokens', 0) logger.debug(f"ClaudeProviderHandler: Thinking signature received")
cache_creation = getattr(usage, 'cache_creation_input_tokens', 0)
if cache_read > 0: elif event_type == 'content_block_stop':
self.cache_stats['cache_hits'] += 1 if current_tool_calls:
self.cache_stats['cache_tokens_read'] += cache_read tool_call = current_tool_calls[-1]
if cache_creation > 0: try:
self.cache_stats['cache_misses'] += 1 args = json.loads(tool_call['function']['arguments']) if tool_call['function']['arguments'] else {}
self.cache_stats['cache_tokens_created'] += cache_creation tool_call['function']['arguments'] = json.dumps(args)
except json.JSONDecodeError:
elif event_type == 'message_stop': logger.warning(f"ClaudeProviderHandler: Invalid tool call arguments JSON")
final_chunk = { tool_call['function']['arguments'] = '{}'
tool_call_chunk = {
'id': completion_id, 'id': completion_id,
'object': 'chat.completion.chunk', 'object': 'chat.completion.chunk',
'created': created_time, 'created': created_time,
'model': f'{self.provider_id}/{model}', 'model': f'{self.provider_id}/{model}',
'choices': [{ 'choices': [{
'index': 0, 'index': 0,
'delta': {}, 'delta': {
'finish_reason': 'stop' 'tool_calls': [{
'index': tool_call['index'],
'id': tool_call['id'],
'type': tool_call['type'],
'function': tool_call['function']
}]
},
'finish_reason': None
}] }]
} }
yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n".encode('utf-8') yield f"data: {json.dumps(tool_call_chunk, ensure_ascii=False)}\n\n".encode('utf-8')
yield b"data: [DONE]\n\n" logger.debug(f"ClaudeProviderHandler: Emitted tool call: {tool_call['function']['name']}")
elif accumulated_thinking:
block_type = "redacted_thinking" if is_redacted_thinking else "thinking"
logger.info(f"ClaudeProviderHandler: {block_type} block completed ({len(accumulated_thinking)} chars)")
accumulated_thinking = ""
is_redacted_thinking = False
thinking_signature = ""
elif event_type == 'message_delta':
usage = getattr(event, 'usage', None)
if usage:
logger.debug(f"ClaudeProviderHandler: Streaming usage update: {usage}")
# Track cache tokens for analytics
cache_read = getattr(usage, 'cache_read_input_tokens', 0)
cache_creation = getattr(usage, 'cache_creation_input_tokens', 0)
if cache_read > 0:
self.cache_stats['cache_hits'] += 1
self.cache_stats['cache_tokens_read'] += cache_read
if cache_creation > 0:
self.cache_stats['cache_misses'] += 1
self.cache_stats['cache_tokens_created'] += cache_creation
elif event_type == 'message_stop':
final_chunk = {
'id': completion_id,
'object': 'chat.completion.chunk',
'created': created_time,
'model': f'{self.provider_id}/{model}',
'choices': [{
'index': 0,
'delta': {},
'finish_reason': 'stop'
}]
}
yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n".encode('utf-8')
yield b"data: [DONE]\n\n"
logger.info(f"ClaudeProviderHandler: SDK streaming completed successfully") logger.info(f"ClaudeProviderHandler: SDK streaming completed successfully")
self.record_success() self.record_success()
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment