Coverage for /home/admin/Documents/AI/applications/lexigram-dev/lexigram/experimental/ai/lexigram-ai-rag/src/lexigram/ai/rag/knowledge_graph/adapter.py: 24%
157 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-25 07:19 +0800
1"""Knowledge-graph adapter backed by ``GraphStoreProtocol``.
3Wraps a ``GraphProtocol`` instance (resolved from the DI container via
4``lexigram-graph``) and exposes the same interface as ``KnowledgeGraph``,
5mapping between the RAG-layer ``Entity``/``Relationship``/``GraphPath``
6domain types and the infrastructure ``GraphNode``/``GraphEdge``/``GraphPath``
7contracts types.
9When ``GraphStoreProtocol`` is registered in the container (e.g. by
10``GraphModule`` / ``GraphProvider``), ``RAGProvider`` automatically creates
11this adapter and registers it as the singleton ``KnowledgeGraph`` — giving
12the RAG pipeline a fully persistent graph back-end. When the infra store
13is absent, ``RAGProvider`` falls back to the in-memory ``KnowledgeGraph``.
14"""
16from __future__ import annotations
18from typing import TYPE_CHECKING, Any
20from lexigram.ai.rag.knowledge_graph.types import (
21 Entity,
22 EntityType,
23 GraphPath,
24 Relationship,
25 RelationshipType,
26)
27from lexigram.contracts.data.graph.enums import EdgeDirection
28from lexigram.contracts.data.graph.types import (
29 EdgeSpec,
30 NodeSpec,
31 StartSpec,
32 TraversalQuery,
33 TraversalStep,
34)
35from lexigram.logging import (
36 get_logger,
37)
39if TYPE_CHECKING:
40 from lexigram.contracts.data.graph.protocols import GraphProtocol
41 from lexigram.contracts.data.graph.types import (
42 GraphEdge,
43 GraphNode,
44 )
45 from lexigram.contracts.data.graph.types import (
46 GraphPath as InfraGraphPath,
47 )
49logger = get_logger(__name__)
51# Property keys stored alongside entity data but excluded from the
52# round-tripped ``Entity.properties`` dict.
53_META_KEY = "__kg_metadata"
54_NAME_KEY = "name"
55_TYPE_KEY = "type"
56_SRC_KEY = "_source_name"
57_TGT_KEY = "_target_name"
58_CONF_KEY = "confidence"
60_RESERVED_NODE_KEYS = frozenset({_META_KEY, _NAME_KEY, _TYPE_KEY})
61_RESERVED_EDGE_KEYS = frozenset({_META_KEY, _SRC_KEY, _TGT_KEY, _CONF_KEY})
64# ── Type-conversion helpers ───────────────────────────────────────────────────
67def _entity_to_node_spec(entity: Entity) -> NodeSpec:
68 """Build a ``NodeSpec`` from an ``Entity``."""
69 properties: dict[str, Any] = {
70 _NAME_KEY: entity.name,
71 _TYPE_KEY: str(entity.type),
72 **entity.properties,
73 }
74 if entity.metadata:
75 properties[_META_KEY] = entity.metadata
76 return NodeSpec(
77 labels=(str(entity.type),),
78 properties=properties,
79 id=entity.name.lower(),
80 )
83def _node_to_entity(node: GraphNode) -> Entity:
84 """Reconstruct an ``Entity`` from a ``GraphNode``."""
85 name: str = node.properties.get(_NAME_KEY, node.id)
86 raw_type: str = (
87 node.properties.get(_TYPE_KEY)
88 or (node.labels[0] if node.labels else None)
89 or str(EntityType.OTHER)
90 )
91 try:
92 entity_type: EntityType | str = EntityType(raw_type)
93 except ValueError:
94 entity_type = raw_type
96 extra = {k: v for k, v in node.properties.items() if k not in _RESERVED_NODE_KEYS}
97 metadata: dict[str, Any] = node.properties.get(_META_KEY, {})
98 return Entity(
99 name=name,
100 type=entity_type,
101 properties=extra,
102 metadata=metadata,
103 )
106def _relationship_to_edge_spec(rel: Relationship) -> EdgeSpec:
107 """Build an ``EdgeSpec`` from a ``Relationship``."""
108 properties: dict[str, Any] = {
109 _CONF_KEY: rel.confidence,
110 _SRC_KEY: rel.source,
111 _TGT_KEY: rel.target,
112 **rel.properties,
113 }
114 if rel.metadata:
115 properties[_META_KEY] = rel.metadata
116 return EdgeSpec(
117 source_id=rel.source.lower(),
118 target_id=rel.target.lower(),
119 type=str(rel.type),
120 properties=properties,
121 )
124def _edge_to_relationship(edge: GraphEdge) -> Relationship:
125 """Reconstruct a ``Relationship`` from a ``GraphEdge``."""
126 source: str = edge.properties.get(_SRC_KEY, edge.source_id)
127 target: str = edge.properties.get(_TGT_KEY, edge.target_id)
128 confidence: float = edge.properties.get(_CONF_KEY, 1.0)
129 raw_type: str = edge.type
130 try:
131 rel_type: RelationshipType | str = RelationshipType(raw_type)
132 except ValueError:
133 rel_type = raw_type
135 extra = {k: v for k, v in edge.properties.items() if k not in _RESERVED_EDGE_KEYS}
136 metadata: dict[str, Any] = edge.properties.get(_META_KEY, {})
137 return Relationship(
138 source=source,
139 target=target,
140 type=rel_type,
141 confidence=confidence,
142 properties=extra,
143 metadata=metadata,
144 )
147def _infra_path_to_rag_path(path: InfraGraphPath) -> GraphPath:
148 """Convert an infrastructure ``GraphPath`` to a RAG ``GraphPath``."""
149 entity_names = [node.properties.get(_NAME_KEY, node.id) for node in path.nodes]
150 rels = [_edge_to_relationship(e) for e in path.edges]
151 confidence_sum = sum(r.confidence for r in rels)
152 score = confidence_sum / len(rels) if rels else 1.0
153 return GraphPath(
154 entities=entity_names,
155 relationships=rels,
156 length=path.length,
157 score=score,
158 )
161# ── Adapter ───────────────────────────────────────────────────────────────────
164class GraphStoreAdapter:
165 """``KnowledgeGraph``-compatible adapter backed by a ``GraphProtocol``.
167 Entities are stored as nodes (label = entity type, id = name.lower()).
168 Relationships are stored as directed edges (type = relationship type).
169 Confidence and original name strings are persisted in edge properties so
170 round-tripping is lossless.
172 Args:
173 graph: A ``GraphProtocol`` instance from ``lexigram-graph`` (or any
174 other backend implementing the contract).
175 """
177 def __init__(self, graph: GraphProtocol) -> None:
178 """Initialise the adapter with a graph protocol instance."""
179 self._graph = graph
181 # ── Entity mutations ──────────────────────────────────────────
183 async def add_entity(self, entity: Entity) -> None:
184 """Add or upsert an entity as a graph node."""
185 spec = _entity_to_node_spec(entity)
186 existing = await self._graph.get_node(entity.name.lower())
187 if existing is None:
188 await self._graph.bulk_create_nodes([spec])
189 else:
190 # Upsert: merge properties
191 await self._graph.update_node(
192 entity.name.lower(),
193 properties={
194 _NAME_KEY: entity.name,
195 _TYPE_KEY: str(entity.type),
196 **entity.properties,
197 **({_META_KEY: entity.metadata} if entity.metadata else {}),
198 },
199 merge=True,
200 )
202 async def add_entities(self, entities: list[Entity]) -> None:
203 """Add or upsert multiple entities."""
204 for entity in entities:
205 await self.add_entity(entity)
207 async def add_relationship(self, relationship: Relationship) -> None:
208 """Add a relationship as a directed edge between two nodes."""
209 spec = _relationship_to_edge_spec(relationship)
210 try:
211 await self._graph.create_edge(
212 source_id=spec.source_id,
213 target_id=spec.target_id,
214 edge_type=spec.type,
215 properties=spec.properties,
216 )
217 except Exception as exc: # noqa: BLE001 — graph adapter relationship; skipped if nodes don't exist; log and continue
218 # Nodes may not exist yet — surface as a warning and skip.
219 logger.warning(
220 "graph_adapter_relationship_skipped",
221 source=relationship.source,
222 target=relationship.target,
223 error=str(exc),
224 )
226 async def add_relationships(self, relationships: list[Relationship]) -> None:
227 """Add multiple relationships."""
228 for rel in relationships:
229 await self.add_relationship(rel)
231 # ── Entity queries ────────────────────────────────────────────
233 async def get_entity(self, name: str) -> Entity | None:
234 """Retrieve an entity by name. Returns ``None`` if not found."""
235 node = await self._graph.get_node(name.lower())
236 if node is None:
237 return None
238 return _node_to_entity(node)
240 async def get_all_entities(self) -> list[Entity]:
241 """Return all entities in the graph."""
242 nodes = await self._graph.find_nodes(limit=10_000)
243 return [_node_to_entity(n) for n in nodes]
245 async def get_entities_by_type(
246 self,
247 entity_type: EntityType | str,
248 ) -> list[Entity]:
249 """Return all entities of the given type."""
250 label = str(entity_type)
251 nodes = await self._graph.find_nodes(labels=[label], limit=10_000)
252 return [_node_to_entity(n) for n in nodes]
254 # ── Relationship queries ──────────────────────────────────────
256 async def get_all_relationships(self) -> list[Relationship]:
257 """Return all relationships in the graph."""
258 nodes = await self._graph.find_nodes(limit=10_000)
259 rels: list[Relationship] = []
260 seen: set[str] = set()
261 for node in nodes:
262 edges = await self._graph.get_edges(
263 node.id,
264 direction=EdgeDirection.OUTGOING,
265 limit=10_000,
266 )
267 for edge in edges:
268 if edge.id not in seen:
269 seen.add(edge.id)
270 rels.append(_edge_to_relationship(edge))
271 return rels
273 async def get_neighbors(
274 self,
275 entity_name: str,
276 direction: str = "outgoing",
277 ) -> list[Entity]:
278 """Return neighbouring entities reachable in one hop."""
279 dir_map: dict[str, EdgeDirection] = {
280 "outgoing": EdgeDirection.OUTGOING,
281 "incoming": EdgeDirection.INCOMING,
282 "both": EdgeDirection.BOTH,
283 }
284 infra_dir = dir_map.get(direction, EdgeDirection.BOTH)
285 nodes = await self._graph.neighbors(
286 node_id=entity_name.lower(),
287 depth=1,
288 direction=infra_dir,
289 )
290 return [_node_to_entity(n) for n in nodes]
292 async def get_relationships(
293 self,
294 entity_name: str,
295 direction: str = "outgoing",
296 ) -> list[Relationship]:
297 """Return relationships attached to the named entity."""
298 dir_map: dict[str, EdgeDirection] = {
299 "outgoing": EdgeDirection.OUTGOING,
300 "incoming": EdgeDirection.INCOMING,
301 "both": EdgeDirection.BOTH,
302 }
303 infra_dir = dir_map.get(direction, EdgeDirection.BOTH)
304 edges = await self._graph.get_edges(
305 node_id=entity_name.lower(),
306 direction=infra_dir,
307 limit=10_000,
308 )
309 return [_edge_to_relationship(e) for e in edges]
311 # ── Path queries ──────────────────────────────────────────────
313 async def find_path(
314 self,
315 source: str,
316 target: str,
317 max_depth: int = 5,
318 relationship_types: list[RelationshipType | str] | None = None,
319 ) -> GraphPath | None:
320 """Find the shortest path between two entities."""
321 edge_types = (
322 [str(t) for t in relationship_types] if relationship_types else None
323 )
324 infra_path = await self._graph.shortest_path(
325 from_id=source.lower(),
326 to_id=target.lower(),
327 max_depth=max_depth,
328 edge_types=edge_types,
329 direction=EdgeDirection.BOTH,
330 )
331 if infra_path is None:
332 return None
333 return _infra_path_to_rag_path(infra_path)
335 async def find_all_paths(
336 self,
337 source: str,
338 target: str,
339 max_depth: int = 5,
340 max_paths: int = 10,
341 ) -> list[GraphPath]:
342 """Find up to ``max_paths`` paths between two entities via traversal."""
343 start = StartSpec(node_ids=(source.lower(),))
344 steps = (
345 TraversalStep(
346 direction=EdgeDirection.BOTH,
347 min_depth=1,
348 max_depth=max_depth,
349 ),
350 )
351 query = TraversalQuery(
352 start=start,
353 steps=steps,
354 limit=max_paths,
355 unique_nodes=True,
356 )
357 all_paths = await self._graph.traverse(query)
358 # Filter to only paths that end at the target node.
359 target_id = target.lower()
360 matching = [p for p in all_paths if p.end_node.id == target_id]
361 rag_paths = [_infra_path_to_rag_path(p) for p in matching[:max_paths]]
362 rag_paths.sort(key=lambda p: p.score, reverse=True)
363 return rag_paths
365 async def query_subgraph(
366 self,
367 entity_name: str,
368 depth: int = 2,
369 ) -> tuple[list[Entity], list[Relationship]]:
370 """Return all entities and relationships within ``depth`` hops."""
371 start_id = entity_name.lower()
372 # Collect reachable nodes at each depth level.
373 entities_found: dict[str, Entity] = {}
374 rels_found: dict[str, Relationship] = {}
376 root = await self._graph.get_node(start_id)
377 if root is None:
378 return [], []
379 entities_found[start_id] = _node_to_entity(root)
381 current_layer: set[str] = {start_id}
382 for _ in range(depth):
383 next_layer: set[str] = set()
384 for nid in current_layer:
385 edges = await self._graph.get_edges(
386 nid,
387 direction=EdgeDirection.BOTH,
388 limit=10_000,
389 )
390 for edge in edges:
391 neighbor_id = (
392 edge.target_id if edge.source_id == nid else edge.source_id
393 )
394 if edge.id not in rels_found:
395 rels_found[edge.id] = _edge_to_relationship(edge)
396 if neighbor_id not in entities_found:
397 neighbor = await self._graph.get_node(neighbor_id)
398 if neighbor is not None:
399 entities_found[neighbor_id] = _node_to_entity(neighbor)
400 next_layer.add(neighbor_id)
401 current_layer = next_layer
403 return list(entities_found.values()), list(rels_found.values())
405 # ── Stats ─────────────────────────────────────────────────────
407 async def get_stats(self) -> dict[str, Any]:
408 """Return basic graph statistics."""
409 node_count = await self._graph.count_nodes()
410 edge_count = await self._graph.count_edges()
411 labels = await self._graph.get_labels()
412 edge_types = await self._graph.get_edge_types()
413 return {
414 "total_entities": node_count,
415 "total_relationships": edge_count,
416 "entity_types": labels,
417 "relationship_types": edge_types,
418 }
420 def __len__(self) -> int:
421 """Return entity count (synchronous best-effort — may be stale)."""
422 # Not async-safe; kept for protocol compatibility.
423 return 0
425 def __repr__(self) -> str:
426 return f"GraphStoreAdapter(graph={self._graph!r})"