orchestrata_ultimate.semantic_graph.spec_trace_enricher

Reference › orchestrata_ultimate › semantic_graph › spec_trace_enricher

Enriches the shared semantic graph (memgraph) with spec, story, requirement, code and test entities, so GraphRAG can answer questions such as which code implements a given requirement or what the requirements for the CXP protocol are. SpecTraceEnricher takes the parsed specs and traces from spec_trace_parser, a project id, and optionally test traces and a test-module map. Its _build step creates the entities and relationships in memory: spec, story and requirement entities, code elements (with stable ids from _code_element_id and labels from _code_element_label), relationships from code to the requirements it implements, TestCase entities with TESTS edges from traced tests, and TestModule entities with TESTS edges from a path heuristic. The entities, relationships and summary properties expose the result as Entity and Relationship models and an EnrichmentSummary of counts. enrich writes everything to a SemanticGraph backend, and the module function enrich_project_with_specs_and_traces runs the whole parse, build and write sequence for a project. The writes use MERGE, so re-running is idempotent. The enrich_semantic_graph_specs script and the package facade import it.

At a glance

Class diagram

Dependency diagram

API

Classes

EnrichmentSummary (dataclass)

Decorators: dataclass · Line 29

Result of a spec/trace enrichment run.

SpecTraceEnricher

Line 73

Build entities and relationships from specs and code traces.

Functions

Source

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

"""Enrich the semantic graph with spec, story, requirement, and trace entities.

The enricher converts parsed specs and @traces annotations into graph nodes and
edges. This lets GraphRAG answer questions such as "which code implements
REQ-004-US-001-001?" or "what are the requirements for the CXP protocol?".
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from pathlib import Path

from orchestrata_ultimate.core.traceability import TraceInfo
from orchestrata_ultimate.semantic_graph.interfaces import SemanticGraph
from orchestrata_ultimate.semantic_graph.models import Entity, Relationship
from orchestrata_ultimate.semantic_graph.spec_trace_parser import (
    SpecInfo,
    collect_code_traces,
    collect_test_traces,
    map_test_modules_to_source,
    parse_all_specs,
)

logger = logging.getLogger(__name__)


@dataclass
class EnrichmentSummary:
    """Result of a spec/trace enrichment run."""

    specs: int = 0
    stories: int = 0
    requirements: int = 0
    code_elements: int = 0
    test_cases: int = 0
    test_modules: int = 0
    relationships: int = 0

    def to_dict(self) -> dict[str, int]:
        return {
            "specs": self.specs,
            "stories": self.stories,
            "requirements": self.requirements,
            "code_elements": self.code_elements,
            "test_cases": self.test_cases,
            "test_modules": self.test_modules,
            "relationships": self.relationships,
        }


def _code_element_id(project_id: str, trace: TraceInfo) -> str:
    """Return a stable entity id for a code element referenced by a trace."""
    parts = [project_id, "code"]
    if trace.module:
        parts.append(trace.module)
    if trace.class_name:
        parts.append(trace.class_name)
    if trace.method:
        parts.append(trace.method)
    return ":".join(parts)


def _code_element_label(trace: TraceInfo) -> str:
    """Return a human-readable label for a code element."""
    if trace.class_name and trace.method:
        return f"{trace.class_name}.{trace.method}"
    if trace.class_name:
        return trace.class_name
    return trace.method or trace.module or "unknown"


class SpecTraceEnricher:
    """Build entities and relationships from specs and code traces."""

    def __init__(
        self,
        specs: list[SpecInfo],
        traces: list[TraceInfo],
        project_id: str,
        test_traces: list[TraceInfo] | None = None,
        test_module_map: dict[str, list[Path]] | None = None,
    ) -> None:
        """Initialize the enricher.

        Args:
            specs: Parsed specifications.
            traces: Collected @traces annotations (source tree).
            project_id: Project tag written into every entity.
            test_traces: Collected @traces annotations from the tests tree
                (audit item 5 -- TESTS edges). Each becomes a ``TestCase``
                entity linked to the requirements/code it traces.
            test_module_map: Heuristic ``source module -> [test file]``
                mapping (see
                :func:`spec_trace_parser.map_test_modules_to_source`). Each
                mapped test file becomes a ``TestModule`` entity linked to
                the ``CodeElement``s of that source module.
        """
        self._specs = specs
        self._traces = traces
        self._project_id = project_id
        self._test_traces = test_traces or []
        self._test_module_map = test_module_map or {}
        self._entities: list[Entity] = []
        self._relationships: list[Relationship] = []
        self._req_to_code_ids: dict[str, set[str]] = {}
        self._summary = EnrichmentSummary()
        self._build()

    def _add_entity(self, entity: Entity) -> None:
        self._entities.append(entity)

    def _add_relationship(self, relationship: Relationship) -> None:
        self._relationships.append(relationship)

    def _build(self) -> None:
        """Construct all entities and relationships."""
        self._build_spec_entities()
        self._build_story_entities()
        self._build_requirement_entities()
        self._build_code_entities()
        self._build_trace_relationships()
        if self._test_traces:
            self._build_test_trace_entities()
        if self._test_module_map:
            self._build_test_module_entities()
        self._summary.specs = len(self._specs)
        self._summary.stories = sum(len(s.stories) for s in self._specs)
        self._summary.requirements = sum(len(s.requirements) for s in self._specs)
        self._summary.relationships = len(self._relationships)

    def _build_spec_entities(self) -> None:
        for spec in self._specs:
            entity_id = f"{self._project_id}:spec:{spec.spec_id}"
            self._add_entity(
                Entity(
                    id=entity_id,
                    labels=["Spec"],
                    properties={
                        "project_id": self._project_id,
                        "spec_id": spec.spec_id,
                        "slug": spec.slug,
                        "title": spec.title,
                        "status": spec.status,
                        "priority": spec.priority,
                        "author": spec.author,
                        "version": spec.version,
                        "components": spec.components,
                    },
                )
            )

    def _build_story_entities(self) -> None:
        for spec in self._specs:
            spec_entity_id = f"{self._project_id}:spec:{spec.spec_id}"
            for story in spec.stories:
                entity_id = f"{self._project_id}:story:{spec.spec_id}:{story.story_id}"
                self._add_entity(
                    Entity(
                        id=entity_id,
                        labels=["Story"],
                        properties={
                            "project_id": self._project_id,
                            "spec_id": spec.spec_id,
                            "story_id": story.story_id,
                            "title": story.title,
                            "role": story.role,
                            "want": story.want,
                            "so_that": story.so_that,
                        },
                    )
                )
                self._add_relationship(
                    Relationship(
                        id=f"{spec_entity_id}:HAS_STORY:{entity_id}",
                        type="HAS_STORY",
                        source_id=spec_entity_id,
                        target_id=entity_id,
                        properties={"project_id": self._project_id},
                    )
                )

    def _build_requirement_entities(self) -> None:
        for spec in self._specs:
            for req in spec.requirements:
                story_entity_id = (
                    f"{self._project_id}:story:{spec.spec_id}:{req.story_id}"
                )
                req_entity_id = f"{self._project_id}:req:{req.req_id}"
                self._add_entity(
                    Entity(
                        id=req_entity_id,
                        labels=["Requirement"],
                        properties={
                            "project_id": self._project_id,
                            "spec_id": spec.spec_id,
                            "story_id": req.story_id,
                            "req_id": req.req_id,
                            "title": req.title,
                            "checked": req.checked,
                        },
                    )
                )
                self._add_relationship(
                    Relationship(
                        id=f"{story_entity_id}:HAS_REQUIREMENT:{req_entity_id}",
                        type="HAS_REQUIREMENT",
                        source_id=story_entity_id,
                        target_id=req_entity_id,
                        properties={"project_id": self._project_id},
                    )
                )

    def _build_code_entities(self) -> None:
        seen: set[str] = set()
        for trace in self._traces:
            entity_id = _code_element_id(self._project_id, trace)
            if entity_id in seen:
                continue
            seen.add(entity_id)
            self._add_entity(
                Entity(
                    id=entity_id,
                    labels=["CodeElement"],
                    properties={
                        "project_id": self._project_id,
                        "module": trace.module,
                        "class_name": trace.class_name,
                        "method": trace.method,
                        "file_path": trace.file_path,
                        "line_number": trace.line_number,
                        "label": _code_element_label(trace),
                    },
                )
            )
        self._summary.code_elements = len(seen)

    def _build_trace_relationships(self) -> None:
        """Link code elements to the requirements they implement."""
        valid_reqs = {r.req_id for spec in self._specs for r in spec.requirements}
        for trace in self._traces:
            code_id = _code_element_id(self._project_id, trace)
            for req_id in trace.reqs:
                if req_id not in valid_reqs:
                    continue
                req_entity_id = f"{self._project_id}:req:{req_id}"
                self._req_to_code_ids.setdefault(req_id, set()).add(code_id)
                self._add_relationship(
                    Relationship(
                        id=f"{code_id}:IMPLEMENTS:{req_entity_id}",
                        type="IMPLEMENTS",
                        source_id=code_id,
                        target_id=req_entity_id,
                        properties={
                            "project_id": self._project_id,
                            "file_path": trace.file_path,
                            "line_number": trace.line_number,
                        },
                    )
                )

    def _build_test_trace_entities(self) -> None:
        """Build ``TestCase`` entities + ``TESTS`` edges from @traces'd tests.

        Each traced test function/method becomes a ``TestCase`` entity linked
        by a ``TESTS`` edge to every Requirement it references, plus (since
        those requirements already tell us which ``CodeElement``s implement
        them) a direct ``TESTS`` edge to those same code elements.
        """
        valid_reqs = {r.req_id for spec in self._specs for r in spec.requirements}
        seen: set[str] = set()

        for trace in self._test_traces:
            name = (
                f"{trace.class_name}.{trace.method}"
                if trace.class_name
                else trace.method
            )
            if not name:
                continue
            test_entity_id = f"{self._project_id}:test:{trace.module}:{name}"

            if test_entity_id not in seen:
                seen.add(test_entity_id)
                self._add_entity(
                    Entity(
                        id=test_entity_id,
                        labels=["TestCase"],
                        properties={
                            "project_id": self._project_id,
                            "module": trace.module,
                            "name": name,
                            "file_path": trace.file_path,
                            "line_number": trace.line_number,
                        },
                    )
                )

            for req_id in trace.reqs:
                if req_id not in valid_reqs:
                    continue
                req_entity_id = f"{self._project_id}:req:{req_id}"
                self._add_relationship(
                    Relationship(
                        id=f"{test_entity_id}:TESTS:{req_entity_id}",
                        type="TESTS",
                        source_id=test_entity_id,
                        target_id=req_entity_id,
                        properties={
                            "project_id": self._project_id,
                            "file_path": trace.file_path,
                            "line_number": trace.line_number,
                        },
                    )
                )
                for code_id in self._req_to_code_ids.get(req_id, ()):
                    self._add_relationship(
                        Relationship(
                            id=f"{test_entity_id}:TESTS:{code_id}",
                            type="TESTS",
                            source_id=test_entity_id,
                            target_id=code_id,
                            properties={"project_id": self._project_id},
                        )
                    )

        self._summary.test_cases = len(seen)

    def _build_test_module_entities(self) -> None:
        """Build ``TestModule`` entities + ``TESTS`` edges via path heuristic.

        Covers test files with no @traces coverage: a ``TestModule`` entity
        per heuristically-mapped test file is linked by ``TESTS`` edges to
        every ``CodeElement`` already known for that source module.
        """
        module_to_code_ids: dict[str, set[str]] = {}
        for trace in self._traces:
            code_id = _code_element_id(self._project_id, trace)
            module_to_code_ids.setdefault(trace.module, set()).add(code_id)

        seen: set[str] = set()
        for source_module, test_files in self._test_module_map.items():
            code_ids = module_to_code_ids.get(source_module)
            if not code_ids:
                continue
            for test_file in test_files:
                module_entity_id = (
                    f"{self._project_id}:testmodule:{test_file.as_posix()}"
                )
                if module_entity_id not in seen:
                    seen.add(module_entity_id)
                    self._add_entity(
                        Entity(
                            id=module_entity_id,
                            labels=["TestModule"],
                            properties={
                                "project_id": self._project_id,
                                "file_path": test_file.as_posix(),
                                "source_module": source_module,
                            },
                        )
                    )
                for code_id in code_ids:
                    self._add_relationship(
                        Relationship(
                            id=f"{module_entity_id}:TESTS:{code_id}",
                            type="TESTS",
                            source_id=module_entity_id,
                            target_id=code_id,
                            properties={"project_id": self._project_id},
                        )
                    )

        self._summary.test_modules = len(seen)

    @property
    def entities(self) -> list[Entity]:
        """All entities built by the enricher."""
        return self._entities

    @property
    def relationships(self) -> list[Relationship]:
        """All relationships built by the enricher."""
        return self._relationships

    @property
    def summary(self) -> EnrichmentSummary:
        """Summary of the enrichment graph."""
        return self._summary

    async def enrich(self, graph: SemanticGraph) -> EnrichmentSummary:
        """Write the spec/trace graph to the backend.

        Args:
            graph: An already-connected semantic graph.

        Returns:
            :class:`EnrichmentSummary` with counts.
        """
        if self._entities:
            await graph.upsert_entities(self._entities)
        if self._relationships:
            await graph.upsert_relationships(self._relationships)
        logger.info(
            "Enriched semantic graph for %s: %s",
            self._project_id,
            self._summary.to_dict(),
        )
        return self._summary


async def enrich_project_with_specs_and_traces(
    graph: SemanticGraph,
    specs_root: Path,
    source_root: Path,
    project_id: str,
    tests_root: Path | None = None,
    with_tests: bool = True,
) -> EnrichmentSummary:
    """High-level helper: parse specs/traces and write them to the graph.

    Args:
        graph: Connected semantic graph.
        specs_root: Root of the specs tree.
        source_root: Root of the source tree.
        project_id: Project tag.
        tests_root: Root of the tests tree (e.g. ``Path("tests")``). Required
            for TESTS-edge enrichment (audit item 5); ignored if
            ``with_tests`` is ``False``.
        with_tests: Whether to scan ``tests_root`` for @traces'd test
            functions/classes and heuristic test-module mappings and emit
            ``TestCase``/``TestModule`` entities + ``TESTS`` edges. Defaults
            to ``True``.

    Returns:
        :class:`EnrichmentSummary`.
    """
    specs = parse_all_specs(specs_root)
    traces = collect_code_traces(source_root)

    test_traces: list[TraceInfo] = []
    test_module_map: dict[str, list[Path]] = {}
    if with_tests and tests_root is not None and tests_root.is_dir():
        test_traces = collect_test_traces(tests_root)
        known_modules = {t.module for t in traces}
        test_module_map = map_test_modules_to_source(tests_root, known_modules)

    enricher = SpecTraceEnricher(
        specs,
        traces,
        project_id,
        test_traces=test_traces,
        test_module_map=test_module_map,
    )
    return await enricher.enrich(graph)