mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-08-14 08:52:06 +00:00
fix: preserve Gemini tool calls while streaming
This commit is contained in:
Dustin Zander
committed by
Elliot Slusky
Elliot Slusky
parent
063dd8ea75
commit
1bfc25a860
@@ -1305,6 +1305,133 @@ class CloudEngine(InferenceEngine):
|
||||
if chunk.text:
|
||||
yield chunk.text
|
||||
|
||||
async def _stream_full_google(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
*,
|
||||
model: str,
|
||||
temperature: float,
|
||||
max_tokens: int,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[StreamChunk]:
|
||||
"""Stream Google text and function-call parts as full chunks."""
|
||||
if self._google_client is None:
|
||||
raise EngineConnectionError("Google client not available")
|
||||
|
||||
system_text = ""
|
||||
contents: List[Dict[str, Any]] = []
|
||||
for message in messages:
|
||||
if message.role.value == "system":
|
||||
system_text = message.content
|
||||
elif message.role.value == "tool":
|
||||
function_response = {
|
||||
"function_response": {
|
||||
"name": message.name or "unknown",
|
||||
"response": {"result": message.content},
|
||||
}
|
||||
}
|
||||
if (
|
||||
contents
|
||||
and contents[-1]["role"] == "user"
|
||||
and contents[-1]["parts"]
|
||||
and "function_response" in contents[-1]["parts"][-1]
|
||||
):
|
||||
contents[-1]["parts"].append(function_response)
|
||||
else:
|
||||
contents.append({"role": "user", "parts": [function_response]})
|
||||
elif message.role.value == "assistant" and message.tool_calls:
|
||||
parts: List[Dict[str, Any]] = []
|
||||
if message.content:
|
||||
parts.append({"text": message.content})
|
||||
for tool_call in message.tool_calls:
|
||||
args = tool_call.arguments
|
||||
if isinstance(args, str):
|
||||
try:
|
||||
args = json.loads(args)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
args = {"input": args}
|
||||
function_call: Dict[str, Any] = {
|
||||
"name": tool_call.name,
|
||||
"args": args if isinstance(args, dict) else {},
|
||||
}
|
||||
signature = self._thought_sigs.get(tool_call.id)
|
||||
if signature is not None:
|
||||
function_call["thought_signature"] = signature
|
||||
parts.append({"function_call": function_call})
|
||||
contents.append({"role": "model", "parts": parts})
|
||||
elif message.role.value == "assistant":
|
||||
contents.append({"role": "model", "parts": [{"text": message.content}]})
|
||||
else:
|
||||
contents.append({"role": "user", "parts": [{"text": message.content}]})
|
||||
|
||||
from google.genai import types as genai_types
|
||||
|
||||
config = genai_types.GenerateContentConfig(
|
||||
temperature=temperature,
|
||||
max_output_tokens=max_tokens,
|
||||
)
|
||||
if system_text:
|
||||
config.system_instruction = system_text
|
||||
|
||||
tools = kwargs.pop("tools", None)
|
||||
if tools:
|
||||
config.tools = [{"function_declarations": _convert_tools_to_google(tools)}]
|
||||
|
||||
tool_ids: Dict[str, int] = {}
|
||||
for chunk in self._google_client.models.generate_content_stream(
|
||||
model=model,
|
||||
contents=contents,
|
||||
config=config,
|
||||
):
|
||||
candidates = getattr(chunk, "candidates", None)
|
||||
parts = []
|
||||
if candidates:
|
||||
parts = getattr(candidates[0].content, "parts", []) or []
|
||||
|
||||
if parts:
|
||||
text_found = False
|
||||
calls: List[Dict[str, Any]] = []
|
||||
for part in parts:
|
||||
text = getattr(part, "text", None)
|
||||
if text:
|
||||
text_found = True
|
||||
yield StreamChunk(content=text)
|
||||
|
||||
function_call = getattr(part, "function_call", None)
|
||||
if function_call:
|
||||
name = getattr(function_call, "name", "")
|
||||
raw_args = getattr(function_call, "args", {})
|
||||
args = dict(raw_args) if hasattr(raw_args, "items") else {}
|
||||
tool_id = f"google_{name}"
|
||||
tool_index = tool_ids.setdefault(tool_id, len(tool_ids))
|
||||
tool_call = {
|
||||
"index": tool_index,
|
||||
"id": tool_id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": name,
|
||||
"arguments": json.dumps(args),
|
||||
},
|
||||
}
|
||||
calls.append(tool_call)
|
||||
signature = getattr(part, "thought_signature", None)
|
||||
if signature is not None:
|
||||
tool_call["thought_signature"] = signature
|
||||
self._thought_sigs[tool_id] = signature
|
||||
if calls:
|
||||
yield StreamChunk(tool_calls=calls)
|
||||
if text_found:
|
||||
continue
|
||||
|
||||
try:
|
||||
text = chunk.text
|
||||
except (AttributeError, ValueError):
|
||||
text = None
|
||||
if text:
|
||||
yield StreamChunk(content=text)
|
||||
|
||||
yield StreamChunk(finish_reason="tool_calls" if tool_ids else "stop")
|
||||
|
||||
async def _stream_openrouter(
|
||||
self,
|
||||
messages: Sequence[Message],
|
||||
@@ -1600,7 +1727,7 @@ class CloudEngine(InferenceEngine):
|
||||
async for chunk in self._stream_full_anthropic(messages, **kw):
|
||||
yield chunk
|
||||
elif _is_google_model(model):
|
||||
async for chunk in super().stream_full(messages, **kw):
|
||||
async for chunk in self._stream_full_google(messages, **kw):
|
||||
yield chunk
|
||||
else:
|
||||
async for chunk in self._stream_full_openai(messages, **kw):
|
||||
|
||||
@@ -3,6 +3,8 @@ and _prepare_anthropic_messages."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from typing import Any, List
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
@@ -63,6 +65,28 @@ def _openai_tool_call_delta(
|
||||
return tc
|
||||
|
||||
|
||||
class _GoogleConfig:
|
||||
def __init__(self, **kwargs: Any) -> None:
|
||||
self.__dict__.update(kwargs)
|
||||
|
||||
|
||||
def _google_stream_chunk(*parts: Any, text: str | None = None) -> Any:
|
||||
candidates = []
|
||||
if parts:
|
||||
candidates = [SimpleNamespace(content=SimpleNamespace(parts=list(parts)))]
|
||||
return SimpleNamespace(text=text, candidates=candidates, usage_metadata=None)
|
||||
|
||||
|
||||
def _google_types_modules() -> dict[str, ModuleType]:
|
||||
types = ModuleType("google.genai.types")
|
||||
types.GenerateContentConfig = _GoogleConfig
|
||||
genai = ModuleType("google.genai")
|
||||
genai.types = types
|
||||
google = ModuleType("google")
|
||||
google.genai = genai
|
||||
return {"google": google, "google.genai": genai, "google.genai.types": types}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _stream_full_openai tests
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -413,6 +437,166 @@ def test_prepare_anthropic_messages_tool_calls():
|
||||
assert blocks[1]["input"] == {"city": "Berlin"}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _stream_full_google tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_text_only(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Google text chunks retain their content and finish normally."""
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[_google_stream_chunk(text="Hello"), _google_stream_chunk(text=" world")]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
messages = [Message(role=Role.USER, content="hi")]
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(messages, model="gemini-2.5-flash")
|
||||
]
|
||||
|
||||
assert [chunk.content for chunk in result[:-1]] == ["Hello", " world"]
|
||||
assert result[-1].finish_reason == "stop"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_preserves_tool_calls(monkeypatch: pytest.MonkeyPatch):
|
||||
"""Google function_call parts become OpenAI-compatible tool call chunks."""
|
||||
function_call = SimpleNamespace(name="get_weather", args={"city": "Berlin"})
|
||||
part = SimpleNamespace(
|
||||
function_call=function_call, text=None, thought_signature=b"sig"
|
||||
)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[_google_stream_chunk(part)]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
messages = [Message(role=Role.USER, content="weather")]
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
messages,
|
||||
model="gemini-2.5-flash",
|
||||
tools=[
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
},
|
||||
}
|
||||
],
|
||||
)
|
||||
]
|
||||
|
||||
tool_call = result[0].tool_calls[0]
|
||||
assert tool_call == {
|
||||
"index": 0,
|
||||
"id": "google_get_weather",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Berlin"}'},
|
||||
"thought_signature": b"sig",
|
||||
}
|
||||
assert engine._thought_sigs["google_get_weather"] == b"sig"
|
||||
assert result[-1].finish_reason == "tool_calls"
|
||||
config = client.models.generate_content_stream.call_args.kwargs["config"]
|
||||
assert config.tools == [
|
||||
{
|
||||
"function_declarations": [
|
||||
{
|
||||
"name": "get_weather",
|
||||
"description": "Get weather",
|
||||
"parameters": {"type": "object", "properties": {}},
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_full_google_preserves_mixed_and_multiple_calls(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
):
|
||||
"""Google streams retain mixed text and multiple deterministic tool calls."""
|
||||
weather = SimpleNamespace(name="get_weather", args={"city": "Berlin"})
|
||||
calendar = SimpleNamespace(name="get_calendar", args={"day": "Monday"})
|
||||
text_part = SimpleNamespace(text="I'll check.", function_call=None)
|
||||
weather_part = SimpleNamespace(
|
||||
function_call=weather, text=None, thought_signature=None
|
||||
)
|
||||
calendar_part = SimpleNamespace(
|
||||
function_call=calendar, text=None, thought_signature=None
|
||||
)
|
||||
client = MagicMock()
|
||||
client.models.generate_content_stream.return_value = iter(
|
||||
[
|
||||
_google_stream_chunk(text_part, weather_part),
|
||||
_google_stream_chunk(weather_part, calendar_part),
|
||||
]
|
||||
)
|
||||
engine = _make_cloud_engine(google_client=client)
|
||||
engine._thought_sigs = {}
|
||||
modules = _google_types_modules()
|
||||
|
||||
with monkeypatch.context() as patch:
|
||||
for name, module in modules.items():
|
||||
patch.setitem(sys.modules, name, module)
|
||||
result = [
|
||||
chunk
|
||||
async for chunk in engine.stream_full(
|
||||
[Message(role=Role.USER, content="plan")], model="gemini-2.5-flash"
|
||||
)
|
||||
]
|
||||
|
||||
assert result[0].content == "I'll check."
|
||||
assert result[1].tool_calls == [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "google_get_weather",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Berlin"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
assert result[2].tool_calls == [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "google_get_weather",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_weather",
|
||||
"arguments": '{"city": "Berlin"}',
|
||||
},
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"id": "google_get_calendar",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "get_calendar",
|
||||
"arguments": '{"day": "Monday"}',
|
||||
},
|
||||
},
|
||||
]
|
||||
assert result[-1].finish_reason == "tool_calls"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# stream_full routing tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
Reference in New Issue
Block a user