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

pycityagent/cityagent/blocks/cognition_block.py

@@ -1,27 +1,31 @@
+import json
+import logging
+
+from pycityagent.environment.simulator import Simulator
 from pycityagent.llm import LLM
-from pycityagent.workflow.block import Block
 from pycityagent.memory import Memory
-from pycityagent.environment.simulator import Simulator
+from pycityagent.workflow.block import Block
 from pycityagent.workflow.prompt import FormatPrompt
-import json
-import logging
+
 logger = logging.getLogger("pycityagent")
 
+
 def extract_json(output_str):
     try:
         # Find the positions of the first '{' and the last '}'
-        start = output_str.find('{')
-        end = output_str.rfind('}')
-        
+        start = output_str.find("{")
+        end = output_str.rfind("}")
+
         # Extract the substring containing the JSON
-        json_str = output_str[start:end+1]
-        
+        json_str = output_str[start : end + 1]
+
         # Convert the JSON string to a dictionary
         return json_str
     except (ValueError, json.JSONDecodeError) as e:
         logger.warning(f"Failed to extract JSON: {e}")
         return None
 
+
 class CognitionBlock(Block):
     configurable_fields = ["top_k"]
     default_values = {"top_k": 20}
@@ -32,12 +36,12 @@ class CognitionBlock(Block):
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("CognitionBlock", llm=llm, memory=memory, simulator=simulator)
         self.top_k = 20
-        
+
     async def set_status(self, status):
         for key in status:
             await self.memory.status.update(key, status[key])
         return
-    
+
     async def attitude_update(self):
         """Cognition - attitude update workflow"""
         attitude = await self.memory.status.get("attitude")
@@ -56,7 +60,7 @@ class CognitionBlock(Block):
             "income": await self.memory.status.get("income"),
             "skill": await self.memory.status.get("skill"),
             "thought": await self.memory.status.get("thought"),
-            "emotion_types": await self.memory.status.get("emotion_types")
+            "emotion_types": await self.memory.status.get("emotion_types"),
         }
         for topic in attitude:
             description_prompt = """
@@ -82,7 +86,7 @@ class CognitionBlock(Block):
                 f"You need to decide your attitude towards topic: {topic}, "
                 f"which you previously rated your attitude towards this topic as: {previous_attitude} "
                 "(0 meaning oppose, 10 meaning support). "
-                "Please return a new attitude rating (0-10, smaller meaning oppose, larger meaning support) in JSON format, and explain, e.g. {{\"attitude\": 5}}"
+                'Please return a new attitude rating (0-10, smaller meaning oppose, larger meaning support) in JSON format, and explain, e.g. {{"attitude": 5}}'
             )
             question_prompt = description_prompt + incident_prompt + problem_prompt
             question_prompt = FormatPrompt(question_prompt)
@@ -99,22 +103,25 @@ class CognitionBlock(Block):
             prompt_data["disgust"] = disgust
             prompt_data["anger"] = anger
             prompt_data["surprise"] = surprise
-            
+
             question_prompt.format(**prompt_data)
             evaluation = True
+            response: dict = {}
             for retry in range(10):
                 try:
-                    response = await self.llm.atext_request(question_prompt.to_dialog(), timeout=300)
-                    response = json.loads(extract_json(response))
+                    _response = await self.llm.atext_request(
+                        question_prompt.to_dialog(), timeout=300
+                    )
+                    response = json.loads(extract_json(_response))  # type:ignore
                     evaluation = False
                     break
                 except:
                     pass
             if evaluation:
-                raise f"Request for attitude:{topic} update failed"
+                raise Exception(f"Request for attitude:{topic} update failed")
             attitude[topic] = response["attitude"]
         await self.memory.status.update("attitude", attitude)
-    
+
     async def thought_update(self):
         """Cognition - thought update workflow"""
         description_prompt = """
@@ -171,33 +178,36 @@ class CognitionBlock(Block):
             surprise=surprise,
             emotion=await self.memory.status.get("emotion"),
             thought=await self.memory.status.get("thought"),
-            emotion_types=await self.memory.status.get("emotion_types")
+            emotion_types=await self.memory.status.get("emotion_types"),
         )
-        
+
         evaluation = True
+        response: dict = {}
         for retry in range(10):
             try:
-                response = await self.llm.atext_request(question_prompt.to_dialog(), timeout=300)
-                response = json.loads(extract_json(response))
+                _response = await self.llm.atext_request(
+                    question_prompt.to_dialog(), timeout=300
+                )
+                response = json.loads(extract_json(_response))  # type:ignore
                 evaluation = False
                 break
             except:
                 pass
         if evaluation:
             raise Exception("Request for cognition update failed")
-        
+
         thought = str(response["thought"])
         await self.memory.status.update("thought", thought)
         await self.memory.stream.add_cognition(description=thought)
         return
-    
+
     async def end_of_day(self):
         """Cognition - end of day workflow"""
         time = await self.simulator.get_simulator_second_from_start_of_day()
         if time >= 86400 - 10 * 60:
             return True
         return False
-    
+
     async def forward(self):
         """Cognition workflow: Daily update"""
         # cognition update: thought and attitude
@@ -220,7 +230,7 @@ class CognitionBlock(Block):
         Joy, Distress, Resentment, Pity, Hope, Fear, Satisfaction, Relief, Disappointment, Pride, Admiration, Shame, Reproach, Liking, Disliking, Gratitude, Anger, Gratification, Remorse, Love, Hate.
         """
 
-        incident_prompt = f"{incident}"  #waiting for incident port
+        incident_prompt = f"{incident}"  # waiting for incident port
         question_prompt = """
             Please reconsider your emotion intensities: 
             sadness, joy, fear, disgust, anger, surprise (0 meaning not at all, 10 meaning very much).
@@ -256,14 +266,17 @@ class CognitionBlock(Block):
             surprise=surprise,
             emotion=await self.memory.status.get("emotion"),
             thought=await self.memory.status.get("thought"),
-            emotion_types=await self.memory.status.get("emotion_types")
+            emotion_types=await self.memory.status.get("emotion_types"),
         )
-        
+
         evaluation = True
+        response: dict = {}
         for retry in range(10):
             try:
-                response = await self.llm.atext_request(question_prompt.to_dialog(), timeout=300)
-                response = json.loads(extract_json(response))
+                _response = await self.llm.atext_request(
+                    question_prompt.to_dialog(), timeout=300
+                )
+                response = json.loads(extract_json(_response))  # type:ignore
                 evaluation = False
                 break
             except Exception as e:
@@ -272,6 +285,16 @@ class CognitionBlock(Block):
         if evaluation:
             raise Exception("Request for cognition update failed")
 
-        await self.memory.status.update("emotion", {"sadness": int(response["sadness"]), "joy": int(response["joy"]), "fear": int(response["fear"]), "disgust": int(response["disgust"]), "anger": int(response["anger"]), "surprise": int(response["surprise"])})
+        await self.memory.status.update(
+            "emotion",
+            {
+                "sadness": int(response["sadness"]),
+                "joy": int(response["joy"]),
+                "fear": int(response["fear"]),
+                "disgust": int(response["disgust"]),
+                "anger": int(response["anger"]),
+                "surprise": int(response["surprise"]),
+            },
+        )
         await self.memory.status.update("emotion_types", str(response["word"]))
         return response["conclusion"]

pycityagent/cityagent/blocks/dispatcher.py

@@ -1,9 +1,11 @@
+import logging
 import random
 from typing import Dict, List
+
 from pycityagent.llm.llm import LLM
 from pycityagent.workflow.block import Block
 from pycityagent.workflow.prompt import FormatPrompt
-import logging
+
 logger = logging.getLogger("pycityagent")
 
 DISPATCHER_PROMPT = """
@@ -14,18 +16,19 @@ Step information:
 {step}
 """
 
+
 class BlockDispatcher:
     def __init__(self, llm: LLM):
         self.llm = llm
         self.blocks: Dict[str, Block] = {}
         self.prompt = FormatPrompt(DISPATCHER_PROMPT)
-        
+
     def register_blocks(self, blocks: List[Block]) -> None:
         """Register multiple blocks at once"""
         for block in blocks:
             block_name = block.__class__.__name__.lower()
             self.blocks[block_name] = block
-        
+
     def _get_function_schema(self) -> dict:
         """Generate function schema for LLM function call"""
         # 创建 block 选项说明
@@ -33,7 +36,7 @@ class BlockDispatcher:
             name: block.description  # type: ignore
             for name, block in self.blocks.items()
         }
-        
+
         return {
             "type": "function",
             "function": {
@@ -45,34 +48,36 @@ class BlockDispatcher:
                         "block_name": {
                             "type": "string",
                             "enum": list(self.blocks.keys()),
-                            "description": f"Available blocks and their descriptions: {block_descriptions}"
+                            "description": f"Available blocks and their descriptions: {block_descriptions}",
                         }
                     },
-                    "required": ["block_name"]
-                }
-            }
+                    "required": ["block_name"],
+                },
+            },
         }
-        
+
     async def dispatch(self, step: dict) -> Block:
         """Dispatch the step to appropriate block based on LLM function call"""
         try:
             function_schema = self._get_function_schema()
             self.prompt.format(step=step)
-            
+
             # Call LLM with tools schema
             function_args = await self.llm.atext_request(
                 self.prompt.to_dialog(),
                 tools=[function_schema],
-                tool_choice={"type": "function", "function": {"name": "select_block"}}
+                tool_choice={"type": "function", "function": {"name": "select_block"}},
             )
-                    
-            selected_block = function_args.get("block_name") # type: ignore
-            
+
+            selected_block = function_args.get("block_name")  # type: ignore
+
             if selected_block not in self.blocks:
-                raise ValueError(f"Selected block '{selected_block}' not found in registered blocks")
-                
+                raise ValueError(
+                    f"Selected block '{selected_block}' not found in registered blocks"
+                )
+
             return self.blocks[selected_block]
-            
+
         except Exception as e:
             logger.warning(f"Failed to dispatch block: {e}")
             return random.choice(list(self.blocks.values()))

pycityagent/cityagent/blocks/economy_block.py

@@ -3,6 +3,7 @@ import random
 from pycityagent.llm.llm import LLM
 from pycityagent.workflow.block import Block
 from pycityagent.memory import Memory
+import pycityproto.city.economy.v2.economy_pb2 as economyv2
 
 from .utils import clean_json_response, TIME_ESTIMATE_PROMPT
 
@@ -21,6 +22,15 @@ from pycityagent.workflow import Block
 from pycityagent.economy import EconomyClient
 from .utils import *
 import numbers
+import numpy as np
+import asyncio
+
+def softmax(x, gamma=1.0):
+    if not isinstance(x, np.ndarray):
+        x = np.array(x)
+    x *= gamma
+    e_x = np.exp(x - np.max(x))
+    return e_x / e_x.sum(axis=-1, keepdims=True)
 
 class WorkBlock(Block):
     """WorkPlace Block"""
@@ -80,10 +90,10 @@ class ConsumptionBlock(Block):
     async def forward(self, step, context):
         self.forward_times += 1
         agent_id = await self.memory.status.get("id") # agent_id
-        firm_id = await self.memory.status.get('firm_id')
+        firms_id = await self.economy_client.get_org_entity_ids(economyv2.ORG_TYPE_FIRM)
         intention = step["intention"]
         month_consumption = await self.memory.status.get('to_consumption_currency')
-        consumption_currency = await self.memory.status.get('consumption_currency')
+        consumption_currency = await self.economy_client.get(agent_id, 'consumption')
         if consumption_currency >= month_consumption:
             node_id = await self.memory.stream.add_economy(description=f"I have passed the monthly consumption limit, so I will not consume.")
             return {
@@ -93,20 +103,20 @@ class ConsumptionBlock(Block):
                 'node_id': node_id
             }
         consumption = min(month_consumption/1, month_consumption-consumption_currency)
-        price = await self.economy_client.get(firm_id, 'price') # agent_id
-        wealth = await self.economy_client.get(agent_id, 'currency')
-        this_demand = int(consumption//price)
-        _, post_consumption_wealth = await self.economy_client.calculate_consumption(firm_id, [agent_id], [this_demand])
-        post_consumption_wealth = post_consumption_wealth[0]
-        await self.memory.status.update('consumption_currency', consumption_currency+wealth-post_consumption_wealth)
-        goods_demand = await self.memory.status.get('goods_demand')
-        goods_consumption = await self.memory.status.get('goods_consumption')
-        await self.memory.status.update('goods_demand', goods_demand+this_demand)
-        await self.memory.status.update('goods_consumption', goods_consumption+int((wealth-post_consumption_wealth)//price))
-        node_id = await self.memory.stream.add_economy(description=f"I bought some goods, and spent {this_demand} on {intention}")
+        prices = []
+        for firm_id in firms_id:
+            price = await self.economy_client.get(firm_id, 'price')
+            prices.append(price)
+        consumption_each_firm = consumption*softmax(prices, gamma=-0.01)
+        demand_each_firm = []
+        for i in range(len(firms_id)):
+            demand_each_firm.append(int(consumption_each_firm//prices[i]))
+
+        real_consumption = await self.economy_client.calculate_consumption(firms_id, agent_id, demand_each_firm)
+        node_id = await self.memory.stream.add_economy(description=f"I bought some goods, and spent {real_consumption:.1f} on {intention}")
         evaluation = {
             'success': True, 
-            'evaluation': f"I bought some goods, and spent {this_demand} on {intention}", 
+            'evaluation': f"I bought some goods, and spent {real_consumption:.1f} on {intention}", 
             'consumed_time': 20, 
             'node_id': node_id
         }
@@ -172,16 +182,22 @@ class MonthPlanBlock(Block):
     async def forward(self):
         if await self.month_trigger():
             agent_id = await self.memory.status.get("id")
+            firms_id = await self.economy_client.get_org_entity_ids(economyv2.ORG_TYPE_FIRM)
             firm_id = await self.memory.status.get('firm_id')
-            bank_id = await self.memory.status.get('bank_id')
+            bank_id = await self.economy_client.get_org_entity_ids(economyv2.ORG_TYPE_BANK)
+            bank_id = bank_id[0]
             name = await self.memory.status.get('name')
             age = await self.memory.status.get('age')
             city = await self.memory.status.get('city')
             job = await self.memory.status.get('occupation')
-            skill = await self.memory.status.get('work_skill')
-            consumption = await self.memory.status.get('consumption_currency')
+            skill = await self.economy_client.get(agent_id, 'skill')
+            consumption = await self.economy_client.get(agent_id, 'consumption')
             tax_paid = await self.memory.status.get('tax_paid')
-            price = await self.economy_client.get(firm_id, 'price')
+            prices = []
+            for firm_id in firms_id:
+                price = await self.economy_client.get(firm_id, 'price')
+                prices.append(price)
+            price = np.mean(prices)
             wealth = await self.economy_client.get(agent_id, 'currency')
             interest_rate = await self.economy_client.get(bank_id, 'interest_rate')
             
@@ -232,27 +248,27 @@ class MonthPlanBlock(Block):
             except:
                 self.llm_error += 1
                 
-            skill = await self.memory.status.get('work_skill')
+            work_skill = await self.economy_client.get(agent_id, 'skill')
             work_propensity = await self.memory.status.get('work_propensity')
             consumption_propensity = await self.memory.status.get('consumption_propensity')
             work_hours = work_propensity * num_labor_hours
-            await self.memory.status.update('income_currency', work_hours * skill)
+            income = await self.economy_client.get(agent_id, 'income')
+            income += work_hours * work_skill
+            
             wealth = await self.economy_client.get(agent_id, 'currency')
-            await self.economy_client.update(agent_id, 'currency', wealth + work_hours * skill)
+            wealth += work_hours * work_skill
+            await self.economy_client.update(agent_id, 'currency', wealth)
             await self.economy_client.add_delta_value(firm_id, 'inventory', int(work_hours*productivity_per_labor))
             
             if UBI and self.forward_times >= 96:
-                income_currency = await self.memory.status.get('income_currency')
-                await self.memory.status.update('income_currency', income_currency + UBI)
-                wealth = await self.economy_client.get(agent_id, 'currency')
-                await self.economy_client.update(agent_id, 'currency', wealth + UBI)
+                income += UBI
+                wealth += UBI
 
-            wealth = await self.economy_client.get(agent_id, 'currency')
             await self.memory.status.update('to_consumption_currency', consumption_propensity*wealth)
             
-            await self.memory.status.update('consumption_currency', 0)
-            await self.memory.status.update('goods_demand', 0)
-            await self.memory.status.update('goods_consumption', 0)
+            await self.economy_client.update(agent_id, 'consumption', 0)
+            await self.economy_client.update(agent_id, 'income', income)
+            await self.economy_client.update(agent_id, 'currency', wealth)
             
             if self.forward_times % 3 == 0:
                 obs_prompt = f'''
@@ -312,6 +328,4 @@ class MonthPlanBlock(Block):
                 obs_prompt = prettify_document(obs_prompt)
                 content = await self.llm.atext_request([{'role': 'user', 'content': obs_prompt}], timeout=300)
                 await self.memory.status.update('ubi_opinion', [content], mode='merge')
-
-            self.forward_times += 1
-            await self.memory.status.update('forward', self.forward_times)
+