Add a redesigned ensemble agent

Author: shuds13

ensemble_agent.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+"""Entry point for ensemble_agent — works as a standalone script."""
+import asyncio
+from ensemble_agent.config import parse_args
+from ensemble_agent.agent import run_agent
+
+asyncio.run(run_agent(parse_args()))

ensemble_agent/__init__.py

@@ -0,0 +1 @@
+"""Modular ensemble agent for generating, running, and fixing simulation scripts."""

ensemble_agent/__main__.py

@@ -0,0 +1,25 @@
+"""Entry point for python -m ensemble_agent."""
+
+import asyncio
+import sys
+
+from .config import parse_args
+from .agent import run_agent
+
+
+def main():
+    config = parse_args()
+    try:
+        asyncio.run(run_agent(config))
+    except KeyboardInterrupt:
+        print("\n\nInterrupted by user")
+        sys.exit(0)
+    except Exception as e:
+        print(f"\nError: {e}")
+        import traceback
+        traceback.print_exc()
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

ensemble_agent/agent.py

@@ -0,0 +1,164 @@
+"""Agent orchestrator — builds the agent, runs autonomous or interactive mode."""
+
+import shutil
+import sys
+from contextlib import AsyncExitStack
+from pathlib import Path
+
+from langchain.agents import create_agent
+from langchain_core.messages import HumanMessage, SystemMessage
+
+from .archive import ArchiveManager
+from .config import AgentConfig, INPUT_MARKER
+from .llm import create_llm
+from .mcp_client import connect_mcp, find_mcp_server
+from .prompts import (
+    AUTONOMOUS_GOAL,
+    INTERACTIVE_GOAL,
+    INTERACTIVE_REVIEW_GOAL,
+    build_system_prompt,
+)
+from .scripts import detect_run_script
+from .skills import load_skills
+from .skills.generator import GeneratorSkill
+
+
+async def run_agent(config: AgentConfig):
+    """Main entry point: build skills, connect MCP, run the agent loop."""
+
+    # Set up archive manager
+    archive = ArchiveManager(config.output_dir)
+
+    # Load skills
+    skills = load_skills(config.skills, config, archive)
+    has_generator = any(isinstance(s, GeneratorSkill) for s in skills)
+
+    # MCP setup for generator skill
+    async with AsyncExitStack() as stack:
+        if has_generator:
+            mcp_server = find_mcp_server(config.mcp_server)
+            print(f"Generator MCP: {mcp_server}")
+            session = await stack.enter_async_context(connect_mcp(mcp_server))
+            print("Connected to MCP server")
+
+            # Get MCP tool schema and inject into generator skill
+            mcp_tools = await session.list_tools()
+            mcp_tool = mcp_tools.tools[0]
+            for skill in skills:
+                if isinstance(skill, GeneratorSkill):
+                    skill.set_mcp_session(session)
+                    skill.set_mcp_tool_schema(mcp_tool)
+
+        # Collect tools from all skills
+        tools = []
+        for skill in skills:
+            tools.extend(skill.get_tools())
+
+        # Create LLM and agent
+        llm = create_llm(config.model, config.temperature, config.base_url)
+        agent = create_agent(llm, tools)
+        print("Agent initialized\n")
+
+        # Build system prompt
+        system_prompt = build_system_prompt(skills, has_generator)
+        messages = [("system", system_prompt)]
+
+        if config.show_prompts:
+            print(f"System prompt:\n{system_prompt}\n")
+
+        # Determine initial message
+        initial_msg = _build_initial_message(config, archive)
+
+        if not config.interactive:
+            await _run_autonomous(agent, messages, initial_msg, config)
+        else:
+            await _run_interactive(agent, messages, initial_msg, config, has_generator)
+
+
+def _build_initial_message(config, archive):
+    """Build the first user message based on config."""
+    if config.scripts_dir:
+        scripts_dir = Path(config.scripts_dir)
+        for f in sorted(scripts_dir.glob("*.py")):
+            shutil.copy(f, archive.work_dir)
+            print(f"Copied: {f.name}")
+        archive.start("copied_scripts")
+        archive.archive_scripts()
+
+        run_scripts = list(archive.work_dir.glob("run_*.py"))
+        run_name = run_scripts[0].name if run_scripts else "run_libe.py"
+        return INTERACTIVE_REVIEW_GOAL.format(run_script_name=run_name)
+
+    user_prompt = config.get_user_prompt()
+    if user_prompt:
+        return user_prompt
+
+    if config.interactive:
+        print("Describe the scripts you want to generate (or press Enter for default demo):", flush=True)
+        print(INPUT_MARKER, flush=True)
+        user_input = input().strip()
+        if user_input:
+            return user_input
+        print("Using default demo prompt")
+
+    return (
+        "Create six_hump_camel APOSMM scripts:\n"
+        "- Executable: six_hump_camel/six_hump_camel.x\n"
+        "- Input: six_hump_camel/input.txt\n"
+        "- Template vars: X0, X1\n"
+        "- 4 workers, 100 sims.\n"
+        "- The output file for each simulation is output.txt\n"
+        "- The bounds should be 0,1 and -1,2 for X0 and X1 respectively"
+    )
+
+
+async def _run_autonomous(agent, messages, initial_msg, config):
+    """Single invocation — agent generates/loads, runs, fixes, reports."""
+    goal = AUTONOMOUS_GOAL.format(initial_msg=initial_msg)
+    messages.append(("user", goal))
+
+    if config.show_prompts:
+        print(f"Goal: {goal}\n")
+    print("Starting agent...\n")
+
+    result = await agent.ainvoke({"messages": messages})
+    print(f"\n{'=' * 60}")
+    print("Agent completed")
+    print(f"{'=' * 60}")
+    print(result["messages"][-1].content)
+
+
+async def _run_interactive(agent, messages, initial_msg, config, has_generator):
+    """Chat loop — agent responds, waits for user input, repeats."""
+    if has_generator:
+        goal = INTERACTIVE_GOAL.format(initial_msg=initial_msg)
+    else:
+        goal = initial_msg
+    messages.append(("user", goal))
+    print("Starting agent...\n")
+
+    while True:
+        try:
+            result = await agent.ainvoke({"messages": messages})
+            messages = result["messages"]
+            response = messages[-1].content
+            if response:
+                print(f"\n{response}", flush=True)
+        except Exception as e:
+            print(f"\nAgent error: {e}", flush=True)
+
+        # Wait for user input
+        print(INPUT_MARKER, flush=True)
+        user_input = input().strip()
+
+        if not user_input or user_input.lower() in ("quit", "exit", "done"):
+            print("\nSession ended")
+            break
+
+        messages.append(
+            SystemMessage(
+                content="STOP. Read the user's next message carefully and respond to exactly what they ask. "
+                "Do not continue previous tasks."
+            )
+        )
+        messages.append(HumanMessage(content=user_input))

ensemble_agent/archive.py

@@ -0,0 +1,70 @@
+"""Versioned script and output archiving — no global state."""
+
+import shutil
+import time
+from pathlib import Path
+
+from .config import ARCHIVE_ITEMS, ARCHIVE_RUNS_DIR
+
+
+class ArchiveManager:
+    """Tracks versioned archives of scripts and run outputs."""
+
+    def __init__(self, work_dir):
+        self.work_dir = Path(work_dir)
+        self.work_dir.mkdir(parents=True, exist_ok=True)
+        self._counter = 1
+        self._current = None
+
+    @property
+    def current_archive(self):
+        return self._current
+
+    def start(self, action):
+        """Begin a new versioned archive (e.g. '1_generated', '2_fix')."""
+        self._current = f"{self._counter}_{action}"
+        (self.work_dir / "versions" / self._current).mkdir(parents=True, exist_ok=True)
+        self._counter += 1
+
+    def archive_scripts(self):
+        """Copy current *.py files into the active version directory."""
+        if not self._current:
+            return
+        dest = self.work_dir / "versions" / self._current
+        for f in self.work_dir.glob("*.py"):
+            shutil.copy(f, dest / f.name)
+
+    def archive_run_output(self, error_msg=""):
+        """Move run artifacts into the active version's output/ subdirectory."""
+        if not self._current:
+            return
+        output_dir = self.work_dir / "versions" / self._current / "output"
+        output_dir.mkdir(parents=True, exist_ok=True)
+
+        if error_msg:
+            (output_dir / "error.txt").write_text(error_msg)
+
+        for item in ARCHIVE_ITEMS:
+            item_path = self.work_dir / item
+            if item_path.exists() and item_path.is_dir():
+                shutil.copytree(str(item_path), str(output_dir / item), dirs_exist_ok=True)
+                shutil.rmtree(str(item_path))
+            else:
+                for fp in self.work_dir.glob(item):
+                    if fp.is_file():
+                        shutil.copy(str(fp), str(output_dir / fp.name))
+                        fp.unlink()
+
+    @staticmethod
+    def archive_existing_output_dir(output_dir, archive_parent=None):
+        """If output_dir exists, move it to archive_parent/output_dir_<unique>, then create fresh."""
+        output_dir = Path(output_dir)
+        archive_dir = Path(archive_parent or ARCHIVE_RUNS_DIR)
+        if not output_dir.exists():
+            output_dir.mkdir(parents=True, exist_ok=True)
+            return
+        archive_dir.mkdir(parents=True, exist_ok=True)
+        dest = archive_dir / f"{output_dir.name}_{hex(time.time_ns())[2:10]}"
+        shutil.move(str(output_dir), str(dest))
+        print(f"Moved existing {output_dir} to {dest}")
+        output_dir.mkdir(parents=True, exist_ok=True)