from __future__ import annotations from dataclasses import dataclass, field from typing import Any, Literal, Mapping, Optional, Sequence, TypedDict import rich from rich.console import Console from rich.panel import Panel from rich.text import Text from kbdebugger.graph import get_graph from kbdebugger.types import GraphRelation, EdgePropertyKey from kbdebugger.utils.json import write_json from kbdebugger.utils.time import now_utc_compact from .utils import normalize_text MatchPattern = Literal["source_label", "target_label", "rel_props"] class RetrievedRelation(TypedDict): relation: GraphRelation match_pattern: MatchPattern @dataclass class KnowledgeGraphRetriever: """ Keyword-guided KG retrieval. MVP: returns 1-hop 'path fragments' (edges) as normalized GraphRelation objects. """ limit_per_pattern: int = 50 console: Console = field(default_factory=Console) # console is a member so we can inject a test console or reuse a global one. def retrieve( self, keyword: str, *, limit_per_pattern: Optional[int] = None, ) -> list[RetrievedRelation]: kw = normalize_text(keyword) limit = int(limit_per_pattern or self.limit_per_pattern) graph = get_graph() results: list[RetrievedRelation] = [] # --- Pattern 1: keyword in source node label --- rels = graph.query_relations( """ MATCH (n:Node)-[r:REL]->(m:Node) WHERE toLower(n.label) CONTAINS $keyword RETURN n.label AS source, m.label AS target, coalesce(r.type, r.label, 'REL') AS predicate, properties(r) AS props, elementId(n) AS source_id, elementId(m) AS target_id, elementId(r) AS rel_id LIMIT $limit """, {"keyword": kw, "limit": limit}, ) results.extend({"relation": rel, "match_pattern": "source_label"} for rel in rels) # --- Pattern 2: keyword in target node label --- rels = graph.query_relations( """ MATCH (n:Node)-[r:REL]->(m:Node) WHERE toLower(m.label) CONTAINS $keyword RETURN n.label AS source, m.label AS target, coalesce(r.type, r.label, 'REL') AS predicate, properties(r) AS props, elementId(n) AS source_id, elementId(m) AS target_id, elementId(r) AS rel_id LIMIT $limit """, {"keyword": kw, "limit": limit}, ) results.extend({"relation": rel, "match_pattern": "target_label"} for rel in rels) # --- Pattern 3: keyword in relationship "semantic" fields --- # We avoid fancy APOC here; just check the usual fields you write. rels = graph.query_relations( """ MATCH (n:Node)-[r:REL]->(m:Node) WHERE toLower(coalesce(r.type, "")) CONTAINS $keyword OR toLower(coalesce(r.label, "")) CONTAINS $keyword OR toLower(coalesce(r.source, "")) CONTAINS $keyword RETURN n.label AS source, m.label AS target, coalesce(r.type, r.label, 'REL') AS predicate, properties(r) AS props, elementId(n) AS source_id, elementId(m) AS target_id, elementId(r) AS rel_id LIMIT $limit """, {"keyword": kw, "limit": limit}, ) results.extend({"relation": rel, "match_pattern": "rel_props"} for rel in rels) # Optional: dedupe identical relations across patterns # (same source/target/predicate + same sentence/source if you want) results = self._dedupe(results) # # TODO: Enable again after we're done with Frontend integration # self.save_results_json( # keyword=keyword, # hits=results, # limit_per_pattern=limit, # ) return results @staticmethod def _dedupe(items: list[RetrievedRelation]) -> list[RetrievedRelation]: seen: set[tuple[str, str, str, str]] = set() out: list[RetrievedRelation] = [] for item in items: rel = item["relation"] props = rel["edge"]["properties"] sentence = str(props.get("sentence", "")) # good lightweight key key = (rel["source"]["label"], rel["target"]["label"], rel["edge"]["label"], sentence) if key in seen: continue seen.add(key) out.append(item) return out @staticmethod def save_results_json( *, keyword: str, hits: Sequence[RetrievedRelation], limit_per_pattern: int | None = None, extra_metadata: Mapping[str, Any] | None = None, ) -> None: """ Save KG retrieval results to a JSON log file. Parameters ---------- keyword: The user keyword used to retrieve the subgraph. hits: The retrieved results (a list of RetrievedRelation). Each element contains: - relation: GraphRelation - match_pattern: str (provenance of how it matched) limit_per_pattern: Optional: include the retriever's `limit_per_pattern` setting in the log. i.e., how many relations were retrieved per `MatchPattern`. extra_metadata: Optional: additional metadata to include in the JSON file (e.g., commit hash, run id, corpus file name, etc.). JSON Schema (high-level) ------------------------ { "keyword": "...", "limit_per_pattern": 50, "num_hits": 123, "hits": [...], "extra": {...} } """ created_at = now_utc_compact() payload: dict[str, Any] = { "keyword": keyword, "limit_per_pattern": limit_per_pattern, "num_hits": len(hits), "hits": list(hits), "created_at": created_at, } if extra_metadata: payload["extra"] = dict(extra_metadata) path = f"logs/02_kg_retrieval_{keyword}_{created_at}.json" write_json(path, payload) rich.print(f"\n[INFO] Wrote KG retrieval log to {path}") def pretty_print( self, hits: Sequence[RetrievedRelation], *, title: str = "Knowledge Graph Retrieval Results", show_props_keys: Optional[Sequence[EdgePropertyKey]] = None, ) -> None: """ Pretty-print RetrievedRelation results using rich. - hits: output of retrieve(..., include_match_pattern=True) - show_props_keys: if set, prints only these keys from edge.properties (in addition to sentence/source/page) """ if not hits: self.console.rich.print("[bold yellow]No matching relations found.[/bold yellow]") return if not show_props_keys: show_props_keys = [ "created_at", "last_updated_at", "original_sentence", ] self.console.rule(f"[bold cyan]{title}[/bold cyan]") for i, hit in enumerate(hits, start=1): rel = hit["relation"] pattern = hit["match_pattern"] src = rel["source"]["label"] tgt = rel["target"]["label"] pred = rel["edge"]["label"] props = rel["edge"]["properties"] sentence = props.get("sentence") source_doc = props.get("source") page = props.get("page_number") header = Text() header.append(f"[{i}] ", style="bold cyan") header.append(f"{src} ", style="bold green") header.append("── ", style="dim") header.append(pred, style="bold magenta") header.append(" ──> ", style="dim") header.append(tgt, style="bold green") body: list[str] = [f"[bold]Matched via:[/bold] {pattern}"] if sentence: body.append(f"[bold]Sentence:[/bold] {sentence}") if source_doc: meta = f"{source_doc}" if page is not None: meta += f", page {page}" body.append(f"[bold]Source:[/bold] {meta}") if show_props_keys: for k in show_props_keys: if k in props and props[k] is not None: # type: ignore body.append(f"[bold]{k}:[/bold] {props[k]}") # type: ignore self.console.rich.print( Panel( "\n".join(body), title=header, border_style="cyan", padding=(1, 2), ) )