pycityagent 2.0.0a49__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a51__cp310-cp310-macosx_11_0_arm64.whl

pycityagent/cityagent/nbsagent.py

@@ -1,30 +1,41 @@
 import asyncio
+import logging
 from typing import Optional
 
 import numpy as np
-from pycityagent import Simulator, InstitutionAgent
-from pycityagent.llm.llm import LLM
+
+from pycityagent import InstitutionAgent, Simulator
 from pycityagent.economy import EconomyClient
-from pycityagent.message import Messager
+from pycityagent.llm.llm import LLM
 from pycityagent.memory import Memory
-from pycityagent.workflow.tool import ExportMlflowMetrics
-import logging
+from pycityagent.message import Messager
+from pycityagent.tools import ExportMlflowMetrics
 
 logger = logging.getLogger("pycityagent")
 
+
 class NBSAgent(InstitutionAgent):
     export_metrics = ExportMlflowMetrics(log_batch_size=3)
-    
-    def __init__(self, 
-                 name: str,
-                 llm_client: Optional[LLM] = None,
-                 simulator: Optional[Simulator] = None,
-                 memory: Optional[Memory] = None,
-                 economy_client: Optional[EconomyClient] = None,
-                 messager: Optional[Messager] = None,
-                 avro_file: Optional[dict] = None,
-                 ) -> None:
-        super().__init__(name=name, llm_client=llm_client, simulator=simulator, memory=memory, economy_client=economy_client, messager=messager, avro_file=avro_file)
+
+    def __init__(
+        self,
+        name: str,
+        llm_client: Optional[LLM] = None,
+        simulator: Optional[Simulator] = None,
+        memory: Optional[Memory] = None,
+        economy_client: Optional[EconomyClient] = None,
+        messager: Optional[Messager] = None,
+        avro_file: Optional[dict] = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            llm_client=llm_client,
+            simulator=simulator,
+            memory=memory,
+            economy_client=economy_client,
+            messager=messager,
+            avro_file=avro_file,
+        )
         self.initailzed = False
         self.last_time_trigger = None
         self.time_diff = 30 * 24 * 60 * 60
@@ -32,6 +43,7 @@ class NBSAgent(InstitutionAgent):
         self.num_labor_hours = 168
         self.productivity_per_labor = 1
         self.price = 1
+
     async def month_trigger(self):
         now_time = await self.simulator.get_time()
         if self.last_time_trigger is None:
@@ -41,46 +53,80 @@ class NBSAgent(InstitutionAgent):
             self.last_time_trigger = now_time
             return True
         return False
-    
+
     async def gather_messages(self, agent_ids, content):
         infos = await super().gather_messages(agent_ids, content)
-        return [info['content'] for info in infos]
+        return [info["content"] for info in infos]
 
     async def forward(self):
         if await self.month_trigger():
             citizens = await self.memory.get("citizens")
             while True:
-                agents_forward = await self.gather_messages(citizens, 'forward')
+                agents_forward = await self.gather_messages(citizens, "forward")
                 if np.all(np.array(agents_forward) > self.forward_times):
                     break
                 await asyncio.sleep(1)
             work_propensity = await self.gather_messages(citizens, "work_propensity")
-            working_hours = np.mean(work_propensity)*self.num_labor_hours
+            working_hours = np.mean(work_propensity) * self.num_labor_hours
             firm_id = await self.memory.get("firm_id")
             price = await self.economy_client.get(firm_id, "price")
             prices = await self.economy_client.get(self._agent_id, "prices")
             initial_price = prices[0]
-            nominal_gdp = working_hours*len(citizens)*self.productivity_per_labor*price
-            real_gdp = working_hours*len(citizens)*self.productivity_per_labor*initial_price
-            await self.economy_client.update(self._agent_id, 'nominal_gdp', [nominal_gdp], mode='merge')
-            await self.economy_client.update(self._agent_id, 'real_gdp', [real_gdp], mode='merge')
-            await self.economy_client.update(self._agent_id, 'working_hours', [working_hours], mode='merge')
-            await self.economy_client.update(self._agent_id, 'prices', [price], mode='merge')
-            depression = await self.gather_messages(citizens, 'depression')
+            nominal_gdp = (
+                working_hours * len(citizens) * self.productivity_per_labor * price
+            )
+            real_gdp = (
+                working_hours
+                * len(citizens)
+                * self.productivity_per_labor
+                * initial_price
+            )
+            await self.economy_client.update(
+                self._agent_id, "nominal_gdp", [nominal_gdp], mode="merge"
+            )
+            await self.economy_client.update(
+                self._agent_id, "real_gdp", [real_gdp], mode="merge"
+            )
+            await self.economy_client.update(
+                self._agent_id, "working_hours", [working_hours], mode="merge"
+            )
+            await self.economy_client.update(
+                self._agent_id, "prices", [price], mode="merge"
+            )
+            depression = await self.gather_messages(citizens, "depression")
             depression = np.mean(depression)
-            await self.economy_client.update(self._agent_id, 'depression', [depression], mode='merge')
-            consumption_currency = await self.gather_messages(citizens, 'consumption_currency')
+            await self.economy_client.update(
+                self._agent_id, "depression", [depression], mode="merge"
+            )
+            consumption_currency = await self.gather_messages(
+                citizens, "consumption_currency"
+            )
             consumption_currency = np.mean(consumption_currency)
-            await self.economy_client.update(self._agent_id, 'consumption_currency', [consumption_currency], mode='merge')
-            income_currency = await self.gather_messages(citizens, 'income_currency')
+            await self.economy_client.update(
+                self._agent_id,
+                "consumption_currency",
+                [consumption_currency],
+                mode="merge",
+            )
+            income_currency = await self.gather_messages(citizens, "income_currency")
             income_currency = np.mean(income_currency)
-            await self.economy_client.update(self._agent_id, 'income_currency', [income_currency], mode='merge')
+            await self.economy_client.update(
+                self._agent_id, "income_currency", [income_currency], mode="merge"
+            )
             self.forward_times += 1
             for uuid in citizens:
-                await self.send_message_to_agent(uuid, f"nbs_forward@{self.forward_times}")
+                await self.send_message_to_agent(
+                    uuid, f"nbs_forward@{self.forward_times}"
+                )
 
-            metrics = {'nominal_gdp': nominal_gdp, 'working_hours': working_hours, 'price': price,
-                        'depression': depression, 'consumption': consumption_currency, 'income': income_currency}
+            metrics = {
+                "nominal_gdp": nominal_gdp,
+                "working_hours": working_hours,
+                "price": price,
+                "depression": depression,
+                "consumption": consumption_currency,
+                "income": income_currency,
+            }
             for k, v in metrics.items():
                 await self.export_metrics(
                     metric={
@@ -89,4 +135,4 @@ class NBSAgent(InstitutionAgent):
                         "step": self.forward_times,
                     },
                     clear_cache=True,
-                )
+                )

pycityagent/cityagent/societyagent.py

@@ -1,22 +1,27 @@
 import asyncio
 import json
+import logging
 from typing import Optional
-from pycityagent import Simulator, CitizenAgent
+
+from pycityagent import CitizenAgent, Simulator
 from pycityagent.agent import Agent
-from pycityagent.llm.llm import LLM
 from pycityagent.economy import EconomyClient
-from pycityagent.message import Messager
+from pycityagent.llm.llm import LLM
 from pycityagent.memory import Memory
-from pycityagent.workflow.tool import UpdateWithSimulator
+from pycityagent.message import Messager
 from pycityagent.workflow import Block
-from .blocks import CognitionBlock, EconomyBlock, MobilityBlock, NeedsBlock, OtherBlock, PlanBlock, SocialBlock
+from pycityagent.tools import UpdateWithSimulator
+
+from .blocks import (CognitionBlock, EconomyBlock, MobilityBlock, NeedsBlock,
+                     OtherBlock, PlanBlock, SocialBlock)
 from .blocks.economy_block import MonthPlanBlock
-import logging
 
 logger = logging.getLogger("pycityagent")
 
+
 class PlanAndActionBlock(Block):
     """主动工作流"""
+
     longTermDecisionBlock: MonthPlanBlock
     needsBlock: NeedsBlock
     planBlock: PlanBlock
@@ -25,48 +30,77 @@ class PlanAndActionBlock(Block):
     economyBlock: EconomyBlock
     otherBlock: OtherBlock
 
-    def __init__(self, agent: Agent, llm: LLM, memory: Memory, simulator: Simulator, economy_client: EconomyClient):
-        super().__init__(name="plan_and_action_block", llm=llm, memory=memory, simulator=simulator)
+    def __init__(
+        self,
+        agent: Agent,
+        llm: LLM,
+        memory: Memory,
+        simulator: Simulator,
+        economy_client: EconomyClient,
+    ):
+        super().__init__(
+            name="plan_and_action_block", llm=llm, memory=memory, simulator=simulator
+        )
         self._agent = agent
-        self.longTermDecisionBlock = MonthPlanBlock(llm=llm, memory=memory, simulator=simulator, economy_client=economy_client)
+        self.longTermDecisionBlock = MonthPlanBlock(
+            llm=llm, memory=memory, simulator=simulator, economy_client=economy_client
+        )
         self.needsBlock = NeedsBlock(llm=llm, memory=memory, simulator=simulator)
         self.planBlock = PlanBlock(llm=llm, memory=memory, simulator=simulator)
         self.mobilityBlock = MobilityBlock(llm=llm, memory=memory, simulator=simulator)
-        self.socialBlock = SocialBlock(agent=self, llm=llm, memory=memory, simulator=simulator)
-        self.economyBlock = EconomyBlock(llm=llm, memory=memory, simulator=simulator, economy_client=economy_client)
+        self.socialBlock = SocialBlock(
+            agent=self, llm=llm, memory=memory, simulator=simulator
+        )
+        self.economyBlock = EconomyBlock(
+            llm=llm, memory=memory, simulator=simulator, economy_client=economy_client
+        )
         self.otherBlock = OtherBlock(llm=llm, memory=memory)
 
     async def check_and_update_step(self):
-        status = await self.memory.get('status')
+        status = await self.memory.get("status")
         if status == 2:
             # 正在运动
             logger.info("Agent is moving")
             await asyncio.sleep(1)
             return False
-        
+
         # 获取上一步信息
         current_step = await self.memory.get("current_step")
-        if current_step['intention'] == "" or current_step['type'] == "":
+        if current_step["intention"] == "" or current_step["type"] == "":
             # 没有上一步，直接返回
             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
+        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:
             # 上一步执行完成
             current_plan = await self.memory.get("current_plan")
-            current_step['evaluation']['consumed_time'] = (time_now - step_start_time)/60
-            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_step["evaluation"]["consumed_time"] = (
+                time_now - step_start_time
+            ) / 60
+            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.update("current_plan", current_plan)
-            if current_step_index is not None and current_step_index + 1 < len(current_plan["steps"]):
+            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.update("current_step", next_step)
             else:
                 # 标记计划完成
                 current_plan["completed"] = True
-                current_plan["end_time"] = await self.simulator.get_time(format_time=True)
+                current_plan["end_time"] = await self.simulator.get_time(
+                    format_time=True
+                )
                 await self.memory.update("current_plan", current_plan)
                 await self.memory.update("current_step", {"intention": "", "type": ""})
                 logger.info("Current plan execution completed.\n")
@@ -80,7 +114,7 @@ class PlanAndActionBlock(Block):
         # 检测上一步是否执行完成
         if not await self.check_and_update_step():
             return
-        
+
         # 长期决策
         await self.longTermDecisionBlock.forward()
 
@@ -102,58 +136,92 @@ class PlanAndActionBlock(Block):
         if current_step and current_step.get("type") and current_step.get("intention"):
             step_type = current_step.get("type")
             position = await self.memory.get("position")
-            if 'aoi_position' in position:
-                current_step['position'] = position['aoi_position']['aoi_id']
-            current_step['start_time'] = int(await self.simulator.get_time())
-            logger.info(f"Executing step: {current_step['intention']} - Type: {step_type}")
+            if "aoi_position" in position:
+                current_step["position"] = position["aoi_position"]["aoi_id"]
+            current_step["start_time"] = int(await self.simulator.get_time())
+            logger.info(
+                f"Executing step: {current_step['intention']} - Type: {step_type}"
+            )
             result = None
             if step_type == "mobility":
-                result = await self.mobilityBlock.forward(current_step, execution_context)
+                result = await self.mobilityBlock.forward(
+                    current_step, execution_context
+                )
             elif step_type == "social":
                 result = await self.socialBlock.forward(current_step, execution_context)
             elif step_type == "economy":
-                result = await self.economyBlock.forward(current_step, execution_context)
+                result = await self.economyBlock.forward(
+                    current_step, execution_context
+                )
             elif step_type == "other":
                 result = await self.otherBlock.forward(current_step, execution_context)
             if result != None:
                 logger.info(f"Execution result: {result}")
-                current_step['evaluation'] = result
-            
+                current_step["evaluation"] = result
+
             # 更新current_step信息，plan信息以及execution_context信息
-            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_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.update("current_step", current_step)
             await self.memory.update("current_plan", current_plan)
             await self.memory.update("execution_context", execution_context)
 
+
 class MindBlock(Block):
     """认知工作流"""
+
     cognitionBlock: CognitionBlock
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__(name="mind_block", llm=llm, memory=memory, simulator=simulator)
-        self.cognitionBlock = CognitionBlock(llm=llm, memory=memory, simulator=simulator)
+        self.cognitionBlock = CognitionBlock(
+            llm=llm, memory=memory, simulator=simulator
+        )
 
     async def forward(self):
         await self.cognitionBlock.forward()
 
+
 class SocietyAgent(CitizenAgent):
     mindBlock: MindBlock
     planAndActionBlock: PlanAndActionBlock
     update_with_sim = UpdateWithSimulator()
 
-    def __init__(self, 
-                 name: str,
-                 llm_client: Optional[LLM] = None,
-                 simulator: Optional[Simulator] = None,
-                 memory: Optional[Memory] = None,
-                 economy_client: Optional[EconomyClient] = None,
-                 ) -> None:
-        super().__init__(name=name, llm_client=llm_client, simulator=simulator, memory=memory, economy_client=economy_client)
-        self.mindBlock = MindBlock(llm=self._llm_client, memory=self._memory, simulator=self._simulator)
-        self.planAndActionBlock = PlanAndActionBlock(agent=self, llm=self._llm_client, memory=self._memory, simulator=self._simulator, economy_client=self._economy_client)
-        self.step_count = -1        
-    
+    def __init__(
+        self,
+        name: str,
+        llm_client: Optional[LLM] = None,
+        simulator: Optional[Simulator] = None,
+        memory: Optional[Memory] = None,
+        economy_client: Optional[EconomyClient] = None,
+    ) -> None:
+        super().__init__(
+            name=name,
+            llm_client=llm_client,
+            simulator=simulator,
+            memory=memory,
+            economy_client=economy_client,
+        )
+        self.mindBlock = MindBlock(
+            llm=self._llm_client, memory=self._memory, simulator=self._simulator
+        )
+        self.planAndActionBlock = PlanAndActionBlock(
+            agent=self,
+            llm=self._llm_client,
+            memory=self._memory,
+            simulator=self._simulator,
+            economy_client=self._economy_client,
+        )
+        self.step_count = -1
+
     # Main workflow
     async def forward(self):
         logger.info(f"Agent {self._uuid} forward")
@@ -166,7 +234,7 @@ class SocietyAgent(CitizenAgent):
         await asyncio.gather(*task_list)
 
     async def process_agent_chat_response(self, payload: dict) -> str:
-        if payload['type'] == 'social':
+        if payload["type"] == "social":
             resp = f"Agent {self._uuid} received agent chat response: {payload}"
             logger.info(resp)
             try:
@@ -174,9 +242,9 @@ class SocietyAgent(CitizenAgent):
                 sender_id = payload.get("from")
                 if not sender_id:
                     return ""
-                    
+
                 raw_content = payload.get("content", "")
-                
+
                 # Parse message content
                 try:
                     message_data = json.loads(raw_content)
@@ -185,32 +253,34 @@ class SocietyAgent(CitizenAgent):
                 except (json.JSONDecodeError, TypeError, KeyError):
                     content = raw_content
                     propagation_count = 1
-                
+
                 if not content:
                     return ""
-                    
+
                 # Get chat histories and ensure proper format
                 chat_histories = await self._memory.get("chat_histories") or {}
                 if not isinstance(chat_histories, dict):
                     chat_histories = {}
-                
+
                 # Update chat history with received message
                 if sender_id not in chat_histories:
                     chat_histories[sender_id] = ""
                 if chat_histories[sender_id]:
                     chat_histories[sender_id] += "，"
                 chat_histories[sender_id] += f"them: {content}"
-                
+
                 # Check propagation limit
                 if propagation_count > 5:
                     await self._memory.update("chat_histories", chat_histories)
-                    logger.info(f"Message propagation limit reached ({propagation_count} > 5), stopping propagation")
+                    logger.info(
+                        f"Message propagation limit reached ({propagation_count} > 5), stopping propagation"
+                    )
                     return ""
-                
+
                 # Get relationship score
                 relationships = await self._memory.get("relationships") or {}
                 relationship_score = relationships.get(sender_id, 50)
-                
+
                 # Decision prompt
                 should_respond_prompt = f"""Based on:
         - Received message: "{content}"
@@ -230,10 +300,15 @@ class SocietyAgent(CitizenAgent):
 
         Answer only YES or NO."""
 
-                should_respond = await self._llm_client.atext_request([
-                    {"role": "system", "content": "You are helping decide whether to respond to a message."},
-                    {"role": "user", "content": should_respond_prompt}
-                ])
+                should_respond = await self._llm_client.atext_request(
+                    [
+                        {
+                            "role": "system",
+                            "content": "You are helping decide whether to respond to a message.",
+                        },
+                        {"role": "user", "content": should_respond_prompt},
+                    ]
+                )
 
                 if should_respond.strip().upper() != "YES":
                     await self._memory.update("chat_histories", chat_histories)
@@ -258,34 +333,42 @@ class SocietyAgent(CitizenAgent):
 
         Response should be ONLY the message text, no explanations."""
 
-                response = await self._llm_client.atext_request([
-                    {"role": "system", "content": "You are helping generate a chat response."},
-                    {"role": "user", "content": response_prompt}
-                ])
+                response = await self._llm_client.atext_request(
+                    [
+                        {
+                            "role": "system",
+                            "content": "You are helping generate a chat response.",
+                        },
+                        {"role": "user", "content": response_prompt},
+                    ]
+                )
 
                 if response:
                     # Update chat history with response
                     chat_histories[sender_id] += f"，me: {response}"
                     await self._memory.update("chat_histories", chat_histories)
-                    
+
                     # Send response
-                    serialized_response = json.dumps({
-                        "content": response,
-                        "propagation_count": propagation_count + 1
-                    }, ensure_ascii=False)
+                    serialized_response = json.dumps(
+                        {
+                            "content": response,
+                            "propagation_count": propagation_count + 1,
+                        },
+                        ensure_ascii=False,
+                    )
                     await self.send_message_to_agent(sender_id, serialized_response)
-                    logger.info('sender_id',sender_id)
-                    logger.info('message',serialized_response)
+                    logger.info("sender_id", sender_id)
+                    logger.info("message", serialized_response)
                 return response
 
             except Exception as e:
                 logger.warning(f"Error in process_agent_chat_response: {str(e)}")
                 return ""
         else:
-            content = payload['content']
+            content = payload["content"]
             key, value = content.split("@")
-            if '.' in value:
+            if "." in value:
                 value = float(value)
             else:
                 value = int(value)
-            await self.memory.update(key, value)  
+            await self.memory.update(key, value)

pycityagent/simulation/agentgroup.py

@@ -48,7 +48,7 @@ class AgentGroup:
         mlflow_run_id: str,
         embedding_model: Embeddings,
         logging_level: int,
-        agent_config_file: Union[str, list[str]] = None,
+        agent_config_file: Optional[Union[str, list[str]]] = None,
     ):
         logger.setLevel(logging_level)
         self._uuid = str(uuid.uuid4())
@@ -163,7 +163,7 @@ class AgentGroup:
                     agent.memory.set_faiss_query(self.faiss_query)
                 if self.embedding_model is not None:
                     agent.memory.set_embedding_model(self.embedding_model)
-                if self.agent_config_file[i]:
+                if self.agent_config_file is not None and self.agent_config_file[i]:
                     agent.load_from_file(self.agent_config_file[i])
                 self.agents.append(agent)
                 self.id2agent[agent._uuid] = agent