mlcompass 0.1.0__py3-none-any.whl

mlcompass/__init__.py

@@ -0,0 +1,3 @@
+"""mlcompass — An LLM agent that sits next to you through your whole ML pipeline."""
+
+__version__ = "0.1.0"

mlcompass/agents/__init__.py

@@ -0,0 +1,5 @@
+"""agentlite-based agents that drive each mlcompass command.
+
+Each module exposes one or more factory functions that return a
+configured ``agentlite.Agent`` ready to be used by the CLI.
+"""

mlcompass/agents/advise.py

@@ -0,0 +1,174 @@
+"""Model + feature engineering advisor.
+
+Consumes the structured output of ``tools.dataset.analyze_dataset`` and
+produces a JSON recommendation: models to try, feature engineering
+suggestions, and pitfalls to mitigate.
+
+The analysis is pre-computed by deterministic Python (no LLM cost), and
+fed to the agent as the user message. The advisor agent itself does not
+need tools — it is a pure reasoner over the structured input. This keeps
+``advise`` fast, cheap, and predictable.
+"""
+
+from __future__ import annotations
+
+import json
+from typing import Any
+
+from agentlite import Agent
+
+ADVISOR_MODEL_DEFAULT = "claude-opus-4-7"
+
+ADVISOR_SYSTEM_PROMPT = """You are a senior data scientist advising a colleague who just analyzed a new dataset. Your job is to recommend the next steps.
+
+You will receive a structured JSON dataset analysis. Based on it, produce:
+
+1. **Top 3 model families** to try, with one-line reasoning and a realistic
+   metric range (AUC for classification, RMSE/MAE for regression).
+2. **Feature engineering suggestions**: per-column or cross-column hints
+   that have a high likelihood of helping.
+3. **Pitfalls**: data-quality or methodology issues the user should mitigate
+   before training.
+
+Format your reply as a single JSON object matching this schema, and only
+that — no preamble, no markdown fence:
+
+{
+  "models": [
+    {"name": "XGBoost", "reason": "...", "expected_metric": "AUC 0.82 - 0.87"}
+  ],
+  "features": [
+    {"column": "signup_date", "suggestion": "derive days_since_signup, month, dayofweek", "reason": "..."}
+  ],
+  "pitfalls": [
+    {"issue": "Class imbalance (12% positive)", "mitigation": "Use AUC/F1, class_weight='balanced', or focal loss"}
+  ]
+}
+
+Rules:
+- Always include at least one interpretable baseline (logistic / linear
+  regression, decision tree, etc.) so the user has a sanity check.
+- Use realistic metric ranges based on the dataset signal — do not promise
+  numbers you can't back up.
+- Cite the column or fact you reasoned from inside the ``reason`` field.
+- When uncertain, prefer conservative, well-established choices."""
+
+
+# --------------------------------------------------------------------------- #
+# Agent construction                                                          #
+# --------------------------------------------------------------------------- #
+
+
+def build_advisor_agent(
+    *,
+    client: Any | None = None,
+    model: str = ADVISOR_MODEL_DEFAULT,
+) -> Agent:
+    """Build the model + feature engineering advisor agent.
+
+    Args:
+        client: Optional Anthropic client or ``MockClient`` for tests.
+            When ``None``, agentlite creates a real client from
+            ``ANTHROPIC_API_KEY``.
+        model: Claude model name. Opus is the default because the advisor's
+            reasoning needs to be sharp; Haiku tends to oversimplify on
+            tabular ML recommendations.
+
+    Returns:
+        A configured ``agentlite.Agent`` ready to accept a dataset analysis
+        as its user message.
+    """
+    return Agent(
+        model=model,
+        system=ADVISOR_SYSTEM_PROMPT,
+        tools=[],  # Pure reasoner — no tool calls needed.
+        client=client,
+        max_turns=2,  # Sanity cap; single-turn is the expected case.
+    )
+
+
+# --------------------------------------------------------------------------- #
+# Public API                                                                  #
+# --------------------------------------------------------------------------- #
+
+
+def get_recommendation(
+    analysis: dict[str, Any],
+    *,
+    client: Any | None = None,
+    model: str = ADVISOR_MODEL_DEFAULT,
+) -> dict[str, Any]:
+    """Run the advisor on a pre-computed dataset analysis.
+
+    Args:
+        analysis: Output of ``tools.dataset.analyze_dataset()``.
+        client: Optional client for testing.
+        model: Claude model name.
+
+    Returns:
+        Parsed recommendation dict with keys ``models``, ``features``,
+        ``pitfalls``.
+
+    Raises:
+        AdvisorParseError: If the agent's response can't be parsed as the
+            expected JSON shape.
+    """
+    agent = build_advisor_agent(client=client, model=model)
+
+    user_message = (
+        "Here is the analysis of a dataset. Please produce model, feature, "
+        "and pitfall recommendations as specified.\n\n"
+        f"```json\n{json.dumps(analysis, indent=2)}\n```"
+    )
+
+    raw_response = agent.run(user_message)
+    return _parse_advisor_response(raw_response)
+
+
+# --------------------------------------------------------------------------- #
+# Internal: response parsing                                                  #
+# --------------------------------------------------------------------------- #
+
+
+class AdvisorParseError(ValueError):
+    """Raised when the advisor returns something other than valid JSON."""
+
+
+def _parse_advisor_response(text: str) -> dict[str, Any]:
+    """Strip optional markdown fences and parse JSON."""
+    stripped = text.strip()
+
+    # The system prompt explicitly forbids fences, but be lenient.
+    if stripped.startswith("```"):
+        lines = stripped.splitlines()
+        # Drop opening fence (``` or ```json) and find the closing fence
+        start = 1
+        end = len(lines)
+        for i, line in enumerate(lines[1:], 1):
+            if line.strip().startswith("```"):
+                end = i
+                break
+        stripped = "\n".join(lines[start:end])
+
+    try:
+        parsed = json.loads(stripped)
+    except json.JSONDecodeError as exc:
+        snippet = text[:200].replace("\n", " ")
+        raise AdvisorParseError(
+            f"Advisor response was not valid JSON: {snippet!r}"
+        ) from exc
+
+    if not isinstance(parsed, dict):
+        raise AdvisorParseError(
+            f"Advisor response was JSON but not an object (got {type(parsed).__name__})"
+        )
+
+    # Light shape validation — keep it loose so the schema can evolve.
+    for required_key in ("models", "features", "pitfalls"):
+        if required_key not in parsed:
+            raise AdvisorParseError(
+                f"Advisor response missing required key '{required_key}'. "
+                f"Got: {sorted(parsed.keys())}"
+            )
+
+    return parsed

mlcompass/cli.py

@@ -0,0 +1,293 @@
+"""Command-line interface for ``mlcompass``.
+
+Subcommands are added incrementally as the project advances through
+its phases. See ARCHITECTURE.md §7 for the full CLI design.
+
+Currently implemented:
+    init    — create a new ``.mlcompass/`` project (Faz 1)
+    advise  — analyze dataset + recommend models / features / pitfalls (Faz 1)
+
+Planned:
+    audit, watch, compare, evaluate, deploy, status
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Callable
+
+import click
+from rich.console import Console
+from rich.panel import Panel
+
+from . import __version__
+from .agents.advise import AdvisorParseError, get_recommendation
+from .context import ProjectContext, ProjectExistsError, ProjectNotFoundError
+from .tools.dataset import analyze_dataset
+from .ui.advise import render_analysis, render_recommendation
+
+
+def _force_utf8_stdio() -> None:
+    """Reconfigure stdout/stderr to UTF-8.
+
+    On Windows the default encoding follows the system code page
+    (e.g. ``cp1254`` for Turkish locales) and cannot represent the
+    Unicode glyphs rich uses for panels, checkmarks, and box drawing.
+    Reconfiguring early avoids ``UnicodeEncodeError`` at print time.
+    """
+    for stream in (sys.stdout, sys.stderr):
+        reconfigure = getattr(stream, "reconfigure", None)
+        if reconfigure is None:
+            continue
+        try:
+            reconfigure(encoding="utf-8", errors="replace")
+        except (AttributeError, ValueError):
+            pass
+
+
+_force_utf8_stdio()
+console = Console()
+
+
+@click.group(
+    help="mlcompass — your AI ML engineer at every pipeline stage.",
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+@click.version_option(__version__, prog_name="mlcompass")
+def cli() -> None:
+    """Root command group."""
+
+
+@cli.command(help="Initialize a new mlcompass project.")
+@click.argument("name")
+@click.option(
+    "--path",
+    "parent_dir",
+    type=click.Path(file_okay=False, dir_okay=True, path_type=Path),
+    default=Path("."),
+    show_default=True,
+    help="Directory in which .mlcompass/ will be created.",
+)
+@click.option(
+    "--default-model",
+    default="claude-opus-4-7",
+    show_default=True,
+    help="Default LLM model for this project.",
+)
+def init(name: str, parent_dir: Path, default_model: str) -> None:
+    """Create a new mlcompass project under ``parent_dir/.mlcompass/``."""
+    try:
+        ctx = ProjectContext.init(
+            name=name,
+            parent_dir=parent_dir,
+            default_model=default_model,
+        )
+    except ProjectExistsError as e:
+        console.print(f"[red]✗[/red] {e}")
+        raise SystemExit(1) from e
+
+    console.print(
+        Panel.fit(
+            f"[green]✓[/green] Project [bold]{name}[/bold] initialized.\n\n"
+            f"Directory:  [cyan]{ctx.path}[/cyan]\n"
+            f"Model:      {default_model}\n\n"
+            f"Next: [yellow]mlcompass advise <data.csv>[/yellow]",
+            title="mlcompass init",
+            border_style="green",
+        )
+    )
+
+
+@cli.command(
+    help="Analyze a dataset and recommend models, features, and pitfalls."
+)
+@click.argument(
+    "dataset_path",
+    type=click.Path(exists=True, dir_okay=False, path_type=Path),
+)
+@click.option(
+    "--target",
+    "target_column",
+    default=None,
+    help="Target column name (auto-detected from name conventions if omitted).",
+)
+@click.option(
+    "--sample-rows",
+    type=int,
+    default=None,
+    help="Limit analysis to the first N rows (useful for very large files).",
+)
+@click.option(
+    "--no-llm",
+    is_flag=True,
+    help="Skip the LLM advisor step; show only the deterministic analysis.",
+)
+@click.option(
+    "--model",
+    "advisor_model",
+    default="claude-opus-4-7",
+    show_default=True,
+    help="Claude model used by the advisor.",
+)
+def advise(
+    dataset_path: Path,
+    target_column: str | None,
+    sample_rows: int | None,
+    no_llm: bool,
+    advisor_model: str,
+) -> None:
+    """Run the deterministic dataset analyzer, then the LLM advisor."""
+    project = _try_load_project()
+
+    with console.status("[cyan]Analyzing dataset...[/cyan]", spinner="dots"):
+        analysis = analyze_dataset(
+            dataset_path,
+            target_column=target_column,
+            sample_rows=sample_rows,
+        )
+
+    render_analysis(console, analysis)
+
+    recommendation: dict[str, Any] | None = None
+
+    if no_llm:
+        console.print("\n[dim](--no-llm specified; skipping advisor)[/dim]")
+    elif not _has_api_key():
+        console.print(
+            "\n[yellow]⚠ ANTHROPIC_API_KEY not set; "
+            "skipping advisor step. Set the variable to enable.[/yellow]"
+        )
+    else:
+        recommendation = _run_advisor(analysis, model=advisor_model)
+        if recommendation is not None:
+            render_recommendation(console, recommendation)
+
+    if project is not None:
+        _persist_advise_result(
+            project=project,
+            dataset_path=dataset_path,
+            analysis=analysis,
+            recommendation=recommendation,
+        )
+
+
+# --------------------------------------------------------------------------- #
+# advise helpers (split out so tests can monkeypatch them)                    #
+# --------------------------------------------------------------------------- #
+
+
+def _try_load_project() -> ProjectContext | None:
+    """Load an existing project context, or warn and return None."""
+    try:
+        return ProjectContext.load()
+    except ProjectNotFoundError:
+        console.print(
+            "[dim](no .mlcompass/ project found in current path — "
+            "running standalone; results will not be persisted)[/dim]\n"
+        )
+        return None
+
+
+def _has_api_key() -> bool:
+    """True iff an Anthropic API key is configured in the environment."""
+    return bool(os.environ.get("ANTHROPIC_API_KEY"))
+
+
+# Indirection point: tests monkeypatch ``cli._advisor_callable`` to inject a
+# fake instead of calling the real ``get_recommendation``.
+_advisor_callable: Callable[..., dict[str, Any]] = get_recommendation
+
+
+def _run_advisor(
+    analysis: dict[str, Any],
+    *,
+    model: str,
+) -> dict[str, Any] | None:
+    """Invoke the LLM advisor, handling parse errors gracefully."""
+    try:
+        with console.status(
+            "[cyan]Consulting model advisor...[/cyan]",
+            spinner="dots",
+        ):
+            return _advisor_callable(analysis, model=model)
+    except AdvisorParseError as exc:
+        console.print(f"\n[red]✗ Advisor returned an invalid response:[/red] {exc}")
+        return None
+
+
+def _persist_advise_result(
+    *,
+    project: ProjectContext,
+    dataset_path: Path,
+    analysis: dict[str, Any],
+    recommendation: dict[str, Any] | None,
+) -> None:
+    """Save dataset metadata + context updates + advice log entry."""
+    fingerprint = project.register_dataset(
+        dataset_path,
+        meta={"analysis": analysis},
+    )
+
+    project.write_context(
+        {
+            "active_dataset": f"datasets/{fingerprint}.json",
+            "project_type": analysis["task_hint"].get("type"),
+            "target_column": analysis["target_hint"].get("column"),
+        }
+    )
+
+    top_model: str
+    if recommendation and recommendation.get("models"):
+        top_model = recommendation["models"][0].get("name", "n/a")
+    else:
+        top_model = "n/a (LLM advisor skipped)"
+
+    project.append_decision(
+        command="advise",
+        summary=f"Top model recommendation: {top_model}",
+        reasoning=(
+            f"Dataset: {dataset_path}, "
+            f"task: {analysis['task_hint'].get('type', 'unknown')}"
+        ),
+    )
+
+    _append_advice_log(project, dataset_path, analysis, recommendation)
+
+    console.print(f"\n[green]✓[/green] Saved to {project.path}")
+
+
+def _append_advice_log(
+    project: ProjectContext,
+    dataset_path: Path,
+    analysis: dict[str, Any],
+    recommendation: dict[str, Any] | None,
+) -> None:
+    """Append a single advice entry to ``.mlcompass/advice.log``."""
+    entry = {
+        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "dataset": str(dataset_path),
+        "task_type": analysis["task_hint"].get("type"),
+        "target": analysis["target_hint"].get("column"),
+        "recommendation": recommendation,
+    }
+    log_path = project.path / "advice.log"
+    with log_path.open("a", encoding="utf-8") as f:
+        f.write(json.dumps(entry) + "\n")
+
+
+# --------------------------------------------------------------------------- #
+# Entry point                                                                 #
+# --------------------------------------------------------------------------- #
+
+
+def main() -> None:
+    """Entry point for the ``mlcompass`` console script."""
+    cli()
+
+
+if __name__ == "__main__":  # pragma: no cover
+    main()

mlcompass/context.py

@@ -0,0 +1,222 @@
+"""Project context: persistent state stored in ``.mlcompass/``.
+
+The directory layout is described in ``ARCHITECTURE.md §2``.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+from . import __version__
+
+DEFAULT_PROJECT_DIR = ".mlcompass"
+
+
+class ProjectExistsError(FileExistsError):
+    """Raised when init is called but ``.mlcompass/`` already exists."""
+
+
+class ProjectNotFoundError(FileNotFoundError):
+    """Raised when load is called but no ``.mlcompass/`` is found."""
+
+
+@dataclass
+class ProjectContext:
+    """Wraps a single ``.mlcompass/`` project directory.
+
+    Layout::
+
+        .mlcompass/
+        ├── project.yaml      # static metadata
+        ├── context.json      # dynamic state (decisions, recommendations)
+        ├── datasets/         # registered datasets
+        ├── runs/             # training run history
+        ├── advice.log        # advisor recommendation history
+        └── cache/            # tool result + LLM prompt cache
+    """
+
+    path: Path  # the ``.mlcompass/`` directory itself
+
+    # ---------------------- Construction ----------------------
+
+    @classmethod
+    def init(
+        cls,
+        name: str,
+        parent_dir: Path | str = ".",
+        *,
+        default_model: str = "claude-opus-4-7",
+    ) -> "ProjectContext":
+        """Create a new ``.mlcompass/`` directory under ``parent_dir``.
+
+        Args:
+            name: Human-readable project name (e.g., ``"churn-model"``).
+            parent_dir: Directory in which ``.mlcompass/`` will be created.
+            default_model: Default LLM model name for the project.
+
+        Returns:
+            The newly created ``ProjectContext``.
+
+        Raises:
+            ProjectExistsError: If a project already exists at the target.
+        """
+        parent_dir = Path(parent_dir).resolve()
+        target = parent_dir / DEFAULT_PROJECT_DIR
+
+        if target.exists():
+            raise ProjectExistsError(f"Project already exists at {target}")
+
+        # Directory tree
+        target.mkdir(parents=True)
+        (target / "datasets").mkdir()
+        (target / "runs").mkdir()
+        (target / "cache").mkdir()
+
+        # Static metadata
+        project_meta = {
+            "name": name,
+            "created": datetime.now(timezone.utc).isoformat(),
+            "mlcompass_version": __version__,
+            "default_model": default_model,
+        }
+        (target / "project.yaml").write_text(
+            yaml.safe_dump(project_meta, sort_keys=False),
+            encoding="utf-8",
+        )
+
+        # Dynamic context starts empty
+        initial_context: dict[str, Any] = {
+            "project_type": None,
+            "target_column": None,
+            "preferred_models": [],
+            "active_dataset": None,
+            "current_run": None,
+            "decisions": [],
+        }
+        (target / "context.json").write_text(
+            json.dumps(initial_context, indent=2),
+            encoding="utf-8",
+        )
+
+        # Advice log starts empty
+        (target / "advice.log").touch()
+
+        # Local .gitignore so cache/ and runs/ don't pollute the user's repo
+        (target / ".gitignore").write_text(
+            "cache/\nruns/\n*.pyc\n__pycache__/\n",
+            encoding="utf-8",
+        )
+
+        return cls(path=target)
+
+    @classmethod
+    def load(cls, search_from: Path | str = ".") -> "ProjectContext":
+        """Find and load an existing project by walking up from ``search_from``.
+
+        Mirrors the behaviour of ``git`` when discovering a repository.
+
+        Raises:
+            ProjectNotFoundError: If no ``.mlcompass/`` is found between
+                ``search_from`` and the filesystem root.
+        """
+        current = Path(search_from).resolve()
+        for candidate in [current, *current.parents]:
+            target = candidate / DEFAULT_PROJECT_DIR
+            if target.is_dir():
+                return cls(path=target)
+        raise ProjectNotFoundError(
+            f"No {DEFAULT_PROJECT_DIR}/ found at or above {current}. "
+            "Run `mlcompass init <name>` to create one."
+        )
+
+    # ---------------------- Read / write ----------------------
+
+    @property
+    def project_meta(self) -> dict[str, Any]:
+        """Static metadata loaded from ``project.yaml``."""
+        return yaml.safe_load(
+            (self.path / "project.yaml").read_text(encoding="utf-8")
+        )
+
+    def read_context(self) -> dict[str, Any]:
+        """Read the dynamic context (``context.json``)."""
+        return json.loads(
+            (self.path / "context.json").read_text(encoding="utf-8")
+        )
+
+    def write_context(self, updates: dict[str, Any]) -> None:
+        """Merge ``updates`` into ``context.json``.
+
+        Top-level keys are replaced. For lists like ``decisions`` prefer
+        ``append_decision`` so timestamps are added automatically.
+        """
+        current = self.read_context()
+        current.update(updates)
+        (self.path / "context.json").write_text(
+            json.dumps(current, indent=2),
+            encoding="utf-8",
+        )
+
+    def append_decision(
+        self,
+        command: str,
+        summary: str,
+        *,
+        reasoning: str = "",
+    ) -> None:
+        """Append a timestamped decision entry to ``decisions``."""
+        ctx = self.read_context()
+        ctx.setdefault("decisions", []).append(
+            {
+                "timestamp": datetime.now(timezone.utc).isoformat(),
+                "command": command,
+                "summary": summary,
+                "reasoning": reasoning,
+            }
+        )
+        (self.path / "context.json").write_text(
+            json.dumps(ctx, indent=2),
+            encoding="utf-8",
+        )
+
+    # ---------------------- Dataset registry ----------------------
+
+    def register_dataset(
+        self,
+        dataset_path: Path | str,
+        meta: dict[str, Any],
+    ) -> str:
+        """Save dataset metadata under ``datasets/<fingerprint>.json``.
+
+        The fingerprint is the truncated SHA-256 of the absolute path plus
+        the file's modification time, so re-registering an unchanged file
+        is idempotent.
+
+        Returns:
+            The fingerprint (16-character hex).
+        """
+        dataset_path = Path(dataset_path).resolve()
+        fingerprint = self._fingerprint(dataset_path)
+        record = {
+            "path": str(dataset_path),
+            "registered_at": datetime.now(timezone.utc).isoformat(),
+            **meta,
+        }
+        (self.path / "datasets" / f"{fingerprint}.json").write_text(
+            json.dumps(record, indent=2),
+            encoding="utf-8",
+        )
+        return fingerprint
+
+    @staticmethod
+    def _fingerprint(p: Path) -> str:
+        """Stable per-file fingerprint based on path + mtime."""
+        marker = f"{p.resolve()}::{p.stat().st_mtime_ns}".encode("utf-8")
+        return hashlib.sha256(marker).hexdigest()[:16]