gradex 0.1.0__py3-none-any.whl

gradex/__init__.py

@@ -0,0 +1,3 @@
+"""gradex: manage and interact with coding host plugins."""
+
+__version__ = "0.1.0"

gradex/ai/__init__.py

@@ -0,0 +1 @@
+"""evo AI layer — LLM client, discover skill, and brief generator."""

gradex/ai/brief.py

@@ -0,0 +1,94 @@
+"""Brief generator: render the per-experiment optimisation prompt template."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import jinja2
+
+from gradex.ai.client import LLMClient
+
+_PROMPTS_DIR = Path(__file__).parent / "prompts"
+
+
+@dataclass
+class ExperimentSummary:
+    """A compact summary of one past experiment for inclusion in a brief."""
+
+    hypothesis: str
+    result: str  # "improved" | "regressed" | "gate_failed" | "failed"
+    reason: str  # e.g. "score went from 41.2 to 45.0 (regression)"
+
+
+class BriefGenerator:
+    """Renders ``optimize_brief.md`` into a Markdown prompt for a subagent.
+
+    ``generate()`` is **synchronous** — briefs are pure Jinja2 template
+    renders, not LLM calls.  The subagent that *receives* the brief is the
+    LLM.
+    """
+
+    def __init__(self, client: LLMClient) -> None:
+        # Client stored for future extensions; not used in synchronous rendering.
+        self._client = client
+
+    def generate(
+        self,
+        optimization_target: str,
+        metric: str,
+        metric_direction: str,
+        baseline_score: float,
+        best_score: float,
+        benchmark_cmd: str,
+        gate_cmd: str,
+        past_experiments: list[ExperimentSummary],
+        shared_notes: str = "",
+        agent_index: int = 0,
+    ) -> str:
+        """Render the optimisation brief for one parallel agent.
+
+        Args:
+            optimization_target: One-sentence description of the target.
+            metric:              Human-readable metric string.
+            metric_direction:    ``"higher"`` or ``"lower"``.
+            baseline_score:      Score at the start of the run.
+            best_score:          Best score seen so far.
+            benchmark_cmd:       Shell command to run the benchmark.
+            gate_cmd:            Shell command to run the gate tests.
+            past_experiments:    History of previous attempts.
+            shared_notes:        Cross-agent knowledge to include.
+            agent_index:         Controls experiment ordering for variety.
+                                 Index 0 → oldest-first;
+                                 Index > 0 → reversed (newest-first) for diversity.
+
+        Returns the rendered Markdown string; never calls the LLM.
+        """
+        template = self._load_template()
+
+        # Vary ordering so parallel agents explore different strategies.
+        experiments = list(past_experiments)
+        if agent_index > 0 and experiments:
+            experiments = list(reversed(experiments))
+
+        return template.render(
+            optimization_target=optimization_target,
+            metric=metric,
+            direction=metric_direction,
+            baseline_score=baseline_score,
+            best_score=best_score,
+            benchmark_cmd=benchmark_cmd,
+            gate_cmd=gate_cmd,
+            failed_experiments=experiments,
+            shared_notes=shared_notes,
+        )
+
+    def _load_template(self) -> jinja2.Template:
+        """Load ``optimize_brief.md`` as a Jinja2 :class:`~jinja2.Template`."""
+        env = jinja2.Environment(
+            loader=jinja2.FileSystemLoader(str(_PROMPTS_DIR)),
+            autoescape=False,
+            trim_blocks=True,
+            lstrip_blocks=True,
+        )
+        return env.get_template("optimize_brief.md")

gradex/ai/client.py

@@ -0,0 +1,232 @@
+"""Unified LLM client supporting Anthropic, OpenAI, and Ollama."""
+
+from __future__ import annotations
+
+import asyncio
+import time
+from dataclasses import dataclass
+from threading import Lock
+
+from gradex.config import LLMConfig, load_llm_config
+
+_rate_limiters: dict[str, "_TokenBucket"] = {}
+_rl_lock = Lock()
+
+
+class _TokenBucket:
+    """Token bucket limiter for max requests per 60 seconds."""
+
+    def __init__(self, max_requests: int = 50) -> None:
+        self._max = max_requests
+        self._tokens = float(max_requests)
+        self._last_refill = time.monotonic()
+        self._lock = Lock()
+
+    def _refill(self) -> None:
+        now = time.monotonic()
+        elapsed = now - self._last_refill
+        refill = elapsed * (self._max / 60.0)
+        self._tokens = min(self._max, self._tokens + refill)
+        self._last_refill = now
+
+    def consume(self) -> float:
+        """Consume one token; return wait seconds if throttled."""
+        with self._lock:
+            self._refill()
+            if self._tokens >= 1.0:
+                self._tokens -= 1.0
+                return 0.0
+            return (1.0 - self._tokens) / (self._max / 60.0)
+
+
+def _get_bucket(provider: str) -> "_TokenBucket":
+    with _rl_lock:
+        if provider not in _rate_limiters:
+            _rate_limiters[provider] = _TokenBucket(max_requests=50)
+        return _rate_limiters[provider]
+
+
+@dataclass
+class LLMResponse:
+    """The result of a single LLM completion call."""
+
+    text: str
+    input_tokens: int
+    output_tokens: int
+    provider: str
+    model: str
+
+
+class LLMClient:
+    """Unified LLM client supporting Anthropic, OpenAI, and Ollama.
+
+    All three backends implement the same interface: a system prompt plus a
+    user prompt produce a text response.  Provider SDKs are imported lazily
+    so the package can be installed without requiring all of them.
+    """
+
+    def __init__(self, config: LLMConfig | None = None) -> None:
+        self._config: LLMConfig = config if config is not None else load_llm_config()
+
+    async def complete(
+        self,
+        system: str,
+        user: str,
+        max_tokens: int | None = None,
+    ) -> LLMResponse:
+        """Send a system+user prompt and return the assistant response.
+
+        Args:
+            system:     System prompt text.
+            user:       User message text.
+            max_tokens: Token budget; defaults to ``config.max_tokens``.
+
+        Raises:
+            ValueError: When the configured provider is not recognised.
+        """
+        cfg = self._config
+        tokens = max_tokens if max_tokens is not None else cfg.max_tokens
+        provider = cfg.provider
+        bucket = _get_bucket(provider)
+
+        for attempt in range(3):
+            wait = bucket.consume()
+            if wait == 0.0:
+                break
+            backoff = wait * (2**attempt)
+            await asyncio.sleep(backoff)
+
+        if provider == "anthropic":
+            return await self._complete_anthropic(system, user, tokens)
+        elif provider == "openai":
+            return await self._complete_openai(system, user, tokens)
+        elif provider == "ollama":
+            return await self._complete_ollama(system, user, tokens)
+        elif provider == "groq":
+            return await self._complete_groq(system, user, tokens)
+        else:
+            raise ValueError(
+                f"Unknown provider: {provider!r}. "
+                f"Choose: anthropic, openai, ollama, groq"
+            )
+
+    # ------------------------------------------------------------------
+    # Provider implementations
+    # ------------------------------------------------------------------
+
+    async def _complete_anthropic(
+        self, system: str, user: str, max_tokens: int
+    ) -> LLMResponse:
+        """Call the Anthropic Messages API."""
+        import anthropic as ant
+
+        client = ant.AsyncAnthropic(api_key=self._config.api_key or None)
+        msg = await client.messages.create(
+            model=self._config.effective_model(),
+            max_tokens=max_tokens,
+            system=system,
+            messages=[{"role": "user", "content": user}],
+        )
+        return LLMResponse(
+            text=msg.content[0].text,  # type: ignore[union-attr]  # always TextBlock in practice
+            input_tokens=msg.usage.input_tokens,
+            output_tokens=msg.usage.output_tokens,
+            provider="anthropic",
+            model=self._config.effective_model(),
+        )
+
+    async def _complete_openai(
+        self, system: str, user: str, max_tokens: int
+    ) -> LLMResponse:
+        """Call the OpenAI Chat Completions API."""
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(api_key=self._config.api_key or None)
+        resp = await client.chat.completions.create(
+            model=self._config.effective_model(),
+            max_tokens=max_tokens,
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+        )
+        choice = resp.choices[0]
+        usage = resp.usage
+        return LLMResponse(
+            text=choice.message.content or "",
+            input_tokens=usage.prompt_tokens if usage else 0,
+            output_tokens=usage.completion_tokens if usage else 0,
+            provider="openai",
+            model=self._config.effective_model(),
+        )
+
+    async def _complete_groq(
+        self, system: str, user: str, max_tokens: int
+    ) -> LLMResponse:
+        """Call Groq's OpenAI-compatible endpoint via the OpenAI SDK.
+
+        Groq free tier: 14,400 requests/day.
+        Get a key at https://console.groq.com.
+        Best free model: ``llama-3.3-70b-versatile``.
+        """
+        from openai import AsyncOpenAI
+
+        client = AsyncOpenAI(
+            api_key=self._config.api_key or None,
+            base_url=self._config.groq_base_url,
+        )
+        resp = await client.chat.completions.create(
+            model=self._config.effective_model(),
+            max_tokens=max_tokens,
+            messages=[
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+            temperature=self._config.temperature,
+        )
+        choice = resp.choices[0]
+        usage = resp.usage
+        return LLMResponse(
+            text=choice.message.content or "",
+            input_tokens=usage.prompt_tokens if usage else 0,
+            output_tokens=usage.completion_tokens if usage else 0,
+            provider="groq",
+            model=self._config.effective_model(),
+        )
+
+    async def _complete_ollama(
+        self, system: str, user: str, max_tokens: int
+    ) -> LLMResponse:
+        """Call Ollama's OpenAI-compatible endpoint via httpx.
+
+        Ollama exposes ``http://localhost:11434/v1`` — no extra SDK required.
+        """
+
+        import httpx
+
+        payload = {
+            "model": self._config.effective_model(),
+            "messages": [
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ],
+            "stream": False,
+            "options": {"num_predict": max_tokens},
+        }
+        async with httpx.AsyncClient(timeout=120.0) as http:
+            resp = await http.post(
+                f"{self._config.ollama_base_url}/chat/completions",
+                json=payload,
+            )
+        resp.raise_for_status()
+        data: dict[str, object] = resp.json()
+        choices = data["choices"]
+        text: str = choices[0]["message"]["content"]  # type: ignore[index]
+        usage: dict[str, int] = data.get("usage", {})  # type: ignore[assignment]  # resp.json() is untyped
+        return LLMResponse(
+            text=text,
+            input_tokens=usage.get("prompt_tokens", 0),
+            output_tokens=usage.get("completion_tokens", 0),
+            provider="ollama",
+            model=self._config.effective_model(),
+        )

gradex/ai/discover.py

@@ -0,0 +1,280 @@
+"""Discover skill: analyse a repo and set up a baseline optimization run."""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Literal
+
+import jinja2
+
+from gradex.ai.client import LLMClient
+from gradex.backends.base import Backend
+
+PROMPTS_DIR = Path(__file__).parent / "prompts"
+
+# ---------------------------------------------------------------------------
+# Result type
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class DiscoverResult:
+    """Everything the discover flow found and set up."""
+
+    optimization_target: str
+    metric: str
+    metric_direction: Literal["higher", "lower"]
+    benchmark_script: str
+    benchmark_path: Path
+    gate_cmds: list[str]
+    baseline_score: float
+    run_id: str
+    notes: str
+
+
+# ---------------------------------------------------------------------------
+# Skill
+# ---------------------------------------------------------------------------
+
+
+class DiscoverSkill:
+    """Analyse a repository, design a benchmark, and establish a baseline score.
+
+    Steps performed by :meth:`run`:
+
+    1. Scan the repo to build a context string.
+    2. Ask the LLM to pick an optimisation target and metric.
+    3. Ask the LLM to write a Python benchmark script; save it.
+    4. Detect existing test files.
+    5. Ask the LLM to identify gate commands.
+    6. Run the benchmark once to capture the baseline score.
+    7. Create a :class:`~evo.state.Run` in the database.
+    """
+
+    def __init__(self, client: LLMClient, backend: Backend) -> None:
+        self._client = client
+        self._backend = backend
+
+    # ------------------------------------------------------------------
+    # Main entry point
+    # ------------------------------------------------------------------
+
+    async def run(
+        self,
+        repo_root: Path,
+        hint: str = "",
+    ) -> DiscoverResult:
+        """Run the full discover flow for *repo_root*.
+
+        Args:
+            repo_root: Root directory of the repository to analyse.
+            hint:      Optional user hint, e.g. "make the parser faster".
+
+        Returns a :class:`DiscoverResult` with all artefacts populated.
+        """
+        # 1 — Repo context
+        repo_context = self.scan_repo(repo_root)
+        test_files = self.detect_test_files(repo_root)
+
+        # 2 — Target + metric
+        analysis_system = (PROMPTS_DIR / "repo_analysis.md").read_text(encoding="utf-8")
+        analysis_user = repo_context
+        if hint:
+            analysis_user += f"\n\nUser hint: {hint}"
+        resp1 = await self._client.complete(analysis_system, analysis_user)
+        optimization_target = self._parse_xml_tag(resp1.text, "optimization_target")
+        metric = self._parse_xml_tag(resp1.text, "metric")
+        direction = self._infer_direction(metric)
+
+        # 3 — Benchmark script
+        bench_system = jinja2.Template(
+            (PROMPTS_DIR / "benchmark_design.md").read_text(encoding="utf-8")
+        ).render(
+            target=optimization_target,
+            metric=metric,
+            repo_context=repo_context,
+        )
+        resp2 = await self._client.complete(
+            bench_system, "Write the benchmark script now."
+        )
+        benchmark_script = self._parse_xml_tag(resp2.text, "benchmark_script")
+        notes = ""
+        try:
+            notes = self._parse_xml_tag(resp2.text, "notes")
+        except ValueError:
+            pass
+
+        # Write benchmark to disk
+        evo_dir = repo_root / ".gradex"
+        evo_dir.mkdir(parents=True, exist_ok=True)
+        benchmark_path = evo_dir / "benchmark.py"
+        benchmark_path.write_text(benchmark_script, encoding="utf-8")
+
+        # 4 + 5 — Gate commands
+        gate_system = jinja2.Template(
+            (PROMPTS_DIR / "gate_design.md").read_text(encoding="utf-8")
+        ).render(
+            target=optimization_target,
+            test_files=test_files,
+        )
+        resp3 = await self._client.complete(gate_system, "Identify the gate commands.")
+        gate_cmds_raw = self._parse_xml_tag(resp3.text, "gate_cmds")
+        gate_cmds: list[str] = json.loads(gate_cmds_raw)
+
+        # 6 — Baseline
+        baseline_score = await self._run_baseline(repo_root, benchmark_path)
+
+        # 7 — Persist Run
+        from gradex.repository import RunRepository
+
+        run = RunRepository().create(
+            benchmark_cmd=f"python {benchmark_path}",
+            metric_direction=direction,
+            gate_cmds=gate_cmds,
+            baseline_score=baseline_score,
+        )
+
+        return DiscoverResult(
+            optimization_target=optimization_target,
+            metric=metric,
+            metric_direction=direction,
+            benchmark_script=benchmark_script,
+            benchmark_path=benchmark_path,
+            gate_cmds=gate_cmds,
+            baseline_score=baseline_score,
+            run_id=run.id,
+            notes=notes,
+        )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def scan_repo(self, repo_root: Path) -> str:
+        """Return a compact text summary of *repo_root* for LLM context.
+
+        Skips ``.git``, ``.evo``, ``__pycache__``, ``node_modules``,
+        ``.venv``, and ``dist``.  Output is capped at 3 000 characters.
+        """
+        _SKIP = {".git", ".gradex", "__pycache__", "node_modules", ".venv", "dist"}
+        entries: list[str] = []
+        ext_counts: dict[str, int] = {}
+        total_size = 0
+
+        def _walk(path: Path, depth: int = 0) -> None:
+            if depth > 3 or len(entries) >= 60:
+                return
+            try:
+                children = sorted(path.iterdir())
+            except PermissionError:
+                return
+            for item in children:
+                if item.name in _SKIP:
+                    continue
+                if len(entries) >= 60:
+                    entries.append("  ... (truncated)")
+                    return
+                indent = "  " * depth
+                if item.is_dir():
+                    entries.append(f"{indent}{item.name}/")
+                    _walk(item, depth + 1)
+                else:
+                    entries.append(f"{indent}{item.name}")
+                    ext = item.suffix.lower() or "(no ext)"
+                    ext_counts[ext] = ext_counts.get(ext, 0) + 1
+                    try:
+                        nonlocal total_size
+                        total_size += item.stat().st_size
+                    except OSError:
+                        pass
+
+        _walk(repo_root)
+        tree_str = "\n".join(entries)
+        ext_str = ", ".join(f"{k}: {v}" for k, v in sorted(ext_counts.items()))
+        result = (
+            f"Directory tree:\n{tree_str}\n\n"
+            f"File types: {ext_str}\n"
+            f"Total size: {total_size} bytes"
+        )
+        return result[:3000]
+
+    def detect_test_files(self, repo_root: Path) -> list[str]:
+        """Return relative paths to test files in *repo_root* (max 20).
+
+        Matches ``test_*.py`` and ``*_test.py`` patterns anywhere under
+        *repo_root*.
+        """
+        seen: set[str] = set()
+        results: list[str] = []
+        for pattern in ("**/test_*.py", "**/*_test.py"):
+            for path in sorted(repo_root.glob(pattern)):
+                rel = str(path.relative_to(repo_root))
+                if rel not in seen:
+                    seen.add(rel)
+                    results.append(rel)
+        return results[:20]
+
+    def _parse_xml_tag(self, text: str, tag: str) -> str:
+        """Extract the content of ``<tag>…</tag>`` from *text*.
+
+        Strips surrounding whitespace.
+
+        Raises:
+            ValueError: If the tag is absent.
+        """
+        match = re.search(f"<{tag}>(.*?)</{tag}>", text, re.DOTALL)
+        if not match:
+            raise ValueError(f"Tag <{tag}> not found in LLM response")
+        return match.group(1).strip()
+
+    def _infer_direction(self, metric: str) -> Literal["higher", "lower"]:
+        """Infer optimisation direction from a metric description string.
+
+        ``"lower"`` keywords: lower, minimize, latency, cost, ms, error, loss.
+        ``"higher"`` keywords: higher, maximize, accuracy, score, %, throughput.
+        Defaults to ``"lower"``.
+        """
+        lower_m = metric.lower()
+        higher_keywords = {"higher", "maximize", "accuracy", "score", "%", "throughput"}
+        lower_keywords = {
+            "lower",
+            "minimize",
+            "latency",
+            "cost",
+            "ms",
+            "millisecond",
+            "error",
+            "loss",
+        }
+        for kw in higher_keywords:
+            if kw in lower_m:
+                return "higher"
+        for kw in lower_keywords:
+            if kw in lower_m:
+                return "lower"
+        return "lower"
+
+    # ------------------------------------------------------------------
+    # Baseline execution
+    # ------------------------------------------------------------------
+
+    async def _run_baseline(self, repo_root: Path, benchmark_path: Path) -> float:
+        """Execute *benchmark_path* once and parse the resulting score.
+
+        Raises:
+            ValueError: If the benchmark times out or yields no parseable score.
+        """
+        from gradex.runner.benchmark import BenchmarkRunner
+
+        runner = BenchmarkRunner(self._backend, timeout=60)
+        result = await runner.run(repo_root, ["python", str(benchmark_path)])
+        if result.timed_out:
+            raise ValueError("Baseline benchmark timed out")
+        if result.score is None:
+            raise ValueError(
+                f"Baseline benchmark returned no parseable score: {result.parse_error!r}"
+            )
+        return result.score

gradex/ai/prompts/benchmark_design.md

@@ -0,0 +1,21 @@
+You are an expert software engineer writing a Python benchmark script.
+
+Target to optimize: {{ target }}
+Metric: {{ metric }}
+Repository context: {{ repo_context }}
+
+Write a Python script that:
+1. Imports and runs the target code
+2. Measures the metric numerically (time it, count it, score it)
+3. Prints ONLY a single float on the last line of stdout (e.g. `print(f"{value:.4f}")`)
+4. Is deterministic enough to compare across runs (warm up if needed, average 3 runs)
+5. Has no hardcoded paths — uses relative imports or Path(__file__).parent
+6. Completes in under 30 seconds
+
+Respond in this exact format:
+<benchmark_script>
+[complete Python script, no markdown fences]
+</benchmark_script>
+<notes>
+[any setup steps needed before running, e.g. "requires sample_data/ directory"]
+</notes>

gradex/ai/prompts/gate_design.md

@@ -0,0 +1,16 @@
+You are an expert software engineer identifying regression tests for a code optimization task.
+
+Target to optimize: {{ target }}
+Repository test files found: {{ test_files }}
+
+Identify the pytest command(s) that must pass after every experiment.
+Prefer: tests directly covering the optimized module.
+If no specific tests exist, use the full test suite.
+
+Respond in this exact format:
+<gate_cmds>
+["pytest tests/test_parser.py", "pytest tests/test_integration.py"]
+</gate_cmds>
+<rationale>
+[why these tests are the right gate]
+</rationale>

gradex/ai/prompts/optimize_brief.md

@@ -0,0 +1,33 @@
+You are a subagent in a parallel code optimization loop.
+
+## Your task
+{{ optimization_target }}
+Metric: {{ metric }} ({{ direction }})
+Current best score: {{ best_score }}
+Baseline score: {{ baseline_score }}
+
+## What has been tried (do NOT repeat these)
+{% for exp in failed_experiments %}
+- Hypothesis: {{ exp.hypothesis }}
+  Result: {{ exp.result }} ({{ exp.reason }})
+{% endfor %}
+
+## Shared knowledge from all agents
+{{ shared_notes }}
+
+## Your hypothesis
+Based on the above, form ONE new hypothesis that has NOT been tried.
+Be specific: name the function, file, and change you will make.
+
+## Instructions
+1. Read the relevant source files
+2. Implement your hypothesis
+3. Run the benchmark: `{{ benchmark_cmd }}`
+   The last line of stdout is your score. Lower/higher is better per the metric above.
+4. Run the gate: `{{ gate_cmd }}`
+5. Write your results to `.gradex/result.json`:
+   {"score": <float>, "hypothesis": "<one sentence>", "change_summary": "<what you changed>"}
+6. Write gate result to `.gradex/gate.json`:
+   {"passed": <bool>, "failures": [<strings>]}
+
+Do not stop until both files are written.