From 8ac30ecdfc0af7ec7bf6a869e38f29b54764d541 Mon Sep 17 00:00:00 2001 From: matteomedioli Date: Thu, 3 Sep 2026 02:22:06 +0200 Subject: [PATCH] feat(llm): GeminiLLM declares supports_structured_output = True MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit GeminiLLM already supports structured output via response_schema/ response_mime_type, but never declared the capability flag — pipeline components checking it fell back to prompt-based JSON. Breaking: SimpleKGPipeline, SchemaFromTextExtractor, and LLMEntityRelationExtractor now enable structured output by default with GeminiLLM. Use use_structured_output=False to keep the previous behavior. Reimplements the intent of #573 (closed, badly stale) fresh on current main. Independent of #572/#621 — different concern, no conflict. --- CHANGELOG.md | 5 +++++ src/neo4j_graphrag/llm/google_genai_llm.py | 2 ++ tests/unit/llm/test_google_genai_llm.py | 10 ++++++++++ 3 files changed, 17 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e03021ba3..aefce4696 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,10 +2,15 @@ ## Next + ### Fixed - Fixed `BaseGeminiLLM`/`GeminiLLM` silently dropping token usage on every call: `invoke`/`ainvoke` (both the string and message-list paths, sync and async) built `LLMResponse` from `response.text` alone, ignoring `response.usage_metadata` entirely. `LLMResponse.usage` is now populated from it, matching `AnthropicLLM`/`OpenAILLM`/`VertexAILLM`. +### Changed + +- (**breaking**) `BaseGeminiLLM`/`GeminiLLM.supports_structured_output` is now `True` — it already supported structured output via `response_schema`/`response_mime_type`, but never declared the capability flag, so pipeline components checking it fell back to prompt-based JSON. `SimpleKGPipeline`, `SchemaFromTextExtractor`, and `LLMEntityRelationExtractor` now enable structured output by default with `GeminiLLM`. To keep the previous behavior, construct the extractors with `use_structured_output=False`. + ## 1.19.0 ### Added diff --git a/src/neo4j_graphrag/llm/google_genai_llm.py b/src/neo4j_graphrag/llm/google_genai_llm.py index 202d4223b..d45ce902e 100644 --- a/src/neo4j_graphrag/llm/google_genai_llm.py +++ b/src/neo4j_graphrag/llm/google_genai_llm.py @@ -87,6 +87,8 @@ class BaseGeminiLLM(LLMBase, abc.ABC): constructing the ``client`` SDK instance. """ + supports_structured_output: bool = True + client: "genai.Client" def __init__( diff --git a/tests/unit/llm/test_google_genai_llm.py b/tests/unit/llm/test_google_genai_llm.py index 19fa2ef59..0e6dda06e 100644 --- a/tests/unit/llm/test_google_genai_llm.py +++ b/tests/unit/llm/test_google_genai_llm.py @@ -208,6 +208,16 @@ def test_gemini_llm_is_base_gemini_llm_subclass() -> None: assert issubclass(GeminiLLM, BaseGeminiLLM) +def test_gemini_llm_supports_structured_output( + mock_genai: Tuple[MagicMock, MagicMock], +) -> None: + """GeminiLLM already supports response_schema/response_mime_type; the + capability flag must say so, or SimpleKGPipeline/SchemaFromTextExtractor/ + LLMEntityRelationExtractor silently fall back to prompt-based JSON.""" + llm = GeminiLLM("gemini-2.0-flash") + assert llm.supports_structured_output is True + + def test_gemini_llm_init_only_constructs_client( mock_genai: Tuple[MagicMock, MagicMock], ) -> None: