orchestrata_ultimate.semantic_graph.adapters.memory

Reference › orchestrata_ultimate › semantic_graph › adapters › memory

InMemorySemanticGraph is the pure-Python implementation of the SemanticGraph protocol, used for fast unit tests and local development; it never touches the memgraph database or any other store. It holds entities, relationships and chunks in dictionaries and implements the same methods as the production memgraph adapter. connect always reports success and close does nothing. upsert_entities, upsert_relationships and get_neighbors (breadth-first to a depth) cover the graph side. add_chunks and search_similar cover the vector side, ranking by _cosine_similarity, which raises ValueError for vectors of different length. The hybrid-retrieval helpers search_keyword, matching_module_names, chunk_prominence (from BELONGS_TO and MENTIONS relationships), chunks_with_id_prefix and delete_chunks mirror the Memgraph adapter so ingestion code behaves the same. retrieve assembles a GraphRAGContext, clear empties the store and dump returns a snapshot for debugging. repository registers it as the memory backend.

At a glance

Class diagram

Dependency diagram

API

Classes

InMemorySemanticGraph

Line 29

In-memory implementation of the SemanticGraph interface.

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/semantic_graph/adapters/memory.py

"""In-memory semantic graph adapter for fast unit tests and local dev."""

from __future__ import annotations

import math
from typing import Any

from orchestrata_ultimate.semantic_graph.config import SemanticGraphConfig
from orchestrata_ultimate.semantic_graph.models import (
    Chunk,
    Entity,
    GraphRAGContext,
    Relationship,
)


def _cosine_similarity(a: list[float], b: list[float]) -> float:
    """Compute cosine similarity between two equal-length vectors."""
    if len(a) != len(b):
        raise ValueError("Vectors must have the same dimension")
    dot = sum(x * y for x, y in zip(a, b))
    norm_a = math.sqrt(sum(x * x for x in a))
    norm_b = math.sqrt(sum(x * x for x in b))
    if norm_a == 0.0 or norm_b == 0.0:
        return 0.0
    return dot / (norm_a * norm_b)


class InMemorySemanticGraph:
    """In-memory implementation of the SemanticGraph interface.

    This adapter is intentionally simple: it stores entities, relationships,
    and chunks in plain dicts and computes vector similarity brute-force.
    It is useful for unit tests and local development without containers.
    """

    def __init__(self, config: SemanticGraphConfig | None = None) -> None:
        self._config = config
        self._entities: dict[str, Entity] = {}
        self._relationships: dict[str, Relationship] = {}
        self._chunks: dict[str, Chunk] = {}

    async def connect(self) -> bool:
        """In-memory store is always available."""
        return True

    async def close(self) -> None:
        """No-op for in-memory store."""
        return

    async def upsert_entities(self, entities: list[Entity]) -> None:
        for entity in entities:
            self._entities[entity.id] = entity

    async def upsert_relationships(self, relationships: list[Relationship]) -> None:
        for rel in relationships:
            self._relationships[rel.id] = rel

    async def get_neighbors(
        self, entity_id: str, depth: int = 1
    ) -> tuple[list[Entity], list[Relationship]]:
        if depth < 1:
            return ([], [])

        visited_ids: set[str] = {entity_id}
        visited_rel_ids: set[str] = set()
        rels_out: list[Relationship] = []
        queue = {entity_id}

        for _ in range(depth):
            next_queue: set[str] = set()
            for rel in self._relationships.values():
                if rel.id in visited_rel_ids:
                    continue
                if rel.source_id in queue or rel.target_id in queue:
                    visited_rel_ids.add(rel.id)
                    rels_out.append(rel)
                    for node_id in (rel.source_id, rel.target_id):
                        if node_id not in visited_ids:
                            visited_ids.add(node_id)
                            next_queue.add(node_id)
            queue = next_queue

        entities_out = [
            self._entities[eid] for eid in visited_ids if eid in self._entities
        ]
        return (entities_out, rels_out)

    async def add_chunks(self, chunks: list[Chunk]) -> None:
        for chunk in chunks:
            self._chunks[chunk.id] = chunk

    async def search_similar(
        self, embedding: list[float], top_k: int = 5
    ) -> list[Chunk]:
        scored: list[tuple[float, Chunk]] = []
        for chunk in self._chunks.values():
            if chunk.embedding is None:
                continue
            try:
                sim = _cosine_similarity(embedding, chunk.embedding)
            except ValueError:
                continue
            scored.append((sim, chunk))
        scored.sort(key=lambda item: item[0], reverse=True)
        return [chunk for _, chunk in scored[:top_k]]

    async def search_keyword(
        self, terms: list[str], limit_per_term: int = 50
    ) -> list[tuple[Chunk, int]]:
        """Return chunks whose content contains any of ``terms``, with a
        per-chunk count of distinct matching terms (pure in-memory mirror of
        :meth:`MemgraphSemanticGraph.search_keyword`, used to unit-test the
        fusion logic without live infrastructure).
        """
        if not terms:
            return []
        lowered_terms = [t.lower() for t in terms if t]
        scored: dict[str, tuple[Chunk, int]] = {}
        for chunk in self._chunks.values():
            content_lower = chunk.content.lower()
            count = sum(1 for term in lowered_terms if term in content_lower)
            if count:
                scored[chunk.id] = (chunk, count)
        return list(scored.values())

    async def matching_module_names(self, terms: list[str]) -> set[str]:
        """Return lower-cased module names (from chunk metadata) matching any term."""
        if not terms:
            return set()
        lowered_terms = {t.lower() for t in terms if t}
        modules = {
            str(chunk.metadata.get("module", "")).lower()
            for chunk in self._chunks.values()
            if chunk.metadata.get("module")
        }
        return modules & lowered_terms

    async def chunk_prominence(self, chunk_ids: list[str]) -> dict[str, float]:
        """Return ``chunk_id -> prominence`` via BELONGS_TO/MENTIONS relationships
        to an Entity carrying a ``prominence`` property.
        """
        if not chunk_ids:
            return {}
        ids = set(chunk_ids)
        prominences: dict[str, float] = {}
        for rel in self._relationships.values():
            if rel.type not in ("BELONGS_TO", "MENTIONS") or rel.source_id not in ids:
                continue
            entity = self._entities.get(rel.target_id)
            if entity is None or "prominence" not in entity.properties:
                continue
            prominence = float(entity.properties["prominence"])
            if (
                rel.source_id not in prominences
                or prominence > prominences[rel.source_id]
            ):
                prominences[rel.source_id] = prominence
        return prominences

    async def chunks_with_id_prefix(self, prefixes: list[str]) -> list[Chunk]:
        """Return chunks whose id starts with any of ``prefixes`` (pure
        in-memory mirror of :meth:`MemgraphSemanticGraph.chunks_with_id_prefix`).
        """
        if not prefixes:
            return []
        return [
            chunk
            for chunk in self._chunks.values()
            if any(chunk.id.startswith(p) for p in prefixes)
        ]

    async def delete_chunks(self, ids: list[str]) -> None:
        """Delete chunks by exact id (pure in-memory mirror of
        :meth:`MemgraphSemanticGraph.delete_chunks`).
        """
        for chunk_id in ids:
            self._chunks.pop(chunk_id, None)

    async def clear(self) -> None:
        self._entities.clear()
        self._relationships.clear()
        self._chunks.clear()

    async def retrieve(
        self,
        query_embedding: list[float],
        top_k: int = 5,
        graph_depth: int = 1,
    ) -> GraphRAGContext:
        chunks = await self.search_similar(query_embedding, top_k)
        entities: dict[str, Entity] = {}
        relationships: dict[str, Relationship] = {}

        for chunk in chunks:
            # Simple heuristic: look for entity ids/names inside chunk content.
            for entity in self._entities.values():
                if entity.id in chunk.content or any(
                    str(value) in chunk.content
                    for value in entity.properties.values()
                    if isinstance(value, str)
                ):
                    entities[entity.id] = entity
                    ents, rels = await self.get_neighbors(entity.id, graph_depth)
                    for e in ents:
                        entities[e.id] = e
                    for r in rels:
                        relationships[r.id] = r

        return GraphRAGContext(
            entities=list(entities.values()),
            relationships=list(relationships.values()),
            chunks=chunks,
        )

    def dump(self) -> dict[str, Any]:
        """Return a snapshot of the in-memory graph for debugging."""
        return {
            "entities": list(self._entities.values()),
            "relationships": list(self._relationships.values()),
            "chunks": list(self._chunks.values()),
        }