pycityagent 2.0.0a66__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a67__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/tools/tool.py

@@ -7,20 +7,45 @@ from typing import Any, Optional, Union
 from mlflow.entities import Metric
 
 from ..agent import Agent
-from ..environment import (LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService,
-                           PersonService)
+from ..environment import AoiService, PersonService
 from ..utils.decorators import lock_decorator
 from ..workflow import Block
 
+__all__ = [
+    "Tool",
+    "ExportMlflowMetrics",
+    "GetMap",
+    "UpdateWithSimulator",
+    "ResetAgentPosition",
+]
+
 
 class Tool:
     """Abstract tool class for callable tools. Can be bound to an `Agent` or `Block` instance.
 
     This class serves as a base for creating various tools that can perform different operations.
     It is intended to be subclassed by specific tool implementations.
+
+    - **Attributes**:
+        - `_instance`: A reference to the instance (`Agent` or `Block`) this tool is bound to.
     """
 
     def __get__(self, instance, owner):
+        """
+        Descriptor method for binding the tool to an instance.
+
+        - **Args**:
+            - `instance`: The instance that the tool is being accessed through.
+            - `owner`: The type of the owner class.
+
+        - **Returns**:
+            - `Tool`: An instance of the tool bound to the given instance.
+
+        - **Description**:
+            - If accessed via the class rather than an instance, returns the descriptor itself.
+            - Otherwise, it checks if the tool has already been instantiated for this instance,
+              and if not, creates and stores a new tool instance specifically for this instance.
+        """
         if instance is None:
             return self
         subclass = type(self)
@@ -36,11 +61,23 @@ class Tool:
         """Invoke the tool's functionality.
 
         This method must be implemented by subclasses to provide specific behavior.
+
+        - **Raises**:
+            - `NotImplementedError`: When called directly on the base class.
         """
         raise NotImplementedError
 
     @property
     def agent(self) -> Agent:
+        """
+        Access the `Agent` this tool is bound to.
+
+        - **Returns**:
+            - `Agent`: The agent instance.
+
+        - **Raises**:
+            - `RuntimeError`: If the tool is not bound to an `Agent`.
+        """
         instance = self._instance  # type:ignore
         if not isinstance(instance, Agent):
             raise RuntimeError(
@@ -50,6 +87,15 @@ class Tool:
 
     @property
     def block(self) -> Block:
+        """
+        Access the `Block` this tool is bound to.
+
+        - **Returns**:
+            - `Block`: The block instance.
+
+        - **Raises**:
+            - `RuntimeError`: If the tool is not bound to a `Block`.
+        """
         instance = self._instance  # type:ignore
         if not isinstance(instance, Block):
             raise RuntimeError(
@@ -71,72 +117,9 @@ class GetMap(Tool):
         return agent.simulator.map
 
 
-class SencePOI(Tool):
-    """Retrieve the Point of Interest (POI) of the current scene.
-
-    This tool computes the POI based on the current `position` stored in memory and returns
-    points of interest (POIs) within a specified radius. Can be bound only to an `Agent` instance.
-
-    Attributes:
-        radius (int): The radius within which to search for POIs.
-        category_prefix (str): The prefix for the categories of POIs to consider.
-        variables (list[str]): A list of variables relevant to the tool's operation.
-
-    Args:
-        radius (int, optional): The circular search radius. Defaults to 100.
-        category_prefix (str, optional): The category prefix to filter POIs. Defaults to LEVEL_ONE_PRE.
-
-    Methods:
-        __call__(radius: Optional[int] = None, category_prefix: Optional[str] = None) -> Union[Any, Callable]:
-            Executes the AOI retrieval operation, returning POIs based on the current state of memory and simulator.
-    """
-
-    def __init__(self, radius: int = 100, category_prefix=LEVEL_ONE_PRE) -> None:
-        self.radius = radius
-        self.category_prefix = category_prefix
-        self.variables = ["position"]
-
-    async def __call__(
-        self, radius: Optional[int] = None, category_prefix: Optional[str] = None
-    ) -> Union[Any, Callable]:
-        """Retrieve the POIs within the specified radius and category prefix.
-
-        If both `radius` and `category_prefix` are None, the method will use the current position
-        from memory to query POIs using the simulator. Otherwise, it will return a new instance
-        of SenceAoi with the specified parameters.
-
-        Args:
-            radius (Optional[int]): A specific radius for the AOI query. If not provided, defaults to the instance's radius.
-            category_prefix (Optional[str]): A specific category prefix to filter POIs. If not provided, defaults to the instance's category_prefix.
-
-        Raises:
-            ValueError: If memory or simulator is not set.
-
-        Returns:
-            Union[Any, Callable]: The query results or a callable for a new SenceAoi instance.
-        """
-        agent = self.agent
-        if agent.memory is None or agent.simulator is None:
-            raise ValueError("Memory or Simulator is not set.")
-        if radius is None and category_prefix is None:
-            position = await agent.status.get("position")
-            resp = []
-            for prefix in self.category_prefix:
-                resp += agent.simulator.map.query_pois(
-                    center=(position["xy_position"]["x"], position["xy_position"]["y"]),
-                    radius=self.radius,
-                    category_prefix=prefix,
-                )
-            # * Map six-digit codes to specific types
-            for poi in resp:
-                cate_str = poi[0]["category"]
-                poi[0]["category"] = POI_TYPE_DICT[cate_str]
-        else:
-            radius_ = radius if radius else self.radius
-            return SencePOI(radius_, category_prefix)
-
-
 class UpdateWithSimulator(Tool):
+    """Automatically update status memory from simulator"""
+
     def __init__(self) -> None:
         self._lock = asyncio.Lock()
 
@@ -180,6 +163,18 @@ class ResetAgentPosition(Tool):
         lane_id: Optional[int] = None,
         s: Optional[float] = None,
     ):
+        """
+        Reset the position of the agent associated with this tool.
+
+        - **Args**:
+            - `aoi_id` (Optional[int], optional): Area of interest ID. Defaults to None.
+            - `poi_id` (Optional[int], optional): Point of interest ID. Defaults to None.
+            - `lane_id` (Optional[int], optional): Lane ID. Defaults to None.
+            - `s` (Optional[float], optional): Position along the lane. Defaults to None.
+
+        - **Description**:
+            - Resets the agent's position based on the provided parameters using the simulator.
+        """
         agent = self.agent
         status = agent.status
         await agent.simulator.reset_person_position(
@@ -192,7 +187,22 @@ class ResetAgentPosition(Tool):
 
 
 class ExportMlflowMetrics(Tool):
+    """
+    A tool for exporting metrics to MLflow in batches.
+
+    - **Attributes**:
+        - `_log_batch_size` (int): The number of metrics to log in each batch.
+        - `metric_log_cache` (Dict[str, List[Metric]]): Cache for storing metrics before batching.
+        - `_lock` (asyncio.Lock): Ensures thread-safe operations when logging metrics.
+    """
+
     def __init__(self, log_batch_size: int = 100) -> None:
+        """
+        Initialize the ExportMlflowMetrics tool with a specified batch size and an asynchronous lock.
+
+        - **Args**:
+            - `log_batch_size` (int, optional): Number of metrics per batch. Defaults to 100.
+        """
         self._log_batch_size = log_batch_size
         # TODO: support other log types
         self.metric_log_cache: dict[str, list[Metric]] = defaultdict(list)
@@ -204,6 +214,17 @@ class ExportMlflowMetrics(Tool):
         metric: Union[Sequence[Union[Metric, dict]], Union[Metric, dict]],
         clear_cache: bool = False,
     ):
+        """
+        Add metrics to the cache and export them to MLflow in batches if the batch size limit is reached.
+
+        - **Args**:
+            - `metric` (Union[Sequence[Union[Metric, dict]], Union[Metric, dict]]): A single metric or a sequence of metrics.
+            - `clear_cache` (bool, optional): Flag indicating whether to clear the cache after logging. Defaults to False.
+
+        - **Description**:
+            - Adds metrics to the cache. If the cache exceeds the batch size, logs a batch of metrics to MLflow.
+            - Optionally clears the entire cache.
+        """
         agent = self.agent
         batch_size = self._log_batch_size
         if not isinstance(metric, Sequence):
@@ -223,7 +244,7 @@ class ExportMlflowMetrics(Tool):
             self.metric_log_cache[metric_key].append(item)
         for metric_key, _cache in self.metric_log_cache.items():
             if len(_cache) > batch_size:
-                client = agent.mlflow_client
+                client = agent.mlflow_client  # type:ignore
                 await client.log_batch(
                     metrics=_cache[:batch_size],
                 )
@@ -234,8 +255,11 @@ class ExportMlflowMetrics(Tool):
     async def _clear_cache(
         self,
     ):
+        """
+        Log any remaining metrics from the cache to MLflow and then clear the cache.
+        """
         agent = self.agent
-        client = agent.mlflow_client
+        client = agent.mlflow_client  # type:ignore
         for metric_key, _cache in self.metric_log_cache.items():
             if len(_cache) > 0:
                 await client.log_batch(

pycityagent/utils/avro_schema.py

@@ -12,9 +12,9 @@ PROFILE_SCHEMA = {
         {"name": "skill", "type": "string"},
         {"name": "occupation", "type": "string"},
         {"name": "family_consumption", "type": "string"},
-        {"name": "consumption", "type": "string"},
+        {"name": "consumption", "type": "float"},
         {"name": "personality", "type": "string"},
-        {"name": "income", "type": "string"},
+        {"name": "income", "type": "float"},
         {"name": "currency", "type": "float"},
         {"name": "residence", "type": "string"},
         {"name": "race", "type": "string"},

pycityagent/utils/parsers/code_block_parser.py

@@ -26,7 +26,7 @@ class CodeBlockParser(ParserBase):
         Parameters:
             response (str): The response string containing the specified language object.
 
-        Returns:
+        - **Returns**:
             str: The parsed `str` object.
         """
         extract_text = self._extract_text_within_tags(

pycityagent/utils/parsers/json_parser.py

@@ -26,7 +26,7 @@ class JsonObjectParser(ParserBase):
         Parameters:
             response (str): The response string containing the JSON object.
 
-        Returns:
+        - **Returns**:
             Any: The parsed JSON object.
         """
         extract_text = self._extract_text_within_tags(
@@ -73,7 +73,7 @@ class JsonDictParser(JsonObjectParser):
         Parameters:
             response (str): The response string containing the JSON object.
 
-        Returns:
+        - **Returns**:
             dict: The parsed JSON object as a dictionary.
         """
         parsed_json = super().parse(response)

pycityagent/utils/parsers/parser_base.py

@@ -21,7 +21,7 @@ class ParserBase(ABC):
         Parameters:
             response (str): The raw string returned by the model.
 
-        Returns:
+        - **Returns**:
             Any: The converted data, the specific type depends on the parsing result.
         """
         pass
@@ -37,7 +37,7 @@ class ParserBase(ABC):
             tag_start (str): The start tag.
             tag_end (str): The end tag.
 
-        Returns:
+        - **Returns**:
             str: The string between the start and end tags.
         """
 

pycityagent/workflow/block.py

@@ -33,11 +33,11 @@ def log_and_check_with_memory(
 
     This decorator is specifically designed to be used with the `block` method. A 'Memory' object is required in method input.
 
-    Args:
-        condition (Callable): A condition function that must be satisfied before the decorated function is executed.
+    - **Args**:
+        - `condition` (Callable): A condition function that must be satisfied before the decorated function is executed.
                              Can be synchronous or asynchronous.
-        trigger_interval (float): The interval (in seconds) to wait between condition checks.
-        record_function_calling (bool): Whether to log the function call information.
+        - `trigger_interval` (float): The interval (in seconds) to wait between condition checks.
+        - `record_function_calling` (bool): Whether to log the function call information.
     """
 
     def decorator(func):
@@ -93,11 +93,11 @@ def log_and_check(
 
     This decorator is specifically designed to be used with the `block` method.
 
-    Args:
-        condition (Callable): A condition function that must be satisfied before the decorated function is executed.
+    - **Args**:
+        - `condition` (Callable): A condition function that must be satisfied before the decorated function is executed.
                              Can be synchronous or asynchronous.
-        trigger_interval (float): The interval (in seconds) to wait between condition checks.
-        record_function_calling (bool): Whether to log the function call information.
+        - `trigger_interval` (float): The interval (in seconds) to wait between condition checks.
+        - `record_function_calling` (bool): Whether to log the function call information.
     """
 
     def decorator(func):
@@ -144,6 +144,15 @@ def trigger_class():
 
 # Define a Block, similar to a layer in PyTorch
 class Block:
+    """
+    A foundational component similar to a layer in PyTorch, used for building complex systems.
+
+    - **Attributes**:
+        - `configurable_fields` (list[str]): A list of fields that can be configured.
+        - `default_values` (dict[str, Any]): Default values for configurable fields.
+        - `fields_description` (dict[str, str]): Descriptions for each configurable field.
+    """
+
     configurable_fields: list[str] = []
     default_values: dict[str, Any] = {}
     fields_description: dict[str, str] = {}
@@ -156,6 +165,17 @@ class Block:
         simulator: Optional[Simulator] = None,
         trigger: Optional[EventTrigger] = None,
     ):
+        """
+        - **Description**:
+            - Initializes a new instance of the Block class with optional LLM, Memory, Simulator, and Trigger components.
+
+        - **Args**:
+            - `name` (str): The name of the block.
+            - `llm` (Optional[LLM], optional): An instance of LLM. Defaults to None.
+            - `memory` (Optional[Memory], optional): An instance of Memory. Defaults to None.
+            - `simulator` (Optional[Simulator], optional): An instance of Simulator. Defaults to None.
+            - `trigger` (Optional[EventTrigger], optional): An event trigger that may be associated with this block. Defaults to None.
+        """
         self.name = name
         self._llm = llm
         self._memory = memory
@@ -167,13 +187,27 @@ class Block:
         self.trigger = trigger
 
     def export_config(self) -> dict[str, Optional[str]]:
+        """
+        - **Description**:
+            - Exports the configuration of the block as a dictionary.
+
+        - **Returns**:
+            - `Dict[str, Optional[str]]`: A dictionary containing the configuration of the block.
+        """
         return {
             field: self.default_values.get(field, "default_value")
             for field in self.configurable_fields
         }
 
     @classmethod
-    def export_class_config(cls) -> dict[str, str]:
+    def export_class_config(cls) -> tuple[dict[str, Any], dict[str, Any]]:
+        """
+        - **Description**:
+            - Exports the default configuration and descriptions for the configurable fields of the class.
+
+        - **Returns**:
+            - `tuple[Dict[str, Any], Dict[str, Any]]`: A tuple containing two dictionaries, one for default values and one for field descriptions.
+        """
         return (
             {
                 field: cls.default_values.get(field, "default_value")
@@ -182,11 +216,21 @@ class Block:
             {
                 field: cls.fields_description.get(field, "")
                 for field in cls.configurable_fields
-            }
+            },
         )
 
     @classmethod
     def import_config(cls, config: dict[str, Union[str, dict]]) -> Block:
+        """
+        - **Description**:
+            - Creates an instance of the Block from a configuration dictionary.
+
+        - **Args**:
+            - `config` (Dict[str, Union[str, dict]]): Configuration dictionary for creating the block.
+
+        - **Returns**:
+            - `Block`: An instance of the Block created from the provided configuration.
+        """
         instance = cls(name=config["name"])  # type: ignore
         assert isinstance(config["config"], dict)
         for field, value in config["config"].items():
@@ -202,7 +246,11 @@ class Block:
 
     def load_from_config(self, config: dict[str, list[dict]]) -> None:
         """
-        使用配置更新当前Block实例的参数，并递归更新子Block。
+        - **Description**:
+            - Updates the current Block instance parameters using a configuration dictionary and recursively updates its children.
+
+        - **Args**:
+            - `config` (Dict[str, List[Dict]]): Configuration dictionary for updating the block.
         """
         # 更新当前Block的参数
         for field in self.configurable_fields:
@@ -233,8 +281,11 @@ class Block:
 
     async def forward(self):
         """
-        Each block performs a specific reasoning task.
-        To be overridden by specific block implementations.
+        - **Description**:
+            - Each block performs a specific reasoning task. This method should be overridden by subclasses.
+
+        - **Raises**:
+            - `NotImplementedError`: Subclasses must implement this method.
         """
         raise NotImplementedError("Subclasses should implement this method")
 

pycityagent/workflow/prompt.py

@@ -1,5 +1,5 @@
-from typing import Optional, Union
 import re
+from typing import Optional, Union
 
 
 class FormatPrompt:
@@ -7,20 +7,21 @@ class FormatPrompt:
     A class to handle the formatting of prompts based on a template,
     with support for system prompts and variable extraction.
 
-    Attributes:
-        template (str): The template string containing placeholders.
-        system_prompt (Optional[str]): An optional system prompt to add to the dialog.
-        variables (list[str]): A list of variable names extracted from the template.
-        formatted_string (str): The formatted string derived from the template and provided variables.
+    - **Attributes**:
+        - `template` (str): The template string containing placeholders.
+        - `system_prompt` (Optional[str]): An optional system prompt to add to the dialog.
+        - `variables` (List[str]): A list of variable names extracted from the template.
+        - `formatted_string` (str): The formatted string derived from the template and provided variables.
     """
 
     def __init__(self, template: str, system_prompt: Optional[str] = None) -> None:
         """
-        Initializes the FormatPrompt with a template and an optional system prompt.
+        - **Description**:
+            - Initializes the FormatPrompt with a template and an optional system prompt.
 
-        Args:
-            template (str): The string template with variable placeholders.
-            system_prompt (Optional[str]): An optional system prompt.
+        - **Args**:
+            - `template` (str): The string template with variable placeholders.
+            - `system_prompt` (Optional[str], optional): An optional system prompt. Defaults to None.
         """
         self.template = template
         self.system_prompt = system_prompt  # Store the system prompt
@@ -29,22 +30,27 @@ class FormatPrompt:
 
     def _extract_variables(self) -> list[str]:
         """
-        Extracts variable names from the template string.
+        - **Description**:
+            - Extracts variable names from the template string using regular expressions.
 
-        Returns:
-            list[str]: A list of variable names found within the template.
+        - **Returns**:
+            - `List[str]`: A list of variable names found within the template.
         """
         return re.findall(r"\{(\w+)\}", self.template)
 
     def format(self, **kwargs) -> str:
         """
-        Formats the template string using the provided keyword arguments.
+        - **Description**:
+            - Formats the template string using the provided keyword arguments.
+
+        - **Args**:
+            - `**kwargs`: Variable names and their corresponding values to format the template.
 
-        Args:
-            **kwargs: Variable names and their corresponding values to format the template.
+        - **Returns**:
+            - `str`: The formatted string.
 
-        Returns:
-            str: The formatted string.
+        - **Raises**:
+            - `KeyError`: If a placeholder in the template does not have a corresponding key in kwargs.
         """
         self.formatted_string = self.template.format(
             **kwargs
@@ -53,10 +59,11 @@ class FormatPrompt:
 
     def to_dialog(self) -> list[dict[str, str]]:
         """
-        Converts the formatted prompt and optional system prompt into a dialog format.
+        - **Description**:
+            - Converts the formatted prompt and optional system prompt into a dialog format suitable for chat systems.
 
-        Returns:
-            list[dict[str, str]]: A list representing the dialog with roles and content.
+        - **Returns**:
+            - `List[Dict[str, str]]`: A list representing the dialog with roles and content.
         """
         dialog = []
         if self.system_prompt:
@@ -70,8 +77,9 @@ class FormatPrompt:
 
     def log(self) -> None:
         """
-        Logs the details of the FormatPrompt, including the template,
-        system prompt, extracted variables, and formatted string.
+        - **Description**:
+            - Logs the details of the FormatPrompt instance, including the template,
+              system prompt, extracted variables, and formatted string.
         """
         print(f"FormatPrompt: {self.template}")
         print(f"System Prompt: {self.system_prompt}")  # Log the system prompt