pycityagent 2.0.0a9__tar.gz → 2.0.0a10__tar.gz

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: pycityagent
-Version: 2.0.0a9
+Version: 2.0.0a10
 Summary: LLM-based城市环境agent构建库
 License: MIT
 Author: Yuwei Yan

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/agent.py

@@ -2,10 +2,12 @@
 
 from abc import ABC, abstractmethod
 import asyncio
+from uuid import UUID
 from copy import deepcopy
 from datetime import datetime
 from enum import Enum
 import logging
+import random
 from typing import Dict, List, Optional
 
 from pycityagent.environment.sim.person_service import PersonService
@@ -80,10 +82,6 @@ class Agent(ABC):
         del state["_llm_client"]
         return state
 
-    async def bind_to_simulator(self):
-        await self._bind_to_simulator()
-        await self._bind_to_economy()
-
     def set_messager(self, messager: Messager):
         """
         Set the messager of the agent.
@@ -114,7 +112,7 @@ class Agent(ABC):
         """
         self._memory = memory
 
-    def set_exp_id(self, exp_id: str):
+    def set_exp_id(self, exp_id: str|UUID):
         """
         Set the exp_id of the agent.
         """
@@ -141,7 +139,7 @@ class Agent(ABC):
                 "bike_attribute",
             }
             simulator = self._simulator
-            memory = self._memory
+            memory = self.memory
             person_id = await memory.get("id")
             # ATTENTION:模拟器分配的id从0开始
             if person_id >= 0:
@@ -178,11 +176,11 @@ class Agent(ABC):
                     await self._economy_client.remove_agents([self._agent_id])
                 except:
                     pass
-                person_id = await self._memory.get("id")
+                person_id = await self.memory.get("id")
                 await self._economy_client.add_agents(
                     {
                         "id": person_id,
-                        "currency": await self._memory.get("currency"),
+                        "currency": await self.memory.get("currency"),
                     }
                 )
                 self._has_bound_to_economy = True
@@ -300,6 +298,9 @@ class Agent(ABC):
     async def forward(self) -> None:
         """智能体行为逻辑"""
         raise NotImplementedError
+    
+    async def handle_gather_message(self, payload: str):
+        raise NotImplementedError
 
     async def run(self) -> None:
         """
@@ -334,6 +335,81 @@ class CitizenAgent(Agent):
             memory,
         )
 
+    async def bind_to_simulator(self):
+        await self._bind_to_simulator()
+        await self._bind_to_economy()
+
+    async def _bind_to_simulator(self):
+        """
+        Bind Agent to Simulator
+
+        Args:
+            person_template (dict, optional): The person template in dict format. Defaults to PersonService.default_dict_person().
+        """
+        if self._simulator is None:
+            logging.warning("Simulator is not set")
+            return
+        if not self._has_bound_to_simulator:
+            FROM_MEMORY_KEYS = {
+                "attribute",
+                "home",
+                "work",
+                "vehicle_attribute",
+                "bus_attribute",
+                "pedestrian_attribute",
+                "bike_attribute",
+            }
+            simulator = self._simulator
+            memory = self._memory
+            person_id = await memory.get("id")
+            # ATTENTION:模拟器分配的id从0开始
+            if person_id >= 0:
+                await simulator.get_person(person_id)
+                logging.debug(f"Binding to Person `{person_id}` already in Simulator")
+            else:
+                dict_person = deepcopy(self._person_template)
+                for _key in FROM_MEMORY_KEYS:
+                    try:
+                        _value = await memory.get(_key)
+                        if _value:
+                            dict_person[_key] = _value
+                    except KeyError as e:
+                        continue
+                resp = await simulator.add_person(
+                    dict2pb(dict_person, person_pb2.Person())
+                )
+                person_id = resp["person_id"]
+                await memory.update("id", person_id, protect_llm_read_only_fields=False)
+                logging.debug(
+                    f"Binding to Person `{person_id}` just added to Simulator"
+                )
+                # 防止模拟器还没有到prepare阶段导致get_person出错
+            self._has_bound_to_simulator = True
+            self._agent_id = person_id
+
+    async def _bind_to_economy(self):
+        if self._economy_client is None:
+            logging.warning("Economy client is not set")
+            return
+        if not self._has_bound_to_economy:
+            if self._has_bound_to_simulator:
+                try:
+                    await self._economy_client.remove_agents([self._agent_id])
+                except:
+                    pass
+                person_id = await self._memory.get("id")
+                await self._economy_client.add_agents(
+                    {
+                        "id": person_id,
+                        "currency": await self._memory.get("currency"),
+                    }
+                )
+                self._has_bound_to_economy = True
+            else:
+                logging.debug(
+                    f"Binding to Economy before binding to Simulator, skip binding to Economy Simulator"
+                )
+
     async def handle_gather_message(self, payload: str):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容
@@ -366,6 +442,57 @@ class InstitutionAgent(Agent):
             memory,
         )
 
+    async def bind_to_simulator(self):
+        await self._bind_to_economy()
+
+    async def _bind_to_economy(self):
+        if self._economy_client is None:
+            logging.debug("Economy client is not set")
+            return
+        if not self._has_bound_to_economy:
+            # TODO: More general id generation
+            _id = random.randint(100000, 999999)
+            self._agent_id = _id
+            await self._memory.update("id", _id, protect_llm_read_only_fields=False)
+            try:
+                await self._economy_client.remove_agents([self._agent_id])
+            except:
+                pass
+            try:
+                id = await self._memory.get("id")
+                type = await self._memory.get("type")
+                nominal_gdp = await self._memory.get("nominal_gdp")
+                real_gdp = await self._memory.get("real_gdp")
+                unemployment = await self._memory.get("unemployment")
+                wages = await self._memory.get("wages")
+                prices = await self._memory.get("prices")
+                inventory = await self._memory.get("inventory")
+                price = await self._memory.get("price")
+                currency = await self._memory.get("currency")
+                interest_rate = await self._memory.get("interest_rate")
+                bracket_cutoffs = await self._memory.get("bracket_cutoffs")
+                bracket_rates = await self._memory.get("bracket_rates")
+                await self._economy_client.add_orgs(
+                    {
+                        "id": id,
+                        "type": type,
+                        "currency": currency,
+                        "nominal_gdp": nominal_gdp,
+                        "real_gdp": real_gdp,
+                        "unemployment": unemployment,
+                        "wages": wages,
+                        "prices": prices,
+                        "inventory": inventory,
+                        "price": price,
+                        "interest_rate": interest_rate,
+                        "bracket_cutoffs": bracket_cutoffs,
+                        "bracket_rates": bracket_rates,
+                    }
+                )
+            except Exception as e:
+                logging.error(f"Failed to bind to Economy: {e}")
+            self._has_bound_to_economy = True
+
     async def handle_gather_message(self, payload: str):
         """处理收到的消息，识别发送者"""
         # 从消息中解析发送者 ID 和消息内容

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/environment/simulator.py

@@ -3,7 +3,6 @@
 import asyncio
 import logging
 import os
-from collections import defaultdict
 from collections.abc import Sequence
 from datetime import datetime, timedelta
 from typing import Any, Optional, Tuple, Union, cast

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/llm/llm.py

@@ -15,7 +15,6 @@ from dashscope import ImageSynthesis
 from PIL import Image
 from io import BytesIO
 from typing import Any, Optional, Union, List, Dict
-import aiohttp
 from .llmconfig import *
 from .utils import *
 
@@ -45,7 +44,13 @@ class LLM:
         elif self.config.text["request_type"] == "deepseek":
             self._aclient = AsyncOpenAI(
                 api_key=self.config.text["api_key"],
-                base_url="https://api.deepseek.com/beta",
+                base_url="https://api.deepseek.com/v1",
+                timeout=300,
+            )
+        elif self.config.text["request_type"] == "qwen":
+            self._aclient = AsyncOpenAI(
+                api_key=self.config.text["api_key"],
+                base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
                 timeout=300,
             )
         elif self.config.text["request_type"] == "zhipuai":
@@ -102,106 +107,6 @@ Token Usage:
             "ratio": rate,
         }
 
-    def text_request(
-        self,
-        dialog: Any,
-        temperature: float = 1,
-        max_tokens: Optional[int] = None,
-        top_p: Optional[float] = None,
-        frequency_penalty: Optional[float] = None,
-        presence_penalty: Optional[float] = None,
-        tools: Optional[List[Dict[str, Any]]] = None,
-        tool_choice: Optional[Dict[str, Any]] = None,
-    ) -> Optional[str]:
-        """
-        文本相关请求
-        Text request
-
-        Args:
-        - dialog (list[dict]): 标准的LLM文本dialog. The standard text LLM dialog
-        - temperature (float): default 1, used in openai
-        - max_tokens (int): default None, used in openai
-        - top_p (float): default None, used in openai
-        - frequency_penalty (float): default None, used in openai
-        - presence_penalty (float): default None, used in openai
-
-        Returns:
-        - (str): the response content
-        """
-        if "api_base" in self.config.text.keys():
-            api_base = self.config.text["api_base"]
-        else:
-            api_base = None
-        if self.config.text["request_type"] == "openai":
-            client = OpenAI(
-                api_key=self.config.text["api_key"],
-                base_url=api_base,
-            )
-            response = 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,
-                presence_penalty=presence_penalty,
-                tools=tools,
-                tool_choice=tool_choice,
-            )
-            self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
-            self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
-            self.request_number += 1
-            if tools != None:
-                return response.tool_calls[0].function.arguments
-            else:
-                return response.choices[0].message.content
-        elif self.config.text["request_type"] == "qwen":
-            response = dashscope.Generation.call(
-                model=self.config.text["model"],
-                api_key=self.config.text["api_key"],
-                messages=dialog,
-                result_format="message",
-            )
-            if response.status_code == HTTPStatus.OK:  # type: ignore
-                return response.output.choices[0]["message"]["content"]  # type: ignore
-            else:
-                return "Error: {}, {}".format(response.status_code, response.message)  # type: ignore
-        elif self.config.text["request_type"] == "deepseek":
-            client = OpenAI(
-                api_key=self.config.text["api_key"],
-                base_url="https://api.deepseek.com/beta",
-            )
-            response = 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,
-                presence_penalty=presence_penalty,
-                stream=False,
-            )
-            self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
-            self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
-            self.request_number += 1
-            return response.choices[0].message.content
-        elif self.config.text["request_type"] == "zhipuai":
-            client = ZhipuAI(api_key=self.config.text["api_key"])
-            response = client.chat.completions.create(
-                model=self.config.text["model"],
-                messages=dialog,
-                temperature=temperature,
-                top_p=top_p,
-                stream=False,
-            )
-            self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
-            self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
-            self.request_number += 1
-            return response.choices[0].message.content  # type: ignore
-        else:
-            print("ERROR: Wrong Config")
-            return "wrong config"
-
     async def atext_request(
         self,
         dialog: Any,
@@ -221,6 +126,7 @@ Token Usage:
         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:
@@ -242,8 +148,12 @@ Token Usage:
                             self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
                             self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
                             self.request_number += 1
-                            if tools != None:
-                                return response.tool_calls[0].function.arguments
+                            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
                     else:
@@ -263,8 +173,12 @@ Token Usage:
                         self.prompt_tokens_used += response.usage.prompt_tokens  # type: ignore
                         self.completion_tokens_used += response.usage.completion_tokens  # type: ignore
                         self.request_number += 1
-                        if tools != None:
-                            return response.tool_calls[0].function.arguments
+                        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:
@@ -327,25 +241,6 @@ Token Usage:
                         await asyncio.sleep(2**attempt)
                     else:
                         raise e
-        elif self.config.text["request_type"] == "qwen":
-            async with aiohttp.ClientSession() as session:
-                api_url = "https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation"
-                headers = {
-                    "Content-Type": "application/json",
-                    "Authorization": f"{self.config.text['api_key']}",
-                }
-                payload = {
-                    "model": self.config.text["model"],
-                    "input": {"messages": dialog},
-                }
-                async with session.post(api_url, json=payload, headers=headers) as resp:
-                    response_json = await resp.json()
-                    if "code" in response_json.keys():
-                        raise Exception(
-                            f"Error: {response_json['code']}, {response_json['message']}"
-                        )
-                    else:
-                        return response_json["output"]["text"]
         else:
             print("ERROR: Wrong Config")
             return "wrong config"

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/memory/memory_base.py

@@ -6,7 +6,7 @@ import asyncio
 import logging
 import time
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, Optional, Sequence, Union
 
 from .const import *
 

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/memory/profile.py

@@ -3,7 +3,7 @@ Agent Profile
 """
 
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/memory/self_define.py

@@ -3,7 +3,7 @@ Self Define Data
 """
 
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/memory/state.py

@@ -3,7 +3,7 @@ Agent State
 """
 
 from copy import deepcopy
-from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Callable, Dict, Optional, Sequence, Union, cast
 
 from ..utils.decorators import lock_decorator
 from .const import *

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/memory/utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Sequence, Union
 
 from .memory_base import MemoryUnit
 

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/message/messager.py

@@ -7,7 +7,7 @@ from aiomqtt import Client
 
 class Messager:
     def __init__(
-        self, hostname, port=1883, username=None, password=None, timeout=math.inf
+        self, hostname:str, port:int=1883, username=None, password=None, timeout=math.inf
     ):
         self.client = Client(
             hostname, port=port, username=username, password=password, timeout=timeout

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/simulation/agentgroup.py

@@ -1,17 +1,18 @@
 import asyncio
 import logging
 import ray
+from uuid import UUID
 from pycityagent.agent import Agent
 from pycityagent.economy.econ_client import EconomyClient
 from pycityagent.environment.simulator import Simulator
 from pycityagent.llm.llm import LLM
 from pycityagent.llm.llmconfig import LLMConfig
 from pycityagent.message import Messager
-
+from typing import Any
 
 @ray.remote
 class AgentGroup:
-    def __init__(self, agents: list[Agent], config: dict, exp_id: str):
+    def __init__(self, agents: list[Agent], config: dict, exp_id: str|UUID):
         self.agents = agents
         self.config = config
         self.exp_id = exp_id
@@ -45,7 +46,8 @@ class AgentGroup:
             agent.set_exp_id(self.exp_id)
             agent.set_llm_client(self.llm)
             agent.set_simulator(self.simulator)
-            agent.set_economy_client(self.economy_client)
+            if self.economy_client is not None:
+                agent.set_economy_client(self.economy_client)
             agent.set_messager(self.messager)
 
     async def init_agents(self):
@@ -69,7 +71,7 @@ class AgentGroup:
             results[agent._agent_id] = await agent.memory.get(content)
         return results
 
-    async def update(self, target_agent_id: str, target_key: str, content: any):
+    async def update(self, target_agent_id: str, target_key: str, content: Any):
         agent = self.id2agent[target_agent_id]
         await agent.memory.update(target_key, content)
 
@@ -123,14 +125,15 @@ class AgentGroup:
         try:
             # 获取开始时间
             start_time = await self.simulator.get_time()
+            assert type(start_time)==int
             # 计算结束时间（秒）
             end_time = start_time + day * 24 * 3600  # 将天数转换为秒
 
             while True:
                 current_time = await self.simulator.get_time()
+                assert type(current_time)==int
                 if current_time >= end_time:
                     break
-
                 await self.step()
 
         except Exception as e:

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pycityagent/simulation/simulation.py

@@ -4,7 +4,7 @@ import logging
 import uuid
 from datetime import datetime
 import random
-from typing import Dict, List, Optional, Callable, Union
+from typing import Dict, List, Optional, Callable, Union,Any
 from mosstool.map._map_util.const import AOI_START_ID
 
 from pycityagent.memory.memory import Memory
@@ -41,7 +41,7 @@ class AgentSimulation:
         self.config = config
         self.agent_prefix = agent_prefix
         self._agents: Dict[str, Agent] = {}
-        self._groups: Dict[str, AgentGroup] = {}
+        self._groups: Dict[str, AgentGroup] = {} # type:ignore
         self._interview_manager = InterviewManager()
         self._interview_lock = asyncio.Lock()
         self._start_time = datetime.now()
@@ -50,14 +50,14 @@ class AgentSimulation:
         self._loop = asyncio.get_event_loop()
         self._blocked_agents: List[str] = []  # 新增：持续阻塞的智能体列表
         self._survey_manager = SurveyManager()
-        self._agentid2group: Dict[str, AgentGroup] = {}
+        self._agentid2group: Dict[str, AgentGroup] = {}# type:ignore
         self._agent_ids: List[str] = []
 
     async def init_agents(
         self,
         agent_count: Union[int, list[int]],
         group_size: int = 1000,
-        memory_config_func: Union[Callable, list[Callable]] = None,
+        memory_config_func: Optional[Union[Callable, list[Callable]]] = None,
     ) -> None:
         """初始化智能体
 
@@ -139,13 +139,13 @@ class AgentSimulation:
                 self._agentid2group[agent_id] = group
 
     async def gather(self, content: str):
-        """收集所有智能体的ID"""
+        """收集智能体的特定信息"""
         gather_tasks = []
         for group in self._groups.values():
             gather_tasks.append(group.gather.remote(content))
         return await asyncio.gather(*gather_tasks)
 
-    async def update(self, target_agent_id: str, target_key: str, content: any):
+    async def update(self, target_agent_id: str, target_key: str, content: Any):
         """更新指定智能体的记忆"""
         group = self._agentid2group[target_agent_id]
         await group.update.remote(target_agent_id, target_key, content)

{pycityagent-2.0.0a9 → pycityagent-2.0.0a10}/pyproject.toml

@@ -3,7 +3,7 @@ in-project = true
 
 [tool.poetry]
 name = "pycityagent"
-version = "2.0.0a9"
+version = "2.0.0a10"
 description = "LLM-based城市环境agent构建库"
 authors = ["Yuwei Yan <pinkgranite86@gmail.com>","Junbo Yan <yanjb20thu@gmali.com>"]
 license = "MIT"