polymath-agent 0.4.0__py3-none-any.whl

polymath/execution_service.py

@@ -0,0 +1,325 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Awaitable, Callable
+
+from polymath.adapters.base import BaseAdapter
+from polymath.model_policy import adapter_for, try_fallback_adapter
+from polymath.orchestrator import AttemptLedger, OutputPolicy, WorkerPool
+from polymath.pipeline import Pipeline, quick_answer
+from polymath.router import select_primary
+from polymath.config import TaskType
+from polymath.domain import PermissionRequest
+from polymath.router import select_reviewer
+from polymath.subagents import run_subagent, select_auto_subagents
+
+
+@dataclass
+class ExecutionCallbacks:
+    on_step: Callable[[Any, str], None]
+    on_token: Callable[[str], None]
+    on_ask: Callable[[str], Awaitable[str] | str]
+    on_permission: Callable[[PermissionRequest], Awaitable[bool] | bool]
+    on_plan_approval: Callable[[str], Awaitable[str] | str]
+    on_model_switch: Callable[[Any], None] | None = None
+
+
+@dataclass
+class ExecutionResult:
+    last_output: str = ""
+    usage: dict[str, tuple[int, int, int]] | None = None
+    summary_line: str = ""
+    notify_message: str = ""
+
+
+def _adapter_usage(adapter: BaseAdapter | None, model_info: Any) -> tuple[int, int, int]:
+    if not adapter or not model_info:
+        return (0, 0, 0)
+    return (
+        getattr(adapter, "_last_input_tokens", 0),
+        getattr(adapter, "_last_output_tokens", 0),
+        getattr(model_info, "context_window", 0),
+    )
+
+
+def _resolve_failover(
+    registry: list,
+    primary_info,
+    primary_adapter: BaseAdapter | None,
+    cfg: dict,
+    profile: str,
+    excluded: set[str] | None = None,
+) -> tuple[object | None, BaseAdapter | None]:
+    """Pick the best non-primary available model as a race partner."""
+    if not primary_info:
+        return None, None
+    skip = set(excluded or set())
+    skip.add(primary_info.id)
+    candidates = [m for m in registry if m.available and m.id not in skip]
+    if not candidates:
+        return None, None
+    pick = select_primary(candidates, TaskType.GENERAL, profile)
+    if not pick or pick.id == primary_info.id:
+        return None, None
+    adapter = adapter_for(pick.provider, cfg)
+    if not adapter:
+        return None, None
+    return pick, adapter
+
+
+def _refresh_reviewer(
+    registry: list[Any],
+    primary_info: Any,
+    primary_adapter: BaseAdapter | None,
+    reviewer_info: Any,
+    reviewer_adapter: BaseAdapter | None,
+    cfg: dict,
+    preferred_profile: str,
+) -> tuple[Any, BaseAdapter | None]:
+    if not primary_info:
+        return reviewer_info, reviewer_adapter
+    if reviewer_info and reviewer_info.id != primary_info.id and reviewer_adapter:
+        return reviewer_info, reviewer_adapter
+    next_reviewer = select_reviewer(registry, primary_info, preferred_profile)
+    if not next_reviewer:
+        return primary_info, primary_adapter
+    from polymath.model_policy import adapter_for
+    next_adapter = adapter_for(next_reviewer.provider, cfg)
+    if not next_adapter:
+        return primary_info, primary_adapter
+    return next_reviewer, next_adapter
+
+
+async def run_ask_flow(
+    *,
+    text: str,
+    session_id: str,
+    active_project: str,
+    task_type: Any,
+    primary_info: Any,
+    primary_adapter: Any,
+    summarizer_info: Any,
+    summarizer_adapter: Any,
+    profile: str,
+    cfg: dict,
+    registry: list[Any],
+    print_line: Callable[[str], None],
+    suppress_echo: Any,
+    token_printer: Callable[[str], None],
+    tried_model_ids: set[str] | None = None,
+    on_model_switch: Callable[[Any], None] | None = None,
+    attempt_ledger: AttemptLedger | None = None,
+    output_policy: OutputPolicy | None = None,
+) -> ExecutionResult:
+    policy = output_policy or OutputPolicy()
+    ledger = attempt_ledger or AttemptLedger()
+    attempted = set(tried_model_ids or set())
+    phase = "chat"
+    while primary_info and primary_adapter:
+        attempted.add(primary_info.id)
+        ledger.note_attempt(phase, primary_info.id)
+        try:
+            failover_info, failover_adapter = _resolve_failover(
+                registry, primary_info, primary_adapter, cfg, profile, excluded=attempted,
+            )
+            with suppress_echo():
+                await quick_answer(
+                    session_id=session_id,
+                    question=text,
+                    primary_adapter=primary_adapter,
+                    primary_model=primary_info,
+                    active_project=active_project,
+                    task_type=task_type,
+                    summarizer_adapter=summarizer_adapter,
+                    summarizer_model=summarizer_info.id if summarizer_info else "",
+                    on_token=token_printer,
+                    cfg=cfg,
+                    failover_adapter=failover_adapter,
+                    failover_model=failover_info,
+                )
+            usage = {
+                primary_info.display_name: (
+                    getattr(primary_adapter, "_last_input_tokens", 0),
+                    getattr(primary_adapter, "_last_output_tokens", 0),
+                    primary_info.context_window,
+                )
+            }
+            return ExecutionResult(usage=usage)
+        except Exception as e:
+            decision = policy.classify_failure(e)
+            ledger.note_failure(phase, primary_info.id, decision.kind, retryable=decision.retryable)
+            if not decision.retryable:
+                raise
+            print_line(f"\n{decision.user_message}")
+            fallback = try_fallback_adapter(registry, primary_info, profile, cfg, attempted)
+            if fallback:
+                primary_info, primary_adapter = fallback
+                if on_model_switch:
+                    on_model_switch(primary_info)
+                print_line(f"[dim]Falling back to {primary_info.display_name}[/dim]\n")
+                continue
+            for line in policy.final_failure_lines(decision, e):
+                if line:
+                    print_line(line)
+            return ExecutionResult()
+    return ExecutionResult()
+
+
+async def run_pipeline_flow(
+    *,
+    text: str,
+    complexity: Any,
+    task_type: Any,
+    verify: bool,
+    simplify: bool,
+    session_id: str,
+    active_project: str,
+    primary_info: Any,
+    primary_adapter: Any,
+    reviewer_info: Any,
+    reviewer_adapter: Any,
+    summarizer_info: Any,
+    summarizer_adapter: Any,
+    profile: str,
+    cfg: dict,
+    registry: list[Any],
+    subs: dict[str, Any],
+    callbacks: ExecutionCallbacks,
+    print_line: Callable[[str], None],
+    render_panel: Callable[[str, str], None],
+    suppress_echo: Any,
+    session_cost: Callable[[str], tuple[int, float]],
+    tried_model_ids: set[str] | None = None,
+    attempt_ledger: AttemptLedger | None = None,
+    output_policy: OutputPolicy | None = None,
+) -> ExecutionResult:
+    policy = output_policy or OutputPolicy()
+    ledger = attempt_ledger or AttemptLedger()
+    workers = WorkerPool()
+    attempted = set(tried_model_ids or set())
+    phase = "pipeline"
+    while primary_info and primary_adapter:
+        attempted.add(primary_info.id)
+        ledger.note_attempt(phase, primary_info.id)
+        try:
+            failover_info, failover_adapter = _resolve_failover(
+                registry, primary_info, primary_adapter, cfg, profile, excluded=attempted,
+            )
+            pipeline = Pipeline(
+                session_id=session_id,
+                primary_adapter=primary_adapter,
+                primary_model=primary_info,
+                reviewer_adapter=reviewer_adapter,
+                reviewer_model=reviewer_info,
+                summarizer_adapter=summarizer_adapter,
+                summarizer_model=summarizer_info.id if summarizer_info else "",
+                review_rounds=cfg.get("review_rounds", 2),
+                active_project=active_project,
+                cfg=cfg,
+                failover_adapter=failover_adapter,
+                failover_model=failover_info,
+                on_step=callbacks.on_step,
+                on_token=callbacks.on_token,
+                on_ask=callbacks.on_ask,
+                on_permission=callbacks.on_permission,
+                on_plan_approval=callbacks.on_plan_approval,
+            )
+            with suppress_echo():
+                result = await pipeline.run(text, complexity, task_type, active_project=active_project)
+            last_output = result.output
+            primary_usage = _adapter_usage(primary_adapter, primary_info)
+            reviewer_usage = _adapter_usage(reviewer_adapter, reviewer_info)
+
+            subagent_flags = {"verify": verify, "simplify": simplify}
+            subagent_configs = select_auto_subagents(task_type, subagent_flags, subs)
+            concurrent_subagents = [sa for sa in subagent_configs if sa.name != "verify" and sa.max_iterations <= 1]
+            sequential_subagents = [sa for sa in subagent_configs if sa not in concurrent_subagents]
+            for sa_config in subagent_configs:
+                print_line(f"\n[dim]▸ subagent: {sa_config.name}[/dim]")
+
+            async def _run_single_subagent(sa_config):
+                return sa_config, await run_subagent(
+                    config=sa_config,
+                    pipeline_output=result.output,
+                    primary_adapter=primary_adapter,
+                    primary_model=primary_info,
+                    session_id=session_id,
+                    on_token=lambda _: None,
+                )
+
+            subagent_results = await workers.gather(
+                [lambda sa_config=sa_config: _run_single_subagent(sa_config) for sa_config in concurrent_subagents]
+            )
+            for item in subagent_results:
+                if isinstance(item, Exception):
+                    print_line(f"[red]Subagent failed: {item}[/red]")
+                    continue
+                sa_config, sa_result = item
+                if sa_result.output:
+                    render_panel(sa_result.output, sa_config.name)
+                if sa_result.iterations > 1:
+                    status = "passed" if sa_result.success else "failed"
+                    print_line(f"[dim]{sa_config.name}: {sa_result.iterations} iterations — {status}[/dim]")
+                last_output = sa_result.output or last_output
+
+            for sa_config in sequential_subagents:
+                try:
+                    _, sa_result = await _run_single_subagent(sa_config)
+                    if sa_result.output:
+                        render_panel(sa_result.output, sa_config.name)
+                    if sa_result.iterations > 1:
+                        status = "passed" if sa_result.success else "failed"
+                        print_line(f"[dim]{sa_config.name}: {sa_result.iterations} iterations — {status}[/dim]")
+                    last_output = sa_result.output or last_output
+                except Exception as e:
+                    print_line(f"[red]Subagent {sa_config.name} failed: {e}[/red]")
+
+            usage: dict[str, tuple[int, int, int]] = {
+                primary_info.display_name: primary_usage
+            }
+            if reviewer_adapter and reviewer_info and reviewer_info.id != primary_info.id:
+                usage[reviewer_info.display_name] = reviewer_usage
+            tokens, _ = session_cost(session_id)
+            token_str = f"  | ~{tokens} tokens" if tokens > 0 else ""
+            summary_line = (
+                f"[dim]steps: {' → '.join(result.steps_run)}  "
+                f"| primary: {result.primary_model}"
+                + (f"  | reviewer: {result.reviewer_model}" if result.reviewer_model else "")
+                + token_str
+                + "[/dim]"
+            )
+            return ExecutionResult(
+                last_output=last_output,
+                usage=usage,
+                summary_line=summary_line,
+                notify_message=f"Done — {result.primary_model}",
+            )
+        except Exception as e:
+            decision = policy.classify_failure(e)
+            ledger.note_failure(phase, primary_info.id, decision.kind, retryable=decision.retryable)
+            if not decision.retryable:
+                raise
+            print_line(f"\n{decision.user_message}")
+            fallback = try_fallback_adapter(registry, primary_info, profile, cfg, attempted)
+            if fallback:
+                primary_info, primary_adapter = fallback
+                reviewer_info, reviewer_adapter = _refresh_reviewer(
+                    registry,
+                    primary_info,
+                    primary_adapter,
+                    reviewer_info,
+                    reviewer_adapter,
+                    cfg,
+                    profile,
+                )
+                if callbacks.on_model_switch:
+                    callbacks.on_model_switch(primary_info)
+                print_line(f"[dim]Falling back to {primary_info.display_name}[/dim]\n")
+                continue
+            for line in policy.final_failure_lines(decision, e):
+                if line:
+                    print_line(line)
+            return ExecutionResult()
+        finally:
+            await workers.cancel_all()
+    return ExecutionResult()