polymath-agent 0.4.0__py3-none-any.whl

polymath/orchestrator/ensemble.py

@@ -0,0 +1,229 @@
+"""
+Ensemble execution: run the same prompt across N models in parallel,
+then optionally merge with a designated merger model.
+
+Bypasses the full pipeline (no clarify / plan / review per worker — that
+would N× the round-trips). Designed for the @all REPL trigger when the
+user wants multiple perspectives or a consensus answer.
+
+Memory retrieval happens once; all workers see the same chunks.
+"""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from typing import Awaitable, Callable
+
+from polymath.adapters.base import BaseAdapter, Message
+from polymath.config import ModelInfo, TaskType
+from polymath.memory.retriever import retrieve_or_fallback
+from polymath.orchestrator.race import RaceCandidate, gather_all
+
+
+DEFAULT_TIMEOUT_SECONDS = 60.0
+
+
+@dataclass
+class EnsembleWorker:
+    adapter: BaseAdapter
+    model: ModelInfo
+
+
+@dataclass
+class WorkerResult:
+    model_id: str
+    display_name: str
+    output: str = ""
+    error: str = ""
+
+    @property
+    def ok(self) -> bool:
+        return not self.error and bool(self.output)
+
+
+@dataclass
+class EnsembleResult:
+    workers: list[WorkerResult] = field(default_factory=list)
+    merged: str = ""
+    merger_model: str = ""
+    disagreement_notes: str = ""
+
+    @property
+    def successful(self) -> list[WorkerResult]:
+        return [w for w in self.workers if w.ok]
+
+
+MERGER_SYSTEM = (
+    "You are a merger / consensus arbiter. You will be shown the same task "
+    "answered independently by multiple models. Produce:\n"
+    "1) CONSENSUS — the best single answer, integrating points the models agreed on, "
+    "preferring better-supported claims. Do not invent new content.\n"
+    "2) DISAGREEMENTS — bullet points where the models gave materially different answers, "
+    "with which model said what. Skip purely cosmetic differences.\n"
+    "Format strictly as:\n"
+    "## CONSENSUS\n<answer>\n\n## DISAGREEMENTS\n- <model A> said X; <model B> said Y\n"
+    "If models substantially agreed, write `(none significant)` under DISAGREEMENTS."
+)
+
+
+async def run_ensemble(
+    *,
+    workers: list[EnsembleWorker],
+    user_input: str,
+    task_type: TaskType,
+    active_project: str,
+    cfg: dict,
+    merger: EnsembleWorker | None = None,
+    extra_system: str = "",
+    timeout_seconds: float = DEFAULT_TIMEOUT_SECONDS,
+    on_worker_started: Callable[[str], None] | None = None,
+    on_worker_done: Callable[[WorkerResult], None] | None = None,
+) -> EnsembleResult:
+    if not workers:
+        return EnsembleResult()
+
+    on_worker_started = on_worker_started or (lambda _name: None)
+    on_worker_done = on_worker_done or (lambda _result: None)
+
+    project_ctx = (
+        await retrieve_or_fallback(active_project, user_input, task_type, cfg)
+        if active_project else ""
+    )
+    base_system = "Answer concisely and directly."
+    if project_ctx:
+        base_system = f"{project_ctx}\n\n{base_system}"
+    if extra_system:
+        base_system = f"{base_system}\n\n{extra_system}"
+
+    candidates: list[RaceCandidate] = []
+    for w in workers:
+        on_worker_started(w.model.display_name)
+        candidates.append(RaceCandidate(
+            id=w.model.id,
+            factory=_make_call(w, user_input, base_system),
+        ))
+
+    raw = await gather_all(candidates, timeout_seconds=timeout_seconds)
+
+    worker_results: list[WorkerResult] = []
+    for w in workers:
+        outcome = raw.get(w.model.id)
+        if isinstance(outcome, Exception):
+            wr = WorkerResult(model_id=w.model.id, display_name=w.model.display_name,
+                              error=str(outcome))
+        else:
+            text = outcome if isinstance(outcome, str) else getattr(outcome, "content", "")
+            wr = WorkerResult(model_id=w.model.id, display_name=w.model.display_name,
+                              output=text or "")
+        worker_results.append(wr)
+        on_worker_done(wr)
+
+    result = EnsembleResult(workers=worker_results)
+
+    successful = result.successful
+    if len(successful) <= 1 or merger is None:
+        if successful:
+            result.merged = successful[0].output
+            result.merger_model = "(no merge: only one successful response)" if len(successful) == 1 else ""
+        return result
+
+    merge_prompt = _build_merge_prompt(user_input, successful)
+    try:
+        resp = await merger.adapter.complete(
+            messages=[Message(role="user", content=merge_prompt)],
+            model_id=merger.model.id,
+            system=MERGER_SYSTEM,
+            temperature=0.3,
+        )
+        merged_text = resp if isinstance(resp, str) else resp.content
+        result.merged, result.disagreement_notes = _split_consensus(merged_text)
+        result.merger_model = merger.model.display_name
+    except Exception as e:
+        # Fall back to the first successful raw output
+        result.merged = successful[0].output
+        result.merger_model = f"(merge failed: {e})"
+
+    return result
+
+
+def _make_call(worker: EnsembleWorker, user_input: str, system: str) -> Callable[[], Awaitable]:
+    async def call():
+        return await worker.adapter.complete(
+            messages=[Message(role="user", content=user_input)],
+            model_id=worker.model.id,
+            system=system,
+            temperature=0.5,
+        )
+    return call
+
+
+def _build_merge_prompt(user_input: str, successful: list[WorkerResult]) -> str:
+    parts = [f"TASK:\n{user_input}\n", "INDEPENDENT ANSWERS:"]
+    for w in successful:
+        parts.append(f"\n--- {w.display_name} ---\n{w.output.strip()}")
+    return "\n".join(parts)
+
+
+def _split_consensus(text: str) -> tuple[str, str]:
+    """Return (consensus, disagreements). Tolerant of header variations."""
+    upper = text.upper()
+    consensus_idx = upper.find("CONSENSUS")
+    disagree_idx = upper.find("DISAGREEMENT")
+    if consensus_idx == -1:
+        return text.strip(), ""
+    consensus_start = text.find("\n", consensus_idx) + 1
+    if disagree_idx == -1 or disagree_idx < consensus_idx:
+        return text[consensus_start:].strip(), ""
+    consensus = text[consensus_start:disagree_idx].strip().rstrip("#").strip()
+    disagree_start = text.find("\n", disagree_idx) + 1
+    disagreements = text[disagree_start:].strip()
+    return consensus, disagreements
+
+
+# ── Worker selection ─────────────────────────────────────────────────────────
+
+def select_ensemble_workers(
+    registry: list,
+    cfg: dict,
+    max_workers: int = 3,
+    exclude_providers: set[str] | None = None,
+) -> tuple[list[EnsembleWorker], EnsembleWorker | None]:
+    """Pick top-N available models, one per provider where possible.
+    Returns (workers, merger). The merger is the highest-quality model."""
+    from polymath.model_policy import adapter_for
+
+    exclude = exclude_providers or set()
+    available = [m for m in registry if m.available and m.provider not in exclude]
+    if not available:
+        return [], None
+
+    by_provider: dict[str, list] = {}
+    for m in available:
+        by_provider.setdefault(m.provider, []).append(m)
+    for provider, ms in by_provider.items():
+        ms.sort(key=lambda x: (-x.quality, x.cost_tier.value))
+
+    chosen: list = []
+    while len(chosen) < max_workers:
+        added = False
+        for provider in sorted(by_provider.keys()):
+            if not by_provider[provider]:
+                continue
+            chosen.append(by_provider[provider].pop(0))
+            added = True
+            if len(chosen) >= max_workers:
+                break
+        if not added:
+            break
+
+    workers: list[EnsembleWorker] = []
+    for m in chosen:
+        ad = adapter_for(m.provider, cfg)
+        if ad:
+            workers.append(EnsembleWorker(adapter=ad, model=m))
+
+    merger: EnsembleWorker | None = None
+    if workers:
+        best = max(workers, key=lambda w: w.model.quality)
+        merger = best
+    return workers, merger

polymath/orchestrator/fanout.py

@@ -0,0 +1,322 @@
+"""
+Specialist fan-out: split a complex task into independent parallel subtasks,
+route each to the best-fit model, run in parallel, integrate results.
+
+Risky: a bad split contaminates everything downstream. Gated behind --fanout
+flag, COMPLEX tasks only. The splitter is conservative — if the task is
+inherently sequential, it returns a single subtask and the executor degrades
+gracefully.
+"""
+from __future__ import annotations
+
+import asyncio
+import json
+import re
+from dataclasses import dataclass, field
+from typing import Awaitable, Callable
+
+from polymath.adapters.base import BaseAdapter, Message
+from polymath.config import ModelInfo, TaskType
+from polymath.orchestrator.race import RaceCandidate, gather_all
+
+
+SPLITTER_SYSTEM = (
+    "You split complex tasks into independent parallel subtasks. Each subtask "
+    "must be executable on its own without seeing the others' outputs. If the "
+    "task is inherently sequential or single-purpose, return one subtask.\n\n"
+    "Output JSON only — an array of objects:\n"
+    '[{"id": "s1", "description": "<what to do>", "type": "code|analysis|creative|math|research|general"}]\n\n'
+    "Rules: 1-5 subtasks max. Each description must be concrete enough to act on alone. "
+    "No prose outside the JSON."
+)
+
+JOINER_SYSTEM = (
+    "You integrate independent subtask outputs into a single coherent answer to the "
+    "original task. Resolve any conflicts between subtask outputs explicitly. Do not "
+    "introduce information that isn't present in the subtask outputs. Match the format "
+    "the original task asked for."
+)
+
+
+@dataclass
+class Subtask:
+    id: str
+    description: str
+    task_type: TaskType
+
+
+@dataclass
+class SubtaskResult:
+    subtask: Subtask
+    model_id: str
+    model_display: str
+    output: str = ""
+    error: str = ""
+
+    @property
+    def ok(self) -> bool:
+        return not self.error and bool(self.output)
+
+
+@dataclass
+class FanoutResult:
+    subtasks: list[Subtask] = field(default_factory=list)
+    results: list[SubtaskResult] = field(default_factory=list)
+    joined: str = ""
+    join_model: str = ""
+    skipped_reason: str = ""
+
+    @property
+    def successful(self) -> list[SubtaskResult]:
+        return [r for r in self.results if r.ok]
+
+
+# ── Public API ──────────────────────────────────────────────────────────────
+
+async def run_fanout(
+    *,
+    task: str,
+    plan: str,
+    registry: list,
+    cfg: dict,
+    splitter_adapter: BaseAdapter,
+    splitter_model: ModelInfo,
+    joiner_adapter: BaseAdapter,
+    joiner_model: ModelInfo,
+    profile: str = "quality-first",
+    on_split: Callable[[list[Subtask]], None] | None = None,
+    on_subtask_started: Callable[[Subtask, ModelInfo], None] | None = None,
+    on_subtask_done: Callable[[SubtaskResult], None] | None = None,
+    timeout_seconds: float = 120.0,
+) -> FanoutResult:
+    on_split = on_split or (lambda _ts: None)
+    on_subtask_started = on_subtask_started or (lambda _s, _m: None)
+    on_subtask_done = on_subtask_done or (lambda _r: None)
+
+    subtasks = await split_task(task, plan, splitter_adapter, splitter_model.id)
+    on_split(subtasks)
+
+    if not subtasks:
+        return FanoutResult(skipped_reason="splitter returned no subtasks")
+    if len(subtasks) == 1:
+        return FanoutResult(
+            subtasks=subtasks,
+            skipped_reason="single subtask — fanout would be a no-op",
+        )
+
+    routed = _route_subtasks(subtasks, registry, cfg, profile)
+    for subtask, route in zip(subtasks, routed):
+        if route:
+            on_subtask_started(subtask, route[0])
+
+    candidates: list[RaceCandidate] = []
+    valid: list[tuple[Subtask, ModelInfo, BaseAdapter]] = []
+    for subtask, route in zip(subtasks, routed):
+        if not route:
+            continue
+        model_info, adapter = route
+        valid.append((subtask, model_info, adapter))
+        candidates.append(RaceCandidate(
+            id=subtask.id,
+            factory=_make_subtask_call(adapter, model_info, subtask, task),
+        ))
+
+    raw = await gather_all(candidates, timeout_seconds=timeout_seconds)
+
+    results: list[SubtaskResult] = []
+    for subtask, model_info, _adapter in valid:
+        outcome = raw.get(subtask.id)
+        if isinstance(outcome, Exception):
+            r = SubtaskResult(
+                subtask=subtask, model_id=model_info.id,
+                model_display=model_info.display_name, error=str(outcome),
+            )
+        else:
+            text = outcome if isinstance(outcome, str) else getattr(outcome, "content", "")
+            r = SubtaskResult(
+                subtask=subtask, model_id=model_info.id,
+                model_display=model_info.display_name, output=text or "",
+            )
+        results.append(r)
+        on_subtask_done(r)
+
+    successful = [r for r in results if r.ok]
+    if not successful:
+        return FanoutResult(
+            subtasks=subtasks, results=results,
+            skipped_reason="all subtasks failed",
+        )
+
+    joined = await join_results(task, successful, joiner_adapter, joiner_model.id)
+    return FanoutResult(
+        subtasks=subtasks,
+        results=results,
+        joined=joined,
+        join_model=joiner_model.display_name,
+    )
+
+
+# ── Splitter ─────────────────────────────────────────────────────────────────
+
+async def split_task(
+    task: str,
+    plan: str,
+    splitter_adapter: BaseAdapter,
+    splitter_model_id: str,
+) -> list[Subtask]:
+    prompt = (
+        f"ORIGINAL TASK:\n{task}\n\n"
+        f"PLAN (already approved):\n{plan or '(no plan provided)'}\n\n"
+        "Now produce the JSON array of independent subtasks."
+    )
+    try:
+        resp = await splitter_adapter.complete(
+            messages=[Message(role="user", content=prompt)],
+            model_id=splitter_model_id,
+            system=SPLITTER_SYSTEM,
+            temperature=0.2,
+        )
+    except Exception:
+        return []
+    text = resp if isinstance(resp, str) else getattr(resp, "content", "")
+    return _parse_subtasks(text)
+
+
+def _parse_subtasks(text: str) -> list[Subtask]:
+    cleaned = text.strip()
+    fence = re.search(r"```(?:json)?\s*([\s\S]+?)\s*```", cleaned)
+    if fence:
+        cleaned = fence.group(1).strip()
+    bracket_start = cleaned.find("[")
+    bracket_end = cleaned.rfind("]")
+    if bracket_start == -1 or bracket_end == -1:
+        return []
+    try:
+        data = json.loads(cleaned[bracket_start : bracket_end + 1])
+    except Exception:
+        return []
+    if not isinstance(data, list):
+        return []
+    out: list[Subtask] = []
+    for i, item in enumerate(data[:5]):
+        if not isinstance(item, dict):
+            continue
+        sid = str(item.get("id") or f"s{i + 1}")
+        desc = str(item.get("description", "")).strip()
+        if not desc:
+            continue
+        type_str = str(item.get("type", "general")).strip().lower()
+        try:
+            task_type = TaskType(type_str)
+        except ValueError:
+            task_type = TaskType.GENERAL
+        out.append(Subtask(id=sid, description=desc, task_type=task_type))
+    return out
+
+
+# ── Routing ──────────────────────────────────────────────────────────────────
+
+def _route_subtasks(
+    subtasks: list[Subtask],
+    registry: list,
+    cfg: dict,
+    profile: str,
+) -> list[tuple[ModelInfo, BaseAdapter] | None]:
+    from polymath.model_policy import adapter_for
+    from polymath.router import select_primary
+
+    available = [m for m in registry if m.available]
+    routes: list[tuple[ModelInfo, BaseAdapter] | None] = []
+    for subtask in subtasks:
+        pick = select_primary(available, subtask.task_type, profile)
+        if not pick:
+            routes.append(None)
+            continue
+        adapter = adapter_for(pick.provider, cfg)
+        if not adapter:
+            routes.append(None)
+            continue
+        routes.append((pick, adapter))
+    return routes
+
+
+# ── Subtask execution ────────────────────────────────────────────────────────
+
+def _make_subtask_call(
+    adapter: BaseAdapter,
+    model_info: ModelInfo,
+    subtask: Subtask,
+    parent_task: str,
+) -> Callable[[], Awaitable]:
+    prompt = (
+        f"You are working on ONE PIECE of a larger task. Your subtask is below. "
+        f"Be self-contained — assume your output will be combined with sibling subtasks "
+        f"by an integrator. Do not reference other subtasks.\n\n"
+        f"PARENT TASK (for context only):\n{parent_task}\n\n"
+        f"YOUR SUBTASK:\n{subtask.description}"
+    )
+    async def call():
+        return await adapter.complete(
+            messages=[Message(role="user", content=prompt)],
+            model_id=model_info.id,
+            system="Answer the subtask only. Be concrete and concise.",
+            temperature=0.4,
+        )
+    return call
+
+
+# ── Joiner ───────────────────────────────────────────────────────────────────
+
+async def join_results(
+    task: str,
+    successful: list[SubtaskResult],
+    joiner_adapter: BaseAdapter,
+    joiner_model_id: str,
+) -> str:
+    parts = [f"ORIGINAL TASK:\n{task}\n", "SUBTASK OUTPUTS:"]
+    for r in successful:
+        parts.append(
+            f"\n--- {r.subtask.id}: {r.subtask.description} (by {r.model_display}) ---\n"
+            f"{r.output.strip()}"
+        )
+    prompt = "\n".join(parts) + "\n\nNow integrate these into a single coherent answer to the original task."
+    try:
+        resp = await joiner_adapter.complete(
+            messages=[Message(role="user", content=prompt)],
+            model_id=joiner_model_id,
+            system=JOINER_SYSTEM,
+            temperature=0.3,
+        )
+    except Exception as e:
+        # On joiner failure, concatenate raw outputs as a graceful degradation
+        return "\n\n".join(
+            f"### {r.subtask.description}\n{r.output}" for r in successful
+        ) + f"\n\n[joiner failed: {e}]"
+    return resp if isinstance(resp, str) else getattr(resp, "content", "")
+
+
+# ── Selection helpers ────────────────────────────────────────────────────────
+
+def select_splitter_and_joiner(
+    registry: list,
+    cfg: dict,
+    profile: str,
+) -> tuple[tuple[ModelInfo, BaseAdapter] | None, tuple[ModelInfo, BaseAdapter] | None]:
+    """Pick high-quality models for split + join (these benefit from intelligence)."""
+    from polymath.model_policy import adapter_for
+    from polymath.router import select_primary
+
+    available = [m for m in registry if m.available]
+    splitter = select_primary(available, TaskType.ANALYSIS, profile)
+    joiner = select_primary(available, TaskType.ANALYSIS, profile)
+    splitter_pair = None
+    joiner_pair = None
+    if splitter:
+        ad = adapter_for(splitter.provider, cfg)
+        if ad:
+            splitter_pair = (splitter, ad)
+    if joiner:
+        ad = adapter_for(joiner.provider, cfg)
+        if ad:
+            joiner_pair = (joiner, ad)
+    return splitter_pair, joiner_pair

polymath/orchestrator/output_policy.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from polymath.adapters.base import AuthExpiredError, is_rate_limit
+
+
+@dataclass(frozen=True)
+class FailureDecision:
+    kind: str
+    retryable: bool
+    user_message: str
+    detail_message: str = ""
+
+
+class OutputPolicy:
+    def classify_failure(self, exc: Exception) -> FailureDecision:
+        if isinstance(exc, AuthExpiredError):
+            return FailureDecision(
+                kind="auth",
+                retryable=True,
+                user_message=f"[yellow]⚠ Auth expired for [bold]{exc.provider}[/bold] — trying next model...[/yellow]",
+                detail_message=f"[dim]Re-login with:[/dim]  [bold cyan]{exc.relogin_cmd}[/bold cyan]",
+            )
+        if is_rate_limit(exc):
+            return FailureDecision(
+                kind="quota",
+                retryable=True,
+                user_message="[yellow]⚠ Primary model failed due to quota/rate limit — trying next model...[/yellow]",
+            )
+        text = str(exc).lower()
+        # No status code distinguishes this one, so phrases are all there is.
+        # Kept narrow: "context" plus "limit" alone matched any error whose
+        # message happened to mention both words.
+        if any(marker in text for marker in (
+            "context window", "context length", "maximum context",
+            "context_length_exceeded", "too many tokens",
+            "prompt is too long",
+        )):
+            return FailureDecision(
+                kind="context",
+                retryable=True,
+                user_message="[yellow]⚠ Context window exhausted — trying next model...[/yellow]",
+            )
+        return FailureDecision(
+            kind="error",
+            retryable=False,
+            user_message=f"[red]Error: {exc}[/red]",
+        )
+
+    def final_failure_lines(self, decision: FailureDecision, exc: Exception) -> list[str]:
+        if decision.kind == "auth":
+            return [
+                f"[red bold]Auth expired:[/red bold] [yellow]{getattr(exc, 'provider', 'provider')}[/yellow] token is no longer valid.",
+                decision.detail_message,
+                "[dim]No other models available to fall back to.[/dim]",
+            ]
+        return [
+            f"[red bold]Primary model failed:[/red bold] {exc}",
+            "[dim]No other models available to fall back to.[/dim]",
+        ]