pycityagent 2.0.0a65__cp311-cp311-macosx_11_0_arm64.whl → 2.0.0a66__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/cityagent/blocks/mobility_block.py

@@ -1,9 +1,12 @@
+import math
 from typing import List
 from .dispatcher import BlockDispatcher
 from pycityagent.environment.simulator import Simulator
 from pycityagent.llm import LLM
 from pycityagent.memory import Memory
 from pycityagent.workflow.block import Block
+import numpy as np
+from operator import itemgetter
 import random
 import logging
 logger = logging.getLogger("pycityagent")
@@ -38,18 +41,83 @@ Current temperature: {temperature}
 Your current emotion: {emotion_types}
 Your current thought: {thought}
 
-Please analyze how these emotions would affect travel willingness and return only a single integer number between 1000-100000 representing the maximum travel radius in meters. A more positive emotional state generally leads to greater willingness to travel further.
+Please analyze how these emotions would affect travel willingness and return only a single integer number between 3000-100000 representing the maximum travel radius in meters. A more positive emotional state generally leads to greater willingness to travel further.
 
 Return only the integer number without any additional text or explanation."""
 
+def gravity_model(pois):
+    N = len(pois)
+    pois_Dis = {
+        "1k": [],
+        "2k": [],
+        "3k": [],
+        "4k": [],
+        "5k": [],
+        "6k": [],
+        "7k": [],
+        "8k": [],
+        "9k": [],
+        "10k": [],
+        "more": [],
+    }
+    for poi in pois:
+        iflt10k = True
+        for d in range(1, 11):
+            if (d - 1) * 1000 <= poi[1] < d * 1000:
+                pois_Dis["{}k".format(d)].append(poi)
+                iflt10k = False
+                break
+        if iflt10k:
+            pois_Dis["more"].append(poi)
+
+    res = []
+    distanceProb = []
+    for poi in pois:
+        iflt10k = True
+        for d in range(1, 11):
+            if (d - 1) * 1000 <= poi[1] < d * 1000:
+                n = len(pois_Dis["{}k".format(d)])
+                S = math.pi * ((d * 1000) ** 2 - ((d - 1) * 1000) ** 2)
+                density = n / S
+                distance = poi[1]
+                distance = distance if distance > 1 else 1
+
+                # distance decay coefficient, use the square of distance to calculate, so that distant places are less likely to be selected
+                weight = density / (distance**2)  # the original weight is reasonable
+                res.append((poi[0]["name"], poi[0]["id"], weight, distance))
+                distanceProb.append(
+                    1 / (math.sqrt(distance))
+                )
+                iflt10k = False
+                break
+
+    distanceProb = np.array(distanceProb)
+    distanceProb = distanceProb / np.sum(distanceProb)
+    distanceProb = list(distanceProb)
+
+    options = list(range(len(res)))
+    sample = list(
+        np.random.choice(options, size=50, p=distanceProb)
+    )  # sample based on the probability value
+
+    get_elements = itemgetter(*sample)
+    random_elements = get_elements(res)
+
+    # normalize the weight to become the true probability value
+    weightSum = sum(item[2] for item in random_elements)
+    final = [
+        (item[0], item[1], item[2] / weightSum, item[3]) for item in random_elements
+    ]
+    return final
+
 class PlaceSelectionBlock(Block):
     """
-    选择目的地
+    Select destination
     PlaceSelectionBlock
     """
     configurable_fields: List[str] = ["search_limit"]
     default_values = {
-        "search_limit": 10
+        "search_limit": 10000
     }
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
@@ -59,7 +127,7 @@ class PlaceSelectionBlock(Block):
         self.secondTypeSelectionPrompt = FormatPrompt(PLACE_SECOND_TYPE_SELECTION_PROMPT)
         self.radiusPrompt = FormatPrompt(RADIUS_PROMPT)
         # configurable fields
-        self.search_limit = 10
+        self.search_limit = 10000
 
     async def forward(self, step, context):
         self.typeSelectionPrompt.format(
@@ -106,9 +174,15 @@ class PlaceSelectionBlock(Block):
                 limit=self.search_limit
             )
         if len(pois) > 0:
-            poi = random.choice(pois)[0]
-            nextPlace = (poi['name'], poi['aoi_id'])
-            # 将地点信息保存到context中
+            pois = gravity_model(pois)
+            probabilities = [item[2] for item in pois]
+            options = list(range(len(pois)))
+            sample = np.random.choice(
+                options, size=1, p=probabilities
+            )  # sample based on the probability value
+            nextPlace = pois[sample[0]]
+            nextPlace = (nextPlace[0], nextPlace[1])
+            # save the destination to context
             context['next_place'] = nextPlace
             node_id = await self.memory.stream.add_mobility(description=f"For {step['intention']}, I selected the destination: {nextPlace}")
             return {
@@ -121,7 +195,7 @@ class PlaceSelectionBlock(Block):
             simmap = self.simulator.map
             poi = random.choice(list(simmap.pois.values()))
             nextPlace = (poi['name'], poi['aoi_id'])
-            # 将地点信息保存到context中
+            # save the destination to context
             context['next_place'] = nextPlace
             node_id = await self.memory.stream.add_mobility(description=f"For {step['intention']}, I selected the destination: {nextPlace}")
             return {
@@ -133,7 +207,7 @@ class PlaceSelectionBlock(Block):
 
 class MoveBlock(Block):
     """
-    移动操作
+    Execute mobility operations
     MoveBlock
     """
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):

pycityagent/cityagent/blocks/needs_block.py

@@ -17,9 +17,6 @@ Profile Information:
 - Age: {age}
 - Monthly Income: {income}
 
-Current Emotion: {emotion_types}
-Current Thought: {thought}
-
 Current Time: {now_time}
 
 Please initialize the agent's satisfaction levels and parameters based on the profile above. Return the values in JSON format with the following structure:
@@ -77,9 +74,6 @@ Current satisfaction:
 - safety_satisfaction: {safety_satisfaction}
 - social_satisfaction: {social_satisfaction}
 
-Current Emotion: {emotion_types}
-Current Thought: {thought}
-
 Please evaluate and adjust the value of {current_need} satisfaction based on the execution results above.
 
 Notes:
@@ -134,8 +128,6 @@ class NeedsBlock(Block):
                 occupation=await self.memory.status.get("occupation"),
                 age=await self.memory.status.get("age"),
                 income=await self.memory.status.get("income"),
-                emotion_types=await self.memory.status.get("emotion_types"),
-                thought=await self.memory.status.get("thought"),
                 now_time=await self.simulator.get_time(format_time=True)
             )
             response = await self.llm.atext_request(
@@ -285,8 +277,6 @@ class NeedsBlock(Block):
             energy_satisfaction=await self.memory.status.get("energy_satisfaction"),
             safety_satisfaction=await self.memory.status.get("safety_satisfaction"),
             social_satisfaction=await self.memory.status.get("social_satisfaction"),
-            emotion_types=await self.memory.status.get("emotion_types"),
-            thought=await self.memory.status.get("thought")
         )
 
         response = await self.llm.atext_request(

pycityagent/cityagent/blocks/social_block.py

@@ -210,16 +210,6 @@ class FindPersonBlock(Block):
 
 class MessageBlock(Block):
     """生成并发送消息"""
-    configurable_fields: List[str] = ["default_message_template", "to_discuss"]
-    default_values = {
-        "default_message_template": """
-        As a {gender} {occupation} with {education} education and {personality} personality,
-        generate a message for a friend (relationship strength: {relationship_score}/100)
-        about {intention}.
-        """,
-        "to_discuss": []
-    }
-    
     def __init__(self, agent, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("MessageBlock", llm=llm, memory=memory, simulator=simulator)
         self.agent = agent

pycityagent/cityagent/initial.py

@@ -73,22 +73,19 @@ async def bind_agent_info(simulation):
     infos = await simulation.gather("id")
     citizen_uuids = await simulation.filter(types=[SocietyAgent])
     firm_uuids = await simulation.filter(types=[FirmAgent])
-    locations = await simulation.gather("location", firm_uuids)
-    locations_plain = {}
-    for info in locations:
-        for k, v in info.items():
-            locations_plain[k] = v
     government_uuids = await simulation.filter(types=[GovernmentAgent])
     bank_uuids = await simulation.filter(types=[BankAgent])
     nbs_uuids = await simulation.filter(types=[NBSAgent])
     citizen_agent_ids = []
     firm_ids = []
+    id2uuid = {}
     for info in infos:
         for k, v in info.items():
             if k in citizen_uuids:
                 citizen_agent_ids.append(v)
             elif k in firm_uuids:
                 firm_ids.append(v)
+                id2uuid[v] = k
             elif k in government_uuids:
                 government_id = v
             elif k in bank_uuids:
@@ -97,9 +94,7 @@ async def bind_agent_info(simulation):
                 nbs_id = v
     for citizen_uuid in citizen_uuids:
         random_firm_id = random.choice(firm_ids)
-        location = locations_plain[random_firm_id]
         await simulation.update(citizen_uuid, "firm_id", random_firm_id)
-        await simulation.update(citizen_uuid, "work", location)
         await simulation.update(citizen_uuid, "government_id", government_id)
         await simulation.update(citizen_uuid, "bank_id", bank_id)
         await simulation.update(citizen_uuid, "nbs_id", nbs_id)

pycityagent/cityagent/memory_config.py

@@ -5,6 +5,7 @@ import numpy as np
 import pycityproto.city.economy.v2.economy_pb2 as economyv2
 from mosstool.map._map_util.const import AOI_START_ID
 
+from .firmagent import FirmAgent
 pareto_param = 8
 payment_max_skill_multiplier = 950
 payment_max_skill_multiplier = float(payment_max_skill_multiplier)
@@ -14,17 +15,24 @@ clipped_skills = np.minimum(pmsm, (pmsm - 1) * pareto_samples + 1)
 sorted_clipped_skills = np.sort(clipped_skills, axis=1)
 agent_skills = list(sorted_clipped_skills.mean(axis=0))
 
+work_locations = [AOI_START_ID + random.randint(1000, 10000) for _ in range(1000)]
+
+async def memory_config_init(simulation):
+    global work_locations
+    number_of_firm = simulation.agent_count[FirmAgent]
+    work_locations = [AOI_START_ID + random.randint(1000, 10000) for _ in range(number_of_firm)]
 
 def memory_config_societyagent():
+    global work_locations
     EXTRA_ATTRIBUTES = {
         "type": (str, "citizen"),
         "city": (str, "New York", True),
 
         # Needs Model
-        "hunger_satisfaction": (float, random.random(), True),  # 饥饿满意度
-        "energy_satisfaction": (float, random.random(), True),  # 精力满意度  
-        "safety_satisfaction": (float, random.random(), True),  # 安全满意度
-        "social_satisfaction": (float, random.random(), True),  # 社交满意度
+        "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),
 
         # Plan Behavior Model
@@ -186,7 +194,7 @@ def memory_config_societyagent():
             "aoi_position": {"aoi_id": AOI_START_ID + random.randint(1000, 10000)}
         },
         "work": {
-            "aoi_position": {"aoi_id": AOI_START_ID + random.randint(1000, 10000)}
+            "aoi_position": {"aoi_id": random.choice(work_locations)}
         },
     }
 
@@ -194,10 +202,11 @@ def memory_config_societyagent():
 
 
 def memory_config_firm():
+    global work_locations
     EXTRA_ATTRIBUTES = {
         "type": (int, economyv2.ORG_TYPE_FIRM),
         "location": {
-            "aoi_position": {"aoi_id": AOI_START_ID + random.randint(1000, 10000)}
+            "aoi_position": {"aoi_id": random.choice(work_locations)}
         },
         "price": (float, float(np.mean(agent_skills))),
         "inventory": (int, 0),

pycityagent/cityagent/societyagent.py

@@ -234,7 +234,8 @@ class SocietyAgent(CitizenAgent):
         await self.update_with_sim()
 
         # check last step
-        if not await self.check_and_update_step():
+        ifpass = await self.check_and_update_step()
+        if not ifpass:
             return
         
         await self.planAndActionBlock.forward()

pycityagent/llm/llm.py

@@ -24,7 +24,6 @@ from .utils import *
 
 os.environ["GRPC_VERBOSITY"] = "ERROR"
 
-
 class LLM:
     """
     大语言模型对象
@@ -46,6 +45,7 @@ class LLM:
             api_keys = [api_keys]
 
         self._aclients = []
+        self._client_usage = []
 
         for api_key in api_keys:
             if self.config.text["request_type"] == "openai":
@@ -69,6 +69,11 @@ class LLM:
                     f"Unsupported `request_type` {self.config.text['request_type']}!"
                 )
             self._aclients.append(client)
+            self._client_usage.append({
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "request_number": 0
+            })
 
     def set_semaphore(self, number_of_coroutine: int):
         self.semaphore = asyncio.Semaphore(number_of_coroutine)
@@ -81,45 +86,62 @@ class LLM:
         clear the storage of used tokens to start a new log message
         Only support OpenAI category API right now, including OpenAI, Deepseek
         """
-        self.prompt_tokens_used = 0
-        self.completion_tokens_used = 0
-        self.request_number = 0
+        for usage in self._client_usage:
+            usage["prompt_tokens"] = 0
+            usage["completion_tokens"] = 0 
+            usage["request_number"] = 0
+
+    def get_consumption(self):
+        consumption = {}
+        for i, usage in enumerate(self._client_usage):
+            consumption[f"api-key-{i+1}"] = {
+                "total_tokens": usage["prompt_tokens"] + usage["completion_tokens"],
+                "request_number": usage["request_number"]
+            }
+        return consumption
 
     def show_consumption(
         self, input_price: Optional[float] = None, output_price: Optional[float] = None
     ):
         """
-        if you give the input and output price of using model, this function will also calculate the consumption for you
+        Show consumption for each API key separately
         """
-        total_token = self.prompt_tokens_used + self.completion_tokens_used
-        if self.completion_tokens_used != 0:
-            rate = self.prompt_tokens_used / self.completion_tokens_used
-        else:
-            rate = "nan"
-        if self.request_number != 0:
-            TcA = total_token / self.request_number
-        else:
-            TcA = "nan"
-        out = f"""Request Number: {self.request_number}
-Token Usage:
-    - Total tokens: {total_token}
-    - Prompt tokens: {self.prompt_tokens_used}
-    - Completion tokens: {self.completion_tokens_used}
-    - Token per request: {TcA}
-    - Prompt:Completion ratio: {rate}:1"""
-        if input_price != None and output_price != None:
-            consumption = (
-                self.prompt_tokens_used / 1000000 * input_price
-                + self.completion_tokens_used / 1000000 * output_price
-            )
-            out += f"\n    - Cost Estimation: {consumption}"
-        print(out)
-        return {
-            "total": total_token,
-            "prompt": self.prompt_tokens_used,
-            "completion": self.completion_tokens_used,
-            "ratio": rate,
+        total_stats = {
+            "total": 0,
+            "prompt": 0,
+            "completion": 0,
+            "requests": 0
         }
+        
+        for i, usage in enumerate(self._client_usage):
+            prompt_tokens = usage["prompt_tokens"]
+            completion_tokens = usage["completion_tokens"]
+            requests = usage["request_number"]
+            total_tokens = prompt_tokens + completion_tokens
+            
+            total_stats["total"] += total_tokens
+            total_stats["prompt"] += prompt_tokens
+            total_stats["completion"] += completion_tokens
+            total_stats["requests"] += requests
+            
+            rate = prompt_tokens / completion_tokens if completion_tokens != 0 else "nan"
+            tokens_per_request = total_tokens / requests if requests != 0 else "nan"
+            
+            print(f"\nAPI Key #{i+1}:")
+            print(f"Request Number: {requests}")
+            print("Token Usage:")
+            print(f"    - Total tokens: {total_tokens}")
+            print(f"    - Prompt tokens: {prompt_tokens}")
+            print(f"    - Completion tokens: {completion_tokens}")
+            print(f"    - Token per request: {tokens_per_request}")
+            print(f"    - Prompt:Completion ratio: {rate}:1")
+            
+            if input_price is not None and output_price is not None:
+                consumption = (prompt_tokens / 1000000 * input_price + 
+                             completion_tokens / 1000000 * output_price)
+                print(f"    - Cost Estimation: {consumption}")
+        
+        return total_stats
 
     def _get_next_client(self):
         """获取下一个要使用的客户端"""
@@ -168,9 +190,9 @@ Token Usage:
                                 tools=tools,
                                 tool_choice=tool_choice,
                             )  # type: ignore
-                            self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
-                            self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
-                            self.request_number += 1
+                            self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
+                            self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
+                            self._client_usage[self._current_client_index]["request_number"] += 1
                             if tools and response.choices[0].message.tool_calls:
                                 return json.loads(
                                     response.choices[0]
@@ -193,9 +215,9 @@ Token Usage:
                             tools=tools,
                             tool_choice=tool_choice,
                         )  # type: ignore
-                        self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
-                        self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
-                        self.request_number += 1
+                        self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["request_number"] += 1
                         if tools and response.choices[0].message.tool_calls:
                             return json.loads(
                                 response.choices[0]
@@ -248,9 +270,9 @@ Token Usage:
                     if task_status != "SUCCESS":
                         raise Exception(f"Task failed with status: {task_status}")
 
-                    self.prompt_tokens_used += result_response.usage.prompt_tokens  # type: ignore
-                    self.completion_tokens_used += result_response.usage.completion_tokens  # type: ignore
-                    self.request_number += 1
+                    self._client_usage[self._current_client_index]["prompt_tokens"] += result_response.usage.prompt_tokens  # type: ignore
+                    self._client_usage[self._current_client_index]["completion_tokens"] += result_response.usage.completion_tokens  # type: ignore
+                    self._client_usage[self._current_client_index]["request_number"] += 1
                     if tools and result_response.choices[0].message.tool_calls:  # type: ignore
                         return json.loads(
                             result_response.choices[0]  # type: ignore

pycityagent/memory/memory.py

@@ -330,7 +330,7 @@ class StreamMemory:
         # 找到所有对应的记忆
         target_memories = []
         for memory in self._memories:
-            if id(memory) in memory_ids:
+            if memory.id in memory_ids:
                 target_memories.append(memory)
 
         if not target_memories:

pycityagent/simulation/agentgroup.py

@@ -186,7 +186,7 @@ class AgentGroup:
         self.message_dispatch_task.cancel()  # type: ignore
         await asyncio.gather(self.message_dispatch_task, return_exceptions=True)  # type: ignore
 
-    async def insert_agents(self):
+    async def insert_agent(self):
         bind_tasks = []
         for agent in self.agents:
             bind_tasks.append(agent.bind_to_simulator())  # type: ignore
@@ -200,7 +200,7 @@ class AgentGroup:
             if day == 0:
                 break
             await asyncio.sleep(1)
-        await self.insert_agents()
+        await self.insert_agent()
         self.id2agent = {agent._uuid: agent for agent in self.agents}
         logger.debug(f"-----Binding Agents to Messager in AgentGroup {self._uuid} ...")
         assert self.messager is not None
@@ -705,6 +705,9 @@ class AgentGroup:
                 )
             )
 
+    def get_llm_consumption(self):
+        return self.llm.get_consumption()
+
     async def step(self):
         try:
             tasks = [agent.run() for agent in self.agents]

pycityagent/simulation/simulation.py

@@ -14,9 +14,10 @@ from langchain_core.embeddings import Embeddings
 
 from ..agent import Agent, InstitutionAgent
 from ..cityagent import (BankAgent, FirmAgent, GovernmentAgent, NBSAgent,
-                         SocietyAgent, memory_config_bank, memory_config_firm,
+                         SocietyAgent)
+from ..cityagent.memory_config import (memory_config_bank, memory_config_firm,
                          memory_config_government, memory_config_nbs,
-                         memory_config_societyagent)
+                         memory_config_societyagent, memory_config_init)
 from ..cityagent.initial import bind_agent_info, initialize_social_network
 from ..cityagent.message_intercept import (EdgeMessageBlock,
                                            MessageBlockListener,
@@ -35,16 +36,15 @@ from .storage.pg import PgWriter, create_pg_tables
 
 logger = logging.getLogger("pycityagent")
 
-
 class AgentSimulation:
-    """城市智能体模拟器"""
+    """Agent Simulation"""
 
     def __init__(
         self,
         config: dict,
         agent_class: Union[None, type[Agent], list[type[Agent]]] = None,
         agent_config_file: Optional[dict] = None,
-        metric_extractor: Optional[list[tuple[int, Callable]]] = None,
+        metric_extractors: Optional[list[tuple[int, Callable]]] = None,
         enable_institution: bool = True,
         agent_prefix: str = "agent_",
         exp_name: str = "default_experiment",
@@ -52,10 +52,14 @@ class AgentSimulation:
     ):
         """
         Args:
-            agent_class: 智能体类
-            config: 配置
-            agent_prefix: 智能体名称前缀
-            exp_name: 实验名称
+            config: Configuration
+            agent_class: Agent class
+            agent_config_file: Agent configuration file
+            metric_extractors: Metric extractor
+            enable_institution: Whether to enable institution
+            agent_prefix: Agent name prefix
+            exp_name: Experiment name
+            logging_level: Logging level
         """
         self.exp_id = str(uuid.uuid4())
         if isinstance(agent_class, list):
@@ -166,11 +170,12 @@ class AgentSimulation:
                 experiment_name=exp_name,
                 run_id=mlflow_run_id,
             )
-            self.metric_extractor = metric_extractor
+            if metric_extractors is not None:
+                self.metric_extractors = metric_extractors
         else:
             logger.warning("Mlflow is not enabled, NO MLFLOW STORAGE")
             self.mlflow_client = None
-            self.metric_extractor = None
+            self.metric_extractors = None
 
         # pg
         _pgsql_config: dict[str, Any] = _storage_config.get("pgsql", {})
@@ -216,8 +221,9 @@ class AgentSimulation:
         - enable_institution: bool, default is True
         - agent_config:
             - agent_config_file: Optional[dict[type[Agent], str]]
+            - memory_config_init_func: Optional[Callable]
             - memory_config_func: Optional[dict[type[Agent], Callable]]
-            - metric_extractor: Optional[list[tuple[int, Callable]]]
+            - metric_extractors: Optional[list[tuple[int, Callable]]]
             - init_func: Optional[list[Callable[AgentSimulation, None]]]
             - group_size: Optional[int]
             - embedding_model: Optional[EmbeddingModel]
@@ -258,7 +264,7 @@ class AgentSimulation:
         simulation = cls(
             config=simulation_config,
             agent_config_file=config["agent_config"].get("agent_config_file", None),
-            metric_extractor=config["agent_config"].get("metric_extractor", None),
+            metric_extractors=config["agent_config"].get("metric_extractors", None),
             enable_institution=config.get("enable_institution", True),
             exp_name=config.get("exp_name", "default_experiment"),
             logging_level=config.get("logging_level", logging.WARNING),
@@ -274,12 +280,16 @@ class AgentSimulation:
         )
         simulation._simulator.set_environment(environment)
         logger.info("Initializing Agents...")
-        agent_count = []
-        agent_count.append(config["agent_config"]["number_of_citizen"])
-        agent_count.append(config["agent_config"]["number_of_firm"])
-        agent_count.append(config["agent_config"]["number_of_government"])
-        agent_count.append(config["agent_config"]["number_of_bank"])
-        agent_count.append(config["agent_config"]["number_of_nbs"])
+        agent_count = {
+            SocietyAgent: config["agent_config"].get("number_of_citizen", 0),
+            FirmAgent: config["agent_config"].get("number_of_firm", 0),
+            GovernmentAgent: config["agent_config"].get("number_of_government", 0),
+            BankAgent: config["agent_config"].get("number_of_bank", 0),
+            NBSAgent: config["agent_config"].get("number_of_nbs", 0),
+        }
+        if agent_count.get(SocietyAgent, 0) == 0:
+            raise ValueError("number_of_citizen is required")
+        
         # support MessageInterceptor
         if "message_intercept" in config:
             _intercept_config = config["message_intercept"]
@@ -318,7 +328,8 @@ class AgentSimulation:
             embedding_model=config["agent_config"].get(
                 "embedding_model", SimpleEmbedding()
             ),
-            memory_config_func=config["agent_config"].get("memory_config_func", None),
+            memory_config_func=config["agent_config"].get("memory_config_func", None),  
+            memory_config_init_func=config["agent_config"].get("memory_config_init_func", None),
             **_message_intercept_kwargs,
             environment=environment,
         )
@@ -465,7 +476,7 @@ class AgentSimulation:
 
     async def init_agents(
         self,
-        agent_count: Union[int, list[int]],
+        agent_count: dict[type[Agent], int],
         group_size: int = 10000,
         pg_sql_writers: int = 32,
         message_interceptors: int = 1,
@@ -473,6 +484,7 @@ class AgentSimulation:
         social_black_list: Optional[list[tuple[str, str]]] = None,
         message_listener: Optional[MessageBlockListenerBase] = None,
         embedding_model: Embeddings = SimpleEmbedding(),
+        memory_config_init_func: Optional[Callable] = None,
         memory_config_func: Optional[dict[type[Agent], Callable]] = None,
         environment: Optional[dict[str, str]] = None,
     ) -> None:
@@ -486,12 +498,13 @@ class AgentSimulation:
             memory_config_func: 返回Memory配置的函数，需要返回(EXTRA_ATTRIBUTES, PROFILE, BASE)元组, 每个元素表示一个智能体类创建的Memory配置函数
             environment: 环境变量，用于更新模拟器的环境变量
         """
-        if not isinstance(agent_count, list):
-            agent_count = [agent_count]
+        self.agent_count = agent_count
 
         if len(self.agent_class) != len(agent_count):
-            raise ValueError("agent_class和agent_count的长度不一致")
+            raise ValueError("The length of agent_class and agent_count does not match")
 
+        if memory_config_init_func is not None:
+            await memory_config_init(self)
         if memory_config_func is None:
             memory_config_func = self.default_memory_config_func  # type:ignore
 
@@ -503,9 +516,11 @@ class AgentSimulation:
         citizen_params = []
 
         # 收集所有参数
+        print(self.agent_class)
+        print(agent_count)
         for i in range(len(self.agent_class)):
             agent_class = self.agent_class[i]
-            agent_count_i = agent_count[i]
+            agent_count_i = agent_count[agent_class]
             assert memory_config_func is not None
             memory_config_func_i = memory_config_func.get(
                 agent_class, self.default_memory_config_func[agent_class]  # type:ignore

{pycityagent-2.0.0a65.dist-info → pycityagent-2.0.0a66.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: pycityagent
-Version: 2.0.0a65
+Version: 2.0.0a66
 Summary: LLM-based city environment agent building library
 Author-email: Yuwei Yan <pinkgranite86@gmail.com>, Junbo Yan <yanjb20thu@gmali.com>, Jun Zhang <zhangjun990222@gmali.com>
 License: MIT License

{pycityagent-2.0.0a65.dist-info → pycityagent-2.0.0a66.dist-info}/RECORD

@@ -1,3 +1,9 @@
+pycityagent-2.0.0a66.dist-info/RECORD,,
+pycityagent-2.0.0a66.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
+pycityagent-2.0.0a66.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
+pycityagent-2.0.0a66.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
+pycityagent-2.0.0a66.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
+pycityagent-2.0.0a66.dist-info/METADATA,sha256=Gd-mcYNBdEIpmg-DYNadjFxVQicXesWi6GK71M85Gfs,9110
 pycityagent/pycityagent-sim,sha256=vskCJGHJEh0B2dUfmYlVyrcy3sDZ3kBNwjqcYUZpmO8,35449490
 pycityagent/__init__.py,sha256=PUKWTXc-xdMG7px8oTNclodsILUgypANj2Z647sY63k,808
 pycityagent/pycityagent-ui,sha256=K2XXJhxIoIk4QWty5i-0FuzZJekkFlbeqrJgPX3tbdE,41225346
@@ -10,10 +16,10 @@ pycityagent/tools/__init__.py,sha256=XtdtGyWeFyK1YOUvWkykBWxemtmwQjWUIuuyU1-gosQ
 pycityagent/tools/tool.py,sha256=D-ESFlX7EESm5mcvs2zRlGEQTzXbVfQc8G7Vpz8TmAw,8651
 pycityagent/llm/llmconfig.py,sha256=4Ylf4OFSBEFy8jrOneeX0HvPhWEaF5jGvy1HkXK08Ro,436
 pycityagent/llm/__init__.py,sha256=iWs6FLgrbRVIiqOf4ILS89gkVCTvS7HFC3vG-MWuyko,205
-pycityagent/llm/llm.py,sha256=4AJlTj9llT909gzx9SfcBvbUJrWqHzEo3DaHWgQ3-3I,15852
+pycityagent/llm/llm.py,sha256=upm246fUurltONXies4I4oiZ7NkG0Xk8nj4xVRTNV9o,17215
 pycityagent/llm/embeddings.py,sha256=2_P4TWm3sJKFdGDx2Q1a2AEapFopDctIXsGuntvmP6E,6816
 pycityagent/llm/utils.py,sha256=hoNPhvomb1u6lhFX0GctFipw74hVKb7bvUBDqwBzBYw,160
-pycityagent/memory/memory.py,sha256=MusbnD-Us5IF16CkYRSlcer-TEbaPsHxQEYzcugo0N4,34589
+pycityagent/memory/memory.py,sha256=-o_W0uC_cOPyWPxdXiCLmU9XiTHF82DQYb53RQyYgeg,34588
 pycityagent/memory/profile.py,sha256=q8ZS9IBmHCg_X1GONUvXK85P6tCepTKQgXKuvuXYNXw,5203
 pycityagent/memory/__init__.py,sha256=_Vfdo1HcLWsuuz34_i8e91nnLVYADpMlHHSVaB3xgIk,297
 pycityagent/memory/memory_base.py,sha256=QG_j3BxZvkadFEeE3uBR_kjl_xcXD1aHUVs8GEF3d6w,5654
@@ -22,9 +28,9 @@ pycityagent/memory/utils.py,sha256=oJWLdPeJy_jcdKcDTo9JAH9kDZhqjoQhhv_zT9qWC0w,8
 pycityagent/memory/const.py,sha256=6zpJPJXWoH9-yf4RARYYff586agCoud9BRn7sPERB1g,932
 pycityagent/memory/faiss_query.py,sha256=V3rIw6d1_xcpNqZBbAYz3qfjVNE7NfJ7xOS5SibPtVU,13180
 pycityagent/memory/state.py,sha256=TYItiyDtehMEQaSBN7PpNrnNxdDM5jGppr9R9Ufv3kA,5134
-pycityagent/simulation/simulation.py,sha256=smWCN3qt6i0C-mWd26WM5yVLsfUAYZIgNsHF0fhYAhM,36949
+pycityagent/simulation/simulation.py,sha256=QRD994g4fRzDTbuCEmYGFdyFxix5a5XnXikgy4L19oc,37765
 pycityagent/simulation/__init__.py,sha256=P5czbcg2d8S0nbbnsQXFIhwzO4CennAhZM8OmKvAeYw,194
-pycityagent/simulation/agentgroup.py,sha256=PBmIp6waJR40Gjki9tTXtKMFJ_c5s8SIsi98FHRg-34,31598
+pycityagent/simulation/agentgroup.py,sha256=BpUWmN_CWR7-PJhq9e4pSsVjXZWcoW3Wt-gh5MIxQjQ,31674
 pycityagent/simulation/storage/pg.py,sha256=xRshSOGttW-p0re0fNBOjOpb-nQ5msIE2LsdT79_E_Y,8425
 pycityagent/message/message_interceptor.py,sha256=w8XTyZStQtMjILpeAX3VMhAWcYAuaxCgSMwXQU1OryM,8951
 pycityagent/message/__init__.py,sha256=f5QH7DKPqEAMyfSlBMnl3uouOKlsoel909STlIe7nUk,276
@@ -75,31 +81,25 @@ pycityagent/environment/sim/clock_service.py,sha256=gBUujvX_vIFMKVfcLRyk1GcpRRL6
 pycityagent/environment/sim/road_service.py,sha256=bKyn3_me0sGmaJVyF6eNeFbdU-9C1yWsa9L7pieDJzg,1285
 pycityagent/environment/interact/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pycityagent/environment/interact/interact.py,sha256=ifxPPzuHeqLHIZ_6zvfXMoBOnBsXNIP4bYp7OJ7pnEQ,6588
-pycityagent/cityagent/memory_config.py,sha256=8PqSeCdIgkp5a9KN2VdyaraMACu93d5W9TLzYTozmfo,10977
+pycityagent/cityagent/memory_config.py,sha256=nBUkLbxmiNtw4J4s34n4brgyFXnhpNYvQxws0c2n-ho,11356
 pycityagent/cityagent/bankagent.py,sha256=lr4GEcqt-iwA7DXoDry0WXkV6benmdaAyLpswqSpKlY,2120
 pycityagent/cityagent/__init__.py,sha256=gcBQ-a50XegFtjigQ7xDXRBZrywBKqifiQFSRnEF8gM,572
 pycityagent/cityagent/firmagent.py,sha256=UVlNN0lpa4cC4PZVqYzQhbc5VJ2oGsA1731mhbCjnR8,4109
 pycityagent/cityagent/nbsagent.py,sha256=WIXW__6dZ5IrqBqDCjvGbrCshpXzuFRV3Ww6gkYw7p4,4387
-pycityagent/cityagent/initial.py,sha256=BXU9ndY-yIASwoRvFosbhUptGukd68PM1h1vmTEaa34,4946
-pycityagent/cityagent/societyagent.py,sha256=RwRS7XaPNFmlTCRJgIpTNkDCWQCG48lpf2SjnZvtcHQ,20288
+pycityagent/cityagent/initial.py,sha256=k9iolPtDj5KiOvg_FaDWPpqhLCpP19T8gmGTKLISa40,4694
+pycityagent/cityagent/societyagent.py,sha256=REGEBWGT7ScKVS9o-x8iJMWSoAPXrMRpBZhAxxYuxS4,20312
 pycityagent/cityagent/message_intercept.py,sha256=1YMOs6-6bbAaTt7RfMn-ALVIcp0frHN7oqGUkWRy5xE,4519
 pycityagent/cityagent/governmentagent.py,sha256=HJLuhvEmllu_1KnFEJsYCIasaBJT0BV9Cn_4Y2QGPqg,2791
 pycityagent/cityagent/blocks/dispatcher.py,sha256=mEa1r3tRS3KI1BMZR_w_sbUGzOj6aUJuiUrsHv1n2n0,2943
-pycityagent/cityagent/blocks/needs_block.py,sha256=s8LikgtKORfo_Sw9SQ5_3biNPTof15QuUs4cDynXCyM,15332
+pycityagent/cityagent/blocks/needs_block.py,sha256=ZAf1cQq1N73YBBPEejYF2vAfEkXXTAVUcvqHIy-9Rhs,14935
 pycityagent/cityagent/blocks/cognition_block.py,sha256=zDbyyLh5GEqje9INJUJA1gMSDPW0wX5lt6yNu97XXn0,14818
-pycityagent/cityagent/blocks/social_block.py,sha256=y46mPK9SLvcOHYB64l6qz5ZgT0dWNmX-C3Cusj_44mE,15540
+pycityagent/cityagent/blocks/social_block.py,sha256=Dust9Tpu145h24gbJzEyy_a-IzbnZ0KtTNwTvteVb6w,15138
 pycityagent/cityagent/blocks/__init__.py,sha256=wydR0s-cCRWgdvQetkfQnD_PU8vC3eTmt2zntcb4fSA,452
 pycityagent/cityagent/blocks/economy_block.py,sha256=m5B67cgGZ9nKWtrYeak5gxMoCoKlRbATAsXpFajYKyg,19129
 pycityagent/cityagent/blocks/utils.py,sha256=8O5p1B8JlreIJTGXKAP03rTcn7MvFSR8qJ1_hhszboU,2065
 pycityagent/cityagent/blocks/other_block.py,sha256=NnDwxQAO5XZ7Uxe-n3qtrfNItHlwFYk2MQsh2GYDKMQ,4338
 pycityagent/cityagent/blocks/plan_block.py,sha256=v04ePs-6b86TyaP3fl9HPMwWh3_lHp4cjyoEQBHkoDU,11280
-pycityagent/cityagent/blocks/mobility_block.py,sha256=f9PlHYX_sCpfOgIhxx3DSA9aViW_e-yNIFpVGwZixUI,12623
+pycityagent/cityagent/blocks/mobility_block.py,sha256=rFRqyVnZG-BMXU0VNAkyrpHFYLv0_QUv4o8vI1wdG5A,15042
 pycityagent/survey/models.py,sha256=YE50UUt5qJ0O_lIUsSY6XFCGUTkJVNu_L1gAhaCJ2fs,3546
 pycityagent/survey/__init__.py,sha256=rxwou8U9KeFSP7rMzXtmtp2fVFZxK4Trzi-psx9LPIs,153
 pycityagent/survey/manager.py,sha256=S5IkwTdelsdtZETChRcfCEczzwSrry_Fly9MY4s3rbk,1681
-pycityagent-2.0.0a65.dist-info/RECORD,,
-pycityagent-2.0.0a65.dist-info/LICENSE,sha256=n2HPXiupinpyHMnIkbCf3OTYd3KMqbmldu1e7av0CAU,1084
-pycityagent-2.0.0a65.dist-info/WHEEL,sha256=NW1RskY9zow1Y68W-gXg0oZyBRAugI1JHywIzAIai5o,109
-pycityagent-2.0.0a65.dist-info/entry_points.txt,sha256=BZcne49AAIFv-hawxGnPbblea7X3MtAtoPyDX8L4OC4,132
-pycityagent-2.0.0a65.dist-info/top_level.txt,sha256=yOmeu6cSXmiUtScu53a3s0p7BGtLMaV0aff83EHCTic,43
-pycityagent-2.0.0a65.dist-info/METADATA,sha256=4zdmqXA1BgOR40yYPnJ1ZC3S3A13TvDy6FChU3-Q9JI,9110