From 18da322e0d07eed3a1297bbaca9161c632c7fb6f Mon Sep 17 00:00:00 2001 From: Mohd Kaif <98801504+KaifAhmad1@users.noreply.github.com> Date: Mon, 11 May 2026 17:38:24 +0530 Subject: [PATCH] Feature/distance intelligence optimization (#550) * Implement embedding cache optimization for Distance Intelligence - Add per-session graph revision-based embedding cache to avoid re-scanning nodes - Update GraphSession with get_cached_embeddings() and automatic cache invalidation - Modify distance matrix and semantic neighborhood endpoints to use cached embeddings - Implement thread-safe caching with proper revision tracking - Add force refresh capability and automatic invalidation on graph modifications - Improve performance for repeated distance intelligence queries Resolves TODO in graph.py: cache embeddings per-session graph revision * Update changelog with Distance Intelligence embedding cache optimization --- CHANGELOG.md | 7 ++ semantica/explorer/routes/graph.py | 31 ++++++--- semantica/explorer/session.py | 103 +++++++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 8 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c1786c67..05bef69f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] ### Added +- **Distance Intelligence Embedding Cache Optimization** by @Assistant + - Implemented per-session graph revision-based embedding cache to avoid re-scanning all nodes on every request + - Added `get_cached_embeddings()` method to GraphSession with thread-safe caching and automatic invalidation + - Updated distance matrix and semantic neighborhood endpoints to use cached embeddings for significant performance improvement + - Added graph revision tracking using hash-based identifiers for cache invalidation + - Implemented force refresh capability and automatic cache invalidation on graph modifications (add_nodes/add_edges) + - Resolved TODO in `graph.py` for embedding caching optimization - **Parquet File Ingestion Support** (#548) by @Luffy2208 - Added ParquetIngestor class with PyArrow backend - Single file and partitioned directory ingestion diff --git a/semantica/explorer/routes/graph.py b/semantica/explorer/routes/graph.py index 018122f9..bdd29741 100644 --- a/semantica/explorer/routes/graph.py +++ b/semantica/explorer/routes/graph.py @@ -102,10 +102,10 @@ def _coerce_embedding_vector(value: object) -> Optional[List[float]]: def _extract_node_embeddings(graph_dict: dict) -> dict[str, List[float]]: + """Extract embeddings from graph dictionary.""" # Top-level keys to probe on each entity (and its metadata/properties dicts). # Priority: generic names first, then KG-extras-specific names. # Must stay aligned with the inner probe list in _coerce_embedding_vector. - # TODO: cache this per-session graph revision to avoid re-scanning all nodes on every request. embedding_keys = ( "embedding", "embeddings", @@ -138,6 +138,11 @@ def _extract_node_embeddings(graph_dict: dict) -> dict[str, List[float]]: return embeddings +def _get_cached_embeddings(session: GraphSession) -> dict[str, List[float]]: + """Get embeddings from session cache for optimal performance.""" + return session.get_cached_embeddings() + + def _node_response(node: dict) -> NodeResponse: return NodeResponse(**node) @@ -480,7 +485,6 @@ async def distance_matrix( ) started = time.perf_counter() - graph_dict = await asyncio.to_thread(session.build_graph_dict) path_finder = session.path_finder n = len(body.node_ids) @@ -493,10 +497,21 @@ async def distance_matrix( src, tgt = body.node_ids[i], body.node_ids[j] try: if body.metric == "semantic" and session.similarity is not None: - sim = await asyncio.to_thread( - session.similarity.cosine_similarity, graph_dict, src, tgt - ) - val = 1.0 - float(sim) if isinstance(sim, (int, float)) else None + # Use cached embeddings for semantic distance calculation + embeddings = _get_cached_embeddings(session) + src_embedding = embeddings.get(src) + tgt_embedding = embeddings.get(tgt) + + if src_embedding is None or tgt_embedding is None: + val = None + else: + # Calculate cosine similarity directly from cached embeddings + import numpy as np + src_vec = np.array(src_embedding) + tgt_vec = np.array(tgt_embedding) + sim = np.dot(src_vec, tgt_vec) / (np.linalg.norm(src_vec) * np.linalg.norm(tgt_vec)) + val = 1.0 - float(sim) if isinstance(sim, (int, float)) else None + matrix[i][j] = val matrix[j][i] = val elif path_finder is not None: @@ -505,6 +520,7 @@ async def distance_matrix( if body.metric == "weighted" else path_finder.bfs_shortest_path ) + graph_dict = await asyncio.to_thread(session.build_graph_dict) result = await asyncio.to_thread(path_fn, graph_dict, src, tgt) path_nodes = result.get("path", []) if isinstance(result, dict) else (result or []) if path_nodes: @@ -550,8 +566,7 @@ async def _semantic_neighborhood_impl( detail="Semantic similarity is unavailable for this graph session.", ) - graph_dict = await asyncio.to_thread(session.build_graph_dict) - embeddings = _extract_node_embeddings(graph_dict) + embeddings = _get_cached_embeddings(session) query_embedding = embeddings.get(node_id) if not embeddings or query_embedding is None: raise HTTPException( diff --git a/semantica/explorer/session.py b/semantica/explorer/session.py index 059fdea7..b2ef8c57 100644 --- a/semantica/explorer/session.py +++ b/semantica/explorer/session.py @@ -52,6 +52,10 @@ class GraphSession: self._similarity: Any = None self._link_predictor: Any = None self._validator: Any = None + + self._graph_revision: int = 0 + self._cached_embeddings: Optional[Dict[str, List[float]]] = None + self._cached_graph_revision: int = -1 self.rebuild_search_index() @classmethod @@ -408,6 +412,20 @@ class GraphSession: def handle_graph_mutation(self, event_type: str, entity_id: str, payload: Dict[str, Any]) -> None: normalized_event = str(event_type or "").upper() + if normalized_event in { + "ADD_NODE", + "UPDATE_NODE", + "REMOVE_NODE", + "DELETE_NODE", + "ADD_EDGE", + "UPDATE_EDGE", + "REMOVE_EDGE", + "DELETE_EDGE", + "RELOAD_GRAPH", + "RESET_GRAPH", + }: + with self._lock: + self._bump_graph_revision_locked() if normalized_event in {"ADD_NODE", "UPDATE_NODE"}: normalized_node = self.normalize_node(payload or {}) if normalized_node.get("id"): @@ -529,6 +547,83 @@ class GraphSession: with self._lock: return self.annotations.pop(annotation_id, None) is not None + def _bump_graph_revision_locked(self) -> None: + self._graph_revision += 1 + self._cached_embeddings = None + self._cached_graph_revision = -1 + + @staticmethod + def _coerce_embedding_vector(value: Any) -> Optional[List[float]]: + if isinstance(value, dict): + for key in ("embedding", "embeddings", "vector", "values", "node2vec", "semantic"): + nested = GraphSession._coerce_embedding_vector(value.get(key)) + if nested is not None: + return nested + return None + + if not isinstance(value, (list, tuple)): + return None + + vector: List[float] = [] + for item in value: + try: + vector.append(float(item)) + except (TypeError, ValueError): + return None + + return vector if vector else None + + def get_cached_embeddings(self, force_refresh: bool = False) -> Dict[str, List[float]]: + with self._lock: + current_revision = self._graph_revision + if ( + not force_refresh + and self._cached_embeddings is not None + and self._cached_graph_revision == current_revision + ): + return self._cached_embeddings + raw_nodes = [ + node.to_dict() if hasattr(node, "to_dict") else node + for node in self.graph.nodes.values() + if node is not None + ] + + embedding_keys = ( + "embedding", + "embeddings", + "vector", + "node_embedding", + "node2vec_embedding", + "semantic_embedding", + "reasoning_embedding", + ) + + embeddings: Dict[str, List[float]] = {} + for raw in raw_nodes: + if not isinstance(raw, dict): + continue + normalized = self.normalize_node(raw) + node_id = normalized.get("id") + if not node_id: + continue + properties = normalized.get("properties") if isinstance(normalized.get("properties"), dict) else {} + for key in embedding_keys: + vector = self._coerce_embedding_vector(normalized.get(key, properties.get(key))) + if vector is not None: + embeddings[str(node_id)] = vector + break + + with self._lock: + if self._graph_revision == current_revision: + self._cached_embeddings = embeddings + self._cached_graph_revision = current_revision + return embeddings + + def invalidate_embedding_cache(self) -> None: + with self._lock: + self._cached_embeddings = None + self._cached_graph_revision = -1 + def build_graph_dict(self, node_ids: Optional[list] = None) -> dict: nodes, _ = self.get_nodes(skip=0, limit=999_999) edges, _ = self.get_edges(skip=0, limit=999_999) @@ -595,6 +690,8 @@ class GraphSession: with self._lock: added = self.graph.add_nodes(nodes) has_mutation_callback = callable(getattr(self.graph, "mutation_callback", None)) + if added and not has_mutation_callback: + self._bump_graph_revision_locked() if added and not has_mutation_callback: self.rebuild_search_index() return added @@ -603,6 +700,8 @@ class GraphSession: with self._lock: added = self.graph.add_edges(edges) has_mutation_callback = callable(getattr(self.graph, "mutation_callback", None)) + if added and not has_mutation_callback: + self._bump_graph_revision_locked() if added and not has_mutation_callback: self.rebuild_search_index() return added @@ -617,6 +716,8 @@ class GraphSession: with self._lock: added = self.graph.add_node(node_id, node_type, content=content, **properties) has_mutation_callback = callable(getattr(self.graph, "mutation_callback", None)) + if added and not has_mutation_callback: + self._bump_graph_revision_locked() if added and not has_mutation_callback: normalized = self.get_node(node_id) if normalized is not None: @@ -640,4 +741,6 @@ class GraphSession: **properties, ) has_mutation_callback = callable(getattr(self.graph, "mutation_callback", None)) + if added and not has_mutation_callback: + self._bump_graph_revision_locked() return added