diff --git a/src/openjarvis/engine/cloud.py b/src/openjarvis/engine/cloud.py index 9e39cf6d..fe5172f6 100644 --- a/src/openjarvis/engine/cloud.py +++ b/src/openjarvis/engine/cloud.py @@ -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): diff --git a/tests/engine/test_cloud_stream_full.py b/tests/engine/test_cloud_stream_full.py index ce73d83a..8adf235e 100644 --- a/tests/engine/test_cloud_stream_full.py +++ b/tests/engine/test_cloud_stream_full.py @@ -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 # ---------------------------------------------------------------------------