From bc8ddcb9369fb702d63adb086f9ca8fadca13769 Mon Sep 17 00:00:00 2001 From: Jean Paul Elisa NIYOKWIZERWA <140616733+Ndevu12@users.noreply.github.com> Date: Tue, 1 Sep 2026 10:35:14 +0000 Subject: [PATCH] feat: enforce LLM output schemas natively via structured outputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Attach pydantic schemas to agents (pydantic-ai output_type) instead of asking for JSON in prose and repairing whatever comes back. The framework validates responses and retries with validation feedback on schema violations; callers receive typed objects or fall back to the existing heuristics. - New src/models/structured.py: create_structured_agent / run_structured / try_run_structured, plus slim role prompts — the 'respond with ONLY JSON' scaffolding is unnecessary when the schema is enforced by the framework - Converted call sites: per-paper extraction, collective synthesis, gap analysis, and LLM query expansion - Extraction pins identity fields (paper ID, title) from retrieval metadata so model output can never relabel a paper, and backfills evidence from grounded passages when the model omits it - Circuit breaker and heuristic fallbacks behave identically on the structured path - Config-gated by llm.structured_outputs (default true); disabling it routes through the legacy prose-JSON repair path unchanged, kept for backends that cannot honor response schemas - Legacy-path tests now pin structured_outputs=false explicitly; new tests cover the structured path end to end using pydantic-ai's TestModel and FunctionModel (no network) - .env.example, the environment-variables reference, and the heuristic-vs-LLM page document the new flag --- .env.example | 5 + config/default.yaml | 1 + docs/configuration/environment-variables.md | 1 + docs/llm/heuristic-vs-llm.md | 12 + src/analysis/gap_analysis.py | 28 ++- src/analysis/synthesis.py | 118 +++++++--- src/config/settings.py | 1 + src/models/__init__.py | 10 + src/models/structured.py | 116 ++++++++++ src/research/query_expansion.py | 27 ++- tests/test_structured_outputs.py | 229 ++++++++++++++++++++ tests/test_synthesis.py | 25 ++- 12 files changed, 535 insertions(+), 38 deletions(-) create mode 100644 src/models/structured.py create mode 100644 tests/test_structured_outputs.py diff --git a/.env.example b/.env.example index 0f1c43c..03fae70 100644 --- a/.env.example +++ b/.env.example @@ -46,6 +46,11 @@ RA_LLM__MODEL=auto # With or without /v1 — Ollama providers normalize via normalize_openai_base_url() RA_LLM__BASE_URL=http://localhost:11434/v1 RA_LLM__API_KEY=ollama + +# Schema-enforced LLM output (native structured outputs). Disable only for +# backends that cannot honor response schemas; the legacy JSON-repair path +# is then used instead. +# RA_LLM__STRUCTURED_OUTPUTS=false # OLLAMA_API_KEY=ollama # Synthesis — see config/ollama_models.yaml for per-model hints diff --git a/config/default.yaml b/config/default.yaml index 329f4f0..00d9182 100644 --- a/config/default.yaml +++ b/config/default.yaml @@ -4,6 +4,7 @@ llm: base_url: http://localhost:11434 temperature: 0.2 timeout_seconds: 120 + structured_outputs: true embedding: model: BAAI/bge-small-en-v1.5 diff --git a/docs/configuration/environment-variables.md b/docs/configuration/environment-variables.md index 23e6a89..4533b52 100644 --- a/docs/configuration/environment-variables.md +++ b/docs/configuration/environment-variables.md @@ -30,6 +30,7 @@ Boolean env values accept standard truthy strings (`true`, `1`, `yes`). | Variable | Default | Description | |----------|---------|-------------| | `RA_LLM__PROVIDER` | `ollama` | `ollama`, `openai`, or `anthropic` | +| `RA_LLM__STRUCTURED_OUTPUTS` | `true` | Schema-enforced LLM output; disable to use the legacy JSON-repair path | | `RA_LLM__MODEL` | `auto` | Model name; `auto` selects from `config/ollama_models.yaml` (Ollama only) | | `RA_LLM__BASE_URL` | `http://localhost:11434` | API base URL (Ollama OpenAI-compatible endpoint) | diff --git a/docs/llm/heuristic-vs-llm.md b/docs/llm/heuristic-vs-llm.md index 2cf0153..57386af 100644 --- a/docs/llm/heuristic-vs-llm.md +++ b/docs/llm/heuristic-vs-llm.md @@ -124,6 +124,18 @@ RA_QUERY_EXPANSION__LLM_ENABLED=true YAML equivalents under `synthesis:` and `query_expansion:` in `config/default.yaml`. See [Stage toggles](../configuration/stage-toggles.md). +## Structured outputs + +When the LLM path runs, output schemas are enforced natively by default +(`RA_LLM__STRUCTURED_OUTPUTS=true`): the pydantic model is attached to the +agent (`output_type`), pydantic-ai validates the response and retries with +validation feedback on schema violations, and the caller receives a typed +object or falls back to heuristics. Disable the flag to route through the +legacy prose-JSON path with the repair/retry machinery in `src/utils/` — +useful only for backends that cannot honor response schemas. Identity +fields (paper ID, title) are always taken from retrieval metadata, never +from model output. + ## Verify LLM is active Check pipeline logs or debug JSON in `logs/debug/`: diff --git a/src/analysis/gap_analysis.py b/src/analysis/gap_analysis.py index db7dccd..9d4744c 100644 --- a/src/analysis/gap_analysis.py +++ b/src/analysis/gap_analysis.py @@ -109,13 +109,6 @@ async def analyze_gaps( llm_config = get_settings().llm - agent = AgentFactory(llm_config).create_agent(AgentRole.GAP_ANALYSIS, config=llm_config) - response_handler = handler or EnhancedResponseHandler(ResponseHandlerConfig()) - context = RequestContext( - user_query=query, - model_name=llm_config.model, - session_id=session_id, - ) prompt = _build_gap_prompt(query, synthesis, clusters) from ..utils.progress_reporter import get_progress_reporter @@ -124,6 +117,27 @@ async def analyze_gaps( if reporter is not None: reporter.set_activity("Analyzing research gaps with AI…") + if llm_config.structured_outputs: + from ..models.structured import try_run_structured + + gap_result = await try_run_structured( + AgentRole.GAP_ANALYSIS, + prompt, + GapAnalysisResult, + llm_config, + ) + if gap_result is not None: + return gap_result + return _heuristic_gap_analysis(synthesis, query, clusters) + + agent = AgentFactory(llm_config).create_agent(AgentRole.GAP_ANALYSIS, config=llm_config) + response_handler = handler or EnhancedResponseHandler(ResponseHandlerConfig()) + context = RequestContext( + user_query=query, + model_name=llm_config.model, + session_id=session_id, + ) + result = await response_handler.process_structured_response( agent, prompt, diff --git a/src/analysis/synthesis.py b/src/analysis/synthesis.py index e0aa140..249f3e9 100644 --- a/src/analysis/synthesis.py +++ b/src/analysis/synthesis.py @@ -343,6 +343,43 @@ def recover_synthesis_output( return synthesis +def _use_structured_outputs(llm_config: LLMConfig | None) -> bool: + if llm_config is not None: + return llm_config.structured_outputs + from ..config.settings import get_settings + + return get_settings().llm.structured_outputs + + +async def _extract_single_paper_structured( + paper: RankedPaper, + query: str, + llm_config: LLMConfig | None, + passages: list[str] | None = None, +) -> tuple[PaperExtraction, bool]: + """Schema-enforced extraction; falls back to heuristics on failure.""" + from ..models.structured import try_run_structured + + prompt = _build_extraction_prompt(paper, query, passages) + extraction = await try_run_structured( + AgentRole.EXTRACTION, + prompt, + PaperExtraction, + llm_config, + ) + if extraction is None: + return _heuristic_extraction(paper, passages), False + + # Identity fields come from our metadata, never from model output. + updates: dict[str, object] = { + "paper_id": paper.paper.paper_id, + "title": paper.paper.title, + } + if not extraction.evidence and passages: + updates["evidence"] = _evidence_from_passages(passages) + return extraction.model_copy(update=updates), True + + async def _extract_single_paper( paper: RankedPaper, query: str, @@ -412,15 +449,21 @@ async def extract_papers( llm_targets = ranked_papers[:max_llm_papers] heuristic_targets = ranked_papers[max_llm_papers:] - agent = AgentFactory(llm_config).create_agent(AgentRole.EXTRACTION, config=llm_config) - response_handler = handler or EnhancedResponseHandler( - _synthesis_handler_config(synthesis_config.extraction_max_retries) - ) - context = RequestContext( - user_query=query, - model_name=llm_config.model if llm_config else "default", - session_id=session_id, - ) + use_structured = _use_structured_outputs(llm_config) + agent = None + response_handler = None + context = None + if not use_structured: + # Legacy prose-JSON path with the repair/retry machinery. + agent = AgentFactory(llm_config).create_agent(AgentRole.EXTRACTION, config=llm_config) + response_handler = handler or EnhancedResponseHandler( + _synthesis_handler_config(synthesis_config.extraction_max_retries) + ) + context = RequestContext( + user_query=query, + model_name=llm_config.model if llm_config else "default", + session_id=session_id, + ) from ..utils.progress_reporter import get_progress_reporter @@ -447,14 +490,22 @@ async def _extract_with_circuit(index: int, paper: RankedPaper) -> PaperExtracti f"Analyzing paper {index}/{len(llm_targets)}: {title_preview}…" ) - extraction, llm_success = await _extract_single_paper( - paper, - query, - agent, - response_handler, - context, - passages=paper_passages, - ) + if use_structured: + extraction, llm_success = await _extract_single_paper_structured( + paper, + query, + llm_config, + passages=paper_passages, + ) + else: + extraction, llm_success = await _extract_single_paper( + paper, + query, + agent, + response_handler, + context, + passages=paper_passages, + ) if llm_success: breaker["consecutive_failures"] = 0 @@ -519,15 +570,6 @@ async def synthesize_collective( llm_config = get_settings().llm - agent = AgentFactory(llm_config).create_agent(AgentRole.SYNTHESIS, config=llm_config) - response_handler = handler or EnhancedResponseHandler( - _synthesis_handler_config(synthesis_config.collective_max_retries) - ) - context = RequestContext( - user_query=query, - model_name=llm_config.model, - session_id=session_id, - ) prompt = _build_synthesis_prompt(query, extractions, clusters, ranked_papers) logger.info("Running LLM collective synthesis across %d paper(s)", len(extractions)) @@ -539,6 +581,30 @@ async def synthesize_collective( f"Synthesizing insights across {len(extractions)} paper(s)…" ) + if _use_structured_outputs(llm_config): + from ..models.structured import try_run_structured + + synthesis = await try_run_structured( + AgentRole.SYNTHESIS, + prompt, + SynthesisResult, + llm_config, + ) + if synthesis is not None: + return synthesis + logger.warning("LLM collective synthesis failed; using heuristic fallback") + return _heuristic_synthesis(query, extractions, clusters, ranked_papers=ranked_papers) + + agent = AgentFactory(llm_config).create_agent(AgentRole.SYNTHESIS, config=llm_config) + response_handler = handler or EnhancedResponseHandler( + _synthesis_handler_config(synthesis_config.collective_max_retries) + ) + context = RequestContext( + user_query=query, + model_name=llm_config.model, + session_id=session_id, + ) + result = await response_handler.process_structured_response( agent, prompt, diff --git a/src/config/settings.py b/src/config/settings.py index cc212f7..2f67efc 100644 --- a/src/config/settings.py +++ b/src/config/settings.py @@ -66,6 +66,7 @@ class LLMConfig(BaseModel): api_key: str | None = None temperature: float = 0.2 timeout_seconds: int = 120 + structured_outputs: bool = True class EmbeddingConfig(BaseModel): diff --git a/src/models/__init__.py b/src/models/__init__.py index 1658c28..31852b7 100644 --- a/src/models/__init__.py +++ b/src/models/__init__.py @@ -2,6 +2,12 @@ """LLM provider abstraction (not paper domain models).""" from .base import ROLE_SYSTEM_PROMPTS, AgentRole, LLMProvider, normalize_openai_base_url +from .structured import ( + STRUCTURED_ROLE_PROMPTS, + create_structured_agent, + run_structured, + try_run_structured, +) from .factory import ( AgentFactory, create_llm_agent, @@ -14,6 +20,10 @@ from .anthropic import AnthropicProviderImpl __all__ = [ + "STRUCTURED_ROLE_PROMPTS", + "create_structured_agent", + "run_structured", + "try_run_structured", "AgentFactory", "AgentRole", "AnthropicProviderImpl", diff --git a/src/models/structured.py b/src/models/structured.py new file mode 100644 index 0000000..ebf9dcf --- /dev/null +++ b/src/models/structured.py @@ -0,0 +1,116 @@ +# -*- coding: utf-8 -*- +"""Schema-enforced LLM calls via pydantic-ai structured outputs. + +Instead of asking the model for JSON in prose and repairing whatever comes +back, the schema is attached to the agent (``output_type``) and pydantic-ai +enforces it natively — invalid output triggers the framework's own retry +with validation feedback, and the caller receives a typed object or an +exception. This replaces the string-repair path for providers that support +structured outputs; heuristic fallbacks remain the safety net. +""" + +from __future__ import annotations + +from typing import TypeVar + +from pydantic import BaseModel +from pydantic_ai import Agent +from pydantic_ai.models import Model + +from ..utils.logging_system import logger +from .base import ROLE_SYSTEM_PROMPTS, AgentRole +from .factory import AgentFactory, create_llm_provider + +OutputT = TypeVar("OutputT", bound=BaseModel) + +# Slim prompts for the structured path: the schema is enforced by the +# framework, so the "respond with ONLY JSON ..." scaffolding the legacy +# prompts carry is unnecessary and only adds noise. +STRUCTURED_ROLE_PROMPTS: dict[AgentRole, str] = { + AgentRole.EXPANSION: ( + "You expand research queries into concise academic search phrases " + "and focused sub-questions." + ), + AgentRole.EXTRACTION: ( + "You are a research paper analyst. Extract the paper's methodology, " + "datasets, benchmarks, limitations, and findings. When full-text " + "passages are provided, ground the extraction in them and quote short " + "verbatim evidence snippets; otherwise leave evidence empty." + ), + AgentRole.SYNTHESIS: ( + "You are a research synthesis expert. From structured per-paper " + "extractions and thematic clusters, derive cross-paper agreements, " + "disagreements, trends, gaps, datasets, and methodologies." + ), + AgentRole.GAP_ANALYSIS: ( + "You are a research strategist. From cross-paper synthesis, identify " + "prioritized research gaps, actionable opportunities, and " + "underexplored areas." + ), +} + + +def create_structured_agent( + role: AgentRole, + output_type: type[OutputT], + llm_config=None, + *, + retries: int = 2, + model: Model | None = None, +) -> Agent: + """Create an agent whose output is validated against ``output_type``.""" + if model is None: + resolved = AgentFactory(llm_config).config + model = create_llm_provider(resolved).create_model(resolved) + prompt = STRUCTURED_ROLE_PROMPTS.get(role) or ROLE_SYSTEM_PROMPTS[role] + return Agent( + model=model, + system_prompt=prompt, + output_type=output_type, + retries=retries, + ) + + +async def run_structured( + role: AgentRole, + prompt: str, + output_type: type[OutputT], + llm_config=None, + *, + retries: int = 2, + model: Model | None = None, +) -> OutputT: + """Run a structured agent and return the validated output. Raises on failure.""" + agent = create_structured_agent( + role, + output_type, + llm_config, + retries=retries, + model=model, + ) + result = await agent.run(prompt) + return result.output + + +async def try_run_structured( + role: AgentRole, + prompt: str, + output_type: type[OutputT], + llm_config=None, + *, + retries: int = 2, + model: Model | None = None, +) -> OutputT | None: + """Like :func:`run_structured`, but returns None on any failure.""" + try: + return await run_structured( + role, + prompt, + output_type, + llm_config, + retries=retries, + model=model, + ) + except Exception as exc: + logger.warning("Structured %s call failed: %s", role.value, exc) + return None diff --git a/src/research/query_expansion.py b/src/research/query_expansion.py index 831d74a..89810e9 100644 --- a/src/research/query_expansion.py +++ b/src/research/query_expansion.py @@ -7,10 +7,19 @@ import time from typing import TYPE_CHECKING +from pydantic import BaseModel, Field + from ..core.context import PipelineContext, StageResult from ..retrieval.models import ExpandedQuerySet, QueryUnderstandingResult from .text_utils import QUERY_STOP_WORDS, extract_core_concepts + +class _ExpansionSuggestions(BaseModel): + """Structured output for LLM query expansion.""" + + variants: list[str] = Field(default_factory=list) + sub_questions: list[str] = Field(default_factory=list) + if TYPE_CHECKING: from ..config.settings import QueryExpansionConfig @@ -271,7 +280,6 @@ async def expand_query_llm( from ..models import AgentFactory, AgentRole settings = get_settings() - agent = AgentFactory(settings.llm).create_agent(AgentRole.EXPANSION, config=settings.llm) prompt = ( f"Expand this research query into {config.max_variants} search variants " f"and {config.max_sub_questions} sub-questions: {query}" @@ -282,6 +290,23 @@ async def expand_query_llm( if reporter is not None: reporter.set_activity("Expanding query with AI…") + if settings.llm.structured_outputs: + from ..models.structured import run_structured + + suggestions = await run_structured( + AgentRole.EXPANSION, + prompt, + _ExpansionSuggestions, + settings.llm, + ) + return ( + [str(item) for item in suggestions.variants][: config.max_variants], + [str(item) for item in suggestions.sub_questions][ + : config.max_sub_questions + ], + ) + + agent = AgentFactory(settings.llm).create_agent(AgentRole.EXPANSION, config=settings.llm) raw_output = await stream_agent_text( agent, prompt, diff --git a/tests/test_structured_outputs.py b/tests/test_structured_outputs.py new file mode 100644 index 0000000..1b046f5 --- /dev/null +++ b/tests/test_structured_outputs.py @@ -0,0 +1,229 @@ +# -*- coding: utf-8 -*- +"""Tests for schema-enforced LLM calls (native structured outputs).""" + +from __future__ import annotations + +from unittest.mock import patch + +import pytest +from pydantic import BaseModel, Field +from pydantic_ai.models.test import TestModel + +from src.analysis.gap_analysis import analyze_gaps +from src.analysis.synthesis import extract_papers, synthesize_collective +from src.config.settings import AppSettings, LLMConfig +from src.models.base import AgentRole +from src.models.structured import ( + STRUCTURED_ROLE_PROMPTS, + run_structured, + try_run_structured, +) +from src.retrieval.models import ( + GapAnalysisResult, + PaperExtraction, + RankedPaper, + RetrievedPaper, + SynthesisResult, +) + + +class _Sample(BaseModel): + items: list[str] = Field(default_factory=list) + note: str = "" + + +def _ranked(title: str, abstract: str = "Findings here.") -> RankedPaper: + return RankedPaper( + paper=RetrievedPaper(title=title, abstract=abstract, provider="test"), + rank_score=0.8, + score_breakdown={}, + ) + + +def _structured_llm() -> LLMConfig: + return LLMConfig(structured_outputs=True) + + +class TestRunStructured: + async def test_returns_schema_valid_output_with_test_model(self) -> None: + result = await run_structured( + AgentRole.EXTRACTION, + "extract this", + _Sample, + model=TestModel(), + ) + + assert isinstance(result, _Sample) + + async def test_try_run_structured_returns_none_on_model_failure(self) -> None: + from pydantic_ai.models.function import FunctionModel + + def exploding(messages, info): + raise RuntimeError("backend down") + + result = await try_run_structured( + AgentRole.SYNTHESIS, + "synthesize", + _Sample, + model=FunctionModel(exploding), + ) + + assert result is None + + def test_all_llm_roles_have_slim_structured_prompts(self) -> None: + for role in ( + AgentRole.EXPANSION, + AgentRole.EXTRACTION, + AgentRole.SYNTHESIS, + AgentRole.GAP_ANALYSIS, + ): + prompt = STRUCTURED_ROLE_PROMPTS[role] + assert "JSON" not in prompt # schema is enforced natively + + +class TestStructuredExtraction: + async def test_extraction_uses_structured_path_and_pins_identity(self) -> None: + ranked = [_ranked("Real Title")] + hallucinated = PaperExtraction( + paper_id="made-up-id", + title="Hallucinated Title", + findings=["A finding"], + ) + + async def fake_structured(role, prompt, output_type, config=None, **kwargs): + assert role is AgentRole.EXTRACTION + return hallucinated + + with patch("src.models.structured.try_run_structured", side_effect=fake_structured): + extractions = await extract_papers( + ranked, + "query", + llm_config=_structured_llm(), + synthesis_config=AppSettings( + synthesis={"llm_enabled": True, "max_llm_papers": 1} + ).synthesis, + ) + + assert extractions[0].paper_id == ranked[0].paper.paper_id + assert extractions[0].title == "Real Title" + assert extractions[0].findings == ["A finding"] + + async def test_structured_failure_falls_back_and_trips_breaker(self) -> None: + ranked = [_ranked(f"Paper {i}") for i in range(4)] + + calls = {"count": 0} + + async def failing_structured(*args, **kwargs): + calls["count"] += 1 + return None + + with patch("src.models.structured.try_run_structured", side_effect=failing_structured): + extractions = await extract_papers( + ranked, + "query", + llm_config=_structured_llm(), + concurrency=1, + synthesis_config=AppSettings( + synthesis={ + "llm_enabled": True, + "max_llm_papers": 4, + "circuit_breaker_failures": 2, + "concurrency": 1, + } + ).synthesis, + ) + + assert len(extractions) == 4 + # Breaker opens after 2 failures; remaining papers skip the LLM. + assert calls["count"] == 2 + assert all("abstract" in e.methodology[0] for e in extractions) + + +class TestStructuredSynthesisAndGaps: + async def test_collective_synthesis_structured_path(self) -> None: + expected = SynthesisResult(agreements=["Agreement"], gaps=["Gap"]) + + async def fake_structured(role, prompt, output_type, config=None, **kwargs): + assert role is AgentRole.SYNTHESIS + assert output_type is SynthesisResult + return expected + + with patch("src.models.structured.try_run_structured", side_effect=fake_structured): + result = await synthesize_collective( + "query", + [PaperExtraction(paper_id="a", title="A", findings=["F"])], + [], + llm_config=_structured_llm(), + synthesis_config=AppSettings( + synthesis={"llm_enabled": True} + ).synthesis, + ) + + assert result is expected + + async def test_gap_analysis_structured_path(self) -> None: + expected = GapAnalysisResult(gaps=["G"], opportunities=["O"]) + + async def fake_structured(role, prompt, output_type, config=None, **kwargs): + assert role is AgentRole.GAP_ANALYSIS + return expected + + with patch("src.models.structured.try_run_structured", side_effect=fake_structured): + result = await analyze_gaps( + "query", + SynthesisResult(gaps=["G"]), + llm_config=_structured_llm(), + ) + + assert result is expected + + async def test_gap_analysis_structured_failure_uses_heuristics(self) -> None: + async def failing(*args, **kwargs): + return None + + with patch("src.models.structured.try_run_structured", side_effect=failing): + result = await analyze_gaps( + "query", + SynthesisResult(gaps=["Known gap"]), + llm_config=_structured_llm(), + ) + + assert "Known gap" in result.gaps + + +class TestStructuredExpansion: + async def test_expansion_uses_structured_output(self, monkeypatch) -> None: + from src.research.query_expansion import _ExpansionSuggestions, expand_query_llm + from src.config.settings import get_settings + + get_settings.cache_clear() + monkeypatch.setenv("RA_LLM__STRUCTURED_OUTPUTS", "true") + + async def fake_run(role, prompt, output_type, config=None, **kwargs): + assert role is AgentRole.EXPANSION + return _ExpansionSuggestions( + variants=["variant one", "variant two"], + sub_questions=["sub question"], + ) + + try: + with patch("src.models.structured.run_structured", side_effect=fake_run): + variants, sub_questions = await expand_query_llm( + "test query", + AppSettings( + query_expansion={ + "llm_enabled": True, + "max_variants": 2, + "max_sub_questions": 1, + } + ).query_expansion, + ) + finally: + get_settings.cache_clear() + + assert variants == ["variant one", "variant two"] + assert sub_questions == ["sub question"] + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/test_synthesis.py b/tests/test_synthesis.py index 2b45376..4ab2ebe 100644 --- a/tests/test_synthesis.py +++ b/tests/test_synthesis.py @@ -25,7 +25,7 @@ synthesize_collective, ) from src.core.stage_recovery import recover_stage_output -from src.config.settings import AppSettings +from src.config.settings import AppSettings, LLMConfig from src.core.context import PipelineContext from src.core.pipeline import ResearchPipeline from src.research.clustering import ClusteringStage @@ -79,6 +79,11 @@ def _mock_handler_failure() -> EnhancedResponseHandler: return handler +def _legacy_llm_config() -> LLMConfig: + """LLM config exercising the legacy prose-JSON handler path.""" + return LLMConfig(structured_outputs=False) + + def _llm_synthesis_config(**overrides: object): config = { "llm_enabled": True, @@ -295,6 +300,7 @@ async def test_extract_papers_uses_handler(self) -> None: extractions = await extract_papers( ranked, "test query", + llm_config=_legacy_llm_config(), handler=handler, concurrency=1, synthesis_config=_llm_synthesis_config(max_llm_papers=1), @@ -312,6 +318,7 @@ async def test_extract_papers_falls_back_on_handler_failure(self) -> None: extractions = await extract_papers( ranked, "test query", + llm_config=_legacy_llm_config(), handler=handler, concurrency=1, synthesis_config=_llm_synthesis_config(max_llm_papers=1), @@ -335,6 +342,7 @@ async def test_synthesize_collective_validates_schema(self) -> None: "transformers", extractions, [], + llm_config=_legacy_llm_config(), handler=handler, synthesis_config=_llm_synthesis_config(), ) @@ -367,6 +375,7 @@ async def test_run_synthesis_two_pass_workflow(self) -> None: "deep learning", ranked, clusters, + llm_config=_legacy_llm_config(), handler=handler, concurrency=2, synthesis_config=_llm_synthesis_config(max_llm_papers=2), @@ -388,7 +397,9 @@ async def test_analyze_gaps_uses_handler(self) -> None: ) handler = _mock_handler_success(expected) - result = await analyze_gaps("query", synthesis, handler=handler) + result = await analyze_gaps( + "query", synthesis, llm_config=_legacy_llm_config(), handler=handler + ) assert result.opportunities == ["Run systematic ablations"] @@ -397,7 +408,9 @@ async def test_analyze_gaps_heuristic_fallback(self) -> None: synthesis = SynthesisResult(gaps=["Understudied domain"]) handler = _mock_handler_failure() - result = await analyze_gaps("query", synthesis, handler=handler) + result = await analyze_gaps( + "query", synthesis, llm_config=_legacy_llm_config(), handler=handler + ) assert "Understudied domain" in result.gaps assert result.opportunities @@ -435,7 +448,10 @@ async def test_gap_analysis_stage(self) -> None: synthesis = SynthesisResult(gaps=["Gap one"]) ctx = PipelineContext.create( "gap stage", - AppSettings(synthesis={"llm_enabled": True}), + AppSettings( + synthesis={"llm_enabled": True}, + llm={"structured_outputs": False}, + ), ) ctx.set_artifact("paper_clusters", [PaperCluster(theme="T", paper_ids=["p1"])]) @@ -570,6 +586,7 @@ async def test_extract_papers_limits_llm_calls(self) -> None: extractions = await extract_papers( ranked, "query", + llm_config=_legacy_llm_config(), handler=handler, synthesis_config=settings.synthesis, )