From e38c274f8d4a4162ab1532014cb3535e0a2d2788 Mon Sep 17 00:00:00 2001 From: Shannon Sands Date: Fri, 30 Jan 2026 15:40:14 +1000 Subject: [PATCH] Added AtroposAIAgent to ovveride standard runner with ManagedServer integration --- atropos_compatible_agent.py | 293 ++++++++++++++++++++++++++++++++++++ model_tools.py | 180 +++++++++++++++++++--- pyproject.toml | 10 +- run_agent.py | 8 +- tools/__init__.py | 147 +++++++++++------- tools/terminal_tool.py | 10 +- 6 files changed, 567 insertions(+), 81 deletions(-) create mode 100644 atropos_compatible_agent.py diff --git a/atropos_compatible_agent.py b/atropos_compatible_agent.py new file mode 100644 index 00000000000..fabe527ece1 --- /dev/null +++ b/atropos_compatible_agent.py @@ -0,0 +1,293 @@ +#!/usr/bin/env python3 +""" +Atropos-compatible Hermes agent runner. + +This is a minimal subclass of Hermes-Agent's `AIAgent` that swaps the OpenAI +function-calling backend for Atroposlib's `ManagedServer`/`ServerManager` backend +and uses Hermes-style XML tool tags: + +- {"name": "...", "arguments": {...}} +- {...} + +Tool observations are appended as `role="user"` messages containing one or more +`` blocks so they survive common chat templates during tokenization. +""" + +from __future__ import annotations + +import asyncio +import json +import re +import time +from contextlib import asynccontextmanager +from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple + +from model_tools import cleanup_vm, handle_function_call +from run_agent import AIAgent + +_TOOL_CALL_RE = re.compile(r"\\s*(.*?)\\s*", re.DOTALL) + + +ATROPOS_TOOL_SYSTEM_PROMPT = """You are a helpful AI assistant with access to tools. + +## Available Tools + +{tool_descriptions} + + +## How to Use Tools +To call a tool, output: +{{"name": "tool_name", "arguments": {{"arg1": "value1"}}}} + +You may include optional reasoning in ... before tool calls. + +After each tool call, you will receive tool results as: +{{...}} + +Continue until finished, then provide a final response with no blocks. +""" + + +class AtroposAIAgent(AIAgent): + """ + Hermes `AIAgent` variant that uses Atroposlib ServerManager/ManagedServer. + + Notes: + - The default Hermes `AIAgent` remains unchanged; this class is opt-in. + - The underlying server must expose `managed_server(tokenizer=...)` OR be a single + APIServer-compatible object usable by Atroposlib's `ManagedServer`. + """ + + def __init__( + self, + *, + server: Any, + tokenizer: Any = None, + model: str = "local", + max_iterations: int = 10, + tool_delay: float = 0.0, + enabled_toolsets: Optional[List[str]] = None, + disabled_toolsets: Optional[List[str]] = None, + save_trajectories: bool = False, + verbose_logging: bool = False, + ephemeral_system_prompt: Optional[str] = None, + log_prefix_chars: int = 100, + log_prefix: str = "", + temperature: float = 0.7, + max_tokens: int = 4096, + ): + # Call parent init mainly to reuse tool selection + trajectory saving utilities. + super().__init__( + base_url="http://unused", + api_key="dummy-key", + model=model, + max_iterations=max_iterations, + tool_delay=tool_delay, + enabled_toolsets=enabled_toolsets, + disabled_toolsets=disabled_toolsets, + save_trajectories=save_trajectories, + verbose_logging=verbose_logging, + ephemeral_system_prompt=ephemeral_system_prompt, + log_prefix_chars=log_prefix_chars, + log_prefix=log_prefix, + ) + + self.server = server + self.tokenizer = tokenizer + self.temperature = temperature + self.max_tokens = max_tokens + + @asynccontextmanager + async def _managed(self) -> AsyncGenerator[Any, None]: + if hasattr(self.server, "managed_server"): + async with self.server.managed_server(tokenizer=self.tokenizer) as managed: + yield managed + return + + # Fall back to directly wrapping a single server object. + from atroposlib.envs.server_handling.managed_server import ManagedServer + + managed = ManagedServer(server=self.server, tokenizer=self.tokenizer) + try: + yield managed + finally: + managed.reset() + + def _tool_descriptions_text(self) -> str: + if not self.tools: + return "(no tools available)" + + parts: List[str] = [] + for tool in self.tools: + fn = (tool or {}).get("function", {}) + name = fn.get("name", "") + desc = (fn.get("description") or "").strip() + if not name: + continue + if desc: + parts.append(f"- {name}: {desc}") + else: + parts.append(f"- {name}") + return "\n".join(parts) if parts else "(no tools available)" + + def _build_system_prompt(self, system_message: Optional[str]) -> Optional[str]: + if system_message is not None: + return system_message + if self.ephemeral_system_prompt: + return self.ephemeral_system_prompt + return ATROPOS_TOOL_SYSTEM_PROMPT.format( + tool_descriptions=self._tool_descriptions_text() + ) + + def _parse_tool_calls(self, content: str) -> Tuple[List[Tuple[str, Dict[str, Any]]], List[str]]: + """ + Returns: + (calls, errors) + """ + calls: List[Tuple[str, Dict[str, Any]]] = [] + errors: List[str] = [] + + for raw in _TOOL_CALL_RE.findall(content or ""): + try: + payload = json.loads(raw) + except json.JSONDecodeError as exc: + errors.append(f"Invalid JSON inside : {exc}") + continue + + name = payload.get("name") + args = payload.get("arguments", {}) + if not isinstance(name, str) or not name: + errors.append("Tool call missing 'name' string") + continue + if not isinstance(args, dict): + errors.append("Tool call 'arguments' must be an object") + continue + + calls.append((name, args)) + + return calls, errors + + async def run_conversation_async( + self, + user_message: str, + system_message: Optional[str] = None, + conversation_history: Optional[List[Dict[str, Any]]] = None, + task_id: Optional[str] = None, + ) -> Dict[str, Any]: + import uuid + + effective_task_id = task_id or str(uuid.uuid4()) + + messages: List[Dict[str, Any]] = conversation_history.copy() if conversation_history else [] + messages.append({"role": "user", "content": user_message}) + + active_system_prompt = self._build_system_prompt(system_message) + + api_call_count = 0 + final_response: Optional[str] = None + managed_state: Optional[Dict[str, Any]] = None + completed = False + + try: + async with self._managed() as managed: + while api_call_count < self.max_iterations: + api_call_count += 1 + + api_messages = messages.copy() + if active_system_prompt: + api_messages = [{"role": "system", "content": active_system_prompt}] + api_messages + + response = await managed.chat_completion( + messages=api_messages, + n=1, + max_tokens=self.max_tokens, + temperature=self.temperature, + ) + + if hasattr(managed, "get_state"): + managed_state = managed.get_state() + + assistant_content = response.choices[0].message.content or "" + messages.append({"role": "assistant", "content": assistant_content}) + + tool_calls, parse_errors = self._parse_tool_calls(assistant_content) + + if parse_errors and not tool_calls: + # Ask the model to retry with valid tool JSON. + err_text = "; ".join(parse_errors[:3]) + messages.append( + { + "role": "user", + "content": ( + f"{json.dumps({'error': err_text}, ensure_ascii=False)}\n" + "The previous blocks were invalid. Please output valid JSON inside ." + ), + } + ) + continue + + if not tool_calls: + # No tool calls: treat as final answer. + final_response = assistant_content + completed = True + break + + tool_responses: List[str] = [] + for tool_name, tool_args in tool_calls: + tool_start = time.time() + tool_result = handle_function_call(tool_name, tool_args, effective_task_id) + tool_duration = time.time() - tool_start + + try: + parsed = json.loads(tool_result) + payload: Any = parsed + except Exception: + payload = tool_result + + tool_payload = { + "name": tool_name, + "duration_s": round(tool_duration, 3), + "result": payload, + } + tool_responses.append( + f"{json.dumps(tool_payload, ensure_ascii=False)}" + ) + + if self.tool_delay and self.tool_delay > 0: + await asyncio.sleep(self.tool_delay) + + messages.append({"role": "user", "content": "\n".join(tool_responses)}) + + if final_response is None: + final_response = "I've reached the maximum number of iterations." + + finally: + try: + cleanup_vm(effective_task_id) + except Exception: + pass + + # Save trajectory using Hermes formatting (optional). + self._save_trajectory(messages, user_message, completed=completed) + + return { + "final_response": final_response, + "messages": messages, + "api_calls": api_call_count, + "completed": completed, + "managed_state": managed_state, + "system_prompt": active_system_prompt, + "task_id": effective_task_id, + } + + def run_conversation(self, *args: Any, **kwargs: Any) -> Dict[str, Any]: + """ + Sync wrapper for convenience. + + If already inside an event loop, call `await run_conversation_async(...)` instead. + """ + try: + asyncio.get_running_loop() + except RuntimeError: + return asyncio.run(self.run_conversation_async(*args, **kwargs)) + raise RuntimeError("AtroposAIAgent.run_conversation() cannot be called from a running event loop; use await run_conversation_async().") diff --git a/model_tools.py b/model_tools.py index 7f752318aa4..9572f1e8e16 100644 --- a/model_tools.py +++ b/model_tools.py @@ -30,30 +30,118 @@ import json import asyncio from typing import Dict, Any, List, Optional -from tools.web_tools import web_search_tool, web_extract_tool, web_crawl_tool, check_firecrawl_api_key -from tools.terminal_tool import terminal_tool, check_terminal_requirements, TERMINAL_TOOL_DESCRIPTION, cleanup_vm -# Hecate/MorphCloud terminal tool (cloud VMs) - available as alternative backend -from tools.terminal_hecate import terminal_hecate_tool, check_hecate_requirements, TERMINAL_HECATE_DESCRIPTION -from tools.vision_tools import vision_analyze_tool, check_vision_requirements -from tools.mixture_of_agents_tool import mixture_of_agents_tool, check_moa_requirements -from tools.image_generation_tool import image_generate_tool, check_image_generation_requirements -from tools.skills_tool import skills_categories, skills_list, skill_view, check_skills_requirements, SKILLS_TOOL_DESCRIPTION -# Browser automation tools (agent-browser + Browserbase) -from tools.browser_tool import ( - browser_navigate, - browser_snapshot, - browser_click, - browser_type, - browser_scroll, - browser_back, - browser_press, - browser_close, - browser_get_images, - browser_vision, - cleanup_browser, - check_browser_requirements, - BROWSER_TOOL_SCHEMAS -) +from tools.terminal_tool import TERMINAL_TOOL_DESCRIPTION, cleanup_vm, check_terminal_requirements, terminal_tool + +# Optional toolsets: keep Hermes importable even when some deps aren't installed. +try: + from tools.web_tools import check_firecrawl_api_key, web_crawl_tool, web_extract_tool, web_search_tool +except ModuleNotFoundError: + web_search_tool = None # type: ignore[assignment] + web_extract_tool = None # type: ignore[assignment] + web_crawl_tool = None # type: ignore[assignment] + + def check_firecrawl_api_key() -> bool: # type: ignore[no-redef] + return False + +try: + # Hecate/MorphCloud terminal tool (cloud VMs) - available as alternative backend + from tools.terminal_hecate import TERMINAL_HECATE_DESCRIPTION, check_hecate_requirements, terminal_hecate_tool +except ModuleNotFoundError: + terminal_hecate_tool = None # type: ignore[assignment] + TERMINAL_HECATE_DESCRIPTION = "" + + def check_hecate_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from tools.vision_tools import check_vision_requirements, vision_analyze_tool +except ModuleNotFoundError: + vision_analyze_tool = None # type: ignore[assignment] + + def check_vision_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from tools.mixture_of_agents_tool import check_moa_requirements, mixture_of_agents_tool +except ModuleNotFoundError: + mixture_of_agents_tool = None # type: ignore[assignment] + + def check_moa_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from tools.image_generation_tool import check_image_generation_requirements, image_generate_tool +except ModuleNotFoundError: + image_generate_tool = None # type: ignore[assignment] + + def check_image_generation_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from tools.skills_tool import ( + SKILLS_TOOL_DESCRIPTION, + check_skills_requirements, + skill_view, + skills_categories, + skills_list, + ) +except ModuleNotFoundError: + SKILLS_TOOL_DESCRIPTION = "" + + def check_skills_requirements() -> bool: # type: ignore[no-redef] + return False + + def skills_categories() -> str: # type: ignore[no-redef] + return json.dumps({"error": "Skills toolset is unavailable (missing dependencies)."}, ensure_ascii=False) + + def skills_list(category: Optional[str] = None) -> str: # type: ignore[no-redef] + _ = category + return json.dumps({"error": "Skills toolset is unavailable (missing dependencies)."}, ensure_ascii=False) + + def skill_view(name: str, file_path: Optional[str] = None) -> str: # type: ignore[no-redef] + _ = (name, file_path) + return json.dumps({"error": "Skills toolset is unavailable (missing dependencies)."}, ensure_ascii=False) + +try: + # Browser automation tools (agent-browser + Browserbase) + from tools.browser_tool import ( + BROWSER_TOOL_SCHEMAS, + browser_back, + browser_click, + browser_close, + browser_get_images, + browser_navigate, + browser_press, + browser_scroll, + browser_snapshot, + browser_type, + browser_vision, + check_browser_requirements, + cleanup_browser, + ) +except ModuleNotFoundError: + BROWSER_TOOL_SCHEMAS: List[Dict[str, Any]] = [] + + def check_browser_requirements() -> bool: # type: ignore[no-redef] + return False + + def cleanup_browser(task_id: Optional[str] = None) -> None: # type: ignore[no-redef] + _ = task_id + return None + + def _browser_unavailable(*_args: Any, **_kwargs: Any) -> str: + return json.dumps({"error": "Browser toolset is unavailable (missing dependencies)."}, ensure_ascii=False) + + browser_navigate = _browser_unavailable # type: ignore[assignment] + browser_snapshot = _browser_unavailable # type: ignore[assignment] + browser_click = _browser_unavailable # type: ignore[assignment] + browser_type = _browser_unavailable # type: ignore[assignment] + browser_scroll = _browser_unavailable # type: ignore[assignment] + browser_back = _browser_unavailable # type: ignore[assignment] + browser_press = _browser_unavailable # type: ignore[assignment] + browser_close = _browser_unavailable # type: ignore[assignment] + browser_get_images = _browser_unavailable # type: ignore[assignment] + browser_vision = _browser_unavailable # type: ignore[assignment] from toolsets import ( get_toolset, resolve_toolset, resolve_multiple_toolsets, get_all_toolsets, get_toolset_names, validate_toolset, @@ -572,6 +660,17 @@ def handle_web_function_call(function_name: str, function_args: Dict[str, Any]) Returns: str: Function result as JSON string """ + if web_search_tool is None or web_extract_tool is None or web_crawl_tool is None: + return json.dumps( + { + "error": ( + "Web toolset is unavailable (missing dependencies and/or FIRECRAWL_API_KEY). " + "Install web tool deps and set FIRECRAWL_API_KEY to enable." + ) + }, + ensure_ascii=False, + ) + if function_name == "web_search": query = function_args.get("query", "") # Always use fixed limit of 5 @@ -624,6 +723,17 @@ def handle_vision_function_call(function_name: str, function_args: Dict[str, Any Returns: str: Function result as JSON string """ + if vision_analyze_tool is None: + return json.dumps( + { + "error": ( + "Vision toolset is unavailable (missing dependencies and/or NOUS_API_KEY). " + "Install vision deps and set NOUS_API_KEY to enable." + ) + }, + ensure_ascii=False, + ) + if function_name == "vision_analyze": image_url = function_args.get("image_url", "") question = function_args.get("question", "") @@ -648,6 +758,17 @@ def handle_moa_function_call(function_name: str, function_args: Dict[str, Any]) Returns: str: Function result as JSON string """ + if mixture_of_agents_tool is None: + return json.dumps( + { + "error": ( + "Mixture-of-Agents toolset is unavailable (missing dependencies and/or NOUS_API_KEY). " + "Install MoA deps and set NOUS_API_KEY to enable." + ) + }, + ensure_ascii=False, + ) + if function_name == "mixture_of_agents": user_prompt = function_args.get("user_prompt", "") @@ -672,6 +793,17 @@ def handle_image_function_call(function_name: str, function_args: Dict[str, Any] Returns: str: Function result as JSON string """ + if image_generate_tool is None: + return json.dumps( + { + "error": ( + "Image generation toolset is unavailable (missing dependencies and/or FAL_KEY). " + "Install image deps and set FAL_KEY to enable." + ) + }, + ensure_ascii=False, + ) + if function_name == "image_generate": prompt = function_args.get("prompt", "") diff --git a/pyproject.toml b/pyproject.toml index 10e257f77ac..f066ef7ff93 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,7 +39,15 @@ dev = ["pytest", "pytest-asyncio"] hermes-agent = "run_agent:main" [tool.setuptools] -py-modules = ["run_agent", "model_tools", "toolsets", "batch_runner", "trajectory_compressor", "toolset_distributions"] +py-modules = [ + "run_agent", + "model_tools", + "toolsets", + "batch_runner", + "trajectory_compressor", + "toolset_distributions", + "atropos_compatible_agent", +] [tool.setuptools.packages.find] include = ["tools"] diff --git a/run_agent.py b/run_agent.py index 10b6e26a01c..0aba84a5f9e 100644 --- a/run_agent.py +++ b/run_agent.py @@ -30,7 +30,6 @@ import threading import uuid from typing import List, Dict, Any, Optional from openai import OpenAI -import fire from datetime import datetime from pathlib import Path @@ -1711,4 +1710,11 @@ def main( if __name__ == "__main__": + try: + import fire # type: ignore + except ModuleNotFoundError as exc: + raise SystemExit( + "Missing optional dependency 'fire'. Install hermes-agent with its CLI extras or add `fire` " + f"to your environment. Original error: {exc}" + ) from exc fire.Fire(main) diff --git a/tools/__init__.py b/tools/__init__.py index 8d2ee3b400f..9bff2593684 100644 --- a/tools/__init__.py +++ b/tools/__init__.py @@ -16,14 +16,6 @@ The tools are imported into model_tools.py which provides a unified interface for the AI agent to access all capabilities. """ -# Export all tools for easy importing -from .web_tools import ( - web_search_tool, - web_extract_tool, - web_crawl_tool, - check_firecrawl_api_key -) - # Primary terminal tool (mini-swe-agent backend: local/docker/singularity/modal) from .terminal_tool import ( terminal_tool, @@ -34,54 +26,106 @@ from .terminal_tool import ( TERMINAL_TOOL_DESCRIPTION ) -# Alternative terminal tool (Hecate/MorphCloud cloud VMs) -from .terminal_hecate import ( - terminal_hecate_tool, - check_hecate_requirements, - TERMINAL_HECATE_DESCRIPTION -) +# Optional toolsets: keep imports soft so users can run subsets of tools without +# installing every dependency (requirements gating lives in model_tools.py). +try: + from .web_tools import check_firecrawl_api_key, web_crawl_tool, web_extract_tool, web_search_tool +except ModuleNotFoundError: # pragma: no cover + web_search_tool = None # type: ignore[assignment] + web_extract_tool = None # type: ignore[assignment] + web_crawl_tool = None # type: ignore[assignment] -from .vision_tools import ( - vision_analyze_tool, - check_vision_requirements -) + def check_firecrawl_api_key() -> bool: # type: ignore[no-redef] + return False -from .mixture_of_agents_tool import ( - mixture_of_agents_tool, - check_moa_requirements -) +try: + # Alternative terminal tool (Hecate/MorphCloud cloud VMs) + from .terminal_hecate import TERMINAL_HECATE_DESCRIPTION, check_hecate_requirements, terminal_hecate_tool +except ModuleNotFoundError: # pragma: no cover + terminal_hecate_tool = None # type: ignore[assignment] + TERMINAL_HECATE_DESCRIPTION = "" -from .image_generation_tool import ( - image_generate_tool, - check_image_generation_requirements -) + def check_hecate_requirements() -> bool: # type: ignore[no-redef] + return False -from .skills_tool import ( - skills_categories, - skills_list, - skill_view, - check_skills_requirements, - SKILLS_TOOL_DESCRIPTION -) +try: + from .vision_tools import check_vision_requirements, vision_analyze_tool +except ModuleNotFoundError: # pragma: no cover + vision_analyze_tool = None # type: ignore[assignment] -# Browser automation tools (agent-browser + Browserbase) -from .browser_tool import ( - browser_navigate, - browser_snapshot, - browser_click, - browser_type, - browser_scroll, - browser_back, - browser_press, - browser_close, - browser_get_images, - browser_vision, - cleanup_browser, - cleanup_all_browsers, - get_active_browser_sessions, - check_browser_requirements, - BROWSER_TOOL_SCHEMAS -) + def check_vision_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from .mixture_of_agents_tool import check_moa_requirements, mixture_of_agents_tool +except ModuleNotFoundError: # pragma: no cover + mixture_of_agents_tool = None # type: ignore[assignment] + + def check_moa_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from .image_generation_tool import check_image_generation_requirements, image_generate_tool +except ModuleNotFoundError: # pragma: no cover + image_generate_tool = None # type: ignore[assignment] + + def check_image_generation_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + from .skills_tool import ( + SKILLS_TOOL_DESCRIPTION, + check_skills_requirements, + skill_view, + skills_categories, + skills_list, + ) +except ModuleNotFoundError: # pragma: no cover + skills_categories = None # type: ignore[assignment] + skills_list = None # type: ignore[assignment] + skill_view = None # type: ignore[assignment] + SKILLS_TOOL_DESCRIPTION = "" + + def check_skills_requirements() -> bool: # type: ignore[no-redef] + return False + +try: + # Browser automation tools (agent-browser + Browserbase) + from .browser_tool import ( + BROWSER_TOOL_SCHEMAS, + browser_back, + browser_click, + browser_close, + browser_get_images, + browser_navigate, + browser_press, + browser_scroll, + browser_snapshot, + browser_type, + browser_vision, + check_browser_requirements, + cleanup_all_browsers, + cleanup_browser, + get_active_browser_sessions, + ) +except ModuleNotFoundError: # pragma: no cover + browser_navigate = None # type: ignore[assignment] + browser_snapshot = None # type: ignore[assignment] + browser_click = None # type: ignore[assignment] + browser_type = None # type: ignore[assignment] + browser_scroll = None # type: ignore[assignment] + browser_back = None # type: ignore[assignment] + browser_press = None # type: ignore[assignment] + browser_close = None # type: ignore[assignment] + browser_get_images = None # type: ignore[assignment] + browser_vision = None # type: ignore[assignment] + cleanup_browser = None # type: ignore[assignment] + cleanup_all_browsers = None # type: ignore[assignment] + get_active_browser_sessions = None # type: ignore[assignment] + BROWSER_TOOL_SCHEMAS = [] + + def check_browser_requirements() -> bool: # type: ignore[no-redef] + return False __all__ = [ # Web tools @@ -132,4 +176,3 @@ __all__ = [ 'check_browser_requirements', 'BROWSER_TOOL_SCHEMAS', ] - diff --git a/tools/terminal_tool.py b/tools/terminal_tool.py index a5058f44cdc..bf0dc6b0915 100644 --- a/tools/terminal_tool.py +++ b/tools/terminal_tool.py @@ -1231,12 +1231,16 @@ def check_terminal_requirements() -> bool: try: if env_type == "local": - from minisweagent.environments.local import LocalEnvironment - return True + # Prefer mini-swe-agent when available, but allow a subprocess fallback. + try: + from minisweagent.environments.local import LocalEnvironment + + return True + except ImportError: + return True elif env_type == "docker": from minisweagent.environments.docker import DockerEnvironment # Check if docker is available - import subprocess result = subprocess.run(["docker", "version"], capture_output=True, timeout=5) return result.returncode == 0 elif env_type == "singularity":