scripts.semantic_search

Reference › scripts › semantic_search

scripts/semantic_search.py is the cheap retrieval command for the orchestrata-ultimate semantic graph (audit #6): a hybrid search over the code chunks that ingest_orchestrata_ultimate stored in Memgraph. A human or an agent runs it instead of grepping. It uses the Memgraph adapter, configuration and embedder factory in orchestrata_ultimate.semantic_graph (see orchestrata_ultimate.semantic_graph.config), and _load_dotenv loads the RUL-019 environment contract first.

Arguments: a positional query (natural language or an identifier), --top-k (default 5), --json and --backend (local or openrouter). The backend must match the one that wrote the index. The two backends produce vectors of different widths, so querying with the wrong one returns nothing rather than failing, and setting SEMANTIC_EMBED_BACKEND in the environment does not help because the profile overlay loaded at startup overrides it. _build_config builds the configuration, run_search embeds the query and returns the ranked results, and _print_results prints them as text or JSON. main runs the async entry point with asyncio.run. It reads the graph and writes nothing.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Functions

Source

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

#!/usr/bin/env python3
"""Cheap retrieval CLI for the orchestrata-ultimate semantic graph (audit #6).

Thin CLI (pattern: scripts/context_pack.py) around the hybrid vector +
keyword + module/entity-prominence search pipeline
(:meth:`SemanticGraphIngester.search`), so any agent can shell out for
context-injection without importing Python:

    python scripts/semantic_search.py "where is escalation handled" [--top-k 5] [--json]

Connection and hybrid-retrieval knobs are read from the environment at call
time via :class:`SemanticGraphConfig.from_env` (``MEMGRAPH_URI``,
``MEMGRAPH_CODESTRAL_CHUNK_LABEL``, ``SEMANTIC_GRAPH_HYBRID_*``, ...) — no
hardcoded connection defaults beyond localhost.
"""

from __future__ import annotations

import argparse
import asyncio
import json
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,
    requires_api_key,
    vector_index_name,
)
from orchestrata_ultimate.semantic_graph.ingestion import (
    CodeChunker,
    SemanticGraphIngester,
)
from orchestrata_ultimate.semantic_graph.rerankers import OpenRouterReranker

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


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.
    """
    load_env_contract()


def _parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Hybrid semantic-graph search over the orchestrata-ultimate codebase.",
    )
    parser.add_argument("query", help="Natural-language or identifier search query")
    parser.add_argument(
        "--top-k", type=int, default=5, help="Number of results to return (default: 5)"
    )
    parser.add_argument("--json", action="store_true", help="Print results as JSON")
    parser.add_argument(
        "--backend",
        choices=("local", "openrouter"),
        default=None,
        help=(
            "Embedding backend for THIS query, overriding the profile. It has "
            "to match whatever WROTE the index: the two backends produce "
            "different widths, so querying a 1536-wide index with the "
            "1024-wide one returns nothing rather than failing. Setting "
            "SEMANTIC_EMBED_BACKEND in the environment does not work here — "
            "the profile overlay loaded at startup overrides it."
        ),
    )
    return parser.parse_args(argv)


def _build_config(env: Mapping[str, str] | None = None) -> SemanticGraphConfig:
    config = SemanticGraphConfig.from_env()
    # 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. Whatever wrote the index has to be what queries it.
    config.chunk_label = chunk_label(env)
    config.vector_index_name = vector_index_name(env)
    return config


async def run_search(
    query: str, top_k: int, env: Mapping[str, str] | None = None
) -> list[dict[str, object]]:
    """Run the hybrid search pipeline and return result dicts.

    Args:
        query: Natural-language or identifier query.
        top_k: Number of results to return.
        env: Environment the backend selection is read from. Defaults to the
            process environment. Passed explicitly so the caller's choice
            cannot be overridden by a profile reload.
    """
    api_key = os.environ.get("OPENROUTER_API_KEY", "")
    if requires_api_key(env) and not api_key:
        raise SystemExit("OPENROUTER_API_KEY is not set")

    config = _build_config(env)
    # 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, env=env)
    reranker = OpenRouterReranker(
        api_key=api_key,
        model=RERANK_MODEL,
        base_url=os.environ.get(
            "OPENROUTER_RERANK_BASE_URL", "https://openrouter.ai/api/v1"
        ),
    )
    graph = MemgraphSemanticGraph(config)

    connected = await graph.connect()
    if not connected:
        raise SystemExit("Failed to connect to Memgraph")

    try:
        ingester = SemanticGraphIngester(
            graph=graph,
            embedder=embedder,
            chunker=CodeChunker(lines_per_chunk=50, overlap_lines=5),
            project_id=PROJECT_ID,
            reranker=reranker,
            rerank_candidates=20,
            config=config,
        )
        chunks, scores = await ingester.search(query, top_k=top_k)
        results: list[dict[str, object]] = [
            {
                "score": round(score, 4),
                "module": chunk.metadata.get("module", ""),
                "file_path": chunk.metadata.get("file_path", "unknown"),
                "start_line": chunk.metadata.get("start_line"),
                "end_line": chunk.metadata.get("end_line"),
                "snippet": chunk.content[:300],
                "explanation": chunk.metadata.get("retrieval_explanation", {}),
            }
            for chunk, score in zip(chunks, scores)
        ]
        return results
    finally:
        await graph.close()


def _print_results(results: list[dict[str, object]], as_json: bool) -> None:
    if as_json:
        print(json.dumps(results, indent=2))
        return
    if not results:
        print("No results.")
        return
    for idx, result in enumerate(results, start=1):
        snippet = str(result["snippet"]).replace("\n", " ")
        print(
            f"{idx}. [{result['score']}] {result['module']} — "
            f"{result['file_path']}:{result['start_line']}-{result['end_line']}"
        )
        print(f"   {snippet}")


async def _main_async(argv: list[str]) -> int:
    _load_dotenv()
    args = _parse_args(argv)
    # Built AFTER the dotenv load and passed by value, so the profile overlay
    # cannot discard it the way it discards an exported variable.
    env: Mapping[str, str] = (
        {**os.environ, "SEMANTIC_EMBED_BACKEND": args.backend}
        if args.backend
        else os.environ
    )
    results = await run_search(args.query, args.top_k, env)
    _print_results(results, args.json)
    return 0


def main() -> int:
    return asyncio.run(_main_async(sys.argv[1:]))


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