orchestrata_ultimate.semantic_graph.ingestion

Reference › orchestrata_ultimate › semantic_graph › ingestion

Source-code ingestion into the shared semantic graph (the memgraph database): it chunks files, embeds the chunks and stores them as vector-indexed nodes. Each chunk carries a project_id, so one graph can host several projects without collision. The context provider semantic_context and semantic_bridge import it for retrieval.

CodeChunker splits text files into overlapping line windows (50 lines with 5 of overlap by default, files over about 2 MB skipped). chunk_file returns the Chunk list for one file and chunk_directory does so for every file matching a glob. Chunk ids are deterministic (_chunk_id), and helper functions infer the module (_module_from_path) and language (_language_from_path).

SemanticGraphIngester takes a graph adapter, an embedders Embedder, a chunker, a project id, a batch size, an optional rerankers reranker and an optional SemanticGraphConfig. ingest_directory chunks, embeds in batches and stores, returning the count. search is hybrid retrieval: vector similarity and keyword matches (terms from _extract_query_terms) are fused with weighted reciprocal-rank fusion (_reciprocal_rank_fusion, k of 60) using the weights in config, then given small module and prominence boosts, and the top candidates can be reranked before the best are returned with their scores.

At a glance

Class diagram

Dependency diagram

API

Classes

CodeChunker

Line 106

Split text files into overlapping line-based chunks.

SemanticGraphIngester

Line 228

Ingest source-code chunks into a semantic graph using an embedder.

Functions

Source

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

"""Source-code ingestion into a semantic graph.

Chunks source files, embeds the chunks, and stores them as vector-indexed
nodes. Each chunk carries a ``project_id`` so one graph can host multiple
projects without collision.
"""

from __future__ import annotations

import logging
import math
import mimetypes
import re
from pathlib import Path
from typing import Any

from orchestrata_ultimate.semantic_graph.config import SemanticGraphConfig
from orchestrata_ultimate.semantic_graph.embedders import Embedder
from orchestrata_ultimate.semantic_graph.models import Chunk

logger = logging.getLogger(__name__)

# Stopwords excluded from keyword-search term extraction (mirrors the
# question-parsing heuristic in scripts/evaluate_semantic_graph_qa.py).
_QUERY_STOPWORDS = {
    "what",
    "does",
    "how",
    "why",
    "when",
    "where",
    "which",
    "must",
    "this",
    "that",
    "with",
    "from",
    "have",
    "has",
    "been",
    "than",
    "only",
    "each",
    "they",
    "them",
    "their",
    "there",
    "would",
    "should",
    "could",
    "about",
    "into",
    "also",
    "such",
    "being",
}


def _extract_query_terms(query: str, max_terms: int) -> list[str]:
    """Extract meaningful, deduplicated search terms from a query.

    Splits compound identifiers (CamelCase and snake_case) into their
    constituent words in addition to the whole identifier, so a query
    mentioning ``SwarmCoordinator`` also matches content containing "swarm"
    or "coordinator" alone. Stopwords and terms of length <= 3 are dropped.
    """
    seen: set[str] = set()
    terms: list[str] = []

    def _add(term: str) -> None:
        low = term.lower()
        if len(low) > 3 and low not in _QUERY_STOPWORDS and low not in seen:
            seen.add(low)
            terms.append(low)

    for word in re.findall(r"[A-Za-z0-9_]+", query):
        _add(word)
        for part in word.split("_"):
            if part:
                _add(part)
        for part in re.findall(r"[A-Z]?[a-z0-9]+|[A-Z]+(?![a-z])", word):
            _add(part)

    return terms[:max_terms]


def _reciprocal_rank_fusion(
    rankings: list[tuple[list[str], float]], rrf_k: float = 60.0
) -> dict[str, float]:
    """Weighted reciprocal-rank fusion across multiple ranked-id lists.

    Each ranking contributes ``weight / (rrf_k + rank + 1)`` to every id it
    contains, where ``rank`` is its 0-based position in that ranking. RRF is
    chosen over raw score blending because vector similarities and
    keyword-match counts live on incomparable scales; rank position is not.
    """
    scores: dict[str, float] = {}
    for ordered_ids, weight in rankings:
        if weight <= 0:
            continue
        for rank, chunk_id in enumerate(ordered_ids):
            scores[chunk_id] = scores.get(chunk_id, 0.0) + weight / (rrf_k + rank + 1)
    return scores


class CodeChunker:
    """Split text files into overlapping line-based chunks.

    The chunker is language-agnostic but records the file extension and MIME
    type in chunk metadata so downstream retrieval can filter by language.
    """

    def __init__(
        self,
        lines_per_chunk: int = 50,
        overlap_lines: int = 5,
        max_file_bytes: int = 2_000_000,
    ) -> None:
        """Initialize the chunker.

        Args:
            lines_per_chunk: Number of source lines in each chunk.
            overlap_lines: Lines shared between adjacent chunks.
            max_file_bytes: Files larger than this are skipped with a warning.
        """
        if lines_per_chunk < 1:
            raise ValueError("lines_per_chunk must be >= 1")
        if overlap_lines < 0 or overlap_lines >= lines_per_chunk:
            raise ValueError("overlap_lines must be >= 0 and < lines_per_chunk")
        self._lines_per_chunk = lines_per_chunk
        self._overlap_lines = overlap_lines
        self._max_file_bytes = max_file_bytes

    def chunk_file(
        self, path: Path, project_id: str, file_id_prefix: str = ""
    ) -> list[Chunk]:
        """Return chunks for a single source file.

        Args:
            path: Path to the file to chunk.
            project_id: Project identifier stored in chunk metadata.
            file_id_prefix: Optional prefix for chunk IDs.

        Returns:
            A list of :class:`Chunk` objects with line-range metadata.
        """
        if not path.is_file():
            raise FileNotFoundError(path)

        try:
            stat = path.stat()
        except OSError as exc:
            logger.warning("Cannot stat %s: %s", path, exc)
            return []

        if stat.st_size > self._max_file_bytes:
            logger.warning(
                "Skipping %s (%s bytes) because it exceeds max_file_bytes",
                path,
                stat.st_size,
            )
            return []

        try:
            text = path.read_text(encoding="utf-8", errors="replace")
        except OSError as exc:
            logger.warning("Cannot read %s: %s", path, exc)
            return []

        lines = text.splitlines()
        language = _language_from_path(path)
        mime_type, _ = mimetypes.guess_type(str(path))
        relative_path = str(path)
        module = _module_from_path(path)
        chunks: list[Chunk] = []
        step = max(1, self._lines_per_chunk - self._overlap_lines)

        for start in range(0, len(lines), step):
            end = min(start + self._lines_per_chunk, len(lines))
            content = "\n".join(lines[start:end])
            if not content.strip():
                continue
            chunk_id = _chunk_id(file_id_prefix, relative_path, start, end)
            metadata: dict[str, Any] = {
                "project_id": project_id,
                "file_path": relative_path,
                "module": module,
                "start_line": start + 1,
                "end_line": end,
                "language": language,
                "mime_type": mime_type or "",
            }
            chunks.append(
                Chunk(id=chunk_id, content=content, embedding=None, metadata=metadata)
            )
            if end == len(lines):
                break

        return chunks

    def chunk_directory(
        self,
        directory: Path,
        project_id: str,
        pattern: str = "**/*.py",
    ) -> list[Chunk]:
        """Chunk all files matching ``pattern`` under ``directory``.

        Args:
            directory: Root directory to search.
            project_id: Project identifier stored in chunk metadata.
            pattern: Glob pattern passed to :meth:`Path.glob`.

        Returns:
            A flat list of chunks from all matching files.
        """
        if not directory.is_dir():
            raise NotADirectoryError(directory)

        chunks: list[Chunk] = []
        for path in directory.glob(pattern):
            if not path.is_file():
                continue
            chunks.extend(self.chunk_file(path, project_id))
        return chunks


class SemanticGraphIngester:
    """Ingest source-code chunks into a semantic graph using an embedder."""

    def __init__(
        self,
        graph: Any,
        embedder: Embedder,
        chunker: CodeChunker,
        project_id: str,
        batch_size: int = 64,
        reranker: Any | None = None,
        rerank_candidates: int = 20,
        config: SemanticGraphConfig | None = None,
    ) -> None:
        """Initialize the ingester.

        Args:
            graph: A semantic graph adapter (e.g. :class:`MemgraphSemanticGraph`).
            embedder: Embedding backend implementing :class:`Embedder`.
            chunker: Chunker that produces :class:`Chunk` objects.
            project_id: Project identifier written into every chunk.
            batch_size: Number of chunks to embed and insert in one batch.
            reranker: Optional reranker to reorder vector-search candidates.
            rerank_candidates: Number of vector-search candidates to rerank.
            config: Hybrid-retrieval knobs (vector/keyword weights, module and
                prominence boosts). Defaults to :meth:`SemanticGraphConfig.from_env`.
        """
        self._graph = graph
        self._embedder = embedder
        self._chunker = chunker
        self._project_id = project_id
        self._batch_size = max(1, batch_size)
        self._reranker = reranker
        self._rerank_candidates = max(1, rerank_candidates)
        config = config or SemanticGraphConfig.from_env()
        self._vector_weight = config.hybrid_vector_weight
        self._keyword_weight = config.hybrid_keyword_weight
        self._keyword_term_cap = config.hybrid_keyword_term_cap
        self._keyword_candidate_cap = config.hybrid_keyword_candidate_cap
        self._keyword_extra_cap = config.hybrid_keyword_extra_cap
        self._module_boost = config.hybrid_module_boost
        self._prominence_boost = config.hybrid_prominence_boost

    async def ingest_directory(
        self,
        directory: Path,
        pattern: str = "**/*.py",
    ) -> int:
        """Chunk, embed, and store all matching files.

        Args:
            directory: Root directory to search.
            pattern: Glob pattern for source files.

        Returns:
            Number of chunks stored.
        """
        chunks = self._chunker.chunk_directory(directory, self._project_id, pattern)
        if not chunks:
            logger.warning("No chunks found in %s for pattern %s", directory, pattern)
            return 0

        for chunk in chunks:
            chunk.metadata.setdefault("project_id", self._project_id)

        total = 0
        for i in range(0, len(chunks), self._batch_size):
            batch = chunks[i : i + self._batch_size]
            texts = [chunk.content for chunk in batch]
            embeddings = await self._embedder.embed(texts, prompt_type="document")
            if len(embeddings) != len(batch):
                raise RuntimeError(
                    f"Embedder returned {len(embeddings)} vectors for {len(batch)} chunks"
                )
            for chunk, embedding in zip(batch, embeddings):
                chunk.embedding = embedding
            await self._graph.add_chunks(batch)
            total += len(batch)
            logger.info(
                "Ingested %s/%s chunks for project %s",
                total,
                len(chunks),
                self._project_id,
            )

        return total

    async def search(
        self, query: str, top_k: int = 5
    ) -> tuple[list[Chunk], list[float]]:
        """Hybrid retrieval: vector + keyword + module/entity-prominence boosts.

        Pipeline:

        1. Vector pass: embed the query and fetch the top ``rerank_candidates``
           (or ``top_k``, whichever is larger) chunks by embedding similarity.
        2. Keyword pass: extract meaningful terms from the query (splitting
           CamelCase/snake_case identifiers) and, if the graph adapter
           supports it, fetch chunks whose content contains any term via a
           per-term ``CONTAINS`` query.
        3. Fusion: the two rankings are combined with weighted reciprocal-rank
           fusion (RRF) — chosen because vector similarity and keyword-match
           counts are on incomparable scales, while rank position is not.
        4. Module boost: chunks belonging to a ``Module`` whose name matches a
           query term get an additive score boost.
        5. Prominence boost: chunks linked (BELONGS_TO/MENTIONS) to an
           ``Entity`` carrying a ``prominence`` property get an additive boost
           proportional to ``log(prominence + 1)``.
        6. Rerank (if configured) is the final stage, applied to the fused
           top pool.

        Every chunk returned carries a ``retrieval_explanation`` metadata key
        (fusion score, module boost, prominence + its boost) for observability.

        Args:
            query: Raw query text.
            top_k: Maximum number of results.

        Returns:
            Tuple of (chunks, scores) in matching order. Scores are the fused
            (pre-rerank) relevance scores; reranking reorders chunks without
            recomputing them.
        """
        embeddings = await self._embedder.embed([query], prompt_type="query")
        if not embeddings:
            return ([], [])

        candidate_k = max(top_k, self._rerank_candidates)
        vector_chunks = await self._graph.search_similar(
            embeddings[0], top_k=candidate_k
        )
        vector_ids = [chunk.id for chunk in vector_chunks]
        vector_id_set = set(vector_ids)

        terms = _extract_query_terms(query, self._keyword_term_cap)
        keyword_pairs: list[tuple[Chunk, int]] = []
        if terms and hasattr(self._graph, "search_keyword"):
            keyword_pairs = await self._graph.search_keyword(
                terms, limit_per_term=self._keyword_candidate_cap
            )
        keyword_pairs.sort(key=lambda pair: pair[1], reverse=True)
        keyword_ids = [chunk.id for chunk, _count in keyword_pairs]

        # The vector-retrieved pool is kept in full — a crude CONTAINS keyword
        # pass can match many chunks on a common word, and letting it fully
        # re-rank (and thus potentially evict) the vector pool would starve
        # the downstream reranker of genuinely relevant candidates. Instead,
        # keyword hits the vector pass missed are added as a small, bounded
        # supplement (``hybrid_keyword_extra_cap``); keyword match counts
        # otherwise only feed the fused *score* used for the no-reranker
        # ordering and the module/prominence boosts below.
        extra_keyword_ids = [cid for cid in keyword_ids if cid not in vector_id_set][
            : self._keyword_extra_cap
        ]
        pool_ids = vector_ids + extra_keyword_ids

        chunk_by_id: dict[str, Chunk] = {chunk.id: chunk for chunk in vector_chunks}
        for chunk, _count in keyword_pairs:
            chunk_by_id.setdefault(chunk.id, chunk)

        fused_scores = _reciprocal_rank_fusion(
            [
                (vector_ids, self._vector_weight),
                (keyword_ids, self._keyword_weight),
            ]
        )
        explanations: dict[str, dict[str, float]] = {
            cid: {"fusion_score": fused_scores.get(cid, 0.0)} for cid in pool_ids
        }

        if terms and pool_ids and hasattr(self._graph, "matching_module_names"):
            matching_modules = await self._graph.matching_module_names(terms)
            if matching_modules:
                for cid in pool_ids:
                    module = str(chunk_by_id[cid].metadata.get("module", "")).lower()
                    if module and module in matching_modules:
                        fused_scores[cid] = (
                            fused_scores.get(cid, 0.0) + self._module_boost
                        )
                        explanations[cid]["module_boost"] = self._module_boost

        if pool_ids and hasattr(self._graph, "chunk_prominence"):
            prominences = await self._graph.chunk_prominence(pool_ids)
            for cid, prominence in prominences.items():
                if prominence > 0 and cid in explanations:
                    boost = self._prominence_boost * math.log(prominence + 1.0)
                    fused_scores[cid] = fused_scores.get(cid, 0.0) + boost
                    explanations[cid]["prominence"] = prominence
                    explanations[cid]["prominence_boost"] = boost

        pool_chunks = [chunk_by_id[cid] for cid in pool_ids]

        if self._reranker is not None and len(pool_chunks) > 1:
            ranked_chunks = await self._reranker.rerank(query, pool_chunks, top_k=top_k)
        else:
            ranked_ids = sorted(
                pool_ids, key=lambda cid: fused_scores.get(cid, 0.0), reverse=True
            )
            ranked_chunks = [chunk_by_id[cid] for cid in ranked_ids[:top_k]]

        for chunk in ranked_chunks:
            chunk.metadata["retrieval_explanation"] = explanations.get(chunk.id, {})

        scores = [fused_scores.get(chunk.id, 0.0) for chunk in ranked_chunks]
        return (ranked_chunks, scores)


def _module_from_path(path: Path) -> str:
    """Infer a module/package name from the file's parent directories.

    For source trees like ``src/orchestrata_ultimate/gateway/server.py`` this
    returns ``gateway``. Files directly under the project root return the
    immediate parent directory name.
    """
    parts = path.parts
    if len(parts) >= 2:
        return parts[-2]
    return ""


def _language_from_path(path: Path) -> str:
    """Infer a simple language tag from the file extension."""
    mapping = {
        ".py": "python",
        ".pyi": "python",
        ".js": "javascript",
        ".ts": "typescript",
        ".tsx": "typescript",
        ".jsx": "javascript",
        ".md": "markdown",
        ".rst": "rst",
        ".yaml": "yaml",
        ".yml": "yaml",
        ".json": "json",
        ".toml": "toml",
        ".sh": "shell",
        ".pl": "prolog",
        ".cpp": "cpp",
        ".c": "c",
        ".h": "c",
        ".rs": "rust",
        ".go": "go",
    }
    return mapping.get(path.suffix.lower(), "text")


def _chunk_id(prefix: str, file_path: str, start_line: int, end_line: int) -> str:
    """Return a deterministic, URL-safe chunk identifier."""
    safe_path = file_path.replace("\\", "/").replace(" ", "%20")
    base = (
        f"{prefix}:{safe_path}:{start_line}:{end_line}"
        if prefix
        else f"{safe_path}:{start_line}:{end_line}"
    )
    # Keep IDs deterministic and avoid very long identifiers by hashing the
    # descriptive base when it exceeds a reasonable length.
    if len(base) > 200:
        import hashlib

        digest = hashlib.sha256(base.encode()).hexdigest()[:16]
        short_path = safe_path.replace("/", "_")[-80:]
        base = (
            f"{prefix}:{short_path}:{start_line}:{end_line}:{digest}"
            if prefix
            else f"{short_path}:{start_line}:{end_line}:{digest}"
        )
    return base