pycityagent 1.0.0__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

pycityagent/utils/parsers/code_block_parser.py

@@ -0,0 +1,37 @@
+import logging
+from typing import Any
+
+from .parser_base import ParserBase
+
+
+class CodeBlockParser(ParserBase):
+    """A parser that extracts specific objects from a response string enclosed within specific tags.
+
+    Attributes:
+        tag_start (str): The start tag used to identify the beginning of the object.
+        tag_end (str): The end tag used to identify the end of the object.
+    """
+
+    tag_start = "```{language_name}"
+    tag_end = "```"
+
+    def __init__(self, language_name: str) -> None:
+        """Initialize the CodeBlockParser with default tags."""
+        super().__init__()
+        self.language_name = language_name
+
+    def parse(self, response: str) -> str:
+        """Parse the response string to extract and return a str.
+
+        Parameters:
+            response (str): The response string containing the specified language object.
+
+        Returns:
+            str: The parsed `str` object.
+        """
+        extract_text = self._extract_text_within_tags(
+            response=response,
+            tag_start=self.tag_start.format(language_name=self.language_name),
+            tag_end=self.tag_end,
+        )
+        return extract_text

pycityagent/utils/parsers/json_parser.py

@@ -0,0 +1,86 @@
+import json
+import logging
+from typing import Any, Dict
+
+from .parser_base import ParserBase
+
+
+class JsonObjectParser(ParserBase):
+    """A parser that extracts and parses JSON objects from a response string enclosed within specific tags.
+
+    Attributes:
+        tag_start (str): The start tag used to identify the beginning of the JSON object.
+        tag_end (str): The end tag used to identify the end of the JSON object.
+    """
+
+    tag_start = "```json"
+    tag_end = "```"
+
+    def __init__(self) -> None:
+        """Initialize the JsonObjectParser with default tags."""
+        super().__init__()
+
+    def parse(self, response: str) -> Any:
+        """Parse the response string to extract and return a JSON object.
+
+        Parameters:
+            response (str): The response string containing the JSON object.
+
+        Returns:
+            Any: The parsed JSON object.
+        """
+        extract_text = self._extract_text_within_tags(
+            response=response,
+            tag_start=self.tag_start,
+            tag_end=self.tag_end,
+        )
+        try:
+            eval_parsed_json = eval(extract_text)
+            return eval_parsed_json
+        except Exception as e:
+            # logging.warning(f"Error when handling text {extract_text} with `eval`")
+            pass
+        try:
+            parsed_json = json.loads(extract_text)
+            return parsed_json
+        except json.decoder.JSONDecodeError as e:
+            raw_response = f"{self.tag_start}{extract_text}{self.tag_end}"
+            raise ValueError(
+                f"The content between {self.tag_start} and {self.tag_end} "
+                f"MUST be a JSON object."
+                f'When parsing "{raw_response}", an error occurred: {e}',
+            ) from None
+
+
+class JsonDictParser(JsonObjectParser):
+    """A parser that extends JsonObjectParser to ensure the parsed JSON is returned as a dictionary.
+
+    Attributes:
+        tag_start (str): The start tag used to identify the beginning of the JSON object.
+        tag_end (str): The end tag used to identify the end of the JSON object.
+    """
+
+    tag_start = "```json"
+    tag_end = "```"
+
+    def __init__(self) -> None:
+        """Initialize the JsonDictParser with default tags."""
+        super().__init__()
+
+    def parse(self, response: str) -> Dict:
+        """Parse the response string to extract and return a JSON object as a dictionary.
+
+        Parameters:
+            response (str): The response string containing the JSON object.
+
+        Returns:
+            Dict: The parsed JSON object as a dictionary.
+        """
+        parsed_json = super().parse(response)
+        if not isinstance(parsed_json, dict):
+            # If not a dictionary, raise an error
+            raise ValueError(
+                "A JSON dictionary object is wanted, "
+                f"but got {type(parsed_json)} instead."
+            )
+        return dict(parsed_json)

pycityagent/utils/parsers/parser_base.py

@@ -0,0 +1,60 @@
+"""
+Base class of parser
+"""
+
+import re
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+
+class ParserBase(ABC):
+    def __init__(
+        self,
+    ) -> None:
+        pass
+
+    @abstractmethod
+    def parse(self, response: str) -> Any:
+        """
+        Parse the string response returned by the model and convert it into a more manageable form.
+
+        Parameters:
+            response (str): The raw string returned by the model.
+
+        Returns:
+            Any: The converted data, the specific type depends on the parsing result.
+        """
+        pass
+
+    def _extract_text_within_tags(
+        self, response: str, tag_start: str, tag_end: str
+    ) -> str:
+        """
+        Return the string between the first occurrence of the start tag and the end tag in the response string.
+
+        Parameters:
+            response (str): The response string.
+            tag_start (str): The start tag.
+            tag_end (str): The end tag.
+
+        Returns:
+            str: The string between the start and end tags.
+        """
+
+        def _extract_substring(
+            text: str, tag_start: str, tag_end: str
+        ) -> Union[str, None]:
+            pattern = re.escape(tag_start) + r"(.*?)" + re.escape(tag_end)
+            match = re.search(pattern, text, re.DOTALL)
+            if match:
+                return match.group(1).strip()
+            else:
+                return None
+
+        sub_response = _extract_substring(response, tag_start, tag_end)
+        if sub_response is not None:
+            return sub_response
+        else:
+            raise ValueError(
+                f"No text between tag {tag_start} and tag {tag_end} in response {response}!"
+            )

pycityagent/workflow/__init__.py

@@ -0,0 +1,22 @@
+# __init__.py
+
+"""
+This module contains classes for creating blocks and running workflows.
+"""
+
+from .block import Block, log_and_check, log_and_check_with_memory
+from .prompt import FormatPrompt
+from .tool import GetMap, SencePOI, Tool
+from .trigger import MemoryChangeTrigger, PortMessageTrigger
+
+__all__ = [
+    "SencePOI",
+    "Tool",
+    "GetMap",
+    "MemoryChangeTrigger",
+    "PortMessageTrigger",
+    "Block",
+    "log_and_check",
+    "log_and_check_with_memory",
+    "FormatPrompt",
+]

pycityagent/workflow/block.py

@@ -0,0 +1,137 @@
+import asyncio
+import functools
+import inspect
+import time
+from typing import Any, Callable, Coroutine, Optional, Union
+
+from ..llm import LLM
+from ..memory import Memory
+from ..utils.decorators import record_call_aio
+
+TRIGGER_INTERVAL = 0.1
+
+
+def log_and_check_with_memory(
+    condition: Union[
+        Callable[[Memory], Coroutine[Any, Any, bool]],
+        Callable[[], Coroutine[Any, Any, bool]],
+        Callable[[Memory], bool],
+        Callable[[], bool],
+    ] = lambda: True,
+    trigger_interval: float = TRIGGER_INTERVAL,
+    record_function_calling: bool = False,
+):
+    """
+    A decorator that logs function calls and optionally checks a condition before executing the function.
+
+    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.
+                             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.
+    """
+
+    def decorator(func):
+        @record_call_aio(record_function_calling)
+        @functools.wraps(func)
+        async def wrapper(self, *args, **kwargs):
+            memory = None
+            for arg in list(args) + list(kwargs.values()):
+                if memory is not None:
+                    break
+                if isinstance(arg, Memory):
+                    memory = arg
+            assert isinstance(memory, Memory), "Input arguments has no `Memory` object!"
+            # Wait until the condition is met
+            sig = inspect.signature(condition)
+            params = list(sig.parameters.values())
+            if len(params) == 1 and params[0].annotation is Memory:
+                if inspect.iscoroutinefunction(condition):
+                    while not await condition(memory):  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+                else:
+                    while not condition(memory):  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+            elif len(params) == 0:
+                if inspect.iscoroutinefunction(condition):
+                    while not await condition():  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+                else:
+                    while not condition():  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+            else:
+                raise RuntimeError(
+                    f"Invalid parameter format in condition function {condition}"
+                )
+            result = await func(self, *args, **kwargs)
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def log_and_check(
+    condition: Union[
+        Callable[[], Coroutine[Any, Any, bool]],
+        Callable[[], bool],
+    ] = lambda: True,
+    trigger_interval: float = TRIGGER_INTERVAL,
+    record_function_calling: bool = False,
+):
+    """
+    A decorator that logs function calls and optionally checks a condition before executing the function.
+
+    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.
+                             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.
+    """
+
+    def decorator(func):
+        @record_call_aio(record_function_calling)
+        @functools.wraps(func)
+        async def wrapper(self, *args, **kwargs):
+            # Wait until the condition is met
+            sig = inspect.signature(condition)
+            params = list(sig.parameters.values())
+            if len(params) == 0:
+                if inspect.iscoroutinefunction(condition):
+                    while not await condition():  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+                else:
+                    while not condition():  # type:ignore
+                        await asyncio.sleep(trigger_interval)
+            else:
+                raise RuntimeError(
+                    f"Invalid parameter format in condition function {condition}"
+                )
+            result = await func(self, *args, **kwargs)
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+# Define a Block, similar to a layer in PyTorch
+class Block:
+    def __init__(
+        self,
+        name: str,
+        llm: Optional[LLM] = None,
+    ):
+        self.name = name
+        self.llm = llm
+
+    async def reason(self, *args, **kwargs):
+        """
+        Each block performs a specific reasoning task.
+        To be overridden by specific block implementations.
+        """
+        raise NotImplementedError("Subclasses should implement this method")

pycityagent/workflow/prompt.py

@@ -0,0 +1,72 @@
+from typing import Any, Callable, Dict, List, Optional, Union
+import re
+
+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.
+    """
+
+    def __init__(self, template: str, system_prompt: Optional[str] = None) -> None:
+        """
+        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.
+        """
+        self.template = template
+        self.system_prompt = system_prompt  # Store the system prompt
+        self.variables = self._extract_variables()
+        self.formatted_string = ''  # To store the formatted string
+
+    def _extract_variables(self) -> List[str]:
+        """
+        Extracts variable names from the template string.
+
+        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.
+
+        Args:
+            **kwargs: Variable names and their corresponding values to format the template.
+
+        Returns:
+            str: The formatted string.
+        """
+        self.formatted_string = self.template.format(**kwargs)  # Store the formatted string
+        return self.formatted_string
+
+    def to_dialog(self) -> List[Dict[str, str]]:
+        """
+        Converts the formatted prompt and optional system prompt into a dialog format.
+
+        Returns:
+            List[Dict[str, str]]: A list representing the dialog with roles and content.
+        """
+        dialog = []
+        if self.system_prompt:
+            dialog.append({'role': 'system', 'content': self.system_prompt})  # Add system prompt if it exists
+        dialog.append({'role': 'user', 'content': self.formatted_string})  # Add user content
+        return dialog
+
+    def log(self) -> None:
+        """
+        Logs the details of the FormatPrompt, 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
+        print(f"Variables: {self.variables}")
+        print(f"Formatted String: {self.formatted_string}")  # Log the formatted string

pycityagent/workflow/tool.py

@@ -0,0 +1,246 @@
+import asyncio
+import logging
+from copy import deepcopy
+from typing import Any, Callable, Dict, List, Optional, Union
+
+from mosstool.util.format_converter import dict2pb
+from pycityproto.city.person.v2 import person_pb2 as person_pb2
+
+from ..agent import Agent
+from ..environment import (LEVEL_ONE_PRE, POI_TYPE_DICT, AoiService,
+                           PersonService)
+from ..workflow import Block
+
+
+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.
+    """
+
+    def __get__(self, instance, owner):
+        if instance is None:
+            return self
+        subclass = type(self)
+        if not hasattr(instance, "_tools"):
+            instance._tools = {}
+        if subclass not in instance._tools:
+            tool_instance = subclass()
+            tool_instance._instance = instance  # type: ignore
+            instance._tools[subclass] = tool_instance
+        return instance._tools[subclass]
+
+    def __call__(self, *args: Any, **kwds: Any) -> Any:
+        """Invoke the tool's functionality.
+
+        This method must be implemented by subclasses to provide specific behavior.
+        """
+        raise NotImplementedError
+
+    @property
+    def agent(self) -> Agent:
+        instance = self._instance  # type:ignore
+        if not isinstance(instance, Agent):
+            raise RuntimeError(
+                f"Tool bind to object `{type(instance).__name__}`, not an `Agent` object!"
+            )
+        return instance
+
+    @property
+    def block(self) -> Block:
+        instance = self._instance  # type:ignore
+        if not isinstance(instance, Block):
+            raise RuntimeError(
+                f"Tool bind to object `{type(instance).__name__}`, not an `Block` object!"
+            )
+        return instance
+
+
+class GetMap(Tool):
+    """Retrieve the map from the simulator. Can be bound only to an `Agent` instance."""
+
+    def __init__(self) -> None:
+        self.variables = []
+
+    async def __call__(self) -> Union[Any, Callable]:
+        agent = self.agent
+        if agent.simulator is None:
+            raise ValueError("Simulator is not set.")
+        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.memory.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):
+    def __init__(
+        self,
+        person_template_func: Callable[[], dict] = PersonService.default_dict_person,
+    ) -> None:
+        self.person_template_func = person_template_func
+
+    async def _bind_to_simulator(
+        self,
+    ):
+        """
+        Bind Agent to Simulator
+
+        Args:
+            person_template (dict, optional): The person template in dict format. Defaults to PersonService.default_dict_person().
+        """
+        agent = self.agent
+        if agent._simulator is None:
+            return
+        if not agent._has_bound_to_simulator:
+            FROM_MEMORY_KEYS = {
+                "attribute",
+                "home",
+                "work",
+                "vehicle_attribute",
+                "bus_attribute",
+                "pedestrian_attribute",
+                "bike_attribute",
+            }
+            simulator = agent.simulator
+            memory = agent.memory
+            person_id = await memory.get("id")
+            # ATTENTION:模拟器分配的id从0开始
+            if person_id >= 0:
+                await simulator.GetPerson(person_id)
+                logging.debug(f"Binding to Person `{person_id}` already in Simulator")
+            else:
+                dict_person = deepcopy(self.person_template_func())
+                for _key in FROM_MEMORY_KEYS:
+                    try:
+                        _value = await memory.get(_key)
+                        if _value:
+                            dict_person[_key] = _value
+                    except KeyError as e:
+                        continue
+                resp = await simulator.AddPerson(
+                    dict2pb(dict_person, person_pb2.Person())
+                )
+                person_id = resp["person_id"]
+                await memory.update("id", person_id, protect_llm_read_only_fields=False)
+                logging.debug(
+                    f"Binding to Person `{person_id}` just added to Simulator"
+                )
+                # 防止模拟器还没有到prepare阶段导致GetPerson出错
+                await asyncio.sleep(5)
+            agent._has_bound_to_simulator = True
+        else:
+            pass
+
+    async def _update_motion_with_sim(
+        self,
+    ):
+        agent = self.agent
+        if agent._simulator is None:
+            return
+        if not agent._has_bound_to_simulator:
+            await self._bind_to_simulator()
+        simulator = agent.simulator
+        memory = agent.memory
+        person_id = await memory.get("id")
+        resp = await simulator.GetPerson(person_id)
+        resp_dict = resp["person"]
+        for k, v in resp_dict.get("motion", {}).items():
+            try:
+                await memory.get(k)
+                await memory.update(
+                    k, v, mode="replace", protect_llm_read_only_fields=False
+                )
+            except KeyError as e:
+                continue
+
+    async def __call__(
+        self,
+    ):
+        agent = self.agent
+        await self._bind_to_simulator()
+        await self._update_motion_with_sim()
+
+
+class ResetAgentPosition(Tool):
+    def __init__(self) -> None:
+        pass
+
+    async def __call__(
+        self,
+        aoi_id: Optional[int] = None,
+        poi_id: Optional[int] = None,
+        lane_id: Optional[int] = None,
+        s: Optional[float] = None,
+    ):
+        agent = self.agent
+        memory = agent.memory
+        await agent.simulator.ResetPersonPosition(
+            person_id=await memory.get("id"),
+            aoi_id=aoi_id,
+            poi_id=poi_id,
+            lane_id=lane_id,
+            s=s,
+        )