diff --git a/src/utils/clients.py b/src/utils/clients.py index 1c042bff..7fcf7915 100644 --- a/src/utils/clients.py +++ b/src/utils/clients.py @@ -16,7 +16,7 @@ from google.genai.types import ( GenerateContentResponse, ) from groq import AsyncGroq -from openai import AsyncOpenAI +from openai import AsyncOpenAI, LengthFinishReasonError from openai.types.chat import ChatCompletion, ChatCompletionChunk from pydantic import BaseModel, Field, ValidationError from sentry_sdk.ai.monitoring import ai_track @@ -606,7 +606,7 @@ async def _stream_final_response( stop_seqs: Stop sequences reasoning_effort: OpenAI reasoning effort (GPT-5 only) verbosity: OpenAI verbosity (GPT-5 only) - thinking_budget_tokens: Anthropic thinking budget + thinking_budget_tokens: Anthropic / Gemini thinking budget Yields: HonchoLLMCallStreamChunk objects containing the streaming response @@ -693,7 +693,7 @@ async def _execute_tool_loop( stop_seqs: Stop sequences reasoning_effort: OpenAI reasoning effort (GPT-5 only) verbosity: OpenAI verbosity (GPT-5 only) - thinking_budget_tokens: Anthropic thinking budget + thinking_budget_tokens: Anthropic / Gemini thinking budget enable_retry: Whether to enable retry with exponential backoff retry_attempts: Number of retry attempts max_input_tokens: Maximum input tokens (for truncation) @@ -1322,7 +1322,7 @@ async def honcho_llm_call( stop_seqs: Stop sequences reasoning_effort: OpenAI reasoning effort (GPT-5 only) verbosity: OpenAI verbosity (GPT-5 only) - thinking_budget_tokens: Anthropic thinking budget + thinking_budget_tokens: Anthropic / Gemini thinking budget enable_retry: Whether to enable retry with exponential backoff retry_attempts: Number of retry attempts stream: Whether to stream the response @@ -1382,7 +1382,7 @@ async def honcho_llm_call( gpt5_verbosity = verbosity # Filter out incompatible parameters when using backup - if provider != "anthropic" and thinking_budget: + if provider not in ("anthropic", "google") and thinking_budget: logger.warning( f"thinking_budget_tokens not supported by {provider}, ignoring" ) @@ -1565,6 +1565,68 @@ async def honcho_llm_call( return result +def _repair_response_model_json( + raw_content: str, + response_model: type[BaseModel], + model: str, +) -> BaseModel: + """Attempt to repair truncated/malformed JSON and validate against response_model. + + Used by all provider paths when structured output parsing fails. + For PromptRepresentation, falls back to an empty instance. + For other models, re-raises ValidationError. + """ + try: + final = validate_and_repair_json(raw_content) + repaired_data = json.loads(final) + + # Schema-aware repair for PromptRepresentation + if ( + response_model is PromptRepresentation + and "deductive" in repaired_data + and isinstance(repaired_data["deductive"], list) + ): + for i, item in enumerate(repaired_data["deductive"]): + if isinstance(item, dict): + if "conclusion" not in item and "premises" in item: + logger.warning( + f"Deductive observation {i} missing conclusion, adding placeholder" + ) + if item["premises"]: + item["conclusion"] = ( + f"[Incomplete reasoning from premises: {item['premises'][0][:100]}...]" + ) + else: + item["conclusion"] = ( + "[Incomplete reasoning - conclusion missing]" + ) + if "premises" not in item: + item["premises"] = [] + + final = json.dumps(repaired_data) + except (json.JSONDecodeError, KeyError, TypeError, ValueError) as repair_err: + final = "" + logger.warning( + f"Could not perform JSON repair on truncated output from {model}: {repair_err}" + ) + + try: + return response_model.model_validate_json(final) + except ValidationError as ve: + logger.error( + f"Validation error after repair of truncated output from {model}: {ve}" + ) + logger.debug(f"Problematic JSON: {final}") + + if response_model is PromptRepresentation: + logger.warning( + "Using fallback empty Representation due to truncated output" + ) + return PromptRepresentation(explicit=[]) + else: + raise + + @overload async def honcho_llm_call_inner( provider: SupportedProviders, @@ -1578,7 +1640,7 @@ async def honcho_llm_call_inner( reasoning_effort: Literal["low", "medium", "high", "minimal"] | None = None, # OpenAI only verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only - thinking_budget_tokens: int | None = None, # Anthropic only + thinking_budget_tokens: int | None = None, # Anthropic / Gemini stream: Literal[False] = False, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -1599,7 +1661,7 @@ async def honcho_llm_call_inner( reasoning_effort: Literal["low", "medium", "high", "minimal"] | None = None, # OpenAI only verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only - thinking_budget_tokens: int | None = None, # Anthropic only + thinking_budget_tokens: int | None = None, # Anthropic / Gemini stream: Literal[False] = False, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -1620,7 +1682,7 @@ async def honcho_llm_call_inner( reasoning_effort: Literal["low", "medium", "high", "minimal"] | None = None, # OpenAI only verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only - thinking_budget_tokens: int | None = None, # Anthropic only + thinking_budget_tokens: int | None = None, # Anthropic / Gemini stream: Literal[True] = ..., tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -1640,7 +1702,7 @@ async def honcho_llm_call_inner( reasoning_effort: Literal["low", "medium", "high", "minimal"] | None = None, # OpenAI only verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only - thinking_budget_tokens: int | None = None, # Anthropic only + thinking_budget_tokens: int | None = None, # Anthropic / Gemini stream: bool = False, tools: list[dict[str, Any]] | None = None, tool_choice: str | dict[str, Any] | None = None, @@ -1829,10 +1891,26 @@ async def honcho_llm_call_inner( thinking_content=thinking_content, thinking_blocks=thinking_full_blocks, ) - except (json.JSONDecodeError, ValidationError, ValueError) as e: - raise ValueError( - f"Failed to parse Anthropic response as {response_model}: {e}. Raw content: {text_content}" - ) from e + except (json.JSONDecodeError, ValidationError, ValueError): + # Attempt JSON repair on truncated/malformed output + logger.warning( + f"Anthropic response_model parse failed for {model}, attempting JSON repair" + ) + raw_content = "{" + text_content + repaired_obj = _repair_response_model_json( + raw_content, response_model, model + ) + return HonchoLLMCallResponse( + content=repaired_obj, + input_tokens=total_input_tokens, + output_tokens=usage.output_tokens if usage else 0, + cache_creation_input_tokens=cache_creation_tokens, + cache_read_input_tokens=cache_read_tokens, + finish_reasons=[stop_reason] if stop_reason else [], + tool_calls_made=[], + thinking_content=thinking_content, + thinking_blocks=thinking_full_blocks, + ) return HonchoLLMCallResponse( content=text_content, @@ -1989,9 +2067,39 @@ async def honcho_llm_call_inner( ) elif response_model: openai_params["response_format"] = response_model - response: ChatCompletion = await client.chat.completions.parse( # pyright: ignore - **openai_params - ) + try: + response: ChatCompletion = await client.chat.completions.parse( # pyright: ignore + **openai_params + ) + except LengthFinishReasonError as e: + # The LLM hit the max token limit before completing valid JSON. + # Extract the truncated content and attempt repair. + logger.warning( + f"LengthFinishReasonError for {model}: attempting JSON repair on truncated output" + ) + truncated_completion = e.completion + raw_content = truncated_completion.choices[0].message.content or "" + usage = truncated_completion.usage + finish_reason = truncated_completion.choices[0].finish_reason + + repaired_obj = _repair_response_model_json( + raw_content, response_model, model + ) + + cache_creation, cache_read = extract_openai_cache_tokens(usage) + return HonchoLLMCallResponse( + content=repaired_obj, + input_tokens=usage.prompt_tokens if usage else 0, + output_tokens=usage.completion_tokens if usage else 0, + cache_creation_input_tokens=cache_creation, + cache_read_input_tokens=cache_read, + finish_reasons=[finish_reason] if finish_reason else [], + tool_calls_made=[], + thinking_content=extract_openai_reasoning_content( + truncated_completion + ), + ) + # Extract the parsed object for structured output parsed_content = response.choices[0].message.parsed if parsed_content is None: @@ -2105,6 +2213,12 @@ async def honcho_llm_call_inner( } } + # Add thinking config for Gemini 2.5 models + if thinking_budget_tokens: + gemini_config["thinking_config"] = { + "thinking_budget": thinking_budget_tokens, + } + if response_model is None: if json_mode and not tools: gemini_config["response_mime_type"] = "application/json" @@ -2273,14 +2387,31 @@ async def honcho_llm_call_inner( finish_reason=finish_reason, ) - # Validate that parsed content matches the response model - if not isinstance(gemini_response.parsed, response_model): - raise ValueError( - f"Parsed content does not match the response model: {gemini_response.parsed} != {response_model}" + # If parsed content is valid and matches model, return directly + if isinstance(gemini_response.parsed, response_model): + return HonchoLLMCallResponse( + content=gemini_response.parsed, + input_tokens=input_token_count, + output_tokens=output_token_count, + finish_reasons=[finish_reason], + tool_calls_made=[], ) + # Parsed content missing or wrong type — attempt JSON repair on raw text + logger.warning( + f"Gemini response_model parse failed for {model} (finish_reason={finish_reason}), attempting JSON repair" + ) + raw_text = "" + if gemini_response.candidates and gemini_response.candidates[0].content: + for part in gemini_response.candidates[0].content.parts or []: + if hasattr(part, "text") and part.text: + raw_text += part.text + + repaired_obj = _repair_response_model_json( + raw_text, response_model, model + ) return HonchoLLMCallResponse( - content=gemini_response.parsed, + content=repaired_obj, input_tokens=input_token_count, output_tokens=output_token_count, finish_reasons=[finish_reason], @@ -2329,10 +2460,24 @@ async def honcho_llm_call_inner( finish_reasons=[finish_reason] if finish_reason else [], tool_calls_made=[], ) - except (json.JSONDecodeError, ValidationError, ValueError) as e: - raise ValueError( - f"Failed to parse Groq response as {response_model}: {e}. Raw content: {response.choices[0].message.content}" # pyright: ignore - ) from e + except (json.JSONDecodeError, ValidationError, ValueError): + # Attempt JSON repair on truncated/malformed output + logger.warning( + f"Groq response_model parse failed for {model}, attempting JSON repair" + ) + raw_content = str(response.choices[0].message.content or "") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] + repaired_obj = _repair_response_model_json( + raw_content, response_model, model + ) + return HonchoLLMCallResponse( + content=repaired_obj, + input_tokens=usage.prompt_tokens if usage else 0, # pyright: ignore + output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore + cache_creation_input_tokens=cache_creation, + cache_read_input_tokens=cache_read, + finish_reasons=[finish_reason] if finish_reason else [], + tool_calls_made=[], + ) else: return HonchoLLMCallResponse( content=response.choices[0].message.content, # pyright: ignore @@ -2361,7 +2506,7 @@ async def handle_streaming_response( client: The LLM client instance params: Request parameters including stream=True json_mode: Whether to use JSON mode - thinking_budget_tokens: Anthropic thinking budget tokens + thinking_budget_tokens: Anthropic / Gemini thinking budget tokens response_model: Pydantic model for structured output reasoning_effort: OpenAI reasoning effort level (GPT-5 only) verbosity: OpenAI verbosity level (GPT-5 only) @@ -2492,6 +2637,12 @@ async def handle_streaming_response( "max_output_tokens": cast(int, params["max_tokens"]), } + # Add thinking config for Gemini 2.5 models (streaming) + if thinking_budget_tokens: + stream_config["thinking_config"] = { + "thinking_budget": thinking_budget_tokens, + } + if response_model is not None: stream_config["response_mime_type"] = "application/json" stream_config["response_schema"] = response_model diff --git a/tests/utils/test_length_finish_reason.py b/tests/utils/test_length_finish_reason.py new file mode 100644 index 00000000..e50b19b7 --- /dev/null +++ b/tests/utils/test_length_finish_reason.py @@ -0,0 +1,454 @@ +""" +Tests for JSON repair handling across all providers in honcho_llm_call_inner, +and Gemini thinking budget support. + +Verifies that when an LLM hits the max token limit or returns malformed JSON, +the truncated output is repaired and returned instead of crashing. +""" + +import json +from typing import Any +from unittest.mock import AsyncMock, Mock, patch + +import pytest +from anthropic import AsyncAnthropic +from anthropic.types import TextBlock, Usage +from openai import AsyncOpenAI, LengthFinishReasonError +from openai.types.chat import ChatCompletion +from openai.types.chat.chat_completion import Choice +from openai.types.chat.chat_completion_message import ChatCompletionMessage +from openai.types.completion_usage import CompletionUsage +from pydantic import BaseModel, ValidationError + +from src.utils.clients import CLIENTS, HonchoLLMCallResponse, honcho_llm_call_inner +from src.utils.representation import PromptRepresentation + +# --- Test models --- + + +class SimpleModel(BaseModel): + """Non-PromptRepresentation model for testing re-raise behavior.""" + + items: list[str] + + +# --- Helpers --- + +VALID_REPR_JSON = { + "explicit": [ + {"content": "hermes is 25 years old"}, + {"content": "hermes has a dog"}, + ] +} + + +def _make_truncated_completion(content: str) -> ChatCompletion: + """Build a ChatCompletion with finish_reason='length' and the given content.""" + return ChatCompletion( + id="test-truncated", + object="chat.completion", + created=1234567890, + model="test-model", + choices=[ + Choice( + index=0, + message=ChatCompletionMessage(role="assistant", content=content), + finish_reason="length", + ) + ], + usage=CompletionUsage( + prompt_tokens=1000, completion_tokens=2000, total_tokens=3000 + ), + ) + + +def _raise_length_error(content: str) -> AsyncMock: + """Return an AsyncMock that raises LengthFinishReasonError with truncated content.""" + completion = _make_truncated_completion(content) + return AsyncMock(side_effect=LengthFinishReasonError(completion=completion)) + + +def _make_anthropic_mock(text: str, stop_reason: str = "end_turn") -> AsyncMock: + """Build a mocked AsyncAnthropic client returning the given text.""" + mock_client = AsyncMock(spec=AsyncAnthropic) + mock_response = Mock() + mock_response.content = [TextBlock(text=text, type="text")] + mock_response.usage = Usage(input_tokens=100, output_tokens=50) + mock_response.stop_reason = stop_reason + mock_client.messages.create = AsyncMock(return_value=mock_response) + return mock_client + + +def _make_gemini_mock( + text: str | None = None, + parsed: Any = None, + finish_reason_name: str = "STOP", +) -> Mock: + """Build a mocked genai.Client returning the given text/parsed content.""" + mock_client = Mock() + + # Build response + mock_response = Mock() + mock_response.parsed = parsed + + # Candidates + mock_candidate = Mock() + mock_finish_reason = Mock() + mock_finish_reason.name = finish_reason_name + mock_candidate.finish_reason = mock_finish_reason + + # Content parts + if text is not None: + mock_part = Mock() + mock_part.text = text + mock_part.function_call = None + mock_content = Mock() + mock_content.parts = [mock_part] + mock_candidate.content = mock_content + else: + mock_candidate.content = None + + mock_response.candidates = [mock_candidate] + + # Usage + mock_usage = Mock() + mock_usage.prompt_token_count = 200 + mock_usage.candidates_token_count = 100 + mock_response.usage_metadata = mock_usage + + mock_client.aio.models.generate_content = AsyncMock(return_value=mock_response) + return mock_client + + +# --------------------------------------------------------------------------- +# OpenAI / Custom provider tests (LengthFinishReasonError path) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestOpenAILengthFinishReasonRepair: + """Tests that LengthFinishReasonError is caught and truncated JSON is repaired.""" + + async def test_truncated_prompt_representation_repaired_openai(self) -> None: + """Truncated but repairable PromptRepresentation JSON should be repaired (openai).""" + truncated_json = json.dumps(VALID_REPR_JSON)[:-2] + + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error(truncated_json) + + with patch.dict(CLIENTS, {"openai": mock_client}): + response = await honcho_llm_call_inner( + provider="openai", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response, HonchoLLMCallResponse) + assert isinstance(response.content, PromptRepresentation) + assert len(response.content.explicit) >= 1 + assert response.finish_reasons == ["length"] + assert response.output_tokens == 2000 + + async def test_truncated_prompt_representation_repaired_custom(self) -> None: + """Truncated but repairable PromptRepresentation JSON should be repaired (custom).""" + truncated_json = json.dumps(VALID_REPR_JSON)[:-2] + + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error(truncated_json) + + with patch.dict(CLIENTS, {"custom": mock_client}): + response = await honcho_llm_call_inner( + provider="custom", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response, HonchoLLMCallResponse) + assert isinstance(response.content, PromptRepresentation) + assert len(response.content.explicit) >= 1 + assert response.finish_reasons == ["length"] + + async def test_completely_broken_json_falls_back_to_empty(self) -> None: + """Completely unrepairable JSON should fall back to empty PromptRepresentation.""" + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error( + "this is not json at all just random text" + ) + + with patch.dict(CLIENTS, {"openai": mock_client}): + response = await honcho_llm_call_inner( + provider="openai", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert response.content.explicit == [] + assert response.finish_reasons == ["length"] + + async def test_empty_content_falls_back_to_empty(self) -> None: + """Empty/null content should fall back to empty PromptRepresentation.""" + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error("") + + with patch.dict(CLIENTS, {"custom": mock_client}): + response = await honcho_llm_call_inner( + provider="custom", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert response.content.explicit == [] + + async def test_non_prompt_representation_reraises_on_unfixable(self) -> None: + """Non-PromptRepresentation with unrepairable JSON should raise ValidationError.""" + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error("not json") + + with ( + patch.dict(CLIENTS, {"openai": mock_client}), + pytest.raises(ValidationError), + ): + await honcho_llm_call_inner( + provider="openai", + model="test-model", + prompt="Generate items", + max_tokens=2000, + response_model=SimpleModel, + json_mode=True, + ) + + async def test_token_counts_preserved(self) -> None: + """Token counts from the truncated completion should be preserved.""" + truncated_json = '{"explicit": [{"content": "fact one"}' + + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error(truncated_json) + + with patch.dict(CLIENTS, {"openai": mock_client}): + response = await honcho_llm_call_inner( + provider="openai", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert response.input_tokens == 1000 + assert response.output_tokens == 2000 + + async def test_valid_json_with_length_finish_reason(self) -> None: + """Valid JSON despite length truncation should parse fine.""" + valid_json = json.dumps(VALID_REPR_JSON) + + mock_client = AsyncMock(spec=AsyncOpenAI) + mock_client.chat.completions.parse = _raise_length_error(valid_json) + + with patch.dict(CLIENTS, {"custom": mock_client}): + response = await honcho_llm_call_inner( + provider="custom", + model="test-model", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert len(response.content.explicit) == 2 + assert response.content.explicit[0].content == "hermes is 25 years old" + + +# --------------------------------------------------------------------------- +# Anthropic provider tests (JSON parse failure -> repair path) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestAnthropicJsonRepair: + """Tests that Anthropic response_model parse failures trigger JSON repair.""" + + async def test_truncated_anthropic_response_repaired(self) -> None: + """Truncated Anthropic JSON response should be repaired.""" + # Anthropic prefills "{" so the response text starts after that + # The code prepends "{" back: json_content = "{" + text_content + truncated_text = json.dumps(VALID_REPR_JSON)[ + 1:-2 + ] # Remove leading { and trailing }] + + mock_client = _make_anthropic_mock(truncated_text, stop_reason="max_tokens") + + with patch.dict(CLIENTS, {"anthropic": mock_client}): + response = await honcho_llm_call_inner( + provider="anthropic", + model="claude-3-sonnet", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert len(response.content.explicit) >= 1 + + async def test_broken_anthropic_response_falls_back_to_empty(self) -> None: + """Completely broken Anthropic JSON should fall back to empty PromptRepresentation.""" + mock_client = _make_anthropic_mock( + "random gibberish that is not json", stop_reason="max_tokens" + ) + + with patch.dict(CLIENTS, {"anthropic": mock_client}): + response = await honcho_llm_call_inner( + provider="anthropic", + model="claude-3-sonnet", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert response.content.explicit == [] + + async def test_non_prompt_representation_reraises(self) -> None: + """Non-PromptRepresentation with broken JSON should raise.""" + mock_client = _make_anthropic_mock("not json", stop_reason="max_tokens") + + with ( + patch.dict(CLIENTS, {"anthropic": mock_client}), + pytest.raises(ValidationError), + ): + await honcho_llm_call_inner( + provider="anthropic", + model="claude-3-sonnet", + prompt="Generate items", + max_tokens=2000, + response_model=SimpleModel, + json_mode=True, + ) + + +# --------------------------------------------------------------------------- +# Gemini provider tests (parsed=None or type mismatch -> repair path) +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestGeminiJsonRepair: + """Tests that Gemini response_model parse failures trigger JSON repair.""" + + async def test_gemini_unparsed_response_repaired(self) -> None: + """Gemini returning text but no parsed object should repair from raw text.""" + from google import genai + + valid_text = json.dumps(VALID_REPR_JSON) + mock_client = _make_gemini_mock( + text=valid_text, parsed=None, finish_reason_name="MAX_TOKENS" + ) + + with ( + patch.dict(CLIENTS, {"google": mock_client}), + patch.object(genai.Client, "__instancecheck__", return_value=True), + ): + # We need the match statement to hit the genai.Client case + mock_client.__class__ = genai.Client # pyright: ignore[reportAttributeAccessIssue] + response = await honcho_llm_call_inner( + provider="google", + model="gemini-2.5-flash", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert len(response.content.explicit) == 2 + + async def test_gemini_broken_text_falls_back_to_empty(self) -> None: + """Gemini with broken text and no parsed content should fall back.""" + from google import genai + + mock_client = _make_gemini_mock( + text="broken json", parsed=None, finish_reason_name="MAX_TOKENS" + ) + mock_client.__class__ = genai.Client # pyright: ignore[reportAttributeAccessIssue] + + with patch.dict(CLIENTS, {"google": mock_client}): + response = await honcho_llm_call_inner( + provider="google", + model="gemini-2.5-flash", + prompt="Analyze messages", + max_tokens=2000, + response_model=PromptRepresentation, + json_mode=True, + ) + + assert isinstance(response.content, PromptRepresentation) + assert response.content.explicit == [] + + +# --------------------------------------------------------------------------- +# Gemini thinking budget tests +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +class TestGeminiThinkingBudget: + """Tests that thinking_budget_tokens is passed to Gemini via ThinkingConfig.""" + + async def test_thinking_budget_passed_to_gemini(self) -> None: + """thinking_budget_tokens should be included in Gemini config.""" + from google import genai + + mock_client = _make_gemini_mock(text="Hello", parsed=None) + mock_client.__class__ = genai.Client # pyright: ignore[reportAttributeAccessIssue] + + with patch.dict(CLIENTS, {"google": mock_client}): + await honcho_llm_call_inner( + provider="google", + model="gemini-2.5-flash", + prompt="Think about this", + max_tokens=2000, + thinking_budget_tokens=4096, + ) + + # Verify generate_content was called with thinking_config + call_args = mock_client.aio.models.generate_content.call_args + config = call_args.kwargs.get("config") or call_args[1].get("config") + assert config is not None + assert "thinking_config" in config + assert config["thinking_config"]["thinking_budget"] == 4096 + + async def test_no_thinking_config_when_budget_is_none(self) -> None: + """When thinking_budget_tokens is None, thinking_config should not be set.""" + from google import genai + + mock_client = _make_gemini_mock(text="Hello", parsed=None) + mock_client.__class__ = genai.Client # pyright: ignore[reportAttributeAccessIssue] + + with patch.dict(CLIENTS, {"google": mock_client}): + await honcho_llm_call_inner( + provider="google", + model="gemini-2.5-flash", + prompt="No thinking needed", + max_tokens=2000, + ) + + call_args = mock_client.aio.models.generate_content.call_args + config = call_args.kwargs.get("config") or call_args[1].get("config") + if config: + assert "thinking_config" not in config