pycityagent 2.0.0a73__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a75__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/agent/agent.py

@@ -8,6 +8,7 @@ from typing import Any, Optional
 
 from mosstool.util.format_converter import dict2pb
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
+import ray
 
 from ..economy import EconomyClient
 from ..environment import Simulator
@@ -120,7 +121,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(dict_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")
@@ -280,7 +281,7 @@ class InstitutionAgent(Agent):
             _id = random.randint(100000, 999999)
             self._agent_id = _id
             self.status.set_agent_id(_id)
-            map_header = self.simulator.map.header
+            map_header = ray.get(self.simulator.map.get_map_header.remote())
             # TODO: remove random position assignment
             await self.status.update(
                 "position",

pycityagent/agent/agent_base.py

@@ -19,7 +19,7 @@ from ..environment import Simulator
 from ..environment.sim.person_service import PersonService
 from ..llm import LLM
 from ..memory import Memory
-from ..message import MessageInterceptor, Messager
+from ..message import Messager
 from ..utils import DIALOG_SCHEMA, SURVEY_SCHEMA, process_survey_for_llm
 from ..workflow import Block
 
@@ -673,6 +673,47 @@ class Agent(ABC):
         )
         print(f"Sent payback message to {self._exp_id}")
 
+    async def save_agent_thought(self, thought: str):
+        """
+        Save the agent's thought to the memory.
+
+        - **Args**:
+            - `thought` (`str`): The thought data to be saved.
+
+        - **Description**:
+            - Saves the thought data to the memory.
+        """
+        _date_time = datetime.now(timezone.utc)
+        _thought_dict = {
+            "id": self._uuid,
+            "day": await self.simulator.get_simulator_day(),
+            "t": await self.simulator.get_simulator_second_from_start_of_day(),
+            "type": 0,
+            "speaker": "",
+            "content": thought,
+            "created_at": int(_date_time.timestamp() * 1000),
+        }
+        auros = [_thought_dict]
+        pg_list = [(_thought_dict, _date_time)]
+        # Avro
+        if self._avro_file is not None:
+            with open(self._avro_file["dialog"], "a+b") as f:
+                fastavro.writer(f, DIALOG_SCHEMA, auros, codec="snappy")
+        # Pg
+        if self._pgsql_writer is not None:
+            if self._last_asyncio_pg_task is not None:
+                await self._last_asyncio_pg_task
+            _keys = ["id", "day", "t", "type", "speaker", "content", "created_at"]
+            _data = [
+                tuple([_dict[k] if k != "created_at" else _date_time for k in _keys])
+                for _dict, _date_time in pg_list
+            ]
+            self._last_asyncio_pg_task = (
+                self._pgsql_writer.async_write_dialog.remote(  # type:ignore
+                    _data
+                )
+            )
+
     async def process_agent_chat_response(self, payload: dict) -> str:
         """
         Log the reception of an agent chat response.
@@ -931,5 +972,4 @@ class Agent(ABC):
         if self._messager is not None:
             await self._messager.ping.remote()  # type:ignore
         if not self._blocked:
-            async with self.llm.semaphore:
-                await self.forward()
+            return await self.forward()

pycityagent/cityagent/blocks/cognition_block.py

@@ -199,7 +199,8 @@ class CognitionBlock(Block):
         thought = str(response["thought"])
         await self.memory.status.update("thought", thought)
         await self.memory.stream.add_cognition(description=thought)
-        return
+
+        return thought
 
     async def end_of_day(self):
         """Cognition - end of day workflow"""

pycityagent/cityagent/blocks/economy_block.py

@@ -36,7 +36,7 @@ class WorkBlock(Block):
     """WorkPlace Block"""
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("WorkBlock", llm=llm, memory=memory, simulator=simulator)
-        self.description = "Do work"
+        self.description = "Do work related tasks"
         self.guidance_prompt = FormatPrompt(template=TIME_ESTIMATE_PROMPT)
     async def forward(self, step, context):
         self.guidance_prompt.format(
@@ -157,7 +157,7 @@ class EconomyBlock(Block):
 
     async def forward(self, step, context):        
         self.trigger_time += 1
-        selected_block = self.consumption_block
+        selected_block = await self.dispatcher.dispatch(step)
         result = await selected_block.forward(step, context) # type: ignore  
         return result
     
@@ -168,7 +168,7 @@ class MonthPlanBlock(Block):
         self.economy_client = economy_client
         self.llm_error = 0
         self.last_time_trigger = None
-        self.time_diff = month_days * 24 * 60 * 60
+        self.time_diff = 30 * 24 * 60 * 60
         self.forward_times = 0
         
     async def month_trigger(self):

pycityagent/cityagent/blocks/mobility_block.py

@@ -1,6 +1,8 @@
 import math
 from typing import List
 
+import ray
+
 from pycityagent.environment.simulator import Simulator
 from pycityagent.llm import LLM
 from pycityagent.memory import Memory
@@ -50,7 +52,6 @@ Return only the integer number without any additional text or explanation."""
 
 
 def gravity_model(pois):
-    N = len(pois)
     pois_Dis = {
         "1k": [],
         "2k": [],
@@ -133,21 +134,22 @@ class PlaceSelectionBlock(Block):
         )
         self.radiusPrompt = FormatPrompt(RADIUS_PROMPT)
         # configurable fields
-        self.search_limit = 10000
+        self.search_limit = 100
 
     async def forward(self, step, context):
+        poi_cate = self.simulator.get_poi_cate()
         self.typeSelectionPrompt.format(
             plan=context["plan"],
             intention=step["intention"],
-            poi_category=list(self.simulator.poi_cate.keys()),
+            poi_category=list(poi_cate.keys()),
         )
         levelOneType = await self.llm.atext_request(self.typeSelectionPrompt.to_dialog())  # type: ignore
         try:
-            sub_category = self.simulator.poi_cate[levelOneType]
+            sub_category = poi_cate[levelOneType]
         except Exception as e:
             logger.warning(f"Wrong type of poi, raw response: {levelOneType}")
-            levelOneType = random.choice(list(self.simulator.poi_cate.keys()))
-            sub_category = self.simulator.poi_cate[levelOneType]
+            levelOneType = random.choice(list(poi_cate.keys()))
+            sub_category = poi_cate[levelOneType]
         self.secondTypeSelectionPrompt.format(
             plan=context["plan"], intention=step["intention"], poi_category=sub_category
         )
@@ -162,21 +164,21 @@ class PlaceSelectionBlock(Block):
         )
         radius = int(await self.llm.atext_request(self.radiusPrompt.to_dialog()))  # type: ignore
         try:
-            pois = self.simulator.map.query_pois(
+            pois = ray.get(self.simulator.map.query_pois.remote(
                 center=center,
                 category_prefix=levelTwoType,
                 radius=radius,
                 limit=self.search_limit,
-            )
+            ))
         except Exception as e:
             logger.warning(f"Error querying pois: {e}")
             levelTwoType = random.choice(sub_category)
-            pois = self.simulator.map.query_pois(
+            pois = ray.get(self.simulator.map.query_pois.remote(
                 center=center,
                 category_prefix=levelTwoType,
                 radius=radius,
                 limit=self.search_limit,
-            )
+            ))
         if len(pois) > 0:
             pois = gravity_model(pois)
             probabilities = [item[2] for item in pois]
@@ -198,8 +200,8 @@ class PlaceSelectionBlock(Block):
                 "node_id": node_id,
             }
         else:
-            simmap = self.simulator.map
-            poi = random.choice(list(simmap.pois.values()))
+            pois = ray.get(self.simulator.map.get_poi.remote())
+            poi = random.choice(pois)
             nextPlace = (poi["name"], poi["aoi_id"])
             # save the destination to context
             context["next_place"] = nextPlace
@@ -309,11 +311,12 @@ class MoveBlock(Block):
                 )
             else:
                 while True:
-                    r_aoi = random.choice(list(self.simulator.map.aois.values()))
+                    aois = ray.get(self.simulator.map.get_aoi.remote())
+                    r_aoi = random.choice(aois)
                     if len(r_aoi["poi_ids"]) > 0:
                         r_poi = random.choice(r_aoi["poi_ids"])
                         break
-                poi = self.simulator.map.pois[r_poi]
+                poi = ray.get(self.simulator.map.get_poi.remote(r_poi))
                 next_place = (poi["name"], poi["aoi_id"])
                 await self.simulator.set_aoi_schedules(
                     person_id=agent_id,
@@ -372,9 +375,6 @@ class MobilityBlock(Block):
 
     async def forward(self, step, context):
         self.trigger_time += 1
-        consumption_start = (
-            self.llm.prompt_tokens_used + self.llm.completion_tokens_used
-        )
 
         # Select the appropriate sub-block using dispatcher
         selected_block = await self.dispatcher.dispatch(step)
@@ -382,7 +382,4 @@ class MobilityBlock(Block):
         # Execute the selected sub-block and get the result
         result = await selected_block.forward(step, context)  # type: ignore
 
-        consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
-        self.token_consumption += consumption_end - consumption_start
-
         return result

pycityagent/cityagent/blocks/needs_block.py

@@ -1,6 +1,8 @@
 import json
 import logging
 
+import ray
+
 from pycityagent import Simulator
 from pycityagent.llm import LLM
 from pycityagent.memory import Memory

pycityagent/cityagent/blocks/plan_block.py

@@ -3,6 +3,8 @@ import logging
 import random
 from typing import Dict, List
 
+import ray
+
 from pycityagent.environment.simulator import Simulator
 from pycityagent.llm import LLM
 from pycityagent.memory import Memory
@@ -196,7 +198,6 @@ class PlanBlock(Block):
         ):
             current_location = "At workplace"
         current_time = await self.simulator.get_time(format_time=True)
-        environment = await self.memory.status.get("environment")
         options = self.guidance_options.get(current_need, [])
         self.guidance_prompt.format(
             weather=self.simulator.sence("weather"),
@@ -205,7 +206,6 @@ class PlanBlock(Block):
             options=options,
             current_location=current_location,
             current_time=current_time,
-            environment=environment,
             emotion_types=await self.memory.status.get("emotion_types"),
             thought=await self.memory.status.get("thought"),
         )
@@ -240,14 +240,12 @@ class PlanBlock(Block):
         ):
             current_location = "At workplace"
         current_time = await self.simulator.get_time(format_time=True)
-        environment = await self.memory.status.get("environment")
         self.detail_prompt.format(
             weather=self.simulator.sence("weather"),
             temperature=self.simulator.sence("temperature"),
             selected_option=selected_option,
             current_location=current_location,
             current_time=current_time,
-            environment=environment,
             emotion_types=await self.memory.status.get("emotion_types"),
             thought=await self.memory.status.get("thought"),
             max_plan_steps=self.max_plan_steps,

pycityagent/cityagent/blocks/utils.py

@@ -24,7 +24,6 @@ Please return the result in JSON format (Do not return any other text):
 """
 
 num_labor_hours = 168
-month_days = 0.03
 productivity_per_labor = 1
 max_price_inflation = 0.1
 max_wage_inflation = 0.05

pycityagent/cityagent/initial.py

@@ -68,8 +68,6 @@ async def initialize_social_network(simulation):
             await simulation.update(
                 agent_id, "interactions", {friend_id: [] for friend_id in friends}
             )
-
-        print("Social network initialization completed!")
         return True
 
     except Exception as e:

pycityagent/cityagent/memory_config.py

@@ -28,16 +28,16 @@ def memory_config_societyagent():
         "type": (str, "citizen"),
         "city": (str, "New York", True),
         # Needs Model
-        "hunger_satisfaction": (float, random.random(), True),  # hunger satisfaction
-        "energy_satisfaction": (float, random.random(), True),  # energy satisfaction  
-        "safety_satisfaction": (float, random.random(), True),  # safety satisfaction
-        "social_satisfaction": (float, random.random(), True),  # social satisfaction
-        "current_need": (str, "none", True),
+        "hunger_satisfaction": (float, random.random(), False),  # hunger satisfaction
+        "energy_satisfaction": (float, random.random(), False),  # energy satisfaction  
+        "safety_satisfaction": (float, random.random(), False),  # safety satisfaction
+        "social_satisfaction": (float, random.random(), False),  # social satisfaction
+        "current_need": (str, "none", False),
         # Plan Behavior Model
-        "current_plan": (list, [], True),
-        "current_step": (dict, {"intention": "", "type": ""}, True),
-        "execution_context": (dict, {}, True),
-        "plan_history": (list, [], True),
+        "current_plan": (list, [], False),
+        "current_step": (dict, {"intention": "", "type": ""}, False),
+        "execution_context": (dict, {}, False),
+        "plan_history": (list, [], False),
         # cognition
         "emotion": (
             dict,
@@ -60,40 +60,39 @@ def memory_config_societyagent():
             random.choice(agent_skills),
             True,
         ),  # work skill
-        "tax_paid": (float, 0.0, True),  # tax paid
-        "consumption_currency": (float, 0.0, True),  # consumption
-        "goods_demand": (int, 0, True),
-        "goods_consumption": (int, 0, True),
-        "work_propensity": (float, 0.0, True),
-        "consumption_propensity": (float, 0.0, True),
-        "income_currency": (float, 0.0, True),  # monthly income
-        "to_income": (float, 0.0, True),
-        "to_consumption_currency": (float, 0.0, True),
-        "firm_id": (int, 0, True),
-        "government_id": (int, 0, True),
-        "bank_id": (int, 0, True),
-        "nbs_id": (int, 0, True),
-        "dialog_queue": (deque(maxlen=3), [], True),
-        "firm_forward": (int, 0, True),
-        "bank_forward": (int, 0, True),
-        "nbs_forward": (int, 0, True),
-        "government_forward": (int, 0, True),
-        "forward": (int, 0, True),
-        "depression": (float, 0.0, True),
-        "ubi_opinion": (list, [], True),
-        "working_experience": (list, [], True),
-        "work_hour_month": (float, 160, True),
-        "work_hour_finish": (float, 0, True),
+        "tax_paid": (float, 0.0, False),  # tax paid
+        "consumption_currency": (float, 0.0, False),  # consumption
+        "goods_demand": (int, 0, False),
+        "goods_consumption": (int, 0, False),
+        "work_propensity": (float, 0.0, False),
+        "consumption_propensity": (float, 0.0, False),
+        "income_currency": (float, 0.0, False),  # monthly income
+        "to_income": (float, 0.0, False),
+        "to_consumption_currency": (float, 0.0, False),
+        "firm_id": (int, 0, False),
+        "government_id": (int, 0, False),
+        "bank_id": (int, 0, False),
+        "nbs_id": (int, 0, False),
+        "dialog_queue": (deque(maxlen=3), [], False),
+        "firm_forward": (int, 0, False),
+        "bank_forward": (int, 0, False),
+        "nbs_forward": (int, 0, False),
+        "government_forward": (int, 0, False),
+        "forward": (int, 0, False),
+        "depression": (float, 0.0, False),
+        "ubi_opinion": (list, [], False),
+        "working_experience": (list, [], False),
+        "work_hour_month": (float, 160, False),
+        "work_hour_finish": (float, 0, False),
         # social
-        "friends": (list, [], True),  # friends list
-        "relationships": (dict, {}, True),  # relationship strength with each friend
-        "relation_types": (dict, {}, True),
-        "chat_histories": (dict, {}, True),  # all chat histories
-        "interactions": (dict, {}, True),  # all interaction records
-        "to_discuss": (dict, {}, True),
+        "friends": (list, [], False),  # friends list
+        "relationships": (dict, {}, False),  # relationship strength with each friend
+        "relation_types": (dict, {}, False),
+        "chat_histories": (dict, {}, False),  # all chat histories
+        "interactions": (dict, {}, False),  # all interaction records
+        "to_discuss": (dict, {}, False),
         # mobility
-        "environment": (str, "The environment outside is good", True),
-        "number_poi_visited": (int, 1, True),
+        "number_poi_visited": (int, 1, False),
     }
 
     PROFILE = {

pycityagent/cityagent/societyagent.py

@@ -1,6 +1,7 @@
 import asyncio
 import json
 import logging
+import time
 from typing import Optional
 
 from pycityagent import CitizenAgent, Simulator
@@ -66,8 +67,7 @@ class PlanAndActionBlock(Block):
     async def plan_generation(self):
         """Generate plan"""
         current_plan = await self.memory.status.get("current_plan")
-        current_need = await self.memory.status.get("current_need")
-        if current_need != "none" and not current_plan:
+        if current_plan is None:
             await self.planBlock.forward()
 
     async def step_execution(self):
@@ -152,7 +152,6 @@ class PlanAndActionBlock(Block):
         # step execution
         await self.step_execution()
 
-
 class MindBlock(Block):
     """Cognition workflow"""
 
@@ -233,6 +232,7 @@ class SocietyAgent(CitizenAgent):
 
     # Main workflow
     async def forward(self):
+        start_time = time.time()
         self.step_count += 1
         # sync agent status with simulator
         await self.update_with_sim()
@@ -242,10 +242,13 @@ class SocietyAgent(CitizenAgent):
         if not ifpass:
             return
 
+        # plan and action
         await self.planAndActionBlock.forward()
 
+        # cognition
         if self.enable_cognition:
             await self.mindBlock.forward()
+        return time.time() - start_time
 
     async def check_and_update_step(self):
         """Check if the previous step has been completed"""
@@ -257,7 +260,7 @@ class SocietyAgent(CitizenAgent):
 
         # Get the previous step information
         current_step = await self.memory.status.get("current_step")
-        if current_step["intention"] == "" or current_step["type"] == "":
+        if current_step["intention"] == "":
             # No previous step, return directly
             return True
         time_now = int(await self.simulator.get_time())
@@ -304,9 +307,7 @@ class SocietyAgent(CitizenAgent):
                         conclusion = await self.mindBlock.cognitionBlock.emotion_update(
                             incident
                         )
-                        await self.memory.stream.add_cognition(
-                            description=conclusion  # type:ignore
-                        )
+                        await self.save_agent_thought(conclusion)
                         await self.memory.stream.add_cognition_to_memory(
                             current_plan["stream_nodes"], conclusion  # type:ignore
                         )
@@ -331,9 +332,7 @@ class SocietyAgent(CitizenAgent):
                     conclusion = await self.mindBlock.cognitionBlock.emotion_update(
                         incident
                     )
-                    await self.memory.stream.add_cognition(
-                        description=conclusion  # type:ignore
-                    )
+                    await self.save_agent_thought(conclusion)
                     await self.memory.stream.add_cognition_to_memory(
                         current_plan["stream_nodes"], conclusion  # type:ignore
                     )

pycityagent/environment/sim/sim_env.py

@@ -13,7 +13,7 @@ from ..utils import encode_to_base64, find_free_port
 __all__ = ["ControlSimEnv"]
 
 
-def _generate_yaml_config(map_file: str, start_step: int, total_step: int) -> str:
+def _generate_yaml_config(map_file: str, max_day: int, start_step: int, total_step: int) -> str:
     map_file = os.path.abspath(map_file)
     return f"""
 input:
@@ -22,6 +22,7 @@ input:
     file: "{map_file}"
 
 control:
+  day: {max_day}
   step:
     # 模拟器起始步
     start: {start_step}
@@ -46,6 +47,7 @@ class ControlSimEnv:
         self,
         task_name: str,
         map_file: str,
+        max_day: int,
         start_step: int,
         total_step: int,
         log_dir: str,
@@ -63,6 +65,7 @@ class ControlSimEnv:
         """
         self._task_name = task_name
         self._map_file = map_file
+        self._max_day = max_day
         self._start_step = start_step
         self._total_step = total_step
         self._log_dir = log_dir
@@ -70,7 +73,7 @@ class ControlSimEnv:
         self._timeout = timeout
         self._max_procs = max_process
 
-        self._sim_config = _generate_yaml_config(map_file, start_step, total_step)
+        self._sim_config = _generate_yaml_config(map_file, max_day, start_step, total_step)
         # sim
         self.sim_port = None
         self._sim_proc = None
@@ -113,6 +116,7 @@ class ControlSimEnv:
                     f":{self.sim_port}",
                     "-run.min_step_time",
                     f"{self._min_step_time}",
+                    "-run.pause_after_one_day",
                     "-output",
                     self._log_dir,
                     "-cache",