orchestrata_ultimate.knowledge_graph.kg_context

Reference › orchestrata_ultimate › knowledge_graph › kg_context

KGCodeContext gives volks their code context from the operational graph (orchkg on orchgraph) instead of from disk; the volk base class is its importer. Code is synced into the graph after generation finishes and every other volk reads it back from there. The neo4j imports are Bolt-client naming only.

The class is built from a Bolt URI and credentials, or from a shared BoltSyncRepository, which is then reused so connections are pooled across the spike components and not owned by this object. connect verifies a shared repository or builds and owns a new one with a 5-second connection timeout, and returns False on failure instead of raising. The read methods are get_file_content (a file's signature summary from its entities), get_project_files, get_source_listing and get_source_listing_for_files, get_classes_and_functions, get_dependencies, get_complexity_hotspots (functions above a complexity threshold), get_call_graph_neighborhood (callers and callees within a number of hops) and get_monorepo_classes. store_design_doc and get_design_docs keep design documents linked to a project. sync_files_to_kg reads source files from disk, writes File nodes and uses _sync_entities, a regex-based extractor, to create class and function nodes under ids from _entity_id that follow canonical_schema. Project ids follow project_resolver. The richer ingestion path is sync.

At a glance

Class diagram

Dependency diagram

API

Classes

KGCodeContext

Line 40

Query code entities from the knowledge graph.

Functions

Source

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

# mypy: ignore-errors
"""Knowledge Graph code context provider.

Replaces direct file reading with knowledge graph queries.
Volks get their code context from the graph, not from disk.
After code generation finishes, sync to KG. All other volks
read from KG.
"""

from __future__ import annotations

import logging
import os
import uuid

from neo4j import Driver
from neo4j.exceptions import ServiceUnavailable

from orchestrata_ultimate.core.graph import BoltSyncRepository, GraphRepositoryConfig
from orchestrata_ultimate.knowledge_graph import canonical_schema as cs

_logger = logging.getLogger(__name__)

# Maximum lines to store per file in the graph
_MAX_CONTENT_LINES = 500


def _entity_id(project_id: str, file_path: str, name: str) -> str:
    """Canonical Class/Function id for the regex-based ``_sync_entities``
    writer: project-qualified ``code_entity_id`` (SPEC-106), the same
    ``<project_id>:<relpath>::<qualname>`` shape ``sync_files_to_kg`` already
    uses for the sibling File node at this exact call site. Replaces the old
    ``{name, file_path}`` MERGE key, which diverged from the canonical
    ``Class.id``/``Function.id`` constraint and collided across projects
    sharing a class/function name (2026-07-22 incident class).
    """
    return f"{project_id}:{cs.code_entity_id(file_path, name)}"


class KGCodeContext:
    """Query code entities from the knowledge graph.

    Provides code context to volks without direct disk reads.
    Falls back to disk reads when the graph is unreachable or nodes
    lack stored content.

    Args:
        graph_uri: Bolt URI for the graph.
        graph_user: Graph username.
        graph_password: Graph password.
    """

    def __init__(
        self,
        graph_uri: str = "",
        graph_user: str = "",
        graph_password: str = "",
        repo: BoltSyncRepository | None = None,
    ) -> None:
        """Initialize KG context.

        Args:
            graph_uri: Bolt URI for the graph.
            graph_user: Graph username.
            graph_password: Graph password.
            repo: Optional shared ``BoltSyncRepository``. If provided,
                uri/user/password are ignored and no new repository/driver
                is created. This enables connection pooling across
                all spike components (backlog #163 phase 5).
        """
        self._uri = graph_uri
        self._user = graph_user
        self._password = graph_password
        self._repo: BoltSyncRepository | None = repo
        self._connected = repo is not None
        self._owns_repo = repo is None

    @property
    def _driver(self) -> Driver | None:
        """Raw ``neo4j.Driver`` escape hatch, kept for external inspection.

        Delegates to the wrapped repository's ``.driver`` property (see
        ``core.graph.adapters.bolt_sync.BoltSyncRepository``). ``None``
        when no repository is held.
        """
        return self._repo.driver if self._repo is not None else None

    def connect(self) -> bool:
        """Connect to the graph.

        If a shared repository was provided at init, verifies it.
        Otherwise builds and owns a new repository.

        Returns:
            True if connection succeeded, False otherwise.
        """
        if self._repo is not None and not self._owns_repo:
            try:
                self._repo.verify_connectivity()
                self._connected = True
                return True
            except Exception as exc:
                _logger.error("Shared driver verify failed: %s", exc)
                self._connected = False
                return False
        try:
            self._repo = BoltSyncRepository(
                GraphRepositoryConfig(
                    uri=self._uri,
                    auth=(self._user, self._password),
                    driver_kwargs={"connection_timeout": 5},
                )
            )
            self._repo.verify_connectivity()
            self._connected = True
            self._owns_repo = True
            return True
        except Exception as exc:
            _logger.error("Connection failed: %s", exc)
            self._connected = False
            return False

    def close(self) -> None:
        """Close the repository if owned.

        Shared repositories are not closed — the kernel manages their
        lifecycle.
        """
        if self._repo is not None and self._owns_repo:
            try:
                self._repo.close()
            except ServiceUnavailable as exc:
                _logger.warning("Close warning: %s", exc)
            self._repo = None
        self._connected = False

    def get_file_content(self, file_path: str) -> str | None:
        """Get file content from KG by path.

        Looks for the content property on the File node. If content
        is not stored, returns class/function signatures as a summary.

        Args:
            file_path: Absolute or workspace-relative file path.

        Returns:
            File content string, or None if not found.
        """
        if not self._connected:
            return None
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            # Try full content first
            rows = self._repo.execute_read(
                "MATCH (f:File) "
                "WHERE f.path = $path OR f.path ENDS WITH $path "
                "RETURN f.content AS content, f.path AS path "
                "LIMIT 1",
                {"path": file_path},
            )
            if rows and rows[0]["content"]:
                return rows[0]["content"]

            # Fallback: return signatures from defined entities
            return self._get_signatures_for_file(
                self._repo,
                file_path,
            )
        except ServiceUnavailable:
            return None
        except Exception as exc:
            _logger.error("get_file_content error: %s", exc)
            return None

    def _get_signatures_for_file(
        self,
        repo: BoltSyncRepository,
        file_path: str,
    ) -> str | None:
        """Build a signature summary for a file from its entities.

        Args:
            repo: Graph repository to query with.
            file_path: File path to query.

        Returns:
            Formatted signature listing, or None if no entities.
        """
        records = repo.execute_read(
            "MATCH (f:File)-[:DEFINES]->(e) "
            "WHERE f.path = $path OR f.path ENDS WITH $path "
            "RETURN labels(e) AS labels, e.name AS name, "
            "e.signature AS signature, e.docstring AS docstring "
            "ORDER BY e.line_start",
            {"path": file_path},
        )
        if not records:
            return None

        lines = [f"# Signatures for {file_path}"]
        for rec in records:
            labels = rec["labels"]
            name = rec["name"]
            sig = rec["signature"] or name
            doc = rec["docstring"] or ""
            entity_type = "class" if "Class" in labels else "def"
            lines.append(f"{entity_type} {sig}")
            if doc:
                doc_preview = doc.split("\n")[0]
                lines.append(f'    """{doc_preview}"""')
        return "\n".join(lines)

    def get_project_files(
        self,
        project_dir: str,
    ) -> list[dict]:
        """Get all files for a project directory from KG.

        Args:
            project_dir: Root directory of the project.

        Returns:
            List of dicts with path, name, and line_count keys.
        """
        if not self._connected:
            return []
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            return self._repo.execute_read(
                "MATCH (f:File) "
                "WHERE f.path STARTS WITH $dir "
                "RETURN f.path AS path, f.name AS name, "
                "f.line_count AS line_count "
                "ORDER BY f.path",
                {"dir": project_dir},
            )
        except Exception as exc:
            _logger.error("get_project_files error: %s", exc)
            return []

    def get_complexity_hotspots(
        self,
        project_dir: str,
        threshold: int = 5,
        limit: int = 10,
    ) -> list[dict]:
        """Return the most complex functions in a project.

        Requires that ``scripts/serialize_codebase_ast.py ast-to-kg`` has been
        run for the project so Function nodes carry ``complexity`` and
        ``branch_count`` properties.

        Args:
            project_dir: Root directory of the project.
            threshold: Minimum cyclomatic complexity to include.
            limit: Maximum number of rows to return.

        Returns:
            List of dicts with ``qualified_name``, ``name``, ``complexity``,
            ``branch_count``, and ``file_path`` keys.
        """
        if not self._connected:
            return []
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            return self._repo.execute_read(
                "MATCH (f:Function) "
                "WHERE f.complexity >= $threshold "
                "  AND (f.file_path STARTS WITH $dir OR f.path STARTS WITH $dir) "
                "RETURN coalesce(f.qualified_name, f.id) AS qualified_name, "
                "       f.name AS name, "
                "       f.complexity AS complexity, "
                "       f.branch_count AS branch_count, "
                "       coalesce(f.file_path, f.path) AS file_path "
                "ORDER BY f.complexity DESC "
                "LIMIT $limit",
                {
                    "dir": project_dir,
                    "threshold": threshold,
                    "limit": limit,
                },
            )
        except Exception as exc:
            _logger.error("get_complexity_hotspots error: %s", exc)
            return []

    def get_source_listing(
        self,
        project_dir: str,
    ) -> str:
        """Get formatted source listing from KG.

        Replaces _read_source_files() and _read_file_contents().
        Returns file contents from KG content property, falling back
        to entity signatures when content is not stored.

        Args:
            project_dir: Root directory of the project.

        Returns:
            Concatenated source listing string.
        """
        if not self._connected:
            return ""
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            # Get files with content
            records = self._repo.execute_read(
                "MATCH (f:File) "
                "WHERE f.path STARTS WITH $dir "
                "AND f.name ENDS WITH '.py' "
                "RETURN f.path AS path, f.name AS name, "
                "f.content AS content "
                "ORDER BY f.path",
                {"dir": project_dir},
            )

            if not records:
                return ""

            listing_parts = []
            for rec in records:
                fpath = rec["path"]
                rel = os.path.relpath(fpath, project_dir)
                content = rec["content"]
                if content:
                    listing_parts.append(f"\n=== {rel} ===\n{content}\n")
                else:
                    # Fallback to signatures
                    sigs = self._get_signatures_for_file(
                        self._repo,
                        fpath,
                    )
                    if sigs:
                        listing_parts.append(f"\n=== {rel} (signatures) ===\n{sigs}\n")
            return "".join(listing_parts)
        except Exception as exc:
            _logger.error("get_source_listing error: %s", exc)
            return ""

    def get_call_graph_neighborhood(
        self,
        qualified_name: str,
        hops: int = 1,
        limit: int = 20,
    ) -> dict:
        """Return callers and callees of a function from the KG call graph.

        Requires that ``scripts/serialize_codebase_ast.py ast-to-kg`` has been
        run so ``CALLS`` relationships exist between :Function nodes.

        Args:
            qualified_name: Function id/qualified_name to center on.
            hops: Maximum path length (1 = direct neighbors).
            limit: Max neighbors per direction.

        Returns:
            Dict with ``callers`` and ``callees`` lists of
            ``{qualified_name, name, distance}``.
        """
        if not self._connected:
            return {"callers": [], "callees": []}
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            # Variable-length pattern bounds must be literals in the graph, so the
            # hop count is interpolated safely into the Cypher string.
            callers_cypher = (
                f"MATCH path = (caller:Function)-[:CALLS*1..{hops}]->"
                f"(f:Function {{id: $qn}}) "
                "RETURN DISTINCT caller.id AS qualified_name, "
                "       caller.name AS name, "
                "       length(path) AS distance "
                "LIMIT $limit"
            )
            callees_cypher = (
                f"MATCH path = (f:Function {{id: $qn}})-[:CALLS*1..{hops}]->"
                "(callee:Function) "
                "RETURN DISTINCT callee.id AS qualified_name, "
                "       callee.name AS name, "
                "       length(path) AS distance "
                "LIMIT $limit"
            )
            callers = self._repo.execute_read(
                callers_cypher,
                {"qn": qualified_name, "limit": limit},
            )
            callees = self._repo.execute_read(
                callees_cypher,
                {"qn": qualified_name, "limit": limit},
            )
            return {
                "callers": [
                    {
                        "qualified_name": r["qualified_name"],
                        "name": r["name"],
                        "distance": r["distance"],
                    }
                    for r in callers
                    if r["qualified_name"]
                ],
                "callees": [
                    {
                        "qualified_name": r["qualified_name"],
                        "name": r["name"],
                        "distance": r["distance"],
                    }
                    for r in callees
                    if r["qualified_name"]
                ],
            }
        except Exception as exc:
            _logger.error("get_call_graph_neighborhood error: %s", exc)
            return {"callers": [], "callees": []}

    def get_source_listing_for_files(
        self,
        project_dir: str,
        files: list[str],
    ) -> str:
        """Get formatted source listing for specific files from KG.

        Replaces _read_file_contents(files). Queries KG for each
        file's content, falls back to signatures.

        Args:
            project_dir: Root directory of the project.
            files: List of relative file paths.

        Returns:
            Concatenated source listing string.
        """
        if not self._connected:
            return ""
        listing_parts = []
        for rel_path in files:
            abs_path = os.path.join(project_dir, rel_path)
            content = self.get_file_content(abs_path)
            if content:
                listing_parts.append(f"\n=== {rel_path} ===\n{content}\n")
            else:
                listing_parts.append(
                    f"\n=== {rel_path} ===\n(not in knowledge graph)\n"
                )
        return "".join(listing_parts)

    def get_classes_and_functions(
        self,
        project_dir: str,
    ) -> str:
        """Get class/function definitions from KG.

        Returns structured listing of classes with their methods
        and standalone functions for all files in the project.

        Args:
            project_dir: Root directory of the project.

        Returns:
            Formatted class and function listing.
        """
        if not self._connected:
            return ""
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            class_records = self._repo.execute_read(
                "MATCH (f:File)-[:DEFINES]->(c:Class) "
                "WHERE f.path STARTS WITH $dir "
                "OPTIONAL MATCH (c)-[:DEFINES]->(m:Function) "
                "RETURN f.path AS file_path, f.name AS file_name, "
                "c.name AS class_name, c.signature AS class_sig, "
                "c.docstring AS class_doc, "
                "collect({"
                "  name: m.name, "
                "  signature: m.signature, "
                "  docstring: m.docstring"
                "}) AS methods "
                # ORDER BY the RETURNed aliases, not the raw bound vars
                # (f.path, c.name): Memgraph scopes ORDER BY to the RETURN
                # projection when an aggregation (collect()) is present and
                # rejects a reference to the pre-aggregation variable
                # ("Unbound variable: f.") — dialect gate fix, plans/
                # neo4j-to-orchgraph-migration-plan.md Step 0.4/2.2. Neo4j
                # accepts both forms, so this is portable.
                "ORDER BY file_path, class_name",
                {"dir": project_dir},
            )

            # Also get standalone functions
            func_records = self._repo.execute_read(
                "MATCH (f:File)-[:DEFINES]->(fn:Function) "
                "WHERE f.path STARTS WITH $dir "
                "AND NOT EXISTS { "
                "  MATCH (:Class)-[:DEFINES]->(fn) "
                "} "
                "RETURN f.path AS file_path, "
                "fn.name AS func_name, "
                "fn.signature AS func_sig, "
                "fn.docstring AS func_doc "
                "ORDER BY f.path, fn.name",
                {"dir": project_dir},
            )

            if not class_records and not func_records:
                return ""

            parts = []
            for rec in class_records:
                rel = os.path.relpath(
                    rec["file_path"],
                    project_dir,
                )
                cls = rec["class_sig"] or rec["class_name"]
                parts.append(f"# {rel}")
                parts.append(f"class {cls}:")
                doc = rec["class_doc"]
                if doc:
                    parts.append(f'    """{doc.split(chr(10))[0]}"""')
                methods = rec["methods"]
                for m in methods:
                    if m.get("name"):
                        sig = m.get("signature") or m["name"]
                        parts.append(f"    def {sig}: ...")
                parts.append("")

            for rec in func_records:
                rel = os.path.relpath(
                    rec["file_path"],
                    project_dir,
                )
                sig = rec["func_sig"] or rec["func_name"]
                parts.append(f"# {rel}")
                parts.append(f"def {sig}: ...")
                parts.append("")

            return "\n".join(parts)
        except Exception as exc:
            _logger.error("get_classes_and_functions error: %s", exc)
            return ""

    def get_dependencies(
        self,
        project_dir: str,
    ) -> list[str]:
        """Get import dependencies from KG.

        Args:
            project_dir: Root directory of the project.

        Returns:
            List of import module names.
        """
        if not self._connected:
            return []
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            rows = self._repo.execute_read(
                "MATCH (f:File)-[:IMPORTS]->(dep) "
                "WHERE f.path STARTS WITH $dir "
                "RETURN DISTINCT dep.name AS name "
                "ORDER BY dep.name",
                {"dir": project_dir},
            )
            return [r["name"] for r in rows if r["name"]]
        except Exception as exc:
            _logger.error("get_dependencies error: %s", exc)
            return []

    def store_design_doc(
        self,
        project_dir: str,
        doc_type: str,
        title: str,
        content: str,
    ) -> str:
        """Store a design document in KG and link it to a project.

        Creates a Document node with the given metadata and content,
        then links it to the Project node whose path or slug matches
        the project directory.

        Args:
            project_dir: Project workspace path.
            doc_type: One of 'architecture', 'design_decision',
                'interface'.
            title: Document title.
            content: Full document content.

        Returns:
            Generated document ID, or empty string on failure.
        """
        if not self._connected:
            return ""
        assert self._repo is not None  # invariant: _connected implies _repo
        doc_id = str(uuid.uuid4())
        try:
            self._repo.execute_write(
                "CREATE (d:Document {"
                "  id: $id, type: $type, title: $title,"
                "  content: $content, path: $path,"
                "  created_at: datetime()"
                "}) "
                "WITH d "
                "OPTIONAL MATCH (p:Project) "
                "WHERE p.path STARTS WITH $path "
                "   OR p.slug CONTAINS $slug "
                "FOREACH (_ IN CASE WHEN p IS NOT NULL "
                "  THEN [1] ELSE [] END | "
                "  MERGE (p)-[:HAS_DOC]->(d)"
                ")",
                {
                    "id": doc_id,
                    "type": doc_type,
                    "title": title,
                    "content": content,
                    "path": project_dir,
                    "slug": os.path.basename(project_dir.rstrip("/")),
                },
            )
            _logger.info("Stored design doc: %s (%s)", title, doc_type)
            return doc_id
        except ServiceUnavailable:
            return ""
        except Exception as exc:
            _logger.error("store_design_doc error: %s", exc)
            return ""

    def get_design_docs(
        self,
        project_dir: str,
        doc_type: str | None = None,
    ) -> list[dict]:
        """Get design documents from KG for a project.

        Queries Document nodes linked to the project or matching
        the project path, optionally filtered by document type.

        Args:
            project_dir: Project workspace path.
            doc_type: Optional filter by type (e.g. 'architecture').

        Returns:
            List of dicts with id, type, title, content, and
            created_at keys.
        """
        if not self._connected:
            return []
        assert self._repo is not None  # invariant: _connected implies _repo
        try:
            if doc_type:
                return self._repo.execute_read(
                    "MATCH (d:Document) "
                    "WHERE d.path = $path "
                    "  AND d.type = $doc_type "
                    "RETURN d.id AS id, d.type AS type, "
                    "d.title AS title, d.content AS content, "
                    "d.created_at AS created_at "
                    "ORDER BY d.created_at DESC",
                    {
                        "path": project_dir,
                        "doc_type": doc_type,
                    },
                )
            return self._repo.execute_read(
                "MATCH (d:Document) "
                "WHERE d.path = $path "
                "RETURN d.id AS id, d.type AS type, "
                "d.title AS title, d.content AS content, "
                "d.created_at AS created_at "
                "ORDER BY d.created_at DESC",
                {"path": project_dir},
            )
        except ServiceUnavailable:
            return []
        except Exception as exc:
            _logger.error("get_design_docs error: %s", exc)
            return []

    def sync_files_to_kg(
        self,
        project_dir: str,
    ) -> int:
        """Sync source files from disk to KG.

        Call after J:Junior generates code to push new/updated files
        into the knowledge graph so other volks can read them.

        Creates/updates File nodes with content property and basic
        entity extraction (class names, function names).

        Args:
            project_dir: Root directory containing source files.

        Returns:
            Number of files synced.
        """
        if not self._connected:
            return 0
        assert self._repo is not None  # invariant: _connected implies _repo

        # Project-qualified id ("{project_id}:{path}"), same shape as the
        # canonical File/Package identity elsewhere in the KG (2026-07-12):
        # a bare-path MERGE key is unsound across projects. project_dir is
        # this writer's only project scope, so — mirroring the same
        # basename-of-root convention used by store_design_doc() in this
        # class and by session/kg_sync's _SyncerBase — derive project_id
        # from it rather than leaving File identity un-qualified.
        project_id = os.path.basename(project_dir.rstrip("/"))

        count = 0
        for root, dirs, filenames in os.walk(project_dir):
            dirs[:] = [
                d
                for d in dirs
                if not d.startswith(".")
                and d != "__pycache__"
                and not d.startswith("Module:")
            ]
            for fname in filenames:
                if not fname.endswith(".py"):
                    continue
                fpath = os.path.join(root, fname)
                try:
                    with open(fpath, encoding="utf-8") as fh:
                        content = fh.read()
                except (OSError, UnicodeDecodeError) as exc:
                    _logger.warning("Cannot read %s: %s", fname, exc)
                    continue

                lines = content.splitlines()
                if len(lines) > _MAX_CONTENT_LINES:
                    content = (
                        "\n".join(lines[:_MAX_CONTENT_LINES]) + f"\n# ... truncated at "
                        f"{_MAX_CONTENT_LINES}/{len(lines)} lines"
                    )

                # gate-ok: the return value is a COUNT of files synced, incremented only on success, so a failed file shows as a shortfall the caller can see
                try:
                    self._repo.execute_write(
                        "MERGE (f:File {id: $id}) "
                        "SET f.path = $path, "
                        "f.project_id = $project_id, "
                        "f.name = $name, "
                        "f.content = $content, "
                        "f.line_count = $line_count, "
                        "f.synced_at = datetime()",
                        {
                            "id": f"{project_id}:{fpath}",
                            "path": fpath,
                            "project_id": project_id,
                            "name": fname,
                            "content": content,
                            "line_count": len(lines),
                        },
                    )
                    # Extract and sync class/function names
                    self._sync_entities(
                        self._repo,
                        fpath,
                        content,
                        project_id,
                    )
                    count += 1
                except Exception as exc:
                    _logger.warning("sync error for %s: %s", fname, exc)
        return count

    def _sync_entities(
        self,
        repo: BoltSyncRepository,
        file_path: str,
        content: str,
        project_id: str,
    ) -> None:
        """Extract and sync class/function entities from file content.

        Uses simple regex-based extraction (not full AST) to create
        Class and Function nodes linked to the File node via DEFINES.

        Args:
            repo: Graph repository to write with.
            file_path: Absolute path to the file.
            content: File content string.
            project_id: This writer's project scope (see ``sync_files_to_kg``)
                -- threaded through so Class/Function ids can be built with
                the same ``_entity_id`` scheme as the sibling File node.
        """
        import re

        # Extract class definitions
        class_pattern = re.compile(
            r"^class\s+(\w+)(\([^)]*\))?:",
            re.MULTILINE,
        )
        class_rows = []
        for match in class_pattern.finditer(content):
            name = match.group(1)
            class_rows.append(
                {
                    "id": _entity_id(project_id, file_path, name),
                    "name": name,
                    "sig": match.group(0).rstrip(":"),
                    "line": content[: match.start()].count("\n") + 1,
                }
            )
        if class_rows:
            repo.execute_write(
                "MATCH (f:File {path: $path}) "
                "UNWIND $rows AS row "
                "MERGE (c:Class {id: row.id}) "
                "SET c.name = row.name, c.file_path = $path, "
                "c.signature = row.sig, c.line_start = row.line "
                "MERGE (f)-[:DEFINES]->(c)",
                {"path": file_path, "rows": class_rows},
            )

        # Extract top-level function definitions
        func_pattern = re.compile(
            r"^def\s+(\w+)\s*\([^)]*\)",
            re.MULTILINE,
        )
        func_rows = []
        for match in func_pattern.finditer(content):
            name = match.group(1)
            func_rows.append(
                {
                    "id": _entity_id(project_id, file_path, name),
                    "name": name,
                    "sig": match.group(0),
                    "line": content[: match.start()].count("\n") + 1,
                }
            )
        if func_rows:
            repo.execute_write(
                "MATCH (f:File {path: $path}) "
                "UNWIND $rows AS row "
                "MERGE (fn:Function {id: row.id}) "
                "SET fn.name = row.name, fn.file_path = $path, "
                "fn.signature = row.sig, fn.line_start = row.line "
                "MERGE (f)-[:DEFINES]->(fn)",
                {"path": file_path, "rows": func_rows},
            )

    def get_monorepo_classes(
        self,
        class_names: set[str],
    ) -> dict[str, tuple[str, str]]:
        """Query KG for orchestrata class sources by name.

        Replaces the filesystem scan in architect._build_monorepo_context.
        Queries the orchestrata project node for Class entities matching
        the given names, returns their file paths and content.

        Args:
            class_names: Set of class names to look up.

        Returns:
            Dict mapping file_path to (relative_path, content) tuples.
        """
        if not self._connected or not class_names:
            return {}
        assert self._repo is not None  # invariant: _connected implies _repo
        loaded: dict[str, tuple[str, str]] = {}
        try:
            rows = self._repo.execute_read(
                "MATCH (proj:Project {name: $proj})"
                "-[:HAS_FILE]->(f:File)"
                "-[:DEFINES]->(c:Class) "
                "WHERE c.name IN $names "
                "RETURN DISTINCT f.path AS path, "
                "f.content AS content",
                {
                    "proj": "orchestrata",
                    "names": list(class_names),
                },
            )
            for rec in rows:
                fpath = rec["path"]
                if fpath in loaded:
                    continue
                content = rec["content"]
                if not content and os.path.isfile(fpath):
                    # Content not in KG — read from disk
                    try:
                        with open(
                            fpath,
                            encoding="utf-8",
                        ) as fh:
                            content = fh.read()
                    except (OSError, UnicodeDecodeError) as exc:
                        _logger.warning("Cannot read %s: %s", fpath, exc)
                        continue
                if not content:
                    continue
                lines = content.splitlines()
                if len(lines) > 200:
                    content = (
                        "\n".join(lines[:200]) + f"\n... (truncated at "
                        f"200/{len(lines)} lines)"
                    )
                rel = (
                    fpath.split("/packages/", 1)[-1] if "/packages/" in fpath else fpath
                )
                loaded[fpath] = (rel, content)
        except Exception as exc:
            _logger.error("get_monorepo_classes error: %s", exc)
        return loaded