Skip to content

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

Fix Google streaming handler to properly extract and send tool calls 

- Added extraction of function_call from Google chunk parts
- Convert Google function_call to OpenAI-compatible tool_calls format
- Track accumulated_tool_calls to calculate deltas
- Include tool_calls in the delta when present
- Send chunks when there are new tool calls, not just text
parent d12485b6
Hide whitespace changes
Inline Side-by-side
Showing with 45 additions and 5 deletions
aisbf/handlers.py
View file @ 0b4afcf6
...@@ -411,6 +411,7 @@ class RequestHandler: ...@@ -411,6 +411,7 @@ class RequestHandler:
# Note: Google returns accumulated text, so we need to track and send only deltas # Note: Google returns accumulated text, so we need to track and send only deltas
chunk_id = 0 chunk_id = 0
accumulated_text = "" # Track text we've already sent accumulated_text = "" # Track text we've already sent
accumulated_tool_calls = [] # Track tool calls we've already sent
last_chunk_id = None # Track the last chunk for finish_reason last_chunk_id = None # Track the last chunk for finish_reason
created_time = int(time.time()) created_time = int(time.time())
response_id = f"google-{request_data['model']}-{created_time}" response_id = f"google-{request_data['model']}-{created_time}"
...@@ -432,8 +433,9 @@ class RequestHandler: ...@@ -432,8 +433,9 @@ class RequestHandler:
logger.debug(f"Google chunk type: {type(chunk)}") logger.debug(f"Google chunk type: {type(chunk)}")
logger.debug(f"Google chunk: {chunk}") logger.debug(f"Google chunk: {chunk}")
# Extract text from Google chunk (this is accumulated text) # Extract text and tool calls from Google chunk (this is accumulated text)
chunk_text = "" chunk_text = ""
chunk_tool_calls = []
finish_reason = None finish_reason = None
try: try:
if hasattr(chunk, 'candidates') and chunk.candidates: if hasattr(chunk, 'candidates') and chunk.candidates:
...@@ -441,8 +443,24 @@ class RequestHandler: ...@@ -441,8 +443,24 @@ class RequestHandler:
if candidate and hasattr(candidate, 'content') and candidate.content: if candidate and hasattr(candidate, 'content') and candidate.content:
if hasattr(candidate.content, 'parts') and candidate.content.parts: if hasattr(candidate.content, 'parts') and candidate.content.parts:
for part in candidate.content.parts: for part in candidate.content.parts:
# Extract text content
if hasattr(part, 'text') and part.text: if hasattr(part, 'text') and part.text:
chunk_text += 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 # Check for finish reason in candidate
if hasattr(candidate, 'finish_reason'): if hasattr(candidate, 'finish_reason'):
google_finish = str(candidate.finish_reason) google_finish = str(candidate.finish_reason)
...@@ -457,12 +475,16 @@ class RequestHandler: ...@@ -457,12 +475,16 @@ class RequestHandler:
delta_text = chunk_text[len(accumulated_text):] if chunk_text.startswith(accumulated_text) else chunk_text 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 accumulated_text = chunk_text # Update accumulated text for next iteration
# Calculate delta tool calls (only new tool calls since last chunk)
delta_tool_calls = chunk_tool_calls[len(accumulated_tool_calls):] if chunk_tool_calls else []
accumulated_tool_calls = chunk_tool_calls # Update accumulated tool calls for next iteration
# Check if this is the last chunk # Check if this is the last chunk
is_last_chunk = (chunk_idx == total_chunks - 1) is_last_chunk = (chunk_idx == total_chunks - 1)
chunk_finish_reason = finish_reason if is_last_chunk else None chunk_finish_reason = finish_reason if is_last_chunk else None
# Only send if there's new content or it's the last chunk with finish_reason # Only send if there's new content, new tool calls, or it's the last chunk with finish_reason
if delta_text or is_last_chunk: if delta_text or delta_tool_calls or is_last_chunk:
# Create OpenAI-compatible chunk with additional fields # Create OpenAI-compatible chunk with additional fields
openai_chunk = { openai_chunk = {
"id": response_id, "id": response_id,
...@@ -479,7 +501,7 @@ class RequestHandler: ...@@ -479,7 +501,7 @@ class RequestHandler:
"content": delta_text if delta_text else "", "content": delta_text if delta_text else "",
"refusal": None, "refusal": None,
"role": "assistant", "role": "assistant",
"tool_calls": None "tool_calls": delta_tool_calls if delta_tool_calls else None
}, },
"finish_reason": chunk_finish_reason, "finish_reason": chunk_finish_reason,
"logprobs": None, "logprobs": None,
...@@ -1511,6 +1533,7 @@ class RotationHandler: ...@@ -1511,6 +1533,7 @@ class RotationHandler:
# Note: Google returns accumulated text, so we need to track and send only deltas # Note: Google returns accumulated text, so we need to track and send only deltas
chunk_id = 0 chunk_id = 0
accumulated_text = "" # Track text we've already sent accumulated_text = "" # Track text we've already sent
accumulated_tool_calls = [] # Track tool calls we've already sent
created_time = int(time.time()) created_time = int(time.time())
response_id = f"google-{model_name}-{created_time}" response_id = f"google-{model_name}-{created_time}"
...@@ -1531,8 +1554,9 @@ class RotationHandler: ...@@ -1531,8 +1554,9 @@ class RotationHandler:
logger.debug(f"Google chunk type: {type(chunk)}") logger.debug(f"Google chunk type: {type(chunk)}")
logger.debug(f"Google chunk: {chunk}") logger.debug(f"Google chunk: {chunk}")
# Extract text from Google chunk (this is accumulated text) # Extract text and tool calls from Google chunk (this is accumulated text)
chunk_text = "" chunk_text = ""
chunk_tool_calls = []
finish_reason = None finish_reason = None
try: try:
if hasattr(chunk, 'candidates') and chunk.candidates: if hasattr(chunk, 'candidates') and chunk.candidates:
...@@ -1540,8 +1564,24 @@ class RotationHandler: ...@@ -1540,8 +1564,24 @@ class RotationHandler:
if candidate and hasattr(candidate, 'content') and candidate.content: if candidate and hasattr(candidate, 'content') and candidate.content:
if hasattr(candidate.content, 'parts') and candidate.content.parts: if hasattr(candidate.content, 'parts') and candidate.content.parts:
for part in candidate.content.parts: for part in candidate.content.parts:
# Extract text content
if hasattr(part, 'text') and part.text: if hasattr(part, 'text') and part.text:
chunk_text += 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 # Check for finish reason in candidate
if hasattr(candidate, 'finish_reason'): if hasattr(candidate, 'finish_reason'):
google_finish = str(candidate.finish_reason) google_finish = str(candidate.finish_reason)
......
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