pycityagent 1.0.0__py3-none-any.whl → 2.0.0a1__py3-none-any.whl

pycityagent/llm/embedding.py

@@ -0,0 +1,136 @@
+"""简单的基于内存的embedding实现"""
+
+import numpy as np
+from typing import List, Dict, Optional
+import hashlib
+import json
+
+class SimpleEmbedding:
+    """简单的基于内存的embedding实现
+    
+    使用简单的词袋模型(Bag of Words)和TF-IDF来生成文本的向量表示。
+    所有向量都保存在内存中，适用于小规模应用。
+    """
+    
+    def __init__(self, vector_dim: int = 128, cache_size: int = 1000):
+        """初始化
+        
+        Args:
+            vector_dim: 向量维度
+            cache_size: 缓存大小，超过此大小将清除最早的缓存
+        """
+        self.vector_dim = vector_dim
+        self.cache_size = cache_size
+        self._cache: Dict[str, np.ndarray] = {}
+        self._vocab: Dict[str, int] = {}  # 词汇表
+        self._idf: Dict[str, float] = {}  # 逆文档频率
+        self._doc_count = 0  # 文档总数
+        
+    def _text_to_hash(self, text: str) -> str:
+        """将文本转换为hash值"""
+        return hashlib.md5(text.encode()).hexdigest()
+        
+    def _tokenize(self, text: str) -> List[str]:
+        """简单的分词"""
+        # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
+        return text.lower().split()
+        
+    def _update_vocab(self, tokens: List[str]):
+        """更新词汇表"""
+        for token in set(tokens):  # 使用set去重
+            if token not in self._vocab:
+                self._vocab[token] = len(self._vocab)
+                
+    def _update_idf(self, tokens: List[str]):
+        """更新IDF值"""
+        self._doc_count += 1
+        unique_tokens = set(tokens)
+        for token in unique_tokens:
+            self._idf[token] = self._idf.get(token, 0) + 1
+            
+    def _calculate_tf(self, tokens: List[str]) -> Dict[str, float]:
+        """计算词频(TF)"""
+        tf = {}
+        total_tokens = len(tokens)
+        for token in tokens:
+            tf[token] = tf.get(token, 0) + 1
+        # 归一化
+        for token in tf:
+            tf[token] /= total_tokens
+        return tf
+        
+    def _calculate_tfidf(self, tokens: List[str]) -> np.ndarray:
+        """计算TF-IDF向量"""
+        vector = np.zeros(self.vector_dim)
+        tf = self._calculate_tf(tokens)
+        
+        for token, tf_value in tf.items():
+            if token in self._idf:
+                idf = np.log(self._doc_count / self._idf[token])
+                idx = self._vocab[token] % self.vector_dim  # 使用取模运算来控制向量维度
+                vector[idx] += tf_value * idf
+                
+        # L2归一化
+        norm = np.linalg.norm(vector)
+        if norm > 0:
+            vector /= norm
+            
+        return vector
+        
+    async def embed(self, text: str) -> np.ndarray:
+        """生成文本的向量表示
+        
+        Args:
+            text: 输入文本
+            
+        Returns:
+            np.ndarray: 文本的向量表示
+        """
+        # 检查缓存
+        text_hash = self._text_to_hash(text)
+        if text_hash in self._cache:
+            return self._cache[text_hash]
+            
+        # 分词
+        tokens = self._tokenize(text)
+        if not tokens:
+            return np.zeros(self.vector_dim)
+            
+        # 更新词汇表和IDF
+        self._update_vocab(tokens)
+        self._update_idf(tokens)
+        
+        # 计算向量
+        vector = self._calculate_tfidf(tokens)
+        
+        # 更新缓存
+        if len(self._cache) >= self.cache_size:
+            # 删除最早的缓存
+            oldest_key = next(iter(self._cache))
+            del self._cache[oldest_key]
+        self._cache[text_hash] = vector
+        
+        return vector
+        
+    def save(self, file_path: str):
+        """保存模型"""
+        state = {
+            'vector_dim': self.vector_dim,
+            'cache_size': self.cache_size,
+            'vocab': self._vocab,
+            'idf': self._idf,
+            'doc_count': self._doc_count
+        }
+        with open(file_path, 'w') as f:
+            json.dump(state, f)
+            
+    def load(self, file_path: str):
+        """加载模型"""
+        with open(file_path, 'r') as f:
+            state = json.load(f)
+        self.vector_dim = state['vector_dim']
+        self.cache_size = state['cache_size']
+        self._vocab = state['vocab']
+        self._idf = state['idf']
+        self._doc_count = state['doc_count']
+        self._cache = {}  # 清空缓存 

pycityagent/llm/llm.py

@@ -0,0 +1,430 @@
+"""UrbanLLM: 智能能力类及其定义"""
+
+import json
+from openai import OpenAI, AsyncOpenAI, APIConnectionError, OpenAIError
+from zhipuai import ZhipuAI
+import logging
+logging.getLogger("zhipuai").setLevel(logging.WARNING)
+
+import asyncio
+from http import HTTPStatus
+import dashscope
+import requests
+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 *
+
+import os
+os.environ["GRPC_VERBOSITY"] = "ERROR"
+
+class LLM:
+    """
+    大语言模型对象
+    The LLM Object used by Agent(Soul)
+    """
+    def __init__(self, config: LLMConfig) -> None:
+        self.config = config
+        if config.text['request_type'] not in ['openai', 'deepseek', 'qwen', 'zhipuai']:
+            raise ValueError("Invalid request type for text request")
+        self.prompt_tokens_used = 0
+        self.completion_tokens_used = 0
+        self.request_number = 0
+        self.semaphore = None
+        if self.config.text['request_type'] == 'openai':
+            self._aclient = AsyncOpenAI(api_key=self.config.text['api_key'], timeout=300)
+        elif self.config.text['request_type'] == 'deepseek':
+            self._aclient = AsyncOpenAI(api_key=self.config.text['api_key'], base_url="https://api.deepseek.com/beta", timeout=300)
+        elif self.config.text['request_type'] == 'zhipuai':
+            self._aclient = ZhipuAI(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:Optional[float]=None, output_price:Optional[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: 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, 
+            temperature:float=1, 
+            max_tokens:Optional[int]=None, 
+            top_p:Optional[float]=None, 
+            frequency_penalty:Optional[float]=None, 
+            presence_penalty:Optional[float]=None, 
+            timeout:int=300, 
+            retries=3,
+            tools:Optional[List[Dict[str, Any]]]=None,
+            tool_choice:Optional[Dict[str, Any]]=None
+            ):
+        """
+        异步版文本请求
+        """
+        if self.config.text['request_type'] == 'openai' or self.config.text['request_type'] == 'deepseek':
+            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, # type: ignore
+                                presence_penalty=presence_penalty, # type: ignore
+                                stream=False,
+                                timeout=timeout,
+                                tools=tools,
+                                tool_choice=tool_choice
+                            ) # type: ignore
+                            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
+                    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, # type: ignore
+                            presence_penalty=presence_penalty, # type: ignore
+                            stream=False,
+                            timeout=timeout,
+                            tools=tools,
+                            tool_choice=tool_choice
+                        ) # type: ignore
+                        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
+                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:
+                    response = self._aclient.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 = self._aclient.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.prompt_tokens_used += result_response.usage.prompt_tokens # type: ignore
+                    self.completion_tokens_used += result_response.usage.completion_tokens # type: ignore
+                    self.request_number += 1
+                    if tools and result_response.choices[0].message.tool_calls:
+                        return json.loads(result_response.choices[0].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
+        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"
+        
+
+    async def img_understand(self, img_path:Union[str, list[str]], prompt:Optional[str]=None) -> str:
+        """
+        图像理解
+        Image understanding
+
+        Args:
+        - img_path (Union[str, list[str]]): 目标图像的路径, 既可以是一个路径也可以是包含多张图片路径的list. The path of selected Image
+        - prompt (str): 理解提示词 - 例如理解方向. The understanding prompts
+
+        Returns:
+        - (str): the understanding content
+        """
+        ppt = "如何理解这幅图像？"
+        if prompt != None:
+            ppt = prompt
+        if self.config.image_u['request_type'] == 'openai':
+            if 'api_base' in self.config.image_u.keys():
+                api_base = self.config.image_u['api_base']
+            else:
+                api_base = None
+            client = OpenAI(
+                api_key=self.config.text['api_key'], 
+                base_url=api_base,
+            )
+            content = []
+            content.append({'type': 'text', 'text': ppt})
+            if isinstance(img_path, str):
+                base64_image = encode_image(img_path)
+                content.append({
+                    'type': 'image_url', 
+                    'image_url': {
+                            'url': f"data:image/jpeg;base64,{base64_image}"
+                        }
+                })
+            elif isinstance(img_path, list) and all(isinstance(item, str) for item in img_path):
+                for item in img_path:
+                    base64_image = encode_image(item)
+                    content.append({
+                        'type': 'image_url',
+                        'image_url': {
+                            'url': f"data:image/jpeg;base64,{base64_image}"
+                        }
+                    })
+            response = client.chat.completions.create(
+                model=self.config.image_u['model'],
+                messages=[{
+                    'role': 'user',
+                    'content': content
+                }]
+            )
+            return response.choices[0].message.content # type: ignore
+        elif self.config.image_u['request_type'] == 'qwen':
+            content = []
+            if isinstance(img_path, str):
+                content.append({'image': 'file://' + img_path})
+                content.append({'text': ppt})
+            elif isinstance(img_path, list) and all(isinstance(item, str) for item in img_path):
+                for item in img_path:
+                    content.append({
+                        'image': 'file://' + item
+                    })
+                content.append({'text': ppt})
+
+            dialog = [{
+                'role': 'user',
+                'content': content
+            }]
+            response = dashscope.MultiModalConversation.call(
+                        model=self.config.image_u['model'],
+                        api_key=self.config.image_u['api_key'],
+                        messages=dialog
+                    )
+            if response.status_code == HTTPStatus.OK: # type: ignore
+                return response.output.choices[0]['message']['content'] # type: ignore
+            else:
+                print(response.code)  # type: ignore # The error code.
+                return "Error"
+        else:
+            print("ERROR: wrong image understanding type, only 'openai' and 'openai' is available")
+            return "Error"
+
+    async def img_generate(self, prompt:str, size:str='512*512', quantity:int = 1):
+        """
+        图像生成
+        Image generation
+
+        Args:
+        - prompt (str): 图像生成提示词. The image generation prompts
+        - size (str): 生成图像尺寸, 默认为'512*512'. The image size, default: '512*512'
+        - quantity (int): 生成图像数量, 默认为1. The quantity of generated images, default: 1
+
+        Returns:
+        - (list[PIL.Image.Image]): 生成的图像列表. The list of generated Images.
+        """
+        rsp = ImageSynthesis.call(
+                model=self.config.image_g['model'],
+                api_key=self.config.image_g['api_key'],
+                prompt=prompt,
+                n=quantity,
+                size=size
+        )
+        if rsp.status_code == HTTPStatus.OK:
+            res = []
+            for result in rsp.output.results:
+                res.append(Image.open(BytesIO(requests.get(result.url).content)))
+            return res
+        else:
+            print('Failed, status_code: %s, code: %s, message: %s' %
+                (rsp.status_code, rsp.code, rsp.message))
+            return None

pycityagent/llm/llmconfig.py

@@ -0,0 +1,15 @@
+class LLMConfig:
+    """
+    大语言模型相关配置
+    The config of LLM
+    """
+    def __init__(
+        self, 
+        config: dict
+    ) -> None:
+        self.config = config
+        self.text = config['text_request']
+        if 'api_base' in self.text.keys() and self.text['api_base'] == 'None':
+            self.text['api_base'] = None
+        self.image_u = config['img_understand_request']
+        self.image_g = config['img_generate_request']

pycityagent/llm/utils.py

@@ -0,0 +1,6 @@
+import base64
+
+
+def encode_image(image_path):
+  with open(image_path, "rb") as image_file:
+    return base64.b64encode(image_file.read()).decode('utf-8')

pycityagent/memory/__init__.py

@@ -0,0 +1,11 @@
+"""Memory."""
+
+from .memory import Memory
+from .memory_base import MemoryBase, MemoryUnit
+from .profile import ProfileMemory, ProfileMemoryUnit
+from .self_define import DynamicMemory
+from .state import StateMemory
+
+__all__ = [
+    "Memory",
+]

pycityagent/memory/const.py

@@ -0,0 +1,41 @@
+from pycityproto.city.person.v2.motion_pb2 import Status
+
+PROFILE_ATTRIBUTES = {
+    "gender": str(),
+    "age": float(),
+    "education": str(),
+    "skill": str(),
+    "occupation": str(),
+    "family_consumption": str(),
+    "consumption": str(),
+    "personality": str(),
+    "income": str(),
+    "residence": str(),
+    "race": str(),
+    "religion": str(),
+    "marital_status": str(),
+}
+
+STATE_ATTRIBUTES = {
+    # base
+    "id": -1,
+    "attribute": dict(),
+    "home": dict(),
+    "work": dict(),
+    "schedules": [],
+    "vehicle_attribute": dict(),
+    "bus_attribute": dict(),
+    "pedestrian_attribute": dict(),
+    "bike_attribute": dict(),
+    # motion
+    "status": Status.STATUS_UNSPECIFIED,
+    "position": dict(),
+    "v": float(),
+    "direction": float(),
+    "activity": str(),
+    "l": float(),
+}
+
+SELF_DEFINE_PREFIX = "self_define_"
+
+TIME_STAMP_KEY = "_timestamp"