pycityagent 1.1.9__py3-none-any.whl → 1.1.11__py3-none-any.whl

pycityagent/simulator.py

@@ -1,6 +1,6 @@
 """Simulator: 城市模拟器类及其定义"""
 
-from typing import Optional, Union
+from typing import Optional, Union, Tuple
 from datetime import datetime, timedelta
 import asyncio
 from pycitysim import *
@@ -58,7 +58,7 @@ class Simulator:
 
         self.map = map.Map(
             mongo_uri = "mongodb://sim:FiblabSim1001@mgo.db.fiblab.tech:8635/",
-            mongo_db = "srt",
+            mongo_db = config['map_request']['mongo_db'],
             mongo_coll = config['map_request']['mongo_coll'],
             cache_dir = config['map_request']['cache_dir'],
         )
@@ -67,6 +67,12 @@ class Simulator:
         - Simulator map object
         """
 
+        self.pois_matrix: dict[str, list[list[list]]] = {}
+        """
+        pois的基于区块的划分——方便快速粗略地查询poi
+        通过Simulator.set_pois_matrix()初始化
+        """
+
         self.routing = RoutingClient(self.config['route_request']['server'])
         """
         - 导航服务grpc客户端
@@ -78,6 +84,20 @@ class Simulator:
         - 模拟城市当前时间
         - The current time of simulator
         """
+        self.poi_cate = {'10': 'eat', 
+                         '13': 'shopping', 
+                         '18': 'sports',
+                         '22': 'excursion',
+                         '16': 'entertainment',
+                         '20': 'medical tratment',
+                         '14': 'trivialities',
+                         '25': 'financial',
+                         '12': 'government and political services',
+                         '23': 'cultural institutions',
+                         '28': 'residence'}
+        self.map_x_gap = None
+        self.map_y_gap = None
+        self.poi_matrix_centers = []
 
     # * Agent相关
     def FindAgentsByArea(self, req: dict, status=None):
@@ -105,7 +125,7 @@ class Simulator:
             resp.motions = motions
             return resp
 
-    async def GetCitizenAgent(self, name:str, id:int):
+    async def GetCitizenAgent(self, name:str, id:int) -> CitizenAgent:
         """
         获取agent
         Get Agent
@@ -119,8 +139,9 @@ class Simulator:
         """
         await self.GetTime()
         resp = await self._client.person_service.GetPerson({"person_id": id})
-        base = resp['base']
-        motion = resp['motion']
+        print(f"Agent {id}: {resp}")
+        base = resp['person']['base']
+        motion = resp['person']['motion']
         agent = CitizenAgent(
             name, 
             self.config['simulator']['server'], 
@@ -132,7 +153,7 @@ class Simulator:
         agent.set_streetview_config(self.config['streetview_request'])
         return agent
 
-    async def GetFuncAgent(self, id:int, name:str):
+    async def GetFuncAgent(self, name:str) -> FuncAgent:
         """
         获取一个Func Agent模板
 
@@ -145,7 +166,6 @@ class Simulator:
         """
         agent = FuncAgent(
                     name,
-                    id+10000000,
                     self.config['simulator']['server'],
                     simulator=self
                 )
@@ -160,6 +180,128 @@ class Simulator:
         """
         print("Not Implemented Yet")
         pass
+
+    def set_poi_matrix(self, map:dict=None, row_number:int=12, col_number:int=10, radius:int=10000):
+        """
+        初始化pois_matrix
+
+        Args:
+        - map (dict): 地图参数
+            east, west, north, south
+        - row_number (int): 行数
+        - col_number (int): 列数
+        - radius (int): 搜索半径, 单位m
+        """
+        if map == None:
+            self.matrix_map = self.map
+        else:
+            self.matrix_map = map
+        print(f"Building Poi searching matrix, Row_number: {row_number}, Col_number: {col_number}, Radius: {radius}m")
+        self.map_x_gap = (self.matrix_map.header['east'] - self.matrix_map.header['west']) / col_number
+        self.map_y_gap = (self.matrix_map.header['north'] - self.matrix_map.header['south']) / row_number
+        for i in range(row_number):
+            self.poi_matrix_centers.append([])
+            for j in range(col_number):
+                center_x = self.matrix_map.header['west'] + self.map_x_gap*j + self.map_x_gap/2
+                center_y = self.matrix_map.header['south'] + self.map_y_gap*i + self.map_y_gap/2
+                self.poi_matrix_centers[i].append((center_x, center_y))
+        
+        for pre in self.poi_cate.keys():
+            print(f"Building matrix for Poi category: {pre}")
+            self.pois_matrix[pre] = []
+            for row_centers in self.poi_matrix_centers:
+                row_pois = []
+                for center in row_centers:
+                    pois = self.map.query_pois(center=center, radius=radius, category_prefix=pre)
+                    row_pois.append(pois)
+                self.pois_matrix[pre].append(row_pois)
+        print("Finished")
+
+    def get_pois_from_matrix(self, center:Tuple[float, float], prefix:str):
+        """
+        从poi搜索矩阵中快速获取poi
+
+        Args:
+        - center (Tuple[float, float]): 位置信息
+        - prefix (str): 类型前缀
+        """
+        if self.map_x_gap == None:
+            print("Set Poi Matrix first")
+            return
+        elif prefix not in self.poi_cate.keys():
+            print(f"Wrong prefix, only {self.poi_cate.keys()} is usable")
+            return
+        elif center[0] > self.matrix_map.header['east'] or center[0] < self.matrix_map.header['west'] or center[1] > self.matrix_map.header['north'] or center[1] < self.matrix_map.header['south']:
+            print("Wrong center")
+            return
+        
+        # 矩阵匹配
+        rows = int((center[1]-self.matrix_map.header['south'])/self.map_y_gap)
+        cols = int((center[0]-self.matrix_map.header['west'])/self.map_x_gap)
+        pois = self.pois_matrix[prefix][rows][cols]
+        return pois
+    
+    def get_cat_from_pois(self, pois:list):
+        cat_2_num = {}
+        for poi in pois:
+            cate = poi['category'][:2]
+            if cate not in self.poi_cate.keys():
+                continue
+            if cate in cat_2_num.keys():
+                cat_2_num[cate] += 1
+            else:
+                cat_2_num[cate] = 1
+        max_cat = ""
+        max_num = 0
+        for key in cat_2_num.keys():
+            if cat_2_num[key] > max_num:
+                max_num = cat_2_num[key]
+                max_cat = self.poi_cate[key]
+        return max_cat
+        
+    def get_poi_matrix_in_rec(self, center:Tuple[float, float], radius:int=2500, rows:int=5, cols:int=5):
+        """
+        获取以center为中心的正方形区域内的poi集合
+
+        Args:
+        - center (Tuple[float, float]): 中心位置点
+        - radius (int): 半径
+        """
+        north = center[1] + radius
+        south = center[1] - radius
+        west = center[0] - radius
+        east = center[0] + radius
+        x_gap = (east-west)/cols
+        y_gap = (north-south)/rows
+        matrix = []
+        for i in range(rows):
+            matrix.append([])
+            for j in range(cols):
+                matrix[i].append([])
+        pois = []
+        for poi in self.map.pois.values():
+            x = poi['position']['x']
+            y = poi['position']['y']
+            if x > west and x < east and y > south and y < north:
+                row_index = int((y-south)/x_gap)
+                col_index = int((x-west)/y_gap)
+                matrix[row_index][col_index].append(poi)
+        matrix_type = []
+        for i in range(rows):
+            for j in range(cols):
+                matrix_type.append(self.get_cat_from_pois(matrix[i][j]))
+        poi_total_number = []
+        poi_type_number = []
+        for i in range(rows):
+            for j in range(cols):
+                poi_total_number.append(len(matrix[i][j]))
+                number = 0
+                for poi in matrix[i][j]:
+                    if poi['category'][:2] in self.poi_cate.keys() and self.poi_cate[poi['category'][:2]] == matrix_type[i*cols+j]:
+                        number += 1
+                poi_type_number.append(number)
+                
+        return matrix, matrix_type, poi_total_number, poi_type_number
         
     async def GetTime(self, format_time:bool=False, format:Optional[str]="%H:%M:%S") -> Union[int, str]:
         """

pycityagent/urbanllm/urbanllm.py

@@ -1,6 +1,8 @@
 """UrbanLLM: 智能能力类及其定义"""
 
-from openai import OpenAI
+from openai import OpenAI, AsyncOpenAI, APIConnectionError, OpenAIError
+import openai
+import asyncio
 from http import HTTPStatus
 import dashscope
 import requests
@@ -9,10 +11,55 @@ from PIL import Image
 from io import BytesIO
 from typing import Union
 import base64
+import aiohttp
+import re
 
 def encode_image(image_path):
   with open(image_path, "rb") as image_file:
     return base64.b64encode(image_file.read()).decode('utf-8')
+  
+class VarSet:
+    def __init__(self, variable_batch, variable_description, K):
+        self.var_batch = variable_batch
+        self.var_des = variable_description
+        self.var_size = len(list(variable_description.keys()))
+        self.batch_size = K
+    
+    def set_variables(self, variables:list[dict]):
+        if len(variables) <= 0 or len(variables) != self.batch_size or (len(variables[0]) != self.var_size):
+            print(f"Your input variables mis-match the size of current Variable Set: Batch size-{self.batch_size}, Variable size-{self.var_size}")
+            return None
+        for i in range(len(self.var_batch)):
+            for var, _ in self.var_batch[i].items():
+                self.var_batch[i][var] = variables[i][var]
+
+    def get_variables(self):
+        return self.var_batch
+    
+    def get_var_des(self):
+        var_des = ""
+        for var, desc in self.var_des.items():
+            var_des += f"batch[x][{var}] -- {desc}\n"
+        return var_des
+
+class PromptBatch:
+    def __init__(self, prompt_batch:str, variable_batch, vasriable_description, K):
+        self.prompt_raw = prompt_batch
+        self.prompt_set = None
+        self.batch_size = K
+        self.var_set = VarSet(variable_batch, vasriable_description, K)
+    
+    def set_variables(self, variables:list[dict]):
+        self.var_set.set_variables(variables)
+        self.prompt_set = self.prompt_raw.format(batch=self.get_variables())
+
+    def get_variables(self):
+        return self.var_set.var_batch
+
+    def get_prompt(self):
+        if self.prompt_set == None:
+            return self.prompt_raw
+        return self.prompt_set
 
 class LLMConfig:
     """
@@ -37,6 +84,53 @@ class UrbanLLM:
     """
     def __init__(self, config: LLMConfig) -> None:
         self.config = config
+        self.prompt_tokens_used = 0
+        self.completion_tokens_used = 0
+        self.request_number = 0
+        self.semaphore = None
+        self._aclient = AsyncOpenAI(api_key=self.config.text['api_key'], timeout=300)
+
+    def set_semaphore(self, number_of_coroutine:int):
+        self.semaphore = asyncio.Semaphore(number_of_coroutine)
+
+    def clear_semaphore(self):
+        self.semaphore = None
+
+    def clear_used(self):
+        """
+        clear the storage of used tokens to start a new log message
+        Only support OpenAI category API right now, including OpenAI, Deepseek
+        """
+        self.prompt_tokens_used = 0
+        self.completion_tokens_used = 0
+        self.request_number = 0
+
+    def show_consumption(self, input_price:float=None, output_price:float=None):
+        """
+        if you give the input and output price of using model, this function will also calculate the consumption for you
+        """
+        total_token = self.prompt_tokens_used + self.completion_tokens_used
+        if self.completion_tokens_used != 0:
+            rate = self.prompt_tokens_used/self.completion_tokens_used
+        else:
+            rate = 'nan'
+        if self.request_number != 0:
+            TcA = total_token/self.request_number
+        else:
+            TcA = 'nan'
+        out = f"""Request Number: {self.request_number}
+Token Usage:
+    - Total tokens: {total_token}
+    - Prompt tokens: {self.prompt_tokens_used}
+    - Completion tokens: {self.completion_tokens_used}
+    - Token per request: {TcA}
+    - Prompt:Completion ratio: {rate}:1"""
+        if input_price != None and output_price != None:
+            consumption = self.prompt_tokens_used/1000000*input_price + self.completion_tokens_used/1000000*output_price
+            out += f"\n    - Cost Estimation: {consumption}"
+        print(out)
+        return {"total": total_token, "prompt": self.prompt_tokens_used, "completion": self.completion_tokens_used, "ratio": rate}
+
 
     def text_request(self, dialog:list[dict], temperature:float=1, max_tokens:int=None, top_p:float=None, frequency_penalty:float=None, presence_penalty:float=None) -> str:
         """
@@ -72,6 +166,9 @@ class UrbanLLM:
                 frequency_penalty=frequency_penalty,
                 presence_penalty=presence_penalty
             )
+            self.prompt_tokens_used += response.usage.prompt_tokens
+            self.completion_tokens_used += response.usage.completion_tokens
+            self.request_number += 1
             return response.choices[0].message.content
         elif self.config.text['request_type'] == 'qwen':
             response = dashscope.Generation.call(
@@ -83,12 +180,128 @@ class UrbanLLM:
             if response.status_code == HTTPStatus.OK:
                 return response.output.choices[0]['message']['content']
             else:
-                return "Error: {}".format(response.status_code)
+                return "Error: {}, {}".format(response.status_code, response.message)
+        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
+            self.completion_tokens_used += response.usage.completion_tokens
+            self.request_number += 1
+            return response.choices[0].message.content
         else:
             print("ERROR: Wrong Config")
             return "wrong config"
+        
+    async def atext_request(self, dialog:list[dict], temperature:float=1, max_tokens:int=None, top_p:float=None, frequency_penalty:float=None, presence_penalty:float=None, timeout:int=300, retries=3):
+        """
+        异步版文本请求
+        """
+        if self.config.text['request_type'] == 'openai':
+            for attempt in range(retries):
+                try:
+                    if self.semaphore != None:
+                        async with self.semaphore:
+                            response = await self._aclient.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,
+                                timeout=timeout,
+                            )
+                            self.prompt_tokens_used += response.usage.prompt_tokens
+                            self.completion_tokens_used += response.usage.completion_tokens
+                            self.request_number += 1
+                            return response.choices[0].message.content
+                    else:
+                        response = await self._aclient.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,
+                            timeout=timeout,
+                        )
+                        self.prompt_tokens_used += response.usage.prompt_tokens
+                        self.completion_tokens_used += response.usage.completion_tokens
+                        self.request_number += 1
+                        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}")
+                    else:
+                        print("An error occurred:", e)
+                    if attempt < retries - 1:
+                        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']
+        elif self.config.text['request_type'] == 'deepseek':
+            client = AsyncOpenAI(
+                api_key=self.config.text['api_key'],
+                base_url="https://api.deepseek.com/beta",
+            )
+            response = await client.chat.completions.create(
+                model="deepseek-chat",
+                messages=dialog,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                top_p=top_p,
+                frequency_penalty=frequency_penalty,
+                presence_penalty=presence_penalty,
+                stream=False,
+                timeout=timeout,
+            )
+            self.prompt_tokens_used += response.usage.prompt_tokens
+            self.completion_tokens_used += response.usage.completion_tokens
+            self.request_number += 1
+            return response.choices[0].message.content
+        else:
+            print("ERROR: Wrong Config")
+            return "wrong config"
+        
 
-    def img_understand(self, img_path:Union[str, list[str]], prompt:str=None) -> str:
+    async def img_understand(self, img_path:Union[str, list[str]], prompt:str=None) -> str:
         """
         图像理解
         Image understanding
@@ -169,7 +382,7 @@ class UrbanLLM:
             print("ERROR: wrong image understanding type, only 'openai' and 'openai' is available")
             return "Error"
 
-    def img_generate(self, prompt:str, size:str='512*512', quantity:int = 1):
+    async def img_generate(self, prompt:str, size:str='512*512', quantity:int = 1):
         """
         图像生成
         Image generation
@@ -197,4 +410,91 @@ class UrbanLLM:
         else:
             print('Failed, status_code: %s, code: %s, message: %s' %
                 (rsp.status_code, rsp.code, rsp.message))
-            return None
+            return None
+
+    async def convert_prompt_to_batch(self, prompt_raw, K, type:str="mr"):
+        """
+        将单个目标的 prompt_raw 转换为适合一次批处理调用的 prompt_batch。
+        prompt_raw: 单目标的prompt，包含占位符!<INPUT X>!
+        K: 批处理的大小
+        type: 批处理prompt类型, 'mr': Multi-request; 'gr': Gather-request
+        返回值: prompt_batch 和 variable_batch
+        """
+        # 分割说明性文本和待格式化的 prompt
+        sections = prompt_raw.split("<commentblockmarker>###</commentblockmarker>")
+        description_section = sections[0].strip()  # 说明文本部分
+        prompt_section = sections[1].strip()  # 待格式化的 prompt
+
+        # 1. 提取变量说明
+        variable_info = {}
+        variable_lines = re.findall(r"!<INPUT (\d+)>! -- (.+)", description_section)
+        for var, desc in variable_lines:
+            variable_info[int(var)] = desc.strip()
+            
+        # 找到所有的占位符，如!<INPUT X>!
+        placeholders = re.findall(r"!<INPUT (\d+)>!", prompt_section)
+        
+        # 构建prompt_batch，使用批次变量占位符，并合并为一个请求
+        prompt_batch = ""
+        for i in range(K):
+            current_prompt = prompt_section
+            for ph in placeholders:
+                current_prompt = current_prompt.replace(f"!<INPUT {ph}>!", f"{{batch[{i}][{ph}]}}")
+            prompt_batch += f"### Request {i + 1} ###\n" + current_prompt + "\n"
+        
+        # 构建variable_batch，假设每个批次有K个值
+        variable_batch = [{int(ph): f"input_{i}_{ph}" for ph in placeholders} for i in range(K)]
+        if type == "mr":
+            return PromptBatch(prompt_batch, variable_batch, variable_info, K)
+        elif type == "gr":
+            messages = [{
+                "role": "system",
+                "content": """
+You are given a batch prompt structured in multiple requests and a variable description. Your task is to simplify and aggregate this prompt by removing redundant request headers and combining variables across different requests into a more compact format.
+
+### Instructions:
+1. Identify all unique sections in the prompt that appear across different requests. These sections will be aggregated into a single instance.
+2. For each variable in the prompt (e.g., `{batch[0][0]}`, `{batch[0][1]}`), group them together across all requests, and present them as a list or in a numbered format.
+3. Remove all redundant request headers (e.g., `### Request 1 ###`) and replace them with a single unified block where variables are grouped by their positions.
+4. Ensure the output is structured clearly for LLM processing, avoiding unnecessary repetition.
+5. Ensure the output only contains the transformed prompt.
+
+Here is an example transformation:
+Original prompt:
+### Request 1 ###
+{batch[0][0]}
+In general, Your Lifestyle is as follows: {batch[0][1]}
+Today is {batch[0][2]}, your wake up hour:
+### Request 2 ###
+{batch[1][0]}
+In general, Your Lifestyle is as follows: {batch[1][1]}
+Today is {batch[1][2]}, your wake up hour:
+Variable description:
+batch[x][0] -- Identity Stable Set
+batch[x][1] -- Lifestyle
+batch[x][2] -- Day
+
+Transformed prompt:
+Your Identity Stable Set:
+1. {batch[0][0]}
+2. {batch[1][0]}
+Your Lifestyle:
+1. {batch[0][1]}
+2. {batch[1][1]}
+The Day:
+1. {batch[0][2]}
+2. {batch[1][2]}
+Today, your wake up hour is:
+""",
+                }]
+            var_des = ""
+            for var, desc in variable_info.items():
+                var_des += f"batch[x][{var}] -- {desc}\n"
+            user_prompt = f"""Here is the original batch prompt:
+{prompt_batch}
+The Variable description:
+{var_des}
+"""
+            messages.append({'role': 'user', 'content': user_prompt})
+            prompt_batch = await self.atext_request(messages)
+            return PromptBatch(prompt_batch, variable_batch, variable_info, K)