codegraph/graph.py

"""
codegraph/graph.py — CodeGraph V2

The innovation that makes SecureCodeEnv unique.
Structured in-memory database of everything the agent has written this episode.
Persisted in Redis between steps via pickle.

V2 changes:
  - tree-sitter replaces ast module → supports Python, JS, TS, TSX
  - 60% threshold for style detection (was 50%) → prevents false penalties
  - "mixed" state added → no penalty when codebase has no clear dominant style
  - compress_graph() added → semantic compression for inference context
"""
from dataclasses import dataclass, field
from collections import Counter
from typing import Dict, Any


@dataclass
class CodeGraph:
    episode_seed: int = 0
    components: Dict[str, Dict[str, Any]] = field(default_factory=dict)
    conventions: Dict[str, Any] = field(default_factory=dict)

    def update(self, filename: str, metadata: Dict[str, Any]) -> None:
        """Add or replace a file's metadata in the graph, then re-derive conventions."""
        if metadata.get("status") == "syntax_error":
            return   # Don't pollute graph with broken code
        name = _file_to_key(filename)
        metadata["file"] = filename
        self.components[name] = metadata
        self._infer_conventions()

    def _infer_conventions(self) -> None:
        """
        Derive dominant codebase style from all components.
        60% threshold: a bare majority (51%) wrongly penalises mixed codebases.
        When no clear style → 'mixed' → consistency grader awards full marks.
        """
        all_fns = [
            f["name"]
            for comp in self.components.values()
            for f in comp.get("functions", [])
        ]
        if all_fns:
            styles = [_naming_style(n) for n in all_fns]
            top, count = Counter(styles).most_common(1)[0]
            self.conventions["naming"] = top if count / len(styles) >= 0.60 else "mixed"
        else:
            self.conventions["naming"] = "unknown"

        uses_try = sum(
            1 for c in self.components.values()
            if c.get("conventions", {}).get("uses_try_catch", False)
        )
        total = len(self.components)
        self.conventions["error_handling"] = "try_catch" if uses_try / max(total, 1) >= 0.5 else "none"

        uses_hints = sum(
            1 for c in self.components.values()
            if c.get("conventions", {}).get("uses_type_hints", False)
        )
        self.conventions["uses_type_hints"] = uses_hints / max(total, 1) >= 0.5

    def to_slim_dict(self, limit: int = 6000) -> str:
        """
        compress_graph() — semantic compression for inference.py context.
        Keeps signatures + conventions, drops function bodies.
        V1 blindly truncated at 2000 chars → agents couldn't see patterns they needed.
        """
        import json
        slim = {
            "conventions": self.conventions,
            "components": {
                name: {
                    "file": comp.get("file", ""),
                    "language": comp.get("language", "py"),
                    "functions": [f["name"] for f in comp.get("functions", [])][:20],
                    "imports": [i.split(".")[0] for i in comp.get("imports", [])][:15],
                    "uses_try_catch": comp.get("conventions", {}).get("uses_try_catch", False),
                    "uses_type_hints": comp.get("conventions", {}).get("uses_type_hints", False),
                }
                for name, comp in self.components.items()
            },
        }
        result = json.dumps(slim, indent=2)
        if len(result) > limit:
            # Further trim: drop imports when still over limit
            for name in slim["components"]:
                slim["components"][name].pop("imports", None)
            result = json.dumps(slim, indent=2)[:limit]
        return result


# ── helpers ──────────────────────────────────────────────────────────────────

def _file_to_key(filename: str) -> str:
    """Convert 'src/auth/UserAuth.py' → 'UserAuth'"""
    base = filename.split("/")[-1]
    for ext in (".py", ".js", ".ts", ".tsx", ".jsx"):
        base = base.replace(ext, "")
    return base


def _naming_style(name: str) -> str:
    if "_" in name:
        return "snake_case"
    if name and name[0].isupper():
        return "PascalCase"
    if any(c.isupper() for c in name[1:]):
        return "camelCase"
    return "snake_case"   # all-lowercase defaults to snake