pycityagent 2.0.0a72__cp310-cp310-macosx_11_0_arm64.whl → 2.0.0a74__cp310-cp310-macosx_11_0_arm64.whl

pycityagent/environment/simulator.py

@@ -5,16 +5,17 @@ import logging
 import os
 from datetime import datetime, timedelta
 import time
-from typing import Any, Optional, Union, cast
+from typing import Optional, Union, cast
 
-from mosstool.type import TripMode
-from mosstool.util.format_converter import coll2pb
+from mosstool.util.format_converter import coll2pb, dict2pb
 from pycitydata.map import Map as SimMap
 from pycityproto.city.map.v2 import map_pb2 as map_pb2
 from pycityproto.city.person.v2 import person_pb2 as person_pb2
 from pycityproto.city.person.v2 import person_service_pb2 as person_service
 from pymongo import MongoClient
+import ray
 from shapely.geometry import Point
+from mosstool.type import TripMode
 
 from .sim import CityClient, ControlSimEnv
 from .utils.const import *
@@ -25,6 +26,42 @@ __all__ = [
     "Simulator",
 ]
 
+@ray.remote
+class CityMap:
+    def __init__(self, mongo_uri: str, mongo_db: str, mongo_coll: str, cache_dir: str):
+        self.map = SimMap(
+            mongo_uri=mongo_uri,
+            mongo_db=mongo_db,
+            mongo_coll=mongo_coll,
+            cache_dir=cache_dir,
+        )
+
+    def get_aoi(self, aoi_id: Optional[int] = None):
+        if aoi_id is None:
+            return list(self.map.aois.values())
+        else:
+            return self.map.aois[aoi_id]
+        
+    def get_poi(self, poi_id: Optional[int] = None):
+        if poi_id is None:
+            return list(self.map.pois.values())
+        else:
+            return self.map.pois[poi_id]
+
+    def query_pois(self, **kwargs):
+        return self.map.query_pois(**kwargs)
+        
+    def get_poi_cate(self):
+        return self.poi_cate
+    
+    def get_map(self):
+        return self.map
+    
+    def get_map_header(self):
+        return self.map.header
+
+    def get_projector(self):
+        return self.map.header["projection"]
 
 class Simulator:
     """
@@ -35,7 +72,7 @@ class Simulator:
         - It reads parameters from a configuration dictionary, initializes map data, and starts or connects to a simulation server as needed.
     """
 
-    def __init__(self, config: dict, secure: bool = False) -> None:
+    def __init__(self, config: dict, create_map: bool = False) -> None:
         self.config = config
         """
         - 模拟器配置
@@ -66,7 +103,8 @@ class Simulator:
                 self._sim_env = sim_env = ControlSimEnv(
                     task_name=config["simulator"].get("task", "citysim"),
                     map_file=_map_pb_path,
-                    start_step=config["simulator"].get("start_step", 0),
+                    max_day=config["simulator"].get("max_day", 1000),
+                    start_step=config["simulator"].get("start_step", 28800),
                     total_step=2147000000,
                     log_dir=config["simulator"].get("log_dir", "./log"),
                     min_step_time=config["simulator"].get("min_step_time", 1000),
@@ -88,16 +126,14 @@ class Simulator:
             logger.warning(
                 "No simulator config found, no simulator client will be used"
             )
-        self.map = SimMap(
-            mongo_uri=_mongo_uri,
-            mongo_db=_mongo_db,
-            mongo_coll=_mongo_coll,
-            cache_dir=_map_cache_dir,
-        )
+        self._map = None
         """
         - 模拟器地图对象
         - Simulator map object
         """
+        if create_map:
+            self._map = CityMap.remote(_mongo_uri, _mongo_db, _mongo_coll, _map_cache_dir)
+            self._create_poi_id_2_aoi_id()
 
         self.time: int = 0
         """
@@ -109,25 +145,41 @@ class Simulator:
         self.map_y_gap = None
         self._bbox: tuple[float, float, float, float] = (-1, -1, -1, -1)
         self._lock = asyncio.Lock()
-        # poi id dict
-        self.poi_id_2_aoi_id: dict[int, int] = {
-            poi["id"]: poi["aoi_id"] for _, poi in self.map.pois.items()
-        }
         self._environment_prompt:dict[str, str] = {}
         self._log_list = []
 
+    def set_map(self, map: CityMap):
+        self._map = map
+        self._create_poi_id_2_aoi_id()
+
+    def _create_poi_id_2_aoi_id(self):
+        pois = ray.get(self._map.get_poi.remote())
+        self.poi_id_2_aoi_id: dict[int, int] = {
+            poi["id"]: poi["aoi_id"] for poi in pois
+        }
+
+    @property
+    def map(self):
+        return self._map
+
     def get_log_list(self):
         return self._log_list
     
     def clear_log_list(self):
         self._log_list = []
 
+    def get_poi_cate(self):
+        return self.poi_cate
+
     @property
     def environment(self) -> dict[str, str]:
         """
         Get the current state of environment variables.
         """
         return self._environment_prompt
+    
+    def get_server_addr(self):
+        return self.server_addr
 
     def set_environment(self, environment: dict[str, str]):
         """
@@ -224,11 +276,11 @@ class Simulator:
         categories: list[str] = []
         if center is None:
             center = (0, 0)
-        _pois: list[dict] = self.map.query_pois(  # type:ignore
+        _pois: list[dict] = ray.get(self.map.query_pois.remote(  # type:ignore
             center=center,
             radius=radius,
             return_distance=False,
-        )
+        ))
         for poi in _pois:
             catg = poi["category"]
             categories.append(catg.split("|")[-1])
@@ -367,13 +419,12 @@ class Simulator:
         self._log_list.append(log)
         return person
 
-    async def add_person(self, person: Any) -> dict:
+    async def add_person(self, dict_person: dict) -> dict:
         """
         Add a new person to the simulation.
 
         - **Args**:
-            - `person` (`Any`): The person object to add. If it's an instance of `person_pb2.Person`,
-              it will be wrapped in an `AddPersonRequest`. Otherwise, `person` is expected to already be a valid request object.
+            - `dict_person` (`dict`): The person object to add.
 
         - **Returns**:
             - `Dict`: Response from adding the person.
@@ -384,6 +435,7 @@ class Simulator:
             "start_time": start_time,
             "consumption": 0
         }
+        person = dict2pb(dict_person, person_pb2.Person())
         if isinstance(person, person_pb2.Person):
             req = person_service.AddPersonRequest(person=person)
         else:
@@ -411,7 +463,7 @@ class Simulator:
               A list of AOI or POI IDs or tuples of (AOI ID, POI ID) that the person will visit.
             - `departure_times` (`Optional[List[float]]`): Departure times for each trip in the schedule.
               If not provided, current time will be used for all trips.
-            - `modes` (`Optional[List[TripMode]]`): Travel modes for each trip.
+            - `modes` (`Optional[List[int]]`): Travel modes for each trip.
               Defaults to `TRIP_MODE_DRIVE_ONLY` if not specified.
         """
         start_time = time.time()
@@ -560,11 +612,11 @@ class Simulator:
                 transformed_poi_type += self.poi_cate[t]
         poi_type_set = set(transformed_poi_type)
         # 获取半径内的poi
-        _pois: list[dict] = self.map.query_pois(  # type:ignore
+        _pois: list[dict] = ray.get(self.map.query_pois.remote(  # type:ignore
             center=center,
             radius=radius,
             return_distance=False,
-        )
+        ))
         # 过滤掉不满足类别前缀的poi
         pois = []
         for poi in _pois:

pycityagent/llm/llm.py

@@ -232,111 +232,112 @@ class LLM:
         """
         start_time = time.time()
         log = {"request_time": start_time}
-        if (
-            self.config.text["request_type"] == "openai"
-            or self.config.text["request_type"] == "deepseek"
-            or self.config.text["request_type"] == "qwen"
-        ):
-            for attempt in range(retries):
-                try:
-                    client = self._get_next_client()
-                    response = await client.chat.completions.create(
-                        model=self.config.text["model"],
-                        messages=dialog,
-                        temperature=temperature,
-                        max_tokens=max_tokens,
-                        top_p=top_p,
-                        frequency_penalty=frequency_penalty,  # type: ignore
-                        presence_penalty=presence_penalty,  # type: ignore
-                        stream=False,
-                        timeout=timeout,
-                        tools=tools,
-                        tool_choice=tool_choice,
-                    )  # type: ignore
-                    self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["request_number"] += 1
-                    end_time = time.time()
-                    log["consumption"] = end_time - start_time
-                    log["input_tokens"] = response.usage.prompt_tokens
-                    log["output_tokens"] = response.usage.completion_tokens
-                    self._log_list.append(log)
-                    if tools and response.choices[0].message.tool_calls:
-                        return json.loads(
-                            response.choices[0]
-                            .message.tool_calls[0]
-                            .function.arguments
-                        )
-                    else:
-                        return response.choices[0].message.content
-                except APIConnectionError as e:
-                    print("API connection error:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-                except OpenAIError as e:
-                    if hasattr(e, "http_status"):
-                        print(f"HTTP status code: {e.http_status}")  # type: ignore
-                    else:
-                        print("An error occurred:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-        elif self.config.text["request_type"] == "zhipuai":
-            for attempt in range(retries):
-                try:
-                    client = self._get_next_client()
-                    response = client.chat.asyncCompletions.create(  # type: ignore
-                        model=self.config.text["model"],
-                        messages=dialog,
-                        temperature=temperature,
-                        top_p=top_p,
-                        timeout=timeout,
-                        tools=tools,
-                        tool_choice=tool_choice,
-                    )
-                    task_id = response.id
-                    task_status = ""
-                    get_cnt = 0
-                    cnt_threshold = int(timeout / 0.5)
-                    while (
-                        task_status != "SUCCESS"
-                        and task_status != "FAILED"
-                        and get_cnt <= cnt_threshold
-                    ):
-                        result_response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)  # type: ignore
-                        task_status = result_response.task_status
-                        await asyncio.sleep(0.5)
-                        get_cnt += 1
-                    if task_status != "SUCCESS":
-                        raise Exception(f"Task failed with status: {task_status}")
-
-                    self._client_usage[self._current_client_index]["prompt_tokens"] += result_response.usage.prompt_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["completion_tokens"] += result_response.usage.completion_tokens  # type: ignore
-                    self._client_usage[self._current_client_index]["request_number"] += 1
-                    end_time = time.time()
-                    log["used_time"] = end_time - start_time
-                    log["token_consumption"] = result_response.usage.prompt_tokens + result_response.usage.completion_tokens
-                    self._log_list.append(log)
-                    if tools and result_response.choices[0].message.tool_calls:  # type: ignore
-                        return json.loads(
-                            result_response.choices[0]  # type: ignore
-                            .message.tool_calls[0]
-                            .function.arguments
+        async with self.semaphore:
+            if (
+                self.config.text["request_type"] == "openai"
+                or self.config.text["request_type"] == "deepseek"
+                or self.config.text["request_type"] == "qwen"
+            ):
+                for attempt in range(retries):
+                    try:
+                        client = self._get_next_client()
+                        response = await client.chat.completions.create(
+                            model=self.config.text["model"],
+                            messages=dialog,
+                            temperature=temperature,
+                            max_tokens=max_tokens,
+                            top_p=top_p,
+                            frequency_penalty=frequency_penalty,  # type: ignore
+                            presence_penalty=presence_penalty,  # type: ignore
+                            stream=False,
+                            timeout=timeout,
+                            tools=tools,
+                            tool_choice=tool_choice,
+                        )  # type: ignore
+                        self._client_usage[self._current_client_index]["prompt_tokens"] += response.usage.prompt_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["completion_tokens"] += response.usage.completion_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["request_number"] += 1
+                        end_time = time.time()
+                        log["consumption"] = end_time - start_time
+                        log["input_tokens"] = response.usage.prompt_tokens
+                        log["output_tokens"] = response.usage.completion_tokens
+                        self._log_list.append(log)
+                        if tools and response.choices[0].message.tool_calls:
+                            return json.loads(
+                                response.choices[0]
+                                .message.tool_calls[0]
+                                .function.arguments
+                            )
+                        else:
+                            return response.choices[0].message.content
+                    except APIConnectionError as e:
+                        print("API connection error:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+                    except OpenAIError as e:
+                        if hasattr(e, "http_status"):
+                            print(f"HTTP status code: {e.http_status}")  # type: ignore
+                        else:
+                            print("An error occurred:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+            elif self.config.text["request_type"] == "zhipuai":
+                for attempt in range(retries):
+                    try:
+                        client = self._get_next_client()
+                        response = client.chat.asyncCompletions.create(  # type: ignore
+                            model=self.config.text["model"],
+                            messages=dialog,
+                            temperature=temperature,
+                            top_p=top_p,
+                            timeout=timeout,
+                            tools=tools,
+                            tool_choice=tool_choice,
                         )
-                    else:
-                        return result_response.choices[0].message.content  # type: ignore
-                except APIConnectionError as e:
-                    print("API connection error:", e)
-                    if attempt < retries - 1:
-                        await asyncio.sleep(2**attempt)
-                    else:
-                        raise e
-        else:
-            print("ERROR: Wrong Config")
-            return "wrong config"
+                        task_id = response.id
+                        task_status = ""
+                        get_cnt = 0
+                        cnt_threshold = int(timeout / 0.5)
+                        while (
+                            task_status != "SUCCESS"
+                            and task_status != "FAILED"
+                            and get_cnt <= cnt_threshold
+                        ):
+                            result_response = client.chat.asyncCompletions.retrieve_completion_result(id=task_id)  # type: ignore
+                            task_status = result_response.task_status
+                            await asyncio.sleep(0.5)
+                            get_cnt += 1
+                        if task_status != "SUCCESS":
+                            raise Exception(f"Task failed with status: {task_status}")
+
+                        self._client_usage[self._current_client_index]["prompt_tokens"] += result_response.usage.prompt_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["completion_tokens"] += result_response.usage.completion_tokens  # type: ignore
+                        self._client_usage[self._current_client_index]["request_number"] += 1
+                        end_time = time.time()
+                        log["used_time"] = end_time - start_time
+                        log["token_consumption"] = result_response.usage.prompt_tokens + result_response.usage.completion_tokens
+                        self._log_list.append(log)
+                        if tools and result_response.choices[0].message.tool_calls:  # type: ignore
+                            return json.loads(
+                                result_response.choices[0]  # type: ignore
+                                .message.tool_calls[0]
+                                .function.arguments
+                            )
+                        else:
+                            return result_response.choices[0].message.content  # type: ignore
+                    except APIConnectionError as e:
+                        print("API connection error:", e)
+                        if attempt < retries - 1:
+                            await asyncio.sleep(2**attempt)
+                        else:
+                            raise e
+            else:
+                print("ERROR: Wrong Config")
+                return "wrong config"
 
     async def img_understand(
         self, img_path: Union[str, list[str]], prompt: Optional[str] = None

pycityagent/memory/memory.py

@@ -341,7 +341,7 @@ class StreamMemory:
         """获取指定ID的记忆"""
         memories = [memory for memory in self._memories if memory.id in memory_ids]
         sorted_results = sorted(memories, key=lambda x: (x.day, x.t), reverse=True)
-        return self.format_memory(sorted_results)
+        return await self.format_memory(sorted_results)
 
     async def search(
         self,

pycityagent/simulation/agentgroup.py

@@ -36,6 +36,7 @@ class AgentGroup:
         number_of_agents: Union[int, list[int]],
         memory_config_function_group: dict[type[Agent], Callable],
         config: dict,
+        map_ref: ray.ObjectRef,
         exp_name: str,
         exp_id: Union[str, UUID],
         enable_avro: bool,
@@ -47,8 +48,8 @@ class AgentGroup:
         embedding_model: Embeddings,
         logging_level: int,
         agent_config_file: Optional[dict[type[Agent], str]] = None,
-        environment: Optional[dict[str, str]] = None,
         llm_semaphore: int = 200,
+        environment: Optional[dict] = None,
     ):
         """
         Represents a group of agents that can be deployed in a Ray distributed environment.
@@ -64,6 +65,7 @@ class AgentGroup:
             - `number_of_agents` (Union[int, List[int]]): Number of instances to create for each agent class.
             - `memory_config_function_group` (Dict[Type[Agent], Callable]): Functions to configure memory for each agent type.
             - `config` (dict): Configuration settings for the agent group.
+            - `map_ref` (ray.ObjectRef): Reference to the map object.
             - `exp_name` (str): Name of the experiment.
             - `exp_id` (str | UUID): Identifier for the experiment.
             - `enable_avro` (bool): Flag to enable AVRO file support.
@@ -143,10 +145,10 @@ class AgentGroup:
         self.llm.set_semaphore(llm_semaphore)
 
         # prepare Simulator
-        logger.info(f"-----Creating Simulator in AgentGroup {self._uuid} ...")
+        logger.info(f"-----Initializing Simulator in AgentGroup {self._uuid} ...")
         self.simulator = Simulator(config["simulator_request"])
-        self.projector = pyproj.Proj(self.simulator.map.header["projection"])
-        self.simulator.set_environment(environment)  # type:ignore
+        self.simulator.set_map(map_ref)
+        self.projector = pyproj.Proj(ray.get(self.simulator.map.get_projector.remote()))
         # prepare Economy client
         logger.info(f"-----Creating Economy client in AgentGroup {self._uuid} ...")
         self.economy_client = EconomyClient(
@@ -430,16 +432,6 @@ class AgentGroup:
         agent = self.id2agent[target_agent_uuid]
         await agent.status.update(target_key, content)
 
-    async def update_environment(self, key: str, value: str):
-        """
-        Updates the environment with a given key-value pair.
-
-        - **Args**:
-            - `key` (str): The key to update in the environment.
-            - `value` (str): The value to set for the specified key.
-        """
-        self.simulator.update_environment(key, value)
-
     async def message_dispatch(self):
         """
         Dispatches messages received via MQTT to the appropriate agents.
@@ -788,12 +780,13 @@ class AgentGroup:
         """
         try:
             tasks = [agent.run() for agent in self.agents]
-            await asyncio.gather(*tasks)
+            agent_time_log =await asyncio.gather(*tasks)
             simulator_log = self.simulator.get_log_list() + self.economy_client.get_log_list()
             group_logs = {
                 "llm_log": self.llm.get_log_list(),
                 "mqtt_log": ray.get(self.messager.get_log_list.remote()),
-                "simulator_log": simulator_log
+                "simulator_log": simulator_log,
+                "agent_time_log": agent_time_log
             }
             self.llm.clear_log_list()
             self.messager.clear_log_list.remote()

pycityagent/simulation/simulation.py

@@ -130,12 +130,16 @@ class AgentSimulation:
         _simulator_config = config["simulator_request"].get("simulator", {})
         if "server" in _simulator_config:
             raise ValueError(f"Passing Traffic Simulation address is not supported!")
-        self._simulator = Simulator(config["simulator_request"])
+        simulator = Simulator(config["simulator_request"], create_map=True)
+        self._simulator = simulator
+        self._map_ref = self._simulator.map
+        server_addr = self._simulator.get_server_addr()
+        config["simulator_request"]["simulator"]["server"] = server_addr
         self._economy_client = EconomyClient(
             config["simulator_request"]["simulator"]["server"]
         )
         if enable_institution:
-            self._economy_addr = economy_addr = self._simulator.server_addr
+            self._economy_addr = economy_addr = server_addr
             if economy_addr is None:
                 raise ValueError(
                     f"`simulator` not provided in `simulator_request`, thus unable to activate economy!"
@@ -382,6 +386,7 @@ class AgentSimulation:
         llm_log_lists = []
         mqtt_log_lists = []
         simulator_log_lists = []
+        agent_time_log_lists = []
         for step in config["workflow"]:
             logger.info(
                 f"Running step: type: {step['type']} - description: {step.get('description', 'no description')}"
@@ -389,17 +394,19 @@ class AgentSimulation:
             if step["type"] not in ["run", "step", "interview", "survey", "intervene", "pause", "resume", "function"]:
                 raise ValueError(f"Invalid step type: {step['type']}")
             if step["type"] == "run":
-                llm_log_list, mqtt_log_list, simulator_log_list = await simulation.run(step.get("days", 1))
+                llm_log_list, mqtt_log_list, simulator_log_list, agent_time_log_list = await simulation.run(step.get("days", 1))
                 llm_log_lists.extend(llm_log_list)
                 mqtt_log_lists.extend(mqtt_log_list)
                 simulator_log_lists.extend(simulator_log_list)
+                agent_time_log_lists.extend(agent_time_log_list)
             elif step["type"] == "step":
                 times = step.get("times", 1)
                 for _ in range(times):
-                    llm_log_list, mqtt_log_list, simulator_log_list = await simulation.step()
+                    llm_log_list, mqtt_log_list, simulator_log_list, agent_time_log_list = await simulation.step()
                     llm_log_lists.extend(llm_log_list)
                     mqtt_log_lists.extend(mqtt_log_list)
                     simulator_log_lists.extend(simulator_log_list)
+                    agent_time_log_lists.extend(agent_time_log_list)
             elif step["type"] == "pause":
                 await simulation.pause_simulator()
             elif step["type"] == "resume":
@@ -407,7 +414,7 @@ class AgentSimulation:
             else:
                 await step["func"](simulation)
         logger.info("Simulation finished")
-        return llm_log_lists, mqtt_log_lists, simulator_log_lists
+        return llm_log_lists, mqtt_log_lists, simulator_log_lists, agent_time_log_lists
     
     @property
     def enable_avro(
@@ -742,6 +749,7 @@ class AgentSimulation:
                 number_of_agents,
                 memory_config_function_group,
                 self.config,
+                self._map_ref,
                 self.exp_name,
                 self.exp_id,
                 self.enable_avro,
@@ -753,8 +761,8 @@ class AgentSimulation:
                 embedding_model,
                 self.logging_level,
                 config_file,
-                environment,
                 llm_semaphore,
+                environment,
             )
             creation_tasks.append((group_name, group))
 
@@ -1018,7 +1026,7 @@ class AgentSimulation:
 
             # step
             simulator_day = await self._simulator.get_simulator_day()
-            simulator_time = int(await self._simulator.get_time())
+            simulator_time = int(await self._simulator.get_simulator_second_from_start_of_day())
             logger.info(
                 f"Start simulation day {simulator_day} at {simulator_time}, step {self._total_steps}"
             )
@@ -1029,14 +1037,15 @@ class AgentSimulation:
             llm_log_list = []
             mqtt_log_list = []
             simulator_log_list = []
+            agent_time_log_list = []
             for log_messages_group in log_messages_groups:
                 llm_log_list.extend(log_messages_group['llm_log'])
                 mqtt_log_list.extend(log_messages_group['mqtt_log'])
                 simulator_log_list.extend(log_messages_group['simulator_log'])
-            
+                agent_time_log_list.extend(log_messages_group['agent_time_log'])
             # save
             simulator_day = await self._simulator.get_simulator_day()
-            simulator_time = int(await self._simulator.get_time())
+            simulator_time = int(await self._simulator.get_simulator_second_from_start_of_day())
             save_tasks = []
             for group in self._groups.values():
                 save_tasks.append(group.save.remote(simulator_day, simulator_time))
@@ -1050,7 +1059,7 @@ class AgentSimulation:
                 ]
                 await self.extract_metric(to_excute_metric)
 
-            return llm_log_list, mqtt_log_list, simulator_log_list
+            return llm_log_list, mqtt_log_list, simulator_log_list, agent_time_log_list
         except Exception as e:
             import traceback
 
@@ -1081,6 +1090,7 @@ class AgentSimulation:
         llm_log_lists = []
         mqtt_log_lists = []
         simulator_log_lists = []
+        agent_time_log_lists = []
         try:
             self._exp_info["num_day"] += day
             await self._update_exp_status(1)  # 更新状态为运行中
@@ -1098,10 +1108,11 @@ class AgentSimulation:
                     current_time = await self._simulator.get_time()
                     if current_time >= end_time:  # type:ignore
                         break
-                    llm_log_list, mqtt_log_list, simulator_log_list = await self.step()
+                    llm_log_list, mqtt_log_list, simulator_log_list, agent_time_log_list = await self.step()
                     llm_log_lists.extend(llm_log_list)
                     mqtt_log_lists.extend(mqtt_log_list)
                     simulator_log_lists.extend(simulator_log_list)
+                    agent_time_log_lists.extend(agent_time_log_list)
             finally:
                 # 设置停止事件
                 stop_event.set()
@@ -1110,7 +1121,7 @@ class AgentSimulation:
 
             # 运行成功后更新状态
             await self._update_exp_status(2)
-            return llm_log_lists, mqtt_log_lists, simulator_log_lists
+            return llm_log_lists, mqtt_log_lists, simulator_log_lists, agent_time_log_lists
         except Exception as e:
             error_msg = f"模拟器运行错误: {str(e)}"
             logger.error(error_msg)

pycityagent/tools/tool.py

@@ -5,6 +5,7 @@ from collections.abc import Callable, Sequence
 from typing import Any, Optional, Union
 
 from mlflow.entities import Metric
+import ray
 
 from ..agent import Agent
 from ..environment import AoiService, PersonService