"""PydanticAI tools and adapters for structured planning outputs."""
from __future__ import annotations
from collections.abc import Callable
from typing import Any, Optional
from lalandre_core.llm.structured import run_structured_agent
from pydantic_ai import Agent, ModelRetry
from ..prompts import render_intent_parser_prompt, render_planner_prompt
from .models import (
DecompositionOutput,
RetrievalEvaluationOutput,
RetrievalPlannerOutput,
RetrievalRefinementOutput,
RoutingIntentOutput,
)
_DECOMP_INSTRUCTIONS = (
"Tu es un assistant de décomposition de questions juridiques complexes.\n"
"Analyse la question et décompose-la en 2 ou 3 sous-questions indépendantes et précises.\n"
"Chaque sous-question doit être autonome et cibler un seul aspect.\n"
"Si la question est simple ou porte sur un seul sujet, retourne une liste vide."
)
_EVAL_INSTRUCTIONS = (
"Tu es un évaluateur de pertinence pour un système RAG juridique.\n"
"Évalue si les extraits fournis permettent de répondre complètement à la question."
)
_REFINE_INSTRUCTIONS = (
"Tu es un assistant de reformulation de requêtes pour un système RAG juridique.\n"
"Produis une requête de recherche ciblée qui comble le manque identifié."
)
_INTENT_AGENT = Agent(
output_type=RoutingIntentOutput,
instructions=(
"You are a legal query parser for FR/EU retrieval routing.\n"
"Choose a profile, a granularity, a top_k, whether relation context is needed, "
"and whether graph support should be enabled."
),
output_retries=2,
defer_model_check=True,
)
@_INTENT_AGENT.output_validator
def _validate_intent_output(result: RoutingIntentOutput) -> RoutingIntentOutput:
if not result.rationale.strip():
raise ModelRetry("A rationale is required.")
return result
_DECOMPOSITION_AGENT = Agent(
output_type=DecompositionOutput,
instructions=_DECOMP_INSTRUCTIONS,
output_retries=2,
defer_model_check=True,
)
@_DECOMPOSITION_AGENT.output_validator
def _validate_decomposition_output(result: DecompositionOutput) -> DecompositionOutput:
if len(result.sub_questions) > 3:
raise ModelRetry("At most 3 sub-questions are allowed.")
return result
_PLANNER_AGENT = Agent(
output_type=RetrievalPlannerOutput,
instructions=(
"Tu es un planificateur de récupération pour un RAG juridique.\n"
"Décide de l'intention, de la requête primaire, d'un éventuel besoin de clarification, "
"des récupérations complémentaires et du besoin de compression."
),
output_retries=2,
defer_model_check=True,
)
@_PLANNER_AGENT.output_validator
def _validate_planner_output(result: RetrievalPlannerOutput) -> RetrievalPlannerOutput:
if result.intent_class != "conversational" and not result.primary_query.strip():
raise ModelRetry("primary_query cannot be empty when retrieval is required.")
return result
_REFINEMENT_AGENT = Agent(
output_type=RetrievalRefinementOutput,
instructions=_REFINE_INSTRUCTIONS,
output_retries=2,
defer_model_check=True,
)
@_REFINEMENT_AGENT.output_validator
def _validate_refinement_output(result: RetrievalRefinementOutput) -> RetrievalRefinementOutput:
if not result.refined_query.strip():
raise ModelRetry("refined_query cannot be empty.")
return result
_EVALUATION_AGENT = Agent(
output_type=RetrievalEvaluationOutput,
instructions=_EVAL_INSTRUCTIONS,
output_retries=2,
defer_model_check=True,
)
[docs]
def run_intent_parser_agent(
*,
question: str,
top_k: int,
requested_granularity: Optional[str],
generate_text: Callable[[str], str],
model_name: str,
) -> tuple[RoutingIntentOutput, int]:
"""Run the structured intent parser agent for one question."""
granularity = requested_granularity if requested_granularity in {"subdivisions", "chunks", "all"} else "auto"
prompt = render_intent_parser_prompt(
question=question,
top_k=max(int(top_k), 1),
requested_granularity=granularity,
)
return run_structured_agent(
agent=_INTENT_AGENT,
prompt=prompt,
llm_or_generate=generate_text,
model_name=model_name,
)
[docs]
def run_decomposition_agent(
*,
question: str,
llm: Any,
model_name: str = "planner:decompose",
) -> tuple[DecompositionOutput, int]:
"""Run the decomposition agent for one complex question."""
prompt = f"Question : {question}"
return run_structured_agent(
agent=_DECOMPOSITION_AGENT,
prompt=prompt,
llm_or_generate=llm,
model_name=model_name,
)
[docs]
def run_planner_agent(
*,
question: str,
llm: Any,
model_name: str = "planner:retrieve",
) -> tuple[RetrievalPlannerOutput, int]:
"""Run the retrieval planner agent."""
return run_structured_agent(
agent=_PLANNER_AGENT,
prompt=render_planner_prompt(question=question),
llm_or_generate=llm,
model_name=model_name,
)
[docs]
def run_refinement_agent(
*,
question: str,
gap_hint: str,
llm: Any,
model_name: str = "planner:refine",
) -> tuple[RetrievalRefinementOutput, int]:
"""Run the corrective refinement agent for weak retrieval results."""
prompt = f"Question originale : {question}\n\nManque identifié dans les résultats actuels : {gap_hint}"
return run_structured_agent(
agent=_REFINEMENT_AGENT,
prompt=prompt,
llm_or_generate=llm,
model_name=model_name,
)
[docs]
def run_evaluation_agent(
*,
question: str,
context_preview: str,
llm: Any,
model_name: str = "planner:evaluate",
) -> tuple[RetrievalEvaluationOutput, int]:
"""Run the sufficiency evaluator on a preview of retrieved context."""
prompt = f"Question : {question}\n\nExtraits récupérés :\n{context_preview}"
return run_structured_agent(
agent=_EVALUATION_AGENT,
prompt=prompt,
llm_or_generate=llm,
model_name=model_name,
)