orchestrata_ultimate.knowledge_graph.query_engine

Reference › orchestrata_ultimate › knowledge_graph › query_engine

QueryEngine is the general read API over the operational graph (orchkg on orchgraph), used by the context assembler, the gateway RPC layer, the kernel and other subsystems (Spec 008, REQ-008-004 and -005, REQ-008-013, -014 and -040, and Spec 016). It wraps a graph_client GraphClient and implements the IQueryEngine protocol from interfaces.

• query runs Cypher through an LRU cache whose keys include a graph version token; invalidate_cache, set_graph_version and the module functions bump_graph_version and get_graph_version invalidate it.
• multi_hop runs a bounded variable-length path match from a start id, optionally filtered by relationship types, and returns a ContextBundle; get_context gives the bundle for an entity at a depth of 2 by default.
• search finds nodes by name substring.
• run_provenance returns the full provenance of a pipeline run by thread id.
• volk_work_context and component_work_context return a volk's owned components, decisions and generated artifacts, and the component-ownership context of SPEC-008.

Every read returns an empty result instead of raising when the client is missing or not connected, and ensure_connected connects lazily. See search for the text-search sibling.

At a glance

Class diagram

Dependency diagram

API

Constants

Classes

QueryEngine

Decorators: traces(spec='008', story='US-001', reqs=['REQ-008-004', 'REQ-008-013', 'REQ-008-014', 'REQ-008-040', 'REQ-008-049', 'RE… · Line 55

Executes queries against the knowledge graph.

Functions

Source

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

# mypy: ignore-errors
"""Query engine for the knowledge graph.

Provides multi-hop traversal, context injection, and
cached query patterns.
"""

from __future__ import annotations

import hashlib
import logging
import time
from collections import OrderedDict
from typing import TYPE_CHECKING, Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.knowledge_graph.types import ContextBundle, Node

if TYPE_CHECKING:
    from orchestrata_ultimate.knowledge_graph.graph_client import GraphClient

logger = logging.getLogger(__name__)

# Module-level graph version used by every QueryEngine cache key. Bumping this
# invalidates cached entries across all engine instances without needing to
# track them centrally. Sync code calls bump_graph_version() after writes.
_GRAPH_VERSION: int = 0


@traces(spec="008", story="US-001", reqs=["REQ-008-015"])
def bump_graph_version() -> None:
    """Invalidate query-engine caches by bumping the global graph version."""
    global _GRAPH_VERSION
    _GRAPH_VERSION += 1


@traces(spec="008", story="US-001", reqs=["REQ-008-015"])
def get_graph_version() -> int:
    """Return the current global graph version used in cache keys."""
    return _GRAPH_VERSION


@traces(
    spec="008",
    story="US-001",
    reqs=[
        "REQ-008-004",
        "REQ-008-013",
        "REQ-008-014",
        "REQ-008-040",
        "REQ-008-049",
        "REQ-008-053",
    ],
)
class QueryEngine:
    """Executes queries against the knowledge graph.

    Implements IQueryEngine protocol with LRU caching.
    """

    def __init__(
        self, graph_client: GraphClient | None = None, cache_size: int = 1000
    ) -> None:
        self._client = graph_client
        self._cache: OrderedDict[str, Any] = OrderedDict()
        self._cache_size = cache_size
        self._query_count = 0

    @traces(spec="008", story="US-001", reqs=["REQ-008-005"])
    async def ensure_connected(self) -> bool:
        """Lazily connect the underlying graph client if needed.

        A wiring seam for callers -- like the gateway -- that construct the
        engine synchronously (no ``await`` available at that point) and need
        a cheap, idempotent way to make sure the client is connected before
        the first real query. Returns False when there is no client at all;
        True immediately if already connected; otherwise attempts
        ``self._client.connect()`` and swallows any connection failure into
        a plain False rather than raising.
        """
        if self._client is None:
            return False
        if self._client.connected:
            return True
        try:
            return await self._client.connect()
        except Exception:
            return False

    @traces(spec="008", story="US-001", reqs=["REQ-008-012"])
    @traces(spec="016", story="US-001", reqs=["REQ-016-238"])
    async def query(
        self, cypher: str, parameters: dict[str, Any] | None = None
    ) -> list[dict[str, Any]]:
        """Execute a Cypher query."""
        cache_key = self._cache_key(cypher, parameters or {})
        cached = self._cache_get(cache_key)
        if cached is not None:
            return cached

        if not self._client or not self._client.connected:
            return []

        self._query_count += 1
        result = await self._client.query(cypher, parameters)
        self._cache_put(cache_key, result)
        return result

    @traces(spec="008", story="US-001", reqs=["REQ-008-005"])
    @traces(spec="016", story="US-001", reqs=["REQ-016-240"])
    async def multi_hop(
        self, start_id: str, depth: int = 3, rel_types: list[str] | None = None
    ) -> ContextBundle:
        """Multi-hop traversal from a starting node."""
        start = time.perf_counter()

        if not self._client or not self._client.connected:
            return ContextBundle(query=f"multi_hop:{start_id}", depth=depth)

        rel_filter = ""
        if rel_types:
            rel_filter = ":" + "|".join(rel_types)

        cypher = (
            f"MATCH path = (start {{id: $start_id}})-[r{rel_filter}*1..{depth}]-(end) "
            "RETURN DISTINCT properties(end) AS node, length(path) AS dist "
            "ORDER BY dist LIMIT 100"
        )
        results = await self._client.query(cypher, {"start_id": start_id})

        nodes = []
        relationships = []
        for record in results:
            node_data = record.get("node", {})
            # Accept dicts from Cypher ``properties()`` and mapping-like
            # Neo4j ``Node`` objects returned directly as ``end``.
            is_mapping = isinstance(node_data, dict) or hasattr(node_data, "get")
            if is_mapping:
                try:
                    nodes.append(
                        Node(
                            id=node_data.get("id", ""),
                            name=node_data.get("name", ""),
                        )
                    )
                except Exception as exc:
                    logger.warning(
                        "Failed to construct Node from record data: %s",
                        exc,
                        exc_info=True,
                    )
            relationships.append(
                {
                    "distance": record.get("dist", 0),
                    "node_id": (node_data.get("id", "") if is_mapping else ""),
                }
            )

        elapsed = (time.perf_counter() - start) * 1000

        return ContextBundle(
            query=f"multi_hop:{start_id}",
            nodes=nodes,
            relationships=relationships,
            depth=depth,
            total_results=len(nodes),
            query_time_ms=elapsed,
        )

    @traces(spec="008", story="US-001", reqs=["REQ-008-004", "REQ-008-005"])
    async def run_provenance(self, thread_id: str) -> dict[str, Any]:
        """Return the full provenance of a pipeline run, keyed by ``thread_id``.

        Traverses the semantic layer the runtime writes — ``(:Intent)-[:TRIGGERED]->
        (:PipelineRun)-[:EXECUTED_BY]->(:Volk)`` and ``(:Artifact)-[:GENERATED_BY]->
        (:Volk)`` — to answer "what did this run do": the originating intent text,
        the volks that executed, and the artifacts they produced. This is the
        run-level counterpart to :meth:`volk_work_context` and gives the
        orchestrator after-the-fact visibility into any thread.

        Empty structure when the graph is unavailable or the run is unknown.
        """
        empty = {
            "thread_id": thread_id,
            "intent": None,
            "volks": [],
            "artifacts": [],
        }
        if not self._client or not self._client.connected:
            return empty

        cypher = (
            "MATCH (r:PipelineRun {id: $run_id}) "
            "OPTIONAL MATCH (i:Intent)-[:TRIGGERED]->(r) "
            "OPTIONAL MATCH (r)-[:EXECUTED_BY]->(v:Volk) "
            "OPTIONAL MATCH (a:Artifact {thread_id: $thread_id}) "
            "RETURN i.text AS intent, "
            "collect(DISTINCT v.alias) AS volks, "
            "collect(DISTINCT properties(a)) AS artifacts"
        )
        results = await self._client.query(
            cypher, {"run_id": f"run:{thread_id}", "thread_id": thread_id}
        )
        if not results:
            return empty
        row = results[0]
        return {
            "thread_id": thread_id,
            "intent": row.get("intent"),
            "volks": [a for a in (row.get("volks") or []) if a],
            "artifacts": [p for p in (row.get("artifacts") or []) if p],
        }

    @traces(spec="008", story="US-001", reqs=["REQ-008-005"])
    async def volk_work_context(self, volk_alias: str) -> dict[str, Any]:
        """Return a volk's work context from the semantic graph (SPEC-008).

        Gathers the components it owns, the decisions it has made, and the
        artifacts it generated — the volk-aware context SPEC-008 specifies for
        injection. Empty lists when the graph is unavailable or the volk has no
        recorded work yet (the W5 semantic layer is populated by the runtime).
        """
        empty = {
            "volk": volk_alias,
            "components": [],
            "decisions": [],
            "artifacts": [],
        }
        if not self._client or not self._client.connected:
            return empty

        cypher = (
            "MATCH (v:Volk {alias: $alias}) "
            "OPTIONAL MATCH (v)-[:OWNS]->(c:Component) "
            "OPTIONAL MATCH (v)-[:MADE]->(d:Decision) "
            "OPTIONAL MATCH (a:Artifact)-[:GENERATED_BY]->(v) "
            "RETURN collect(DISTINCT properties(c)) AS components, "
            "collect(DISTINCT properties(d)) AS decisions, "
            "collect(DISTINCT properties(a)) AS artifacts"
        )
        results = await self._client.query(cypher, {"alias": volk_alias})
        if not results:
            return empty
        row = results[0]
        return {
            "volk": volk_alias,
            "components": [p for p in (row.get("components") or []) if p],
            "decisions": [p for p in (row.get("decisions") or []) if p],
            "artifacts": [p for p in (row.get("artifacts") or []) if p],
        }

    @traces(spec="008", story="US-003", reqs=["REQ-008-012"])
    async def component_work_context(
        self, volk_alias: str, component_name: str
    ) -> dict[str, Any]:
        """Return SPEC-008's flagship component-ownership context query.

        Everything a volk needs to resume work on a component it owns: the
        decisions that apply to it, its in-progress tasks, and the sessions
        those tasks belong to. Empty shape when the graph is unavailable or
        the volk does not (yet) ``OWN`` the named component.

        ``component_name`` is the bare component name (e.g. ``"kernel"``).
        Component's canonical merge key is now the project-qualified ``id =
        "{project}:{name}"`` (ratified 2026-07-19, plans/architecture-
        ontology-draft.md §9.2 ruling 2), but this method has no project
        context of its own to qualify it with — it resolves against the
        ``"orchestrata-ultimate"`` default project, the same fallback
        ``kernel/pipeline_kg.py``'s writers use. A caller with real
        multi-project context should resolve/pass the qualified id itself
        once one exists.
        """
        empty = {
            "volk": volk_alias,
            "component_name": component_name,
            "component": None,
            "decisions": [],
            "tasks": [],
            "sessions": [],
        }
        if not self._client or not self._client.connected:
            return empty

        component_id = f"orchestrata-ultimate:{component_name}"
        cypher = (
            "MATCH (v:Volk {alias: $alias})-[:OWNS]->(c:Component {id: $component_id}) "
            "OPTIONAL MATCH (c)<-[:APPLIES_TO]-(d:Decision) "
            "OPTIONAL MATCH (c)<-[:REQUIRES]-(t:Task {status: 'in_progress'}) "
            "OPTIONAL MATCH (s:Session)-[:CONTAINS]->(t) "
            "RETURN properties(c) AS component, "
            "collect(DISTINCT properties(d)) AS decisions, "
            "collect(DISTINCT properties(t)) AS tasks, "
            "collect(DISTINCT properties(s)) AS sessions"
        )
        results = await self._client.query(
            cypher, {"alias": volk_alias, "component_id": component_id}
        )
        if not results:
            return empty
        row = results[0]
        return {
            "volk": volk_alias,
            "component_name": component_name,
            "component": row.get("component") or None,
            "decisions": [p for p in (row.get("decisions") or []) if p],
            "tasks": [p for p in (row.get("tasks") or []) if p],
            "sessions": [p for p in (row.get("sessions") or []) if p],
        }

    async def search(
        self, query_text: str, node_type: str | None = None, limit: int = 100
    ) -> list[Node]:
        """Search nodes by name substring.

        Without ``node_type`` the search is scoped to the code-symbol labels
        (Class|Function|Predicate — the same universe as SearchEngine's
        ``search_names`` fulltext index) via per-label UNION branches so each
        branch can use its TEXT index; a label-less ``MATCH (n) WHERE n.name
        CONTAINS`` cannot use ANY index and full-scans the node store.
        """
        if not self._client or not self._client.connected:
            return []

        if node_type:
            from orchestrata_ultimate.knowledge_graph.schema import get_node_label

            label = get_node_label(node_type)
            cypher = (
                f"MATCH (n:{label}) "
                "WHERE n.name CONTAINS $query "
                "RETURN properties(n) AS props "
                "LIMIT $limit"
            )
        else:
            from orchestrata_ultimate.knowledge_graph.schema import (
                LABEL_CLASS,
                LABEL_FUNCTION,
                LABEL_PREDICATE,
            )

            cypher = (
                "CALL () { "
                f"MATCH (n:{LABEL_CLASS}) WHERE n.name CONTAINS $query RETURN n "
                f"UNION MATCH (n:{LABEL_FUNCTION}) WHERE n.name CONTAINS $query RETURN n "
                f"UNION MATCH (n:{LABEL_PREDICATE}) WHERE n.name CONTAINS $query RETURN n "
                "} WITH n "
                "RETURN properties(n) AS props LIMIT $limit"
            )

        results = await self._client.query(
            cypher, {"query": query_text, "limit": limit}
        )

        nodes = []
        for record in results:
            props = record.get("props", {})
            if props:
                from orchestrata_ultimate.knowledge_graph.nodes import (
                    cypher_record_to_node,
                )

                node = cypher_record_to_node(props)
                if node:
                    nodes.append(node)
        return nodes

    @traces(
        spec="008",
        story="US-001",
        reqs=[
            "REQ-008-004",
            "REQ-008-013",
            "REQ-008-041",
            "REQ-008-046",
        ],
    )
    async def get_context(self, entity_id: str, depth: int = 2) -> ContextBundle:
        """Get context bundle for an entity (for LLM injection)."""
        return await self.multi_hop(entity_id, depth=depth)

    @traces(spec="008", story="US-001", reqs=["REQ-008-015", "REQ-008-039"])
    @traces(spec="016", story="US-001", reqs=["REQ-016-245"])
    def invalidate_cache(self) -> None:
        """Clear the query cache."""
        self._cache.clear()

    def set_graph_version(self, version: str) -> None:
        """Set the graph version token used in cache keys.

        Deprecated alias for :func:`bump_graph_version`. Prefer the module-level
        bump function because the cache key uses the global version, not a
        per-engine value.
        """
        bump_graph_version()

    @property
    def query_count(self) -> int:
        return self._query_count

    @property
    @traces(
        spec="002",
        story="US-003",
        reqs=["REQ-002-US-003-016", "REQ-002-US-009-008", "REQ-002-US-009-019"],
    )
    @traces(spec="016", story="US-003", reqs=["REQ-016-277"])
    def cache_size(self) -> int:
        return len(self._cache)

    def _cache_key(self, cypher: str, params: dict[str, Any]) -> str:
        """Generate a cache key that includes the current graph version.

        The graph version is bumped by callers (e.g. after a sync) so cached
        entries from a stale graph state are not reused.
        """
        raw = f"{_GRAPH_VERSION}:{cypher}:{sorted(params.items())}"
        return hashlib.md5(raw.encode()).hexdigest()

    @traces(spec="008", story="US-001", reqs=["REQ-008-015", "REQ-008-036"])
    @traces(spec="016", story="US-001", reqs=["REQ-016-279"])
    def _cache_get(self, key: str) -> Any | None:
        """Get from cache (LRU)."""
        if key in self._cache:
            self._cache.move_to_end(key)
            return self._cache[key]
        return None

    @traces(spec="008", story="US-001", reqs=["REQ-008-015", "REQ-008-039"])
    @traces(spec="016", story="US-001", reqs=["REQ-016-280"])
    def _cache_put(self, key: str, value: Any) -> None:
        """Put to cache (LRU eviction)."""
        self._cache[key] = value
        self._cache.move_to_end(key)
        while len(self._cache) > self._cache_size:
            self._cache.popitem(last=False)