Source code for lalandre_rag.agentic.deps

"""Dependency container for the PydanticAI planning runtime."""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Protocol

from langchain_core.messages import BaseMessage
from langchain_core.prompts import ChatPromptTemplate




class RetrievalServiceProtocol(Protocol):
    """Protocol for the retrieval service consumed by the planning graph."""

    last_retrieval_stats: Any



    def retrieve(self, **_: Any) -> list[Any]:
        """Run retrieval and return raw retrieval results."""
        ...






class ContextServiceProtocol(Protocol):
    """Protocol for retrieval-result enrichment services."""



    def enrich_results(
        self,
        results: list[Any],
        include_relations: bool = False,
        include_subjects: bool = False,
    ) -> list[Any]:
        """Enrich retrieval results with act metadata and relationships."""
        ...






class QueryRouterProtocol(Protocol):
    """Protocol for deterministic or LLM-assisted retrieval routing."""

    intent_parser: Any



    def route(
        self,
        *,
        question: str,
        top_k: int,
        requested_granularity: Optional[str],
    ) -> Any:
        """Build a retrieval plan for the incoming user question."""
        ...






class GraphRAGServiceProtocol(Protocol):
    """Protocol for optional graph-retrieval services."""

    neo4j: Any





class CommunityEnricherProtocol(Protocol):
    """Protocol for community-summary enrichers used by global mode."""

    available: bool



    def build_context(
        self,
        *,
        _seed_act_ids: set[int],
    ) -> tuple[str, Dict[str, Any]]:
        """Build a community-level context block and its metadata."""
        ...




ProgressCallback = Callable[[Dict[str, Any]], None]
PreambleCallback = Callable[[Optional[Dict[str, Any]], Dict[str, Any]], None]
TokenCallback = Callable[[str], None]
FinalAnswerCallback = Callable[[str], None]




@dataclass
class AgenticPlanningDeps:
    """Runtime dependencies injected into the planning graph."""

    retrieval_service: RetrievalServiceProtocol
    context_service: ContextServiceProtocol
    llm: Any
    lightweight_llm: Any
    rag_prompt: ChatPromptTemplate
    query_router: QueryRouterProtocol
    graph_rag_service: Optional[GraphRAGServiceProtocol]
    community_enricher: Optional[CommunityEnricherProtocol]
    question: str
    top_k: int
    score_threshold: Optional[float]
    filters: Optional[Dict[str, Any]]
    include_relations: bool
    include_subjects: bool
    include_full_content: bool
    return_sources: bool
    collections: Optional[List[str]]
    granularity: Optional[str]
    graph_depth: Optional[int]
    use_graph: Optional[bool]
    embedding_preset: Optional[str]
    retrieval_depth: Optional[str]
    chat_history: Optional[List[BaseMessage]]
    progress_callback: Optional[ProgressCallback] = None
    preamble_callback: Optional[PreambleCallback] = None
    token_callback: Optional[TokenCallback] = None
    final_answer_callback: Optional[FinalAnswerCallback] = None