orchestrata_ultimate.session.kg_context

Reference › orchestrata_ultimate › session › kg_context

The read side of the knowledge graph for session work: KnowledgeGraphContext queries the graph and formats context for injection into an agent. The write side is the KnowledgeGraphSyncer, and the CLI query commands in cli/kg call these methods.

connect returns whether the graph is reachable, connected reports the state and close ends it. The context methods are get_file_context (a file's code entities and related specs), get_spec_context (a spec's requirements and completion), get_class_context (file, methods and traces) and get_session_context (purpose, actions and modified files). The audit queries are get_undocumented, get_untraced and get_unchecked_requirements. search finds classes, functions and predicates by name, and multi_hop follows relationships from an entity up to a given depth (default 3). format_for_injection renders a context result as a concise text block.

At a glance

Class diagram

Dependency diagram

API

Classes

KnowledgeGraphContext

Decorators: traces(spec='006', story='FR-018', reqs=['REQ-006-FR18-006']), traces(spec='021', story='US-001', reqs=['REQ-021-004']) · Line 27

Queries the graph and formats context for agent injection.

Source

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

# mypy: ignore-errors
"""Knowledge Graph Context Provider.

Queries graph for relevant nodes and formats them for injection
into agent context. Used on every request to provide graph-informed
context to Claude and other agents.

Spec: 006 - Session Management
"""

import logging
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.graph import BoltSyncRepository, GraphRepositoryConfig
from orchestrata_ultimate.core.graph_factory import (
    resolve_kg_database,
    resolve_kg_dialect,
    resolve_uri,
)

logger = logging.getLogger(__name__)


@traces(spec="006", story="FR-018", reqs=["REQ-006-FR18-006"])
@traces(spec="021", story="US-001", reqs=["REQ-021-004"])
class KnowledgeGraphContext:
    """Queries the graph and formats context for agent injection.

    Provides read-only access to the graph for context retrieval.
    All methods return plain dicts/lists suitable for JSON serialization.
    """

    def __init__(
        self,
        graph_uri: str | None = None,
        graph_user: str | None = None,
        graph_password: str | None = None,
        graph_database: str | None = None,
    ) -> None:
        self._uri = graph_uri or resolve_uri()
        self._user = graph_user or ""
        self._password = graph_password or ""
        self._database = resolve_kg_database(graph_database)
        self._repo: BoltSyncRepository | None = None

    @property
    def _driver(self) -> Any | None:
        """Raw-driver read-only escape hatch, kept for external inspection."""
        return self._repo.driver if self._repo is not None else None

    def connect(self) -> bool:
        """Connect to graph. Returns True on success."""
        try:
            repo = BoltSyncRepository(
                GraphRepositoryConfig(
                    uri=self._uri,
                    auth=(self._user, self._password),
                    database=self._database,
                )
            )
            self._repo = repo
            repo.verify_connectivity()
            return True
        except Exception as e:
            logger.warning("Graph connection failed: %s", e)
            self._repo = None
            return False

    def close(self) -> None:
        """Close graph connection."""
        if self._repo:
            self._repo.close()
            self._repo = None

    @property
    def connected(self) -> bool:
        return self._repo is not None

    def _run(self, cypher: str, **params) -> list[dict]:
        """Execute a read-only Cypher query."""
        if not self._repo:
            return []
        return self._repo.execute_read(cypher, params)

    # ── Query methods ─────────────────────────────────────────────

    @traces(spec="021", story="US-001", reqs=["REQ-021-004"])
    def get_file_context(self, file_path: str) -> dict:
        """Get context for a specific file: its code entities and related specs.

        Anchors on ``Module {file_path}`` — the label the KG sync actually
        writes ``DEFINES`` edges from (``knowledge_graph/sync.py``; live graph
        has thousands of ``Module-[:DEFINES]->Class/Function`` edges and ZERO
        ``File``-anchored ones — the previous ``(f:File {path:...})`` match
        returned empty for every file, discovered during the SPEC-021 W2
        port). ``file_path`` is repo-relative, e.g.
        ``src/orchestrata_ultimate/agents/__init__.py``.
        """
        rows = self._run(
            """
            MATCH (m:Module {file_path: $path})-[:DEFINES]->(c)
            OPTIONAL MATCH (c)-[:TRACES_TO]->(ri)<-[:HAS_ITEM]-(r)<-[:DEFINES_REQUIREMENT]-(s:Spec)
            RETURN m.file_path AS file, m.line_count AS lines,
                   collect(DISTINCT {type: labels(c)[0], name: c.name,
                           has_docstring: c.has_docstring}) AS entities,
                   collect(DISTINCT s.id) AS specs
            """,
            path=file_path,
        )
        if not rows:
            return {"file": file_path, "entities": [], "specs": []}
        r = rows[0]
        return {
            "file": r.get("file", file_path),
            "lines": r.get("lines", 0),
            "entities": r.get("entities", []),
            "specs": [s for s in r.get("specs", []) if s],
        }

    def get_spec_context(self, spec_id: str) -> dict:
        """Get context for a spec: its requirements and completion status.

        Matches on ``natural_id`` (the bare spec number, e.g. ``"008"``), not
        the canonical (project-qualified) ``id``: this class has no project
        parameter anywhere in its constructor or call surface (2026-07-12
        multi-project identity migration), so it cannot build a qualified
        id here. When more than one project shares ``spec_id`` in the graph,
        this returns whichever one graph matches first — a known ambiguity
        for a future wave to resolve by threading a project scope through
        this class, not introduced by this change.
        """
        rows = self._run(
            """
            MATCH (s:Spec {natural_id: $id})-[:DEFINES_REQUIREMENT]->(r)-[:HAS_ITEM]->(ri)
            RETURN s.natural_id AS spec_id, s.title AS title, s.status AS status,
                   r.id AS req_id, r.title AS req_title,
                   collect({id: ri.id, text: ri.text, checked: ri.checked}) AS items
            """,
            id=spec_id,
        )
        if not rows:
            return {"spec_id": spec_id, "requirements": []}
        reqs = []
        for r in rows:
            items = r.get("items", [])
            total = len(items)
            checked = sum(1 for i in items if i.get("checked"))
            reqs.append(
                {
                    "id": r.get("req_id", ""),
                    "title": r.get("req_title", ""),
                    "progress": f"{checked}/{total}",
                    "items": items,
                }
            )
        return {
            "spec_id": rows[0].get("spec_id", spec_id),
            "title": rows[0].get("title", ""),
            "status": rows[0].get("status", ""),
            "requirements": reqs,
        }

    def get_class_context(self, class_name: str) -> dict:
        """Get context for a class: file, methods, traces."""
        rows = self._run(
            """
            MATCH (f:File)-[:DEFINES]->(c:Class {name: $name})
            OPTIONAL MATCH (c)-[:DEFINES]->(m)
            OPTIONAL MATCH (c)-[:TRACES_TO]->(ri)
            RETURN f.path AS file, c.name AS class_name, c.docstring AS docstring,
                   collect(DISTINCT m.name) AS methods,
                   collect(DISTINCT ri.id) AS traced_reqs
            """,
            name=class_name,
        )
        if not rows:
            return {
                "class_name": class_name,
                "file": "",
                "methods": [],
                "traced_reqs": [],
            }
        r = rows[0]
        return {
            "class_name": r.get("class_name", class_name),
            "file": r.get("file", ""),
            "docstring": r.get("docstring", ""),
            "methods": r.get("methods", []),
            "traced_reqs": [t for t in r.get("traced_reqs", []) if t],
        }

    @traces(spec="006", story="US-003", reqs=["REQ-006-US-003-001"])
    def get_session_context(self, session_seq: int) -> dict:
        """Get context for a session: purpose, actions, modified files."""
        rows = self._run(
            """
            MATCH (s:Session {seq: $seq})
            OPTIONAL MATCH (s)-[:MODIFIED|PRODUCED]->(f:File)
            OPTIONAL MATCH (s)-[:LOGGED_ACTION]->(a:Action)
            RETURN s.purpose AS purpose, s.ai_model AS model,
                   collect(DISTINCT f.path) AS files,
                   collect(DISTINCT {type: a.type, desc: a.description}) AS actions
            """,
            seq=session_seq,
        )
        if not rows:
            return {"seq": session_seq, "purpose": "", "files": [], "actions": []}
        r = rows[0]
        return {
            "seq": session_seq,
            "purpose": r.get("purpose", ""),
            "model": r.get("model", ""),
            "files": [f for f in r.get("files", []) if f],
            "actions": [a for a in r.get("actions", []) if a.get("type")],
        }

    def get_undocumented(self) -> list[dict]:
        """List code entities missing docstrings."""
        rows = self._run("""
            MATCH (p:Project)-[:HAS_ROOT_DIR]->()-[:CONTAINS*]->(f:File)-[:DEFINES]->(c)
            WHERE c.has_docstring = false
            RETURN f.path AS file, labels(c)[0] AS type, c.name AS name
            ORDER BY f.path
            """)
        return [dict(r) for r in rows]

    def get_untraced(self) -> list[dict]:
        """List code entities with no requirement traceability."""
        rows = self._run("""
            MATCH (f:File)-[:DEFINES]->(c)
            WHERE NOT (c)-[:TRACES_TO]->()
            RETURN f.path AS file, labels(c)[0] AS type, c.name AS name
            ORDER BY f.path
            """)
        return [dict(r) for r in rows]

    def get_unchecked_requirements(self) -> list[dict]:
        """List requirement items that are not checked."""
        rows = self._run("""
            MATCH (s:Spec)-[:DEFINES_REQUIREMENT]->(r)-[:HAS_ITEM]->(ri)
            WHERE ri.checked = false
            RETURN coalesce(s.natural_id, s.id) AS spec, r.title AS requirement,
                   ri.id AS item_id, ri.text AS text
            ORDER BY spec, ri.id
            """)
        return [dict(r) for r in rows]

    #: Cypher for the substring branch of :meth:`search`. Scoped to the same
    #: labels the ``search_names`` fulltext index covers, one UNION branch per
    #: label so each can use its own ``*_name`` TEXT index — a label-less
    #: ``MATCH (n) WHERE n.name CONTAINS`` uses no index and full-scans the
    #: store. ``score`` is substring coverage (term length / name length, 1.0
    #: on an exact name), which is orderable and deterministic; it is NOT a
    #: relevance score and does not claim to match graph's.
    _SUBSTRING_SEARCH_CYPHER = """
        CALL () {
            MATCH (n:Class) WHERE n.name CONTAINS $term RETURN n
            UNION MATCH (n:Function) WHERE n.name CONTAINS $term RETURN n
            UNION MATCH (n:Predicate) WHERE n.name CONTAINS $term RETURN n
        }
        WITH n, toFloat(size($term)) / size(n.name) AS score
        RETURN labels(n)[0] AS type, n.name AS name, n.id AS id, score
        ORDER BY score DESC
        LIMIT $limit
    """

    _FULLTEXT_SEARCH_CYPHER = """
        CALL db.index.fulltext.queryNodes('search_names', $term)
        YIELD node, score
        RETURN labels(node)[0] AS type, node.name AS name,
               node.id AS id, score
        ORDER BY score DESC
        LIMIT $limit
    """

    def search(self, query: str, limit: int = 20) -> list[dict]:
        """Search classes, functions, and predicates by name.

        Engine-aware: ``db.index.fulltext.queryNodes`` is a graph-only
        procedure — orchgraph answers "There is no procedure named
        'db.index.fulltext.queryNodes'" (measured live 2026-07-26), and this
        method has no fallback, so after the cutover every ``session query
        search`` raised ``ClientError``. On Memgraph it runs the substring
        branch instead of a call that cannot succeed.
        """
        cypher = (
            self._SUBSTRING_SEARCH_CYPHER
            if resolve_kg_dialect() == "memgraph"
            else self._FULLTEXT_SEARCH_CYPHER
        )
        rows = self._run(cypher, term=query, limit=limit)
        return [dict(r) for r in rows]

    def multi_hop(self, entity_id: str, depth: int = 3) -> list[dict]:
        """Traverse relationships from an entity up to N hops.

        Plain variable-length Cypher, no APOC: the live server has no APOC
        installed, so the previous ``apoc.path.subgraphAll`` primary path
        raised ``ProcedureNotFound`` (``_run`` does not swallow errors, which
        made the old "if not rows" fallback unreachable). ``*0..`` keeps the
        start node in the result like ``subgraphAll`` did. Cypher forbids
        parameters in variable-length bounds, so the int-coerced depth is
        interpolated.
        """
        hops = max(1, int(depth))
        rows = self._run(
            f"""
            MATCH path = (start {{id: $id}})-[*0..{hops}]-()
            UNWIND nodes(path) AS n
            RETURN DISTINCT labels(n)[0] AS type, n.id AS id, n.name AS name
            """,
            id=entity_id,
        )
        return [dict(r) for r in rows]

    # ── Formatting ────────────────────────────────────────────────

    @staticmethod
    def format_for_injection(context: Any) -> str:
        """Format context data as a concise text block for agent injection.

        Args:
            context: Dict or list from any query method; anything else is
                stringified as-is (annotation is ``Any`` so the defensive
                final branch is honest rather than type-unreachable).

        Returns:
            Formatted text suitable for system prompt injection.
        """
        if isinstance(context, list):
            return KnowledgeGraphContext._format_list(context)
        if isinstance(context, dict):
            return KnowledgeGraphContext._format_dict(context)
        return str(context)

    @staticmethod
    def _format_list(context: list) -> str:
        """Format a list of row dicts as one ``k=v | ...`` line each."""
        lines = []
        for item in context:
            parts = [f"{k}={v}" for k, v in item.items() if v]
            lines.append(" | ".join(parts))
        return "\n".join(lines)

    @staticmethod
    def _format_dict(context: dict) -> str:
        """Format a dict, rendering list values as comma-joined (first 10) entries."""
        lines = []
        for key, value in context.items():
            if isinstance(value, list):
                if value:
                    lines.append(f"{key}: {', '.join(str(v) for v in value[:10])}")
            elif value:
                lines.append(f"{key}: {value}")
        return "\n".join(lines)