scripts/demo_research_agent.py

#!/usr/bin/env python3
"""Run 1–2 episodes of the research LLM agent with verbose terminal output."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from lab_env.env import LabEnv
from agents.research_llm_agent import ResearchLLMAgent


def main() -> None:
    parser = argparse.ArgumentParser(description="Demo: Research LLM agent (verbose)")
    parser.add_argument("--episodes", type=int, default=2, help="Number of episodes to run")
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--max-trials", type=int, default=5, help="Max trials per episode")
    args = parser.parse_args()

    env = LabEnv()
    agent = ResearchLLMAgent(max_trials=args.max_trials)

    print("=" * 60)
    print("  Research LLM Agent — Self-Improving Lab Scientist Demo")
    print("=" * 60)

    for ep in range(1, args.episodes + 1):
        print(f"\n--- Episode {ep}/{args.episodes} (seed={args.seed + ep}) ---")
        callback: list[dict] = []
        result = agent.run_episode(env, seed=args.seed + ep, verbose=True, episode_callback=callback)
        for step in callback:
            print(f"  Trial {step['trial']}: hypothesis {step['hypothesis']} -> ran {step['params_used']} -> {step['result']}")
        print(f"  Outcome: {'SUCCESS' if result['success'] else 'partial' if result['partial'] else 'fail'}")
        print(f"  Reward: {result['reward']:.1f}  Cost: ${result['cost']:.1f}  Steps: {result['steps']}")
        print(f"  Knowledge: temp_range={agent.knowledge['temp_range']}, cycle_range={agent.knowledge['cycle_range']}")

    env.close()
    print("\nDone.")


if __name__ == "__main__":
    main()