econagents 0.0.1__py3-none-any.whl

econagents/core/state/game.py

@@ -0,0 +1,222 @@
+from typing import Any, Callable, Optional, Protocol, Type, TypeVar
+
+from pydantic import BaseModel, ConfigDict
+
+from econagents.core.events import Message
+from econagents.core.state.fields import EventField
+
+EventHandler = Callable[[str, dict[str, Any]], None]
+T = TypeVar("T", bound="GameState")
+
+
+class PropertyMapping(BaseModel):
+    """Mapping between event data and state properties
+
+    Args:
+        event_key: Key in the event data
+        state_key: Key in the state object
+        state_type: Whether to update private or public information ("private" or "public")
+        phases: Optional list of phases where this mapping should be applied. If None, applies to all phases.
+        exclude_phases: Optional list of phases where this mapping should not be applied.
+                      Cannot be used together with phases.
+    """
+
+    event_key: str
+    state_key: str
+    state_type: str = "private"
+    events: list[str] | None = None
+    exclude_events: list[str] | None = None
+
+    def model_post_init(self, __context: Any) -> None:
+        """Validate that events and exclude_events are not both specified"""
+        if self.events is not None and self.exclude_events is not None:
+            raise ValueError("Cannot specify both events and exclude_events")
+
+    def should_apply_in_event(self, current_event: str) -> bool:
+        """Determine if this mapping should be applied in the current event"""
+        if self.events is not None:
+            return current_event in self.events
+        if self.exclude_events is not None:
+            return current_event not in self.exclude_events
+        return True
+
+
+class PrivateInformation(BaseModel):
+    """Private information for each agent in the game"""
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=False)
+
+
+class PublicInformation(BaseModel):
+    """Public information for the game"""
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=False)
+
+
+class MetaInformation(BaseModel):
+    """Meta information for the game"""
+
+    model_config = ConfigDict(extra="allow", arbitrary_types_allowed=False)
+
+    game_id: int = EventField(default=0)
+    """ID of the game"""
+    player_name: Optional[str] = EventField(default=None)
+    """Name of the player"""
+    player_number: Optional[int] = EventField(default=None)
+    """Number of the player"""
+    players: list[dict[str, Any]] = EventField(default_factory=list)
+    """List of players in the game"""
+    phase: int = EventField(default=0)
+    """Current phase of the game"""
+
+
+class GameStateProtocol(Protocol):
+    meta: MetaInformation
+    private_information: PrivateInformation
+    public_information: PublicInformation
+
+    def model_dump(self) -> dict[str, Any]: ...
+
+    def model_dump_json(self) -> str: ...
+
+
+class GameState(BaseModel):
+    """Game state for a given game"""
+
+    meta: MetaInformation = EventField(default_factory=MetaInformation)
+    """Meta information for the game"""
+    private_information: PrivateInformation = EventField(default_factory=PrivateInformation)
+    """Private information for each agent in the game"""
+    public_information: PublicInformation = EventField(default_factory=PublicInformation)
+    """Public information for the game"""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._property_mappings = self._get_property_mappings()
+
+    def update(self, event: Message) -> None:
+        """
+        Generic state update method that handles both property mappings and custom event handlers.
+
+        Args:
+            event (Message): The event message containing event_type and data
+
+        This method will:
+        1. Check for custom event handlers first
+        2. Fall back to property mappings if no custom handler exists
+        3. Update state based on property mappings, considering phase restrictions
+        """
+        # Get custom event handlers from child class
+        custom_handlers = self.get_custom_handlers()
+
+        # Check if there's a custom handler for this event type
+        if event.event_type in custom_handlers:
+            custom_handlers[event.event_type](event.event_type, event.data)
+            return
+
+        # Update state based on mappings
+        for mapping in self._property_mappings:
+            # Skip if mapping shouldn't be applied in current event
+            if not mapping.should_apply_in_event(event.event_type):
+                continue
+
+            # Skip if the event key isn't in the event data
+            if mapping.event_key not in event.data:
+                continue
+
+            value = event.data[mapping.event_key]
+
+            # Update the appropriate state object based on state_type
+            if mapping.state_type == "meta":
+                setattr(self.meta, mapping.state_key, value)
+            elif mapping.state_type == "private":
+                setattr(self.private_information, mapping.state_key, value)
+            elif mapping.state_type == "public":
+                setattr(self.public_information, mapping.state_key, value)
+
+    def _get_property_mappings(self) -> list[PropertyMapping]:
+        """
+        Default implementation that generates property mappings from EventField metadata.
+
+        Returns:
+            list[PropertyMapping]: List of PropertyMapping objects generated from field metadata.
+        """
+        mappings = []
+
+        # Generate mappings from meta information fields
+        mappings.extend(self._generate_mappings_from_model(self.meta.__class__, "meta"))
+
+        # Generate mappings from private information fields
+        mappings.extend(self._generate_mappings_from_model(self.private_information.__class__, "private"))
+
+        # Generate mappings from public information fields
+        mappings.extend(self._generate_mappings_from_model(self.public_information.__class__, "public"))
+
+        return mappings
+
+    def _generate_mappings_from_model(self, model_class: Type, state_type: str) -> list[PropertyMapping]:
+        """
+        Generate property mappings from a Pydantic model class.
+
+        Args:
+            model_class (Type): The Pydantic model class to inspect
+            state_type (str): The state type ("meta", "private", or "public")
+
+        Returns:
+            list[PropertyMapping]: List of PropertyMapping objects
+        """
+        mappings = []
+
+        for field_name, field_info in model_class.model_fields.items():
+            # Skip fields with exclude_from_mapping=True
+            exclude_from_mapping = False
+            if hasattr(field_info, "json_schema_extra") and "event_metadata" in field_info.json_schema_extra:
+                exclude_from_mapping = field_info.json_schema_extra["event_metadata"]["exclude_from_mapping"]
+            elif hasattr(field_info, "exclude_from_mapping"):  # For backward compatibility
+                exclude_from_mapping = field_info.exclude_from_mapping
+
+            if exclude_from_mapping:
+                continue
+
+            # Get event key from event_key if provided, otherwise use field name
+            event_key = None
+            if hasattr(field_info, "json_schema_extra") and "event_metadata" in field_info.json_schema_extra:
+                event_key = field_info.json_schema_extra["event_metadata"]["event_key"]
+            elif hasattr(field_info, "event_key"):  # For backward compatibility
+                event_key = field_info.event_key
+
+            if event_key is None:
+                event_key = field_name
+
+            # Get events and exclude_events if provided
+            events = None
+            exclude_events = None
+
+            if hasattr(field_info, "json_schema_extra") and "event_metadata" in field_info.json_schema_extra:
+                events = field_info.json_schema_extra["event_metadata"]["events"]
+                exclude_events = field_info.json_schema_extra["event_metadata"]["exclude_events"]
+            else:
+                # For backward compatibility
+                events = getattr(field_info, "events", None)
+                exclude_events = getattr(field_info, "exclude_events", None)
+
+            mappings.append(
+                PropertyMapping(
+                    event_key=event_key,
+                    state_key=field_name,
+                    state_type=state_type,
+                    events=events,
+                    exclude_events=exclude_events,
+                )
+            )
+
+        return mappings
+
+    def get_custom_handlers(self) -> dict[str, EventHandler]:
+        """
+        Override this method to provide custom event handlers.
+
+        Returns:
+            dict[str, EventHandler]: A mapping of event types to handler functions.
+        """
+        return {}

econagents/core/state/market.py

@@ -0,0 +1,124 @@
+from typing import Optional
+
+from pydantic import BaseModel, Field, computed_field
+
+
+class Order(BaseModel):
+    id: int
+    sender: int
+    price: float
+    quantity: float
+    type: str
+    condition: int
+    now: bool = False
+
+
+class Trade(BaseModel):
+    from_id: int
+    to_id: int
+    price: float
+    quantity: float
+    condition: int
+    median: Optional[float] = None
+
+
+class MarketState(BaseModel):
+    """
+    Represents the current state of the market:
+    - Active orders in an order book
+    - History of recent trades
+    """
+
+    orders: dict[int, Order] = Field(default_factory=dict)
+    trades: list[Trade] = Field(default_factory=list)
+
+    @computed_field
+    def order_book(self) -> str:
+        asks = sorted(
+            [order for order in self.orders.values() if order.type == "ask"],
+            key=lambda x: x.price,
+            reverse=True,
+        )
+        bids = sorted(
+            [order for order in self.orders.values() if order.type == "bid"],
+            key=lambda x: x.price,
+            reverse=True,
+        )
+        sorted_orders = asks + bids
+        return "\n".join([str(order) for order in sorted_orders])
+
+    def process_event(self, event_type: str, data: dict):
+        """
+        Update the MarketState based on the eventType and
+        event data from the server.
+        """
+        if event_type == "add-order":
+            self._on_add_order(data["order"])
+
+        elif event_type == "update-order":
+            self._on_update_order(data["order"])
+
+        elif event_type == "delete-order":
+            self._on_delete_order(data["order"])
+
+        elif event_type == "contract-fulfilled":
+            self._on_contract_fulfilled(data)
+
+    def get_orders_from_player(self, player_id: int) -> list[Order]:
+        """Get all orders from a specific player."""
+        return [order for order in self.orders.values() if order.sender == player_id]
+
+    def _on_add_order(self, order_data: dict):
+        """
+        The server is telling us a new order has been added.
+        We'll store it in self.orders by ID.
+        """
+        order_id = order_data["id"]
+        new_order = Order(
+            id=order_id,
+            sender=order_data["sender"],
+            price=order_data["price"],
+            quantity=order_data["quantity"],
+            type=order_data["type"],
+            condition=order_data["condition"],
+            now=order_data.get("now", False),
+        )
+        self.orders[order_id] = new_order
+
+    def _on_update_order(self, order_data: dict):
+        """
+        The server is telling us the order's quantity or other fields
+        have changed (often due to partial fills).
+        """
+        order_id = order_data["id"]
+        if order_id in self.orders:
+            existing = self.orders[order_id]
+            existing.quantity = order_data.get("quantity", existing.quantity)
+            self.orders[order_id] = existing
+
+    def _on_delete_order(self, order_data: dict):
+        """
+        The server is telling us this order is removed
+        from the order book (fully filled or canceled).
+        """
+        order_id = order_data["id"]
+        if order_id in self.orders:
+            del self.orders[order_id]
+
+    def _on_contract_fulfilled(self, data: dict):
+        """
+        This indicates a trade has happened between 'from' and 'to'.
+        The server might also send update-order or delete-order events
+        to reflect the fill on the order book.
+        We track the trade in self.trades, but we typically rely
+        on update-order or delete-order to fix the order's quantity.
+        """
+        new_trade = Trade(
+            from_id=data["from"],
+            to_id=data["to"],
+            price=data["price"],
+            quantity=data.get("quantity", 1.0),
+            condition=data["condition"],
+            median=data.get("median"),
+        )
+        self.trades.append(new_trade)

econagents/core/transport.py

@@ -0,0 +1,132 @@
+from abc import ABC, abstractmethod
+import asyncio
+import json
+import logging
+from typing import Any, Callable, Optional
+
+import websockets
+from websockets.asyncio.client import ClientConnection
+from websockets.exceptions import ConnectionClosed
+
+from econagents.core.logging_mixin import LoggerMixin
+
+
+class AuthenticationMechanism(ABC):
+    """Abstract base class for authentication mechanisms."""
+
+    @abstractmethod
+    async def authenticate(self, transport: "WebSocketTransport", **kwargs) -> bool:
+        """Authenticate the transport."""
+        pass
+
+    @classmethod
+    def __get_pydantic_core_schema__(cls, _source_type, _handler):
+        from pydantic_core import core_schema
+
+        return core_schema.is_instance_schema(AuthenticationMechanism)
+
+
+class SimpleLoginPayloadAuth(AuthenticationMechanism):
+    """Authentication mechanism that sends a login payload as the first message."""
+
+    async def authenticate(self, transport: "WebSocketTransport", **kwargs) -> bool:
+        """Send the login payload as a JSON message."""
+        initial_message = json.dumps(kwargs)
+        await transport.send(initial_message)
+        return True
+
+
+class WebSocketTransport(LoggerMixin):
+    """
+    Responsible for connecting to a WebSocket, sending/receiving messages,
+    and reporting received messages to a callback function.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        logger: Optional[logging.Logger] = None,
+        auth_mechanism: Optional[AuthenticationMechanism] = None,
+        auth_mechanism_kwargs: Optional[dict[str, Any]] = None,
+        on_message_callback: Optional[Callable[[str], Any]] = None,
+    ):
+        """
+        Initialize the WebSocket transport.
+
+        Args:
+            url: WebSocket server URL
+            logger: (Optional) Logger instance
+            auth_mechanism: (Optional) Authentication mechanism
+            auth_mechanism_kwargs: (Optional) Keyword arguments to pass to auth_mechanism during authentication
+            on_message_callback: Callback function that receives raw message strings.
+                               Can be synchronous or asynchronous.
+        """
+        self.url = url
+        self.auth_mechanism = auth_mechanism
+        self.auth_mechanism_kwargs = auth_mechanism_kwargs
+        if logger:
+            self.logger = logger
+        self.on_message_callback = on_message_callback
+        self.ws: Optional[ClientConnection] = None
+        self._running = False
+
+    async def connect(self) -> bool:
+        """Establish the WebSocket connection and authenticate."""
+        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 = {}
+                auth_success = await self.auth_mechanism.authenticate(self, **self.auth_mechanism_kwargs)
+                if not auth_success:
+                    self.logger.error("Authentication failed")
+                    await self.stop()
+                    self.ws = None  # Ensure ws is set to None after stopping
+                    return False
+
+        except Exception as e:
+            self.logger.exception(f"Transport connection error: {e}")
+            return False
+        else:
+            return True
+
+    async def start_listening(self):
+        """Begin receiving messages in a loop."""
+        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
+
+    async def send(self, message: str):
+        """Send a raw string message to the WebSocket."""
+        if self.ws:
+            try:
+                self.logger.debug(f"--> Transport sending: {message}")
+                await self.ws.send(message)
+            except Exception:
+                self.logger.exception("Error sending message.")
+
+    async def stop(self):
+        """Gracefully close the WebSocket connection."""
+        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

econagents/llm/__init__.py

@@ -0,0 +1,3 @@
+from econagents.llm.openai import ChatOpenAI
+
+__all__: list[str] = ["ChatOpenAI"]

econagents/llm/openai.py

@@ -0,0 +1,61 @@
+from typing import Any, Optional
+
+from langsmith import traceable
+from langsmith.wrappers import wrap_openai
+from openai import AsyncOpenAI
+
+
+class ChatOpenAI:
+    """
+    A simple wrapper for LLM queries, e.g. using OpenAI and LangSmith.
+    """
+
+    def __init__(
+        self,
+        model_name: str = "gpt-4o",
+        api_key: Optional[str] = None,
+    ) -> None:
+        """Initialize the LLM interface."""
+        self.model_name = model_name
+        self.api_key = api_key
+
+    def build_messages(self, system_prompt: str, user_prompt: str):
+        """Build messages for the LLM.
+
+        Args:
+            system_prompt (str): The system prompt for the LLM.
+            user_prompt (str): The user prompt for the LLM.
+
+        Returns:
+            list[dict[str, Any]]: The messages for the LLM.
+        """
+        return [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+
+    @traceable
+    async def get_response(
+        self,
+        messages: list[dict[str, Any]],
+        tracing_extra: dict[str, Any],
+        **kwargs: Any,
+    ):
+        """Get a response from the LLM.
+
+        Args:
+            messages (list[dict[str, Any]]): The messages for the LLM.
+            tracing_extra (dict[str, Any]): The extra tracing information.
+
+        Returns:
+            str: The response from the LLM.
+        """
+        client = wrap_openai(AsyncOpenAI(api_key=self.api_key))
+        response = await client.chat.completions.create(
+            messages=messages,  # type: ignore
+            model=self.model_name,
+            response_format={"type": "json_object"},
+            langsmith_extra=tracing_extra,
+            **kwargs,
+        )
+        return response.choices[0].message.content

econagents-0.0.1.dist-info/METADATA

@@ -0,0 +1,90 @@
+Metadata-Version: 2.3
+Name: econagents
+Version: 0.0.1
+Summary: 
+License: Apache-2.0
+Author: Dylan
+Requires-Python: >=3.10,<3.13
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: langsmith (>=0.3.13,<0.4.0)
+Requires-Dist: numpy (>=2.2.3,<3.0.0)
+Requires-Dist: openai (>=1.65.5,<2.0.0)
+Requires-Dist: pydantic (>=2.10.6,<3.0.0)
+Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: typing-extensions (>=4.12.2,<5.0.0)
+Requires-Dist: websockets (>=15.0,<16.0)
+Project-URL: Homepage, https://github.com/iwanalabs/econagents
+Project-URL: Repository, https://github.com/iwanalabs/econagents
+Description-Content-Type: text/markdown
+
+<div align="center">
+  <img src="https://raw.githubusercontent.com/iwanalabs/econagents/main/assets/logo_200w.png">
+</div>
+
+<div align="center">
+
+![Python compat](https://img.shields.io/badge/%3E=python-3.10-blue.svg)
+[![PyPi](https://img.shields.io/pypi/v/econagents.svg)](https://pypi.python.org/pypi/econagents)
+[![GHA Status](https://github.com/iwanalabs/econagents/actions/workflows/tests.yaml/badge.svg?branch=main)](https://github.com/iwanalabs/econagents/actions?query=workflow%3Atests)
+[![Documentation Status](https://readthedocs.org/projects/econagents/badge/?version=latest)](https://econagents.readthedocs.io/en/latest/?badge=latest)
+
+</div>
+
+---
+
+# econagents
+
+econagents is a Python library that lets you use LLM agents in economic experiments. The framework connects LLM agents to game servers through WebSockets and provides a flexible architecture for designing, customizing, and running economic simulations.
+
+## Installation
+
+```shell
+# Install from PyPI
+pip install econagents
+
+# Or install directly from GitHub
+pip install git+https://github.com/iwanalabs/econagents.git
+```
+
+## Framework Components
+
+econagents consists of four key components:
+
+1. **Agent Roles**: Define player roles with customizable behaviors using a flexible prompt system.
+2. **Game State**: Hierarchical state management with automatic event-driven updates.
+3. **Agent Managers**: Manage agent connections to game servers and handle event processing.
+4. **Game Runner**: Orchestrates experiments by gluing together the other components.
+
+## Example Experiments
+
+The repository includes two example experiments:
+
+1. **`prisoner`**: An iterated Prisoner's Dilemma game with 5 rounds and 2 LLM agents.
+2. **`tudeflt/harberger`**: A Harberger Tax simulation with LLM agents.
+3. **`tudeflt/futarchy`**: A Futarchy simulation with LLM agents.
+
+### Running the Prisoner's Dilemma Experiment
+
+```shell
+# Run the server
+python examples/server/prisoner/server.py
+
+# Run the experiment (in a separate terminal)
+python examples/prisoner/run_game.py
+```
+
+## Key Features
+
+- **Flexible Agent Customization**: Customize agent behavior with Jinja templates or custom Python methods
+- **Event-Driven State Management**: Automatically update game state based on server events
+- **Turn-Based and Continuous Action Support**: Handle both turn-based games and continuous action phases
+- **LangChain Integration**: Built-in support for LangChain's agent capabilities
+
+## Documentation
+
+For detailed guides and API reference, visit [the documentation](https://econagents.readthedocs.io/en/latest/).
+

econagents-0.0.1.dist-info/RECORD

@@ -0,0 +1,21 @@
+econagents/__init__.py,sha256=7oAI7W8akjmKDPArgmCjA7SDpWqvI7eLPbGeBRTiEzc,1050
+econagents/_c_extension.pyi,sha256=evVvDNUCGqyMPrNViPF7QXfGUNNIMbUdY5HemRNQ1_o,113
+econagents/core/__init__.py,sha256=QZoOp6n5CX1j-Ob6PZgyCNY78vi2kWmd_LVLrJUj1TU,393
+econagents/core/agent_role.py,sha256=viT7V6U9AmwDplgr4xGIVQvm3vLMtIFpGaMcGKM4oWE,15216
+econagents/core/events.py,sha256=hx-Ru_NoSISuN--7ZFC3CIql5hry3AATSnHZJJv3Kds,294
+econagents/core/game_runner.py,sha256=CToHQWIWQ1dwFLWRBxgFGlNecOjtlQUx3YbuS4cfOFQ,12869
+econagents/core/logging_mixin.py,sha256=tYsRc5ngW-hzfElrb838KO-9-BGOPyUv2v5LLuJToBE,1421
+econagents/core/manager/__init__.py,sha256=bDpCQlFcw_E-js575X3Xl6iwZ1uILC18An1vt6oE7S4,284
+econagents/core/manager/base.py,sha256=IMGkyCrghHlkJnkQLGUVfUSSr7sqZsKHYsmgf4UtGlI,16186
+econagents/core/manager/phase.py,sha256=M7s7jyA99BESXC1V9VNBuAOcvVncGoN5lF1nK8dh0Eo,19632
+econagents/core/state/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+econagents/core/state/fields.py,sha256=YxVOqdriaHRHoyeXsIB8ZDHygneMJD1OikOyeILK_oA,1854
+econagents/core/state/game.py,sha256=Ux0s7WhOxu0aFvwgX_LM0Aiho-aa1N3yh1ManwWBRP4,8681
+econagents/core/state/market.py,sha256=Jg-X9mYH6B3cYOwxzjFDV5PDbCIYxipx2UN4ecfyyDE,3909
+econagents/core/transport.py,sha256=7eq31nb2KY67RuL5i2kxJrcGtwfcVm5qy0eVj4_xWQw,5063
+econagents/llm/__init__.py,sha256=-tgv6qf77EdceWENIX6pDWXxu2AumhuUCjLiv4FmGKk,82
+econagents/llm/openai.py,sha256=1w8nHr8Ge2aEzO8lEsxKO3tUgLPEaGhYLwYu12653GY,1782
+econagents/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+econagents-0.0.1.dist-info/METADATA,sha256=N9v2yhoCGJ_2JdYq-lZQzoE7QKtMHqDamEDVAQVjAno,3431
+econagents-0.0.1.dist-info/WHEEL,sha256=XbeZDeTWKc1w7CSIyre5aMDU_-PohRwTQceYnisIYYY,88
+econagents-0.0.1.dist-info/RECORD,,

econagents-0.0.1.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: poetry-core 2.1.1
+Root-Is-Purelib: true
+Tag: py3-none-any