econagents 0.0.5__py3-none-any.whl → 0.0.8__py3-none-any.whl

econagents/config_parser/basic.py

@@ -0,0 +1,45 @@
+import json
+from typing import Any, Dict, Optional
+
+from econagents import AgentRole
+from econagents.core.events import Message
+from econagents.core.manager.phase import PhaseManager
+from econagents.core.state.game import GameState
+from econagents.config_parser.base import BaseConfigParser
+
+
+class BasicConfigParser(BaseConfigParser):
+    """
+    Basic configuration parser that adds a custom event handler for sending a
+    player-is-ready message when it receives a certain message from the server.
+    """
+
+    def create_manager(
+        self, game_id: int, state: GameState, agent_role: Optional[AgentRole], auth_kwargs: Dict[str, Any]
+    ) -> PhaseManager:
+        """
+        Create a manager instance with a custom event handler for the assign-name event.
+
+        Args:
+            game_id: The game ID
+            state: The game state instance
+            agent_role: The agent role instance
+            auth_kwargs: Authentication mechanism keyword arguments
+
+        Returns:
+            A PhaseManager instance with custom event handlers
+        """
+        # Get the base manager
+        manager = super().create_manager(game_id=game_id, state=state, agent_role=agent_role, auth_kwargs=auth_kwargs)
+
+        # Register custom event handler for assign-name event
+        async def handle_name_assignment(message: Message) -> None:
+            """Handle the name assignment event."""
+            # Include the agent ID from auth_kwargs in the ready message
+            agent_id = auth_kwargs.get("agent_id")
+            ready_msg = {"gameId": game_id, "type": "player-is-ready", "agentId": agent_id}
+            await manager.send_message(json.dumps(ready_msg))
+
+        manager.register_event_handler("assign-name", handle_name_assignment)
+
+        return manager

econagents/config_parser/ibex_tudelft.py

@@ -0,0 +1,243 @@
+import json
+from pathlib import Path
+from typing import Any, Dict, Optional, Type, List, Callable, cast
+from pydantic import create_model
+
+from econagents import AgentRole
+from econagents.core.events import Message
+from econagents.core.manager.phase import PhaseManager
+from econagents.core.state.game import GameState
+from econagents.core.game_runner import GameRunner
+from econagents.config_parser.base import BaseConfigParser
+from econagents.llm.observability import get_observability_provider
+
+
+def handle_market_event_impl(self: GameState, event_type: str, data: dict[str, Any]) -> None:
+    """Handles market-related events by delegating to the MarketState instance."""
+    try:
+        getattr(self.public_information, self.meta._market_state_variable_name).process_event(  # type: ignore
+            event_type=event_type, data=data
+        )
+    except Exception as e:
+        raise ValueError(f"Error processing market event: {e}") from e
+
+
+def handle_asset_movement_event_impl(self: GameState, event_type: str, data: dict[str, Any]) -> None:
+    """Handles asset-movement events by delegating to the MarketState instance."""
+    try:
+        winning_condition = self.public_information.winning_condition  # type: ignore
+        self.private_information.wallet[winning_condition]["balance"] = data["balance"]  # type: ignore
+        self.private_information.wallet[winning_condition]["shares"] = data["shares"]  # type: ignore
+    except Exception as e:
+        raise ValueError(f"Error processing asset-movement event: {e}") from e
+
+
+def get_custom_handlers_impl(self: GameState) -> Dict[str, Callable[[Message], None]]:
+    """Returns custom handlers for market events."""
+    market_events = ["add-order", "update-order", "delete-order", "contract-fulfilled"]
+    asset_movement_event_handlers: Dict[str, Callable[[Message], None]] = {
+        "asset-movement": self._handle_asset_movement_event  # type: ignore
+    }
+    market_event_handlers: Dict[str, Callable[[Message], None]] = {
+        event: self._handle_market_event  # type: ignore
+        for event in market_events
+    }
+    return {**asset_movement_event_handlers, **market_event_handlers}
+
+
+class IbexTudelftConfigParser(BaseConfigParser):
+    """
+    IBEX-TUDelft configuration parser that extends the BasicConfigParser
+    and adds role assignment functionality.
+    """
+
+    def __init__(self, config_path: Path):
+        """
+        Initialize the IBEX-TUDelft config parser.
+
+        Args:
+            config_path: Path to the YAML configuration file
+        """
+        super().__init__(config_path)
+        self._role_classes: Dict[int, Type[AgentRole]] = {}
+
+    def register_role_class(self, role_id: int, role_class: Type[AgentRole]) -> None:
+        """
+        Register a role class for a specific role ID.
+
+        Args:
+            role_id: The role ID
+            role_class: The agent role class
+        """
+        self._role_classes[role_id] = role_class
+
+    def create_manager(
+        self, game_id: int, state: GameState, agent_role: Optional[AgentRole], auth_kwargs: Dict[str, Any]
+    ) -> PhaseManager:
+        """
+        Create a manager instance with custom event handlers for assign-name and assign-role events.
+        The manager won't have a role initially, but will be assigned one during the game.
+
+        Args:
+            game_id: The game ID
+            state: The game state instance
+            agent_role: The agent role instance (may be None)
+            auth_kwargs: Authentication mechanism keyword arguments
+
+        Returns:
+            A PhaseManager instance with custom event handlers
+        """
+        # Get the manager with the name assignment handler
+        manager = super().create_manager(
+            game_id=game_id,
+            state=state,
+            agent_role=None,
+            auth_kwargs=auth_kwargs,
+        )
+
+        # Register custom event handler for assign-name event
+        async def handle_name_assignment(message: Message) -> None:
+            """Handle the name assignment event."""
+            # Include the agent ID from auth_kwargs in the ready message
+            agent_id = auth_kwargs.get("agent_id")
+            ready_msg = {"gameId": game_id, "type": "player-is-ready", "agentId": agent_id}
+            await manager.send_message(json.dumps(ready_msg))
+
+        # Register custom event handler for assign-role event
+        async def handle_role_assignment(message: Message) -> None:
+            """Handle the role assignment event."""
+            role_id = int(message.data.get("role", 0))
+            manager.logger.info(f"Role assigned: {role_id}")
+
+            # Initialize the agent based on the assigned role
+            self._initialize_agent(manager, role_id)
+
+        manager.register_event_handler("assign-role", handle_role_assignment)
+        manager.register_event_handler("assign-name", handle_name_assignment)
+
+        return manager
+
+    def _detect_market_state_in_config(self) -> tuple[bool, Optional[tuple[str, str]]]:
+        """
+        Detects if MarketState is used in the state configuration.
+
+        Returns:
+            A tuple: (has_market_state_field, market_state_details).
+            market_state_details is (field_name_on_section, section_attribute_name_on_gamestate)
+            or None if not found.
+        """
+        state_conf = self.config.state
+        # Check public_information first, then private, then meta
+        for field_conf in state_conf.public_information:
+            if field_conf.type == "MarketState":
+                return True, (field_conf.name, "public_information")
+        for field_conf in state_conf.private_information:
+            if field_conf.type == "MarketState":
+                return True, (field_conf.name, "private_information")
+        for field_conf in state_conf.meta_information:
+            if field_conf.type == "MarketState":
+                return True, (field_conf.name, "meta")
+        return False, None
+
+    def _create_enhanced_state_class(self, base_class: Type[GameState]) -> Type[GameState]:
+        """
+        Creates an enhanced GameState class by subclassing base_class and injecting market event handlers.
+        """
+        enhanced_class_name = f"Enhanced{base_class.__name__}"
+        enhanced_class = create_model(
+            enhanced_class_name,
+            __base__=base_class,
+        )
+        setattr(enhanced_class, "_handle_market_event", handle_market_event_impl)
+        setattr(enhanced_class, "_handle_asset_movement_event", handle_asset_movement_event_impl)
+        setattr(enhanced_class, "get_custom_handlers", get_custom_handlers_impl)
+        return cast(Type[GameState], enhanced_class)
+
+    def _check_additional_required_fields(self, base_dynamic_state_class: Type[GameState]) -> None:
+        public_fields = base_dynamic_state_class().public_information.model_json_schema()["properties"]  # type: ignore
+        private_fields = base_dynamic_state_class().private_information.model_json_schema()["properties"]  # type: ignore
+
+        if (
+            "winning_condition" not in public_fields.keys()  # type: ignore
+            or "wallet" not in private_fields.keys()  # type: ignore
+        ):
+            raise ValueError("Winning condition or wallet is not present in the config")
+
+    async def run_experiment(self, login_payloads: List[Dict[str, Any]], game_id: int) -> None:
+        """
+        Run the experiment from this configuration, potentially enhancing the GameState
+        class with market event handlers if MarketState is specified in the config.
+        """
+        # Step 1: Get the base state class from the original StateConfig
+        base_dynamic_state_class = self.config.state.create_state_class()
+
+        # Step 2: Detect if MarketState is used and get details
+        has_market_state_field, market_state_details = self._detect_market_state_in_config()
+
+        # Step 3: If MarketState is present, create an enhanced state class
+        if has_market_state_field and market_state_details:
+            self._check_additional_required_fields(base_dynamic_state_class)
+            final_state_class = self._create_enhanced_state_class(base_dynamic_state_class)
+        else:
+            final_state_class = base_dynamic_state_class
+
+        if not self.config.agent_roles:
+            raise ValueError("Configuration has no 'agent_roles'.")
+
+        # Create managers for each agent
+        agents_for_runner = []
+        for payload in login_payloads:
+            current_agent_manager = self.create_manager(
+                game_id=game_id,
+                state=final_state_class(),
+                agent_role=None,
+                auth_kwargs=payload,
+            )
+            current_agent_manager.state.meta.game_id = game_id
+            if market_state_details:
+                setattr(current_agent_manager.state.meta, "_market_state_variable_name", market_state_details[0])
+            agents_for_runner.append(current_agent_manager)
+
+        # Create runner config
+        runner_config_instance = self.config.runner.create_runner_config()
+        runner_config_instance.state_class = final_state_class
+        runner_config_instance.game_id = game_id
+
+        if any(hasattr(role, "prompts") and role.prompts for role in self.config.agent_roles):
+            prompts_dir = self.config._compile_inline_prompts()
+            runner_config_instance.prompts_dir = prompts_dir
+
+        runner = GameRunner(config=runner_config_instance, agents=agents_for_runner)
+        await runner.run_game()
+
+        if self.config._temp_prompts_dir and self.config._temp_prompts_dir.exists():
+            import shutil
+
+            shutil.rmtree(self.config._temp_prompts_dir)
+
+    def _initialize_agent(self, manager: PhaseManager, role_id: int) -> None:
+        """
+        Initialize the agent instance based on the assigned role.
+
+        Args:
+            manager: The phase manager instance
+            role_id: The role ID
+        """
+        agent_roles = self.config.agent_roles
+        agent_role = next((role for role in agent_roles if role.role_id == role_id), None)
+        if agent_role:
+            manager.agent_role = agent_role.create_agent_role()
+            manager.agent_role.logger = manager.logger  # type: ignore
+            if self.config.runner.observability_provider:
+                manager.agent_role.llm.observability = get_observability_provider(
+                    self.config.runner.observability_provider
+                )
+        else:
+            manager.logger.error("Invalid role assigned; cannot initialize agent.")
+            raise ValueError("Invalid role for agent initialization.")
+
+
+async def run_experiment_from_yaml(yaml_path: Path, login_payloads: List[Dict[str, Any]], game_id: int) -> None:
+    """Run an experiment from a YAML configuration file."""
+    parser = IbexTudelftConfigParser(yaml_path)
+    await parser.run_experiment(login_payloads, game_id)

econagents/core/game_runner.py

@@ -4,7 +4,7 @@ import queue
 from contextvars import ContextVar
 from logging.handlers import QueueHandler, QueueListener
 from pathlib import Path
-from typing import Literal, Optional, Type
+from typing import Literal, Optional, Type, List
 
 from pydantic import BaseModel, Field
 
@@ -63,6 +63,11 @@ class GameRunnerConfig(BaseModel):
     observability_provider: Optional[Literal["langsmith", "langfuse"]] = None
     """Name of the observability provider to use. Options: 'langsmith' or 'langfuse'"""
 
+    max_game_duration: int = Field(
+        default=600,
+        description="Maximum game duration in seconds. Default is 600 (10 minutes). Set to 0 or a negative value to disable the timeout.",
+    )
+
     # Agent stop configuration
     end_game_event: str = "game-over"
     """Event type that signals the end of the game and should stop the agent."""
@@ -301,6 +306,8 @@ class GameRunner:
 
         if not agent_manager.state and self.config.state_class:
             agent_manager.state = self.config.state_class()
+            # Set the game_id in the state
+            agent_manager.state.meta.game_id = self.config.game_id
             agent_manager.logger.debug(f"Injected default state: {agent_manager.state}")
 
         if not agent_manager.auth_mechanism:
@@ -338,8 +345,6 @@ class GameRunner:
             agent_id (int): Agent identifier
         """
         agent_logger = self.get_agent_logger(agent_id, self.config.game_id)
-        ctx_agent_id.set(str(agent_id))  # Convert int to str for context variable
-
         agent_manager.logger = agent_logger
 
     async def spawn_agent(self, agent_manager: PhaseManager, agent_id: int) -> None:
@@ -350,32 +355,126 @@ class GameRunner:
             agent_manager (PhaseManager): Agent manager to spawn
             agent_id (int): Agent identifier
         """
+        token = ctx_agent_id.set(str(agent_id))
         try:
             self._inject_agent_logger(agent_manager, agent_id)
             self._inject_default_config(agent_manager)
 
             agent_manager.logger.info(f"Connecting to WebSocket URL: {agent_manager.url}")
             await agent_manager.start()
+        except asyncio.CancelledError:
+            agent_manager.logger.info(f"Agent {agent_id} spawn_agent task was cancelled.")
+            raise
         except Exception:
             agent_manager.logger.exception(f"Error in game for Agent {agent_id}")
             raise
+        finally:
+            agent_manager.logger.info(
+                f"Agent {agent_id} spawn_agent task finished. Agent running state: {agent_manager.running}"
+            )
+            ctx_agent_id.reset(token)
+
+    async def _timeout_watchdog(self, game_logger: logging.Logger, agent_tasks: List[asyncio.Task]) -> None:
+        """
+        Timeout watchdog that monitors game duration and initiates shutdown when exceeded.
+
+        Args:
+            game_logger: Logger instance for the game
+            agent_tasks: List of agent tasks to cancel if timeout occurs
+        """
+        try:
+            await asyncio.sleep(self.config.max_game_duration)
+            game_logger.warning(
+                f"Game {self.config.game_id} reached maximum duration of {self.config.max_game_duration}s. Initiating shutdown."
+            )
+
+            stop_agent_manager_tasks = []
+            for idx, ag_mgr in enumerate(self.agents):
+                if ag_mgr.running:
+                    game_logger.info(f"Timeout: Stopping agent {idx + 1} for game {self.config.game_id}.")
+                    stop_agent_manager_tasks.append(
+                        asyncio.create_task(ag_mgr.stop(), name=f"TimeoutStopAgent-{self.config.game_id}-{idx + 1}")
+                    )
+            if stop_agent_manager_tasks:
+                await asyncio.gather(*stop_agent_manager_tasks, return_exceptions=True)
+
+            game_logger.info(f"Timeout: Checking for lingering agent tasks for game {self.config.game_id}.")
+            for agent_task_instance in agent_tasks:
+                if not agent_task_instance.done():
+                    game_logger.warning(
+                        f"Timeout: Agent task {agent_task_instance.get_name()} still running. Cancelling."
+                    )
+                    agent_task_instance.cancel()
+                    # Await the cancellation to ensure it's processed
+                    try:
+                        await agent_task_instance
+                    except asyncio.CancelledError:
+                        game_logger.info(f"Agent task {agent_task_instance.get_name()} cancelled by timeout watchdog.")
+                    except Exception as e_cancel:
+                        game_logger.error(
+                            f"Error awaiting cancelled agent task {agent_task_instance.get_name()}: {e_cancel}"
+                        )
+
+        except asyncio.CancelledError:
+            game_logger.info(f"Timeout watchdog for game {self.config.game_id} was cancelled.")
+            raise
+        except Exception as e_watchdog:
+            game_logger.exception(f"Error in timeout watchdog for game {self.config.game_id}: {e_watchdog}")
 
     async def run_game(self) -> None:
         """Run a game using provided game data."""
 
         game_logger = self.get_game_logger(self.config.game_id)
-        game_logger.info(f"Running game with ID: {self.config.game_id}")
+        game_logger.info(f"Running game with ID: {self.config.game_id}. Max duration: {self.config.max_game_duration}s")
+
+        agent_tasks: List[asyncio.Task] = []
+        timeout_monitor_task: Optional[asyncio.Task] = None
 
         try:
-            tasks = []
-            game_logger.info("Starting game")
-
-            for i, agent_manager in enumerate(self.agents, start=1):
-                tasks.append(self.spawn_agent(agent_manager, i))
-            await asyncio.gather(*tasks)
-        except Exception as e:
-            game_logger.exception(f"Failed to run game: {e}")
+            game_logger.info("Configuring and starting agents...")
+            for i, agent_manager_instance in enumerate(self.agents, start=1):
+                task = asyncio.create_task(
+                    self.spawn_agent(agent_manager_instance, i), name=f"AgentTask-{self.config.game_id}-{i}"
+                )
+                agent_tasks.append(task)
+
+            if self.config.max_game_duration is not None and self.config.max_game_duration > 0:
+                timeout_monitor_task = asyncio.create_task(
+                    self._timeout_watchdog(game_logger, agent_tasks), name=f"TimeoutWatchdog-{self.config.game_id}"
+                )
+
+            if agent_tasks:
+                results = await asyncio.gather(*agent_tasks, return_exceptions=True)
+                for i, result in enumerate(results):
+                    if isinstance(result, Exception):
+                        game_logger.error(f"Agent task {agent_tasks[i].get_name()} failed with: {result}")
+            else:
+                game_logger.debug("No agent tasks to run.")
+
+        except Exception as e_run_game:
+            game_logger.exception(f"GameRunner.run_game failed: {e_run_game}")
             raise
         finally:
-            game_logger.info("Game over")
+            if timeout_monitor_task and not timeout_monitor_task.done():
+                game_logger.debug(
+                    f"Game {self.config.game_id} finished or errored before timeout. Cancelling timeout watchdog."
+                )
+                timeout_monitor_task.cancel()
+
+            game_logger.info(f"Game {self.config.game_id}: Final cleanup - ensuring all agents are stopped.")
+            final_stop_tasks = []
+            for idx, agent_manager_instance in enumerate(self.agents):
+                if agent_manager_instance.running:
+                    game_logger.info(f"Final cleanup: Stopping agent {idx + 1} for game {self.config.game_id}.")
+                    final_stop_tasks.append(
+                        asyncio.create_task(
+                            agent_manager_instance.stop(), name=f"FinalStopAgent-{self.config.game_id}-{idx + 1}"
+                        )
+                    )
+
+            if final_stop_tasks:
+                await asyncio.gather(*final_stop_tasks, return_exceptions=True)
+                game_logger.info(f"Game {self.config.game_id}: Final agent stop tasks completed.")
+
             self.cleanup_logging()
+            game_logger.info(f"Game {self.config.game_id} finished and cleaned up.")

econagents/core/manager/base.py

@@ -142,7 +142,7 @@ class AgentManager(LoggerMixin):
             auth_mechanism_kwargs=self._auth_mechanism_kwargs,
         )
 
-    def _raw_message_received(self, raw_message: str):
+    async def _raw_message_received(self, raw_message: str):
         """Process raw message from the transport layer"""
         msg = self._extract_message_data(raw_message)
         if msg:
@@ -218,19 +218,15 @@ class AgentManager(LoggerMixin):
                 raise ValueError("URL must be set before starting the agent manager")
 
         self.running = True
-        connected = await self.transport.connect()
-        if connected:
-            self.logger.info("Connected to WebSocket server. Receiving messages...")
-            await self.transport.start_listening()
-        else:
-            self.logger.error("Failed to connect to WebSocket server")
+        self.logger.info("Starting agent manager. Receiving messages...")
+        await self.transport.start_listening()
 
     async def stop(self):
         """Stop the agent manager and close the connection."""
         self.running = False
         if self.transport:
             await self.transport.stop()
-            self.logger.info("Agent manager stopped and connection closed.")
+        self.logger.info("Agent manager stopped and connection closed.")
 
     async def on_event(self, message: Message):
         """

econagents/core/transport.py

@@ -69,14 +69,11 @@ class WebSocketTransport(LoggerMixin):
         self.on_message_callback = on_message_callback
         self.ws: Optional[ClientConnection] = None
         self._running = False
+        self._authenticated = False
 
-    async def connect(self) -> bool:
-        """Establish the WebSocket connection and authenticate."""
+    async def _authenticate(self) -> bool:
+        """Authenticate the connection."""
         try:
-            self.ws = await websockets.connect(self.url, ping_interval=30, ping_timeout=10)
-            self.logger.info("WebSocketTransport: connection opened.")
-
-            # Perform authentication using the callback
             if self.auth_mechanism:
                 if not self.auth_mechanism_kwargs:
                     self.auth_mechanism_kwargs = {}
@@ -84,9 +81,9 @@ class WebSocketTransport(LoggerMixin):
                 if not auth_success:
                     self.logger.error("Authentication failed")
                     await self.stop()
-                    self.ws = None  # Ensure ws is set to None after stopping
+                    self.ws = None
                     return False
-
+            self._authenticated = True
         except Exception as e:
             self.logger.exception(f"Transport connection error: {e}")
             return False
@@ -95,24 +92,55 @@ class WebSocketTransport(LoggerMixin):
 
     async def start_listening(self):
         """Begin receiving messages in a loop."""
+        self.logger.info("WebSocketTransport: starting to listen.")
         self._running = True
-        while self._running and self.ws:
-            try:
-                message_str = await self.ws.recv()
-                if self.on_message_callback:
-                    # Call the callback, supporting both sync and async functions
-                    self.logger.debug(f"<-- Transport received: {message_str}")
-                    result = self.on_message_callback(message_str)
-                    # If the callback is a coroutine function, await it
-                    if asyncio.iscoroutine(result):
-                        asyncio.create_task(result)
-            except ConnectionClosed:
-                self.logger.info("WebSocket connection closed by remote.")
-                break
-            except Exception:
-                self.logger.exception("Error in receive loop.")
-                break
-        self._running = False
+
+        try:
+            async for websocket in websockets.connect(self.url):
+                if not self._running:
+                    self.logger.info("WebSocketTransport: stopping as requested.")
+                    break
+
+                try:
+                    self.ws = websocket
+
+                    if not self._authenticated:
+                        await self._authenticate()
+                        if not self._authenticated:
+                            self.logger.error("Authentication failed. Stopping transport.")
+                            break
+
+                    async for message in self.ws:
+                        if not self._running:
+                            break
+                        if self.on_message_callback:
+                            self.logger.info(f"<-- Transport received: {message}")
+                            await self.on_message_callback(message)
+
+                except ConnectionClosed as e:
+                    self.logger.info(f"WebSocketTransport: connection closed: ({e.code}) {e.reason}")
+                    if not self._running:
+                        self.logger.info("WebSocketTransport: connection closed by client. Stopping transport.")
+                        break
+                    self.logger.info("WebSocketTransport: reconnecting...")
+                    continue
+                except Exception as e:
+                    self.logger.exception(f"Error in receive loop: {e}")
+                    break
+                finally:
+                    if self.ws:
+                        try:
+                            await self.ws.close()
+                            self.logger.info("WebSocketTransport: connection closed.")
+                        except Exception as e:
+                            self.logger.debug(f"Error closing websocket: {e}")
+                        finally:
+                            self.ws = None
+        except Exception as e:
+            self.logger.exception(f"Error in start_listening: {e}")
+        finally:
+            self._running = False
+            self.logger.info("WebSocketTransport: stopped listening.")
 
     async def send(self, message: str):
         """Send a raw string message to the WebSocket."""
@@ -125,8 +153,14 @@ class WebSocketTransport(LoggerMixin):
 
     async def stop(self):
         """Gracefully close the WebSocket connection."""
+        self.logger.info("WebSocketTransport: stopping...")
         self._running = False
         if self.ws:
-            await self.ws.close()
-            self.logger.info("WebSocketTransport: connection closed.")
-            self.ws = None  # Set ws to None after closing
+            try:
+                await self.ws.close()
+                self.logger.info("WebSocketTransport: connection closed.")
+            except Exception as e:
+                self.logger.debug(f"Error during stop: {e}")
+            finally:
+                self.ws = None
+                self._authenticated = False