pycityagent 2.0.0a43__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/memory/profile.py

@@ -0,0 +1,165 @@
+"""
+Agent Profile
+"""
+
+from collections.abc import Callable, Sequence
+from copy import deepcopy
+from typing import Any, Optional, Union, cast
+
+from ..utils.decorators import lock_decorator
+from .const import *
+from .memory_base import MemoryBase, MemoryUnit
+from .utils import convert_msg_to_sequence
+
+
+class ProfileMemoryUnit(MemoryUnit):
+    def __init__(
+        self,
+        content: Optional[dict] = None,
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__(
+            content=content,
+            required_attributes=PROFILE_ATTRIBUTES,
+            activate_timestamp=activate_timestamp,
+        )
+
+
+class ProfileMemory(MemoryBase):
+    def __init__(
+        self,
+        msg: Optional[
+            Union[ProfileMemoryUnit, Sequence[ProfileMemoryUnit], dict, Sequence[dict]]
+        ] = None,
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__()
+        if msg is None:
+            msg = deepcopy(PROFILE_ATTRIBUTES)
+        self.activate_timestamp = activate_timestamp
+        msg = convert_msg_to_sequence(
+            msg,
+            sequence_type=ProfileMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            self._memories[unit] = {}
+
+    @lock_decorator
+    async def add(
+        self, msg: Union[ProfileMemoryUnit, Sequence[ProfileMemoryUnit]]
+    ) -> None:
+        _memories = self._memories
+        msg = convert_msg_to_sequence(
+            msg,
+            sequence_type=ProfileMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in _memories:
+                _memories[unit] = {}
+
+    @lock_decorator
+    async def pop(self, index: int) -> ProfileMemoryUnit:
+
+        _memories = self._memories
+        try:
+            pop_unit = list(_memories.keys())[index]
+            _memories.pop(pop_unit)
+
+            return pop_unit
+        except IndexError as e:
+
+            raise ValueError(f"Index {index} not in memory!")
+
+    @lock_decorator
+    async def load(
+        self,
+        snapshots: Union[dict, Sequence[dict]],
+        reset_memory: bool = False,
+    ) -> None:
+        if reset_memory:
+            self._memories = {}
+        msg = convert_msg_to_sequence(
+            snapshots,
+            sequence_type=ProfileMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in self._memories:
+                self._memories[unit] = {}
+
+    @lock_decorator
+    async def export(
+        self,
+    ) -> Sequence[dict]:
+        _res = []
+        for m in self._memories.keys():
+            m = cast(ProfileMemoryUnit, m)
+            _dict_values = await m.dict_values()
+            _res.append(_dict_values)
+
+        return _res
+
+    async def reset(self) -> None:
+        self._memories = {}
+
+    # interact
+    @lock_decorator
+    async def get(self, key: Any):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: ProfileMemoryUnit = _latest_memories[-1]
+        return _latest_memory[key]
+
+    @lock_decorator
+    async def get_top_k(
+        self,
+        key: Any,
+        metric: Callable[[Any], Any],
+        top_k: Optional[int] = None,
+        preserve_order: bool = True,
+    ) -> Union[Sequence[Any], Any]:
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: ProfileMemoryUnit = _latest_memories[-1]
+        _top_k = await _latest_memory.top_k_values(key, metric, top_k, preserve_order)
+
+        return _top_k
+
+    @lock_decorator
+    async def update(self, key: Any, value: Any, store_snapshot: bool = False):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: ProfileMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update({key: value})
+
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update({key: value})
+            msg = ProfileMemoryUnit(_content, self.activate_timestamp)
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}
+
+    @lock_decorator
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: ProfileMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update(to_update_dict)
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update(to_update_dict)
+            msg = ProfileMemoryUnit(_content, self.activate_timestamp)
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}

pycityagent/memory/self_define.py

@@ -0,0 +1,165 @@
+"""
+Self Define Data
+"""
+
+from collections.abc import Callable, Sequence
+from copy import deepcopy
+from typing import Any, Optional, Union, cast
+
+from ..utils.decorators import lock_decorator
+from .const import *
+from .memory_base import MemoryBase, MemoryUnit
+from .utils import convert_msg_to_sequence
+
+
+class DynamicMemoryUnit(MemoryUnit):
+    def __init__(
+        self,
+        content: Optional[dict] = None,
+        required_attributes: Optional[dict] = None,
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__(
+            content=content,
+            required_attributes=required_attributes,
+            activate_timestamp=activate_timestamp,
+        )
+
+
+class DynamicMemory(MemoryBase):
+
+    def __init__(
+        self,
+        required_attributes: dict[Any, Any],
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__()
+        self._required_attributes = deepcopy(required_attributes)
+        self.activate_timestamp = activate_timestamp
+        msg = DynamicMemoryUnit(
+            self._required_attributes, None, self.activate_timestamp
+        )
+        self._memories[msg] = {}
+
+    @lock_decorator
+    async def add(
+        self, msg: Union[DynamicMemoryUnit, Sequence[DynamicMemoryUnit]]
+    ) -> None:
+        _memories = self._memories
+        msg = convert_msg_to_sequence(
+            msg,
+            sequence_type=DynamicMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in _memories:
+                _memories[unit] = {}
+
+    @lock_decorator
+    async def pop(self, index: int) -> DynamicMemoryUnit:
+        _memories = self._memories
+        try:
+            pop_unit = list(_memories.keys())[index]
+            _memories.pop(pop_unit)
+
+            return pop_unit
+        except IndexError as e:
+
+            raise ValueError(f"Index {index} not in memory!")
+
+    @lock_decorator
+    async def load(
+        self,
+        snapshots: Union[dict, Sequence[dict]],
+        reset_memory: bool = False,
+    ) -> None:
+        if reset_memory:
+            self._memories = {}
+        msg = convert_msg_to_sequence(
+            snapshots,
+            sequence_type=DynamicMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in self._memories:
+                self._memories[unit] = {}
+
+    @lock_decorator
+    async def export(
+        self,
+    ) -> Sequence[dict]:
+        _res = []
+        for m in self._memories.keys():
+            m = cast(DynamicMemoryUnit, m)
+            _dict_values = await m.dict_values()
+            _res.append(_dict_values)
+
+        return _res
+
+    async def reset(self) -> None:
+        self._memories = {}
+
+    # interact
+    @lock_decorator
+    async def get(self, key: Any):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: DynamicMemoryUnit = _latest_memories[-1]
+
+        return _latest_memory[key]
+
+    @lock_decorator
+    async def get_top_k(
+        self,
+        key: Any,
+        metric: Callable[[Any], Any],
+        top_k: Optional[int] = None,
+        preserve_order: bool = True,
+    ) -> Union[Sequence[Any], Any]:
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: DynamicMemoryUnit = _latest_memories[-1]
+        _top_k = await _latest_memory.top_k_values(key, metric, top_k, preserve_order)
+
+        return _top_k
+
+    @lock_decorator
+    async def update(self, key: Any, value: Any, store_snapshot: bool = False):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: DynamicMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update({key: value})
+
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update({key: value})
+            msg = DynamicMemoryUnit(
+                _content, self._required_attributes, self.activate_timestamp
+            )
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}
+
+    @lock_decorator
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: DynamicMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update(to_update_dict)
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update(to_update_dict)
+            msg = DynamicMemoryUnit(
+                _content, self._required_attributes, self.activate_timestamp
+            )
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}

pycityagent/memory/state.py

@@ -0,0 +1,173 @@
+"""
+Agent State
+"""
+
+from collections.abc import Callable, Sequence
+from copy import deepcopy
+from typing import Any, Optional, Union, cast
+
+from ..utils.decorators import lock_decorator
+from .const import *
+from .memory_base import MemoryBase, MemoryUnit
+from .utils import convert_msg_to_sequence
+
+
+class StateMemoryUnit(MemoryUnit):
+    def __init__(
+        self,
+        content: Optional[dict] = None,
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__(
+            content=content,
+            required_attributes=STATE_ATTRIBUTES,
+            activate_timestamp=activate_timestamp,
+        )
+
+
+class StateMemory(MemoryBase):
+    def __init__(
+        self,
+        msg: Optional[
+            Union[MemoryUnit, Sequence[MemoryUnit], dict, Sequence[dict]]
+        ] = None,
+        activate_timestamp: bool = False,
+    ) -> None:
+        super().__init__()
+        if msg is None:
+            msg = deepcopy(STATE_ATTRIBUTES)
+        self.activate_timestamp = activate_timestamp
+        msg = convert_msg_to_sequence(
+            msg,
+            sequence_type=StateMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            self._memories[unit] = {}
+
+    @lock_decorator
+    async def add(self, msg: Union[MemoryUnit, Sequence[MemoryUnit]]) -> None:
+
+        _memories = self._memories
+        msg = convert_msg_to_sequence(
+            msg,
+            sequence_type=StateMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in _memories:
+                _memories[unit] = {}
+
+    @lock_decorator
+    async def pop(self, index: int) -> MemoryUnit:
+
+        _memories = self._memories
+        try:
+            pop_unit = list(_memories.keys())[index]
+            _memories.pop(pop_unit)
+
+            return pop_unit
+        except IndexError as e:
+
+            raise ValueError(f"Index {index} not in memory!")
+
+    @lock_decorator
+    async def load(
+        self,
+        snapshots: Union[dict, Sequence[dict]],
+        reset_memory: bool = False,
+    ) -> None:
+
+        if reset_memory:
+            self._memories = {}
+        msg = convert_msg_to_sequence(
+            snapshots,
+            sequence_type=StateMemoryUnit,
+            activate_timestamp=self.activate_timestamp,
+        )
+        for unit in msg:
+            if unit not in self._memories:
+                self._memories[unit] = {}
+
+    @lock_decorator
+    async def export(
+        self,
+    ) -> Sequence[dict]:
+
+        _res = []
+        for m in self._memories.keys():
+            m = cast(StateMemoryUnit, m)
+            _dict_values = await m.dict_values()
+            _res.append(_dict_values)
+
+        return _res
+
+    async def reset(self) -> None:
+
+        self._memories = {}
+
+    # interact
+    @lock_decorator
+    async def get(self, key: Any):
+
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: StateMemoryUnit = _latest_memories[-1]
+
+        return _latest_memory[key]
+
+    @lock_decorator
+    async def get_top_k(
+        self,
+        key: Any,
+        metric: Callable[[Any], Any],
+        top_k: Optional[int] = None,
+        preserve_order: bool = True,
+    ) -> Union[Sequence[Any], Any]:
+
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: StateMemoryUnit = _latest_memories[-1]
+        _top_k = await _latest_memory.top_k_values(key, metric, top_k, preserve_order)
+
+        return _top_k
+
+    @lock_decorator
+    async def update(self, key: Any, value: Any, store_snapshot: bool = False):
+
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: StateMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update({key: value})
+
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update({key: value})
+            msg = StateMemoryUnit(_content, self.activate_timestamp)
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}
+
+    @lock_decorator
+    async def update_dict(self, to_update_dict: dict, store_snapshot: bool = False):
+
+        _latest_memories = self._fetch_recent_memory()
+        _latest_memory: StateMemoryUnit = _latest_memories[-1]
+        if not store_snapshot:
+            # write in place
+            await _latest_memory.update(to_update_dict)
+
+        else:
+            # insert new unit
+            _dict_values = await _latest_memory.dict_values()
+            _content = deepcopy(
+                {k: v for k, v in _dict_values.items() if k not in {TIME_STAMP_KEY}}
+            )
+            _content.update(to_update_dict)
+            msg = StateMemoryUnit(_content, self.activate_timestamp)
+            for unit in [msg]:
+                if unit not in self._memories:
+                    self._memories[unit] = {}

pycityagent/memory/utils.py

@@ -0,0 +1,28 @@
+from collections.abc import Sequence
+from typing import Any, Union
+
+from .memory_base import MemoryUnit
+
+
+def convert_msg_to_sequence(
+    msg: Union[Any, Sequence[Any]],
+    sequence_type=MemoryUnit,
+    activate_timestamp: bool = False,
+) -> Sequence[Any]:
+    # Convert the input message to a sequence if it is not already one
+    if not isinstance(msg, Sequence):
+        _sequence_msg = [msg]
+    else:
+        _sequence_msg = msg
+
+    # Initialize an empty list to store the converted MemoryUnit objects
+    _sequence_unit = []
+
+    # Iterate over each unit in the sequence
+    for unit in _sequence_msg:
+        # If the unit is not already a MemoryUnit, convert it to one
+        if not isinstance(unit, sequence_type):
+            unit = sequence_type(content=unit, activate_timestamp=activate_timestamp)
+        _sequence_unit.append(unit)
+
+    return _sequence_unit

pycityagent/message/__init__.py

@@ -0,0 +1,3 @@
+from .messager import Messager
+
+__all__ = ["Messager"]

pycityagent/message/messager.py

@@ -0,0 +1,88 @@
+import asyncio
+import json
+import logging
+import math
+from typing import Any, List, Union
+from aiomqtt import Client
+import ray
+
+logger = logging.getLogger("pycityagent")
+
+@ray.remote
+class Messager:
+    def __init__(
+        self, hostname:str, port:int=1883, username=None, password=None, timeout=60
+    ):
+        self.client = Client(
+            hostname, port=port, username=username, password=password, timeout=timeout
+        )
+        self.connected = False  # 是否已连接标志
+        self.message_queue = asyncio.Queue()  # 用于存储接收到的消息
+        self.receive_messages_task = None
+    
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        await self.stop()
+
+    async def connect(self):
+        for i in range(3):
+            try:
+                await self.client.__aenter__()
+                self.connected = True
+                logger.info("Connected to MQTT Broker")
+                return
+            except Exception as e:
+                logger.error(f"Attempt {i+1}: Failed to connect to MQTT Broker: {e}")
+                await asyncio.sleep(10)
+        self.connected = False
+        logger.error("All connection attempts failed.")
+
+    async def disconnect(self):
+        await self.client.__aexit__(None, None, None)
+        self.connected = False
+        logger.info("Disconnected from MQTT Broker")
+
+    def is_connected(self):
+        """检查是否成功连接到 Broker"""
+        return self.connected
+
+    async def subscribe(self, topics: Union[str, List[str]], agents: Union[Any, List[Any]]):
+        if not self.is_connected():
+            logger.error(
+                f"Cannot subscribe to {topics} because not connected to the Broker."
+            )
+            return
+        if not isinstance(topics, list):
+            topics = [topics]
+        if not isinstance(agents, list):
+            agents = [agents]
+        await self.client.subscribe(topics, qos=1)
+
+    async def receive_messages(self):
+        """监听并将消息存入队列"""
+        async for message in self.client.messages:
+            await self.message_queue.put(message)
+
+    async def fetch_messages(self):
+        """从队列中批量获取消息"""
+        messages = []
+        while not self.message_queue.empty():
+            messages.append(await self.message_queue.get())
+        return messages
+
+    async def send_message(self, topic: str, payload: dict):
+        """通过 Messager 发送消息"""
+        message = json.dumps(payload, default=str)
+        await self.client.publish(topic=topic, payload=message, qos=1)
+        logger.info(f"Message sent to {topic}: {message}")
+
+    async def start_listening(self):
+        """启动消息监听任务"""
+        if self.is_connected():
+            self.receive_messages_task = asyncio.create_task(self.receive_messages())
+        else:
+            logger.error("Cannot start listening because not connected to the Broker.")
+
+    async def stop(self):
+        self.receive_messages_task.cancel()
+        await asyncio.gather(self.receive_messages_task, return_exceptions=True)
+        await self.disconnect()

pycityagent/metrics/__init__.py

@@ -0,0 +1,6 @@
+from .mlflow_client import MlflowClient,init_mlflow_connection
+
+__all__ = [
+    "MlflowClient",
+    "init_mlflow_connection",
+]