scripts.ingest_orchestrata_ultimate

Reference › scripts › ingest_orchestrata_ultimate

scripts/ingest_orchestrata_ultimate.py ingests the orchestrata-ultimate source code into the Memgraph semantic graph: it chunks the source, embeds the chunks and stores them for semantic search (see semantic_search). A human runs it; it is long-running. It is built on orchestrata_ultimate.semantic_graph.adapters.memgraph, orchestrata_ultimate.semantic_graph.config, orchestrata_ultimate.semantic_graph.embedder_factory and orchestrata_ultimate.core.env_bootstrap.

main is async. _load_dotenv loads the RUL-019 two-phase environment contract first. Options: --backend chooses the embedding backend for this run (local or openrouter), overriding the profile; it is passed as an explicit environment mapping, because setting SEMANTIC_EMBED_BACKEND in the process environment loses to the profile overlay loaded at startup. --prune-stale-labels scans for stale chunk nodes (labels SemanticChunk, QwenChunk, Chunk carrying the orchestrata chunk-id shape) and reports them without touching unrelated wiki-graph Chunk nodes; it is a dry run unless --apply is added.

An ingest run exits 1 when the source directory is missing, or when the hosted backend is selected and OPENROUTER_API_KEY is not set; the local backend needs no key. It logs the chosen backend, model and dimension before starting, because an earlier run lost fifteen minutes on the wrong model. _ingest does the work and _verify checks the result.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Constants

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/ingest_orchestrata_ultimate.py

"""Ingest orchestrata-ultimate source code into the Memgraph semantic graph.

Uses the highest-quality online models for code RAG:

* Embedding: ``mistralai/codestral-embed-2505`` via OpenRouter (1536 dims).
* Reranking: ``cohere/rerank-4-pro`` via OpenRouter.

Each chunk is tagged with ``project_id: orchestrata-ultimate`` so the graph can
host multiple projects.
"""

from __future__ import annotations

import argparse
import asyncio
import logging
import os
import sys
from collections.abc import Mapping
from pathlib import Path

from orchestrata_ultimate.core.env_bootstrap import load_env_contract
from orchestrata_ultimate.semantic_graph.adapters.memgraph import (
    MemgraphSemanticGraph,
)
from orchestrata_ultimate.semantic_graph.config import SemanticGraphConfig
from orchestrata_ultimate.semantic_graph.embedder_factory import (
    chunk_label,
    create_embedder,
    embed_backend,
    embed_dimension,
    embed_model,
    requires_api_key,
    vector_index_name,
)
from orchestrata_ultimate.semantic_graph.embedders import Embedder
from orchestrata_ultimate.semantic_graph.ingestion import (
    CodeChunker,
    SemanticGraphIngester,
)
from orchestrata_ultimate.semantic_graph.rerankers import OpenRouterReranker

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)-8s %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

PROJECT_ROOT = Path(__file__).resolve().parent.parent
SOURCE_DIR = PROJECT_ROOT / "src" / "orchestrata_ultimate"
PROJECT_ID = "orchestrata-ultimate"
RERANK_MODEL = "cohere/rerank-4-pro"

# File patterns that make sense for code RAG on this project.
PATTERNS = [
    "**/*.py",
    "**/*.pyi",
    "**/*.md",
]

# Additional documentation roots outside ``src/orchestrata_ultimate``. These
# answer architecture, process, and standards questions that source-code chunks
# alone cannot cover. Each entry is ``(Path, [patterns])``.
EXTRA_DOC_ROOTS: list[tuple[Path, list[str]]] = [
    (PROJECT_ROOT, ["AGENTS.md", "README.md"]),
    (PROJECT_ROOT / "docs", ["**/*.md"]),
    (PROJECT_ROOT / "architecture", ["**/*.md"]),
    (PROJECT_ROOT / "design", ["**/*.md"]),
    (
        PROJECT_ROOT / "plans",
        ["execution-plan.md", "implementation-plan.md", "class-mapping.json"],
    ),
    (PROJECT_ROOT / "tests" / "semantic_graph", ["**/*.py"]),
]

# Sample queries used to verify retrieval quality after ingestion.
SAMPLE_QUERIES = [
    "How does the gateway handle WebSocket connections?",
    "What is the Memgraph semantic graph adapter?",
    "How are agents created and configured?",
    "Explain the session controller and checkpoints.",
]


def _graph_config(
    base: SemanticGraphConfig, env: Mapping[str, str] | None = None
) -> SemanticGraphConfig:
    return SemanticGraphConfig(
        backend=base.backend,
        uri=base.uri,
        username=base.username,
        password=base.password,
        vector_index_name=vector_index_name(env),
        vector_index_metric=base.vector_index_metric,
        vector_index_capacity=base.vector_index_capacity,
        # Label and index follow the SELECTED backend: they are what say
        # "these vectors are comparable to each other", and this graph is
        # shared with the knowledge-wiki.
        chunk_label=chunk_label(env),
        ha_controller_url=base.ha_controller_url,
        uris=list(base.uris),
    )


async def _ingest(graph: MemgraphSemanticGraph, embedder: Embedder) -> int:
    connected = await graph.connect()
    if not connected:
        logger.error("Failed to connect to Memgraph")
        return 0

    await graph.clear()

    try:
        chunker = CodeChunker(lines_per_chunk=50, overlap_lines=5)
        ingester = SemanticGraphIngester(
            graph=graph,
            embedder=embedder,
            chunker=chunker,
            project_id=PROJECT_ID,
            batch_size=64,
        )
        total = 0
        for pattern in PATTERNS:
            count = await ingester.ingest_directory(SOURCE_DIR, pattern=pattern)
            total += count
            logger.info("Ingested %s chunks matching %s", count, pattern)

        for root, patterns in EXTRA_DOC_ROOTS:
            if not root.exists():
                logger.warning("Skipping missing doc root: %s", root)
                continue
            for pattern in patterns:
                count = await ingester.ingest_directory(root, pattern=pattern)
                total += count
                logger.info(
                    "Ingested %s chunks from %s matching %s", count, root, pattern
                )

        if total > 0:
            logger.info("Linking chunks to module nodes...")
            await graph.link_chunks_to_modules()
        logger.info("Total chunks: %s", total)
        return total
    finally:
        await graph.close()


async def _verify(
    graph: MemgraphSemanticGraph,
    embedder: Embedder,
    reranker: OpenRouterReranker | None,
) -> None:
    """Sample-query the freshly written index.

    ``reranker`` is optional because it needs an OpenRouter account and the
    local embedding backend does not. Without it the sample searches still
    run: reranking reorders candidates, it does not produce them.
    """
    connected = await graph.connect()
    if not connected:
        logger.error("Cannot verify: Memgraph not reachable")
        return

    try:
        ingester = SemanticGraphIngester(
            graph=graph,
            embedder=embedder,
            chunker=CodeChunker(lines_per_chunk=50, overlap_lines=5),
            project_id=PROJECT_ID,
            batch_size=64,
            reranker=reranker,
            rerank_candidates=20,
        )
        for query in SAMPLE_QUERIES:
            chunks, _ = await ingester.search(query, top_k=3)
            logger.info("Query: %s", query)
            for idx, chunk in enumerate(chunks, start=1):
                file_path = chunk.metadata.get("file_path", "unknown")
                logger.info(
                    "  %s. %s:%s-%s",
                    idx,
                    file_path,
                    chunk.metadata.get("start_line"),
                    chunk.metadata.get("end_line"),
                )
    finally:
        await graph.close()


def _load_dotenv() -> None:
    """Load the RUL-019 two-phase environment contract.

    Delegates to ``core.env_bootstrap.load_env_contract`` — THE canonical
    loader. What stood here loaded ``.env`` and stopped, and ``.env`` is the
    prod-BASE layer: it carries ``CXP_GRPC_HOST=localhost`` for the host
    user-units, so a script that stops after phase 1 resolves PRODUCTION
    endpoints while believing it is on the test fleet.

    ``override=False`` was the right instinct aimed at the wrong half. It was
    added on 2026-07-26 after a bootstrap intended for the test node wrote 592
    chunks into the production cluster, because ``.env``'s value had replaced
    an explicitly-exported one. The canonical loader keeps that property for
    phase 1 (``.env`` only fills keys the environment has not already set) and
    then applies the ``ORCHESTRATA_ENV_PROFILE`` overlay, which is the half
    that actually points the process at the right infrastructure — and honours
    ``declare_owned`` so it cannot reset a test worker's per-slot isolation.
    """
    # Rooted at this module's PROJECT_ROOT rather than letting the loader
    # resolve the repo itself. The missing-key tests redirect PROJECT_ROOT
    # at a tmp_path precisely so this call cannot re-read the real .env and
    # put OPENROUTER_API_KEY back; a self-resolving loader walked past the
    # redirect, reloaded the key, and turned a no-key assertion into a live
    # billed request to OpenRouter.
    load_env_contract(start=PROJECT_ROOT)


def _parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Ingest orchestrata-ultimate source code into the Memgraph "
            "semantic graph, or prune stale chunk-label nodes (audit #3)."
        ),
    )
    parser.add_argument(
        "--prune-stale-labels",
        action="store_true",
        help=(
            "Instead of ingesting, scan for stale chunk nodes (SemanticChunk/"
            "QwenChunk/Chunk carrying the orchestrata chunk-id shape) and "
            "report them. Never touches the unrelated wiki-graph Chunk nodes. "
            "Default is a dry-run report; pass --apply to actually delete."
        ),
    )
    parser.add_argument(
        "--apply",
        action="store_true",
        help="With --prune-stale-labels, actually delete the stale nodes found.",
    )
    parser.add_argument(
        "--backend",
        choices=("local", "openrouter"),
        default=None,
        help=(
            "Embedding backend for THIS run, overriding the profile. Setting "
            "SEMANTIC_EMBED_BACKEND in the environment does not work here: "
            "the profile overlay loaded at startup overrides the process "
            "environment, so the profile always won and the run silently "
            "used a backend nobody asked for. Defaults to the profile."
        ),
    )
    return parser.parse_args(argv)


async def _prune_stale_labels(graph: MemgraphSemanticGraph, apply: bool) -> int:
    connected = await graph.connect()
    if not connected:
        logger.error("Failed to connect to Memgraph")
        return 1

    try:
        stale = await graph.find_stale_chunk_nodes()
        by_label: dict[str, int] = {}
        for node in stale:
            by_label[node["label"]] = by_label.get(node["label"], 0) + 1

        if not stale:
            logger.info("No stale chunk nodes found.")
            return 0

        logger.info("Stale chunk nodes found (scoped away from the wiki graph):")
        for label, count in by_label.items():
            logger.info("  %s: %s node(s)", label, count)

        if apply:
            deleted = await graph.prune_stale_chunk_nodes(apply=True)
            logger.info("Deleted %s stale node(s).", deleted)
        else:
            logger.info(
                "Dry run only (%s node(s) found) — pass --apply to delete them.",
                len(stale),
            )
        return 0
    finally:
        await graph.close()


async def main() -> int:
    _load_dotenv()
    args = _parse_args(sys.argv[1:])

    # Passed explicitly rather than exported: every reader below takes an
    # env mapping, so there is no process-wide mutation and no "must happen
    # after the dotenv load" ordering rule to get wrong. Setting
    # SEMANTIC_EMBED_BACKEND in the real environment would not have worked
    # anyway — the profile overlay overrides it during startup.
    env: Mapping[str, str] = (
        {**os.environ, "SEMANTIC_EMBED_BACKEND": args.backend}
        if args.backend
        else os.environ
    )

    if args.prune_stale_labels:
        base_config = SemanticGraphConfig.from_env()
        graph = MemgraphSemanticGraph(base_config)
        return await _prune_stale_labels(graph, apply=args.apply)

    if not SOURCE_DIR.is_dir():
        logger.error("Source directory not found: %s", SOURCE_DIR)
        return 1

    base_config = SemanticGraphConfig.from_env()
    api_key = os.environ.get("OPENROUTER_API_KEY", "")
    # Only the hosted backend needs an account. Demanding a key regardless
    # made the local backend — whose whole point is needing none — refuse to
    # start on a machine that never had one.
    if requires_api_key(env) and not api_key:
        logger.error("OPENROUTER_API_KEY is not set")
        return 1

    # Say which backend is about to run, before the work starts. The run this
    # guards against burned 15 minutes of a multi-hour job on the wrong model
    # and reported nothing until it was killed.
    logger.info(
        "Embedding backend: %s (model=%s, dim=%s) -> label=%s index=%s",
        embed_backend(env),
        embed_model(env),
        embed_dimension(env),
        chunk_label(env),
        vector_index_name(env),
    )

    # Conservative token ceiling: markdown docs can contain long tables and
    # prose paragraphs, and OpenRouter rejects batches whose total token count
    # is too high. The embedder further subdivides by estimated tokens.
    #
    # OPENROUTER_EMBED_BASE_URL/OPENROUTER_RERANK_BASE_URL: read at call
    # time, defaulting to the real OpenRouter API (the classes' own
    # defaults) -- lets tests redirect these two HTTP round trips at local
    # test servers without touching the real API (policy ruling 2026-08-04:
    # tests make no LLM/embedding network calls).
    # The backend is a choice, not a constant (semantic_graph.embedder_factory):
    # dev/test embed locally for free, prod keeps codestral via OpenRouter. The
    # model and dimension come from the same place so an index is always
    # queried by whatever wrote it.
    embedder = create_embedder(
        api_key=api_key, batch_size=64, max_tokens_per_batch=8_000, env=env
    )
    # Optional, and for the same reason the key check above is conditional:
    # the reranker needs an account, the local embedding backend does not,
    # and an unconditional constructor here re-imposed the requirement the
    # check had just lifted. Absent is announced, never silent.
    reranker: OpenRouterReranker | None = None
    if api_key:
        reranker = OpenRouterReranker(
            api_key=api_key,
            model=RERANK_MODEL,
            base_url=os.environ.get(
                "OPENROUTER_RERANK_BASE_URL", "https://openrouter.ai/api/v1"
            ),
        )
    else:
        logger.info("No OPENROUTER_API_KEY: verification searches will run unreranked.")
    graph = MemgraphSemanticGraph(_graph_config(base_config, env))

    count = await _ingest(graph, embedder)
    if count == 0:
        logger.error("No chunks were ingested")
        return 1

    logger.info("Ingestion complete: %s chunks", count)
    await _verify(graph, embedder, reranker)
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))