pycityagent 2.0.0a47__cp39-cp39-macosx_11_0_arm64.whl → 2.0.0a49__cp39-cp39-macosx_11_0_arm64.whl

Sign up to get free protection for your applications and to get access to all the features.
Files changed (33) hide show
  1. pycityagent/__init__.py +3 -2
  2. pycityagent/agent.py +109 -4
  3. pycityagent/cityagent/__init__.py +20 -0
  4. pycityagent/cityagent/bankagent.py +54 -0
  5. pycityagent/cityagent/blocks/__init__.py +20 -0
  6. pycityagent/cityagent/blocks/cognition_block.py +304 -0
  7. pycityagent/cityagent/blocks/dispatcher.py +78 -0
  8. pycityagent/cityagent/blocks/economy_block.py +356 -0
  9. pycityagent/cityagent/blocks/mobility_block.py +258 -0
  10. pycityagent/cityagent/blocks/needs_block.py +305 -0
  11. pycityagent/cityagent/blocks/other_block.py +103 -0
  12. pycityagent/cityagent/blocks/plan_block.py +309 -0
  13. pycityagent/cityagent/blocks/social_block.py +345 -0
  14. pycityagent/cityagent/blocks/time_block.py +116 -0
  15. pycityagent/cityagent/blocks/utils.py +66 -0
  16. pycityagent/cityagent/firmagent.py +75 -0
  17. pycityagent/cityagent/governmentagent.py +60 -0
  18. pycityagent/cityagent/initial.py +98 -0
  19. pycityagent/cityagent/memory_config.py +202 -0
  20. pycityagent/cityagent/nbsagent.py +92 -0
  21. pycityagent/cityagent/societyagent.py +291 -0
  22. pycityagent/memory/memory.py +0 -18
  23. pycityagent/message/messager.py +6 -3
  24. pycityagent/simulation/agentgroup.py +123 -37
  25. pycityagent/simulation/simulation.py +311 -316
  26. pycityagent/workflow/block.py +66 -1
  27. pycityagent/workflow/tool.py +9 -4
  28. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/METADATA +2 -2
  29. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/RECORD +33 -14
  30. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/LICENSE +0 -0
  31. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/WHEEL +0 -0
  32. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/entry_points.txt +0 -0
  33. {pycityagent-2.0.0a47.dist-info → pycityagent-2.0.0a49.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,356 @@
1
+ import json
2
+ import random
3
+ from pycityagent.llm.llm import LLM
4
+ from pycityagent.workflow.block import Block
5
+ from pycityagent.memory import Memory
6
+
7
+ from .utils import clean_json_response, TIME_ESTIMATE_PROMPT
8
+
9
+ from .dispatcher import BlockDispatcher
10
+ from pycityagent.environment.simulator import Simulator
11
+ from pycityagent.llm import LLM
12
+ from pycityagent.memory import Memory
13
+ from pycityagent.workflow import Block
14
+ import random
15
+ import logging
16
+ logger = logging.getLogger("pycityagent")
17
+ import pickle as pkl
18
+
19
+ from pycityagent.workflow.prompt import FormatPrompt
20
+ from pycityagent.workflow import Block
21
+ from pycityagent.economy import EconomyClient
22
+ from .utils import *
23
+ import numbers
24
+ import asyncio
25
+
26
+
27
+
28
+
29
+ class WorkBlock(Block):
30
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
31
+ super().__init__("WorkBlock", llm, memory, simulator)
32
+ self.description = "Do work"
33
+ self.guidance_prompt = FormatPrompt(template=TIME_ESTIMATE_PROMPT)
34
+ async def forward(self, step, context):
35
+ self.guidance_prompt.format(
36
+ plan=context['plan'],
37
+ intention=step['intention']
38
+ )
39
+ result = await self.llm.atext_request(self.guidance_prompt.to_dialog())
40
+ result = clean_json_response(result)
41
+ try:
42
+ result = json.loads(result)
43
+ time = result['time']
44
+ start_time = await self.simulator.get_time(format_time=True)
45
+ await self.memory.update("working_experience", [f"Start from {start_time}, worked {time} minutes on {step['intention']}"], mode="merge")
46
+ work_hour_finish = await self.memory.get("work_hour_finish")
47
+ work_hour_finish += float(time/60)
48
+ await self.memory.update("work_hour_finish", work_hour_finish)
49
+ return {
50
+ 'success': True,
51
+ 'evaluation': f'工作:{step["intention"]}',
52
+ 'consumed_time': time
53
+ }
54
+ except Exception as e:
55
+ logger.warning(f"解析时间评估响应时发生错误: {str(e)}, 原始结果: {result}")
56
+ time = random.randint(1, 5)*60
57
+ start_time = await self.simulator.get_time(format_time=True)
58
+ await self.memory.update("working_experience", [f"Start from {start_time}, worked {time} minutes on {step['intention']}"], mode="merge")
59
+ work_hour_finish = await self.memory.get("work_hour_finish")
60
+ work_hour_finish += float(time/60)
61
+ await self.memory.update("work_hour_finish", work_hour_finish)
62
+ return {
63
+ 'success': True,
64
+ 'evaluation': f'工作:{step["intention"]}',
65
+ 'consumed_time': time
66
+ }
67
+
68
+ class ConsumptionBlock(Block):
69
+ """
70
+ determine the consumption place, time, amount, and items
71
+ """
72
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator, economy_client: EconomyClient):
73
+ super().__init__("ConsumptionBlock", llm, memory, simulator)
74
+ self.economy_client = economy_client
75
+ self.forward_times = 0
76
+ self.description = "Used to determine the consumption place, time, amount, and items"
77
+
78
+ async def forward(self, step, context):
79
+ self.forward_times += 1
80
+ agent_id = await self.memory.get("id") # agent_id
81
+ firm_id = await self.memory.get('firm_id')
82
+ intention = step["intention"]
83
+ nowPlace = await self.memory.get('position')
84
+ nowPlace = nowPlace['xy_position']
85
+ nowPlace = (nowPlace['x'], nowPlace['y'])
86
+ all_poi_types = self.simulator.get_poi_categories(nowPlace, 2000)
87
+ if len(all_poi_types) == 0:
88
+ all_poi_types.extend(['supermarket'])
89
+ prompt = 'Your intention is to {intention}. Which type of POI do you want to go to? Select one and only one type from the following types: {all_poi_types}. Any other output is NOT acceptable.'
90
+ prompt = FormatPrompt(prompt)
91
+ prompt.format(intention=intention, all_poi_types=all_poi_types)
92
+ poi_type = await self.llm.atext_request(prompt.to_dialog(), timeout=300)
93
+ poi_type = poi_type.strip('[').strip(']')
94
+ if poi_type not in all_poi_types:
95
+ poi_type = random.choice(all_poi_types)
96
+ around_pois = self.simulator.get_around_poi(nowPlace, 2000, poi_type)
97
+ if len(around_pois) == 0:
98
+ around_pois.extend(['沃尔玛', '家乐福', '华润万家', '物美'])
99
+ prompt = """
100
+ You're a {gender} and your education level is {education}.
101
+ Your intention is to {intention} and you plan to consume for ${consumption}.
102
+ The follwing are the places you can go to consume:
103
+ {around_pois}
104
+ Where should you go to make a purchase, when should you go, and what items will you spend on?
105
+ Respond with this format: [place, time, items to buy]. Any other output is NOT acceptable.
106
+ """
107
+ prompt = FormatPrompt(prompt)
108
+ data = {key: await self.memory.get(key) for key in ['gender', 'education']}
109
+ data['around_pois'] = around_pois
110
+ data['intention'] = intention
111
+ month_consumption = await self.memory.get('to_consumption_currency')
112
+ consumption_currency = await self.memory.get('consumption_currency')
113
+ if consumption_currency >= month_consumption:
114
+ return {'Success': True, 'Evaluation': f"Consumption Done", 'consumed_time': 0}
115
+ data['consumption'] = min(month_consumption/1, month_consumption-consumption_currency)
116
+ price = await self.economy_client.get(firm_id, 'price') # agent_id
117
+ wealth = await self.economy_client.get(agent_id, 'currency')
118
+ this_demand = int(data['consumption']//price)
119
+ _, post_consumption_wealth = await self.economy_client.calculate_consumption(firm_id, [agent_id], [this_demand])
120
+ post_consumption_wealth = post_consumption_wealth[0]
121
+ await self.memory.update('consumption_currency', consumption_currency+wealth-post_consumption_wealth)
122
+ goods_demand = await self.memory.get('goods_demand')
123
+ goods_consumption = await self.memory.get('goods_consumption')
124
+ await self.memory.update('goods_demand', goods_demand+this_demand)
125
+ await self.memory.update('goods_consumption', goods_consumption+int((wealth-post_consumption_wealth)//price))
126
+ prompt.format(**data)
127
+ response = await self.llm.atext_request(prompt.to_dialog(), timeout=300)
128
+ evaluation = {'success': True, 'evaluation': f"Consumption: {response}", 'consumed_time': 100}
129
+ return evaluation
130
+
131
+ class EconomyNoneBlock(Block):
132
+ """
133
+ Do nothing
134
+ NoneBlock
135
+ """
136
+ def __init__(self, llm: LLM, memory: Memory):
137
+ super().__init__("NoneBlock", llm, memory)
138
+ self.description = "Do nothing"
139
+ self.forward_times = 0
140
+
141
+ async def forward(self, step, context):
142
+ self.forward_times += 1
143
+ return {
144
+ 'success': True,
145
+ 'evaluation': f'完成执行{step["intention"]}',
146
+ 'consumed_time': 0
147
+ }
148
+
149
+ class EconomyBlock(Block):
150
+ work_block: WorkBlock
151
+ consumption_block: ConsumptionBlock
152
+ none_block: EconomyNoneBlock
153
+
154
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator, economy_client: EconomyClient):
155
+ super().__init__("EconomyBlock", llm, memory, simulator)
156
+ self.economy_client = economy_client
157
+ # 初始化所有块
158
+ self.work_block = WorkBlock(llm, memory, simulator)
159
+ self.consumption_block = ConsumptionBlock(llm, memory, simulator, economy_client)
160
+ self.none_block = EconomyNoneBlock(llm, memory)
161
+ self.trigger_time = 0
162
+ self.token_consumption = 0
163
+ # 初始化调度器
164
+ self.dispatcher = BlockDispatcher(llm)
165
+ # 注册所有块
166
+ self.dispatcher.register_blocks([self.work_block, self.consumption_block, self.none_block])
167
+
168
+ async def forward(self, step, context):
169
+ self.trigger_time += 1
170
+ consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
171
+
172
+ # Select the appropriate sub-block using dispatcher
173
+ # selected_block = await self.dispatcher.dispatch(step)
174
+ selected_block = self.consumption_block
175
+
176
+ # Execute the selected sub-block and get the result
177
+ result = await selected_block.forward(step, context) # type: ignore
178
+
179
+ return result
180
+
181
+ # @trigger_class()
182
+ class MonthPlanBlock(Block):
183
+ """Monthly Planning"""
184
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator, economy_client: EconomyClient):
185
+ super().__init__("MonthPlanBlock", llm, memory, simulator)
186
+ self.economy_client = economy_client
187
+ self.llm_error = 0
188
+ self.last_time_trigger = None
189
+ self.time_diff = month_days * 24 * 60 * 60
190
+ self.forward_times = 0
191
+
192
+ async def month_trigger(self):
193
+ now_time = await self.simulator.get_time()
194
+ if self.last_time_trigger is None or now_time - self.last_time_trigger >= self.time_diff:
195
+ self.last_time_trigger = now_time
196
+ return True
197
+ return False
198
+
199
+ async def forward(self):
200
+ if await self.month_trigger():
201
+ # while True:
202
+ # if self.forward_times == 0:
203
+ # break
204
+ # firm_forward = await self.memory.get('firm_forward')
205
+ # bank_forward = await self.memory.get('bank_forward')
206
+ # nbs_forward = await self.memory.get('nbs_forward')
207
+ # government_forward = await self.memory.get('government_forward')
208
+ # if self.forward_times <= firm_forward and self.forward_times <= bank_forward and self.forward_times <= nbs_forward and self.forward_times <= government_forward:
209
+ # break
210
+ # await asyncio.sleep(1)
211
+ agent_id = await self.memory.get("id")
212
+ firm_id = await self.memory.get('firm_id')
213
+ bank_id = await self.memory.get('bank_id')
214
+ name = await self.memory.get('name')
215
+ age = await self.memory.get('age')
216
+ city = await self.memory.get('city')
217
+ job = await self.memory.get('occupation')
218
+ skill = await self.memory.get('work_skill')
219
+ consumption = await self.memory.get('consumption_currency')
220
+ tax_paid = await self.memory.get('tax_paid')
221
+ price = await self.economy_client.get(firm_id, 'price')
222
+ wealth = await self.economy_client.get(agent_id, 'currency')
223
+ interest_rate = await self.economy_client.get(bank_id, 'interest_rate')
224
+
225
+ problem_prompt = f'''
226
+ You're {name}, a {age}-year-old individual living in {city}. As with all Americans, a portion of your monthly income is taxed by the federal government. This taxation system is tiered, income is taxed cumulatively within defined brackets, combined with a redistributive policy: after collection, the government evenly redistributes the tax revenue back to all citizens, irrespective of their earnings.
227
+ '''
228
+ job_prompt = f'''
229
+ In the previous month, you worked as a(an) {job}. If you continue working this month, your expected hourly income will be ${skill:.2f}.
230
+ '''
231
+ consumption_propensity = await self.memory.get('consumption_propensity')
232
+ if (consumption <= 0) and (consumption_propensity > 0):
233
+ consumption_prompt = f'''
234
+ Besides, you had no consumption due to shortage of goods.
235
+ '''
236
+ else:
237
+ consumption_prompt = f'''
238
+ Besides, your consumption was ${consumption:.2f}.
239
+ '''
240
+ tax_prompt = f'''Your tax deduction amounted to ${tax_paid:.2f}, and the government uses the tax revenue to provide social services to all citizens.'''
241
+ if UBI:
242
+ tax_prompt = f'{tax_prompt} Specifically, the government directly provides ${UBI} per capita in each month.'
243
+ price_prompt = f'''Meanwhile, in the consumption market, the average price of essential goods is now at ${price:.2f}.'''
244
+ job_prompt = prettify_document(job_prompt)
245
+ obs_prompt = f'''
246
+ {problem_prompt} {job_prompt} {consumption_prompt} {tax_prompt} {price_prompt}
247
+ Your current savings account balance is ${wealth:.2f}. Interest rates, as set by your bank, stand at {interest_rate*100:.2f}%.
248
+ Your goal is to maximize your utility by deciding how much to work and how much to consume. Your utility is determined by your consumption, income, saving, social service recieved and leisure time. You will spend the time you do not work on leisure activities.
249
+ With all these factors in play, and considering aspects like your living costs, any future aspirations, and the broader economic trends, how is your willingness to work this month? Furthermore, how would you plan your expenditures on essential goods, keeping in mind good price?
250
+ Please share your decisions in a JSON format as follows:
251
+ {{'work': a value between 0 and 1, indicating the propensity to work,
252
+ 'consumption': a value between 0 and 1, indicating the proportion of all your savings and income you intend to spend on essential goods
253
+ }}
254
+ Any other output words are NOT allowed.
255
+ '''
256
+ obs_prompt = prettify_document(obs_prompt)
257
+ await self.memory.update('dialog_queue', [{'role': 'user', 'content': obs_prompt}], mode='merge')
258
+ dialog_queue = await self.memory.get('dialog_queue')
259
+ content = await self.llm.atext_request(list(dialog_queue), timeout=300)
260
+ await self.memory.update('dialog_queue', [{'role': 'assistant', 'content': content}], mode='merge')
261
+ try:
262
+ propensity_dict = extract_dict_from_string(content)[0]
263
+ work_propensity, consumption_propensity = propensity_dict['work'], propensity_dict['consumption']
264
+ if isinstance(work_propensity, numbers.Number) and isinstance(consumption_propensity, numbers.Number):
265
+ await self.memory.update('work_propensity', work_propensity)
266
+ await self.memory.update('consumption_propensity', consumption_propensity)
267
+ else:
268
+ self.llm_error += 1
269
+ except:
270
+ self.llm_error += 1
271
+
272
+ skill = await self.memory.get('work_skill')
273
+ work_propensity = await self.memory.get('work_propensity')
274
+ consumption_propensity = await self.memory.get('consumption_propensity')
275
+ work_hours = work_propensity * num_labor_hours
276
+ await self.memory.update('income_currency', work_hours * skill)
277
+ wealth = await self.economy_client.get(agent_id, 'currency')
278
+ await self.economy_client.update(agent_id, 'currency', wealth + work_hours * skill)
279
+ await self.economy_client.add_delta_value(firm_id, 'inventory', int(work_hours*productivity_per_labor))
280
+
281
+ income_currency = await self.memory.get('income_currency')
282
+ await self.memory.update('income_currency', income_currency + UBI)
283
+ wealth = await self.economy_client.get(agent_id, 'currency')
284
+ await self.economy_client.update(agent_id, 'currency', wealth + UBI)
285
+
286
+ wealth = await self.economy_client.get(agent_id, 'currency')
287
+ await self.memory.update('to_consumption_currency', consumption_propensity*wealth)
288
+
289
+ await self.memory.update('consumption_currency', 0)
290
+ await self.memory.update('goods_demand', 0)
291
+ await self.memory.update('goods_consumption', 0)
292
+
293
+ if self.forward_times % 3 == 0:
294
+ obs_prompt = f'''
295
+ {problem_prompt} {job_prompt} {consumption_prompt} {tax_prompt} {price_prompt}
296
+ Your current savings account balance is ${wealth:.2f}. Interest rates, as set by your bank, stand at {interest_rate*100:.2f}%.
297
+ Please fill in the following questionnaire:
298
+ Indicate how often you have felt this way during the last week by choosing one of the following options:
299
+ "Rarely" means Rarely or none of the time (less than 1 day),
300
+ "Some" means Some or a little of the time (1-2 days),
301
+ "Occasionally" means Occasionally or a moderate amount of the time (3-4 days),
302
+ "Most" means Most or all of the time (5-7 days).
303
+ Statement 1: I was bothered by things that usually don't bother me.
304
+ Statement 2: I did not feel like eating; my appetite was poor.
305
+ Statement 3: I felt that I could not shake off the blues even with help from my family or friends.
306
+ Statement 4: I felt that I was just as good as other people.
307
+ Statement 5: I had trouble keeping my mind on what I was doing.
308
+ Statement 6: I felt depressed.
309
+ Statement 7: I felt that everything I did was an effort.
310
+ Statement 8: I felt hopeful about the future.
311
+ Statement 9: I thought my life had been a failure.
312
+ Statement 10: I felt fearful.
313
+ Statement 11: My sleep was restless.
314
+ Statement 12: I was happy.
315
+ Statement 13: I talked less than usual.
316
+ Statement 14: I felt lonely.
317
+ Statement 15: People were unfriendly.
318
+ Statement 16: I enjoyed life.
319
+ Statement 17: I had crying spells.
320
+ Statement 18: I felt sad.
321
+ Statement 19: I felt that people disliked me.
322
+ Statement 20: I could not get "going".
323
+ Please response with json format with keys being numbers 1-20 and values being one of "Rarely", "Some", "Occasionally", "Most".
324
+ Any other output words are NOT allowed.
325
+ '''
326
+ obs_prompt = prettify_document(obs_prompt)
327
+ content = await self.llm.atext_request([{'role': 'user', 'content': obs_prompt}], timeout=300)
328
+ inverse_score_items = [3, 8, 12, 16]
329
+ category2score = {'rarely': 0, 'some': 1, 'occasionally': 2, 'most': 3}
330
+ try:
331
+ content = extract_dict_from_string(content)[0]
332
+ for k in content:
333
+ if k in inverse_score_items:
334
+ content[k] = 3 - category2score[content[k].lower()]
335
+ else:
336
+ content[k] = category2score[content[k].lower()]
337
+ depression = sum(list(content.values()))
338
+ await self.memory.update('depression', depression)
339
+ except:
340
+ self.llm_error += 1
341
+
342
+ if UBI and self.forward_times % 12:
343
+ obs_prompt = f'''
344
+ {problem_prompt} {job_prompt} {consumption_prompt} {tax_prompt} {price_prompt}
345
+ Your current savings account balance is ${wealth:.2f}. Interest rates, as set by your bank, stand at {interest_rate*100:.2f}%.
346
+ What's your opinion on the UBI policy, including the advantages and disadvantages?
347
+ '''
348
+ obs_prompt = prettify_document(obs_prompt)
349
+ content = await self.llm.atext_request([{'role': 'user', 'content': obs_prompt}], timeout=300)
350
+ await self.memory.update('ubi_opinion', [content], mode='merge')
351
+ ubi_opinion = await self.memory.get('ubi_opinion')
352
+ with open(f'/data1/linian/SocietySim/socialcity-agent-zoo/ubi_opinions/{agent_id}.pkl', 'wb') as f:
353
+ pkl.dump(ubi_opinion, f)
354
+
355
+ self.forward_times += 1
356
+ await self.memory.update('forward', self.forward_times)
@@ -0,0 +1,258 @@
1
+ from typing import List
2
+ from .dispatcher import BlockDispatcher
3
+ from pycityagent.environment.simulator import Simulator
4
+ from pycityagent.llm.llm import LLM
5
+ from pycityagent.memory.memory import Memory
6
+ from pycityagent.workflow.block import Block
7
+ import random
8
+ import logging
9
+ logger = logging.getLogger("pycityagent")
10
+ from pycityagent.workflow.prompt import FormatPrompt
11
+
12
+ PLACE_TYPE_SELECTION_PROMPT = """
13
+ As an intelligent decision system, please determine the type of place the user needs to visit based on their input requirement.
14
+ User Plan: {plan}
15
+ User requirement: {intention}
16
+ Your output must be a single selection from {poi_category} without any additional text or explanation.
17
+ """
18
+
19
+ PLACE_SECOND_TYPE_SELECTION_PROMPT = """
20
+ As an intelligent decision system, please determine the type of place the user needs to visit based on their input requirement.
21
+ User Plan: {plan}
22
+ User requirement: {intention}
23
+ Your output must be a single selection from {poi_category} without any additional text or explanation.
24
+ """
25
+
26
+ PLACE_ANALYSIS_PROMPT = """
27
+ As an intelligent analysis system, please determine the type of place the user needs to visit based on their input requirement.
28
+ User Plan: {plan}
29
+ User requirement: {intention}
30
+
31
+ Your output must be a single selection from ['home', 'workplace', 'other'] without any additional text or explanation.
32
+ """
33
+
34
+ class PlaceSelectionBlock(Block):
35
+ """
36
+ 选择目的地
37
+ PlaceSelectionBlock
38
+ """
39
+ configurable_fields: List[str] = ["search_radius", "search_limit"]
40
+ default_values = {
41
+ "search_radius": 100000,
42
+ "search_limit": 10
43
+ }
44
+
45
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
46
+ super().__init__("PlaceSelectionBlock", llm, memory, simulator)
47
+ self.description = "Used to select and determine destinations for unknown locations (excluding home and workplace), such as choosing specific shopping malls, restaurants and other places"
48
+ self.typeSelectionPrompt = FormatPrompt(PLACE_TYPE_SELECTION_PROMPT)
49
+ self.secondTypeSelectionPrompt = FormatPrompt(PLACE_SECOND_TYPE_SELECTION_PROMPT)
50
+
51
+ # configurable fields
52
+ self.search_radius = 100000
53
+ self.search_limit = 10
54
+
55
+ async def forward(self, step, context):
56
+ self.typeSelectionPrompt.format(
57
+ plan = context['plan'],
58
+ intention = step['intention'],
59
+ poi_category = list(self.simulator.poi_cate.keys())
60
+ )
61
+ levelOneType = await self.llm.atext_request(self.typeSelectionPrompt.to_dialog()) # type: ignore
62
+ try:
63
+ sub_category = self.simulator.poi_cate[levelOneType]
64
+ except Exception as e:
65
+ logger.warning(f"Wrong type of poi, raw response: {levelOneType}")
66
+ levelOneType = random.choice(list(self.simulator.poi_cate.keys()))
67
+ sub_category = self.simulator.poi_cate[levelOneType]
68
+ self.secondTypeSelectionPrompt.format(
69
+ plan = context['plan'],
70
+ intention = step['intention'],
71
+ poi_category = sub_category
72
+ )
73
+ levelTwoType = await self.llm.atext_request(self.secondTypeSelectionPrompt.to_dialog()) # type: ignore
74
+ center = await self.memory.get('position')
75
+ center = (center['xy_position']['x'], center['xy_position']['y'])
76
+ try:
77
+ pois = self.simulator.map.query_pois(
78
+ center = center,
79
+ category_prefix = levelTwoType,
80
+ radius=self.search_radius,
81
+ limit=self.search_limit
82
+ )
83
+ except Exception as e:
84
+ logger.warning(f"Error querying pois: {e}")
85
+ levelTwoType = random.choice(sub_category)
86
+ pois = self.simulator.map.query_pois(
87
+ center = center,
88
+ category_prefix = levelTwoType,
89
+ radius=self.search_radius,
90
+ limit=self.search_limit
91
+ )
92
+ if len(pois) > 0:
93
+ poi = random.choice(pois)[0]
94
+ nextPlace = (poi['name'], poi['aoi_id'])
95
+ # 将地点信息保存到context中
96
+ context['next_place'] = nextPlace
97
+ # 这里应该添加选择地点的具体逻辑
98
+ return {
99
+ 'success': True,
100
+ 'evaluation': f'Successfully selected the destination: {nextPlace}',
101
+ 'consumed_time': 5
102
+ }
103
+ else:
104
+ simmap = self.simulator.map
105
+ poi = random.choice(list(simmap.pois.values()))
106
+ nextPlace = (poi['name'], poi['aoi_id'])
107
+ # 将地点信息保存到context中
108
+ context['next_place'] = nextPlace
109
+ # 这里应该添加选择地点的具体逻辑
110
+ return {
111
+ 'success': True,
112
+ 'evaluation': f'Successfully selected the destination: {nextPlace}',
113
+ 'consumed_time': 5
114
+ }
115
+
116
+ class MoveBlock(Block):
117
+ """
118
+ 移动操作
119
+ MoveBlock
120
+ """
121
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
122
+ super().__init__("MoveBlock", llm, memory, simulator)
123
+ self.description = "Used to execute specific mobility operations, such as returning home, going to work, or visiting a specific location"
124
+ self.placeAnalysisPrompt = FormatPrompt(PLACE_ANALYSIS_PROMPT)
125
+
126
+ async def forward(self, step, context):
127
+ # 这里应该添加移动的具体逻辑
128
+ agent_id = await self.memory.get("id")
129
+ self.placeAnalysisPrompt.format(
130
+ plan = context['plan'],
131
+ intention = step["intention"]
132
+ )
133
+ response = await self.llm.atext_request(self.placeAnalysisPrompt.to_dialog()) # type: ignore
134
+ if response == 'home':
135
+ # 返回到家
136
+ home = await self.memory.get('home')
137
+ home = home['aoi_position']['aoi_id']
138
+ nowPlace = await self.memory.get('position')
139
+ if 'aoi_position' in nowPlace and nowPlace['aoi_position']['aoi_id'] == home:
140
+ return {
141
+ 'success': True,
142
+ 'evaluation': f'Successfully returned home (already at home)',
143
+ 'from_place': home,
144
+ 'to_place': home,
145
+ 'consumed_time': 0
146
+ }
147
+ await self.simulator.set_aoi_schedules(
148
+ person_id=agent_id,
149
+ target_positions=home,
150
+ )
151
+ return {
152
+ 'success': True,
153
+ 'evaluation': f'Successfully returned home',
154
+ 'from_place': nowPlace['aoi_position']['aoi_id'],
155
+ 'to_place': home,
156
+ 'consumed_time': 45
157
+ }
158
+ elif response == 'workplace':
159
+ # 返回到工作地点
160
+ work = await self.memory.get('work')
161
+ work = work['aoi_position']['aoi_id']
162
+ nowPlace = await self.memory.get('position')
163
+ if 'aoi_position' in nowPlace and nowPlace['aoi_position']['aoi_id'] == work:
164
+ return {
165
+ 'success': True,
166
+ 'evaluation': f'Successfully reached the workplace (already at the workplace)',
167
+ 'from_place': work,
168
+ 'to_place': work,
169
+ 'consumed_time': 0
170
+ }
171
+ await self.simulator.set_aoi_schedules(
172
+ person_id=agent_id,
173
+ target_positions=work,
174
+ )
175
+ return {
176
+ 'success': True,
177
+ 'evaluation': f'Successfully reached the workplace',
178
+ 'from_place': nowPlace['aoi_position']['aoi_id'],
179
+ 'to_place': work,
180
+ 'consumed_time': 45
181
+ }
182
+ else:
183
+ # 移动到其他地点
184
+ next_place = context.get('next_place', None)
185
+ nowPlace = await self.memory.get('position')
186
+ if next_place != None:
187
+ await self.simulator.set_aoi_schedules(
188
+ person_id=agent_id,
189
+ target_positions=next_place[1],
190
+ )
191
+ else:
192
+ while True:
193
+ r_aoi = random.choice(list(self.simulator.map.aois.values()))
194
+ if len(r_aoi['poi_ids']) > 0:
195
+ r_poi = random.choice(r_aoi['poi_ids'])
196
+ break
197
+ poi = self.simulator.map.pois[r_poi]
198
+ next_place = (poi['name'], poi['aoi_id'])
199
+ await self.simulator.set_aoi_schedules(
200
+ person_id=agent_id,
201
+ target_positions=next_place[1],
202
+ )
203
+ return {
204
+ 'success': True,
205
+ 'evaluation': f'Successfully reached the destination: {next_place}',
206
+ 'from_place': nowPlace['aoi_position']['aoi_id'],
207
+ 'to_place': next_place[1],
208
+ 'consumed_time': 45
209
+ }
210
+
211
+ class MobilityNoneBlock(Block):
212
+ """
213
+ 空操作
214
+ MobilityNoneBlock
215
+ """
216
+ def __init__(self, llm: LLM, memory: Memory):
217
+ super().__init__("MobilityNoneBlock", llm, memory)
218
+ self.description = "Used to handle other cases"
219
+
220
+ async def forward(self, step, context):
221
+ return {
222
+ 'success': True,
223
+ 'evaluation': f'Finished executing {step["intention"]}',
224
+ 'consumed_time': 0
225
+ }
226
+
227
+ class MobilityBlock(Block):
228
+ place_selection_block: PlaceSelectionBlock
229
+ move_block: MoveBlock
230
+ mobility_none_block: MobilityNoneBlock
231
+
232
+ def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
233
+ super().__init__("MobilityBlock", llm, memory, simulator)
234
+ # 初始化所有块
235
+ self.place_selection_block = PlaceSelectionBlock(llm, memory, simulator)
236
+ self.move_block = MoveBlock(llm, memory, simulator)
237
+ self.mobility_none_block = MobilityNoneBlock(llm, memory)
238
+ self.trigger_time = 0
239
+ self.token_consumption = 0
240
+ # 初始化调度器
241
+ self.dispatcher = BlockDispatcher(llm)
242
+ # 注册所有块
243
+ self.dispatcher.register_blocks([self.place_selection_block, self.move_block, self.mobility_none_block])
244
+
245
+ async def forward(self, step, context):
246
+ self.trigger_time += 1
247
+ consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
248
+
249
+ # Select the appropriate sub-block using dispatcher
250
+ selected_block = await self.dispatcher.dispatch(step)
251
+
252
+ # Execute the selected sub-block and get the result
253
+ result = await selected_block.forward(step, context) # type: ignore
254
+
255
+ consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
256
+ self.token_consumption += consumption_end - consumption_start
257
+
258
+ return result