diff --git a/README.md b/README.md index e32d7c5a..6bc5e607 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,9 @@ In benchmark `task` mode, the planning layer (1) is bypassed so workflow synthes - [Installation](#installation) - [Configuration](#configuration) - [Running Mimosa](#running-mimosa) + - [Interactive Onboarding (recommended for first-time setup)](#interactive-onboarding-recommended-for-first-time-setup) + - [Goal mode — multi-step scientific objective](#goal-mode--multi-step-scientific-objective) + - [Task mode — single granular operation](#task-mode--single-granular-operation) - [Workspace and Audit Trail](#workspace-and-audit-trail) - [Learning through Evolution of Multi-Agent Workflows](#learning-through-evolution-of-multi-agent-workflows) - [Transparency](#transparency) @@ -184,7 +187,25 @@ Custom MCP tools can be added via the [Toolomics docs](https://github.com/Holobi --- -## Configuration + +## Running Mimosa + +### Interactive Onboarding (recommended for first-time setup) + +> **If you are new to Mimosa, start here.** + +Running Mimosa with **no arguments** launches an interactive, step-by-step onboarding wizard that guides you through everything before the first execution: + +```bash +uv run main.py +``` +Once you complete setup once, subsequent runs remember your workspace path via `config_default.json` — no re-configuration needed. + +--- + +### Manual onboarding: + +**1. Start by editing the config:** ```bash cp config_default.json my_config.json @@ -204,13 +225,9 @@ Edit `my_config.json`. Key parameters: | `learned_score_threshold` | Minimum score to accept a result and stop iterating | | `max_learning_evolve_iterations` | Maximum self-improvement iterations before accepting the result | ---- - -## Running Mimosa - -Mimosa supports two execution modes: **Goal** and **Task**. +**2. Choose a mode `task` or `goal` depending on the complexity of your objective.** -### Goal mode — multi-step scientific objective +**2.1 Goal mode — multi-step scientific objective** Use this when your objective requires planning across multiple distinct operations (e.g., reproducing a paper, building an ML pipeline). @@ -229,7 +246,7 @@ uv run main.py \ --config my_config.json ``` -### Task mode — single granular operation +**2.2 Task mode — single granular operation** Use this for a focused, self-contained operation without long-term planning. @@ -257,7 +274,7 @@ uv run main.py --task "Conduct a literature review on graph neural networks for During execution, Mimosa reads and writes files inside the Toolomics workspace configured by `workspace_dir`. When a run finishes, the workspace contents are copied into a timestamped folder under `runs_capsule/` so the final state is preserved as an archive. - **Toolomics `workspace/`** — live working directory: intermediate files, scripts, downloads, generated outputs -- **`sources/workflows//`** — generated workflow and execution metadata: `state_result.json`, `evaluation.txt`, `reward_progress.png`, `memory/` traces +- **`sources/workflows//`** — generated workflow and execution metadata: `state_result.json`, `evaluation.txt`, `reward_progress.png` - **`runs_capsule//`** — archived snapshot of the run for later inspection, comparison, or sharing - **`memory_explorer.py `** — replay a workflow execution step-by-step to inspect agent traces, tool calls, and outputs diff --git a/config.py b/config.py index 5cd5536e..79f4f760 100755 --- a/config.py +++ b/config.py @@ -85,6 +85,7 @@ def __init__(self): self.runner_default_max_cpu_percent: int = 100 self.runner_temp_dir: str = "./tmp" self.runner_requirements: list[str] = [ + "setuptools>=70.0", "python-dotenv", "fastmcp==2.8.1", "requests>=2.31.0", diff --git a/main.py b/main.py index 1f164fd0..64e6363a 100755 --- a/main.py +++ b/main.py @@ -15,10 +15,13 @@ # Prevent tokenizers parallelism warnings when forking processes os.environ["TOKENIZERS_PARALLELISM"] = "false" +from sources.cli.pretty_print import print_ok, print_warn, print_err, print_info + from config import Config from sources.core.dgm import DarwinMachine from sources.core.planner import Planner from sources.extensibility.human_mode import HumanMode +from sources.cli import OnboardCLI from sources.evaluation.csv_mode import CsvEvaluationMode from sources.evaluation.scenario_loader import ScenarioLoader from sources.evaluation.eval_workflow_generation import WorkflowEval @@ -35,11 +38,11 @@ def validate_environment() -> None: envs_key = ['ANTHROPIC_API_KEY', 'MISTRAL_API_KEY', 'DEEPSEEK_API_KEY', 'OPENAI_API_KEY', 'HF_TOKEN', 'OPENROUTER_API_KEY'] for key in envs_key: if os.getenv(key): - print(f"✅ Found environment variable: {key}") + print_ok(f"Found environment variable: {key}") key_found = True if not key_found: raise ValueError( - "⚠️ No valid API key environment variable found. Please set one of the supported API keys. Supported keys: " + ", ".join(envs_key) + "No valid API key environment variable found. Please set one of the supported API keys. Supported keys: " + ", ".join(envs_key) ) def add_config_arguments(parser: argparse.ArgumentParser, config: Config) -> None: @@ -86,7 +89,7 @@ def apply_config_overrides(args: argparse.Namespace, config: Config) -> None: def setup_signal_handlers(): """Setup signal handlers for graceful shutdown.""" def signal_handler(signum, frame): - print(f"\n⚠️ Received signal {signum}. Shutting down gracefully...") + print_warn(f"Received signal {signum}. Shutting down gracefully…") for task in asyncio.all_tasks(): if not task.done(): task.cancel() @@ -110,7 +113,7 @@ async def papers_mode(args, config): async def science_bench_papers_mode(args, config): papers = CsvEvaluationMode(config, csv_runs_limit=args.csv_runs_limit) if args.single_agent: - print(f"⚠️ Starting in single agent mode") + print_info("Starting in single agent mode") await papers.start_evaluation(dataset_type="science_agent_bench", dataset_path="datasets/ScienceAgentBench.csv", learning=args.learn, @@ -138,11 +141,11 @@ def load_goal_from_file_or_string(goal_input: str) -> str: try: with open(goal_input, 'r', encoding='utf-8') as f: content = f.read().strip() - print(f"✅ Loaded goal from file: {goal_input}") + print_ok(f"Loaded goal from file: {goal_input}") return content except Exception as e: - print(f"⚠️ Failed to read file '{goal_input}': {e}") - print(f"Using input as a literal string instead.") + print_warn(f"Failed to read file '{goal_input}': {e}") + print_info("Using input as a literal string instead.") return goal_input return goal_input @@ -156,7 +159,7 @@ async def normal_execution_mode(args, config): # Load goal from file if args.task is a file path goal_content = load_goal_from_file_or_string(args.task) if args.single_agent: - print(f"⚠️ Starting in single agent mode") + print_info("Starting in single agent mode") await dgm.start_dgm(goal=goal_content, judge=not args.disable_judge, scenario_rubric=args.scenario, @@ -223,7 +226,10 @@ async def main(): "--scenario", type=str, help="Use scenario benchmark (eg: datasets/scenarios/X.json) with criterions for workflow evaluation and auto-improvement" ) parser.add_argument( - "--debug", action="store_true", help="Enable debug logging to console" + "--debug", action="store_true", help="Enable advanced debug logging to console" + ) + parser.add_argument( + "--verbose", action="store_true", help="Enable verbose logging to console" ) parser.add_argument( "--max_evolve_iterations", type=int, default=1, help="Maximum number of learning iterations. Used for retrying/learning a task." @@ -240,8 +246,29 @@ async def main(): # security check PackageCheck().run() # Setup logging with debug flag - setup_logging(debug=args.debug) + setup_logging(debug=args.debug, disable=not args.verbose) + + # Detect interactive (no-argument) mode early so we can skip pre-checks + no_mode_selected = not any([ + args.manual, + args.papers, + args.science_agent_bench, + args.task, + args.goal, + args.scenario, + args.workflow_eval_mode, + ]) + + if no_mode_selected: + # Interactive onboarding CLI for full setup flow. + try: + cli = OnboardCLI(config) + await cli.run() + except KeyboardInterrupt: + print("\n\n Interrupted. Goodbye!\n") + return + # ── Normal (argument-driven) execution path ─────────────────────────── # Apply CLI argument overrides (these override config file values) apply_config_overrides(args, config) @@ -263,8 +290,6 @@ async def main(): await normal_execution_mode(args, config) elif args.workflow_eval_mode: await workflow_generation_evals(args, config) - else: - raise ValueError("No goal provided. Use --task, --goal, --papers to start.") except KeyboardInterrupt: raise except Exception as e: diff --git a/sources/cache/openrouter_pricing.json b/sources/cache/openrouter_pricing.json index f0bde056..58a21ccf 100644 --- a/sources/cache/openrouter_pricing.json +++ b/sources/cache/openrouter_pricing.json @@ -1,9 +1,65 @@ { - "timestamp": "2026-03-26T14:56:47.062467", + "timestamp": "2026-04-08T15:44:21.411840", "pricing": { - "reka/reka-edge": { - "input": 0.19999999999999998, - "output": 0.19999999999999998 + "anthropic/claude-opus-4.6-fast": { + "input": 30.0, + "output": 150.0 + }, + "z-ai/glm-5.1": { + "input": 1.26, + "output": 3.9600000000000004 + }, + "google/gemma-4-26b-a4b-it:free": { + "input": 0.0, + "output": 0.0 + }, + "google/gemma-4-26b-a4b-it": { + "input": 0.13, + "output": 0.39999999999999997 + }, + "google/gemma-4-31b-it:free": { + "input": 0.0, + "output": 0.0 + }, + "google/gemma-4-31b-it": { + "input": 0.14, + "output": 0.39999999999999997 + }, + "qwen/qwen3.6-plus": { + "input": 0.325, + "output": 1.95 + }, + "z-ai/glm-5v-turbo": { + "input": 1.2, + "output": 4.0 + }, + "arcee-ai/trinity-large-thinking": { + "input": 0.22, + "output": 0.85 + }, + "x-ai/grok-4.20-multi-agent": { + "input": 2.0, + "output": 6.0 + }, + "x-ai/grok-4.20": { + "input": 2.0, + "output": 6.0 + }, + "google/lyria-3-pro-preview": { + "input": 0.0, + "output": 0.0 + }, + "google/lyria-3-clip-preview": { + "input": 0.0, + "output": 0.0 + }, + "kwaipilot/kat-coder-pro-v2": { + "input": 0.3, + "output": 1.2 + }, + "rekaai/reka-edge": { + "input": 0.09999999999999999, + "output": 0.09999999999999999 }, "xiaomi/mimo-v2-omni": { "input": 0.39999999999999997, @@ -33,14 +89,6 @@ "input": 1.2, "output": 4.0 }, - "x-ai/grok-4.20-multi-agent-beta": { - "input": 2.0, - "output": 6.0 - }, - "x-ai/grok-4.20-beta": { - "input": 2.0, - "output": 6.0 - }, "nvidia/nemotron-3-super-120b-a12b:free": { "input": 0.0, "output": 0.0 @@ -138,8 +186,8 @@ "output": 0.0 }, "minimax/minimax-m2.5": { - "input": 0.19999999999999998, - "output": 1.17 + "input": 0.118, + "output": 0.9900000000000001 }, "z-ai/glm-5": { "input": 0.72, @@ -174,8 +222,8 @@ "output": 0.0 }, "moonshotai/kimi-k2.5": { - "input": 0.44999999999999996, - "output": 2.2 + "input": 0.3827, + "output": 1.72 }, "upstage/solar-pro-3": { "input": 0.15, @@ -241,10 +289,6 @@ "input": 0.09999999999999999, "output": 0.3 }, - "allenai/olmo-3.1-32b-think": { - "input": 0.15, - "output": 0.5 - }, "xiaomi/mimo-v2-flash": { "input": 0.09, "output": 0.29 @@ -373,13 +417,9 @@ "input": 0.25, "output": 2.0 }, - "kwaipilot/kat-coder-pro": { - "input": 0.207, - "output": 0.828 - }, "moonshotai/kimi-k2-thinking": { - "input": 0.47, - "output": 2.0 + "input": 0.6, + "output": 2.5 }, "amazon/nova-premier-v1": { "input": 2.5, @@ -413,14 +453,6 @@ "input": 0.10400000000000001, "output": 0.41600000000000004 }, - "liquid/lfm2-8b-a1b": { - "input": 0.01, - "output": 0.02 - }, - "liquid/lfm-2.2-6b": { - "input": 0.01, - "output": 0.02 - }, "ibm-granite/granite-4.0-h-micro": { "input": 0.017, "output": 0.11 @@ -841,10 +873,6 @@ "input": 0.25, "output": 0.75 }, - "qwen/qwen3-4b:free": { - "input": 0.0, - "output": 0.0 - }, "meta-llama/llama-guard-4-12b": { "input": 0.18, "output": 0.18 @@ -937,10 +965,6 @@ "input": 150.0, "output": 600.0 }, - "mistralai/mistral-small-3.1-24b-instruct:free": { - "input": 0.0, - "output": 0.0 - }, "mistralai/mistral-small-3.1-24b-instruct": { "input": 0.03, "output": 0.11 @@ -977,6 +1001,10 @@ "input": 2.5, "output": 10.0 }, + "rekaai/reka-flash-3": { + "input": 0.09999999999999999, + "output": 0.19999999999999998 + }, "google/gemma-3-27b-it:free": { "input": 0.0, "output": 0.0 @@ -1169,10 +1197,6 @@ "input": 0.7999999999999999, "output": 4.0 }, - "anthropic/claude-3.5-sonnet": { - "input": 6.0, - "output": 30.0 - }, "anthracite-org/magnum-v4-72b": { "input": 3.0, "output": 5.0 @@ -1197,10 +1221,6 @@ "input": 0.16999999999999998, "output": 0.43 }, - "meta-llama/llama-3.2-1b-instruct": { - "input": 0.027, - "output": 0.19999999999999998 - }, "meta-llama/llama-3.2-3b-instruct:free": { "input": 0.0, "output": 0.0 @@ -1209,6 +1229,10 @@ "input": 0.051, "output": 0.33999999999999997 }, + "meta-llama/llama-3.2-1b-instruct": { + "input": 0.027, + "output": 0.19999999999999998 + }, "meta-llama/llama-3.2-11b-vision-instruct": { "input": 0.049, "output": 0.049 @@ -1217,14 +1241,14 @@ "input": 0.12, "output": 0.39 }, - "cohere/command-r-08-2024": { - "input": 0.15, - "output": 0.6 - }, "cohere/command-r-plus-08-2024": { "input": 2.5, "output": 10.0 }, + "cohere/command-r-08-2024": { + "input": 0.15, + "output": 0.6 + }, "sao10k/l3.1-euryale-70b": { "input": 0.85, "output": 0.85 @@ -1261,11 +1285,11 @@ "input": 0.02, "output": 0.04 }, - "openai/gpt-4o-mini": { + "openai/gpt-4o-mini-2024-07-18": { "input": 0.15, "output": 0.6 }, - "openai/gpt-4o-mini-2024-07-18": { + "openai/gpt-4o-mini": { "input": 0.15, "output": 0.6 }, @@ -1285,6 +1309,10 @@ "input": 0.14, "output": 0.14 }, + "openai/gpt-4o-2024-05-13": { + "input": 5.0, + "output": 15.0 + }, "openai/gpt-4o": { "input": 2.5, "output": 10.0 @@ -1293,10 +1321,6 @@ "input": 6.0, "output": 18.0 }, - "openai/gpt-4o-2024-05-13": { - "input": 5.0, - "output": 15.0 - }, "meta-llama/llama-3-8b-instruct": { "input": 0.03, "output": 0.04 @@ -1325,14 +1349,14 @@ "input": 2.0, "output": 6.0 }, - "openai/gpt-3.5-turbo-0613": { - "input": 1.0, - "output": 2.0 - }, "openai/gpt-4-turbo-preview": { "input": 10.0, "output": 30.0 }, + "openai/gpt-3.5-turbo-0613": { + "input": 1.0, + "output": 2.0 + }, "mistralai/mixtral-8x7b-instruct": { "input": 0.54, "output": 0.54 @@ -1349,14 +1373,14 @@ "input": 10.0, "output": 30.0 }, - "mistralai/mistral-7b-instruct-v0.1": { - "input": 0.11, - "output": 0.19 - }, "openai/gpt-3.5-turbo-instruct": { "input": 1.5, "output": 2.0 }, + "mistralai/mistral-7b-instruct-v0.1": { + "input": 0.11, + "output": 0.19 + }, "openai/gpt-3.5-turbo-16k": { "input": 3.0, "output": 4.0 @@ -1373,17 +1397,17 @@ "input": 0.06, "output": 0.06 }, - "openai/gpt-3.5-turbo": { - "input": 0.5, - "output": 1.5 - }, - "openai/gpt-4": { + "openai/gpt-4-0314": { "input": 30.0, "output": 60.0 }, - "openai/gpt-4-0314": { + "openai/gpt-4": { "input": 30.0, "output": 60.0 + }, + "openai/gpt-3.5-turbo": { + "input": 0.5, + "output": 1.5 } } } \ No newline at end of file diff --git a/sources/cli/__init__.py b/sources/cli/__init__.py new file mode 100644 index 00000000..cd4939df --- /dev/null +++ b/sources/cli/__init__.py @@ -0,0 +1,40 @@ +""" +Interactive onboarding CLI for Mimosa-AI. +Guides users through setup and launches the appropriate execution mode. +""" + +from .onboard_cli import OnboardCLI +from .pretty_print import ( + print_ok, + print_warn, + print_err, + print_info, + print_step, + print_phase, + print_section, + print_rule, + print_iteration_header, + print_box, + print_kv_row, + print_summary, + print_agent_answers, + CYAN, GREEN, YELLOW, RED, BLUE, MAGENTA, BOLD, DIM, RESET, +) + +__all__ = [ + "OnboardCLI", + "print_ok", + "print_warn", + "print_err", + "print_info", + "print_step", + "print_phase", + "print_section", + "print_rule", + "print_iteration_header", + "print_box", + "print_kv_row", + "print_summary", + "print_agent_answers", + "CYAN", "GREEN", "YELLOW", "RED", "BLUE", "MAGENTA", "BOLD", "DIM", "RESET", +] diff --git a/sources/cli/onboard_cli.py b/sources/cli/onboard_cli.py new file mode 100644 index 00000000..d95d1916 --- /dev/null +++ b/sources/cli/onboard_cli.py @@ -0,0 +1,1032 @@ +""" +Interactive onboarding CLI for Mimosa-AI. + +Guides new users through setup step-by-step (Claude-code style), +checks that Toolomics is online, clarifies and refines the user's +objective using an LLM conversation loop, classifies it as Goal-mode +or Task-mode, then hands off to the appropriate execution entry-point +(planner.start_planner or dgm.start_dgm). +""" + +from __future__ import annotations + +import asyncio +import json +import time +import os +import sys +import textwrap +from typing import Literal + +from config import Config +from sources.core.llm_provider import LLMConfig, LLMProvider, extract_model_pattern +from sources.core.tools_manager import ToolManager +from sources.utils.list_files import list_files +from sources.utils.transfer_toolomics import LocalTransfer + + +# --------------------------------------------------------------------------- +# Terminal helpers +# --------------------------------------------------------------------------- + +CYAN = "\033[96m" +GREEN = "\033[92m" +YELLOW = "\033[93m" +RED = "\033[91m" +BOLD = "\033[1m" +DIM = "\033[2m" +RESET = "\033[0m" + +MIMOSA_BANNER = f""" +{CYAN}{BOLD} + ███╗ ███╗██╗███╗ ███╗ ██████╗ ███████╗ █████╗ + ████╗ ████║██║████╗ ████║██╔═══██╗██╔════╝██╔══██╗ + ██╔████╔██║██║██╔████╔██║██║ ██║███████╗███████║ + ██║╚██╔╝██║██║██║╚██╔╝██║██║ ██║╚════██║██╔══██║ + ██║ ╚═╝ ██║██║██║ ╚═╝ ██║╚██████╔╝███████║██║ ██║ + ╚═╝ ╚═╝╚═╝╚═╝ ╚═╝ ╚═════╝ ╚══════╝╚═╝ ╚═╝ +{RESET} +{DIM} Self-evolving AI Framework for Autonomous Scientific Research{RESET} +""" + +MIMOSA_START_BANNER = f""" +{GREEN}{BOLD} + ╔══════════════════════════════════════════════════════════════╗ + ║ ║ + ║ 🌱 M I M O S A — S T A R T I N G U P 🌱 ║ + ║ ║ + ╚══════════════════════════════════════════════════════════════╝ +{RESET} +""" + +TOTAL_STEPS = 9 + +# --------------------------------------------------------------------------- +# Model presets — ordered by quality/preference +# (env_key, display_label, litellm_model_id) +# --------------------------------------------------------------------------- +_MODEL_PRESETS: list[tuple[str, str, str]] = [ + ("ANTHROPIC_API_KEY", "Claude Sonnet 4.5 (Anthropic)", "anthropic/claude-sonnet-4-5"), + ("DEEPSEEK_API_KEY", "DeepSeek Chat (DeepSeek)", "deepseek/deepseek-chat"), + ("OPENROUTER_API_KEY", "GLM-5 via OpenRouter (z-ai)", "openrouter/z-ai/glm-5"), + ("OPENAI_API_KEY", "GPT-4o (OpenAI)", "openai/gpt-4o"), + ("MISTRAL_API_KEY", "Mistral Large (Mistral)", "mistral/mistral-large-latest"), +] +# Config keys that all share the same "main" LLM selection +_MODEL_CFG_KEYS = [ + "planner_llm_model", + "prompts_llm_model", + "workflow_llm_model", + "judge_model", +] + + +def _print_step(step: int, total: int, title: str, no_count: bool = False) -> None: + bar = "─" * 60 + print(f"\n{CYAN}{bar}{RESET}") + if not no_count: + print(f"{CYAN} Step {step}/{total} · {title}{RESET}") + else: + print(f"{CYAN} {title}{RESET}") + print(f"{CYAN}{bar}{RESET}") + + +def _ok(msg: str) -> None: + print(f"{GREEN} ✅ {msg}{RESET}") + + +def _warn(msg: str) -> None: + print(f"{YELLOW} ⚠️ {msg}{RESET}") + + +def _err(msg: str) -> None: + print(f"{RED} ❌ {msg}{RESET}") + + +def _info(msg: str) -> None: + print(f"{DIM} ℹ️ {msg}{RESET}") + + +def _ask(prompt: str, default: str = "") -> str: + """Print a prompt and return stripped user input. Empty → *default*.""" + suffix = f" [{default}]" if default else "" + try: + answer = input(f"\n{BOLD} ➤ {prompt}{suffix}: {RESET}").strip() + except (EOFError, KeyboardInterrupt): + print() + sys.exit(0) + return answer if answer else default + + +def _ask_yn(prompt: str, default: bool = True) -> bool: + """Ask a yes/no question and return a boolean.""" + hint = "Y/n" if default else "y/N" + raw = _ask(f"{prompt} ({hint})", default="y" if default else "n").lower() + return raw in ("y", "yes", "1", "true") + + +def _wrap(text: str, width: int = 72, indent: int = 4) -> str: + return textwrap.fill(text, width=width, initial_indent=" " * indent, + subsequent_indent=" " * indent) + + +def _build_llm(config: Config, temperature: float = 0.0, + max_tokens: int = 512) -> LLMProvider: + """Build a lightweight LLMProvider from the planner model config.""" + provider, model = extract_model_pattern(config.planner_llm_model) + llm_config = LLMConfig( + model=model, + provider=provider, + temperature=temperature, + reasoning_effort="low", + max_tokens=max_tokens, + ) + return LLMProvider( + agent_name=None, + memory_path=None, + system_msg=None, # system msg set per call below + config=llm_config, + ) + + +def _call_llm(llm: LLMProvider, system: str, user: str) -> str: + """Override the provider's system message and call it.""" + llm.sys_msg = system + return llm(user, use_cache=False) + + +def _parse_json_response(raw: str) -> dict: + """Strip markdown fences and parse JSON.""" + raw = raw.strip() + if raw.startswith("```"): + raw = "\n".join( + line for line in raw.splitlines() + if not line.strip().startswith("```") + ).strip() + return json.loads(raw) + + +# --------------------------------------------------------------------------- +# LLM prompts +# --------------------------------------------------------------------------- + +_CLARIFIER_SYSTEM = """\ +You are an expert scientific research assistant helping users formulate \ +their research objectives clearly for the Mimosa-AI autonomous research framework. + +Given the user's current objective (and any additional context they provided), decide: +1. Is the objective sufficiently clear and actionable for an AI to execute autonomously? +2. If NOT clear: identify the single most important missing piece of information and \ + formulate one concise clarifying question. +3. If CLEAR: produce a polished, detailed, self-contained restatement that an AI agent \ + can act on directly (include dataset names, metrics, file paths, or any specifics \ + already mentioned). + +Return ONLY valid JSON (no markdown fences) in this exact shape: +{ + "is_clear": true | false, + "question": "", + "refined_prompt": "" +} + +Rules: +- Ask at most ONE question per turn. +- Only mark is_clear=true when you have enough detail to write a rich refined_prompt. +- The refined_prompt must incorporate ALL context provided so far. +- Do not ask for information that is not strictly necessary for execution. +""" + +_CLASSIFIER_SYSTEM = """\ +You are an expert assistant for the Mimosa-AI scientific research framework. +Your job is to classify a user's research objective into one of two execution modes. + +MODES: +• "task" — A single, focused, self-contained operation that does NOT require multi-step + planning. Examples: training a model on one dataset, running a literature review, + producing a specific figure, downloading/processing a file. + +• "goal" — A high-level, multi-step scientific objective that benefits from autonomous + decomposition into sub-tasks before execution. Examples: reproducing a full paper, + building an end-to-end ML pipeline from scratch, running a complete bioinformatics + analysis across several tools. + +Return ONLY valid JSON (no markdown fences) like: +{ + "mode": "task" | "goal", + "confidence": 0.0-1.0, + "reasoning": "", + "suggested_label": "" +} +""" + + +# --------------------------------------------------------------------------- +# Main onboarding class +# --------------------------------------------------------------------------- + +ModeType = Literal["task", "goal"] + + +class OnboardCLI: + """Interactive setup wizard that guides the user through Mimosa-AI setup.""" + + def __init__(self, config: Config) -> None: + self.config = config + self._objective: str = "" + self._mode: ModeType = "task" + self._learn: bool = False + + # ------------------------------------------------------------------ + # Public entry-point + # ------------------------------------------------------------------ + + async def run(self) -> None: + """Run the full onboarding flow, then launch the selected mode.""" + print(MIMOSA_BANNER) + print(_wrap( + "Welcome to Mimosa-AI!" + "Press Ctrl-C at any time to quit.", + width=70, indent=2, + )) + + # Step 1 – API keys + _print_step(1, TOTAL_STEPS, "API Key Check") + self._check_api_keys() + + # Step 2 – Config file + _print_step(2, TOTAL_STEPS, "Configuration") + self._load_config() + + # Step 3 – LLM model selection + _print_step(3, TOTAL_STEPS, "LLM Model Selection") + self._choose_models() + + # Step 4 – Toolomics / MCP connectivity (loops until online or skipped) + _print_step(4, TOTAL_STEPS, "Toolomics MCP Connectivity") + await self._check_toolomics() + + # Step 5 – Workspace file setup (select/clean files, optionally import) + _print_step(5, TOTAL_STEPS, "Workspace Setup") + self._setup_workspace_files() + + first_pass = True + while True: + # Infine loop for conversation to continue + # Step 6 – Initial objective + step_6_text = "Your Research Objective" if first_pass else "Keep working on the same objective" + _print_step(6, TOTAL_STEPS, "Your Research Objective", no_count=not first_pass) + self._collect_objective() + + # Step 7 – LLM clarification + prompt refinement loop + _print_step(7, TOTAL_STEPS, "Objective Clarification & Refinement", no_count=not first_pass) + self._clarify_and_refine() + + # Step 8 – Mode classification + _print_step(8, TOTAL_STEPS, "Mode Selection (Goal vs Task)", no_count=not first_pass) + self._classify_and_confirm() + + # Step 9 – Extra options then launch + _print_step(9, TOTAL_STEPS, "Options & Launch", no_count=not first_pass) + self._collect_options() + + await self._launch() + first_pass = False + + # ------------------------------------------------------------------ + # Step implementations + # ------------------------------------------------------------------ + + def _check_api_keys(self) -> None: + """Check for at least one known LLM API key in the environment.""" + known_keys = [ + "ANTHROPIC_API_KEY", + "OPENAI_API_KEY", + "DEEPSEEK_API_KEY", + "MISTRAL_API_KEY", + "HF_TOKEN", + "OPENROUTER_API_KEY", + ] + found = [k for k in known_keys if os.getenv(k)] + + if found: + for k in found: + _ok(f"Found {k}") + return + + _warn("No LLM API key found in environment.") + print(_wrap( + "Mimosa needs at least one API key to call an LLM. " + "Supported variables: " + ", ".join(known_keys), + width=70, indent=2, + )) + print() + for k in known_keys: + value = _ask(f"Enter {k} (leave blank to skip)") + if value: + os.environ[k] = value + _ok(f"{k} set for this session.") + break + else: + _err("No API key provided. Mimosa cannot run without one.") + sys.exit(1) + + def _load_config(self) -> None: + """Optionally load a JSON config file. + + When the user leaves the path blank, *config_default.json* is loaded + automatically (if it exists) so that any previously saved settings — + including the Toolomics workspace_dir — are picked up without asking. + """ + _info( + "A config file lets you override LLM models, workspace paths, " + "port ranges, etc. (see config_default.json for reference)." + ) + path = _ask( + f"Path to config file (leave blank to auto-load {self._CONFIG_DEFAULT_PATH})" + ) + if path: + if not os.path.isfile(path): + _warn(f"File not found: {path}. Using default configuration.") + else: + try: + self.config.load(path) + _ok(f"Configuration loaded from {path}") + except Exception as exc: + _warn(f"Failed to load config ({exc}). Using defaults.") + else: + # Auto-load config_default.json if it exists + if os.path.isfile(self._CONFIG_DEFAULT_PATH): + try: + self.config.load(self._CONFIG_DEFAULT_PATH) + _ok(f"Loaded {self._CONFIG_DEFAULT_PATH} (workspace: {self.config.workspace_dir})") + except Exception as exc: + _warn(f"Failed to load {self._CONFIG_DEFAULT_PATH} ({exc}). Using built-in defaults.") + else: + _info("Using built-in default configuration.") + + # Ensure internal directories exist before later steps need them + self.config.create_paths() + + async def _check_toolomics(self) -> None: + """Discover MCP servers; loop until at least one is found or user skips.""" + print(_wrap( + "Mimosa requires Toolomics (the companion MCP server) to be running " + "before execution. Scanning your configured discovery addresses …", + width=70, indent=2, + )) + + tool_manager = ToolManager(config=self.config) + + while True: + mcps = await self._discover_once(tool_manager) + + if mcps: + tool_manager.mcps = mcps + for mcp in mcps: + _ok(f"MCP server online: {mcp}") + bash_ok = await tool_manager.verify_tools() + if not bash_ok: + _warn( + "No 'execute_command' tool found.\n" + "Make sure the shell MCP is deployed in Toolomics.\n" + "Retrying soon..." + ) + time.sleep(15) + continue + else: + _ok("Shell tool (execute_command) is available.") + + # ── Workspace directory check ────────────────────────── + self._verify_workspace_dir() + return # ← success, exit loop + + # No MCPs found — ask the user what to do + _err("No MCP/Toolomics servers found.") + print(_wrap( + "Please start Toolomics on the configured port range " + f"({self.config.discovery_addresses}).", + width=70, indent=2, + )) + print(f"\n {BOLD}Options:{RESET}") + print(f" {CYAN}Enter{RESET} – retry scan") + print(f" {CYAN}skip{RESET} – continue without Toolomics " + f"(execution will fail later)") + choice = _ask("Retry or skip?").lower() + if choice == "skip": + _warn("Skipping Toolomics check. Execution may fail at runtime.") + return + # Any other input (including blank/Enter) → retry + + async def _discover_once(self, tool_manager: ToolManager) -> list: + """Run a single MCP discovery pass, returning the list (may be empty).""" + try: + return await tool_manager.discover_mcp_servers() + except Exception as exc: + _warn(f"Discovery error: {exc}") + return [] + + _CONFIG_DEFAULT_PATH = "config_default.json" + + def _verify_workspace_dir(self) -> None: + """Check that config.workspace_dir exists; prompt the user until it does. + + When the user supplies a valid path it is written back to + *config_default.json* so that subsequent runs don't ask again. + """ + while True: + workspace = self.config.workspace_dir + if os.path.isdir(workspace): + _ok(f"Workspace directory found: {workspace}") + return + + _err(f"Workspace directory not found: {workspace}") + print(_wrap( + "This path must point to the Toolomics workspace folder — the shared " + "directory where Mimosa reads and writes task artifacts. " + "Please enter the correct absolute path, or press Enter to skip.", + width=70, indent=2, + )) + new_path = _ask("Workspace directory path (Enter to skip)") + if not new_path: + _warn( + "Skipping workspace check. " + "Execution will fail unless workspace_dir is set correctly." + ) + return + new_path = os.path.expanduser(new_path.strip()) + if os.path.isdir(new_path): + self.config.workspace_dir = new_path + _ok(f"Workspace directory set to: {new_path}") + self._persist_workspace_dir(new_path) + return + _err(f"Directory does not exist: {new_path}. Please try again.") + + def _persist_workspace_dir(self, path: str) -> None: + """Write *path* as workspace_dir into config_default.json.""" + cfg_path = self._CONFIG_DEFAULT_PATH + try: + # Read existing config (or start from empty dict) + if os.path.isfile(cfg_path): + with open(cfg_path, encoding="utf-8") as fh: + data = json.load(fh) + else: + data = {} + + data["workspace_dir"] = path + + with open(cfg_path, "w", encoding="utf-8") as fh: + json.dump(data, fh, indent=2) + fh.write("\n") + + _ok(f"Saved workspace_dir to {cfg_path}") + except Exception as exc: + _warn(f"Could not persist workspace path to {cfg_path}: {exc}") + + # ------------------------------------------------------------------ + # Workspace file setup + # ------------------------------------------------------------------ + + def _setup_workspace_files(self) -> None: + """Step 5 – Let the user curate workspace contents before execution. + + • If the workspace contains files, list them and let the user choose + which ones to keep (the rest are deleted), or delete everything. + • If the workspace is empty (or became empty after cleanup), offer to + copy files from a user-supplied source directory. + """ + import shutil + from pathlib import Path + + workspace = self.config.workspace_dir + if not os.path.isdir(workspace): + _warn(f"Workspace directory not found ({workspace}). Skipping setup.") + return + + # ── List current workspace files ────────────────────────────── + raw_listing = list_files(path=workspace, max_depth=2) + file_list: list[str] = [ + f for f in raw_listing.splitlines() if f.strip() + ] + + workspace_was_empty = len(file_list) == 0 + kept_files: list[str] = [] + + if file_list: + print(_wrap( + f"The workspace ({workspace}) currently contains " + f"{len(file_list)} file(s):", + width=70, indent=2, + )) + print() + + # Show numbered file list + for idx, fname in enumerate(file_list, start=1): + print(f" {CYAN}[{idx}]{RESET} {fname}") + + print() + print(f" {BOLD}Options:{RESET}") + print(f" {CYAN}Enter numbers{RESET} – comma-separated list of files to " + f"{GREEN}keep{RESET} (others will be deleted)") + print(f" {CYAN}all{RESET} – keep all files") + print(f" {CYAN}none{RESET} – {RED}delete all{RESET} files in the workspace") + print() + + choice = _ask("Files to keep").strip().lower() + + if choice == "all": + kept_files = list(file_list) + _ok(f"Keeping all {len(kept_files)} file(s).") + elif choice == "none" or choice == "": + # Delete everything + kept_files = [] + else: + # Parse comma-separated indices + selected_indices: set[int] = set() + for part in choice.replace(" ", "").split(","): + # Support ranges like "1-5" + if "-" in part: + bounds = part.split("-", 1) + try: + lo, hi = int(bounds[0]), int(bounds[1]) + selected_indices.update(range(lo, hi + 1)) + except ValueError: + _warn(f"Ignoring invalid range: {part}") + else: + try: + selected_indices.add(int(part)) + except ValueError: + _warn(f"Ignoring invalid number: {part}") + + for idx in sorted(selected_indices): + if 1 <= idx <= len(file_list): + kept_files.append(file_list[idx - 1]) + + if kept_files: + _ok(f"Keeping {len(kept_files)} file(s).") + else: + _warn("No valid files selected — all files will be deleted.") + + # ── Perform deletion of un-kept files ───────────────────── + if kept_files and len(kept_files) < len(file_list): + kept_set = set(kept_files) + deleted = 0 + for fname in file_list: + if fname not in kept_set: + full_path = os.path.join(workspace, fname) + try: + if os.path.isfile(full_path): + os.remove(full_path) + deleted += 1 + elif os.path.isdir(full_path): + shutil.rmtree(full_path) + deleted += 1 + except OSError as exc: + _warn(f"Could not delete {fname}: {exc}") + # Clean up empty parent directories left behind + self._prune_empty_dirs(workspace) + if deleted: + _ok(f"Deleted {deleted} file(s) from workspace.") + elif not kept_files and file_list: + # Delete everything + for item in Path(workspace).iterdir(): + try: + if item.is_dir(): + shutil.rmtree(item) + else: + item.unlink() + except OSError as exc: + _warn(f"Could not delete {item.name}: {exc}") + _ok("All workspace files deleted.") + + # ── Offer to import files if workspace is (now) empty ───────── + # Re-check after potential deletions + fresh_listing = list_files(path=workspace, max_depth=2) + remaining = [f for f in fresh_listing.splitlines() if f.strip()] + workspace_is_empty = len(remaining) == 0 + + if workspace_is_empty: + if workspace_was_empty: + _info("Workspace is empty.") + print() + want_import = _ask_yn( + "Copy files from a source directory into the workspace?", + default=False, + ) + if want_import: + self._import_files_to_workspace() + else: + _info("Workspace will remain empty — agents can create files at runtime.") + + def _import_files_to_workspace(self) -> None: + """Ask for a source directory and copy its contents into the workspace + using ``LocalTransfer.transfer_files_to_workspace``. + """ + while True: + src_path = _ask("Path to source directory") + if not src_path: + _info("No path provided — skipping import.") + return + src_path = os.path.expanduser(src_path.strip()) + if os.path.isdir(src_path): + break + _err(f"Directory not found: {src_path}. Please try again.") + + try: + transfer = LocalTransfer( + config=self.config, + workspace_path=self.config.workspace_dir, + runs_capsule_dir=self.config.runs_capsule_dir, + ) + copied = transfer.transfer_files_to_workspace(src_path) + _ok(f"Copied {copied} file(s) into workspace.") + except Exception as exc: + _err(f"File transfer failed: {exc}") + + @staticmethod + def _prune_empty_dirs(root: str) -> None: + """Remove empty sub-directories under *root* (bottom-up).""" + for dirpath, dirnames, filenames in os.walk(root, topdown=False): + if dirpath == root: + continue + if not filenames and not dirnames: + try: + os.rmdir(dirpath) + except OSError: + pass + + # ------------------------------------------------------------------ + # Model selection helpers + # ------------------------------------------------------------------ + + def _model_menu( + self, + prompt_desc: str, + current_value: str, + available: list[tuple[str, str]], + ) -> str: + """Generic numbered model-selection menu. + + Args: + prompt_desc: One-line description shown to the user (what this model + controls). + current_value: Value already in config (may be empty string). + available: List of (display_label, litellm_model_id) for presets whose + API key is present. + + Returns: + The chosen model ID (may equal *current_value* if the user just + pressed Enter). + """ + suggested = current_value or (available[0][1] if available else "") + + if current_value: + _info(f"Current value (from config): {current_value}") + + print(_wrap(prompt_desc, width=70, indent=2)) + print() + + if available: + print(f" {BOLD}Available presets:{RESET}") + for idx, (label, model_id) in enumerate(available, start=1): + is_default = (model_id == suggested) + tag = f"{GREEN}← default{RESET}" if is_default else "" + num_color = GREEN if is_default else CYAN + print(f" {num_color}[{idx}]{RESET} {label} {tag}") + print(f" {DIM}{model_id}{RESET}") + print(f" {CYAN}[c]{RESET} Enter a custom model ID") + else: + _warn("No matching API key found — enter a model ID manually.") + + print() + if suggested: + choice = _ask( + "Select number, 'c' for custom, or Enter to keep current", + default="", + ) + else: + choice = _ask("Select number or 'c' for custom") + + if not choice and suggested: + return suggested + if choice.lower() == "c" or (not available): + custom = _ask( + "Enter model ID (e.g. openai/gpt-4o, " + "anthropic/claude-3-5-sonnet-20241022)" + ) + return custom.strip() if custom.strip() else suggested + # Numbered selection + try: + idx = int(choice) - 1 + if 0 <= idx < len(available): + return available[idx][1] + _warn(f"Invalid selection '{choice}'. Using default.") + except ValueError: + _warn(f"Unrecognised input '{choice}'. Using default.") + return suggested or (available[0][1] if available else "") + + def _choose_models(self) -> None: + """Step 3 – model selection. + + Sub-step 3a: orchestration model (planner, prompts, workflow, judge). + Sub-step 3b: agent execution model (smolagent_model_id). + + Both choices are persisted to *config_default.json*. + """ + available: list[tuple[str, str]] = [ + (label, model_id) + for env_key, label, model_id in _MODEL_PRESETS + if os.getenv(env_key) + ] + + # ── 3a · Orchestration model ────────────────────────────────── + print(f"\n{BOLD} 3a · Orchestration model{RESET}") + print(f" {DIM}Used for planning, workflow generation, and evaluation.{RESET}") + orch_model = self._model_menu( + prompt_desc=( + "Choose the main LLM Mimosa will use for orchestration " + "(planning, workflow generation, and evaluation). Applied to " + "planner, prompts, workflow, and judge roles." + ), + current_value=self.config.planner_llm_model or "", + available=available, + ) + + if orch_model: + for key in _MODEL_CFG_KEYS: + setattr(self.config, key, orch_model) + _ok(f"Orchestration model: {orch_model}") + else: + _warn("No orchestration model chosen — keeping existing config values.") + + # ── 3b · Agent execution model (smolagent_model_id) ────────── + print(f"\n{BOLD} 3b · Agent execution model (SmolAgents){RESET}") + print(f" {DIM}Used by the code-executing agents inside each workflow.{RESET}") + print(f" {DIM}Can be the same as the orchestration model or a faster/cheaper one.{RESET}") + agent_model = self._model_menu( + prompt_desc=( + "Choose the LLM for agent execution (SmolAgents tasks). " + "A fast, cost-effective model works well here." + ), + current_value=self.config.smolagent_model_id or "", + available=available, + ) + + if agent_model: + self.config.smolagent_model_id = agent_model + _ok(f"Agent execution model: {agent_model}") + else: + _warn("No agent model chosen — keeping existing config values.") + + # Persist both choices at once + self._persist_models(orch_model or "", agent_model or "") + + def _persist_models(self, orch_model_id: str, agent_model_id: str) -> None: + """Write both model choices to config_default.json.""" + cfg_path = self._CONFIG_DEFAULT_PATH + try: + if os.path.isfile(cfg_path): + with open(cfg_path, encoding="utf-8") as fh: + data = json.load(fh) + else: + data = {} + + if orch_model_id: + for key in _MODEL_CFG_KEYS: + data[key] = orch_model_id + + if agent_model_id: + data["smolagent_model_id"] = agent_model_id + + with open(cfg_path, "w", encoding="utf-8") as fh: + json.dump(data, fh, indent=2) + fh.write("\n") + + _ok(f"Saved model choices to {cfg_path}") + except Exception as exc: + _warn(f"Could not persist model choices to {cfg_path}: {exc}") + + def _collect_objective(self) -> None: + """Prompt the user for their initial research objective.""" + print(_wrap( + "Describe what you want Mimosa to do. This can be a high-level " + "scientific goal (e.g. 'Reproduce Figure 3 from paper X') or a " + "focused task (e.g. 'Train a toxicity model on the ClinTox dataset'). " + "Don't worry about being too vague — we'll refine it together next.", + width=70, indent=2, + )) + while True: + objective = _ask("Your objective") + if len(objective.strip()) >= 10: + self._objective = objective.strip() + break + _warn("Please enter a more descriptive objective (at least 10 characters).") + + def _clarify_and_refine(self) -> None: + """LLM conversation loop: clarify missing info, then refine the prompt.""" + print(_wrap( + "The assistant will now check whether your objective is clear enough " + "for Mimosa to execute and may ask one or more follow-up questions. " + "Once complete, it will produce a refined, actionable prompt.", + width=70, indent=2, + )) + + llm = _build_llm(self.config, temperature=0.3, max_tokens=768) + + # Accumulate context: original objective + Q&A pairs + context_lines: list[str] = [f"Objective: {self._objective}"] + max_clarification_rounds = 5 + + for round_num in range(max_clarification_rounds): + full_context = "\n".join(context_lines) + + print(f"\n{DIM} [Clarification round {round_num + 1}/{max_clarification_rounds}]{RESET}") + + try: + raw = _call_llm(llm, _CLARIFIER_SYSTEM, full_context) + result = _parse_json_response(raw) + except Exception as exc: + _warn(f"LLM clarification failed ({exc}). Skipping refinement.") + return + + is_clear = result.get("is_clear", False) + question = result.get("question", "").strip() + refined_prompt = result.get("refined_prompt", "").strip() + + if not is_clear and question: + # Ask the clarifying question + print() + print(f" {BOLD}Assistant:{RESET} {question}") + answer = _ask("Your answer") + if answer: + context_lines.append(f"Q: {question}") + context_lines.append(f"A: {answer}") + else: + _info("No answer provided — skipping this question.") + continue # loop for next round + + if is_clear and refined_prompt: + # Show the refined prompt and ask for confirmation + print() + print(f" {BOLD}Refined objective:{RESET}") + print() + # Print wrapped refined prompt with colour + for line in textwrap.wrap(refined_prompt, width=64): + print(f" {CYAN}{line}{RESET}") + print() + confirmed = _ask_yn("Accept this refined objective?", default=True) + if confirmed: + self._objective = refined_prompt + _ok("Objective accepted.") + return + else: + # Let the user correct it manually + correction = _ask( + "Edit the objective (or press Enter to keep the original)" + ) + if correction: + self._objective = correction.strip() + context_lines = [f"Objective: {self._objective}"] + _ok(f"Continuing with: {self._objective[:80]}") + return + + # Exhausted rounds without clarity — keep whatever we have + _warn( + f"Clarification loop completed ({max_clarification_rounds} rounds). " + "Using current objective as-is." + ) + + def _classify_and_confirm(self) -> None: + """Use LLM to classify objective as goal or task, confirm with user.""" + print(_wrap( + "Asking the LLM to classify your objective as Goal-mode " + "(multi-step planning) or Task-mode (single focused operation) …", + width=70, indent=2, + )) + + classification: dict | None = None + llm = _build_llm(self.config, temperature=0.0, max_tokens=256) + + try: + raw = _call_llm( + llm, + _CLASSIFIER_SYSTEM, + f"Classify this research objective:\n\n{self._objective}", + ) + classification = _parse_json_response(raw) + except Exception as exc: + _warn(f"LLM classification failed ({exc}). Falling back to manual selection.") + + if classification: + mode = classification.get("mode", "task") + confidence = float(classification.get("confidence", 0.0)) + reasoning = classification.get("reasoning", "") + label = classification.get("suggested_label", self._objective[:40]) + + print() + print(f" {BOLD}Suggested mode:{RESET} {CYAN}{mode.upper()}{RESET} " + f"(confidence: {confidence:.0%})") + print(f" {BOLD}Reasoning:{RESET} {reasoning}") + print(f" {BOLD}Label:{RESET} {label}") + print() + _info( + "Goal mode → Mimosa decomposes the objective into a plan of tasks " + "and executes them sequentially (planner).\n" + " ℹ️ Task mode → Mimosa directly synthesises and runs a single " + "multi-agent workflow for the objective (DGM)." + ) + + confirmed = _ask_yn(f"Accept '{mode}' mode?", default=True) + if confirmed: + self._mode = mode # type: ignore[assignment] + return + + # Manual fallback / override + print() + print(f" {BOLD}Available modes:{RESET}") + print(f" {CYAN}goal{RESET} – high-level research objective (planner + DGM)") + print(f" {CYAN}task{RESET} – single focused operation (DGM only)") + choice = _ask("Choose mode", default="task").lower() + self._mode = "goal" if choice.startswith("g") else "task" + _ok(f"Mode set to: {self._mode.upper()}") + + def _collect_options(self) -> None: + """Ask about learning mode and other options.""" + print(_wrap( + "Learning mode enables Mimosa to iteratively improve its workflow " + "through Darwinian self-evolution until a quality threshold is met " + "(recommended for first-time runs on a new objective).", + width=70, indent=2, + )) + self._learn = _ask_yn("Enable learning mode?", default=False) + if self._learn: + _ok("Learning mode enabled.") + else: + _info("Learning mode disabled (single-pass execution).") + + # Summary + print() + print(f" {BOLD}{'─'*54}{RESET}") + print(f" {BOLD}LAUNCH SUMMARY{RESET}") + print(f" {'─'*54}") + print(f" Mode: {CYAN}{self._mode.upper()}{RESET}") + print(f" Learning: {'Yes' if self._learn else 'No'}") + print(f" Objective: {self._objective[:60]}{'…' if len(self._objective) > 60 else ''}") + print(f" {'─'*54}") + print() + go = _ask_yn("Launch Mimosa now?", default=True) + if not go: + print("\n Exiting without launching. Run again when ready.\n") + sys.exit(0) + + async def _launch(self) -> None: + """Validate config paths and start the selected execution mode.""" + try: + self.config.validate_paths() + except AssertionError as exc: + _err(f"Configuration validation failed: {exc}") + _info( + "Check that your workspace_dir and other paths in config are correct. " + "Make sure Toolomics is running and the workspace exists and try again." + ) + sys.exit(1) + + print(MIMOSA_START_BANNER) + + if self._mode == "goal": + await self._launch_goal() + else: + await self._launch_task() + # Archive workspace after completion + trs = LocalTransfer( + config=self.config, + workspace_path=self.config.workspace_dir, + runs_capsule_dir=self.config.runs_capsule_dir, + ) + capsule = trs.transfer_workspace_files_to_capsule(self._objective) + print(f"\n{GREEN}{BOLD} Workspace files archived to capsule: {capsule}{RESET}\n") + + async def _launch_goal(self) -> None: + """Start planner mode (multi-step goal).""" + from sources.core.planner import Planner + from sources.utils.transfer_toolomics import LocalTransfer + + print(f"\n{GREEN}{BOLD} Launching in GOAL mode …{RESET}\n") + planner = Planner(self.config) + await planner.start_planner( + goal=self._objective, + judge=True, + max_evolve_iteration=self.config.max_learning_evolve_iterations if self._learn else 1, + ) + + async def _launch_task(self) -> None: + """Start DGM task mode (single operation).""" + from sources.core.dgm import DarwinMachine + + print(f"\n{GREEN}{BOLD} Launching in TASK mode …{RESET}\n") + dgm = DarwinMachine(self.config) + await dgm.start_dgm( + goal=self._objective, + judge=True, + learning_mode=self._learn, + max_iteration=self.config.max_learning_evolve_iterations if self._learn else 1, + ) diff --git a/sources/cli/pretty_print.py b/sources/cli/pretty_print.py new file mode 100644 index 00000000..27a47abe --- /dev/null +++ b/sources/cli/pretty_print.py @@ -0,0 +1,256 @@ +""" +Pretty-print utilities for Mimosa-AI CLI output. +Modern Claude-code–style terminal output with consistent, clean styling. + +Usage:: + + from sources.cli.pretty_print import ( + print_ok, print_warn, print_err, print_info, + print_phase, print_section, print_rule, + print_iteration_header, + print_box, + print_kv_row, print_summary, + print_agent_answers, + print_step, + CYAN, GREEN, YELLOW, RED, BLUE, MAGENTA, BOLD, DIM, RESET, + ) +""" + +from __future__ import annotations + +import textwrap + +# ── ANSI colour constants ────────────────────────────────────────────────────── +CYAN = "\033[96m" +GREEN = "\033[92m" +YELLOW = "\033[93m" +RED = "\033[91m" +BLUE = "\033[94m" +MAGENTA = "\033[95m" +BOLD = "\033[1m" +DIM = "\033[2m" +RESET = "\033[0m" + +# Default column width used for banners/summaries +_W = 80 + + +# ── Status lines ────────────────────────────────────────────────────────────── + +def print_ok(msg: str) -> None: + """Print a green success line ✓ """ + print(f"{GREEN} ✓ {msg}{RESET}") + + +def print_warn(msg: str) -> None: + """Print a yellow warning line ⚠ """ + print(f"{YELLOW} ⚠ {msg}{RESET}") + + +def print_err(msg: str) -> None: + """Print a red error line ✗ """ + print(f"{RED} ✗ {msg}{RESET}") + + +def print_info(msg: str) -> None: + """Print a dim informational line · """ + print(f"{DIM} · {msg}{RESET}") + + +# ── Step header (used by onboarding CLI) ────────────────────────────────────── + +def print_step(step: int, total: int, title: str, width: int = 60) -> None: + """ + Numbered step header used by the onboarding wizard. + + Example:: + + ──────────────────────────────────────────────────────────── + Step 1/8 · API Key Check + ──────────────────────────────────────────────────────────── + """ + bar = "─" * width + print(f"\n{CYAN}{bar}{RESET}") + print(f"{CYAN} Step {step}/{total} · {title}{RESET}") + print(f"{CYAN}{bar}{RESET}") + + +# ── Phase / section banners ────────────────────────────────────────────────── + +def print_phase( + title: str, + icon: str = "", + width: int = _W, + color: str = CYAN, +) -> None: + """ + Full-width phase banner with centred title and horizontal rules. + """ + label = f"{icon} {title}" if icon else title + bar = "─" * width + print(f"\n{color}{bar}{RESET}") + print(f"{color}{label:^{width}}{RESET}") + print(f"{color}{bar}{RESET}") + + +def print_section( + title: str, + color: str = CYAN, + width: int = 60, +) -> None: + """ + Compact inline section label. + """ + label = f" {title} " + remaining = max(0, width - len(label) - 2) + print(f"\n{color}──{label}{'─' * remaining}{RESET}") + + +def print_rule(width: int = _W, color: str = CYAN) -> None: + """Print a plain horizontal rule.""" + print(f"{color}{'─' * width}{RESET}") + + +# ── Iteration banner ───────────────────────────────────────────────────────── + +def print_iteration_header( + current: int, + total: int, + subtitle: str = "Self-Improvement Loop", + width: int = _W, +) -> None: + """ + Prominent iteration counter banner shown at the start of each DGM loop. + """ + bar = "═" * width + print(f"\n{CYAN}{bar}{RESET}") + print(f"{CYAN}{BOLD} ITERATION {current}/{total} · {subtitle}{RESET}") + print(f"{DIM} Mimosa will now learn how to build the workflow for the task.{RESET}") + print(f"{CYAN}{bar}{RESET}") + + +# ── Content box ────────────────────────────────────────────────────────────── + +def print_box( + content: str, + title: str = "", + color: str = CYAN, + width: int = 60, + truncate: int = 256, +) -> None: + """ + Render content inside a Unicode-bordered box. + + Long lines are word-wrapped to fit inside the box. Lines longer than + *truncate* characters are hard-truncated and annotated with a + ``…(N chars not shown)`` note. + + Example:: + + ╭─ CURRENT TASK ──────────────────────────────────────────╮ + │ Create a simple test workflow that demonstrates basic │ + │ functionality by outputting the text 'Hello World'. │ + ╰──────────────────────────────────────────────────────────────╯ + """ + inner_w = width - 4 # 2 border chars + 2 spaces padding on each side + + # Top border + if title: + top_label = f"─ {title} " + top_fill = max(0, width - len(top_label) - 2) + top = f"╭{top_label}{'─' * top_fill}╮" + else: + top = f"╭{'─' * (width - 2)}╮" + + print(f"{color}{top}{RESET}") + + for raw_line in content.splitlines(): + if len(raw_line) <= inner_w: + print(f"{color}│ {raw_line.ljust(inner_w)} │{RESET}") + elif len(raw_line) <= truncate: + # word-wrap + wrapped_lines = textwrap.wrap(raw_line, inner_w) or [""] + for wrapped in wrapped_lines: + print(f"{color}│ {wrapped.ljust(inner_w)} │{RESET}") + else: + # hard-truncate + annotate + shown = raw_line[:truncate] + note = f"…({len(raw_line) - truncate} chars not shown)" + for wrapped in textwrap.wrap(shown, inner_w) or [""]: + print(f"{color}│ {wrapped.ljust(inner_w)} │{RESET}") + print(f"{DIM}│ {note.ljust(inner_w)} │{RESET}") + + print(f"{color}╰{'─' * (width - 2)}╯{RESET}") + + +# ── Key-value rows / summary ────────────────────────────────────────────────── + +def print_kv_row( + key: str, + value: str, + color: str = CYAN, + key_width: int = 20, +) -> None: + """Print a single ``key → value`` row with aligned columns.""" + print(f" {BOLD}{key:<{key_width}}{RESET} {color}{value}{RESET}") + + +def print_summary( + title: str, + items: list[tuple[str, str]], + color: str = CYAN, + width: int = _W, +) -> None: + """ + Print a titled summary block with aligned key-value pairs. + + Example:: + + ──────────────────────────────────────────────────────────────────────────────── + ✨ WORKFLOW COMPLETION SUMMARY + ──────────────────────────────────────────────────────────────────────────────── + UUID 20260408_145530_9a908c10 + Total time 12.340s + Generation 3.120s + Dependencies 1.200s + Execution 8.020s + ──────────────────────────────────────────────────────────────────────────────── + """ + bar = "─" * width + key_width = max((len(k) for k, _ in items), default=12) + 2 + print(f"\n{color}{bar}{RESET}") + print(f"{color} {title}{RESET}") + print(f"{color}{bar}{RESET}") + for key, value in items: + print(f" {BOLD}{key:<{key_width}}{RESET} {color}{value}{RESET}") + print(f"{color}{bar}{RESET}\n") + + +# ── Agent answers ───────────────────────────────────────────────────────────── + +def print_agent_answers( + answers_text: str, + color: str = CYAN, + width: int = _W, +) -> None: + """ + Display workflow agent answers inside a labelled section. + + Example:: + + ─────────────────────── WORKFLOW AGENTS ANSWERS ─────────────────────────────── + agent 0: Task completed successfully. + agent 1: All assertions passed. + ───────────────────────────────────────────────────────────────────────────────── + """ + if not answers_text or not answers_text.strip(): + return + + label = " WORKFLOW AGENTS ANSWERS " + side = max(0, width - len(label)) + left = side // 2 + right = side - left + print(f"\n{color}{'─' * left}{label}{'─' * right}{RESET}") + for line in answers_text.splitlines(): + print(f"{color} {line}{RESET}") + print(f"{color}{'─' * width}{RESET}\n") diff --git a/sources/core/dgm.py b/sources/core/dgm.py index 72fd6558..ffbc9dfc 100644 --- a/sources/core/dgm.py +++ b/sources/core/dgm.py @@ -21,6 +21,13 @@ from .workflow_selection import WorkflowSelector from .schema import IndividualRun, ImprovementLog from .improvement_validator import ImprovementValidator +from sources.cli.pretty_print import ( + print_ok, print_warn, print_err, print_info, + print_phase, print_section, + print_iteration_header, print_box, + print_summary, print_agent_answers, + CYAN, GREEN, YELLOW, RED, DIM, RESET, BOLD, +) def check_answer_success(answer: str) -> bool: @@ -142,10 +149,7 @@ def get_flow_answers(self, wf_state: any) -> str: return flow_answers def show_answers(self, flow_answers): - print(f"\n\033[96m{'> WORKFLOW AGENTS ANSWERS':^60}\033[0m") - print(f"\033[96m{'─' * 60}\033[0m") - print(f"\033[96m{flow_answers}\033[0m") - print(f"\033[96m{'─' * 60}\033[0m\n") + print_box(flow_answers, title="Workflow Agents Answers", color=YELLOW) def improvement_prompt( self, @@ -233,11 +237,9 @@ def select_workflow_template(self, goal, template_uuid: str = None) -> WorkflowI threshold_similary=0.8, threshod_score=0.0, ) - print(f"\n\033[96m{'🎯 WORKFLOW SELECTION':^60}\033[0m") - print(f"\033[96m{'─' * 60}\033[0m") - print(f"\033[96mSelected {len(candidates)} candidates.\033[0m") - print(f"\033[96mTop candidate: {candidates[0].uuid if candidates else str(None)}\033[0m") - print(f"\033[96m{'─' * 60}\033[0m\n") + print_section("🎯 WORKFLOW SELECTION") + print_info(f"Selected {len(candidates)} candidate(s)") + print_info(f"Top candidate: {candidates[0].uuid if candidates else 'None'}") return WorkflowInfo(candidates[0].uuid, Path(f"{self.workflow_dir}/{candidates[0].uuid}")) if candidates else None return WorkflowInfo(template_uuid, Path(f"{self.workflow_dir}/{template_uuid}")) @@ -275,7 +277,7 @@ async def start_dgm( instead of calling orchestrate_workflow. Useful for testing and debugging. """ if learning_mode: - max_iteration = max(3, self.config.max_learning_evolve_iterations) + max_iteration = self.config.max_learning_evolve_iterations wf = self.select_workflow_template( goal, template_uuid=template_uuid @@ -286,10 +288,8 @@ async def start_dgm( if wf is None: raise ValueError("❌ Mockup mode requires a valid workflow template. " "Please provide a template_uuid or ensure workflows exist in the workflow directory.") - print(f"\n\033[93m{'MOCKUP MODE':^60}\033[0m") - print(f"\033[93m{'─' * 60}\033[0m") - print(f"\033[93mUsing existing workflow data from: {wf.uuid}\033[0m") - print(f"\033[93m{'─' * 60}\033[0m\n") + print_phase("MOCKUP MODE", color="\033[93m") + print_info(f"Using existing workflow data from: {wf.uuid}") mock_run = IndividualRun( goal=wf.goal or goal, prompt=wf.code or goal, @@ -307,7 +307,7 @@ async def start_dgm( ) flow_answers = self.get_flow_answers(wf.state_result) self.show_answers(flow_answers) - print(f"\n\033[93mMockup run completed with reward: {wf.overall_score:.1f}\033[0m") + print_ok(f"Mockup run completed with reward: {wf.overall_score:.1f}") return [mock_run] craft_instructions = self.get_craft_instructions(goal, wf) @@ -347,11 +347,12 @@ async def recursive_self_improvement( self._log_iteration_start(runs[-1].goal, runs[-1].iteration_count, runs[-1].max_depth) iteration_start_time = time.time() + on_error = False uuid = None current_iteration_cost = 0.0 # Cost for this iteration only, not cumulative # Execute workflow - print(f"\nCurrently at run: {runs[-1].iteration_count}. max depth: {runs[-1].max_depth}.\n") + print_info(f"Run {runs[-1].iteration_count + 1} of {runs[-1].max_depth}") run_stdout, uuid, workflow_code, executed = await self.orchestrator.orchestrate_workflow( goal=runs[-1].goal, craft_instructions=runs[-1].prompt, @@ -359,12 +360,14 @@ async def recursive_self_improvement( single_agent_mode=single_agent_mode ) wf_info = WorkflowInfo(uuid, Path(f"{self.workflow_dir}/{uuid}")) + if "WORKFLOW_GENERATION_ERROR" in run_stdout: + print_err(f"Workflow generation failed:\n{run_stdout[256:]}") + on_error = True if workflow_code: # Evaluate and calculate costs eval_type, current_iteration_cost = await self._evaluate_and_calculate_cost( executed, runs[-1].judge, uuid, runs[-1].answers, runs[-1].scenario_rubric, assertion_history ) - # Don't overwrite runs[-1].cost - it contains cumulative from previous iterations runs[-1].reward = wf_info.overall_score runs[-1].current_uuid = uuid @@ -414,12 +417,12 @@ async def recursive_self_improvement( all_success = evaluate_workflow_success(wf_info, runs[-1].answers) # Check termination conditions - if runs[-1].iteration_count >= runs[-1].max_depth-1: - print("\nmax recursive depth reached.\n") + if runs[-1].iteration_count >= runs[-1].max_depth-1 and not on_error: + print_info("Maximum recursive depth reached.") return runs if learning_mode and wf_info.overall_score > self.config.learned_score_threshold: # reach learning threshold - print("\nDGM done learning task.\n") + print_ok("DGM done learning task.") self._save_final_plots(assertion_history, rewards_history, uuid) self.notifier.send_message( f"Done learning task: {wf_info.goal[:256]} \n" @@ -429,19 +432,20 @@ async def recursive_self_improvement( priority=0 ) return runs - elif not learning_mode and all_success: - self._save_final_plots(assertion_history, rewards_history, uuid) - print("\nDGM completed task.\n") - self.notifier.send_message( - f"Evolution completed successfully!\n" - f"Goal: {runs[-1].goal[:128]}...\n" - f"Final UUID: {uuid}\n" - f"Iterations: {runs[-1].iteration_count + 1}/{runs[-1].max_depth}\n" - f"All workflows successful!", - title=f"Evolution success - {uuid}", - priority=0 - ) - return runs + elif not on_error: + if not learning_mode and all_success: + self._save_final_plots(assertion_history, rewards_history, uuid) + print_ok("DGM completed task successfully.") + self.notifier.send_message( + f"Task completed successfully!\n" + f"Goal: {runs[-1].goal[:128]}...\n" + f"Final UUID: {uuid}\n" + f"Iterations: {runs[-1].iteration_count + 1}/{runs[-1].max_depth}\n" + f"All workflows successful!", + title=f"Evolution success - {uuid}", + priority=0 + ) + return runs # select and use best scoring workflow wf_info_best = self.select_workflow_template( @@ -468,7 +472,6 @@ async def recursive_self_improvement( original_task=runs[-1].original_task # PRESERVE original_task for workflow selection )) - time.sleep(5) runs = await self.recursive_self_improvement( runs, rewards_history=rewards_history, @@ -484,29 +487,15 @@ def _get_human_validation(self) -> bool: """Get human validation for continuing the workflow.""" human_validation = input("Attempt to retry task? (yes/no): ").strip().lower() if human_validation not in ["yes", "y"]: - print("Exiting self-improvement loop.\n") + print("Exiting self-improvement loop.") return False return True def _log_iteration_start(self, goal: str, iteration_count: int, max_depth: int): """Log the start of an iteration.""" logger = logging.getLogger(__name__) - - print(f"\n\033[94m{'=' * 60}\033[0m") - print(f"\033[94mITERATION {iteration_count + 1}/{max_depth} - Self-Improvement Loop.\n\033[0m" - f"\033[94mDGM Will attempt to retry and improve workflow on same task.\033[0m") - print(f"\033[94m{'=' * 60}\033[0m") - print(f"\n\033[94m{'📋 CURRENT TASK':^60}\033[0m") - print(f"\033[94m{'─' * 60}\033[0m") - goal_lines = goal.split('\n') - for line in goal_lines: - if len(line) <= 256: - print(f"\033[94m {line}\033[0m") - else: - truncated = line[:256] - remaining = len(line) - 256 - print(f"\033[94m {truncated}...({remaining} remaining characters not displayed)\033[0m") - print(f"\033[94m{'─' * 60}\033[0m\n") + print_iteration_header(iteration_count + 1, max_depth) + print_box(goal, title="📋 CURRENT TASK", truncate=256) logger.info(f"[ITERATION START] {iteration_count + 1}/{max_depth} - {goal[:50]}...") async def _evaluate_and_calculate_cost( @@ -534,22 +523,16 @@ async def _evaluate_workflow( ) -> str: """Evaluate the workflow and update assertion history.""" logger = logging.getLogger(__name__) - - print(f"\n\033[94m{'⚖️ WORKFLOW EVALUATION PHASE':^80}\033[0m") - print(f"\033[94m{'=' * 80}\033[0m") - + print_phase("⚖️ WORKFLOW EVALUATION PHASE") eval_start = time.time() eval_result = self.judge.evaluate(uuid=uuid, answer=answer, scenario_rubric=scenario_rubric) eval_type = 'scenario' if scenario_rubric else 'generic' eval_time = time.time() - eval_start - logger.info(f"[WORKFLOW EVALUATION] {uuid}:\n{json.dumps(eval_result, indent=2)}") - print(f"\033[94m✅ Workflow evaluation completed in {eval_time:.3f}s\033[0m") - + print_ok(f"Workflow evaluation completed in {eval_time:.3f}s") # Track assertion progress for scenario evaluation if scenario_rubric and isinstance(eval_result, dict) and assertion_history is not None: self._update_assertion_history(eval_result, assertion_history) - return eval_type def _update_assertion_history(self, eval_result: dict, assertion_history: list): @@ -557,15 +540,14 @@ def _update_assertion_history(self, eval_result: dict, assertion_history: list): passed = eval_result.get('passed_assertions', eval_result.get('earned_points', 0)) total = eval_result.get('total_assertions', eval_result.get('total_points', 100)) assertion_history.append([passed, total]) - print(f"\033[94m📊 Assertions Progress: {passed}/{total} " - f"({passed/total*100 if total > 0 else 0:.0f}%)\033[0m") + pct = passed / total * 100 if total > 0 else 0 + print_info(f"📊 Assertions progress: {passed}/{total} ({pct:.0f}%)") def _update_visualizations( self, rewards_history: list, assertion_history: list, goal: str, scenario_rubric: str, uuid: str ): """Update all visualizations with current data.""" - # Update assertion plot if available if assertion_history: self._update_assertion_plot(assertion_history, scenario_rubric, uuid) elif rewards_history: @@ -580,18 +562,12 @@ def _update_assertion_plot( ): """Update assertion progress plot.""" from sources.evaluation.scenario_loader import ScenarioLoader - scenario = ScenarioLoader().load_scenario(scenario_rubric) total_assertions = len(scenario.get("assertions", [])) if scenario else 0 - - self.viz_utils.update_assertion_progress_plot( - assertion_history, total_assertions - ) - - # Save plot after each update for real-time monitoring + self.viz_utils.update_assertion_progress_plot(assertion_history, total_assertions) plot_filename = f"{self.workflow_dir}/{uuid}/assertion_progress.png" self.viz_utils.save_plot(plot_filename) - print(f"\033[94m📊 Assertion progress plot updated: {plot_filename}\033[0m") + print_info(f"📊 Assertion progress plot updated: {plot_filename}") def _log_iteration_completion( self, iteration_count: int, max_depth: int, iteration_start_time: float, @@ -601,18 +577,18 @@ def _log_iteration_completion( """Log iteration completion and send notification.""" logger = logging.getLogger(__name__) iteration_time = time.time() - iteration_start_time - logger.info( f"[ITERATION END] {iteration_count + 1}/{max_depth} completed in {iteration_time:.3f}s - " f"Rewards: {wf_rewards:.1f}, Cost: {exec_cost:.3f} USD" ) - - print(f"\n\033[94m{'-' * 60}\033[0m") - print(f"\033[94mTotal rewards: {wf_rewards:.1f}\033[0m") - print(f"\033[94mTotal cost: {exec_cost:.6f} USD\033[0m") - print(f"\033[94mIteration time: {iteration_time:.3f}s\033[0m") - print(f"\033[94m{'-' * 60}\033[0m\n") - + print_summary( + f"ITERATION {iteration_count + 1}/{max_depth} COMPLETE", + [ + ("Rewards", f"{wf_rewards:.1f}"), + ("Cost", f"${exec_cost:.6f}"), + ("Time", f"{iteration_time:.3f}s"), + ], + ) self.notifier.send_message( f"Iteration {iteration_count + 1} completed.\n" f"Goal: {goal[:128]}...\n" @@ -628,5 +604,5 @@ def _save_final_plots(self, assertion_history: list, reward_history: list, uuid: if assertion_history or reward_history: plot_filename = f"{self.workflow_dir}/{uuid}/reward_progress.png" self.viz_utils.save_plot(plot_filename) - print(f"\033[94m📊 Assertion progress plot saved to: {plot_filename}\033[0m") + print_info(f"📊 Reward progress plot saved: {plot_filename}") return plot_filename diff --git a/sources/core/orchestrator.py b/sources/core/orchestrator.py index 6dfe1a58..bf65c040 100644 --- a/sources/core/orchestrator.py +++ b/sources/core/orchestrator.py @@ -6,6 +6,10 @@ import time from sources.utils.notify import PushNotifier +from sources.cli.pretty_print import ( + print_ok, print_err, print_info, + print_phase, print_summary, +) from .workflow_factory import WorkflowFactory from .workflow_runner import ExecutionStatus, RuntimeConfig, WorkflowRunner @@ -37,7 +41,7 @@ def __init__(self, config) -> None: async def workflow_requirements_install(self): deps = self.config.runner_requirements - print(f"\033[96m📦 Installing workflow dependencies: {deps}\033[0m") + print_info(f"📦 Installing workflow dependencies: {deps}") dep_result = await self.workflow_runner.install_dependencies(deps) if dep_result.status != ExecutionStatus.COMPLETED: raise RuntimeError(f"Dependency installation failed: {dep_result.stderr}") @@ -49,19 +53,17 @@ async def workflow_sandbox_run(self, workflow_code: str) -> str: def progress_handler(line: str): print(line) - print("\033[96m▶ Executing workflow in Python sandbox...\033[0m") + print_info("▶ Executing workflow in Python sandbox…") result = await self.workflow_runner.execute( workflow_code, progress_callback=progress_handler ) if result.status == ExecutionStatus.COMPLETED: - print( - f"\033[96m✅ Workflow execution completed successfully in {result.execution_time:.3f}s\033[0m" - ) + print_ok(f"Workflow execution completed in {result.execution_time:.3f}s") return ( result.stdout or result.stderr or "No output from workflow execution." ) else: - print(f"\033[91m❌ Workflow execution failed: {result.stderr}\033[0m") + print_err(f"Workflow execution failed: {result.stderr}") raise Exception(f"Workflow execution failed: {result.stderr}") async def orchestrate_workflow( @@ -87,8 +89,7 @@ async def orchestrate_workflow( execution_output = "" logger.info(f"[WORKFLOW START] Orchestrating workflow - {goal[:50]}...") - print(f"\n\033[96m{'🏗️ WORKFLOW GENERATION PHASE':^80}\033[0m") - print(f"\033[96m{'=' * 80}\033[0m") + print_phase("🏗️ WORKFLOW GENERATION PHASE") # Workflow generation timing generation_start = time.time() @@ -113,8 +114,7 @@ async def orchestrate_workflow( uuid_part, actual_error = error_msg.split("|", 1) workflow_uuid = uuid_part.replace("UUID:", "") logger.warning(f"[WORKFLOW_GENERATION_ERROR]\n{actual_error}\n") - - # Send notification for workflow generation error + self.notifier.send_message( f"Workflow {workflow_uuid} generation failed after {generation_time:.1f}s\n" f"Goal: {goal[:128]}...\n" @@ -125,8 +125,7 @@ async def orchestrate_workflow( return f"WORKFLOW_GENERATION_ERROR: {actual_error}", workflow_uuid, "error", False else: logger.warning(f"[WORKFLOW_GENERATION_ERROR]\n{error_msg}\n") - - # Send notification for workflow generation error + self.notifier.send_message( f"Workflow generation failed after {generation_time:.1f}s\n" f"Goal: {goal[:128]}...\n" @@ -135,51 +134,42 @@ async def orchestrate_workflow( priority=1 ) return f"WORKFLOW_GENERATION_ERROR: {error_msg}", "generation_failed", "error", False - + generation_time = time.time() - generation_start logger.info(f"[WORKFLOW GENERATION] {uuid} generated in {generation_time:.3f}s") - print( - f"\033[96m✅ Workflow {uuid} generated successfully in {generation_time:.3f}s\033[0m" - ) + print_ok(f"Workflow {uuid} generated in {generation_time:.3f}s") if no_run: return "", uuid, workflow_code, True try: # Dependencies installation phase - print(f"\n\033[96m{'📦 DEPENDENCIES INSTALLATION PHASE':^80}\033[0m") - print(f"\033[96m{'=' * 80}\033[0m") + print_phase("📦 DEPENDENCIES INSTALLATION PHASE") deps_start = time.time() await self.workflow_requirements_install() deps_time = time.time() - deps_start logger.info( f"[WORKFLOW DEPS] {uuid} dependencies installed in {deps_time:.3f}s" ) - print( - f"\033[96m✅ Dependencies installed successfully in {deps_time:.3f}s\033[0m" - ) + print_ok(f"Dependencies installed in {deps_time:.3f}s") # Execution phase - print(f"\n\033[96m{'▶ WORKFLOW EXECUTION PHASE':^80}\033[0m") - print(f"\033[96m{'=' * 80}\033[0m") + print_phase("▶ WORKFLOW EXECUTION PHASE") exec_start = time.time() execution_output = await self.workflow_sandbox_run(complete_code) exec_time = time.time() - exec_start logger.info(f"[WORKFLOW EXECUTION] {uuid} executed in {exec_time:.3f}s") - print( - f"\033[96m✅ Workflow executed successfully in {exec_time:.3f}s\033[0m" - ) + print_ok(f"Workflow executed in {exec_time:.3f}s") except Exception as e: workflow_time = time.time() - workflow_start_time logger.info( f"[WORKFLOW ERROR] {uuid} failed after {workflow_time:.3f}s - {str(e)}" ) - print(f"❌ Error during {uuid} workflow execution: {e}") + print_err(f"Error during {uuid} workflow execution: {e}") import traceback traceback.print_exc() - - # Send notification for workflow execution failure + self.notifier.send_message( f"Workflow {uuid} execution failed after {workflow_time:.1f}s\n" f"Goal: {goal[:128]}...\n" @@ -189,19 +179,21 @@ async def orchestrate_workflow( ) return str(e), uuid, workflow_code, False finally: - print("\nCleaning up sandbox...") + print_info("Cleaning up sandbox…") workflow_time = time.time() - workflow_start_time logger.info(f"[WORKFLOW END] {uuid} completed in {workflow_time:.3f}s") - print(f"\n\033[96m{'✨ WORKFLOW COMPLETION SUMMARY':^80}\033[0m") - print(f"\033[96m{'=' * 80}\033[0m") - print(f"\033[96m📋 Workflow UUID: {uuid}\033[0m") - print(f"\033[96m⏱️ Total Time: {workflow_time:.3f}s\033[0m") - print(f"\033[96m • Generation: {generation_time:.3f}s\033[0m") - print(f"\033[96m • Dependencies: {deps_time:.3f}s\033[0m") - print(f"\033[96m • Execution: {exec_time:.3f}s\033[0m") - print(f"\033[96m{'=' * 80}\033[0m\n") + print_summary( + "✨ WORKFLOW COMPLETION SUMMARY", + [ + ("UUID", uuid), + ("Total time", f"{workflow_time:.3f}s"), + ("Generation", f"{generation_time:.3f}s"), + ("Dependencies", f"{deps_time:.3f}s"), + ("Execution", f"{exec_time:.3f}s"), + ], + ) output = ( execution_output.strip() @@ -219,7 +211,7 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): try: await self.workflow_runner.cleanup() except Exception as e: - print(f"❌ Error during cleanup: {e}") + print_err(f"Error during cleanup: {e}") import traceback traceback.print_exc() diff --git a/sources/core/planner.py b/sources/core/planner.py index c89b4838..e41b37b0 100644 --- a/sources/core/planner.py +++ b/sources/core/planner.py @@ -13,6 +13,11 @@ from sources.utils.planner_visualization import PlannerVisualizer from sources.utils.list_files import list_files from sources.extensibility.text_to_speech import create_tts_service +from sources.cli.pretty_print import ( + print_ok, print_warn, print_err, print_info, + print_phase, print_section, print_summary, + CYAN, BOLD, DIM, RESET, +) class PlanValidationError(Exception): @@ -79,7 +84,7 @@ def make_plan(self, system_prompt: str, goal_prompt: str, max_retries: int = 3) prompt = f"You must generate a plan for goal:\n{goal_prompt}\nImportant: Every task description should be very detailled and specific with the full path of all input output files specified." for attempt in range(1, max_retries + 1): try: - print(f"🔄 Plan generation attempt {attempt}/{max_retries}") + print_info(f"Plan generation attempt {attempt}/{max_retries}") memory_path = getattr(self.config, 'memory_path', 'sources/memory') raw_plan = LLMProvider("plan_creator", memory_path=memory_path, system_msg=system_prompt, config=self.config_llm, use_flat_cache=True)(prompt, use_cache=True) @@ -87,41 +92,40 @@ def make_plan(self, system_prompt: str, goal_prompt: str, max_retries: int = 3) if not raw_plan or not isinstance(raw_plan, str): raise ValueError("LLM returned empty or invalid response") - print(f"📝 Received plan response ({len(raw_plan)} characters)") + print_info(f"Received plan response ({len(raw_plan)} characters)") print(raw_plan) print("---") plan_dict = self._extract_json_from_code_block(raw_plan) if plan_dict is None: raise ValueError("Failed to extract valid JSON from LLM response\n") plan = self._parse_and_validate_plan(plan_dict, goal_prompt) - print(f"✅ Successfully generated and validated plan with {len(plan.steps)} steps") + print_ok(f"Plan generated and validated — {len(plan.steps)} step(s)") return plan except (ValueError, PlanValidationError, json.JSONDecodeError) as e: last_error = e error_msg = str(e) - # Check if this might be a truncation error (unterminated string at end of response) is_truncation = False if "Unterminated string" in error_msg or "Unexpected end of data" in error_msg: is_truncation = True error_msg = f"{error_msg} (This often indicates the response was truncated due to max_tokens limit)" - print(f"⚠️ Attempt {attempt} failed: {error_msg}") + print_warn(f"Attempt {attempt} failed: {error_msg}") if is_truncation: - print(f"💡 Tip: Consider increasing max_tokens in your config (current: {getattr(self.config, 'max_tokens', 'not set')})") + print_info(f"Tip: consider increasing max_tokens (current: {getattr(self.config, 'max_tokens', 'not set')})") if attempt < max_retries: wait_time = 2 ** attempt - print(f"⏳ Waiting {wait_time} seconds before retry...") + print_info(f"Waiting {wait_time}s before retry…") time.sleep(wait_time) if attempt > 1: goal_prompt = self._enhance_prompt_with_error(goal_prompt, error_msg) else: - print(f"❌ All {max_retries} attempts failed") + print_err(f"All {max_retries} attempts failed") except Exception as e: last_error = e - print(f"❌ Unexpected error in attempt {attempt}: {str(e)}") + print_err(f"Unexpected error in attempt {attempt}: {str(e)}") if attempt >= max_retries: break time.sleep(2 ** attempt) @@ -287,18 +291,18 @@ def _display_plan(self, plan: Plan) -> None: Args: plan: The plan to display """ - print(f"\n📋 Execution Plan: {plan.goal}") - print("=" * 80) + print_phase("📋 EXECUTION PLAN") + print(f" {BOLD}Goal:{RESET} {plan.goal}\n") for i, step in enumerate(plan.steps, 1): - print(f"\n{i}. {step.name.upper()} [{step.complexity}]") - print(f" Task: {step.task}") + print(f" {CYAN}{BOLD}{i}. {step.name.upper()}{RESET} {DIM}[{step.complexity}]{RESET}") + print(f" {step.task}") if step.depends_on: - print(f" Dependencies: {', '.join(step.depends_on)}") + print(f" {DIM}Depends on: {', '.join(step.depends_on)}{RESET}") if step.required_inputs: - print(f" Required Inputs: {', '.join(step.required_inputs)}") + print(f" {DIM}Inputs: {', '.join(step.required_inputs)}{RESET}") if step.expected_outputs: - print(f" Expected Outputs: {', '.join(step.expected_outputs)}") - print("\n" + "=" * 80) + print(f" {DIM}Outputs: {', '.join(step.expected_outputs)}{RESET}") + print() def _request_human_plan_validation(self, plan: Plan) -> tuple[bool, str]: """ @@ -308,22 +312,17 @@ def _request_human_plan_validation(self, plan: Plan) -> tuple[bool, str]: - is_approved: True if human pressed Enter (approve), False otherwise - feedback: User's correction/feedback if plan not approved """ - print("\n" + "_" * 40) - print("👤 HUMAN VALIDATION REQUIRED") - print("_" * 40) - print("\nPlease review the plan above.") - print("\nOptions:") - print(" • Press [ENTER] to approve and execute the plan") - print(" • Type your corrections/feedback and press [ENTER] to regenerate") - print("\n" + "─" * 80) + print_section("👤 HUMAN VALIDATION REQUIRED") + print(f" {DIM}Please review the plan above.{RESET}") + print(f" {DIM}Press [ENTER] to approve · Type feedback and [ENTER] to regenerate{RESET}\n") - user_input = input("\n👉 Your decision: ").strip() + user_input = input(f" {BOLD}➤ Your decision: {RESET}").strip() if not user_input: - print("\n✅ Plan approved by human. Proceeding with execution...") + print_ok("Plan approved. Proceeding with execution…") return True, "" else: - print(f"\n📝 Feedback received: {user_input}") - print("🔄 Will regenerate plan based on your feedback...") + print_info(f"Feedback received: {user_input}") + print_info("Regenerating plan based on your feedback…") return False, user_input def _generate_plan_with_human_validation(self, goal: str, human_approve = False) -> Plan: @@ -369,7 +368,7 @@ def _init_visualization(self, plan: Plan) -> None: try: self.visualizer = PlannerVisualizer(plan) - print("🎨 Visualization window initialized") + print_ok("Visualization window initialized") # On Linux and Windows, use a separate thread for event handling. # On macOS, event handling must be done from main thread (will be called periodically) @@ -384,13 +383,13 @@ def visualization_loop(): self.visualizer_thread = threading.Thread(target=visualization_loop, daemon=True) self.visualizer_thread.start() platform_name = "Windows" if self.is_windows else "Linux" - print(f"🎨 Visualization running in separate thread ({platform_name})") + print_info(f"Visualization running in separate thread ({platform_name})") else: - print("🎨 Visualization will update from main thread (macOS)") + print_info("Visualization will update from main thread (macOS)") except Exception as e: - print(f"⚠️ Could not initialize visualization: {str(e)}") - print("⚠️ Continuing without visualization...") + print_warn(f"Could not initialize visualization: {str(e)}") + print_warn("Continuing without visualization…") self.use_visualization = False self.visualizer = None @@ -410,7 +409,7 @@ def _update_visualization(self, total_cost: float = 0.0) -> None: self.visualizer.update_tasks(self.task_history, total_cost=total_cost) except Exception as e: - print(f"⚠️ Error updating visualization: {str(e)}") + print_warn(f"Error updating visualization: {str(e)}") def _cleanup_visualization(self) -> None: """ @@ -428,9 +427,9 @@ def _cleanup_visualization(self) -> None: self.visualizer_thread.join(timeout=0.5) self.visualizer = None self.visualizer_thread = None - print("🎨 Visualization window closed") + print_ok("Visualization window closed") except Exception as e: - print(f"⚠️ Error closing visualization: {str(e)}") + print_warn(f"Error closing visualization: {str(e)}") def _get_workspace_files(self) -> list[str]: @@ -448,7 +447,7 @@ def _get_workspace_files(self) -> list[str]: try: workspace_path = Path(self.workspace_path) if not workspace_path.exists(): - print(f"⚠️ Workspace path does not exist: {self.workspace_path}") + print_warn(f"Workspace path does not exist: {self.workspace_path}") return files for root, dirs, filenames in os.walk(workspace_path): @@ -464,7 +463,7 @@ def _get_workspace_files(self) -> list[str]: except ValueError: continue except Exception as e: - print(f"⚠️ Error scanning workspace files: {str(e)}") + print_warn(f"Error scanning workspace files: {str(e)}") return files @@ -473,7 +472,7 @@ def _capture_workspace_snapshot(self) -> None: Capture a snapshot of current workspace files before step execution. """ self._workspace_files_before_step = self._get_workspace_files() - print(f"📸 Captured workspace snapshot: {len(self._workspace_files_before_step)} files") + print_info(f"Workspace snapshot: {len(self._workspace_files_before_step)} file(s)") return self._workspace_files_before_step def _verify_required_inputs(self, step: PlanStep) -> tuple[bool, list[str]]: @@ -570,9 +569,9 @@ async def evolve_runs(self, task, judge, max_evolve_iteration, cached_wf_allow=T if max_evolve_iteration is None or max_evolve_iteration < 1: max_evolve_iteration = 1 - print(f"⚠️ Invalid max_evolve_iteration, using default: {max_evolve_iteration}") + print_warn(f"Invalid max_evolve_iteration, using default: {max_evolve_iteration}") - print(f"🎯 Starting Iterative-Learning for task: {task[:50]}...") + print_info(f"Starting Iterative-Learning for task: {task[:60]}…") try: # Use original_task for lookup to avoid knowledge wrapper interference @@ -586,10 +585,10 @@ async def evolve_runs(self, task, judge, max_evolve_iteration, cached_wf_allow=T if past_wf_lookups and len(past_wf_lookups) > 0: best_match = past_wf_lookups[0] if best_match is None: - print("⚠️ Best match is None, proceeding with new Evolution run") + print_warn("Best match is None, proceeding with new Evolution run") #elif self._get_evolve_success(best_match): elif best_match.is_success: - print(f"🔁 Using previously run workflow result with UUID: {getattr(best_match, 'uuid', 'N/A')}") + print_ok(f"Using cached workflow result UUID: {getattr(best_match, 'uuid', 'N/A')}") run = IndividualRun( goal=best_match.goal, @@ -603,7 +602,7 @@ async def evolve_runs(self, task, judge, max_evolve_iteration, cached_wf_allow=T return [run] # Generate new workflows via Evolution - print(f"🔄 No cached run found, starting task learning (max_iter: {max_evolve_iteration})") + print_info(f"No cached run found, starting task learning (max_iter: {max_evolve_iteration})") if self.dgm is None: raise ValueError("❌ Planner: instance is None") @@ -618,7 +617,7 @@ async def evolve_runs(self, task, judge, max_evolve_iteration, cached_wf_allow=T ) if runs is None: - print("⚠️ Runs is None, returning empty list") + print_warn("Runs is None, returning empty list") return [] return runs @@ -644,7 +643,7 @@ async def run_attempts(self, attempt_counts, max_attempts, step, judge, max_evol PlanStep: The updated step with execution status """ if step is None: - print("❌ Step is None, cannot execute") + print_err("Step is None, cannot execute") return step if attempt_counts is None: @@ -652,7 +651,7 @@ async def run_attempts(self, attempt_counts, max_attempts, step, judge, max_evol if max_attempts is None or max_attempts < 1: max_attempts = 1 - print(f"⚠️ Invalid max_attempts, using default: {max_attempts}") + print_warn(f"Invalid max_attempts, using default: {max_attempts}") step_name = getattr(step, 'name', 'unknown_step') goal = getattr(step, 'goal_context', '') @@ -665,7 +664,7 @@ async def run_attempts(self, attempt_counts, max_attempts, step, judge, max_evol attempt += 1 attempt_counts[step_name] = attempt - print(f"🔄 Attempt {attempt}/{max_attempts} for task: {step_name}") + print_info(f"Attempt {attempt}/{max_attempts} for task: {step_name}") if self.tts: self.tts.speak(f"now starting task {step_name}", voice_index=0) @@ -716,13 +715,13 @@ async def run_attempts(self, attempt_counts, max_attempts, step, judge, max_evol outputs_produced, missing_outputs = self._verify_expected_outputs(step) step.status = TaskStatus.COMPLETED if outputs_produced: - print(f"✅ Task '{step_name}' completed successfully") + print_ok(f"Task '{step_name}' completed successfully") break else: - print(f"⚠️ Task '{step_name}' completed but missing expected outputs: {missing_outputs}") + print_warn(f"Task '{step_name}' completed but missing expected outputs: {missing_outputs}") break else: - print(f"❌ Task {step_name} (uuid: {final_uuid}) failed with score {attempt_score}\n") + print_err(f"Task {step_name} (uuid: {final_uuid}) failed with score {attempt_score}") if self.tts: self.tts.speak(f"Task {step_name} failure, retrying...", voice_index=0) continue @@ -763,7 +762,7 @@ async def start_planner( raise ValueError("❌ Planner: Goal must be a non-empty string") goal = "\nAvailable files:\n" + list_files(self.config.workspace_dir) + "\n" + goal - print(f"▶ Starting planner with goal: {goal}") + print_info(f"Starting planner with goal: {goal[:80]}…") try: # Generate plan with human validation loop @@ -776,7 +775,7 @@ async def start_planner( self._init_visualization(self.current_plan) # Check for stop condition if self._check_stop_condition(self.current_plan): - print("⏹️ Stop condition found in plan. Exiting.") + print_info("Stop condition found in plan. Exiting.") return self.task_history # Validate plan has steps @@ -790,17 +789,17 @@ async def start_planner( total_cost = 0 for step_idx, step in enumerate(self.current_plan.steps): if step is None: - print(f"⚠️ Step {step_idx + 1} is None, skipping") + print_warn(f"Step {step_idx + 1} is None, skipping") continue step_name = getattr(step, 'name', f'step_{step_idx}') - print(f"\n{'='*60}") - print(f"📋 Executing Step {step_idx + 1}/{len(self.current_plan.steps)}: {step_name}") - print(f"{'='*60}") + print_phase( + f"📋 STEP {step_idx + 1}/{len(self.current_plan.steps)} · {step_name}", + ) # Check if step can be executed (dependencies satisfied) if lst_step: can_execute, missing_deps = self._can_execute_step(lst_step) if not can_execute: - self.request_user_exit(f"⚠️ Cannot execute step '{step_name}' - missing dependencies: {missing_deps}") + self.request_user_exit(f"Cannot execute step '{step_name}' — missing dependencies: {missing_deps}") # Execute the step with retry logic step.status = TaskStatus.RUNNING @@ -819,7 +818,6 @@ async def start_planner( if step.status != TaskStatus.COMPLETED: step.status = TaskStatus.FAILED - # Send notification for task failure self.notifier.send_message( f"Task '{step_name}' failed after {max_attempts} attempts\n" f"Goal: {goal[:128]}...\n" @@ -829,10 +827,15 @@ async def start_planner( ) raise Exception(f"❌ Giving up on task '{step_name}' after {max_attempts} attempts") - print(f"\n🏁 Planner execution completed. Executed {len(self.task_history)} tasks. Cost: {total_cost}") - - # Send success notification completed_tasks = sum(1 for t in self.task_history if t.status == TaskStatus.COMPLETED) + print_summary( + "🏁 PLANNER EXECUTION COMPLETE", + [ + ("Tasks executed", str(len(self.task_history))), + ("Tasks completed", str(completed_tasks)), + ("Total cost", f"${total_cost:.6f}"), + ], + ) self.notifier.send_message( f"Planner completed successfully!\n" f"Goal: {goal[:128]}...\n" @@ -846,7 +849,7 @@ async def start_planner( return self.task_history except Exception as e: - print(f"❌ Critical error in planner execution: {str(e)}") + print_err(f"Critical error in planner execution: {str(e)}") self.notifier.send_message(str(e), title="error during Mimosa execution.") self._cleanup_visualization() raise ValueError(f"❌ Planner: Execution failed: {str(e)}") from e diff --git a/sources/core/workflow_factory.py b/sources/core/workflow_factory.py index 4b43b741..f662ea8b 100644 --- a/sources/core/workflow_factory.py +++ b/sources/core/workflow_factory.py @@ -12,6 +12,7 @@ from .llm_provider import LLMConfig, LLMProvider, extract_model_pattern from .tools_manager import ToolManager +from sources.cli.pretty_print import print_ok, print_info, print_warn, print_err class WorkflowFactory: @@ -94,7 +95,7 @@ async def load_tools_code(self) -> tuple[str, str]: client_prompt = tool_manager.get_client_prompt(mcp) tools_code += client_code + "\n" existing_tool_prompt += client_prompt + "\n" - print(f"🔧 Discovered {len(mcps)} MCP servers capabilities. Workflow generation can start.") + print_ok(f"Discovered {len(mcps)} MCP server(s) — workflow generation can start.") return tools_code, existing_tool_prompt def remove_imports(self, code: str) -> str: @@ -209,22 +210,22 @@ def create_workflow_code( self.logger.info("Generating workflow code with LLM...") system_prompt = self.get_system_prompt() try: - print("📝 Step 1/2: Generating prompts code...") + print_info("Step 1/2: Generating prompts code…") llm_output = self.llm_make_prompts( system_prompt, craft_instructions, existing_tool_prompt, path, allow_cache ) prompts_code = self.extract_python_code(llm_output) commentary = llm_output.replace(prompts_code, "").split("```python")[0] - print("💬 LLM commentary on prompt:") + print_info("LLM commentary on prompts:") print(commentary) - print("🔧 Step 2/2: Generating workflow code...") + print_info("Step 2/2: Generating workflow code…") llm_output = self.llm_make_workflow( system_prompt, craft_instructions, existing_tool_prompt, path, prompts_code, allow_cache ) workflow_code = self.extract_python_code(llm_output) commentary = llm_output.replace(workflow_code, "").split("```python")[0] - print("💬 LLM commentary on workflow:") + print_info("LLM commentary on workflow:") print(commentary) workflow_code = prompts_code + "\n\n" + workflow_code @@ -295,7 +296,7 @@ def validate_workflow_structure(self, workflow_code: str) -> None: entry_node = start_match.group(1) if entry_node not in nodes: raise ValueError(f"START targets non-existent node '{entry_node}'") - self.logger.debug(f"🚀 Workflow entry point: START → {entry_node}") + self.logger.debug(f"Workflow entry point: START → {entry_node}") self.logger.info("✅ Workflow structure validation passed") @@ -382,7 +383,6 @@ def assemble_workflow( # Generated workflow {workflow_code} -print("worflow run: compiling workflow...") app = workflow.compile() # Initialize and execute workflow @@ -392,18 +392,14 @@ def assemble_workflow( if WORKFLOW_PATH: try: png = app.get_graph().draw_mermaid_png() - print("workflow run: saving workflow graph as PNG at ", WORKFLOW_PATH) with open(os.path.join(WORKFLOW_PATH, "workflow_{uuid_str}.png"), "wb") as f: - print("workflow run: writing PNG file...") f.write(png) - print("PNG saved at ", os.path.join(WORKFLOW_PATH, "workflow_{uuid_str}.png")) except Exception as e: RuntimeError(f"Could not save workflow graph:" + str(e)) except Exception as e: print(f"❌ Error saving PNG workflow:" + str(e)) pass -print("workflow run: invoking workflow...") try: result_state = app.invoke(initial_state) except KeyboardInterrupt: @@ -587,13 +583,11 @@ def get_engine_code(self) -> str: engine_name = {self.config.engine_name!r} engine = None if engine_name == "mlx": - print("Using MLXModel for local execution.") engine = MLXModel( model_id=model_id, max_tokens=max_tokens, ) elif engine_name == "inference_client": - print("Using InferenceClientModel for inference client execution.") if not token: raise ValueError("Hugging Face token is required. Please set the HF_TOKEN environment variable or pass a token.") engine = InferenceClientModel( @@ -710,7 +704,6 @@ async def craft_single_agent(self, goal: str, original_task: str = None): ) def save_agent_memories(agent, memory_path: str, agent_name: str): - print(f"Saving agent memory to: {{{{memory_path}}}}") try: memories = [] for idx, step in enumerate(agent.memory.steps): diff --git a/sources/core/workflow_runner.py b/sources/core/workflow_runner.py index 13064ac1..09891306 100644 --- a/sources/core/workflow_runner.py +++ b/sources/core/workflow_runner.py @@ -3,8 +3,11 @@ """ import asyncio +import fcntl import logging import os +import pty +import sys import time from collections.abc import Callable from dataclasses import dataclass @@ -197,37 +200,89 @@ async def execute( script_path = os.path.abspath(os.path.join(self.config.temp_dir, f"{execution_id}.py")) with open(script_path, "w") as f: f.write(code) - print(f"Executing script: {script_path}") cmd = [*self._python_cmd, script_path] return await self._run_command(cmd, execution_id, progress_callback) + @staticmethod + def _build_color_env() -> dict[str, str]: + """Build an environment dict that encourages color output. + + Copies the host environment and adds variables commonly checked by + CLI tools and Python libraries (rich, click, tqdm, pytest, …) to + force colored output even when stdout is not a real TTY. + """ + env = dict(os.environ) + env.setdefault("TERM", "xterm-256color") + env["FORCE_COLOR"] = "1" # chalk / supports-color (Node & Python) + env["PY_COLORS"] = "1" # pytest, tox, … + env["CLICOLOR_FORCE"] = "1" # BSD / GNU convention + env["PYTHONUNBUFFERED"] = "1" # disable Python output buffering + return env + + def _pty_available(self) -> bool: + """Return True when pseudo-terminal support can be used.""" + return sys.platform != "win32" + async def _run_command( self, cmd: list[str], execution_id: str | None = None, progress_callback: Callable[[str], None] | None = None, ) -> ExecutionResult: - """Core async command execution with monitoring.""" + """Core async command execution with monitoring. + + On platforms that support PTYs (Linux / macOS) the subprocess stdout + is connected to a pseudo-terminal so that child processes see + ``isatty(1) == True`` and emit ANSI colour codes. Stderr is still + captured via a regular pipe. + """ start_time = asyncio.get_event_loop().time() + master_fd = slave_fd = -1 try: - process = await asyncio.create_subprocess_exec( - *cmd, - stdout=asyncio.subprocess.PIPE, - stderr=asyncio.subprocess.PIPE, - limit=1024 * 1024, # 1MB buffer limit - env=dict(os.environ), # Pass host environment variables - cwd=self.execution_dir, # Set working directory for execution - ) + env = self._build_color_env() + + if self._pty_available(): + # --- PTY path: child stdout goes through a pseudo-terminal --- + master_fd, slave_fd = pty.openpty() + + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=slave_fd, + stderr=asyncio.subprocess.PIPE, + env=env, + cwd=self.execution_dir, + ) + # Parent no longer needs the slave side; the child inherited it. + os.close(slave_fd) + slave_fd = -1 - if execution_id: - self._active_processes[execution_id] = process + if execution_id: + self._active_processes[execution_id] = process - stdout_data, stderr_data = await asyncio.wait_for( - self._stream_output(process, progress_callback), - timeout=self.config.timeout, - ) + stdout_data, stderr_data = await asyncio.wait_for( + self._stream_output_pty(process, master_fd, progress_callback), + timeout=self.config.timeout, + ) + else: + # --- Pipe fallback (Windows or if PTY unavailable) --- + process = await asyncio.create_subprocess_exec( + *cmd, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, + limit=1024 * 1024, + env=env, + cwd=self.execution_dir, + ) + + if execution_id: + self._active_processes[execution_id] = process + + stdout_data, stderr_data = await asyncio.wait_for( + self._stream_output_pipe(process, progress_callback), + timeout=self.config.timeout, + ) await process.wait() execution_time = asyncio.get_event_loop().time() - start_time @@ -258,27 +313,83 @@ async def _run_command( return ExecutionResult(ExecutionStatus.FAILED, -1, "", str(e), 0.0) finally: + if slave_fd >= 0: + os.close(slave_fd) + if master_fd >= 0: + os.close(master_fd) if execution_id and execution_id in self._active_processes: del self._active_processes[execution_id] - async def _stream_output( + # ------------------------------------------------------------------ + # Output streaming helpers + # ------------------------------------------------------------------ + + async def _stream_output_pty( self, process: asyncio.subprocess.Process, + master_fd: int, progress_callback: Callable[[str], None] | None = None, ) -> tuple[str, str]: - """Stream process output with real-time callbacks.""" + """Stream stdout from a PTY master fd and stderr from a pipe. - stdout_lines = [] - stderr_lines = [] + The PTY preserves ANSI escape sequences (colours, bold, …) because + the child process sees a real terminal on its stdout. + """ + loop = asyncio.get_event_loop() + stdout_chunks: list[str] = [] + stderr_lines: list[str] = [] + stdout_done = asyncio.Event() - async def read_stdout(): + # Make the master fd non-blocking so we can use add_reader. + flags = fcntl.fcntl(master_fd, fcntl.F_GETFL) + fcntl.fcntl(master_fd, fcntl.F_SETFL, flags | os.O_NONBLOCK) + + def _on_master_readable() -> None: + """Called by the event loop when data is available on the PTY.""" + try: + data = os.read(master_fd, 65536) + if not data: + loop.remove_reader(master_fd) + stdout_done.set() + return + text = data.decode("utf-8", errors="replace") + stdout_chunks.append(text) + if progress_callback: + for line in text.splitlines(): + progress_callback(line) + except OSError: + # EIO is expected when the slave side is closed (child exited). + loop.remove_reader(master_fd) + stdout_done.set() + + loop.add_reader(master_fd, _on_master_readable) + + async def _read_stderr() -> None: + async for raw_line in process.stderr: + stderr_lines.append(raw_line.decode("utf-8", errors="replace")) + + # Wait for both stdout (PTY) and stderr (pipe) to finish. + await asyncio.gather(stdout_done.wait(), _read_stderr()) + + return "".join(stdout_chunks), "".join(stderr_lines) + + async def _stream_output_pipe( + self, + process: asyncio.subprocess.Process, + progress_callback: Callable[[str], None] | None = None, + ) -> tuple[str, str]: + """Fallback: stream stdout/stderr when both are plain pipes.""" + stdout_lines: list[str] = [] + stderr_lines: list[str] = [] + + async def read_stdout() -> None: async for line in process.stdout: line_str = line.decode("utf-8", errors="replace") stdout_lines.append(line_str) if progress_callback: progress_callback(line_str.rstrip()) - async def read_stderr(): + async def read_stderr() -> None: async for line in process.stderr: stderr_lines.append(line.decode("utf-8", errors="replace")) @@ -324,7 +435,12 @@ async def main(): runner = WorkflowRunner(config) await runner.install_dependencies(["requests", "numpy"]) code = """ -print("Hello from the workflow runner!") +import sys + +# Demonstrate that color is preserved through the PTY +print("\\033[32m✔ Hello from the workflow runner (green)\\033[0m") +print("\\033[1;34mBold blue text\\033[0m") +print(f"stdout is a TTY: {sys.stdout.isatty()}") """ def progress_handler(line: str): diff --git a/sources/core/workflow_selection.py b/sources/core/workflow_selection.py index a22c1c8b..eb0e542f 100644 --- a/sources/core/workflow_selection.py +++ b/sources/core/workflow_selection.py @@ -12,7 +12,7 @@ def __init__(self, config: Config) -> None: self.config = config self.workflows_folder = Path(config.workflow_dir) self.workflows_info = self.discover_workflows() - self.model = SentenceTransformer("all-MiniLM-L6-v2") + self.model = SentenceTransformer("all-MiniLM-L6-v2", token=False) def discover_workflows(self) -> dict[str, WorkflowInfo]: workflows = {} @@ -20,29 +20,27 @@ def discover_workflows(self) -> dict[str, WorkflowInfo]: if not self.workflows_folder.exists(): print(f"Workflows directory {self.workflows_folder} does not exist.") return workflows - + for workflow_folder in self.workflows_folder.iterdir(): if not workflow_folder.is_dir(): continue - + uuid = workflow_folder.name workflow_info = WorkflowInfo(uuid, workflow_folder) - + if not workflow_info.is_valid(): continue - + # Check if state_result is empty if not workflow_info.load_state_result(): - print(f"Skipping workflow {uuid}: empty state_result.json") continue - + workflow_info.load_code() if not workflow_info.code: - print(f"Skipping workflow {uuid}: unable to load code") continue - + workflows[uuid] = workflow_info - + return workflows def cosine_similarity(self, a: str, b: str) -> float: @@ -61,12 +59,12 @@ def sort_similar_workflows( self, goal: str, threshold=0.8, debug=False ) -> list[WorkflowInfo]: """Find workflows with similar goals using original unwrapped tasks. - + Args: goal: The task to match against (will be compared with original_task of workflows) threshold: Minimum similarity score (0.0-1.0) debug: Whether to print debug information - + Returns: list[WorkflowInfo]: Workflows sorted by similarity, filtered by threshold """ @@ -75,14 +73,14 @@ def sort_similar_workflows( if not self.workflows_info: print("No workflows found.") return [] - + # Use original_task for comparison to avoid knowledge wrapper interference similar_workflows = sorted( self.workflows_info.values(), key=lambda wf: self.cosine_similarity(wf.original_task[-512:], goal[-512:]), reverse=True, ) - + if debug: for wf in similar_workflows: sim = self.cosine_similarity(wf.original_task[-512:], goal[-512:]) @@ -90,13 +88,13 @@ def sort_similar_workflows( f"Original Task:\n{wf.original_task[:512]}\n" f"Target:\n{goal[:512]}\n" f"Similarity: {sim:.4f}\n---\n") - + return [ wf for wf in similar_workflows if self.cosine_similarity(wf.original_task[-512:], goal[-512:]) >= threshold ] - + def sort_workflows_by_score( self, workflows_info: list[WorkflowInfo], threshold: float ) -> list[WorkflowInfo]: diff --git a/sources/evaluation/csv_mode.py b/sources/evaluation/csv_mode.py index 48f20b94..e93535b1 100644 --- a/sources/evaluation/csv_mode.py +++ b/sources/evaluation/csv_mode.py @@ -17,6 +17,10 @@ from sources.evaluation.capsule_evaluator import CapsuleEvaluator from sources.utils.transfer_toolomics import LocalTransfer from sources.utils.list_files import list_files +from sources.cli.pretty_print import ( + print_ok, print_warn, print_err, print_info, + print_phase, print_summary, +) class CsvEvaluationMode: """ @@ -172,7 +176,7 @@ def _restore_execution_history_from_cache(self, cached_notes: dict) -> None: f"[CACHE RECOVERY] Restored {total_eval} evaluations: " f"VER={ver_success}, SR={sr_success}, CBS={avg_cbs:.3f}" ) - print(f"\033[95mRestored {total_eval} previous evaluations from cache\033[0m") + print_info(f"Restored {total_eval} previous evaluations from cache") def _save_run_notes(self, capsule_name: str, goal: str, analysis: dict, execution_time: float) -> None: @@ -312,12 +316,12 @@ def sab_files_transfer(self, sab_loader, file_transfer, row): file_transfer.clean_workspace() task_dataset_path = sab_loader.get_dataset_path(row) self.logger.info(f"[PAPERS DATASET MODE] Transferring dataset from: {task_dataset_path}") - print(f"\033[95m📁 Transferring dataset: {task_dataset_path.name}\033[0m") + print_info(f"📁 Transferring dataset: {task_dataset_path.name}") files_transferred = file_transfer.transfer_files_to_workspace(str(task_dataset_path)) time.sleep(0.5) # Give filesystem a moment to sync workspace_files_after = file_transfer.count_files_recursive(Path(file_transfer.workspace_path)) - print(f"\033[95m✓ Transferred {files_transferred} files to workspace\033[0m") - print(f"\033[95m📊 Verification: {workspace_files_after} files present in workspace\033[0m") + print_ok(f"Transferred {files_transferred} file(s) to workspace") + print_info(f"Verification: {workspace_files_after} file(s) present in workspace") if workspace_files_after == 0: raise ValueError( @@ -349,7 +353,7 @@ def _evaluate_with_science_agent_bench( Updated execution_data dictionary with evaluation metrics """ try: - print("\033[95m📊 Evaluating results with ScienceAgentBench metrics...\033[0m") + print_info("📊 Evaluating results with ScienceAgentBench metrics…") api_cost = runs[-1].cost if runs and hasattr(runs[-1], 'cost') else 0.0 evaluator = CapsuleEvaluator( @@ -370,7 +374,7 @@ def _evaluate_with_science_agent_bench( 'eval_cost': eval_results['cost'], 'runs': runs }) - print(f"\033[95m{eval_results['summary']}\033[0m") + print_ok(eval_results['summary']) self.logger.info( f"[SAB EVAL] Task {row.get('instance_id')}: " @@ -382,7 +386,7 @@ def _evaluate_with_science_agent_bench( except Exception as eval_error: self.logger.error(f"[SAB EVAL] Evaluation error: {str(eval_error)}") - print(f"\033[91m⚠️ Evaluation failed: {str(eval_error)}\033[0m") + print_err(f"Evaluation failed: {str(eval_error)}") execution_data.update({ 'VER': False, 'SR': False, @@ -426,19 +430,18 @@ async def run_autonomous_eval_loop(self, dataset_type: str, dataset_path: str, l csvfile.seek(0) reader = csv.DictReader(csvfile) self.logger.info(f"[PAPERS DATASET MODE] Starting autonomous loop for {total_rows} CSV entry") - print(f"\n\033[95m{'🤖 RUN ON PAPERS DATASETS':^80}\033[0m") - print(f"\033[95m{'=' * 80}\033[0m") + print_phase("🤖 RUN ON PAPERS DATASETS") for i, row in enumerate(reader): if i < start_row: - print("Skipping evaluation (using cache) for :", i+1) + print_info(f"Skipping evaluation (using cache) for row {i + 1}") continue if i >= self.csv_runs_limit: break try: iteration_start_time = time.time() goal, scenario_id, scenario_rubric_filename = self._generate_next_task(row, dataset_type) - print(f"\033[95m📋 GOAL: {goal}\033[0m") - print(f"\033[95m📄 Scenario Rubric: {scenario_rubric_filename}\033[0m") + print_info(f"📋 GOAL: {goal[:120]}…" if len(goal) > 120 else f"📋 GOAL: {goal}") + print_info(f"📄 Scenario Rubric: {scenario_rubric_filename}") if dataset_type == "science_agent_bench" and sab_loader: self.sab_files_transfer(sab_loader, file_transfer, row) @@ -457,10 +460,10 @@ async def run_autonomous_eval_loop(self, dataset_type: str, dataset_path: str, l max_task_retry=3 ) results_str = self._format_goal_mode_results(tasks_data) - print("\033[95m📊 Transfering results files...\033[0m") + print_info("📦 Transferring results files…") trs = LocalTransfer(config=self.config, workspace_path=self.config.workspace_dir, runs_capsule_dir=self.config.runs_capsule_dir) capsule_name = trs.transfer_workspace_files_to_capsule(goal) - print("\033[95m📊 Analyzing results...\033[0m") + print_info("📊 Analyzing results…") execution_time = time.time() - iteration_start_time analysis = self._analyze_results(goal, results_str, execution_time) execution_data = { @@ -486,12 +489,12 @@ async def run_autonomous_eval_loop(self, dataset_type: str, dataset_path: str, l analysis, execution_time ) - print(f"\033[95m✅ Iteration {i + 1} completed\033[0m") - print(f"\033[95m Success Level: {analysis.get('success_level', 'Unknown')}\033[0m") - print(f"\033[95m Time: {execution_time:.2f}s\033[0m") + print_ok(f"Iteration {i + 1} completed") + print_info(f" Success Level: {analysis.get('success_level', 'Unknown')}") + print_info(f" Time: {execution_time:.2f}s") except Exception as e: self.logger.error(f"[PAPERS DATASET MODE] Error in csv row {i + 1}: {str(e)}") - print(f"\033[91m❌ Error in csv row {i + 1}: {str(e)}\033[0m") + print_err(f"Error in csv row {i + 1}: {str(e)}") continue self._print_final_summary() @@ -504,38 +507,43 @@ def _print_final_summary(self) -> None: successful_runs = [exec_data for exec_data in current_runs if exec_data.get("success_level") in ["High", "Medium"]] - print(f"\n\033[95mSUMMARY step: {len(current_runs)}\033[0m") - print(f"\033[95m{'=' * 80}\033[0m") - print(f"\033[95mSuccessful runs: {len(successful_runs)}\033[0m") - if len(current_runs) > 0: - print(f"\033[95mSuccess rate: {len(successful_runs)/len(current_runs)*100:.1f}%\033[0m") + + success_rate = ( + f"{len(successful_runs)/len(current_runs)*100:.1f}%" + if current_runs else "N/A" + ) + rows = [ + ("Steps evaluated", str(len(current_runs))), + ("Successful runs", str(len(successful_runs))), + ("Success rate", success_rate), + ] # For SAB metrics, also exclude cached entries - sab_runs = [exec_data for exec_data in current_runs - if 'VER' in exec_data] + sab_runs = [exec_data for exec_data in current_runs if 'VER' in exec_data] if sab_runs: - print(f"\033[95m\n{'ScienceAgentBench Metrics':^80}\033[0m") - print(f"\033[95m{'-' * 80}\033[0m") ver_success = sum(1 for run in sab_runs if run.get('VER', False)) sr_success = sum(1 for run in sab_runs if run.get('SR', False)) avg_cbs = sum(run.get('CBS', 0.0) for run in sab_runs) / len(sab_runs) total_cost = sum(run.get('eval_cost', 0.0) for run in sab_runs) + rows += [ + ("── ScienceAgentBench ──", ""), + ("VER (Valid Exec Rate)", f"{ver_success}/{len(sab_runs)} ({ver_success/len(sab_runs)*100:.1f}%)"), + ("SR (Success Rate)", f"{sr_success}/{len(sab_runs)} ({sr_success/len(sab_runs)*100:.1f}%)"), + ("CBS (CodeBERT avg)", f"{avg_cbs:.3f}"), + ("Total API Cost", f"${total_cost:.4f}"), + ("Avg cost/task", f"${total_cost/len(sab_runs):.4f}"), + ] - print(f"\033[95mVER (Valid Execution Rate): {ver_success}/{len(sab_runs)} ({ver_success/len(sab_runs)*100:.1f}%)\033[0m") - print(f"\033[95mSR (Success Rate): {sr_success}/{len(sab_runs)} ({sr_success/len(sab_runs)*100:.1f}%)\033[0m") - print(f"\033[95mCBS (CodeBERT Score) Average: {avg_cbs:.3f}\033[0m") - print(f"\033[95mTotal API Cost: ${total_cost:.4f}\033[0m") - print(f"\033[95mAverage API Cost per Task: ${total_cost/len(sab_runs):.4f}\033[0m") - print(f"\033[95m{'=' * 80}\033[0m\n") + print_summary("📊 EVALUATION SUMMARY", rows) async def start_evaluation(self, dataset_type: str = "default", dataset_path = "datasets/our_benchmark.csv", learning=False, single_agent_mode=False) -> None: """Public method to start the autonomous mode.""" try: await self.run_autonomous_eval_loop(dataset_type, dataset_path, learning, single_agent_mode) except KeyboardInterrupt: - print("\n\033[95m⚠️ Autonomous mode interrupted by user\033[0m") + print_warn("Autonomous mode interrupted by user") self._print_final_summary() except Exception as e: self.logger.error(f"[PAPERS DATASET MODE] Fatal error: {str(e)}") - print(f"\033[91m❌ Fatal error in autonomous mode: {str(e)}\033[0m") + print_err(f"Fatal error in autonomous mode: {str(e)}") raise diff --git a/sources/modules/smolagent_factory.py b/sources/modules/smolagent_factory.py index e21402d2..03fb113a 100755 --- a/sources/modules/smolagent_factory.py +++ b/sources/modules/smolagent_factory.py @@ -205,13 +205,11 @@ def extend_system_prompt(self, added_prompt: str): def get_engine(self): if self.engine_name == "mlx": - print("Using MLXModel for local execution.") return MLXModel( model_id=self.model_id, max_tokens=self.max_tokens, ) elif self.engine_name == "inference_client": - print("Using InferenceClientModel for inference client execution.") if not self.token: raise ValueError("Hugging Face token is required. Please set the HF_TOKEN environment variable or pass a token.") return InferenceClientModel( @@ -221,7 +219,6 @@ def get_engine(self): max_tokens=self.max_tokens, ) elif self.engine_name == "litellm": - print(f"Using LiteLLM for {self.model_id} execution.") return LiteLLMModel( model_id=self.model_id, temperature=1.0, diff --git a/sources/utils/logging.py b/sources/utils/logging.py index 2c6f0657..401d84ed 100644 --- a/sources/utils/logging.py +++ b/sources/utils/logging.py @@ -1,33 +1,37 @@ -def setup_logging(debug=False): - """Configure logging with timing, line numbers, and log rotation.""" +def setup_logging(debug=False, disable=False): + """Configure logging with timing, line numbers, and log rotation. Set disable=True to disable all logging.""" import logging.handlers import os - + # Create logs directory logs_dir = "logs/" os.makedirs(logs_dir, exist_ok=True) - + # Configure root logger logger = logging.getLogger() logger.setLevel(logging.DEBUG if debug else logging.INFO) - + + if disable: + logger.setLevel(logging.CRITICAL + 1) + return + # Remove existing handlers to avoid duplication for handler in logger.handlers[:]: logger.removeHandler(handler) - + # Create formatter formatter = logging.Formatter( '%(asctime)s [%(levelname)8s] %(name)s:%(lineno)d - %(funcName)s() - %(message)s', datefmt='%Y-%m-%d %H:%M:%S' ) - + # Console handler console_handler = logging.StreamHandler() console_handler.setLevel(logging.DEBUG if debug else logging.INFO) console_handler.setFormatter(formatter) logger.addHandler(console_handler) - + # Rotating file handler for general logs file_handler = logging.handlers.RotatingFileHandler( os.path.join(logs_dir, 'mimosa.log'), @@ -37,7 +41,7 @@ def setup_logging(debug=False): file_handler.setLevel(logging.DEBUG) file_handler.setFormatter(formatter) logger.addHandler(file_handler) - + # Separate handler for workflow execution logs workflow_handler = logging.handlers.RotatingFileHandler( os.path.join(logs_dir, 'workflows.log'), @@ -46,16 +50,16 @@ def setup_logging(debug=False): ) workflow_handler.setLevel(logging.INFO) workflow_handler.setFormatter(formatter) - + # Add workflow handler to specific loggers workflow_loggers = [ 'sources.core.dgm', - 'sources.core.orchestrator', + 'sources.core.orchestrator', 'sources.core.workflow_factory', 'sources.core.workflow_runner', 'sources.evaluation.evaluator' ] - + for logger_name in workflow_loggers: workflow_logger = logging.getLogger(logger_name) workflow_logger.addHandler(workflow_handler) \ No newline at end of file diff --git a/sources/utils/pricing.py b/sources/utils/pricing.py index 194f5289..90075090 100644 --- a/sources/utils/pricing.py +++ b/sources/utils/pricing.py @@ -29,10 +29,10 @@ def __init__(self, config): # Common routing prefixes that should be stripped for matching ROUTING_PREFIXES = ['openrouter/', 'litellm/', 'together/', 'anyscale/'] - + def _strip_routing_prefix(self, model_name: str) -> str: """Strip common routing prefixes from model name. - + Handles cases like: - openrouter/mistralai/mistral-large-2407 -> mistralai/mistral-large-2407 - litellm/anthropic/claude-3.5-sonnet -> anthropic/claude-3.5-sonnet @@ -45,7 +45,7 @@ def _strip_routing_prefix(self, model_name: str) -> str: def _normalize_model_name(self, model_name: str) -> str: """Normalize model name for flexible matching. - + Handles variations like: - claude-haiku-4-5 vs claude-haiku-4.5 (hyphen vs dot in versions) - claude-haiku-4-5-20251001 (strips date suffixes) @@ -55,26 +55,26 @@ def _normalize_model_name(self, model_name: str) -> str: # First strip any routing prefixes normalized = self._strip_routing_prefix(model_name) normalized = normalized.lower() - + # Remove common date/version suffixes (e.g., -20251001, -2024-11-20, -v1, -001) # Match patterns like: -YYYYMMDD, -YYYY-MM-DD, -MMDD, -vX, -XXX (3+ digits at end) normalized = re.sub(r'-\d{8}$', '', normalized) # -20251001 normalized = re.sub(r'-\d{4}-\d{2}-\d{2}$', '', normalized) # -2024-11-20 normalized = re.sub(r'-\d{4}$', '', normalized) # -2501 (YYMM format) normalized = re.sub(r'-\d{3}$', '', normalized) # -001 - + # Normalize version separators: replace dots and underscores with hyphens # This makes "4.5" become "4-5" and "4_5" become "4-5" normalized = normalized.replace('.', '-').replace('_', '-') - + # Collapse multiple consecutive hyphens into one normalized = re.sub(r'-+', '-', normalized) - + return normalized.strip('-') def _find_model_by_substring(self, target_model: str) -> str | None: """Find best matching model from pricing data using flexible matching. - + Uses a multi-strategy approach: 1. Direct substring match (original behavior) 2. Normalized name matching (handles version format variations) @@ -84,7 +84,7 @@ def _find_model_by_substring(self, target_model: str) -> str | None: return None matches = [] - + # Strategy 1: Direct substring match (original logic) for available_model in self.model_pricing: if available_model in target_model: @@ -94,12 +94,12 @@ def _find_model_by_substring(self, target_model: str) -> str | None: valid_end = end_idx == len(target_model) or target_model[end_idx] in ['/', '-', ':'] if valid_start and valid_end: matches.append((available_model, len(available_model), 1)) # priority 1 (best) - + # Strategy 2: Normalized name matching normalized_target = self._normalize_model_name(target_model) for available_model in self.model_pricing: normalized_available = self._normalize_model_name(available_model) - + # Check if normalized available model is contained in normalized target if normalized_available in normalized_target: idx = normalized_target.find(normalized_available) @@ -110,26 +110,26 @@ def _find_model_by_substring(self, target_model: str) -> str | None: # Avoid duplicates from strategy 1 if not any(m[0] == available_model for m in matches): matches.append((available_model, len(normalized_available), 2)) # priority 2 - + # Strategy 3: Provider + base model name matching (more lenient) # Extract provider (e.g., "anthropic") and base model name target_provider = target_model.split('/')[0] target_base = target_model.split('/')[-1] if '/' in target_model else target_model - + for available_model in self.model_pricing: if '/' not in available_model: continue available_provider = available_model.split('/')[0] available_base = available_model.split('/')[-1] - + # Must match provider if target_provider != available_provider: continue - + # Normalize base names and check for significant overlap norm_target_base = self._normalize_model_name(target_base) norm_avail_base = self._normalize_model_name(available_base) - + # Check if one contains the other (after normalization) if norm_avail_base in norm_target_base or norm_target_base in norm_avail_base: # Avoid duplicates @@ -138,34 +138,34 @@ def _find_model_by_substring(self, target_model: str) -> str | None: common_len = len(os.path.commonprefix([norm_target_base, norm_avail_base])) if common_len >= 5: # Require at least 5 chars of common prefix matches.append((available_model, common_len, 3)) # priority 3 (lowest) - + if not matches: return None - + # Sort by: priority (ascending), then length (descending) # This prefers direct matches, then longer matches within same priority matches.sort(key=lambda x: (x[2], -x[1])) return matches[0][0] - + def _get_model_pricing_with_fallback(self, model_name: str) -> dict: """Get model pricing with intelligent fallback.""" - + # 1. Try exact match if model_name in self.model_pricing: return self.model_pricing[model_name] - + # 2. Try exact match after stripping routing prefix stripped_name = self._strip_routing_prefix(model_name) if stripped_name != model_name and stripped_name in self.model_pricing: print(f"📊 Using pricing for {stripped_name} (stripped prefix from {model_name})") return self.model_pricing[stripped_name] - + # 3. Try substring matching (includes normalization and prefix stripping) pattern_match = self._find_model_by_substring(model_name) if pattern_match: print(f"📊 Using pricing for {pattern_match} (pattern matched from {model_name})") return self.model_pricing[pattern_match] - + print(f"⚠️ No match found for {model_name}, please enter model cost manually:") try: input_str = input("Input cost per 1M tokens: ") @@ -193,8 +193,6 @@ def calculate_cost(self, uuid: str) -> float: float: The total cost in USD """ - print("\n📊 Calculating final cost...") - memory_path = Path(self.memory_dir) / uuid if not memory_path.exists(): @@ -221,7 +219,7 @@ def calculate_cost(self, uuid: str) -> float: json_call["usage"]["total_tokens"], ) ) - + # Check for single agent mode (no orchestrator calls but has task files) if not orchestrator_calls_found: print("📊 Single agent mode detected - calculating agent execution costs only") @@ -274,20 +272,57 @@ def calculate_cost(self, uuid: str) -> float: print("📊 Skipping SmolAgent cost calculation (workflow execution failed)") total_cost = 0.0 - print("\n💰 Cost Breakdown:") - print("=" * 60) + total_input_tokens = 0 + total_output_tokens = 0 + total_all_tokens = 0 + + # Pre-calculate costs + call_costs = [] for call in llm_calls: pricing = self._get_model_pricing_with_fallback(call.model) cost = ( call.input_tokens * pricing["input"] + call.output_tokens * pricing["output"] ) / 1_000_000 - print("Agent:", call.agent) - print(f" Model: {call.model}") - print(f" Tokens: {call.total_tokens:,}") - print(f" Cost: {cost:.3f} USD") - print("-" * 40) + call_costs.append((call, cost)) total_cost += cost + total_input_tokens += call.input_tokens + total_output_tokens += call.output_tokens + total_all_tokens += call.total_tokens + + from sources.cli.pretty_print import ( + BOLD, CYAN, DIM, GREEN, MAGENTA, RESET, YELLOW, + ) + + W = 64 + + # Header + print(f"\n{CYAN}{'─' * W}{RESET}") + print(f"{CYAN}{'💰 COST BREAKDOWN':^{W}}{RESET}") + print(f"{CYAN}{'─' * W}{RESET}") + + for call, cost in call_costs: + # Agent name row + label = call.agent.replace("_", " ").title() + print(f" {BOLD}{MAGENTA}▸ {label}{RESET}") + # Details as aligned key-value pairs + print(f" {DIM}Model{RESET} {call.model}") + print( + f" {DIM}Tokens{RESET} " + f"{call.input_tokens:>9,} in │ " + f"{call.output_tokens:>9,} out │ " + f"{BOLD}{call.total_tokens:>10,}{RESET} total" + ) + cost_color = GREEN if cost < 0.01 else (YELLOW if cost < 0.10 else CYAN) + print(f" {DIM}Cost{RESET} {cost_color}${cost:.4f}{RESET}") + print(f" {DIM}{'·' * (W - 4)}{RESET}") + + # Totals + print(f"\n {BOLD}{'Total Tokens':<14}{RESET} {total_all_tokens:>10,} " + f"{DIM}({total_input_tokens:,} in / {total_output_tokens:,} out){RESET}") + total_color = GREEN if total_cost < 0.05 else (YELLOW if total_cost < 0.50 else CYAN) + print(f" {BOLD}{'Total Cost':<14}{RESET} {total_color}{BOLD}${total_cost:.4f} USD{RESET}") + print(f"{CYAN}{'─' * W}{RESET}\n") return total_cost