orchestrata_ultimate.context.providers.ast_context

Reference › orchestrata_ultimate › context › providers › ast_context

This provider gives the volk deterministic, AST-derived signals about the files it is about to change. No LLM takes part in selecting them. The assembler calls it for the ast layer. Tests are tests/context/providers/test_ast_context.py and test_ast_context_coverage.py. render_ast_context carries @traces for REQ-107-014. The underlying AST library is orchestrata_ast.

_load_ast_tool loads the serializer script scripts/serialize_codebase_ast.py by file path at run time and caches the module per process, so the package has no static dependency on scripts/. The script path is found with resolve_project_root. _resolve_sources turns the target files into existing .py files or directories, trying the raw path and the path under the project root. render_ast_context then asks the toolkit for public API signatures, complexity hotspots above a threshold of five, branch points, functions missing @traces, and a call-graph neighbourhood counting definitions and uses of up to three public symbols. Each section is limited to five entries and wrapped in its own try block that logs at debug level, so one failure drops one section. The block is capped at max_lines (default 30). An optional project id is added to the header. If the toolkit cannot be loaded, a warning is logged and the result is an empty string.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

API

Constants

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/src/orchestrata_ultimate/context/providers/ast_context.py

"""AST-derived context provider for volk prompts.

Wraps the AST toolkit (``scripts/serialize_codebase_ast.py``) so the unified
context assembler can inject compact, deterministic signals about the target
code: complexity hotspots, branch counts, missing ``@traces`` coverage, and a
small call-graph neighborhood. All signals are computed from real source files;
no LLM is involved in selection.

When the toolkit cannot be loaded or the target files cannot be parsed, the
provider degrades to an empty string.
"""

from __future__ import annotations

import importlib.util
import logging
from pathlib import Path
from typing import Any

from orchestrata_ultimate.core.base.traces import traces
from orchestrata_ultimate.core.paths import resolve_project_root

_logger = logging.getLogger(__name__)

#: Resolved through ``resolve_project_root``, never a fixed ``__file__`` parent
#: count: since the proper install (dd919d55c) this module sits in
#: site-packages, where four parents up is ``.venv/lib/python3.14t``.
_TOOL_PATH = resolve_project_root() / "scripts" / "serialize_codebase_ast.py"

# Module-level cache: the AST toolkit is ~10k LOC; loading it once per process
# avoids repeated import/exec cost on every volk call. The cache key is the
# resolved toolkit path so tests can pass an alternate path without collision.
_TOOL_CACHE: dict[Path, Any] = {}


def _load_ast_tool(tool_path: Path | None = None) -> Any:
    """Load the AST toolkit module on demand.

    Dynamic import avoids making ``scripts/`` a static dependency of the
    ``knowledge_graph`` package, while still reusing the existing analysis
    implementation rather than duplicating it. The loaded module is cached
    per process.

    Args:
        tool_path: Optional override path to the toolkit module. Tests use this
            to exercise load failures without patching global state.
    """
    path = tool_path or _TOOL_PATH
    path = path.resolve()
    if path in _TOOL_CACHE:
        return _TOOL_CACHE[path]
    spec = importlib.util.spec_from_file_location("ast_toolkit", path)
    if spec is None or spec.loader is None:
        raise ImportError(f"cannot load AST toolkit from {path}")
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    _TOOL_CACHE[path] = mod
    return mod


def _resolve_sources(target_files: list[str], project_root: str) -> list[str]:
    """Return parseable source paths for the toolkit.

    Prefers paths that exist under *project_root*; falls back to the target
    path itself so the provider still works when callers pass absolute paths.
    """
    root = Path(project_root).resolve()
    sources: list[str] = []
    for raw in target_files:
        candidates = [raw]
        try:
            rel = root / raw
            candidates.append(str(rel))
        except Exception:
            pass
        for cand in candidates:
            p = Path(cand)
            if p.is_file() and p.suffix == ".py":
                sources.append(str(p.resolve()))
                break
            if p.is_dir():
                sources.append(str(p.resolve()))
                break
    return sources


def _public_symbols(tool: Any, sources: list[str]) -> list[dict[str, Any]]:
    """Top-level public classes/functions in the target sources."""
    rows: list[dict[str, Any]] = tool.list_symbols(sources)
    return [r for r in rows if "." not in r["name"] and not r["name"].startswith("_")]


def _public_api_signatures(
    tool: Any, sources: list[str], max_symbols: int = 5
) -> list[str]:
    """Return compact signature lines for the top-level public API."""
    symbols = _public_symbols(tool, sources)[:max_symbols]
    lines: list[str] = []
    for sym in symbols:
        sig = sym.get("signature", "")
        if sig:
            lines.append(f"- `{sig}` @ {sym.get('file', '?')}:{sym.get('line', 0)}")
    return lines


@traces(spec="107", reqs=["REQ-107-014"])
def render_ast_context(
    target_files: list[str],
    subject: str = "",
    project_root: str = ".",
    max_lines: int = 30,
    tool_path: Path | None = None,
    project_id: str | None = None,
) -> str:
    """Render a compact, budgeted AST context block.

    Args:
        target_files: Files the volk is being asked to act on.
        subject: CXP subject (currently unused for filtering, reserved).
        project_root: Root used to resolve relative target paths.
        max_lines: Hard line cap to keep the block inside its token budget.
        tool_path: Optional override path to the AST toolkit module. Primarily
            for tests; production callers omit this.
        project_id: Optional project scope. When provided, the rendered block
            identifies the project so AST signals are not confused across
            multi-project workspaces.

    Returns:
        Markdown AST block or "" when no signals are available.
    """
    if not target_files:
        return ""

    try:
        tool = _load_ast_tool(tool_path)
    except Exception as exc:
        _logger.warning("AST context toolkit unavailable: %s", exc)
        return ""

    sources = _resolve_sources(target_files, project_root)
    if not sources:
        return ""

    lines: list[str] = ["# AST CONTEXT"]
    if project_id:
        lines.append(f"Project: `{project_id}`")

    try:
        sig_lines = _public_api_signatures(tool, sources, max_symbols=5)
        if sig_lines:
            lines.append("## Public API signatures")
            lines.extend(sig_lines)
    except Exception as exc:
        _logger.debug("AST public API extraction failed: %s", exc)

    try:
        hotspots = tool.measure_complexity(sources, threshold=5)
        if hotspots:
            lines.append("## Complexity hotspots")
            for row in hotspots[:5]:
                lines.append(
                    f"- `{row['name']}` complexity={row['complexity']} "
                    f"@ {row['file']}:{row['line']}"
                )
    except Exception as exc:
        _logger.debug("AST complexity extraction failed: %s", exc)

    try:
        branches = tool.list_branches(None, sources)
        if branches:
            lines.append("## Branch points")
            for row in branches[:5]:
                count = len(row.get("branches", []))
                lines.append(
                    f"- `{row['function']}` {count} branches "
                    f"@ {row['file']}:{row['line']}"
                )
    except Exception as exc:
        _logger.debug("AST branch extraction failed: %s", exc)

    try:
        untraced = tool.find_untraced(sources)
        if untraced:
            lines.append("## Missing @traces")
            for row in untraced[:5]:
                lines.append(f"- `{row['name']}` @ {row['file']}:{row['line']}")
    except Exception as exc:
        _logger.debug("AST trace-lint extraction failed: %s", exc)

    # Call-graph neighborhood: for a few public symbols, report internal uses.
    try:
        symbols = _public_symbols(tool, sources)[:3]
        if symbols:
            cg_lines: list[str] = []
            for sym in symbols:
                name = sym["name"]
                refs = tool.find_references(name, sources)
                defs = refs.get("definitions", [])
                uses = refs.get("uses", [])
                if defs or uses:
                    cg_lines.append(
                        f"- `{name}` defined {len(defs)}x, used {len(uses)}x"
                    )
            if cg_lines:
                lines.append("## Call-graph neighborhood")
                lines.extend(cg_lines)
    except Exception as exc:
        _logger.debug("AST call-graph extraction failed: %s", exc)

    while len(lines) > max_lines:
        lines.pop()

    if len(lines) == 1:
        return ""

    return "\n".join(lines)