Skip to content

A
aisbf
Commit 42e06a0c authored by Stefy Lanza (nextime / spora )'s avatar Stefy Lanza (nextime / spora )
Browse files
Options

Streaming support: anthropic + ollama, and true incremental streaming for google 

- anthropic: previously ignored the stream flag entirely (always non-stream).
  Add _handle_streaming_request converting Anthropic SSE events into
  OpenAI-compatible chunks (text deltas, tool_use/input_json deltas, usage).
- ollama: previously hard-coded stream=False. Add incremental streaming over
  /api/generate yielding OpenAI chunks.
- google: was buffering the entire response before yielding. Stream chunks
  incrementally as they arrive when no tools are requested (the common case);
  keep the buffer-and-parse path only when tools are present, since the
  text-encoded tool-call detection needs the complete response.
Co-Authored-By: 's avatarClaude Opus 4.8 <noreply@anthropic.com>
parent afef7809
Hide whitespace changes
Inline Side-by-side
Showing with 552 additions and 236 deletions
aisbf/handlers.py
View file @ 42e06a0c
......@@ -4409,202 +4409,239 @@ class RotationHandler:
completion_tokens = 0
accumulated_response_text = "" # Track full response for token counting
# Collect all chunks first to know when we're at the last one
chunks_list = []
async for chunk in response:
chunks_list.append(chunk)
total_chunks = len(chunks_list)
chunk_idx = 0
for chunk in chunks_list:
try:
logger.debug(f"Google chunk type: {type(chunk)}")
logger.debug(f"Google chunk: {chunk}")
# Extract text and tool calls from Google chunk (this is accumulated text)
# Whether the client requested tools. Text-encoded tool-call
# detection (see the buffered path below) needs the COMPLETE
# response, so we only stream incrementally when no tools were
# requested -- the dominant chat case.
request_has_tools = bool(request_data.get('tools'))
if not request_has_tools:
# TRUE INCREMENTAL streaming: forward each text delta to
# the client as soon as Google yields it.
sent_role = False
async for chunk in response:
chunk_text = ""
chunk_tool_calls = []
finish_reason = None
chunk_finish = None
try:
if hasattr(chunk, 'candidates') and chunk.candidates:
candidate = chunk.candidates[0] if chunk.candidates else None
if candidate and hasattr(candidate, 'content') and candidate.content:
if hasattr(candidate.content, 'parts') and candidate.content.parts:
for part in candidate.content.parts:
# Extract text content
if hasattr(part, 'text') and part.text:
chunk_text += part.text
# Extract function calls (Google's format)
if hasattr(part, 'function_call') and part.function_call:
function_call = part.function_call
# Convert Google function call to OpenAI format
import json
openai_tool_call = {
"id": f"call_{len(chunk_tool_calls)}",
"type": "function",
"function": {
"name": function_call.name,
"arguments": json.dumps(function_call.args) if hasattr(function_call, 'args') else "{}"
}
}
chunk_tool_calls.append(openai_tool_call)
logger.info(f"Extracted tool call from Google chunk: {openai_tool_call}")
# Check for finish reason in candidate
if hasattr(candidate, 'finish_reason'):
google_finish = str(candidate.finish_reason)
if google_finish in ('STOP', 'END_TURN', 'FINISH_REASON_UNSPECIFIED'):
finish_reason = "stop"
elif google_finish == 'MAX_TOKENS':
finish_reason = "length"
if hasattr(candidate, 'finish_reason'):
google_finish = str(candidate.finish_reason)
if google_finish in ('STOP', 'END_TURN', 'FINISH_REASON_UNSPECIFIED'):
chunk_finish = "stop"
elif google_finish == 'MAX_TOKENS':
chunk_finish = "length"
except Exception as e:
logger.error(f"Error extracting text from Google chunk: {e}")
# Calculate the delta (only the new text since last chunk)
delta_text = chunk_text[len(accumulated_text):] if chunk_text.startswith(accumulated_text) else chunk_text
accumulated_text = chunk_text # Update accumulated text for next iteration
# Track completion tokens for Google responses
logger.error(f"Error extracting text from Google stream chunk: {e}")
continue
# Google chunk text is the incremental delta; guard
# against providers that send accumulated text.
delta_text = chunk_text[len(accumulated_text):] if chunk_text.startswith(accumulated_text) and accumulated_text else chunk_text
if delta_text:
accumulated_text += delta_text
accumulated_response_text += delta_text
chunk_idx += 1
except Exception as chunk_error:
error_msg = str(chunk_error)
logger.error(f"Error processing Google chunk: {error_msg}")
logger.error(f"Chunk type: {type(chunk)}")
logger.error(f"Chunk content: {chunk}")
chunk_idx += 1
continue
# After collecting all chunks, check if the accumulated text contains a tool call pattern
# This handles models that return tool calls as text instead of using function_call attributes
tool_calls = None
final_text = accumulated_response_text
# Check for tool call patterns in the accumulated text
if accumulated_response_text:
import re as re_module
# Simple approach: just look for "tool: {...}" pattern and extract the JSON
# This avoids complex nested parsing issues
tool_pattern = r'tool:\s*(\{[^{}]*\{[^{}]*\}[^{}]*\}|\{[^{}]+\})'
tool_match = re_module.search(tool_pattern, accumulated_response_text, re_module.DOTALL)
if tool_match:
try:
# Extract the tool JSON using brace counting for robustness
tool_start = accumulated_response_text.find('tool:')
if tool_start != -1:
json_start = accumulated_response_text.find('{', tool_start)
if json_start != -1:
brace_count = 0
json_end = json_start
for i, c in enumerate(accumulated_response_text[json_start:], json_start):
if c == '{':
brace_count += 1
elif c == '}':
brace_count -= 1
if brace_count == 0:
json_end = i + 1
break
tool_json_str = accumulated_response_text[json_start:json_end]
logger.debug(f"Extracted tool JSON: {tool_json_str[:200]}...")
try:
parsed_tool = json.loads(tool_json_str)
except json.JSONDecodeError:
# Try fixing common issues: single quotes, trailing commas
fixed_json = tool_json_str.replace("'", '"')
fixed_json = re_module.sub(r',\s*}', '}', fixed_json)
fixed_json = re_module.sub(r',\s*]', ']', fixed_json)
parsed_tool = json.loads(fixed_json)
# Convert to OpenAI tool_calls format
tool_calls = [{
"id": f"call_0",
"type": "function",
"function": {
"name": parsed_tool.get('action', parsed_tool.get('name', 'unknown')),
"arguments": json.dumps({k: v for k, v in parsed_tool.items() if k not in ['action', 'name']})
}
}]
logger.info(f"Converted streaming tool call to OpenAI format: {tool_calls}")
# Extract final assistant text after the tool JSON
# Look for pattern: }\\nassistant: [{'type': 'text', 'text': "..."}]
# or just return empty since the tool call is the main content
after_tool = accumulated_response_text[json_end:]
assistant_pattern = r"assistant:\s*\[.*'text':\s*['\"](.+?)['\"].*\]\s*\]?\s*$"
assistant_match = re_module.search(assistant_pattern, after_tool, re_module.DOTALL)
if assistant_match:
final_text = assistant_match.group(1)
# Unescape common escape sequences
final_text = final_text.replace("\\n", "\n").replace("\\'", "'").replace('\\"', '"')
else:
final_text = ""
except (json.JSONDecodeError, ValueError, SyntaxError, Exception) as e:
logger.debug(f"Failed to parse tool JSON in streaming: {e}")
# Now send the response chunks
# If we detected tool calls, send them in the first chunk with role
if tool_calls:
# First chunk with tool_calls
tool_chunk = {
delta = {"content": delta_text, "refusal": None, "role": None, "tool_calls": None}
if not sent_role:
delta["role"] = "assistant"
sent_role = True
out_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": None,
"provider": provider_id,
"choices": [{
"index": 0,
"delta": delta,
"finish_reason": None,
"logprobs": None,
"native_finish_reason": None
}]
}
yield f"data: {json.dumps(out_chunk)}\n\n".encode('utf-8')
await asyncio.sleep(0)
# Final chunk with finish reason and usage statistics.
if accumulated_response_text:
completion_tokens = count_messages_tokens([{"role": "assistant", "content": accumulated_response_text}], model_name)
total_tokens = effective_context + completion_tokens
final_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": None,
"usage": {
"prompt_tokens": effective_context,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
"effective_context": effective_context
},
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": None,
"refusal": None,
"role": "assistant",
"tool_calls": tool_calls
},
"finish_reason": None,
"delta": {"content": "", "function_call": None, "refusal": None, "role": None, "tool_calls": None},
"finish_reason": "stop",
"logprobs": None,
"native_finish_reason": None
"native_finish_reason": "stop"
}]
}
yield f"data: {json.dumps(tool_chunk)}\n\n".encode('utf-8')
# If there's final assistant text, send it
if final_text:
text_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": None,
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": final_text,
"refusal": None,
"role": None,
"tool_calls": None
},
"finish_reason": None,
"logprobs": None,
"native_finish_reason": None
}]
}
yield f"data: {json.dumps(text_chunk)}\n\n".encode('utf-8')
yield f"data: {json.dumps(final_chunk)}\n\n".encode('utf-8')
await asyncio.sleep(0)
else:
# No tool calls detected, send text normally
# Send the accumulated text as a single chunk
# Collect all chunks first to know when we're at the last one
chunks_list = []
async for chunk in response:
chunks_list.append(chunk)
total_chunks = len(chunks_list)
chunk_idx = 0
for chunk in chunks_list:
try:
logger.debug(f"Google chunk type: {type(chunk)}")
logger.debug(f"Google chunk: {chunk}")
# Extract text and tool calls from Google chunk (this is accumulated text)
chunk_text = ""
chunk_tool_calls = []
finish_reason = None
try:
if hasattr(chunk, 'candidates') and chunk.candidates:
candidate = chunk.candidates[0] if chunk.candidates else None
if candidate and hasattr(candidate, 'content') and candidate.content:
if hasattr(candidate.content, 'parts') and candidate.content.parts:
for part in candidate.content.parts:
# Extract text content
if hasattr(part, 'text') and part.text:
chunk_text += part.text
# Extract function calls (Google's format)
if hasattr(part, 'function_call') and part.function_call:
function_call = part.function_call
# Convert Google function call to OpenAI format
import json
openai_tool_call = {
"id": f"call_{len(chunk_tool_calls)}",
"type": "function",
"function": {
"name": function_call.name,
"arguments": json.dumps(function_call.args) if hasattr(function_call, 'args') else "{}"
}
}
chunk_tool_calls.append(openai_tool_call)
logger.info(f"Extracted tool call from Google chunk: {openai_tool_call}")
# Check for finish reason in candidate
if hasattr(candidate, 'finish_reason'):
google_finish = str(candidate.finish_reason)
if google_finish in ('STOP', 'END_TURN', 'FINISH_REASON_UNSPECIFIED'):
finish_reason = "stop"
elif google_finish == 'MAX_TOKENS':
finish_reason = "length"
except Exception as e:
logger.error(f"Error extracting text from Google chunk: {e}")
# Calculate the delta (only the new text since last chunk)
delta_text = chunk_text[len(accumulated_text):] if chunk_text.startswith(accumulated_text) else chunk_text
accumulated_text = chunk_text # Update accumulated text for next iteration
# Track completion tokens for Google responses
if delta_text:
accumulated_response_text += delta_text
chunk_idx += 1
except Exception as chunk_error:
error_msg = str(chunk_error)
logger.error(f"Error processing Google chunk: {error_msg}")
logger.error(f"Chunk type: {type(chunk)}")
logger.error(f"Chunk content: {chunk}")
chunk_idx += 1
continue
# After collecting all chunks, check if the accumulated text contains a tool call pattern
# This handles models that return tool calls as text instead of using function_call attributes
tool_calls = None
final_text = accumulated_response_text
# Check for tool call patterns in the accumulated text
if accumulated_response_text:
text_chunk = {
import re as re_module
# Simple approach: just look for "tool: {...}" pattern and extract the JSON
# This avoids complex nested parsing issues
tool_pattern = r'tool:\s*(\{[^{}]*\{[^{}]*\}[^{}]*\}|\{[^{}]+\})'
tool_match = re_module.search(tool_pattern, accumulated_response_text, re_module.DOTALL)
if tool_match:
try:
# Extract the tool JSON using brace counting for robustness
tool_start = accumulated_response_text.find('tool:')
if tool_start != -1:
json_start = accumulated_response_text.find('{', tool_start)
if json_start != -1:
brace_count = 0
json_end = json_start
for i, c in enumerate(accumulated_response_text[json_start:], json_start):
if c == '{':
brace_count += 1
elif c == '}':
brace_count -= 1
if brace_count == 0:
json_end = i + 1
break
tool_json_str = accumulated_response_text[json_start:json_end]
logger.debug(f"Extracted tool JSON: {tool_json_str[:200]}...")
try:
parsed_tool = json.loads(tool_json_str)
except json.JSONDecodeError:
# Try fixing common issues: single quotes, trailing commas
fixed_json = tool_json_str.replace("'", '"')
fixed_json = re_module.sub(r',\s*}', '}', fixed_json)
fixed_json = re_module.sub(r',\s*]', ']', fixed_json)
parsed_tool = json.loads(fixed_json)
# Convert to OpenAI tool_calls format
tool_calls = [{
"id": f"call_0",
"type": "function",
"function": {
"name": parsed_tool.get('action', parsed_tool.get('name', 'unknown')),
"arguments": json.dumps({k: v for k, v in parsed_tool.items() if k not in ['action', 'name']})
}
}]
logger.info(f"Converted streaming tool call to OpenAI format: {tool_calls}")
# Extract final assistant text after the tool JSON
# Look for pattern: }\\nassistant: [{'type': 'text', 'text': "..."}]
# or just return empty since the tool call is the main content
after_tool = accumulated_response_text[json_end:]
assistant_pattern = r"assistant:\s*\[.*'text':\s*['\"](.+?)['\"].*\]\s*\]?\s*$"
assistant_match = re_module.search(assistant_pattern, after_tool, re_module.DOTALL)
if assistant_match:
final_text = assistant_match.group(1)
# Unescape common escape sequences
final_text = final_text.replace("\\n", "\n").replace("\\'", "'").replace('\\"', '"')
else:
final_text = ""
except (json.JSONDecodeError, ValueError, SyntaxError, Exception) as e:
logger.debug(f"Failed to parse tool JSON in streaming: {e}")
# Now send the response chunks
# If we detected tool calls, send them in the first chunk with role
if tool_calls:
# First chunk with tool_calls
tool_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
......@@ -4616,53 +4653,106 @@ class RotationHandler:
"choices": [{
"index": 0,
"delta": {
"content": accumulated_response_text,
"content": None,
"refusal": None,
"role": "assistant",
"tool_calls": None
"tool_calls": tool_calls
},
"finish_reason": None,
"logprobs": None,
"native_finish_reason": None
}]
}
yield f"data: {json.dumps(text_chunk)}\n\n".encode('utf-8')
yield f"data: {json.dumps(tool_chunk)}\n\n".encode('utf-8')
# If there's final assistant text, send it
if final_text:
text_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": None,
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": final_text,
"refusal": None,
"role": None,
"tool_calls": None
},
"finish_reason": None,
"logprobs": None,
"native_finish_reason": None
}]
}
yield f"data: {json.dumps(text_chunk)}\n\n".encode('utf-8')
else:
# No tool calls detected, send text normally
# Send the accumulated text as a single chunk
if accumulated_response_text:
text_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": None,
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": accumulated_response_text,
"refusal": None,
"role": "assistant",
"tool_calls": None
},
"finish_reason": None,
"logprobs": None,
"native_finish_reason": None
}]
}
yield f"data: {json.dumps(text_chunk)}\n\n".encode('utf-8')
# Send final chunk with finish reason and usage statistics
if accumulated_response_text:
completion_tokens = count_messages_tokens([{"role": "assistant", "content": accumulated_response_text}], model_name)
total_tokens = effective_context + completion_tokens
final_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": {
"prompt_tokens": effective_context,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
"effective_context": effective_context
},
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": "",
"function_call": None,
"refusal": None,
"role": None,
"tool_calls": None
# Send final chunk with finish reason and usage statistics
if accumulated_response_text:
completion_tokens = count_messages_tokens([{"role": "assistant", "content": accumulated_response_text}], model_name)
total_tokens = effective_context + completion_tokens
final_chunk = {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": model_name,
"service_tier": None,
"system_fingerprint": system_fingerprint,
"usage": {
"prompt_tokens": effective_context,
"completion_tokens": completion_tokens,
"total_tokens": total_tokens,
"effective_context": effective_context
},
"finish_reason": "stop",
"logprobs": None,
"native_finish_reason": "stop"
}]
}
yield f"data: {json.dumps(final_chunk)}\n\n".encode('utf-8')
# Yield control to event loop to ensure final chunk is flushed to client
await asyncio.sleep(0)
"provider": provider_id,
"choices": [{
"index": 0,
"delta": {
"content": "",
"function_call": None,
"refusal": None,
"role": None,
"tool_calls": None
},
"finish_reason": "stop",
"logprobs": None,
"native_finish_reason": "stop"
}]
}
yield f"data: {json.dumps(final_chunk)}\n\n".encode('utf-8')
# Yield control to event loop to ensure final chunk is flushed to client
await asyncio.sleep(0)
elif is_kilo_provider or is_coderai_provider:
# Handle Kilo/KiloCode streaming response
# Kilo/CoderAI return async generators that yield OpenAI-compatible SSE bytes
......
aisbf/providers/anthropic.py
View file @ 42e06a0c
......@@ -316,6 +316,14 @@ class AnthropicProviderHandler(BaseProviderHandler):
logging.info(f"Full payload: {_json.dumps(debug_params, indent=2, default=str)}")
logging.info(f"=== END ANTHROPIC API REQUEST PAYLOAD ===")
# Streaming mode: return an async generator that converts Anthropic
# SSE events into OpenAI-compatible chat.completion.chunk dicts. The
# caller records success after the stream is primed/consumed, so we
# do NOT record success here (the upstream call has not run yet).
if stream:
logging.info(f"AnthropicProviderHandler: Using streaming mode")
return self._handle_streaming_request(api_params, model)
response = self.client.messages.create(**api_params)
logging.info(f"AnthropicProviderHandler: Response received: {response}")
self.record_success()
......@@ -458,6 +466,123 @@ class AnthropicProviderHandler(BaseProviderHandler):
self.record_failure()
raise e
async def _handle_streaming_request(self, api_params: Dict, model: str):
"""Async generator that streams an Anthropic request and yields
OpenAI-compatible chat.completion.chunk dicts.
The Anthropic SDK client is synchronous, so the underlying event stream
is iterated with a regular ``for`` loop (matching how the OpenAI provider
forwards its sync SDK stream). Records failure on error and success once
the stream finishes.
"""
logging.info(f"AnthropicProviderHandler: Starting streaming request for model {model}")
response_id = f"anthropic-{model}-{int(time.time())}"
created_time = int(time.time())
stop_reason_map = {
'end_turn': 'stop',
'max_tokens': 'length',
'stop_sequence': 'stop',
'tool_use': 'tool_calls',
}
def _base_chunk():
return {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": f"{self.provider_id}/{model}",
"choices": [{"index": 0, "delta": {}, "finish_reason": None}],
}
try:
stream_params = dict(api_params)
stream_params["stream"] = True
prompt_tokens = 0
completion_tokens = 0
tool_block_index = -1 # OpenAI tool_calls index, incremented per tool_use block
finish_reason = "stop"
# First chunk announces the assistant role (OpenAI convention).
first = _base_chunk()
first["choices"][0]["delta"] = {"role": "assistant"}
yield first
stream = self.client.messages.create(**stream_params)
for event in stream:
etype = getattr(event, "type", None)
if etype == "message_start":
usage = getattr(getattr(event, "message", None), "usage", None)
if usage is not None:
prompt_tokens = getattr(usage, "input_tokens", 0) or 0
elif etype == "content_block_start":
block = getattr(event, "content_block", None)
if block is not None and getattr(block, "type", None) == "tool_use":
tool_block_index += 1
chunk = _base_chunk()
chunk["choices"][0]["delta"] = {
"tool_calls": [{
"index": tool_block_index,
"id": getattr(block, "id", "") or f"call_{tool_block_index}",
"type": "function",
"function": {"name": getattr(block, "name", "") or "", "arguments": ""},
}]
}
yield chunk
elif etype == "content_block_delta":
delta = getattr(event, "delta", None)
dtype = getattr(delta, "type", None)
if dtype == "text_delta":
text = getattr(delta, "text", "") or ""
if text:
chunk = _base_chunk()
chunk["choices"][0]["delta"] = {"content": text}
yield chunk
elif dtype == "input_json_delta":
partial = getattr(delta, "partial_json", "") or ""
if partial and tool_block_index >= 0:
chunk = _base_chunk()
chunk["choices"][0]["delta"] = {
"tool_calls": [{
"index": tool_block_index,
"function": {"arguments": partial},
}]
}
yield chunk
elif etype == "message_delta":
delta = getattr(event, "delta", None)
sr = getattr(delta, "stop_reason", None)
if sr:
finish_reason = stop_reason_map.get(sr, "stop")
usage = getattr(event, "usage", None)
if usage is not None:
completion_tokens = getattr(usage, "output_tokens", 0) or completion_tokens
# Final chunk with finish_reason and usage.
final = _base_chunk()
final["choices"][0]["finish_reason"] = finish_reason
final["usage"] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
yield final
yield b"data: [DONE]\n\n"
self.record_success()
logging.info(f"AnthropicProviderHandler: Streaming completed successfully")
except Exception as e:
logging.error(f"AnthropicProviderHandler: Streaming error: {str(e)}", exc_info=True)
self.record_failure()
raise
async def get_models(self) -> List[Model]:
"""
Return list of available Anthropic models.
......
aisbf/providers/google.py
View file @ 42e06a0c
......@@ -173,36 +173,44 @@ class GoogleProviderHandler(BaseProviderHandler):
from google import genai
stream_client = genai.Client(api_key=self.api_key)
chunks = []
for chunk in stream_client.models.generate_content_stream(
model=model,
contents=content,
config=config_params
):
chunks.append(chunk)
logging.info(f"GoogleProviderHandler: Streaming response received (total chunks: {len(chunks)})")
self.record_success()
# After successful streaming response, create cached content if pending
if hasattr(self, '_pending_cache_key') and self._pending_cache_key:
cache_key, cache_ttl, cache_messages = self._pending_cache_key
# Return an async generator that forwards each chunk to the caller
# AS IT ARRIVES (true incremental streaming) rather than buffering
# the whole response first. Success/cache-creation happen after the
# upstream stream is fully consumed. record_success is intentionally
# NOT called eagerly here: the caller records success after priming/
# consuming the stream (the upstream request has not run yet).
provider_self = self
async def async_generator():
try:
new_cached_name = self._create_cached_content(cache_messages, model, cache_ttl)
if new_cached_name:
expiry_time = time.time() + cache_ttl
self._cached_content_refs[cache_key] = (new_cached_name, expiry_time)
logging.info(f"GoogleProviderHandler: Cached content stored (streaming): {new_cached_name}, expires in {cache_ttl}s")
for chunk in stream_client.models.generate_content_stream(
model=model,
contents=content,
config=config_params
):
yield chunk
logging.info(f"GoogleProviderHandler: Streaming response fully received")
provider_self.record_success()
# After a successful stream, create cached content if pending.
if hasattr(provider_self, '_pending_cache_key') and provider_self._pending_cache_key:
cache_key, cache_ttl, cache_messages = provider_self._pending_cache_key
try:
new_cached_name = provider_self._create_cached_content(cache_messages, model, cache_ttl)
if new_cached_name:
expiry_time = time.time() + cache_ttl
provider_self._cached_content_refs[cache_key] = (new_cached_name, expiry_time)
logging.info(f"GoogleProviderHandler: Cached content stored (streaming): {new_cached_name}, expires in {cache_ttl}s")
except Exception as e:
logging.warning(f"GoogleProviderHandler: Failed to create cache after streaming: {e}")
provider_self._pending_cache_key = None
except Exception as e:
logging.warning(f"GoogleProviderHandler: Failed to create cache after streaming: {e}")
self._pending_cache_key = None
async def async_generator():
for chunk in chunks:
yield chunk
logging.error(f"GoogleProviderHandler: Streaming error: {str(e)}", exc_info=True)
provider_self.record_failure()
raise
return async_generator()
else:
# Non-streaming request
......
aisbf/providers/ollama.py
View file @ 42e06a0c
......@@ -111,7 +111,14 @@ class OllamaProviderHandler(BaseProviderHandler):
options = {"temperature": temperature}
if max_tokens is not None:
options["num_predict"] = max_tokens
# Streaming mode: return an async generator yielding OpenAI-compatible
# chunks. The caller records success after the stream is consumed, so
# we don't record success here.
if stream:
logger.info("OllamaProviderHandler: Using streaming mode")
return self._handle_streaming_request(model, prompt, options)
request_data = {
"model": model,
"prompt": prompt,
......@@ -213,6 +220,92 @@ class OllamaProviderHandler(BaseProviderHandler):
self.record_failure()
raise e
async def _handle_streaming_request(self, model: str, prompt: str, options: dict):
"""Async generator that streams an Ollama /api/generate request and
yields OpenAI-compatible chat.completion.chunk dicts."""
logger = logging.getLogger(__name__)
logger.info(f"OllamaProviderHandler: Starting streaming request for model {model}")
response_id = f"ollama-{model}-{int(time.time())}"
created_time = int(time.time())
headers = {}
if self.api_key:
headers["Authorization"] = f"Bearer {self.api_key}"
request_data = {
"model": model,
"prompt": prompt,
"options": options,
"stream": True,
}
def _base_chunk():
return {
"id": response_id,
"object": "chat.completion.chunk",
"created": created_time,
"model": f"{self.provider_id}/{model}",
"choices": [{"index": 0, "delta": {}, "finish_reason": None}],
}
try:
prompt_tokens = 0
completion_tokens = 0
first = _base_chunk()
first["choices"][0]["delta"] = {"role": "assistant"}
yield first
async with self.client.stream("POST", "/api/generate", json=request_data, headers=headers) as response:
if response.status_code == 429:
body = await response.aread()
try:
response_data = json.loads(body)
except Exception:
response_data = body.decode("utf-8", errors="replace")
self.handle_429_error(response_data, dict(response.headers))
response.raise_for_status()
response.raise_for_status()
async for line in response.aiter_lines():
if not line.strip():
continue
try:
obj = json.loads(line)
except json.JSONDecodeError:
logger.warning(f"OllamaProviderHandler: Skipping unparseable stream line: {line[:200]}")
continue
text = obj.get("response", "")
if text:
chunk = _base_chunk()
chunk["choices"][0]["delta"] = {"content": text}
yield chunk
if obj.get("done"):
prompt_tokens = obj.get("prompt_eval_count", 0) or 0
completion_tokens = obj.get("eval_count", 0) or 0
final = _base_chunk()
final["choices"][0]["finish_reason"] = "stop"
final["usage"] = {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
yield final
yield b"data: [DONE]\n\n"
self.record_success()
logger.info("OllamaProviderHandler: Streaming completed successfully")
except Exception as e:
logger.error(f"OllamaProviderHandler: Streaming error: {str(e)}", exc_info=True)
self.record_failure()
raise
async def get_models(self) -> List[Model]:
await self.apply_rate_limit()
......
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