autoforge-trainer 1.0.2__py3-none-any.whl

autoforge/__init__.py

@@ -0,0 +1,3 @@
+"""AutoForge: AI that trains AI."""
+
+__version__ = "1.0.2"

autoforge/agents/loop.py

@@ -0,0 +1,168 @@
+"""Built-in agent loop (P4b): an LLM drives AutoForge through the same tools
+the MCP server exposes. The tools run in-process via an MCP client, so the
+privacy guard, audit log and error envelopes are identical to P4a.
+
+Human approval gate: by default the person confirms before training starts
+(`run_start`, `run_retry`); --yes turns it off. A declined call returns
+USER_DECLINED with the person's feedback so the model can revise the plan.
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Awaitable, Callable
+
+from mcp import Client
+
+from autoforge import config
+from autoforge.agents.providers import Provider, ToolCall, ToolResult, ToolSpec
+from autoforge.schemas.core import BrainKind
+from autoforge.serve.mcp_server import INSTRUCTIONS, build_server
+
+GATED_TOOLS = {"run_start", "run_retry", "hpo_start"}
+MAX_RESULT_CHARS = 20_000
+
+SYSTEM_PROMPT = """\
+You are AutoForge's training agent. You analyse what the user needs, choose a training method that fits
+their data and hardware, train the model on their machine, and report the result.
+
+{instructions}
+
+How to work:
+- Start with hardware_profile and data_profile. Base every choice on what the tools return.
+- Search packs for the task and read pack_info's knowledge before choosing a recipe; follow its guidance
+  (e.g. start with the simple baseline recipe).
+- Validate the plan, start the run, then poll with run_wait until it finishes. If it fails, read run_logs,
+  fix the plan and try again. Compare against the baseline before trying anything heavier.
+- If a tool returns USER_DECLINED, read the user's feedback in the message and revise.
+- If a tool returns PRIVACY_DENIED, stop and explain; never look for another way to read that data.
+- After a working baseline, use run_diagnose for problems and hpo_start to tune (trial 0 is the baseline);
+  report whether tuning improved the metric.
+- When the user wants to use the model, register the best run with model_register and show a
+  model_predict example.
+- Finish with a short report: the method and why, the metrics, the run_id and the model path.
+"""
+
+
+@dataclass
+class AgentEvent:
+    kind: str  # text | tool_call | tool_result | approval | info | done
+    data: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class AgentResult:
+    final_text: str
+    steps: int
+    stop: str  # end | refusal | max_steps | error
+    runs: list[str]
+    transcript: str
+
+
+ApproveFn = Callable[[ToolCall], Awaitable[tuple[bool, str]]]
+EventFn = Callable[[AgentEvent], None]
+
+
+def _truncate(text: str) -> str:
+    if len(text) <= MAX_RESULT_CHARS:
+        return text
+    return text[:MAX_RESULT_CHARS] + f"\n... [truncated {len(text) - MAX_RESULT_CHARS} chars]"
+
+
+class Transcript:
+    """JSONL record of the whole agent session, next to the runs it started."""
+
+    def __init__(self) -> None:
+        d = config.home() / "agent-sessions"
+        d.mkdir(parents=True, exist_ok=True)
+        self.path = d / f"{time.strftime('%Y%m%d-%H%M%S')}.jsonl"
+
+    def write(self, kind: str, **data: Any) -> None:
+        with self.path.open("a", encoding="utf-8") as f:
+            f.write(json.dumps({"t": time.time(), "kind": kind, **data}, ensure_ascii=False, default=str) + "\n")
+
+
+async def run_agent(
+    goal: str,
+    provider: Provider,
+    actor: str,
+    brain: BrainKind,
+    approve: ApproveFn | None = None,
+    on_event: EventFn | None = None,
+    max_steps: int = 40,
+    data_path: str | None = None,
+) -> AgentResult:
+    emit = on_event or (lambda e: None)
+    transcript = Transcript()
+    runs: list[str] = []
+    server = build_server(actor, log_level="WARNING")  # the agent reports tool errors itself
+
+    user_msg = goal
+    if data_path:
+        user_msg += f"\n\nData: {Path(data_path).expanduser().resolve()}"
+    transcript.write("start", goal=goal, actor=actor, brain=brain.value, data_path=data_path)
+
+    async with Client(server) as mcp:
+        listed = await mcp.list_tools()
+        tools = [ToolSpec(t.name, t.description or "", t.input_schema) for t in listed.tools]
+        provider.start(SYSTEM_PROMPT.format(instructions=INSTRUCTIONS), tools)
+        provider.add_user(user_msg)
+
+        final_text, stop = "", "max_steps"
+        for step in range(1, max_steps + 1):
+            turn = await provider.complete()
+            if turn.text:
+                final_text = turn.text
+                emit(AgentEvent("text", {"text": turn.text}))
+                transcript.write("text", text=turn.text)
+            if turn.stop == "refusal":
+                stop = "refusal"
+                emit(AgentEvent("info", {"message": f"model declined: {turn.detail or ''}"}))
+                break
+            if not turn.tool_calls:
+                if turn.stop == "pause":
+                    continue  # server paused a long turn; resend to let it continue
+                stop = "end"
+                break
+
+            results = []
+            for call in turn.tool_calls:
+                emit(AgentEvent("tool_call", {"name": call.name, "arguments": call.arguments}))
+                transcript.write("tool_call", name=call.name, arguments=call.arguments)
+                if call.parse_error:
+                    content = json.dumps({"ok": False, "error": {"code": "BAD_TOOL_CALL", "message": call.parse_error}})
+                    results.append(ToolResult(call.id, content, is_error=True))
+                    continue
+                if call.name in GATED_TOOLS and approve is not None:
+                    ok, feedback = await approve(call)
+                    emit(AgentEvent("approval", {"name": call.name, "approved": ok, "feedback": feedback}))
+                    transcript.write("approval", name=call.name, approved=ok, feedback=feedback)
+                    if not ok:
+                        msg = "The user declined this action." + (f" Feedback: {feedback}" if feedback else "")
+                        content = json.dumps({"ok": False, "error": {"code": "USER_DECLINED", "message": msg}})
+                        results.append(ToolResult(call.id, content, is_error=True))
+                        continue
+                res = await mcp.call_tool(call.name, call.arguments)
+                envelope = res.structured_content
+                if envelope is None:  # unstructured (e.g. argument validation error from the MCP layer)
+                    text = "\n".join(getattr(c, "text", "") for c in res.content)
+                    envelope = {"ok": False, "error": {"code": "BAD_TOOL_CALL", "message": text}}
+                if call.name in ("run_start", "run_retry") and envelope.get("ok"):
+                    runs.append(envelope["data"]["run_id"])
+                if call.name == "hpo_start" and envelope.get("ok"):
+                    runs.append(envelope["data"]["study_id"])
+                content = _truncate(json.dumps(envelope, ensure_ascii=False, default=str))
+                emit(AgentEvent("tool_result", {"name": call.name, "ok": envelope.get("ok"),
+                                                "error": (envelope.get("error") or {}).get("code")}))
+                transcript.write("tool_result", name=call.name, ok=envelope.get("ok"), content=content[:2000])
+                results.append(ToolResult(call.id, content, is_error=not envelope.get("ok", False)))
+            provider.add_tool_results(results)
+        else:
+            step = max_steps
+
+    transcript.write("done", stop=stop, steps=step, runs=runs)
+    emit(AgentEvent("done", {"stop": stop, "steps": step, "runs": runs}))
+    return AgentResult(final_text=final_text, steps=step, stop=stop, runs=runs, transcript=str(transcript.path))

autoforge/agents/providers.py

@@ -0,0 +1,272 @@
+"""LLM providers for the built-in agent (P4b).
+
+Each provider owns its own conversation state, so the agent loop stays
+provider-neutral:
+    p.start(system, tools); p.add_user(text)
+    turn = await p.complete()          # text + tool calls
+    p.add_tool_results([...])          # one result per call, all at once
+
+Brain specs:
+    anthropic:<model>   Claude via the official SDK        (cloud)
+    openai:<model>      OpenAI or any cloud OpenAI-compatible API (--base-url)
+    ollama:<model>      local Ollama (OpenAI-compatible endpoint)
+    local:<model>       any local OpenAI-compatible server (vLLM, LM Studio, llama.cpp) with --base-url
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any
+
+from autoforge.errors import AutoForgeError
+from autoforge.schemas.core import BrainKind
+
+DEFAULT_BRAIN = "anthropic:claude-opus-5-5"
+# Models that accept server-side refusal fallbacks (fallbacks: "default").
+CLAUDE_FALLBACK_MODELS = {"claude-fable-5-1", "claude-opus-5-5", "claude-opus-5", "claude-sonnet-5-5"}
+
+
+class ProviderError(AutoForgeError):
+    code = "PROVIDER_ERROR"
+
+
+@dataclass
+class ToolSpec:
+    name: str
+    description: str
+    input_schema: dict[str, Any]
+
+
+@dataclass
+class ToolCall:
+    id: str
+    name: str
+    arguments: dict[str, Any]
+    parse_error: str | None = None
+
+
+@dataclass
+class ToolResult:
+    call_id: str
+    content: str
+    is_error: bool = False
+
+
+@dataclass
+class Turn:
+    text: str = ""
+    tool_calls: list[ToolCall] = field(default_factory=list)
+    stop: str = "end"  # end | tool_use | pause | max_tokens | refusal
+    detail: str | None = None
+
+
+@dataclass
+class BrainSpec:
+    provider: str
+    model: str
+    base_url: str | None = None
+    api_key_env: str | None = None
+
+    @property
+    def is_local(self) -> bool:
+        return self.provider in ("ollama", "local")
+
+    @property
+    def brain_kind(self) -> BrainKind:
+        return BrainKind.local if self.is_local else BrainKind.api
+
+    @property
+    def actor(self) -> str:
+        # 'ai:local:...' is what the privacy guard recognises as a local model.
+        return f"ai:local:{self.provider}:{self.model}" if self.is_local else f"ai:{self.provider}:{self.model}"
+
+
+def parse_brain(spec: str | None, base_url: str | None = None, api_key_env: str | None = None) -> BrainSpec:
+    spec = spec or os.environ.get("AUTOFORGE_BRAIN") or DEFAULT_BRAIN
+    provider, sep, model = spec.partition(":")
+    if not sep or not model:
+        raise ProviderError(f"bad brain '{spec}'", hint="Use provider:model, e.g. anthropic:claude-opus-5-5 or ollama:qwen3")
+    if provider not in ("anthropic", "openai", "ollama", "local"):
+        raise ProviderError(f"unknown provider '{provider}'", hint="Choose anthropic, openai, ollama or local.")
+    if provider == "local" and not base_url:
+        raise ProviderError("local:<model> needs --base-url (an OpenAI-compatible endpoint, e.g. http://localhost:8000/v1)")
+    return BrainSpec(provider, model, base_url, api_key_env)
+
+
+def make_provider(spec: BrainSpec) -> "Provider":
+    if spec.provider == "anthropic":
+        return AnthropicProvider(spec.model)
+    if spec.provider == "ollama":
+        host = os.environ.get("OLLAMA_HOST", "http://localhost:11434").rstrip("/")
+        if not host.startswith("http"):
+            host = "http://" + host
+        return OpenAICompatProvider(spec.model, spec.base_url or f"{host}/v1", api_key=None)
+    if spec.provider == "local":
+        return OpenAICompatProvider(spec.model, spec.base_url, api_key=None)
+    key_env = spec.api_key_env or "OPENAI_API_KEY"
+    key = os.environ.get(key_env)
+    if not key:
+        raise ProviderError(f"{key_env} is not set", hint=f"export {key_env}=... (or pass --api-key-env)")
+    return OpenAICompatProvider(spec.model, spec.base_url or "https://api.openai.com/v1", api_key=key)
+
+
+class Provider:
+    def start(self, system: str, tools: list[ToolSpec]) -> None: ...
+    def add_user(self, text: str) -> None: ...
+    async def complete(self) -> Turn: ...
+    def add_tool_results(self, results: list[ToolResult]) -> None: ...
+
+
+# ---------------------------------------------------------------- Claude
+
+
+class AnthropicProvider(Provider):
+    """Claude via the official Anthropic SDK (manual tool-use loop).
+
+    - Effort 'high' (agentic work); adaptive thinking is the model default.
+    - Server-side refusal fallback ('default') on models that support it.
+    - The assistant content is appended back unchanged (thinking blocks included),
+      so the history stays append-only.
+    """
+
+    def __init__(self, model: str, client: Any = None, effort: str = "high", max_tokens: int = 16000):
+        if client is None:
+            try:
+                import anthropic
+            except ImportError:
+                raise ProviderError("the anthropic package is not installed", hint='pip install "autoforge[anthropic]"')
+            client = anthropic.AsyncAnthropic()
+        self.client = client
+        self.model = model
+        self.effort = effort
+        self.max_tokens = max_tokens
+        self.messages: list[dict[str, Any]] = []
+
+    def start(self, system: str, tools: list[ToolSpec]) -> None:
+        self.system = system
+        self.tools = [{"name": t.name, "description": t.description, "input_schema": t.input_schema} for t in tools]
+
+    def add_user(self, text: str) -> None:
+        self.messages.append({"role": "user", "content": text})
+
+    async def complete(self) -> Turn:
+        kwargs: dict[str, Any] = dict(
+            model=self.model,
+            max_tokens=self.max_tokens,
+            system=self.system,
+            tools=self.tools,
+            messages=self.messages,
+            output_config={"effort": self.effort},
+            cache_control={"type": "ephemeral"},  # tools + system + history are a stable prefix
+        )
+        if self.model in CLAUDE_FALLBACK_MODELS:
+            kwargs |= {"betas": ["server-side-fallback-2026-07-01"], "fallbacks": "default"}
+        try:
+            response = await self.client.beta.messages.create(**kwargs)
+        except Exception as e:  # SDK errors: auth, rate limit, network, bad request
+            raise ProviderError(f"Claude API call failed: {type(e).__name__}: {e}") from e
+
+        self.messages.append({"role": "assistant", "content": response.content})
+        text = "\n".join(b.text for b in response.content if b.type == "text" and b.text)
+        calls = [ToolCall(b.id, b.name, dict(b.input or {})) for b in response.content if b.type == "tool_use"]
+        stop = {"end_turn": "end", "tool_use": "tool_use", "pause_turn": "pause",
+                "max_tokens": "max_tokens", "refusal": "refusal"}.get(response.stop_reason, "end")
+        detail = None
+        if stop == "refusal" and getattr(response, "stop_details", None):
+            detail = f"{response.stop_details.category}: {response.stop_details.explanation}"
+        if stop == "max_tokens":
+            # A tool_use block cut off by max_tokens may have incomplete input: don't run it,
+            # but every tool_use still needs a tool_result, so report it as an error.
+            for c in calls:
+                c.parse_error = "response hit max_tokens; tool input may be truncated, call the tool again"
+        return Turn(text=text, tool_calls=calls, stop=stop, detail=detail)
+
+    def add_tool_results(self, results: list[ToolResult]) -> None:
+        self.messages.append({
+            "role": "user",
+            "content": [
+                {"type": "tool_result", "tool_use_id": r.call_id, "content": r.content, "is_error": r.is_error}
+                for r in results
+            ],
+        })
+
+
+# ---------------------------------------------------------------- OpenAI-compatible (Ollama, OpenAI, vLLM ...)
+
+
+def simplify_schema(schema: Any) -> Any:
+    """Make JSON Schemas easier for small local models: `X | None` (anyOf with null) becomes X,
+    and pydantic's 'title' noise is dropped. Defaults and descriptions are kept."""
+    if isinstance(schema, list):
+        return [simplify_schema(s) for s in schema]
+    if not isinstance(schema, dict):
+        return schema
+    out = {k: simplify_schema(v) for k, v in schema.items() if k != "title"}
+    any_of = out.get("anyOf")
+    if isinstance(any_of, list):
+        non_null = [s for s in any_of if s != {"type": "null"}]
+        if len(non_null) == 1 and len(non_null) < len(any_of):
+            merged = {k: v for k, v in out.items() if k != "anyOf"}
+            return non_null[0] | merged
+    return out
+
+
+class OpenAICompatProvider(Provider):
+    """Chat Completions with function calling over plain HTTP (httpx2)."""
+
+    def __init__(self, model: str, base_url: str, api_key: str | None, transport: Any = None, timeout: float = 600):
+        import httpx2
+
+        headers = {"Authorization": f"Bearer {api_key}"} if api_key else {}
+        self.client = httpx2.AsyncClient(base_url=base_url.rstrip("/") + "/", headers=headers,
+                                         timeout=timeout, transport=transport)
+        self.model = model
+        self.base_url = base_url
+        self.messages: list[dict[str, Any]] = []
+
+    def start(self, system: str, tools: list[ToolSpec]) -> None:
+        self.messages = [{"role": "system", "content": system}]
+        self.tools = [
+            {"type": "function", "function": {"name": t.name, "description": t.description,
+                                              "parameters": simplify_schema(t.input_schema)}}
+            for t in tools
+        ]
+
+    def add_user(self, text: str) -> None:
+        self.messages.append({"role": "user", "content": text})
+
+    async def complete(self) -> Turn:
+        import httpx2
+
+        body = {"model": self.model, "messages": self.messages, "tools": self.tools}
+        try:
+            r = await self.client.post("chat/completions", json=body)
+            r.raise_for_status()
+            data = r.json()
+        except httpx2.HTTPStatusError as e:
+            raise ProviderError(f"{self.base_url} returned {e.response.status_code}: {e.response.text[:500]}") from e
+        except httpx2.RequestError as e:
+            raise ProviderError(f"cannot reach {self.base_url}: {e}",
+                                hint="Is the model server running? For Ollama: `ollama serve` and `ollama pull <model>`.") from e
+        choice = data["choices"][0]
+        msg = choice["message"]
+        self.messages.append({k: v for k, v in msg.items() if k in ("role", "content", "tool_calls") and v is not None})
+        calls = []
+        for i, tc in enumerate(msg.get("tool_calls") or []):
+            fn = tc.get("function", {})
+            raw = fn.get("arguments") or "{}"
+            try:
+                args = raw if isinstance(raw, dict) else json.loads(raw)
+                err = None
+            except json.JSONDecodeError as e:
+                args, err = {}, f"arguments are not valid JSON: {e}"
+            calls.append(ToolCall(tc.get("id") or f"call_{i}", fn.get("name", ""), args, err))
+        finish = choice.get("finish_reason")
+        stop = "tool_use" if calls else ("max_tokens" if finish == "length" else "end")
+        return Turn(text=msg.get("content") or "", tool_calls=calls, stop=stop)
+
+    def add_tool_results(self, results: list[ToolResult]) -> None:
+        for r in results:
+            self.messages.append({"role": "tool", "tool_call_id": r.call_id, "content": r.content})