vasiuuu/DGX_AI

1from __future__ import annotations2 3import ast4import importlib5import importlib.util6import logging7from typing import Literal8 9from pydantic import BaseModel, ConfigDict10 11_log = logging.getLogger(__name__)12 13 14# ---------------------------------------------------------------------------15# Models16# ---------------------------------------------------------------------------17 18 19class Symbol(BaseModel):20    """A single symbol extracted from source code by AST walking."""21 22    model_config = ConfigDict(frozen=True)23    module: str24    attr: str | None25    kind: Literal["import", "attribute"]26    resolved: bool27    line: int28 29 30class GroundingReport(BaseModel):31    """Result of grounding analysis on source code."""32 33    model_config = ConfigDict(frozen=True)34    total_symbols: int35    grounded: tuple[Symbol, ...]36    ungrounded: tuple[Symbol, ...]37    groundedness: float38 39 40# ---------------------------------------------------------------------------41# Helpers42# ---------------------------------------------------------------------------43 44 45def _module_spec(name: str) -> bool:46    """Return True if the module can be found by the import system."""47    try:48        return importlib.util.find_spec(name) is not None49    except (ImportError, ValueError, ModuleNotFoundError):50        return False51 52 53def _has_attr(module_name: str, attr: str) -> bool:54    """Check if *module_name* exposes *attr*.55 56    Uses the FULL module path (e.g. ``os.path``) — not just57    the top-level package.  This is the fix for SYSTEM_DESIGN §4.8.358    bug #3.59    """60    try:61        mod = importlib.import_module(module_name)62    except Exception:63        return False64    return hasattr(mod, attr)65 66 67# ---------------------------------------------------------------------------68# Public API69# ---------------------------------------------------------------------------70 71 72def ground(73    source: str,74    *,75    local_modules: frozenset[str] = frozenset(),76) -> GroundingReport:77    """AST-parse *source*, check every import and attribute access resolves.78 79    Three fixes baked in from day one (SYSTEM_DESIGN §4.8.3):80    1. SyntaxError → groundedness=0.0  (was 1.0)81    2. Zero symbols → groundedness=0.5 (was 1.0)82    3. Attribute resolution against full module path (was top-level only)83 84    *local_modules*: set of module names (e.g. ``{"core", "main"}``) that are85    local to the agent's project and should be treated as grounded even though86    ``importlib.util.find_spec`` cannot resolve them from the grader process.87    """88    # ----- parse --------------------------------------------------------89    try:90        tree = ast.parse(source)91    except SyntaxError:92        # FIX 1: unparseable code → 0.0, not 1.093        return GroundingReport(94            total_symbols=0,95            grounded=(),96            ungrounded=(),97            groundedness=0.0,98        )99 100    symbols: list[Symbol] = []101    import_to_module: dict[str, str] = {}102 103    # ----- walk imports -------------------------------------------------104    for node in ast.walk(tree):105        if isinstance(node, ast.Import):106            for alias in node.names:107                pkg = alias.name.split(".")[0]108                # Local modules are always treated as grounded109                resolved = (110                    pkg in local_modules or _module_spec(alias.name)111                )112                symbols.append(113                    Symbol(114                        module=alias.name,115                        attr=None,116                        kind="import",117                        resolved=resolved,118                        line=node.lineno,119                    )120                )121                import_to_module[alias.asname or pkg] = alias.name122 123        elif isinstance(node, ast.ImportFrom):124            if node.level != 0 or node.module is None:125                continue126            mod_top = node.module.split(".")[0]127            is_local = mod_top in local_modules128            resolved_mod = is_local or _module_spec(node.module)129            for alias in (node.names or []):130                attr_resolved = resolved_mod if is_local else (131                    resolved_mod and _has_attr(node.module, alias.name)132                )133                symbols.append(134                    Symbol(135                        module=node.module,136                        attr=alias.name,137                        kind="import",138                        resolved=attr_resolved,139                        line=node.lineno,140                    )141                )142 143    # ----- walk attribute accesses --------------------------------------144    for node in ast.walk(tree):145        if not isinstance(node, ast.Attribute):146            continue147 148        # Resolve the chain: e.g. os.path.join → base="os", chain=["path"], attr="join"149        chain: list[str] = []150        cursor: ast.expr = node.value151        while isinstance(cursor, ast.Attribute):152            chain.append(cursor.attr)153            cursor = cursor.value154        if not isinstance(cursor, ast.Name):155            continue156 157        base = cursor.id158        mod_name = import_to_module.get(base)159        if mod_name is None:160            continue161 162        # Build the full module path for chained access:163        # import os.path → import_to_module["os"] = "os.path"164        # os.path.join → chain=["path"], we need to resolve "join" against "os.path"165        # The chain intermediates are sub-module parts already covered by mod_name.166        # We check the final attr against the deepest resolvable module.167        if chain:168            # chain was built bottom-up, reverse to get top-down order169            chain.reverse()170            # Build candidate module: mod_name + chain parts171            full_mod = mod_name + "." + ".".join(chain)172            # Try the full module first; fall back to mod_name if it doesn't exist173            check_mod = full_mod if _module_spec(full_mod) else mod_name174        else:175            check_mod = mod_name176 177        # FIX 3: resolve against full module path, not just top-level178        resolved = _has_attr(check_mod, node.attr)179        symbols.append(180            Symbol(181                module=check_mod,182                attr=node.attr,183                kind="attribute",184                resolved=resolved,185                line=node.lineno,186            )187        )188 189    # ----- compute groundedness -----------------------------------------190    grounded = tuple(s for s in symbols if s.resolved)191    ungrounded = tuple(s for s in symbols if not s.resolved)192    total = len(symbols)193 194    # FIX 2: zero symbols → 0.5 (neutral), not 1.0195    groundedness = 0.5 if total == 0 else len(grounded) / total196 197    return GroundingReport(198        total_symbols=total,199        grounded=grounded,200        ungrounded=ungrounded,201        groundedness=groundedness,202    )203