huzzle-labs/visual_memory

1"""2Gym-agnostic Agent Runner — connects an LLM to any OpenEnv environment.3 4This module is the CORE of the evaluation platform. It:5  1. Receives a pre-connected OpenEnv client (from AutoEnv discovery)6  2. Discovers tools via list_tools()7  3. Gives the LLM a scenario prompt + available tools8  4. Loops: LLM reasons → agent calls env.step() → observation → LLM reasons again9  5. Collects an EpisodeLog with timestamps for reward calculation + trajectory logging10 11Usage:12    from openenv import AutoEnv13    env = AutoEnv.from_env("visual_memory", base_url="http://localhost:8000")14    runner = AgentRunner(model="gpt-4o", env_client=env)15    episode, breakdown = runner.run_scenario(scenario, checker)16"""17 18import json19import logging20import time21from datetime import datetime, timezone, timedelta22from typing import Any, Dict, List, Tuple23 24IST = timezone(timedelta(hours=5, minutes=30))25 26from openenv.core.mcp_client import MCPToolClient27from openenv.core.env_server.mcp_types import CallToolAction, CallToolObservation, Tool28 29from rewards.base import (30    EpisodeLog,31    RewardBreakdown,32    RewardCalculator,33    Scenario,34    OpenEnvRewardCalculator,35)36from .llm import LLMClient37 38logger = logging.getLogger(__name__)39 40 41SYSTEM_PROMPT = """\42You are an AI agent interacting with an environment through tools.43 44Your job:451. Read the task description carefully.462. Use the available tools to complete the task.473. Call tools one at a time. Wait for each result before deciding the next step.484. When the task is complete, respond with a plain text summary of what you did.49   Do NOT call any more tools after you're done.50 51Rules:52- Only use tools that are listed as available.53- Provide all required arguments for each tool call.54- If a tool call fails, read the error and decide how to recover.55- Be efficient — complete the task in as few steps as possible.56- When you're done, clearly state what you accomplished.57"""58 59 60def mcp_tools_to_openai(tools: List[Tool]) -> List[Dict[str, Any]]:61    """Convert OpenEnv MCP tool definitions to OpenAI function-calling format."""62    openai_tools = []63    for tool in tools:64        schema = tool.input_schema or {"type": "object", "properties": {}}65        if "type" not in schema:66            schema["type"] = "object"67        if "properties" not in schema:68            schema["properties"] = {}69 70        openai_tools.append({71            "type": "function",72            "function": {73                "name": tool.name,74                "description": tool.description or "",75                "parameters": schema,76            },77        })78    return openai_tools79 80 81def _observation_to_str(step_result) -> str:82    """Convert an OpenEnv step result to a string the LLM can read."""83    obs = step_result.observation84    if isinstance(obs, CallToolObservation):85        if obs.error:86            return json.dumps({"error": obs.error.message}, indent=2)87        result = obs.result88        if hasattr(result, "data"):89            result = result.data90        elif isinstance(result, dict) and "data" in result:91            result = result["data"]92        try:93            return json.dumps(result, indent=2, default=str)94        except (TypeError, ValueError):95            return str(result)96    if hasattr(obs, "metadata") and obs.metadata:97        return json.dumps(obs.metadata, indent=2, default=str)98    return str(obs)99 100 101class AgentRunner:102    """103    Gym-agnostic agent that connects an LLM to any OpenEnv environment.104 105    Reward modes:106      - "custom"  (default): Episode-level reward via RewardCalculator107      - "openenv": Per-step reward via Transform + ground truth108    """109 110    def __init__(111        self,112        model: str,113        env_client: MCPToolClient,114        temperature: float = 0.0,115        max_tokens: int = 1024,116        reward_mode: str = "custom",117        transform=None,118    ):119        self.llm = LLMClient(120            model=model,121            temperature=temperature,122            max_tokens=max_tokens,123        )124        self.env_client = env_client125        self.reward_mode = reward_mode126        self.transform = transform127 128        self.calculator = RewardCalculator()129 130        if reward_mode == "openenv":131            self.openenv_calculator = OpenEnvRewardCalculator()132 133    def run_scenario(134        self,135        scenario: Scenario,136        checker: Any,137    ) -> Tuple[EpisodeLog, RewardBreakdown]:138        """Run a single scenario through the LLM agent."""139        return self._execute(scenario, checker, self.env_client)140 141    def _execute(142        self,143        scenario: Scenario,144        checker: Any,145        env: MCPToolClient,146    ) -> Tuple[EpisodeLog, RewardBreakdown]:147 148        env.reset()149 150        session_id = None151        try:152            session_result = env.step(153                CallToolAction(tool_name="get_session_info", arguments={})154            )155            obs = session_result.observation156            if isinstance(obs, CallToolObservation) and obs.result:157                result_data = obs.result158                if hasattr(result_data, "data"):159                    result_data = result_data.data160                elif isinstance(result_data, dict) and "data" in result_data:161                    result_data = result_data["data"]162                if isinstance(result_data, dict):163                    session_id = result_data.get("session_id")164                elif isinstance(result_data, str):165                    import json as _json166                    try:167                        parsed = _json.loads(result_data)168                        session_id = parsed.get("session_id")169                    except (ValueError, TypeError):170                        pass171        except Exception as e:172            logger.warning(f"Could not get session_id: {e}")173 174        if session_id and hasattr(checker, "set_session"):175            checker.set_session(session_id)176            logger.info(f"Session-scoped checker -> {session_id}")177 178        if self.transform and hasattr(self.transform, "set_scenario"):179            self.transform.set_scenario(scenario)180 181        all_tools = env.list_tools(use_cache=False)182        tools = [t for t in all_tools if t.name != "get_session_info"]183        openai_tools = mcp_tools_to_openai(tools)184        tool_names = [t.name for t in tools]185        logger.info(f"Discovered {len(tools)} agent tools: {tool_names}")186 187        messages = [188            {"role": "system", "content": SYSTEM_PROMPT},189            {"role": "user", "content": scenario.prompt},190        ]191 192        episode = EpisodeLog()193        step_rewards = []194        final_answer = None195 196        for step_num in range(1, scenario.max_steps + 1):197            logger.info(f"Step {step_num}/{scenario.max_steps}")198 199            response = self.llm.chat(messages, tools=openai_tools)200            tool_calls = LLMClient.extract_tool_calls(response)201 202            if not tool_calls:203                final_answer = LLMClient.get_text_response(response)204                logger.info(f"Agent done. Final answer: {(final_answer or '')[:100]}...")205                break206 207            messages.append(response.choices[0].message.model_dump())208 209            for tc in tool_calls:210                tool_name = tc["name"]211                arguments = tc["arguments"]212                call_id = tc["id"]213 214                logger.info(f"  Tool: {tool_name}({json.dumps(arguments, default=str)[:100]})")215 216                step_ts = datetime.now(IST).isoformat()217                step_start = time.time()218                error_msg = None219                try:220                    step_result = env.step(221                        CallToolAction(tool_name=tool_name, arguments=arguments)222                    )223                    obs = step_result.observation224                    is_error = (225                        isinstance(obs, CallToolObservation)226                        and obs.error is not None227                    )228                    result_str = _observation_to_str(step_result)229                    if is_error and isinstance(obs, CallToolObservation):230                        error_msg = obs.error.message231                except Exception as exc:232                    is_error = True233                    error_msg = str(exc)234                    result_str = json.dumps({"error": error_msg})235                    obs = None236 237                step_elapsed = time.time() - step_start238 239                if self.reward_mode == "openenv" and self.transform and obs is not None:240                    transformed = self.transform(obs)241                    step_rewards.append(242                        transformed.reward if transformed.reward is not None else 0.0243                    )244 245                episode.add_step(246                    tool_name=tool_name,247                    arguments=arguments,248                    success=not is_error,249                    result=result_str,250                    error=error_msg,251                    timestamp=step_ts,252                    elapsed=step_elapsed,253                )254 255                logger.info(f"    -> success={not is_error} ({step_elapsed:.2f}s)")256 257                messages.append({258                    "role": "tool",259                    "tool_call_id": call_id,260                    "content": result_str,261                })262 263        if hasattr(checker, "set_episode"):264            checker.set_episode(episode)265 266        outcome_results = checker.check_all(scenario.outcome_checks)267 268        if self.reward_mode == "openenv":269            breakdown = self.openenv_calculator.calculate(270                step_rewards=step_rewards,271                outcome_results=outcome_results,272                max_steps=scenario.max_steps,273                actual_steps=len(episode.steps),274            )275        else:276            breakdown = self.calculator.calculate(277                episode=episode,278                scenario=scenario,279                outcome_results=outcome_results,280            )281 282        return episode, breakdown283