polymath-agent 0.4.0__py3-none-any.whl

polymath/orchestrator/race.py

@@ -0,0 +1,311 @@
+"""
+Concurrent execution patterns: failover (primary + delayed backup, first wins)
+and first-success race (N candidates in parallel, first wins).
+
+Both cancel losers cleanly. Both return (winner_id, result) so the caller can
+log to the AttemptLedger which adapter actually produced the output.
+"""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+
+@dataclass
+class RaceCandidate:
+    """One contender in a race. The factory is invoked at race start; the id
+    is purely for telemetry / ledger logging."""
+    id: str
+    factory: Callable[[], Awaitable[Any]]
+
+
+@dataclass
+class RaceOutcome:
+    winner_id: str
+    result: Any
+    cancelled: list[str]
+    failed: dict[str, Exception]
+
+
+class AllCandidatesFailed(RuntimeError):
+    def __init__(self, failures: dict[str, Exception]) -> None:
+        super().__init__(
+            "All candidates failed: " +
+            "; ".join(f"{cid}: {exc}" for cid, exc in failures.items())
+        )
+        self.failures = failures
+
+
+async def with_failover(
+    primary: RaceCandidate,
+    fallback: RaceCandidate | None,
+    delay_seconds: float = 5.0,
+) -> RaceOutcome:
+    """Start primary; after delay (or on primary failure), also start fallback.
+    First success wins. If both fail, raises AllCandidatesFailed."""
+    if fallback is None:
+        result = await primary.factory()
+        return RaceOutcome(winner_id=primary.id, result=result, cancelled=[], failed={})
+
+    primary_task = asyncio.create_task(primary.factory())
+    failed: dict[str, Exception] = {}
+
+    if delay_seconds > 0:
+        try:
+            await asyncio.wait_for(asyncio.shield(primary_task), timeout=delay_seconds)
+            return RaceOutcome(
+                winner_id=primary.id,
+                result=primary_task.result(),
+                cancelled=[],
+                failed={},
+            )
+        except asyncio.TimeoutError:
+            pass
+        except Exception as e:
+            failed[primary.id] = e
+
+    fallback_task: asyncio.Task | None = None
+    if not primary_task.done() or primary_task.cancelled() or primary_task.exception() is not None:
+        fallback_task = asyncio.create_task(fallback.factory())
+
+    pending: set[asyncio.Task] = set()
+    if not primary_task.done():
+        pending.add(primary_task)
+    if fallback_task is not None:
+        pending.add(fallback_task)
+
+    while pending:
+        done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
+        for task in done:
+            cid = primary.id if task is primary_task else fallback.id
+            if task.cancelled():
+                continue
+            exc = task.exception()
+            if exc is not None:
+                failed[cid] = exc
+                continue
+            cancelled = []
+            for other in pending:
+                other.cancel()
+                cancelled.append(primary.id if other is primary_task else fallback.id)
+            await asyncio.gather(*pending, return_exceptions=True)
+            return RaceOutcome(
+                winner_id=cid,
+                result=task.result(),
+                cancelled=cancelled,
+                failed=failed,
+            )
+    raise AllCandidatesFailed(failed)
+
+
+async def race_first_success(
+    candidates: list[RaceCandidate],
+    timeout_seconds: float | None = None,
+) -> RaceOutcome:
+    """Fire all candidates in parallel; first non-error result wins; cancel rest."""
+    if not candidates:
+        raise AllCandidatesFailed({})
+    if len(candidates) == 1:
+        result = await candidates[0].factory()
+        return RaceOutcome(winner_id=candidates[0].id, result=result, cancelled=[], failed={})
+
+    tasks: dict[asyncio.Task, str] = {}
+    for c in candidates:
+        tasks[asyncio.create_task(c.factory())] = c.id
+
+    failed: dict[str, Exception] = {}
+    pending = set(tasks.keys())
+    deadline = (asyncio.get_event_loop().time() + timeout_seconds) if timeout_seconds else None
+
+    while pending:
+        timeout = max(0.0, deadline - asyncio.get_event_loop().time()) if deadline else None
+        done, pending = await asyncio.wait(
+            pending,
+            return_when=asyncio.FIRST_COMPLETED,
+            timeout=timeout,
+        )
+        if not done:  # timeout
+            for t in pending:
+                t.cancel()
+            raise AllCandidatesFailed({**failed, "*": TimeoutError("race timed out")})
+        for task in done:
+            cid = tasks[task]
+            if task.cancelled():
+                continue
+            exc = task.exception()
+            if exc is not None:
+                failed[cid] = exc
+                continue
+            cancelled: list[str] = []
+            for other in pending:
+                other.cancel()
+                cancelled.append(tasks[other])
+            await asyncio.gather(*pending, return_exceptions=True)
+            return RaceOutcome(
+                winner_id=cid,
+                result=task.result(),
+                cancelled=cancelled,
+                failed=failed,
+            )
+    raise AllCandidatesFailed(failed)
+
+
+@dataclass
+class StreamCandidate:
+    """A streaming contender. The factory must return an async iterator of strings."""
+    id: str
+    factory: Callable[[], Any]   # returns AsyncIterator[str] (sync return, async iterate)
+
+
+@dataclass
+class StreamOutcome:
+    winner_id: str
+    output: str
+    cancelled: list[str]
+    failed: dict[str, Exception]
+
+
+async def race_streams(
+    candidates: list[StreamCandidate],
+    on_token: Callable[[str], None] | None = None,
+    timeout_seconds: float | None = None,
+) -> StreamOutcome:
+    """Race multiple async streams. First to emit a token claims the terminal;
+    the rest are silently cancelled. Loser tokens never reach on_token.
+    Raises AllCandidatesFailed if every candidate errors before any token."""
+    if not candidates:
+        raise AllCandidatesFailed({})
+    on_token = on_token or (lambda _t: None)
+    if len(candidates) == 1:
+        output = ""
+        async for token in candidates[0].factory():
+            if isinstance(token, str):
+                output += token
+                on_token(token)
+        return StreamOutcome(winner_id=candidates[0].id, output=output, cancelled=[], failed={})
+
+    queues: dict[str, asyncio.Queue] = {c.id: asyncio.Queue() for c in candidates}
+
+    async def runner(c: StreamCandidate) -> None:
+        try:
+            async for token in c.factory():
+                if isinstance(token, str):
+                    await queues[c.id].put(("tok", token))
+            await queues[c.id].put(("done", ""))
+        except asyncio.CancelledError:
+            raise
+        except Exception as e:
+            await queues[c.id].put(("err", e))
+
+    runner_tasks: dict[str, asyncio.Task] = {
+        c.id: asyncio.create_task(runner(c)) for c in candidates
+    }
+
+    winner_id: str | None = None
+    output = ""
+    failures: dict[str, Exception] = {}
+    cancelled: list[str] = []
+    finished_ids: set[str] = set()
+
+    waiters: dict[asyncio.Task, str] = {}
+    for cid, q in queues.items():
+        waiters[asyncio.create_task(q.get())] = cid
+
+    deadline = (asyncio.get_event_loop().time() + timeout_seconds) if timeout_seconds else None
+
+    try:
+        while waiters:
+            timeout = max(0.0, deadline - asyncio.get_event_loop().time()) if deadline else None
+            done, _ = await asyncio.wait(
+                waiters.keys(),
+                return_when=asyncio.FIRST_COMPLETED,
+                timeout=timeout,
+            )
+            if not done:
+                raise AllCandidatesFailed({**failures, "*": TimeoutError("stream race timed out")})
+
+            for finished_waiter in done:
+                cid = waiters.pop(finished_waiter)
+                kind, payload = finished_waiter.result()
+
+                if kind == "err":
+                    failures[cid] = payload
+                    finished_ids.add(cid)
+                    if winner_id == cid:
+                        raise payload
+                    continue
+
+                if kind == "done":
+                    finished_ids.add(cid)
+                    if winner_id == cid:
+                        return StreamOutcome(
+                            winner_id=winner_id, output=output,
+                            cancelled=cancelled, failed=failures,
+                        )
+                    continue
+
+                # kind == "tok"
+                if winner_id is None:
+                    winner_id = cid
+                    output += payload
+                    on_token(payload)
+                    waiters[asyncio.create_task(queues[cid].get())] = cid
+                    for other_id, task in runner_tasks.items():
+                        if other_id != cid and not task.done():
+                            task.cancel()
+                            cancelled.append(other_id)
+                elif cid == winner_id:
+                    output += payload
+                    on_token(payload)
+                    waiters[asyncio.create_task(queues[cid].get())] = cid
+                # else: loser token, drop silently
+    finally:
+        for task in runner_tasks.values():
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*runner_tasks.values(), return_exceptions=True)
+        for waiter in waiters:
+            waiter.cancel()
+        await asyncio.gather(*waiters.keys(), return_exceptions=True)
+
+    if winner_id is None:
+        raise AllCandidatesFailed(failures)
+    return StreamOutcome(
+        winner_id=winner_id, output=output,
+        cancelled=cancelled, failed=failures,
+    )
+
+
+async def gather_all(
+    candidates: list[RaceCandidate],
+    timeout_seconds: float | None = None,
+) -> dict[str, Any | Exception]:
+    """Run all in parallel, wait for all (or timeout). Used by ensemble.
+    Returns {candidate_id: result_or_exception}."""
+    if not candidates:
+        return {}
+    tasks: dict[str, asyncio.Task] = {
+        c.id: asyncio.create_task(c.factory()) for c in candidates
+    }
+    try:
+        if timeout_seconds is not None:
+            await asyncio.wait_for(
+                asyncio.gather(*tasks.values(), return_exceptions=True),
+                timeout=timeout_seconds,
+            )
+    except asyncio.TimeoutError:
+        for task in tasks.values():
+            if not task.done():
+                task.cancel()
+        await asyncio.gather(*tasks.values(), return_exceptions=True)
+
+    results: dict[str, Any | Exception] = {}
+    for cid, task in tasks.items():
+        if task.cancelled():
+            results[cid] = TimeoutError("cancelled by ensemble timeout")
+        elif task.exception() is not None:
+            results[cid] = task.exception()
+        else:
+            results[cid] = task.result()
+    return results

polymath/orchestrator/run_controller.py

@@ -0,0 +1,91 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+from polymath.config import TaskComplexity, TaskType
+from polymath.orchestrator.attempt_ledger import AttemptLedger
+from polymath.orchestrator.output_policy import OutputPolicy
+from polymath.orchestrator.state_responder import StateResponder
+
+
+class ExecutionStrategy(Enum):
+    STATE = "state"
+    CHAT = "chat"
+    PIPELINE = "pipeline"
+    ENSEMBLE = "ensemble"
+    FANOUT = "fanout"
+
+
+@dataclass
+class RunRequest:
+    text: str
+    explicit_ask: bool
+    task_type: TaskType
+    complexity: TaskComplexity
+    verify: bool
+    simplify: bool
+    profile: str
+    registry: list
+    ensemble: bool = False
+    fanout: bool = False
+    force_provider: str = ""
+
+
+@dataclass
+class RunDecision:
+    strategy: ExecutionStrategy
+    local_response: str | None = None
+    attempt_ledger: AttemptLedger = field(default_factory=AttemptLedger)
+
+
+class RunController:
+    def __init__(
+        self,
+        state_responder: StateResponder | None = None,
+        output_policy: OutputPolicy | None = None,
+    ) -> None:
+        self.state_responder = state_responder or StateResponder()
+        self.output_policy = output_policy or OutputPolicy()
+
+    def decide(self, request: RunRequest) -> RunDecision:
+        local_response = self.state_responder.answer(
+            text=request.text,
+            registry=request.registry,
+            profile=request.profile,
+        )
+        if local_response is not None:
+            return RunDecision(
+                strategy=ExecutionStrategy.STATE,
+                local_response=local_response
+            )
+
+        # User overrides
+        if request.ensemble:
+            return RunDecision(strategy=ExecutionStrategy.ENSEMBLE)
+        if request.fanout:
+            return RunDecision(strategy=ExecutionStrategy.FANOUT)
+
+        # 3. CHAT for simple general tasks or forced provider
+        use_direct_chat = request.explicit_ask or request.force_provider != "" or (
+            request.task_type == TaskType.GENERAL and
+            request.complexity == TaskComplexity.SIMPLE and
+            not request.verify and
+            not request.simplify
+        )
+
+        if use_direct_chat:
+            return RunDecision(strategy=ExecutionStrategy.CHAT)
+
+        # Ensemble and fanout are never chosen for the user. Both multiply
+        # what a question costs — ensemble runs N models plus a merger — and
+        # the trigger would have been a heuristic classification of the
+        # prompt text, so the same question could cost one model or four
+        # with nothing to distinguish them beforehand. They stay on the
+        # explicit flags above: `@all` / --ensemble, and --fanout. That is
+        # also why the ensemble flow carries a cost warning; the warning is
+        # written for someone who asked for it.
+        return RunDecision(
+            strategy=ExecutionStrategy.PIPELINE,
+            attempt_ledger=AttemptLedger(),
+        )

polymath/orchestrator/speculative_review.py

@@ -0,0 +1,120 @@
+"""
+Speculative concurrent reviewer.
+
+While primary's agent loop runs, the reviewer is fed partial output snapshots
+and emits interim concerns in the background. When primary finishes, the
+final review call uses those accumulated notes as priming — typically faster
+than a cold review of the whole output.
+
+Opt-in via cfg["memory"]["speculative_review"] (default off until proven on
+real workloads). Failures are silent — they degrade to sequential review.
+"""
+from __future__ import annotations
+
+import asyncio
+
+from polymath.adapters.base import BaseAdapter, Message
+
+
+SPECULATIVE_SYSTEM = (
+    "You are a critical reviewer reading partial output from another model that is "
+    "still working. Note ONLY concrete issues you can already see (factual errors, "
+    "missing requirements, broken logic). Skip style/wording. Limit to 3 bullet "
+    "points. If you see no issues yet, output exactly: NONE_YET"
+)
+
+FINAL_REVIEW_SYSTEM = (
+    "You are a critical reviewer. Format strictly as:\n"
+    "VERDICT: APPROVED | NEEDS_WORK\n"
+    "ISSUES:\n- ...\n"
+    "SUGGESTIONS:\n- ...\n"
+    "If APPROVED with no issues, just say: VERDICT: APPROVED"
+)
+
+
+class SpeculativeReviewer:
+    """Background reviewer that observes partial output during primary execution."""
+
+    def __init__(
+        self,
+        adapter: BaseAdapter,
+        model_id: str,
+        task: str,
+        max_inflight: int = 1,
+    ) -> None:
+        self.adapter = adapter
+        self.model_id = model_id
+        self.task = task
+        self._notes: list[str] = []
+        self._inflight: asyncio.Task | None = None
+        self._max_inflight = max_inflight
+        self.last_partial_seen = ""
+
+    def kick_off(self, partial_output: str) -> None:
+        """Schedule a speculative review on the current partial output. Non-blocking.
+        Skips if a previous speculation is still in flight (only 1 at a time)."""
+        if not partial_output or partial_output == self.last_partial_seen:
+            return
+        if self._inflight is not None and not self._inflight.done():
+            return
+        self.last_partial_seen = partial_output
+        self._inflight = asyncio.create_task(self._speculate(partial_output))
+
+    async def _speculate(self, partial_output: str) -> None:
+        try:
+            resp = await self.adapter.complete(
+                messages=[Message(role="user", content=(
+                    f"ORIGINAL TASK:\n{self.task}\n\n"
+                    f"PARTIAL OUTPUT (in progress):\n{partial_output}\n\n"
+                    f"Concerns so far?"
+                ))],
+                model_id=self.model_id,
+                system=SPECULATIVE_SYSTEM,
+                temperature=0.2,
+            )
+            text = resp if isinstance(resp, str) else getattr(resp, "content", "")
+            if text and "NONE_YET" not in text.upper()[:30]:
+                self._notes.append(text.strip())
+        except Exception:
+            pass
+
+    async def finalize(self, final_output: str, timeout_seconds: float = 2.0) -> str:
+        """Wait briefly for any in-flight speculation, then make the final verdict call."""
+        if self._inflight is not None and not self._inflight.done():
+            try:
+                await asyncio.wait_for(asyncio.shield(self._inflight), timeout=timeout_seconds)
+            except (asyncio.TimeoutError, Exception):
+                pass
+
+        notes_block = ""
+        if self._notes:
+            joined_notes = "\n\n".join(self._notes[-3:])
+            notes_block = (
+                f"\nINTERIM CONCERNS YOU NOTED EARLIER (use these as priming, "
+                f"verify against the final output):\n{joined_notes}\n"
+            )
+
+        prompt = (
+            f"ORIGINAL TASK:\n{self.task}\n\n"
+            f"FINAL RESPONSE TO REVIEW:\n{final_output}\n"
+            f"{notes_block}"
+        )
+        resp = await self.adapter.complete(
+            messages=[Message(role="user", content=prompt)],
+            model_id=self.model_id,
+            system=FINAL_REVIEW_SYSTEM,
+            temperature=0.3,
+        )
+        return resp if isinstance(resp, str) else getattr(resp, "content", "")
+
+    async def cancel(self) -> None:
+        if self._inflight is not None and not self._inflight.done():
+            self._inflight.cancel()
+            try:
+                await self._inflight
+            except (asyncio.CancelledError, Exception):
+                pass
+
+    @property
+    def notes(self) -> list[str]:
+        return list(self._notes)