Source code for lalandre_rag.modes.hybrid_graph

"""
Graph enrichment for HybridMode — Neo4j expansion and community context.

Extracted from hybrid_mode.py to keep the orchestrator focused on the pipeline.
"""

import logging
import time
from typing import Any, Dict, List, Optional, Set

from lalandre_core.config import get_config

from ..graph import CommunityContextEnricher, GraphRAGService
from .hybrid_helpers import GraphFetchResult

logger = logging.getLogger(__name__)




def fetch_graph_context(
    *,
    act_ids: Set[int],
    graph_rag_service: GraphRAGService,
    community_enricher: Optional[CommunityContextEnricher],
    max_depth: Optional[int] = None,
) -> Optional[GraphFetchResult]:
    """Fetch graph data for acts found in retrieval results.

    Uses act_ids from hybrid retrieval (semantic + BM25) as seeds for
    Neo4j graph traversal. Returns structured data for ranking/budgeting.
    Non-fatal: returns None on failure.
    """
    if not act_ids:
        return None

    started_at = time.perf_counter()
    try:
        config = get_config()
        effective_depth = max_depth if max_depth is not None else config.graph.hybrid_enrichment_depth

        graph_result = graph_rag_service.neo4j.expand_from_acts(
            act_ids=list(act_ids),
            max_depth=effective_depth,
        )

        nodes = list(graph_result.nodes or [])
        expanded_act_ids: set[int] = {int(n["id"]) for n in nodes if n.get("id") is not None}

        relationships: List[Dict[str, Any]] = []
        for rel in graph_result.relationships or []:
            relationships.append(
                {
                    "type": rel.get("type", "RELATED_TO"),
                    "start_node": rel.get("start_node"),
                    "end_node": rel.get("end_node"),
                    "description": rel.get("properties", {}).get("description", ""),
                }
            )

        community_block = ""
        community_meta: Dict[str, Any] = {}
        if config.graph.use_communities_in_rag and community_enricher is not None and community_enricher.available:
            community_block, community_meta = community_enricher.build_context(
                seed_act_ids=act_ids | expanded_act_ids,
            )

        duration_ms = round((time.perf_counter() - started_at) * 1000.0, 1)
        return GraphFetchResult(
            nodes=nodes,
            relationships=relationships,
            seed_act_ids=act_ids,
            expanded_act_ids=expanded_act_ids,
            community_block=community_block,
            community_meta=community_meta,
            duration_ms=duration_ms,
        )
    except Exception as exc:
        logger.warning("Graph enrichment failed (non-fatal), continuing without: %s", exc)
        return None