pycityagent 2.0.0a66__cp39-cp39-macosx_11_0_arm64.whl → 2.0.0a67__cp39-cp39-macosx_11_0_arm64.whl

pycityagent/cityagent/initial.py

@@ -1,55 +1,67 @@
 import random
-
+import numpy as np
 from pycityagent.cityagent import (BankAgent, FirmAgent, GovernmentAgent,
                                    NBSAgent, SocietyAgent)
 
+import logging
+
+logger = logging.getLogger("pycityagent")
 
 async def initialize_social_network(simulation):
     """
-    初始化智能体之间的社交网络，包括好友关系类型、好友关系和社交关系强度
+    Initializes the social network between agents.
+
+    - **Description**:
+        - Creates friendship relationships between agents
+        - Assigns relationship types (family, colleague, friend)
+        - Sets relationship strengths based on type
+        - Initializes chat histories and interaction records
+
+    - **Returns**:
+        - None
     """
     try:
-        print("Initializing social network...")
+        logger.info("Initializing social network...")
 
-        # 定义可能的关系类型
+        # Define possible relationship types
         relation_types = ["family", "colleague", "friend"]
 
-        # 获取所有智能体ID
+        # Get all agent IDs
         agent_ids = simulation.agent_uuids
         for agent_id in agent_ids:
-            # 为每个智能体随机选择2-5个好友
+            # Randomly select 2-5 friends for each agent
             num_friends = random.randint(2, 5)
             possible_friends = [aid for aid in agent_ids if aid != agent_id]
             friends = random.sample(
                 possible_friends, min(num_friends, len(possible_friends))
             )
 
-            # 初始化好友关系
+            # Initialize friend relationships
             await simulation.update(agent_id, "friends", friends)
 
-            # 初始化与每个好友的关系类型和关系强度
+            # Initialize relationship types and strengths with each friend
             relationships = {}
             relation_type_map = {}
 
             for friend_id in friends:
-                # 随机选择关系类型
+                # Randomly select relationship type
                 relation_type = random.choice(relation_types)
-                # 根据关系类型设置初始关系强度范围
+                # Set initial relationship strength range based on type
                 if relation_type == "family":
-                    strength = random.randint(60, 90)  # 家人关系强度较高
+                    strength = random.randint(60, 90)  # Higher strength for family
                 elif relation_type == "colleague":
-                    strength = random.randint(40, 70)  # 同事关系强度中等
+                    strength = random.randint(40, 70)  # Medium strength for colleagues
                 else:  # friend
-                    strength = random.randint(30, 80)  # 朋友关系强度范围较广
+                    strength = random.randint(30, 80)  # Wide range for friends
 
                 relationships[friend_id] = strength
                 relation_type_map[friend_id] = relation_type
 
-            # 更新关系强度和类型
+            # Update relationship strengths and types
             await simulation.update(agent_id, "relationships", relationships)
             await simulation.update(agent_id, "relation_types", relation_type_map)
 
-            # 初始化空的聊天历史和互动记录
+            # Initialize empty chat histories and interaction records
             await simulation.update(
                 agent_id, "chat_histories", {friend_id: [] for friend_id in friends}
             )
@@ -64,12 +76,47 @@ async def initialize_social_network(simulation):
         print(f"Error initializing social network: {str(e)}")
         return False
 
+def zipf_distribution(N, F, s=1.0):
+    """
+    Generates employee counts for F companies following Zipf's law, with total employees N.
+
+    - **Description**:
+        - Uses Zipf's law to distribute N total employees across F companies
+        - The distribution follows a power law where employee count is proportional to 1/rank^s
+        - Normalizes the distribution to ensure total employees equals N
+
+    - **Parameters**:
+        - `N`: Total number of employees across all companies
+        - `F`: Number of companies to distribute employees across  
+        - `s`: Power law exponent for Zipf's law, typically close to 1
+
+    - **Returns**:
+        - List of integer employee counts for each company, summing to N
+    """
+    # Calculate employee count for each rank (following Zipf's law distribution)
+    ranks = np.arange(1, F + 1)  # Ranks from 1 to F
+    sizes = 1 / (ranks ** s)  # Calculate employee count ratio according to Zipf's law
+
+    # Calculate normalization coefficient to make total employees equal N
+    total_size = np.sum(sizes)
+    normalized_sizes = sizes / total_size * N  # Normalize to total employees N
+
+    # Return employee count for each company (integers)
+    return np.round(normalized_sizes).astype(int)
 
 async def bind_agent_info(simulation):
     """
-    绑定智能体的信息，包括公民、公司、政府、银行和NBS的ID
+    Binds agent information including IDs for citizens, firms, government, banks and NBS.
+
+    - **Description**:
+        - Gathers all agent IDs and maps them between UUID and agent ID
+        - Assigns employees to firms following Zipf's law distribution
+        - Links citizens to government and bank systems
+
+    - **Returns**:
+        - None
     """
-    print("Binding agent info...")
+    logger.info("Binding agent info...")
     infos = await simulation.gather("id")
     citizen_uuids = await simulation.filter(types=[SocietyAgent])
     firm_uuids = await simulation.filter(types=[FirmAgent])
@@ -77,36 +124,24 @@ async def bind_agent_info(simulation):
     bank_uuids = await simulation.filter(types=[BankAgent])
     nbs_uuids = await simulation.filter(types=[NBSAgent])
     citizen_agent_ids = []
-    firm_ids = []
-    id2uuid = {}
+    uid2agent, agent2uid = dict(), dict()
     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:
-                bank_id = v
-            elif k in nbs_uuids:
-                nbs_id = v
-    for citizen_uuid in citizen_uuids:
-        random_firm_id = random.choice(firm_ids)
-        await simulation.update(citizen_uuid, "firm_id", random_firm_id)
-        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)
-    for firm_uuid in firm_uuids:
-        await simulation.update(firm_uuid, "employees", citizen_uuids)
-        await simulation.update(firm_uuid, "employees_agent_id", citizen_agent_ids)
+            uid2agent[k] = v
+            agent2uid[v] = k
+    citizen_agent_ids_cp = citizen_agent_ids.copy()
+    random.shuffle(citizen_agent_ids_cp)
+    employee_sizes = zipf_distribution(len(citizen_agent_ids_cp), len(firm_uuids))
+    for firm_uuid, size in zip(firm_uuids, employee_sizes):
+        await simulation.economy_update(uid2agent[firm_uuid], "employees", citizen_agent_ids_cp[:size])
+        for citizen_agent_id in citizen_agent_ids_cp[:size]:
+            await simulation.update(agent2uid[citizen_agent_id], "firm_id", uid2agent[firm_uuid])
+        citizen_agent_ids_cp = citizen_agent_ids_cp[size:]
     for government_uuid in government_uuids:
-        await simulation.update(government_uuid, "citizens", citizen_uuids)
-        await simulation.update(government_uuid, "citizens_agent_id", citizen_agent_ids)
+        await simulation.economy_update(uid2agent[government_uuid], "citizens", citizen_agent_ids)
     for bank_uuid in bank_uuids:
-        await simulation.update(bank_uuid, "citizens", citizen_uuids)
-        await simulation.update(bank_uuid, "citizens_agent_id", citizen_agent_ids)
-    for nbs_uuid in nbs_uuids:
-        await simulation.update(nbs_uuid, "firm_id", random.choice(firm_ids))
-    print("Agent info binding completed!")
+        await simulation.economy_update(uid2agent[bank_uuid], "citizens", citizen_agent_ids)
+    logger.info("Agent info binding completed!")
+

pycityagent/cityagent/memory_config.py

@@ -7,8 +7,8 @@ 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)
+payment_max_skill_multiplier_base = 950
+payment_max_skill_multiplier = float(payment_max_skill_multiplier_base)
 pmsm = payment_max_skill_multiplier
 pareto_samples = np.random.pareto(pareto_param, size=(1000, 10))
 clipped_skills = np.minimum(pmsm, (pmsm - 1) * pareto_samples + 1)
@@ -27,20 +27,17 @@ def memory_config_societyagent():
     EXTRA_ATTRIBUTES = {
         "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),
-
         # Plan Behavior Model
         "current_plan": (list, [], True),
         "current_step": (dict, {"intention": "", "type": ""}, True),
         "execution_context": (dict, {}, True),
         "plan_history": (list, [], True),
-
         # cognition
         "emotion": (
             dict,
@@ -57,7 +54,6 @@ def memory_config_societyagent():
         "attitude": (dict, {}, True),
         "thought": (str, "Currently nothing good or bad is happening", True),
         "emotion_types": (str, "Relief", True),
-
         # economy
         "work_skill": (
             float,
@@ -88,7 +84,6 @@ def memory_config_societyagent():
         "working_experience": (list, [], True),
         "work_hour_month": (float, 160, True),
         "work_hour_finish": (float, 0, True),
-
         # social
         "friends": (list, [], True),  # friends list
         "relationships": (dict, {}, True),  # relationship strength with each friend
@@ -96,97 +91,122 @@ def memory_config_societyagent():
         "chat_histories": (dict, {}, True),  # all chat histories
         "interactions": (dict, {}, True),  # all interaction records
         "to_discuss": (dict, {}, True),
-
         # mobility
         "environment": (str, "The environment outside is good", True),
         "number_poi_visited": (int, 1, True),
     }
 
     PROFILE = {
-        "name": (str,random.choice(
-            [
-                "Alice",
-                "Bob",
-                "Charlie",
-                "David",
-                "Eve",
-                "Frank",
-                "Grace",
-                "Helen",
-                "Ivy",
-                "Jack",
-                "Kelly",
-                "Lily",
-                "Mike",
-                "Nancy",
-                "Oscar",
-                "Peter",
-                "Queen",
-                "Rose",
-                "Sam",
-                "Tom",
-                "Ulysses",
-                "Vicky",
-                "Will",
-                "Xavier",
-                "Yvonne",
-                "Zack",
-            ]
-        ), True),
+        "name": (
+            str,
+            random.choice(
+                [
+                    "Alice",
+                    "Bob",
+                    "Charlie",
+                    "David",
+                    "Eve",
+                    "Frank",
+                    "Grace",
+                    "Helen",
+                    "Ivy",
+                    "Jack",
+                    "Kelly",
+                    "Lily",
+                    "Mike",
+                    "Nancy",
+                    "Oscar",
+                    "Peter",
+                    "Queen",
+                    "Rose",
+                    "Sam",
+                    "Tom",
+                    "Ulysses",
+                    "Vicky",
+                    "Will",
+                    "Xavier",
+                    "Yvonne",
+                    "Zack",
+                ]
+            ),
+            True,
+        ),
         "gender": (str, random.choice(["male", "female"]), True),
         "age": (int, random.randint(18, 65), True),
-        "education": (str, random.choice(
-            ["Doctor", "Master", "Bachelor", "College", "High School"]
-        ), True),
-        "skill": (str, random.choice(
-            [
-                "Good at problem-solving",
-                "Good at communication",
-                "Good at creativity",
-                "Good at teamwork",
-                "Other",
-            ]
-        ), True),
-        "occupation": (str, random.choice(
-            [
-                "Student",
-                "Teacher",
-                "Doctor",
-                "Engineer",
-                "Manager",
-                "Businessman",
-                "Artist",
-                "Athlete",
-                "Other",
-            ]
-        ), True),
+        "education": (
+            str,
+            random.choice(["Doctor", "Master", "Bachelor", "College", "High School"]),
+            True,
+        ),
+        "skill": (
+            str,
+            random.choice(
+                [
+                    "Good at problem-solving",
+                    "Good at communication",
+                    "Good at creativity",
+                    "Good at teamwork",
+                    "Other",
+                ]
+            ),
+            True,
+        ),
+        "occupation": (
+            str,
+            random.choice(
+                [
+                    "Student",
+                    "Teacher",
+                    "Doctor",
+                    "Engineer",
+                    "Manager",
+                    "Businessman",
+                    "Artist",
+                    "Athlete",
+                    "Other",
+                ]
+            ),
+            True,
+        ),
         "family_consumption": (str, random.choice(["low", "medium", "high"]), True),
-        "consumption": (str, random.choice(["sightly low", "low", "medium", "high"]), True),
-        "personality": (str, random.choice(
-            ["outgoint", "introvert", "ambivert", "extrovert"]
-        ), True),
-        "income": "0",
-        "currency": random.randint(1000, 100000),
+        "consumption": (float, 0, True),
+        "personality": (
+            str,
+            random.choice(["outgoint", "introvert", "ambivert", "extrovert"]),
+            True,
+        ),
+        "income": (float, 0, True),
+        "currency": (float, random.randint(1000, 100000), True),
         "residence": (str, random.choice(["city", "suburb", "rural"]), True),
-        "race": (str, random.choice(
-            [
-                "Chinese",
-                "American",
-                "British",
-                "French",
-                "German",
-                "Japanese",
-                "Korean",
-                "Russian",
-                "Other",
-            ]
-        ), True),
-        "religion": (str, random.choice(
-            ["none", "Christian", "Muslim", "Buddhist", "Hindu", "Other"]
-        ), True),
-        "marital_status": (str, random.choice(
-            ["not married", "married", "divorced", "widowed"]
-        ), True),
+        "race": (
+            str,
+            random.choice(
+                [
+                    "Chinese",
+                    "American",
+                    "British",
+                    "French",
+                    "German",
+                    "Japanese",
+                    "Korean",
+                    "Russian",
+                    "Other",
+                ]
+            ),
+            True,
+        ),
+        "religion": (
+            str,
+            random.choice(
+                ["none", "Christian", "Muslim", "Buddhist", "Hindu", "Other"]
+            ),
+            True,
+        ),
+        "marital_status": (
+            str,
+            random.choice(["not married", "married", "divorced", "widowed"]),
+            True,
+        ),
     }
 
     BASE = {

pycityagent/cityagent/message_intercept.py

@@ -8,10 +8,6 @@ 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}")
     is_valid = True
     prompt = f"""

pycityagent/cityagent/metrics.py

@@ -0,0 +1,41 @@
+import pycityproto.city.economy.v2.economy_pb2 as economyv2
+from pycityagent.cityagent import SocietyAgent
+
+async def mobility_metric(simulation):
+    # 使用函数属性来存储计数
+    if not hasattr(mobility_metric, 'step_count'):
+        mobility_metric.step_count = 0
+        
+    # 统计访问地点数量
+    citizen_agents = await simulation.filter(types = [SocietyAgent])
+    poi_visited_info = await simulation.gather( "number_poi_visited", citizen_agents)
+    poi_visited_sum = 0
+    for group_gather in poi_visited_info:
+        for agent_uuid, poi_visited in group_gather.items():
+            poi_visited_sum += poi_visited
+    average_poi_visited = float(poi_visited_sum / len(citizen_agents))
+    print(f"Metric: Average POIs visited: {average_poi_visited}")
+    await simulation.mlflow_client.log_metric(key="average_poi_visited", value=average_poi_visited, step=mobility_metric.step_count)
+    await simulation.mlflow_client.log_metric(key="poi_visited_sum", value=poi_visited_sum, step=mobility_metric.step_count)
+    mobility_metric.step_count += 1
+
+async def economy_metric(simulation):
+    # 使用函数属性来存储计数
+    if not hasattr(economy_metric, 'step_count'):
+        economy_metric.step_count = 0
+    
+    nbs_id = await simulation.economy_client.get_org_entity_ids(economyv2.ORG_TYPE_NBS)
+    nbs_id = nbs_id[0]
+    try:
+        real_gdp = await simulation.economy_client.get(nbs_id, 'real_gdp')
+    except:
+        real_gdp = []
+    if len(real_gdp) > 0:
+        real_gdp = real_gdp[0]
+        await simulation.mlflow_client.log_metric(key="real_gdp", value=real_gdp, step=economy_metric.step_count)
+        other_metrics = ['prices', 'working_hours', 'depression', 'consumption_currency', 'income_currency']
+        other_metrics_names = ['price', 'working_hours', 'depression', 'consumption', 'income']
+        for metric, metric_name in zip(other_metrics, other_metrics_names):
+            metric_value = (await simulation.economy_client.get(nbs_id, metric))[-1]
+            await simulation.mlflow_client.log_metric(key=metric_name, value=metric_value, step=economy_metric.step_count)
+    economy_metric.step_count += 1

pycityagent/cityagent/nbsagent.py

@@ -9,11 +9,24 @@ from pycityagent.economy import EconomyClient
 from pycityagent.llm.llm import LLM
 from pycityagent.memory import Memory
 from pycityagent.message import Messager
+import pycityproto.city.economy.v2.economy_pb2 as economyv2
 
 logger = logging.getLogger("pycityagent")
 
 
 class NBSAgent(InstitutionAgent):
+    configurable_fields = ["time_diff", "num_labor_hours", "productivity_per_labor"]
+    default_values = {
+        "time_diff": 30 * 24 * 60 * 60,
+        "num_labor_hours": 168,
+        "productivity_per_labor": 1,
+    }
+    fields_description = {
+        "time_diff": "Time difference between each forward, day * hour * minute * second",
+        "num_labor_hours": "Number of labor hours per week",
+        "productivity_per_labor": "Productivity per labor hour",
+    }
+
     def __init__(
         self,
         name: str,
@@ -57,45 +70,25 @@ class NBSAgent(InstitutionAgent):
 
     async def forward(self):
         if await self.month_trigger():
-            citizens = await self.memory.status.get("citizens")
-            agents_forward = []
-            if not np.all(np.array(agents_forward) > self.forward_times):
-                return
-            work_propensity = await self.gather_messages(citizens, "work_propensity")
+            await self.economy_client.calculate_real_gdp()
+            citizens_uuid = await self.memory.status.get("citizens")
+            citizens = await self.economy_client.get(self._agent_id, "citizens")
+            work_propensity = await self.gather_messages(citizens_uuid, "work_propensity")
             working_hours = np.mean(work_propensity) * self.num_labor_hours
-            firm_id = await self.memory.status.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"
             )
+            firms_id = await self.economy_client.get_org_entity_ids(economyv2.ORG_TYPE_FIRM)
+            prices = await self.economy_client.get(firms_id, "price")
             await self.economy_client.update(
-                self._agent_id, "prices", [price], mode="merge"
+                self._agent_id, "prices", [float(np.mean(prices))], mode="merge"
             )
-            depression = await self.gather_messages(citizens, "depression")
+            depression = await self.gather_messages(citizens_uuid, "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"
-            )
+            consumption_currency = await self.economy_client.get(citizens, "consumption")
             consumption_currency = np.mean(consumption_currency)
             await self.economy_client.update(
                 self._agent_id,
@@ -103,13 +96,8 @@ class NBSAgent(InstitutionAgent):
                 [consumption_currency],
                 mode="merge",
             )
-            income_currency = await self.gather_messages(citizens, "income_currency")
+            income_currency = await self.economy_client.get(citizens, "income")
             income_currency = np.mean(income_currency)
             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}", "economy"
-                )
+            )

pycityagent/cityagent/societyagent.py

@@ -62,6 +62,7 @@ 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")
@@ -151,8 +152,10 @@ class PlanAndActionBlock(Block):
         # step execution
         await self.step_execution()
 
+
 class MindBlock(Block):
     """Cognition workflow"""
+
     cognitionBlock: CognitionBlock
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
@@ -164,11 +167,12 @@ class MindBlock(Block):
     async def forward(self):
         await self.cognitionBlock.forward()
 
+
 class SocietyAgent(CitizenAgent):
     update_with_sim = UpdateWithSimulator()
     mindBlock: MindBlock
     planAndActionBlock: PlanAndActionBlock
-    update_with_sim = UpdateWithSimulator()
+    
     configurable_fields = [
         "enable_cognition",
         "enable_mobility",
@@ -237,7 +241,7 @@ class SocietyAgent(CitizenAgent):
         ifpass = await self.check_and_update_step()
         if not ifpass:
             return
-        
+
         await self.planAndActionBlock.forward()
 
         if self.enable_cognition:
@@ -362,7 +366,7 @@ class SocietyAgent(CitizenAgent):
 
                 if not content:
                     return ""
-                
+
                 # 添加记忆
                 description = f"You received a social message: {content}"
                 await self.memory.stream.add_social(description=description)

pycityagent/cli/wrapper.py

@@ -1,7 +1,7 @@
 import os
-import sys
-import subprocess
 import signal
+import subprocess
+import sys
 
 _script_dir = os.path.dirname(os.path.abspath(__file__))
 _parent_dir = os.path.dirname(_script_dir)