Aman Nalakath
7503a441c2
Some checks failed
tests / core (macos-latest, py3.11) (push) Has been cancelled
tests / core (macos-latest, py3.12) (push) Has been cancelled
tests / core (macos-latest, py3.13) (push) Has been cancelled
tests / core (ubuntu-latest, py3.11) (push) Has been cancelled
tests / core (ubuntu-latest, py3.12) (push) Has been cancelled
tests / core (ubuntu-latest, py3.13) (push) Has been cancelled
tests / doctor CLI smoke test (push) Has been cancelled
90 lines
3.3 KiB
Python
90 lines
3.3 KiB
Python
"""Traversal queries over the knowledge graph: entity/citation -> chunks, neighbors."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from .builder import _normalize_entity
|
|
|
|
|
|
def _entity_key(text: str) -> str:
|
|
return f"entity:{_normalize_entity(text)}"
|
|
|
|
|
|
def _chunk_payload(graph: Any, chunk_node: str) -> dict[str, Any]:
|
|
data = graph.nodes[chunk_node]
|
|
doc_node = next(
|
|
(src for src, _, e in graph.in_edges(chunk_node, data=True) if e.get("kind") == "CONTAINS"),
|
|
None,
|
|
)
|
|
doc = graph.nodes[doc_node] if doc_node is not None else {}
|
|
return {
|
|
"chunk": chunk_node,
|
|
"chunk_index": data.get("chunk_index", 0),
|
|
"text": data.get("text", ""),
|
|
"source_path": doc.get("source_path", ""),
|
|
"ocr_engine": doc.get("ocr_engine", ""),
|
|
"ocr_confidence_mean": doc.get("ocr_confidence_mean", 0.0),
|
|
"timestamp": doc.get("timestamp", ""),
|
|
}
|
|
|
|
|
|
def chunks_for_entity(graph: Any, entity_text: str) -> list[dict[str, Any]]:
|
|
"""All chunks mentioning the given entity (case-insensitive)."""
|
|
key = _entity_key(entity_text)
|
|
if key not in graph:
|
|
return []
|
|
chunks = [
|
|
src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "MENTIONS"
|
|
]
|
|
return [_chunk_payload(graph, c) for c in sorted(chunks)]
|
|
|
|
|
|
def chunks_for_citation(graph: Any, identifier: str) -> list[dict[str, Any]]:
|
|
"""All chunks citing the given identifier (e.g. '10.1038/nature12345')."""
|
|
key = identifier if identifier.startswith("citation:") else f"citation:{identifier}"
|
|
if key not in graph:
|
|
return []
|
|
chunks = [src for src, _, edge in graph.in_edges(key, data=True) if edge.get("kind") == "CITES"]
|
|
return [_chunk_payload(graph, c) for c in sorted(chunks)]
|
|
|
|
|
|
def related_entities(graph: Any, entity_text: str, limit: int = 10) -> list[tuple[str, int]]:
|
|
"""Entities co-occurring with the given one, as (text, shared-chunk weight)."""
|
|
key = _entity_key(entity_text)
|
|
if key not in graph:
|
|
return []
|
|
neighbors: dict[str, int] = {}
|
|
for _, target, edge in graph.out_edges(key, data=True):
|
|
if edge.get("kind") == "CO_OCCURS":
|
|
neighbors[target] = edge.get("weight", 1)
|
|
for source, _, edge in graph.in_edges(key, data=True):
|
|
if edge.get("kind") == "CO_OCCURS":
|
|
neighbors[source] = max(neighbors.get(source, 0), edge.get("weight", 1))
|
|
ranked = sorted(
|
|
((graph.nodes[n].get("text", n), w) for n, w in neighbors.items()),
|
|
key=lambda item: item[1],
|
|
reverse=True,
|
|
)
|
|
return ranked[:limit]
|
|
|
|
|
|
def expand_context(graph: Any, entity_text: str, limit: int = 5) -> list[dict[str, Any]]:
|
|
"""Chunks mentioning the entity or its strongest co-occurring neighbors.
|
|
|
|
This is the graph-RAG primitive: one hop of CO_OCCURS expansion, so a query
|
|
for 'BRCA1' also surfaces chunks that only mention its interactors.
|
|
"""
|
|
seen: set[str] = set()
|
|
results: list[dict[str, Any]] = []
|
|
|
|
seeds = [entity_text] + [text for text, _ in related_entities(graph, entity_text, limit=limit)]
|
|
for seed in seeds:
|
|
for payload in chunks_for_entity(graph, seed):
|
|
if payload["chunk"] in seen:
|
|
continue
|
|
seen.add(payload["chunk"])
|
|
payload = {**payload, "via_entity": seed}
|
|
results.append(payload)
|
|
return results
|