pycityagent 2.0.0a53__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a55__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/agent/agent.py

@@ -13,7 +13,8 @@ from ..economy import EconomyClient
 from ..environment import Simulator
 from ..llm import LLM
 from ..memory import Memory
-from ..message.messager import Messager
+from ..message import MessageInterceptor, Messager
+from ..metrics import MlflowClient
 from .agent_base import Agent, AgentType
 
 logger = logging.getLogger("pycityagent")
@@ -32,6 +33,7 @@ class CitizenAgent(Agent):
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,  # type:ignore
+        message_interceptor: Optional[MessageInterceptor] = None,  # type:ignore
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
@@ -40,10 +42,12 @@ class CitizenAgent(Agent):
             llm_client=llm_client,
             economy_client=economy_client,
             messager=messager,
+            message_interceptor=message_interceptor,
             simulator=simulator,
             memory=memory,
             avro_file=avro_file,
         )
+        self._mlflow_client = None
 
     async def bind_to_simulator(self):
         await self._bind_to_simulator()
@@ -78,9 +82,7 @@ class CitizenAgent(Agent):
                     dict_person[_key] = _value
             except KeyError as e:
                 continue
-        resp = await simulator.add_person(
-            dict2pb(dict_person, person_pb2.Person())
-        )
+        resp = await simulator.add_person(dict2pb(dict_person, person_pb2.Person()))
         person_id = resp["person_id"]
         await status.update("id", person_id, protect_llm_read_only_fields=False)
         logger.debug(f"Binding to Person `{person_id}` just added to Simulator")
@@ -123,6 +125,21 @@ class CitizenAgent(Agent):
         }
         await self._send_message(sender_id, payload, "gather")
 
+    @property
+    def mlflow_client(self) -> MlflowClient:
+        """The Agent's MlflowClient"""
+        if self._mlflow_client is None:
+            raise RuntimeError(
+                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
+            )
+        return self._mlflow_client
+
+    def set_mlflow_client(self, mlflow_client: MlflowClient):
+        """
+        Set the mlflow_client of the agent.
+        """
+        self._mlflow_client = mlflow_client
+
 
 class InstitutionAgent(Agent):
     """
@@ -137,6 +154,7 @@ class InstitutionAgent(Agent):
         memory: Optional[Memory] = None,
         economy_client: Optional[EconomyClient] = None,
         messager: Optional[Messager] = None,  # type:ignore
+        message_interceptor: Optional[MessageInterceptor] = None,  # type:ignore
         avro_file: Optional[dict] = None,
     ) -> None:
         super().__init__(
@@ -145,10 +163,12 @@ class InstitutionAgent(Agent):
             llm_client=llm_client,
             economy_client=economy_client,
             messager=messager,
+            message_interceptor=message_interceptor,
             simulator=simulator,
             memory=memory,
             avro_file=avro_file,
         )
+        self._mlflow_client = None
         # 添加响应收集器
         self._gather_responses: dict[str, asyncio.Future] = {}
 
@@ -308,3 +328,18 @@ class InstitutionAgent(Agent):
             # 清理Future
             for key in futures:
                 self._gather_responses.pop(key, None)
+
+    @property
+    def mlflow_client(self) -> MlflowClient:
+        """The Agent's MlflowClient"""
+        if self._mlflow_client is None:
+            raise RuntimeError(
+                f"MlflowClient access before assignment, please `set_mlflow_client` first!"
+            )
+        return self._mlflow_client
+
+    def set_mlflow_client(self, mlflow_client: MlflowClient):
+        """
+        Set the mlflow_client of the agent.
+        """
+        self._mlflow_client = mlflow_client

pycityagent/agent/agent_base.py

@@ -13,14 +13,13 @@ from typing import Any, Optional, Union, get_type_hints
 import fastavro
 import ray
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
-from pyparsing import Dict
 
 from ..economy import EconomyClient
 from ..environment import Simulator
 from ..environment.sim.person_service import PersonService
 from ..llm import LLM
 from ..memory import Memory
-from ..message.messager import Messager
+from ..message import MessageInterceptor, Messager
 from ..metrics import MlflowClient
 from ..utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
 from ..workflow import Block
@@ -56,7 +55,8 @@ class Agent(ABC):
         type: AgentType = AgentType.Unspecified,
         llm_client: Optional[LLM] = None,
         economy_client: Optional[EconomyClient] = None,
-        messager: Optional[Messager] = None,  # type:ignore
+        messager: Optional[ray.ObjectRef] = None,
+        message_interceptor: Optional[ray.ObjectRef] = None,
         simulator: Optional[Simulator] = None,
         memory: Optional[Memory] = None,
         avro_file: Optional[dict[str, str]] = None,
@@ -82,6 +82,7 @@ class Agent(ABC):
         self._llm_client = llm_client
         self._economy_client = economy_client
         self._messager = messager
+        self._message_interceptor = message_interceptor
         self._simulator = simulator
         self._memory = memory
         self._exp_id = -1
@@ -102,12 +103,12 @@ class Agent(ABC):
         return state
 
     @classmethod
-    def export_class_config(cls) -> Dict[str, Dict]:
+    def export_class_config(cls) -> dict[str, dict]:
         result = {
             "agent_name": cls.__name__,
             "config": {},
             "description": {},
-            "blocks": []
+            "blocks": [],
         }
         config = {
             field: cls.default_values.get(field, "default_value")
@@ -123,16 +124,18 @@ class Agent(ABC):
         for attr_name, attr_type in hints.items():
             if inspect.isclass(attr_type) and issubclass(attr_type, Block):
                 block_config = attr_type.export_class_config()
-                result["blocks"].append({
-                    "name": attr_name,
-                    "config": block_config[0],
-                    "description": block_config[1],
-                    "children": cls._export_subblocks(attr_type)
-                })
+                result["blocks"].append(
+                    {
+                        "name": attr_name,
+                        "config": block_config[0],  # type:ignore
+                        "description": block_config[1],  # type:ignore
+                        "children": cls._export_subblocks(attr_type),
+                    }
+                )
         return result
 
     @classmethod
-    def _export_subblocks(cls, block_cls: type[Block]) -> list[Dict]:
+    def _export_subblocks(cls, block_cls: type[Block]) -> list[dict]:
         children = []
         hints = get_type_hints(block_cls)  # 获取类的注解
         for attr_name, attr_type in hints.items():
@@ -141,8 +144,8 @@ class Agent(ABC):
                 children.append(
                     {
                         "name": attr_name,
-                        "config": block_config[0],
-                    "description": block_config[1],
+                        "config": block_config[0],  # type:ignore
+                        "description": block_config[1],  # type:ignore
                         "children": cls._export_subblocks(attr_type),
                     }
                 )
@@ -253,6 +256,12 @@ class Agent(ABC):
         """
         self._pgsql_writer = pgsql_writer
 
+    def set_message_interceptor(self, message_interceptor: ray.ObjectRef):
+        """
+        Set the PostgreSQL copy writer of the agent.
+        """
+        self._message_interceptor = message_interceptor
+
     @property
     def uuid(self):
         """The Agent's UUID"""
@@ -289,24 +298,24 @@ class Agent(ABC):
                 f"Memory access before assignment, please `set_memory` first!"
             )
         return self._memory
-    
+
     @property
     def status(self):
         """The Agent's Status Memory"""
-        if self._memory.status is None:
+        if self.memory.status is None:
             raise RuntimeError(
                 f"Status access before assignment, please `set_memory` first!"
             )
-        return self._memory.status
-    
+        return self.memory.status
+
     @property
     def stream(self):
         """The Agent's Stream Memory"""
-        if self._memory.stream is None:
+        if self.memory.stream is None:
             raise RuntimeError(
                 f"Stream access before assignment, please `set_memory` first!"
             )
-        return self._memory.stream
+        return self.memory.stream
 
     @property
     def simulator(self):
@@ -335,7 +344,7 @@ class Agent(ABC):
     async def messager_ping(self):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
-        return await self._messager.ping()
+        return await self._messager.ping.remote()  # type:ignore
 
     async def generate_user_survey_response(self, survey: dict) -> str:
         """生成回答 —— 可重写
@@ -421,7 +430,7 @@ class Agent(ABC):
                     _data_tuples
                 )
             )
-        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})# type:ignore
 
     async def generate_user_chat_response(self, question: str) -> str:
         """生成回答 —— 可重写
@@ -508,7 +517,7 @@ class Agent(ABC):
                     _data
                 )
             )
-        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})
+        await self.messager.send_message.remote(f"exps/{self._exp_id}/user_payback", {"count": 1})# type:ignore
         print(f"Sent payback message to {self._exp_id}")
 
     async def process_agent_chat_response(self, payload: dict) -> str:
@@ -579,7 +588,12 @@ class Agent(ABC):
         if self._messager is None:
             raise RuntimeError("Messager is not set")
         topic = f"exps/{self._exp_id}/agents/{to_agent_uuid}/{sub_topic}"
-        await self._messager.send_message.remote(topic, payload)
+        await self._messager.send_message.remote(  # type:ignore
+            topic,
+            payload,
+            self._uuid,
+            to_agent_uuid,
+        )
 
     async def send_message_to_agent(
         self, to_agent_uuid: str, content: str, type: str = "social"
@@ -643,6 +657,6 @@ class Agent(ABC):
         当_blocked为True时，不执行forward方法
         """
         if self._messager is not None:
-            await self._messager.ping.remote()
+            await self._messager.ping.remote()  # type:ignore
         if not self._blocked:
             await self.forward()

pycityagent/cityagent/blocks/plan_block.py

@@ -294,4 +294,4 @@ Execution Steps: \n{formated_steps}
     def clean_json_response(self, response: str) -> str:
         """Clean special characters in LLM response"""
         response = response.replace('```json', '').replace('```', '')
-        return response.strip() 
+        return response.strip() 

pycityagent/cityagent/message_intercept.py

@@ -0,0 +1,99 @@
+import asyncio
+
+from pycityagent.llm import LLM
+from pycityagent.message import MessageBlockBase, MessageBlockListenerBase
+
+
+async def check_message(
+    from_uuid: str, to_uuid: str, llm_client: LLM, content: str
+) -> bool:
+    """
+    使用LLM检查消息是否合规
+    返回: (是否合规, from_uuid, to_uuid)
+    """
+    print(f"\n检查消息: {from_uuid} -> {to_uuid}: {content}")
+
+    prompt = f"""
+        请判断以下消息是否具有情绪煽动性：
+        发送者ID: {from_uuid}
+        接收者ID: {to_uuid}
+        消息内容: {content}
+        
+        如果消息具有情绪煽动性，请返回 False；如果消息正常，请返回 True。
+        """
+
+    if "test" in content.lower():
+        is_valid = False
+    else:
+        is_valid = True
+    print(f"消息检查结果: {'合规' if is_valid else '不合规'}")
+    return is_valid
+
+
+class EdgeMessageBlock(MessageBlockBase):
+    async def forward(  # type:ignore
+        self,
+        from_uuid: str,
+        to_uuid: str,
+        msg: str,
+        black_list: list[tuple[str, str]],
+    ):
+        if (from_uuid, to_uuid) in set(black_list):
+            # 可选同时返回入队的信息(False,err) 如果只返回bool值则默认报错信息入队
+            return False
+        else:
+            is_valid = await check_message(
+                from_uuid=from_uuid,
+                to_uuid=to_uuid,
+                llm_client=self.llm,
+                content=msg,
+            )
+            if not is_valid:
+                # 直接添加即可 在框架内部的异步锁保证不会冲突
+                black_list.append((from_uuid, to_uuid))
+            return is_valid
+
+
+class PointMessageBlock(MessageBlockBase):
+    async def forward(  # type:ignore
+        self,
+        from_uuid: str,
+        to_uuid: str,
+        msg: str,
+        violation_counts: dict[str, int],
+        black_list: list[tuple[str, str]],
+    ):
+        if (from_uuid, to_uuid) in set(black_list):
+            # 可选同时返回入队的信息(False,err) 如果只返回bool值则默认报错信息入队
+            return False
+        else:
+            # violation count在框架内自动维护 这里不用管
+            is_valid = await check_message(
+                from_uuid=from_uuid,
+                to_uuid=to_uuid,
+                llm_client=self.llm,
+                content=msg,
+            )
+            if not is_valid and violation_counts[from_uuid] >= 3 - 1:
+                # 直接添加即可 在框架内部的异步锁保证不会冲突
+                black_list.append((from_uuid, to_uuid))
+            return is_valid
+
+
+class MessageBlockListener(MessageBlockListenerBase):
+    def __init__(
+        self, save_queue_values: bool = False, get_queue_period: float = 0.1
+    ) -> None:
+        super().__init__(save_queue_values, get_queue_period)
+
+    async def forward(
+        self,
+    ):
+        while True:
+            if self.has_queue:
+                value = await self.queue.get_async()  # type: ignore
+                if self._save_queue_values:
+                    self._values_from_queue.append(value)
+                print(f"get `{value}` from queue")
+                # do something with the value
+            await asyncio.sleep(self._get_queue_period)

pycityagent/cityagent/societyagent.py

@@ -8,8 +8,8 @@ from pycityagent.agent import Agent
 from pycityagent.economy import EconomyClient
 from pycityagent.llm.llm import LLM
 from pycityagent.memory import Memory
-from pycityagent.workflow import Block
 from pycityagent.tools import UpdateWithSimulator
+from pycityagent.workflow import Block
 
 from .blocks import (CognitionBlock, EconomyBlock, MobilityBlock, NeedsBlock,
                      OtherBlock, PlanBlock, SocialBlock)
@@ -54,7 +54,6 @@ class PlanAndActionBlock(Block):
             llm=llm, memory=memory, simulator=simulator, economy_client=economy_client
         )
         self.otherBlock = OtherBlock(llm=llm, memory=memory)
-    
     async def plan_generation(self):
         """Generate plan"""
         current_plan = await self.memory.status.get("current_plan")
@@ -79,45 +78,43 @@ class PlanAndActionBlock(Block):
             )
             result = None
             if step_type == "mobility":
-                if self.enable_mobility:
+                if self.enable_mobility:  # type:ignore
                     result = await self.mobilityBlock.forward(
                         current_step, execution_context
                     )
                 else:
                     result = {
-                        'success': False,
-                        'evaluation': f'Mobility Behavior is disabled',
-                        'consumed_time': 0,
-                        'node_id': None
+                        "success": False,
+                        "evaluation": f"Mobility Behavior is disabled",
+                        "consumed_time": 0,
+                        "node_id": None,
                     }
             elif step_type == "social":
-                if self.enable_social:
+                if self.enable_social:  # type:ignore
                     result = await self.socialBlock.forward(
                         current_step, execution_context
                     )
                 else:
                     result = {
-                        'success': False,
-                        'evaluation': f'Social Behavior is disabled',
-                        'consumed_time': 0,
-                        'node_id': None
+                        "success": False,
+                        "evaluation": f"Social Behavior is disabled",
+                        "consumed_time": 0,
+                        "node_id": None,
                     }
             elif step_type == "economy":
-                if self.enable_economy:
+                if self.enable_economy:  # type:ignore
                     result = await self.economyBlock.forward(
                         current_step, execution_context
                     )
                 else:
                     result = {
-                        'success': False,
-                        'evaluation': f'Economy Behavior is disabled',
-                        'consumed_time': 0,
-                        'node_id': None
+                        "success": False,
+                        "evaluation": f"Economy Behavior is disabled",
+                        "consumed_time": 0,
+                        "node_id": None,
                     }
             elif step_type == "other":
-                result = await self.otherBlock.forward(
-                    current_step, execution_context
-                )
+                result = await self.otherBlock.forward(current_step, execution_context)
             if result != None:
                 logger.info(f"Execution result: {result}")
                 current_step["evaluation"] = result
@@ -156,7 +153,9 @@ class MindBlock(Block):
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__(name="mind_block", llm=llm, memory=memory, simulator=simulator)
-        self.cognitionBlock = CognitionBlock(llm=self.llm, memory=self.memory, simulator=simulator)
+        self.cognitionBlock = CognitionBlock(
+            llm=self.llm, memory=self.memory, simulator=simulator
+        )
 
     async def forward(self):
         await self.cognitionBlock.forward()
@@ -165,7 +164,12 @@ class SocietyAgent(CitizenAgent):
     mindBlock: MindBlock
     planAndActionBlock: PlanAndActionBlock
     update_with_sim = UpdateWithSimulator()
-    configurable_fields = ["enable_cognition", "enable_mobility", "enable_social", "enable_economy"]
+    configurable_fields = [
+        "enable_cognition",
+        "enable_mobility",
+        "enable_social",
+        "enable_economy",
+    ]
     default_values = {
         "enable_cognition": True,
         "enable_mobility": True,
@@ -195,14 +199,14 @@ class SocietyAgent(CitizenAgent):
             economy_client=economy_client,
         )
         self.mindBlock = MindBlock(
-            llm=self._llm_client, memory=self.memory, simulator=self._simulator
+            llm=self.llm, memory=self.memory, simulator=self.simulator
         )
         self.planAndActionBlock = PlanAndActionBlock(
             agent=self,
-            llm=self._llm_client,
+            llm=self.llm,
             memory=self.memory,
-            simulator=self._simulator,
-            economy_client=self._economy_client,
+            simulator=self.simulator,
+            economy_client=self.economy_client,
         )
         self.step_count = -1
         self.cognition_update = -1
@@ -305,7 +309,114 @@ class SocietyAgent(CitizenAgent):
         # The previous step has not been completed
         return False
 
-    async def process_agent_chat_response(self, payload: dict) -> str:
+        # check last step
+        if not await self.check_and_update_step():
+            return
+
+        await self.planAndActionBlock.forward()
+
+        if self.enable_cognition:
+            await self.mindBlock.forward()
+
+    async def check_and_update_step(self):
+        """Check if the previous step has been completed"""
+        status = await self.memory.status.get("status")
+        if status == 2:
+            # Agent is moving
+            logger.info("Agent is moving")
+            await asyncio.sleep(1)
+            return False
+
+        # Get the previous step information
+        current_step = await self.memory.status.get("current_step")
+        if current_step["intention"] == "" or current_step["type"] == "":
+            # No previous step, return directly
+            return True
+        time_now = int(await self.simulator.get_time())
+        step_start_time = current_step["start_time"]
+        step_consumed_time = current_step["evaluation"]["consumed_time"]
+        time_end_plan = step_start_time + int(step_consumed_time) * 60
+        if time_now >= time_end_plan:
+            # The previous step has been completed
+            current_plan = await self.memory.status.get("current_plan")
+            current_step["evaluation"]["consumed_time"] = (
+                time_now - step_start_time
+            ) / 60
+            current_plan["stream_nodes"].append(current_step["evaluation"]["node_id"])
+            if current_step["evaluation"]["success"]:
+                # Last step is completed
+                current_step_index = next(
+                    (
+                        i
+                        for i, step in enumerate(current_plan["steps"])
+                        if step["intention"] == current_step["intention"]
+                        and step["type"] == current_step["type"]
+                    ),
+                    None,
+                )
+                current_plan["steps"][current_step_index] = current_step
+                await self.memory.status.update("current_plan", current_plan)
+                if current_step_index is not None and current_step_index + 1 < len(
+                    current_plan["steps"]
+                ):
+                    next_step = current_plan["steps"][current_step_index + 1]
+                    await self.memory.status.update("current_step", next_step)
+                else:
+                    # Whole plan is completed
+                    current_plan["completed"] = True
+                    current_plan["end_time"] = await self.simulator.get_time(
+                        format_time=True
+                    )
+                    if self.enable_cognition:
+                        # Update emotion for the plan
+                        related_memories = await self.memory.stream.get_by_ids(
+                            current_plan["stream_nodes"]
+                        )
+                        incident = f"You have successfully completed the plan: {related_memories}"
+                        conclusion = await self.mindBlock.cognitionBlock.emotion_update(
+                            incident
+                        )
+                        await self.memory.stream.add_cognition(
+                            description=conclusion  # type:ignore
+                        )
+                        await self.memory.stream.add_cognition_to_memory(
+                            current_plan["stream_nodes"], conclusion  # type:ignore
+                        )
+                    await self.memory.status.update("current_plan", current_plan)
+                    await self.memory.status.update(
+                        "current_step", {"intention": "", "type": ""}
+                    )
+                return True
+            else:
+                current_plan["failed"] = True
+                current_plan["end_time"] = await self.simulator.get_time(
+                    format_time=True
+                )
+                if self.enable_cognition:
+                    # Update emotion for the plan
+                    related_memories = await self.memory.stream.get_by_ids(
+                        current_plan["stream_nodes"]
+                    )
+                    incident = (
+                        f"You have failed to complete the plan: {related_memories}"
+                    )
+                    conclusion = await self.mindBlock.cognitionBlock.emotion_update(
+                        incident
+                    )
+                    await self.memory.stream.add_cognition(
+                        description=conclusion  # type:ignore
+                    )
+                    await self.memory.stream.add_cognition_to_memory(
+                        current_plan["stream_nodes"], conclusion  # type:ignore
+                    )
+                await self.memory.status.update("current_plan", current_plan)
+                await self.memory.status.update(
+                    "current_step", {"intention": "", "type": ""}
+                )
+        # The previous step has not been completed
+        return False
+
+    async def process_agent_chat_response(self, payload: dict) -> str:  # type:ignore
         if payload["type"] == "social":
             resp = f"Agent {self._uuid} received agent chat response: {payload}"
             logger.info(resp)
@@ -380,7 +491,7 @@ class SocietyAgent(CitizenAgent):
 
         Answer only YES or NO."""
 
-                should_respond = await self._llm_client.atext_request(
+                should_respond = await self._llm_client.atext_request(  # type:ignore
                     [
                         {
                             "role": "system",
@@ -390,7 +501,7 @@ class SocietyAgent(CitizenAgent):
                     ]
                 )
 
-                if should_respond.strip().upper() != "YES":
+                if should_respond.strip().upper() != "YES":  # type:ignore
                     await self.memory.status.update("chat_histories", chat_histories)
                     return ""
 
@@ -414,7 +525,7 @@ class SocietyAgent(CitizenAgent):
 
         Response should be ONLY the message text, no explanations."""
 
-                response = await self._llm_client.atext_request(
+                response = await self.llm.atext_request(
                     [
                         {
                             "role": "system",
@@ -440,7 +551,7 @@ class SocietyAgent(CitizenAgent):
                     await self.send_message_to_agent(sender_id, serialized_response)
                     logger.info("sender_id", sender_id)
                     logger.info("message", serialized_response)
-                return response
+                return response  # type:ignore
 
             except Exception as e:
                 logger.warning(f"Error in process_agent_chat_response: {str(e)}")
@@ -454,6 +565,8 @@ class SocietyAgent(CitizenAgent):
                 value = int(value)
             description = f"You received a economic message: Your {key} has changed from {await self.memory.status.get(key)} to {value}"
             await self.memory.status.update(key, value)
-            await self.memory.stream.add_economic(description=description)
+            await self.memory.stream.add_economic(  # type:ignore
+                description=description
+            )
             if self.enable_cognition:
                 await self.mindBlock.cognitionBlock.emotion_update(description)