diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ee2e8dd62..d709bf006 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -27,7 +27,7 @@ jobs: run: uv python install ${{ matrix.python-version }} - name: Run tests - run: uv run --python ${{ matrix.python-version }} --all-extras pytest --cov --cov-report='' + run: uv run --python ${{ matrix.python-version }} --all-extras --no-extra browseruse pytest --cov --cov-report='' lint-ruff: runs-on: ubuntu-latest diff --git a/docs/v6/reference/agents.mdx b/docs/v6/reference/agents.mdx index 62a93f03b..764d0b71f 100644 --- a/docs/v6/reference/agents.mdx +++ b/docs/v6/reference/agents.mdx @@ -57,7 +57,7 @@ agent = ClaudeAgent(ClaudeConfig(model="claude-sonnet-4-5", max_steps=30)) |-------|--------|---------------| | `ClaudeAgent` | `ClaudeConfig` | `claude-sonnet-4-6` | | `OpenAIAgent` | `OpenAIConfig` | `gpt-5.6` | -| `GeminiAgent` | `GeminiConfig` | `gemini-3-pro-preview` | +| `GeminiAgent` | `GeminiConfig` | `gemini-3.1-pro-preview` | | `OpenAIChatAgent` | `OpenAIChatConfig` | `gpt-5.4-mini` | | `ClaudeCLIAgent` | `ClaudeCLIConfig` | `claude-sonnet-5` | | `CodexCLIAgent` | `CodexCLIConfig` | `gpt-5.6-sol` | @@ -72,6 +72,10 @@ config; `__call__(run)` takes only the run. `OpenAIConfig.max_output_tokens` and `GeminiConfig.max_output_tokens` are unset by default, so each request gets the model's own output limit. Set them to cap output per turn. Gemini counts thinking tokens toward the cap, so a low cap can truncate a thinking model before it answers. +`GeminiConfig.temperature`, `top_p`, and `top_k` are unset by default, so each request uses the model's own sampling; Gemini 4 models ignore them. `GeminiConfig.thinking_level` accepts `minimal`, `low`, `medium`, or `high`, subject to the model (Gemini 4 Argon rejects `minimal`); unset uses the model's default. + +`GeminiAgent` streams each turn. For Gemini 4 long decode, set `max_output_tokens` above 65,536 and install `google-genai` 2.30.1 or later; the turn then runs past the per-request output cap and is recorded as one turn. Otherwise a turn that reaches the cap ends the run with `stop_reason="length"`. Keep `include_thoughts` on for long-thinking runs, since without it the model sends nothing while it thinks and a long think can hit the stream's first-chunk deadline. + `OpenAIConfig` generates a stable `prompt_cache_key` for the agent and sends it with every Responses request. Pass your own key when multiple agent instances should share the same cached prompt prefix. diff --git a/hud/agents/claude/tools/hosted.py b/hud/agents/claude/tools/hosted.py index d7076acb3..9f4b33e23 100644 --- a/hud/agents/claude/tools/hosted.py +++ b/hud/agents/claude/tools/hosted.py @@ -6,7 +6,6 @@ from typing import TYPE_CHECKING, Any from anthropic.types.beta import ( - BetaCitationsConfigParam, BetaToolSearchToolBm25_20251119Param, BetaToolUnionParam, BetaWebFetchTool20250910Param, @@ -75,7 +74,7 @@ def to_params(self) -> BetaWebFetchTool20250910Param: if self.max_content_tokens is not None: params["max_content_tokens"] = self.max_content_tokens if self.citations_enabled: - params["citations"] = BetaCitationsConfigParam(enabled=True) + params["citations"] = {"enabled": True} return params diff --git a/hud/agents/gemini/agent.py b/hud/agents/gemini/agent.py index 649ce27a8..b508329d1 100644 --- a/hud/agents/gemini/agent.py +++ b/hud/agents/gemini/agent.py @@ -201,22 +201,40 @@ async def get_response( thinking_config=thinking_config, ) - api_response = await self.gemini_client.aio.models.generate_content( + stream = await self.gemini_client.aio.models.generate_content_stream( model=self.config.model, contents=cast("Any", messages), config=generate_config, ) - if not api_response.candidates: + # The turn keeps every streamed part as received: thought signatures belong to the + # parts they arrive on, including empty text parts. + parts: list[genai_types.Part] = [] + candidate_seen = False + model_version: str | None = None + usage_meta: genai_types.GenerateContentResponseUsageMetadata | None = None + grounding_meta: genai_types.GroundingMetadata | None = None + finish_reason: genai_types.FinishReason | None = None + async for chunk in stream: + model_version = chunk.model_version or model_version + if chunk.usage_metadata is not None: + usage_meta = chunk.usage_metadata + if not chunk.candidates: + continue + candidate_seen = True + candidate = chunk.candidates[0] + if candidate.content is not None and candidate.content.parts: + parts.extend(candidate.content.parts) + if candidate.grounding_metadata is not None: + grounding_meta = candidate.grounding_metadata + if candidate.finish_reason is not None: + finish_reason = candidate.finish_reason + if not candidate_seen: raise RuntimeError(f"Gemini returned no candidates for model {self.config.model}") - - candidate = api_response.candidates[0] - content = candidate.content - if content is not None: - messages.append(content) + if parts: + messages.append(genai_types.Content(role="model", parts=parts)) result = AgentStep(content="", done=True) - result.model = api_response.model_version or self.config.model - usage_meta = api_response.usage_metadata + result.model = model_version or self.config.model if usage_meta is not None: result.usage = Usage( prompt_tokens=usage_meta.prompt_token_count, @@ -226,7 +244,7 @@ async def get_response( text_parts: list[str] = [] thought_parts: list[str] = [] - for part in (content.parts or []) if content else []: + for part in parts: function_call = part.function_call if function_call is not None: tc = self._make_tool_call(function_call, computer_tool) @@ -241,14 +259,13 @@ async def get_response( result.content = "".join(text_parts) if thought_parts: - result.reasoning = "\n".join(thought_parts) + result.reasoning = "".join(thought_parts) - grounding_meta = candidate.grounding_metadata if grounding_meta is not None: result.citations = _grounding_citations(grounding_meta) - if candidate.finish_reason is not None: - result.finish_reason = candidate.finish_reason.name + if finish_reason is not None: + result.finish_reason = finish_reason.name if result.finish_reason == "MALFORMED_FUNCTION_CALL": result.error = "Provider returned a malformed function call" result.stop_reason = "malformed_tool_call" diff --git a/hud/agents/tests/test_claude_agent.py b/hud/agents/tests/test_claude_agent.py index 38571b214..c9a07b4c7 100644 --- a/hud/agents/tests/test_claude_agent.py +++ b/hud/agents/tests/test_claude_agent.py @@ -104,10 +104,10 @@ def _state(agent: ClaudeAgent) -> Any: def _inline_error(type_: str, *, headers: dict[str, str] | None = None) -> APIStatusError: body = {"type": "error", "error": {"type": type_, "message": "stream failed"}} - response = httpx.Response( + response = httpx2.Response( 200, headers=headers, - request=httpx.Request("POST", "https://api.anthropic.com/v1/messages"), + request=httpx2.Request("POST", "https://api.anthropic.com/v1/messages"), ) return APIStatusError(str(body), response=response, body=body) diff --git a/hud/agents/tests/test_gemini_agent.py b/hud/agents/tests/test_gemini_agent.py index 30b672dda..45d5e0c2c 100644 --- a/hud/agents/tests/test_gemini_agent.py +++ b/hud/agents/tests/test_gemini_agent.py @@ -1,13 +1,17 @@ -"""``GeminiAgent`` — ``get_response`` parsing over a fake Generate Content client, +"""``GeminiAgent`` — ``get_response`` over a fake streaming Generate Content client, plus ``_make_tool_call`` mapping and ``_grounding_citations``. """ from __future__ import annotations +import json from types import SimpleNamespace from typing import TYPE_CHECKING, Any, cast +import httpx import mcp.types as mcp_types +from google import genai +from google.genai import types as genai_types from hud.agents.gemini.agent import GeminiAgent, _grounding_citations from hud.agents.tool_agent import RunState @@ -15,15 +19,22 @@ from hud.types import MCPToolResult, Step, Trace if TYPE_CHECKING: + from collections.abc import AsyncIterator + from hud.eval.run import Run +async def _chunks(*chunks: Any) -> AsyncIterator[Any]: + for chunk in chunks: + yield chunk + + class FakeModels: def __init__(self, response: Any) -> None: self._response = response - async def generate_content(self, **_kwargs: Any) -> Any: - return self._response + async def generate_content_stream(self, **_kwargs: Any) -> AsyncIterator[Any]: + return _chunks(self._response) class FakeGenai: @@ -109,7 +120,7 @@ async def test_get_response_done_text_only() -> None: async def test_get_response_no_candidates_raises() -> None: - agent = _agent(SimpleNamespace(candidates=[])) + agent = _agent(SimpleNamespace(candidates=[], model_version=None, usage_metadata=None)) try: await agent.get_response(_state(agent)) except RuntimeError: @@ -151,6 +162,209 @@ def test_grounding_citations() -> None: assert cites[0].type == "grounding" +def _chunk(*parts: genai_types.Part, finish_reason: str | None = None) -> Any: + return genai_types.GenerateContentResponse( + candidates=[ + genai_types.Candidate( + content=genai_types.Content(role="model", parts=list(parts)), + finish_reason=genai_types.FinishReason(finish_reason) if finish_reason else None, + ) + ], + model_version="gemini-test-v2", + ) + + +async def test_streamed_chunks_form_one_model_turn() -> None: + """Every streamed part joins one model turn as received; text and thoughts concatenate.""" + parts = [ + genai_types.Part(text="Checking ", thought=True), + genai_types.Part(text="divisors.", thought=True), + genai_types.Part(text="It is 17 *"), + genai_types.Part(text=" 23."), + genai_types.Part( + function_call=genai_types.FunctionCall(name="bash", args={"command": "ls"}), + thought_signature=b"sig-1", + ), + genai_types.Part(text="", thought_signature=b"sig-2"), + ] + stream = _chunks( + _chunk(*parts[:2]), + _chunk(*parts[2:4]), + _chunk(*parts[4:], finish_reason="STOP"), + ) + + async def generate_content_stream(**_kwargs: Any) -> AsyncIterator[Any]: + return stream + + agent = GeminiAgent( + GeminiConfig( + model="gemini-test", + model_client=SimpleNamespace( + aio=SimpleNamespace( + models=SimpleNamespace(generate_content_stream=generate_content_stream) + ) + ), + ) + ) + state = _state(agent) + + result = await agent.get_response(state) + + assert result.content == "It is 17 * 23." + assert result.reasoning == "Checking divisors." + assert [tc.name for tc in result.tool_calls] == ["bash"] + assert result.finish_reason == "STOP" + assert result.model == "gemini-test-v2" + turn = state.messages[-1] + assert turn.role == "model" + assert turn.parts == parts + + +def _sse_client(*hops: list[dict[str, Any]]) -> tuple[genai.Client, list[httpx.Request]]: + """A ``genai.Client`` whose n-th request streams the n-th hop's events over SSE.""" + requests: list[httpx.Request] = [] + + def handler(request: httpx.Request) -> httpx.Response: + requests.append(request) + events = hops[len(requests) - 1] + return httpx.Response( + 200, + content=b"".join( + b"data: " + json.dumps(event).encode() + b"\r\n\r\n" for event in events + ), + headers={"content-type": "text/event-stream"}, + ) + + client = genai.Client( + api_key="hud-key", + http_options=genai_types.HttpOptions( + base_url="https://gateway.test", + api_version="v1beta", + async_client_args={"transport": httpx.MockTransport(handler)}, + ), + ) + return client, requests + + +async def test_turn_streams_through_the_sdk() -> None: + """A turn arrives over the SDK's SSE stream and is recorded as one model turn.""" + events: list[dict[str, Any]] = [ + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "Thinking.", "thought": True}]}} + ] + }, + {"candidates": [{"content": {"role": "model", "parts": [{"text": "39"}]}}]}, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "1"}, {"text": "", "thoughtSignature": "c2ln"}], + }, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 2, + "thoughtsTokenCount": 5, + "totalTokenCount": 14, + }, + "modelVersion": "gemini-4-argon", + }, + ] + client, requests = _sse_client(events) + agent = GeminiAgent(GeminiConfig(model="gemini-4-argon", model_client=client)) + state = _state(agent) + + result = await agent.get_response(state) + + [request] = requests + assert request.url.path == "/v1beta/models/gemini-4-argon:streamGenerateContent" + assert request.url.params["alt"] == "sse" + assert result.content == "391" + assert result.reasoning == "Thinking." + assert result.finish_reason == "STOP" + assert result.usage is not None + assert result.usage.prompt_tokens == 7 + turn = state.messages[-1] + assert [(part.text, part.thought, part.thought_signature) for part in turn.parts or []] == [ + ("Thinking.", True, None), + ("39", None, None), + ("1", None, None), + ("", None, b"sig"), + ] + + +async def test_turn_past_the_per_request_cap_continues_as_one_turn() -> None: + """A stream that stops at the per-request cap resumes with its continuation token, and + both hops form one model turn.""" + first_hop: list[dict[str, Any]] = [ + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "Working it out.", "thought": True}], + } + } + ] + }, + { + "candidates": [ + { + "content": { + "role": "model", + "parts": [{"text": "", "thoughtSignature": "c2lnLTE="}], + }, + "finishReason": "CONTINUATION", + "continuationToken": "dG9rZW4tMQ==", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "thoughtsTokenCount": 65526, + "totalTokenCount": 65536, + }, + }, + ] + second_hop: list[dict[str, Any]] = [ + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "42"}]}, "finishReason": "STOP"} + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 2, + "totalTokenCount": 12, + }, + "modelVersion": "gemini-4-argon", + } + ] + client, requests = _sse_client(first_hop, second_hop) + agent = GeminiAgent(GeminiConfig(model="gemini-4-argon", model_client=client)) + state = _state(agent) + + result = await agent.get_response(state) + + first, second = (json.loads(request.content) for request in requests) + assert second == {**first, "continuationToken": "dG9rZW4tMQ=="} + assert result.content == "42" + assert result.reasoning == "Working it out." + assert result.done is True + assert result.finish_reason == "STOP" + assert result.usage is not None + assert result.usage.prompt_tokens == 20 + turns = [content for content in state.messages if content.role == "model"] + assert len(turns) == 1 + assert [(part.text, part.thought_signature) for part in turns[0].parts or []] == [ + ("Working it out.", None), + ("", b"sig-1"), + ("42", None), + ] + + async def test_malformed_function_call_is_a_normalized_error() -> None: response = _api_response( SimpleNamespace( @@ -172,12 +386,12 @@ def __init__(self, response: Any) -> None: self._response = response self.configs: list[Any] = [] - async def generate_content(self, **kwargs: Any) -> Any: + async def generate_content_stream(self, **kwargs: Any) -> AsyncIterator[Any]: self.configs.append(kwargs["config"]) - return self._response + return _chunks(self._response) -async def _sent_output_cap(**config: Any) -> int | None: +async def _sent_config(**config: Any) -> Any: resp_content = SimpleNamespace( role="model", parts=[SimpleNamespace(function_call=None, text="answer", thought=None)], @@ -196,15 +410,25 @@ async def _sent_output_cap(**config: Any) -> int | None: ) ) await agent.get_response(_state(agent)) - return models.configs[0].max_output_tokens + return models.configs[0] async def test_output_cap_defaults_to_the_model_limit() -> None: - assert await _sent_output_cap() is None + assert (await _sent_config()).max_output_tokens is None async def test_configured_output_cap_is_sent() -> None: - assert await _sent_output_cap(max_output_tokens=4096) == 4096 + assert (await _sent_config(max_output_tokens=4096)).max_output_tokens == 4096 + + +async def test_sampling_defaults_to_the_model() -> None: + sent = await _sent_config() + assert (sent.temperature, sent.top_p, sent.top_k) == (None, None, None) + + +async def test_configured_sampling_is_sent() -> None: + sent = await _sent_config(temperature=0.5, top_p=0.9, top_k=20) + assert (sent.temperature, sent.top_p, sent.top_k) == (0.5, 0.9, 20) class RecordingModels: @@ -214,10 +438,10 @@ def __init__(self, responses: list[Any]) -> None: self._responses = list(responses) self.requests: list[Any] = [] - async def generate_content(self, **kwargs: Any) -> Any: + async def generate_content_stream(self, **kwargs: Any) -> AsyncIterator[Any]: # Snapshot: the agent goes on mutating the list it passed. self.requests.append(list(kwargs["contents"])) - return self._responses.pop(0) + return _chunks(self._responses.pop(0)) class RecordingGenai: diff --git a/hud/agents/tests/test_tool_agent.py b/hud/agents/tests/test_tool_agent.py index 19ff303b8..1ca164622 100644 --- a/hud/agents/tests/test_tool_agent.py +++ b/hud/agents/tests/test_tool_agent.py @@ -725,7 +725,13 @@ async def test_loop_marks_length_finish_as_truncated() -> None: # A final turn cut off at the provider token cap (e.g. mid-tool-call) ends the # rollout normally but is a truncation, not a natural finish — across every # provider's finish-reason vocabulary. - for finish_reason in ("length", "max_output_tokens", "max_tokens", "MAX_TOKENS"): + for finish_reason in ( + "length", + "max_output_tokens", + "max_tokens", + "MAX_TOKENS", + "CONTINUATION", + ): agent = DictAgent([AgentStep(content="partial", done=True, finish_reason=finish_reason)]) run = cast("Run", _FakeRun()) diff --git a/hud/agents/tool_agent.py b/hud/agents/tool_agent.py index 198776e0a..d909fe424 100644 --- a/hud/agents/tool_agent.py +++ b/hud/agents/tool_agent.py @@ -31,8 +31,11 @@ logger = logging.getLogger(__name__) # Output-token-cap finish reasons: OpenAI chat "length", Responses "max_output_tokens", -# Claude "max_tokens", Gemini "MAX_TOKENS". -TRUNCATION_FINISH_REASONS = frozenset({"length", "max_output_tokens", "max_tokens", "MAX_TOKENS"}) +# Claude "max_tokens", Gemini "MAX_TOKENS", and Gemini "CONTINUATION" (a per-request +# cap hit by a turn that was not continued). +TRUNCATION_FINISH_REASONS = frozenset( + {"length", "max_output_tokens", "max_tokens", "MAX_TOKENS", "CONTINUATION"} +) MAX_CONSECUTIVE_SSH_FAILURES = 3 MessageT = TypeVar("MessageT") diff --git a/hud/agents/types.py b/hud/agents/types.py index 5be380bf8..a40b06494 100644 --- a/hud/agents/types.py +++ b/hud/agents/types.py @@ -146,9 +146,9 @@ class GeminiConfig(ToolAgentConfig): model_name: str = "Gemini" model: str = Field(default="gemini-3.1-pro-preview", validation_alias=_model_alias) - temperature: float = 1.0 - top_p: float = 0.95 - top_k: int = 40 + temperature: float | None = None + top_p: float | None = None + top_k: int | None = None max_output_tokens: int | None = None excluded_predefined_functions: list[str] = Field(default_factory=list) thinking_level: Literal["minimal", "low", "medium", "high"] | None = None diff --git a/hud/utils/tests/test_gateway.py b/hud/utils/tests/test_gateway.py index f452994f6..4ba61f9cf 100644 --- a/hud/utils/tests/test_gateway.py +++ b/hud/utils/tests/test_gateway.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx +import httpx2 import pytest from hud.settings import settings @@ -30,13 +31,13 @@ async def test_openai_client_resolves_trace_id_per_request( ) -> None: seen_headers: dict[str, str] = {} - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: trace_id = request.headers["Trace-Id"] await asyncio.sleep(0) seen_headers[request.url.path] = trace_id - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, @@ -69,12 +70,12 @@ async def test_openai_client_trace_context_overrides_empty_header( ) -> None: seen_trace_id: str | None = None - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: nonlocal seen_trace_id seen_trace_id = request.headers["Trace-Id"] - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, @@ -103,12 +104,12 @@ async def test_openai_client_sends_child_and_parent_trace_ids( ) -> None: seen_headers: dict[str, str] = {} - async def handler(request: httpx.Request) -> httpx.Response: + async def handler(request: httpx2.Request) -> httpx2.Response: seen_headers["Trace-Id"] = request.headers["Trace-Id"] seen_headers["X-HUD-Parent-Trace-Id"] = request.headers["X-HUD-Parent-Trace-Id"] - return httpx.Response(200, json={"object": "list", "data": []}) + return httpx2.Response(200, json={"object": "list", "data": []}) - transport = httpx.MockTransport(handler) + transport = httpx2.MockTransport(handler) real_client_factory = gateway.DefaultAsyncHttpxClient monkeypatch.setattr( gateway, diff --git a/pyproject.toml b/pyproject.toml index b4ccdbfe5..25caa69d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,6 +129,8 @@ bedrock = [ dev = [ "hud[agents]", # Include agents for dev "hud[robot]", # the robot end-to-end gates (hud/tests/test_robot.py) run the real wire + # GeminiAgent's long-decode tests need the SDK's automatic continuation. + "google-genai>=2.30.1", "dotenv>=0.9.9", # Testing and linting "ruff==0.16.1", @@ -173,6 +175,9 @@ train = [ [tool.uv] required-version = "==0.12.2" +# Every browser-use release pins google-genai below 2.0, so the browseruse extra +# resolves apart from dev, which needs google-genai 2.30.1 or later. +conflicts = [[{ extra = "browseruse" }, { extra = "dev" }]] [tool.ruff] target-version = "py311"