millforge 0.1.0__py3-none-any.whl

millforge/_forge/base_runner.py

@@ -0,0 +1,121 @@
+"""Private Forge graph assembly for the public Millforge base runner facade."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Callable
+
+from millforge._forge.adapter import ForgeContextFactory, ForgeGuardrailBackend
+from millforge._forge.errors import ForgeError
+from millforge.base.composition import MillforgeBaseComponents
+from millforge.compiled_plan import CompiledHarnessPlan
+from millforge.contracts import (
+    ArtifactRef,
+    CompiledHarnessRef,
+    HarnessExecutionRequest,
+    HarnessExecutionResult,
+)
+from millforge.exceptions import BackendTranslationError
+from millforge.protocols import (
+    CancellationResolver,
+    ModelClient,
+    RuntimeArtifactWriter,
+    RuntimeClock,
+)
+from millforge.runtime import DefaultHarnessRuntime
+
+
+@dataclass(frozen=True, slots=True)
+class _PinnedCompiledHarnessLoader:
+    plan: CompiledHarnessPlan
+
+    async def load(self, ref: CompiledHarnessRef) -> CompiledHarnessPlan:
+        del ref
+        return self.plan
+
+
+class _LazyRuntimeArtifactWriter:
+    """Defer writer construction until runtime run-directory preparation succeeds."""
+
+    __slots__ = ("_error", "_factory", "_run_directory", "_writer")
+
+    def __init__(
+        self,
+        *,
+        factory: Callable[[Path], RuntimeArtifactWriter],
+        run_directory: Path,
+    ) -> None:
+        self._factory = factory
+        self._run_directory = run_directory
+        self._writer: RuntimeArtifactWriter | None = None
+        self._error: Exception | None = None
+
+    def _resolve(self) -> RuntimeArtifactWriter:
+        if self._error is not None:
+            raise self._error
+        if self._writer is None:
+            try:
+                self._writer = self._factory(self._run_directory)
+            except Exception as exc:
+                self._error = exc
+                raise
+        return self._writer
+
+    async def write_terminal_result(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_terminal_result(ref, data)
+
+    async def write_execution_summary(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_execution_summary(ref, data)
+
+    async def write_events(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_events(ref, data)
+
+    async def write_tool_trace(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_tool_trace(ref, data)
+
+    async def write_metrics(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_metrics(ref, data)
+
+    async def write_artifact_manifest(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_artifact_manifest(ref, data)
+
+    async def write_diagnostic(self, ref: ArtifactRef, data: object) -> None:
+        await self._resolve().write_diagnostic(ref, data)
+
+
+async def execute_base_invocation(
+    *,
+    components: MillforgeBaseComponents,
+    request: HarnessExecutionRequest,
+    model_client: ModelClient,
+    clock: RuntimeClock,
+    cancellation_resolver: CancellationResolver,
+    artifact_writer_factory: Callable[[Path], RuntimeArtifactWriter],
+) -> HarnessExecutionResult:
+    """Build and execute one fresh private adapter graph."""
+    try:
+        loader = _PinnedCompiledHarnessLoader(components.compiled_plan)
+        writer = _LazyRuntimeArtifactWriter(
+            factory=artifact_writer_factory,
+            run_directory=request.run_directory.path,
+        )
+        tool_executor = components.tool_executor.fork_for_invocation()
+        backend = ForgeGuardrailBackend(
+            model_client=model_client,
+            tool_executor=tool_executor,
+            plan_loader=loader,
+            context_factory=ForgeContextFactory(),
+            clock=clock,
+            cancellation_resolver=cancellation_resolver,
+        )
+        runtime = DefaultHarnessRuntime(
+            backend=backend,
+            plan_loader=loader,
+            artifact_writer=writer,
+            clock=clock,
+            cancellation_resolver=cancellation_resolver,
+        )
+        return await runtime.execute(request)
+    except ForgeError:
+        raise BackendTranslationError("Private Forge backend failed") from None

millforge/_forge/clients/__init__.py

@@ -0,0 +1,10 @@
+"""Private transport-free Forge client protocol helpers."""
+
+from millforge._forge.clients.base import ChunkType, LLMClient, StreamChunk, TokenUsage
+
+__all__ = [
+    "ChunkType",
+    "LLMClient",
+    "StreamChunk",
+    "TokenUsage",
+]

millforge/_forge/clients/base.py

@@ -0,0 +1,200 @@
+"""Streaming types and LLM client protocol."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import AsyncIterator
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, Protocol, runtime_checkable
+
+from millforge._forge.core.workflow import LLMResponse, ToolSpec
+
+# Verbatim OpenAI-shape payloads forwarded by the proxy. The proxy hands the
+# client the user's original ``tools`` array so the backend sees the exact
+# schema the client authored, instead of forge's reconstructed ToolSpec.
+RawOpenAITools = list[dict[str, Any]]
+RawOpenAIMessages = list[dict[str, Any]]
+
+
+@dataclass(frozen=True)
+class TokenUsage:
+    """Token counts from a single LLM response.
+
+    Populated from the server's ``usage`` field when available (e.g.
+    llama-server).  Backends that don't report usage leave the client's
+    ``last_usage`` empty and the context manager falls back to heuristic
+    estimation.
+    """
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+# Both Ollama and llama-server use the OpenAI tool schema format today.
+# If a backend diverges, move this back into the relevant client module.
+def format_tool(spec: ToolSpec) -> dict[str, Any]:
+    """Convert a ToolSpec into the OpenAI-compatible tool schema."""
+    return {
+        "type": "function",
+        "function": {
+            "name": spec.name,
+            "description": spec.description,
+            "parameters": spec.get_json_schema(),
+        },
+    }
+
+
+def decode_tool_args(raw: Any) -> Any:
+    """Decode a tool-call ``arguments`` payload, fail-loud.
+
+    JSON-string args are parsed; on malformed JSON the raw string is returned
+    unchanged (a non-dict). ``ResponseValidator``'s args-shape check then routes
+    it through the tool-error channel instead of crashing the parser or coercing
+    to ``{}`` — so a structural arg failure rides the same lane as a runtime
+    tool error rather than a trailing retry nudge.
+
+    Non-string payloads (an already-decoded dict from Ollama / the Anthropic
+    SDK, or any other shape) pass through untouched for the validator to judge.
+    A missing or empty payload is a no-arg call (``{}``).
+    """
+    if raw is None:
+        return {}
+    if not isinstance(raw, str):
+        return raw
+    if not raw:
+        return {}
+    try:
+        return json.loads(raw)
+    except json.JSONDecodeError:
+        return raw
+
+
+class ChunkType(str, Enum):
+    """What kind of partial data a stream chunk carries."""
+
+    TEXT_DELTA = "text_delta"
+    TOOL_CALL_DELTA = "tool_call_delta"
+    FINAL = "final"
+    RETRY = "retry"
+
+
+@dataclass(frozen=True)
+class StreamChunk:
+    """A single chunk from a streaming LLM response.
+
+    Consumers (UI, logging) process TEXT_DELTA and TOOL_CALL_DELTA as they
+    arrive. The runner ignores all chunks except FINAL, which carries the
+    resolved response. On RETRY, consumers should discard the partial output
+    from the failed attempt.
+    """
+
+    type: ChunkType
+    content: str = ""
+    response: LLMResponse | None = None
+
+
+@runtime_checkable
+class LLMClient(Protocol):
+    """Interface that client adapters implement.
+
+    The client is responsible for:
+    1. Sending messages to the LLM backend
+    2. Parsing the response into ToolCall or TextResponse
+    3. Handling native FC or prompt-injected calling internally
+    4. Optionally streaming partial responses via send_stream()
+
+    The client does NOT retry. Retry logic lives in the WorkflowRunner.
+    """
+
+    api_format: str
+    """Wire format for Message.to_api_dict(): 'ollama' or 'openai'."""
+
+    model: str
+    """The backend model identity, sent verbatim as the wire "model" field
+    (the served-model-name, gguf stem, or model tag depending on backend).
+    Distinct from any sampling-registry lookup key a client also derives."""
+
+    async def send(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[ToolSpec] | None = None,
+        sampling: dict[str, Any] | None = None,
+        passthrough: dict[str, Any] | None = None,
+        inbound_anthropic_body: dict[str, Any] | None = None,
+        raw_openai_tools: RawOpenAITools | None = None,
+    ) -> LLMResponse:
+        """Send messages and return a parsed response.
+
+        Returns list[ToolCall] if the model produced valid tool invocations.
+        Returns TextResponse if the model produced text (reasoning, refusal,
+        or malformed output that couldn't be parsed as a tool call).
+
+        The runner inspects the response and decides whether to retry.
+
+        Args:
+            messages: API-format messages to send.
+            tools: Tool specs to include with the request.
+            sampling: Optional per-call sampling overrides
+                (``temperature``, ``top_p``, ``top_k``, ``min_p``,
+                ``repeat_penalty``, ``presence_penalty``, ``seed``).
+                Per-call values win over instance state for this call only;
+                the client's instance fields are not mutated.
+            passthrough: Optional dict of inbound body fields forge doesn't
+                own. The client merges these into the outbound body before
+                overlaying its own fields (model, messages, tools, sampling).
+                Used by the proxy to preserve user intent (max_tokens, stop,
+                tool_choice, etc.) without forge having to enumerate every
+                supported field. None = no extras to merge.
+            inbound_anthropic_body: Path-1 only — when set, the AnthropicClient
+                will send this body verbatim (bypassing its deconstruct/rebuild
+                path) to preserve block-level Anthropic fields like
+                ``cache_control``. The runner clears this kwarg on any
+                forge-mutation (retry / compaction / context warning) so
+                only the clean first-attempt call rides verbatim. Other
+                clients accept and ignore. See ADR-015.
+            raw_openai_tools: Proxy-only — the client's verbatim OpenAI
+                ``tools`` array. When set, LlamafileClient's native path sends
+                it as-is instead of re-emitting ``format_tool(spec)``, so the
+                backend sees the original schema (no name/schema drift). Other
+                clients accept and ignore.
+        """
+        ...
+
+    async def send_stream(
+        self,
+        messages: list[dict[str, str]],
+        tools: list[ToolSpec] | None = None,
+        sampling: dict[str, Any] | None = None,
+        passthrough: dict[str, Any] | None = None,
+        inbound_anthropic_body: dict[str, Any] | None = None,
+        raw_openai_tools: RawOpenAITools | None = None,
+    ) -> AsyncIterator[StreamChunk]:
+        """Send messages and yield streaming chunks.
+
+        Yields TEXT_DELTA or TOOL_CALL_DELTA chunks as they arrive.
+        The final chunk has type FINAL and carries the resolved LLMResponse
+        (same list[ToolCall] | TextResponse as send() would return).
+
+        The runner forwards chunks to its on_chunk callback for UI/logging,
+        then inspects the FINAL chunk and decides whether to retry.
+
+        Args:
+            messages: API-format messages to send.
+            tools: Tool specs to include with the request.
+            sampling: Optional per-call sampling overrides (see ``send``).
+                Per-call values win over instance state without mutating self.
+            passthrough: Optional inbound-body extras dict (see ``send``).
+            inbound_anthropic_body: Optional path-1 verbatim body (see ``send``).
+            raw_openai_tools: Optional verbatim OpenAI tools array (see ``send``).
+        """
+        ...
+
+    async def get_context_length(self) -> int | None:
+        """Query the backend for its configured context window size."""
+        ...
+
+    async def aclose(self) -> None:
+        """Release held network resources (e.g. the httpx connection pool)."""
+        ...

millforge/_forge/context/__init__.py

@@ -0,0 +1,23 @@
+"""Private Forge context management subset without hardware discovery."""
+
+from millforge._forge.context.manager import (
+    CompactEvent,
+    ContextManager,
+    default_context_warning,
+)
+from millforge._forge.context.strategies import (
+    CompactStrategy,
+    NoCompact,
+    SlidingWindowCompact,
+    TieredCompact,
+)
+
+__all__ = [
+    "CompactEvent",
+    "CompactStrategy",
+    "ContextManager",
+    "NoCompact",
+    "SlidingWindowCompact",
+    "TieredCompact",
+    "default_context_warning",
+]

millforge/_forge/context/manager.py

@@ -0,0 +1,178 @@
+"""Context manager for budget tracking and compaction triggering."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from millforge._forge.context.strategies import CompactStrategy
+from millforge._forge.core.messages import Message
+
+
+@dataclass(frozen=True)
+class CompactEvent:
+    """Emitted by ContextManager when compaction fires."""
+
+    step_index: int
+    tokens_before: int
+    tokens_after: int
+    budget_tokens: int
+    messages_before: int
+    messages_after: int
+    phase_reached: int
+
+
+# ── Default context warning ──────────────────────────────────────
+
+
+def default_context_warning(tokens: int, budget: int, pct: float) -> str | None:
+    """Default context threshold callback.
+
+    Returns an escalating warning based on how full the context is.
+    """
+    if pct >= 0.80:
+        return (
+            f"[Context usage: {pct:.0%} ({tokens:,} / {budget:,} tokens). "
+            "Context is nearly full. Older tool results and reasoning will be "
+            "compacted soon — key information may be lost. Summarize critical "
+            "findings now and prioritize completing the current task.]"
+        )
+    if pct >= 0.65:
+        return (
+            f"[Context usage: {pct:.0%} ({tokens:,} / {budget:,} tokens). "
+            "Context is filling up. When compaction triggers, older tool results "
+            "and reasoning will be condensed. Be concise in your responses and "
+            "front-load important information.]"
+        )
+    return (
+        f"[Context usage: {pct:.0%} ({tokens:,} / {budget:,} tokens). "
+        "Be mindful of context usage.]"
+    )
+
+
+class ContextManager:
+    """Manages context window budget and triggers compaction."""
+
+    def __init__(
+        self,
+        strategy: CompactStrategy,
+        budget_tokens: int,
+        on_compact: Callable[[CompactEvent], None] | None = None,
+        context_thresholds: list[float] | None = None,
+        on_context_threshold: Callable[[int, int, float], str | None] | None = None,
+    ) -> None:
+        """
+        Args:
+            strategy: Compaction strategy to use. The strategy owns its own
+                compaction thresholds (e.g. ``TieredCompact(compact_threshold=0.75)``
+                or ``TieredCompact(phase_thresholds=(0.6, 0.75, 0.9))``).
+            budget_tokens: Maximum context budget in tokens.
+            on_compact: Callback invoked when compaction fires. Receives a
+                CompactEvent with before/after token counts, phase reached,
+                and which messages were affected. Use for logging, debugging,
+                or surfacing compaction to a UI.
+            context_thresholds: Sorted list of budget fractions (e.g.
+                ``[0.5, 0.65, 0.8]``) at which to fire the context
+                threshold callback. Each threshold fires at most once per
+                session (resets if usage drops below it after compaction).
+                Defaults to None (disabled).
+            on_context_threshold: Callback invoked when a context threshold
+                is crossed. Receives ``(tokens, budget, pct)`` and returns
+                an optional string to inject as a transient system message
+                before the next inference call. Return None to skip
+                injection. Defaults to None (disabled).
+        """
+        self.strategy = strategy
+        self.budget_tokens = budget_tokens
+        self.on_compact = on_compact
+        self._context_thresholds = (
+            sorted(context_thresholds) if context_thresholds else []
+        )
+        self._on_context_threshold = on_context_threshold
+        self._fired_thresholds: set[float] = set()
+        self._last_known_tokens: int | None = None
+
+    def update_token_count(self, total_tokens: int) -> None:
+        """Record actual token count from the backend.
+
+        Called after each LLM response when the backend reports usage.
+        Subsequent calls to ``estimate_tokens`` will return this value
+        until the next update.
+        """
+        self._last_known_tokens = total_tokens
+
+    def estimate_tokens(self, messages: list[Message]) -> int:
+        """Return actual token count if available, else char/4 heuristic."""
+        if self._last_known_tokens is not None:
+            return self._last_known_tokens
+        return (
+            sum(
+                len(message.content) + len(message.reasoning_content or "")
+                for message in messages
+            )
+            // 4
+        )
+
+    def check_thresholds(self, messages: list[Message]) -> str | None:
+        """Check context thresholds and return an optional warning to inject.
+
+        Fires the ``on_context_threshold`` callback when usage crosses a
+        configured threshold for the first time. Thresholds reset if usage
+        drops below them (e.g. after compaction).
+
+        Returns:
+            A string to inject as a transient system message, or None.
+        """
+        if not self._context_thresholds or not self._on_context_threshold:
+            return None
+
+        tokens = self.estimate_tokens(messages)
+        if self.budget_tokens <= 0:
+            return None
+
+        pct = tokens / self.budget_tokens
+
+        # Reset thresholds that usage has dropped below (after compaction)
+        self._fired_thresholds = {t for t in self._fired_thresholds if pct >= t}
+
+        # Find the highest unfired threshold that has been crossed
+        highest_crossed: float | None = None
+        for threshold in self._context_thresholds:
+            if pct >= threshold and threshold not in self._fired_thresholds:
+                highest_crossed = threshold
+
+        if highest_crossed is None:
+            return None
+
+        self._fired_thresholds.add(highest_crossed)
+        return self._on_context_threshold(tokens, self.budget_tokens, pct)
+
+    def maybe_compact(
+        self,
+        messages: list[Message],
+        step_index: int = 0,
+        step_hint: str = "",
+    ) -> list[Message]:
+        """Delegate to the strategy, which owns threshold logic."""
+        tokens_before = self.estimate_tokens(messages)
+
+        result, phase = self.strategy.compact(
+            messages, self.budget_tokens, step_hint=step_hint
+        )
+
+        if phase == 0:
+            return messages
+
+        if self.on_compact is not None:
+            event = CompactEvent(
+                step_index=step_index,
+                tokens_before=tokens_before,
+                tokens_after=self.estimate_tokens(result),
+                budget_tokens=self.budget_tokens,
+                messages_before=len(messages),
+                messages_after=len(result),
+                phase_reached=phase,
+            )
+            self.on_compact(event)
+
+        return result