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

pycityagent/cityagent/blocks/mobility_block.py

@@ -1,6 +1,6 @@
 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
@@ -9,8 +9,11 @@ import numpy as np
 from operator import itemgetter
 import random
 import logging
-logger = logging.getLogger("pycityagent")
 from pycityagent.workflow.prompt import FormatPrompt
+from .dispatcher import BlockDispatcher
+
+logger = logging.getLogger("pycityagent")
+
 
 PLACE_TYPE_SELECTION_PROMPT = """
 As an intelligent decision system, please determine the type of place the user needs to visit based on their input requirement.
@@ -45,6 +48,7 @@ Please analyze how these emotions would affect travel willingness and return onl
 
 Return only the integer number without any additional text or explanation."""
 
+
 def gravity_model(pois):
     N = len(pois)
     pois_Dis = {
@@ -85,9 +89,7 @@ def gravity_model(pois):
                 # 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))
-                )
+                distanceProb.append(1 / (math.sqrt(distance)))
                 iflt10k = False
                 break
 
@@ -110,32 +112,36 @@ def gravity_model(pois):
     ]
     return final
 
+
 class PlaceSelectionBlock(Block):
     """
     Select destination
     PlaceSelectionBlock
     """
+
     configurable_fields: List[str] = ["search_limit"]
-    default_values = {
-        "search_limit": 10000
-    }
+    default_values = {"search_limit": 10000}
 
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
-        super().__init__("PlaceSelectionBlock", llm=llm, memory=memory, simulator=simulator)
+        super().__init__(
+            "PlaceSelectionBlock", llm=llm, memory=memory, simulator=simulator
+        )
         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"
         self.typeSelectionPrompt = FormatPrompt(PLACE_TYPE_SELECTION_PROMPT)
-        self.secondTypeSelectionPrompt = FormatPrompt(PLACE_SECOND_TYPE_SELECTION_PROMPT)
+        self.secondTypeSelectionPrompt = FormatPrompt(
+            PLACE_SECOND_TYPE_SELECTION_PROMPT
+        )
         self.radiusPrompt = FormatPrompt(RADIUS_PROMPT)
         # configurable fields
         self.search_limit = 10000
 
     async def forward(self, step, context):
         self.typeSelectionPrompt.format(
-            plan = context['plan'],
-            intention = step['intention'],
-            poi_category = list(self.simulator.poi_cate.keys())
+            plan=context["plan"],
+            intention=step["intention"],
+            poi_category=list(self.simulator.poi_cate.keys()),
         )
-        levelOneType = await self.llm.atext_request(self.typeSelectionPrompt.to_dialog()) # type: ignore
+        levelOneType = await self.llm.atext_request(self.typeSelectionPrompt.to_dialog())  # type: ignore
         try:
             sub_category = self.simulator.poi_cate[levelOneType]
         except Exception as e:
@@ -143,35 +149,33 @@ class PlaceSelectionBlock(Block):
             levelOneType = random.choice(list(self.simulator.poi_cate.keys()))
             sub_category = self.simulator.poi_cate[levelOneType]
         self.secondTypeSelectionPrompt.format(
-            plan = context['plan'],
-            intention = step['intention'],
-            poi_category = sub_category
+            plan=context["plan"], intention=step["intention"], poi_category=sub_category
         )
-        levelTwoType = await self.llm.atext_request(self.secondTypeSelectionPrompt.to_dialog()) # type: ignore
-        center = await self.memory.status.get('position')
-        center = (center['xy_position']['x'], center['xy_position']['y'])
+        levelTwoType = await self.llm.atext_request(self.secondTypeSelectionPrompt.to_dialog())  # type: ignore
+        center = await self.memory.status.get("position")
+        center = (center["xy_position"]["x"], center["xy_position"]["y"])
         self.radiusPrompt.format(
             emotion_types=await self.memory.status.get("emotion_types"),
             thought=await self.memory.status.get("thought"),
             weather=self.simulator.sence("weather"),
-            temperature=self.simulator.sence("temperature")
+            temperature=self.simulator.sence("temperature"),
         )
-        radius = int(await self.llm.atext_request(self.radiusPrompt.to_dialog())) # type: ignore
+        radius = int(await self.llm.atext_request(self.radiusPrompt.to_dialog()))  # type: ignore
         try:
             pois = self.simulator.map.query_pois(
-                center = center,
-                category_prefix = levelTwoType,
+                center=center,
+                category_prefix=levelTwoType,
                 radius=radius,
-                limit=self.search_limit
+                limit=self.search_limit,
             )
         except Exception as e:
             logger.warning(f"Error querying pois: {e}")
             levelTwoType = random.choice(sub_category)
             pois = self.simulator.map.query_pois(
-                center = center,
-                category_prefix = levelTwoType,
+                center=center,
+                category_prefix=levelTwoType,
                 radius=radius,
-                limit=self.search_limit
+                limit=self.search_limit,
             )
         if len(pois) > 0:
             pois = gravity_model(pois)
@@ -183,33 +187,39 @@ class PlaceSelectionBlock(Block):
             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}")
+            context["next_place"] = nextPlace
+            node_id = await self.memory.stream.add_mobility(
+                description=f"For {step['intention']}, I selected the destination: {nextPlace}"
+            )
             return {
-                'success': True,
-                'evaluation': f'Successfully selected the destination: {nextPlace}',
-                'consumed_time': 5,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f"Successfully selected the destination: {nextPlace}",
+                "consumed_time": 5,
+                "node_id": node_id,
             }
         else:
             simmap = self.simulator.map
             poi = random.choice(list(simmap.pois.values()))
-            nextPlace = (poi['name'], poi['aoi_id'])
+            nextPlace = (poi["name"], poi["aoi_id"])
             # 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}")
+            context["next_place"] = nextPlace
+            node_id = await self.memory.stream.add_mobility(
+                description=f"For {step['intention']}, I selected the destination: {nextPlace}"
+            )
             return {
-                'success': True,
-                'evaluation': f'Successfully selected the destination: {nextPlace}',
-                'consumed_time': 5,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f"Successfully selected the destination: {nextPlace}",
+                "consumed_time": 5,
+                "node_id": node_id,
             }
 
+
 class MoveBlock(Block):
     """
     Execute mobility operations
     MoveBlock
     """
+
     def __init__(self, llm: LLM, memory: Memory, simulator: Simulator):
         super().__init__("MoveBlock", llm=llm, memory=memory, simulator=simulator)
         self.description = "Used to execute specific mobility operations, such as returning home, going to work, or visiting a specific location"
@@ -219,68 +229,79 @@ class MoveBlock(Block):
         # 这里应该添加移动的具体逻辑
         agent_id = await self.memory.status.get("id")
         self.placeAnalysisPrompt.format(
-            plan = context['plan'],
-            intention = step["intention"]
+            plan=context["plan"], intention=step["intention"]
         )
         number_poi_visited = await self.memory.status.get("number_poi_visited")
         number_poi_visited += 1
         await self.memory.status.update("number_poi_visited", number_poi_visited)
-        response = await self.llm.atext_request(self.placeAnalysisPrompt.to_dialog()) # type: ignore
-        if response == 'home':
+        response = await self.llm.atext_request(self.placeAnalysisPrompt.to_dialog())  # type: ignore
+        if response == "home":
             # 返回到家
-            home = await self.memory.status.get('home')
-            home = home['aoi_position']['aoi_id']
-            nowPlace = await self.memory.status.get('position')
-            node_id = await self.memory.stream.add_mobility(description=f"I returned home")
-            if 'aoi_position' in nowPlace and nowPlace['aoi_position']['aoi_id'] == home:
+            home = await self.memory.status.get("home")
+            home = home["aoi_position"]["aoi_id"]
+            nowPlace = await self.memory.status.get("position")
+            node_id = await self.memory.stream.add_mobility(
+                description=f"I returned home"
+            )
+            if (
+                "aoi_position" in nowPlace
+                and nowPlace["aoi_position"]["aoi_id"] == home
+            ):
                 return {
-                    'success': True,
-                    'evaluation': f'Successfully returned home (already at home)',
-                    'to_place': home,
-                    'consumed_time': 0,
-                    'node_id': node_id
+                    "success": True,
+                    "evaluation": f"Successfully returned home (already at home)",
+                    "to_place": home,
+                    "consumed_time": 0,
+                    "node_id": node_id,
                 }
             await self.simulator.set_aoi_schedules(
                 person_id=agent_id,
                 target_positions=home,
             )
             return {
-                'success': True,
-                'evaluation': f'Successfully returned home',
-                'to_place': home,
-                'consumed_time': 45,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f"Successfully returned home",
+                "to_place": home,
+                "consumed_time": 45,
+                "node_id": node_id,
             }
-        elif response == 'workplace':
+        elif response == "workplace":
             # 返回到工作地点
-            work = await self.memory.status.get('work')
-            work = work['aoi_position']['aoi_id']
-            nowPlace = await self.memory.status.get('position')
-            node_id = await self.memory.stream.add_mobility(description=f"I went to my workplace")
-            if 'aoi_position' in nowPlace and nowPlace['aoi_position']['aoi_id'] == work:
+            work = await self.memory.status.get("work")
+            work = work["aoi_position"]["aoi_id"]
+            nowPlace = await self.memory.status.get("position")
+            node_id = await self.memory.stream.add_mobility(
+                description=f"I went to my workplace"
+            )
+            if (
+                "aoi_position" in nowPlace
+                and nowPlace["aoi_position"]["aoi_id"] == work
+            ):
                 return {
-                    'success': True,
-                    'evaluation': f'Successfully reached the workplace (already at the workplace)',
-                    'to_place': work,
-                    'consumed_time': 0,
-                    'node_id': node_id
+                    "success": True,
+                    "evaluation": f"Successfully reached the workplace (already at the workplace)",
+                    "to_place": work,
+                    "consumed_time": 0,
+                    "node_id": node_id,
                 }
             await self.simulator.set_aoi_schedules(
                 person_id=agent_id,
                 target_positions=work,
             )
             return {
-                'success': True,
-                'evaluation': f'Successfully reached the workplace',
-                'to_place': work,
-                'consumed_time': 45,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f"Successfully reached the workplace",
+                "to_place": work,
+                "consumed_time": 45,
+                "node_id": node_id,
             }
         else:
             # 移动到其他地点
-            next_place = context.get('next_place', None)
-            nowPlace = await self.memory.status.get('position')
-            node_id = await self.memory.stream.add_mobility(description=f"I went to {next_place}")
+            next_place = context.get("next_place", None)
+            nowPlace = await self.memory.status.get("position")
+            node_id = await self.memory.stream.add_mobility(
+                description=f"I went to {next_place}"
+            )
             if next_place != None:
                 await self.simulator.set_aoi_schedules(
                     person_id=agent_id,
@@ -289,41 +310,46 @@ class MoveBlock(Block):
             else:
                 while True:
                     r_aoi = random.choice(list(self.simulator.map.aois.values()))
-                    if len(r_aoi['poi_ids']) > 0:
-                        r_poi = random.choice(r_aoi['poi_ids'])
+                    if len(r_aoi["poi_ids"]) > 0:
+                        r_poi = random.choice(r_aoi["poi_ids"])
                         break
                 poi = self.simulator.map.pois[r_poi]
-                next_place = (poi['name'], poi['aoi_id'])
+                next_place = (poi["name"], poi["aoi_id"])
                 await self.simulator.set_aoi_schedules(
                     person_id=agent_id,
                     target_positions=next_place[1],
                 )
             return {
-                'success': True,
-                'evaluation': f'Successfully reached the destination: {next_place}',
-                'to_place': next_place[1],
-                'consumed_time': 45,
-                'node_id': node_id
+                "success": True,
+                "evaluation": f"Successfully reached the destination: {next_place}",
+                "to_place": next_place[1],
+                "consumed_time": 45,
+                "node_id": node_id,
             }
-    
+
+
 class MobilityNoneBlock(Block):
     """
     空操作
     MobilityNoneBlock
     """
+
     def __init__(self, llm: LLM, memory: Memory):
         super().__init__("MobilityNoneBlock", llm=llm, memory=memory)
         self.description = "Used to handle other cases"
 
     async def forward(self, step, context):
-        node_id = await self.memory.stream.add_mobility(description=f"I finished {step['intention']}")
+        node_id = await self.memory.stream.add_mobility(
+            description=f"I finished {step['intention']}"
+        )
         return {
-            'success': True,
-            'evaluation': f'Finished executing {step["intention"]}',
-            'consumed_time': 0,
-            'node_id': node_id
+            "success": True,
+            "evaluation": f'Finished executing {step["intention"]}',
+            "consumed_time": 0,
+            "node_id": node_id,
         }
 
+
 class MobilityBlock(Block):
     place_selection_block: PlaceSelectionBlock
     move_block: MoveBlock
@@ -340,19 +366,23 @@ class MobilityBlock(Block):
         # 初始化调度器
         self.dispatcher = BlockDispatcher(llm)
         # 注册所有块
-        self.dispatcher.register_blocks([self.place_selection_block, self.move_block, self.mobility_none_block])
+        self.dispatcher.register_blocks(
+            [self.place_selection_block, self.move_block, self.mobility_none_block]
+        )
 
-    async def forward(self, step, context):       
+    async def forward(self, step, context):
         self.trigger_time += 1
-        consumption_start = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
+        consumption_start = (
+            self.llm.prompt_tokens_used + self.llm.completion_tokens_used
+        )
 
         # Select the appropriate sub-block using dispatcher
         selected_block = await self.dispatcher.dispatch(step)
-        
+
         # Execute the selected sub-block and get the result
-        result = await selected_block.forward(step, context) # type: ignore
+        result = await selected_block.forward(step, context)  # type: ignore
 
         consumption_end = self.llm.prompt_tokens_used + self.llm.completion_tokens_used
         self.token_consumption += consumption_end - consumption_start
-        
-        return result
+
+        return result