pycityagent 2.0.0a65__cp312-cp312-macosx_11_0_arm64.whl → 2.0.0a67__cp312-cp312-macosx_11_0_arm64.whl

pycityagent/llm/embeddings.py

@@ -15,6 +15,14 @@ __all__ = [
 
 
 class SentenceEmbedding(Embeddings):
+    """
+    Main class for generating sentence embeddings using a pre-trained language model.
+
+    - **Description**:
+        - This class initializes a tokenizer and a pre-trained model to generate embeddings for input texts.
+        - It supports automatic CUDA device allocation and handles caching of models locally.
+    """
+
     def __init__(
         self,
         pretrained_model_name_or_path: Union[str, os.PathLike] = "BAAI/bge-m3",
@@ -24,6 +32,17 @@ class SentenceEmbedding(Embeddings):
         cache_dir: str = "./cache",
         proxies: Optional[dict] = None,
     ):
+        """
+        Initializes the SentenceEmbedding instance.
+
+        - **Parameters**:
+            - `pretrained_model_name_or_path`: Name or path of the pre-trained model. Default is "BAAI/bge-m3".
+            - `max_seq_len`: Maximum sequence length for inputs. Default is 8192.
+            - `auto_cuda`: Automatically move the model to available CUDA device if True. Default is False.
+            - `local_files_only`: Only try to load model from local files if True. Default is False.
+            - `cache_dir`: Directory to cache models. Default is "./cache".
+            - `proxies`: Proxy settings for HTTP/HTTPS. Default is None.
+        """
         os.makedirs(cache_dir, exist_ok=True)
         self.tokenizer = AutoTokenizer.from_pretrained(
             pretrained_model_name_or_path,
@@ -46,6 +65,15 @@ class SentenceEmbedding(Embeddings):
         self.max_seq_len = max_seq_len
 
     def _embed(self, texts: list[str]) -> list[list[float]]:
+        """
+        Internal method to compute embeddings for a list of input texts.
+
+        - **Parameters**:
+            - `texts`: A list of strings representing the input texts.
+
+        - **Returns**:
+            - A list of lists containing floating-point numbers representing the embeddings.
+        """
         # Tokenize sentences
         encoded_input = self.tokenizer(
             texts, padding=True, truncation=True, return_tensors="pt"
@@ -73,27 +101,46 @@ class SentenceEmbedding(Embeddings):
         return sentence_embeddings.tolist()
 
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        """Embed documents."""
+        """
+        Embeds a list of documents.
+
+        - **Parameters**:
+            - `texts`: The documents to embed.
+
+        - **Returns**:
+            - A list of document embeddings.
+        """
         return self._embed(texts)
 
     def embed_query(self, text: str) -> list[float]:
-        """Embed query text."""
+        """
+        Embeds a single query text.
+
+        - **Parameters**:
+            - `text`: The query text to embed.
+
+        - **Returns**:
+            - The embedding of the query text.
+        """
         return self._embed([text])[0]
 
 
 class SimpleEmbedding(Embeddings):
-    """简单的基于内存的embedding实现
+    """
+    A simple in-memory embedding implementation using Bag of Words and TF-IDF for generating text vectors.
 
-    使用简单的词袋模型(Bag of Words)和TF-IDF来生成文本的向量表示。
-    所有向量都保存在内存中，适用于小规模应用。
+    - **Description**:
+        - This class provides a basic approach to creating text embeddings based on term frequency-inverse document frequency (TF-IDF).
+        - It maintains an in-memory cache of computed embeddings and updates its vocabulary dynamically as new texts are processed.
     """
 
     def __init__(self, vector_dim: int = 128, cache_size: int = 1000):
-        """初始化
+        """
+        Initializes the SimpleEmbedding instance.
 
-        Args:
-            vector_dim: 向量维度
-            cache_size: 缓存大小，超过此大小将清除最早的缓存
+        - **Parameters**:
+            - `vector_dim`: Dimensionality of the vectors. Default is 128.
+            - `cache_size`: Cache size; oldest entries are removed when exceeded. Default is 1000.
         """
         self.vector_dim = vector_dim
         self.cache_size = cache_size
@@ -103,29 +150,73 @@ class SimpleEmbedding(Embeddings):
         self._doc_count = 0  # 文档总数
 
     def _text_to_hash(self, text: str) -> str:
-        """将文本转换为hash值"""
+        """
+        Converts text into a hash value.
+
+        - **Parameters**:
+            - `text`: The input text to hash.
+
+        - **Returns**:
+            - A string representing the MD5 hash of the input text.
+        """
         return hashlib.md5(text.encode()).hexdigest()
 
     def _tokenize(self, text: str) -> list[str]:
-        """简单的分词"""
+        """
+        Performs simple tokenization on the input text.
+
+        - **Parameters**:
+            - `text`: The input text to tokenize.
+
+        - **Returns**:
+            - A list of tokens extracted from the input text.
+        """
         # 这里使用简单的空格分词，实际应用中可以使用更复杂的分词方法
         return text.lower().split()
 
     def _update_vocab(self, tokens: list[str]):
-        """更新词汇表"""
+        """
+        Updates the vocabulary with new tokens.
+
+        - **Description**:
+            - This method adds unique tokens from the input list to the vocabulary.
+            - It ensures that each token is only added once using a set for deduplication.
+
+        - **Parameters**:
+            - `tokens`: A list of strings representing the tokens to add to the vocabulary.
+        """
         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值"""
+        """
+        Updates the IDF values based on the tokens.
+
+        - **Description**:
+            - Increases the document count and updates the inverse document frequency (IDF) for each unique token.
+
+        - **Parameters**:
+            - `tokens`: A list of strings representing the tokens from a single document.
+        """
         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)"""
+        """
+        Calculates term frequency (TF) for the tokens.
+
+        - **Description**:
+            - Computes the frequency of each token within the provided list and normalizes it by the total number of tokens.
+
+        - **Parameters**:
+            - `tokens`: A list of strings representing the tokens to calculate TF for.
+
+        - **Returns**:
+            - A dictionary mapping each token to its normalized term frequency.
+        """
         tf = {}
         total_tokens = len(tokens)
         for token in tokens:
@@ -136,7 +227,18 @@ class SimpleEmbedding(Embeddings):
         return tf
 
     def _calculate_tfidf(self, tokens: list[str]) -> list[float]:
-        """计算TF-IDF向量"""
+        """
+        Calculates the TF-IDF vector for the tokens.
+
+        - **Description**:
+            - Generates a vector representation of the input tokens using the Term Frequency-Inverse Document Frequency (TF-IDF) weighting scheme.
+
+        - **Parameters**:
+            - `tokens`: A list of strings representing the tokens to generate the TF-IDF vector for.
+
+        - **Returns**:
+            - A list of floating-point numbers representing the TF-IDF vector.
+        """
         vector = np.zeros(self.vector_dim)
         tf = self._calculate_tf(tokens)
 
@@ -154,13 +256,19 @@ class SimpleEmbedding(Embeddings):
         return list(vector)
 
     def _embed(self, text: str) -> list[float]:
-        """生成文本的向量表示
+        """
+        Generates a vector representation for the input text.
 
-        Args:
-            text: 输入文本
+        - **Description**:
+            - Creates an embedding for the given text by first checking if it's already cached,
+              then tokenizing, updating the vocabulary and IDF, calculating the TF-IDF vector,
+              and finally caching the result.
 
-        Returns:
-            np.ndarray: 文本的向量表示
+        - **Parameters**:
+            - `text`: The input text to generate the vector for.
+
+        - **Returns**:
+            - A list of floating-point numbers representing the vector of the input text.
         """
         # 检查缓存
         text_hash = self._text_to_hash(text)
@@ -189,19 +297,30 @@ class SimpleEmbedding(Embeddings):
         return list(vector)
 
     def embed_documents(self, texts: list[str]) -> list[list[float]]:
-        """Embed documents."""
+        """
+        Embeds a list of documents.
+
+        - **Parameters**:
+            - `texts`: A list of strings representing the documents to embed.
+
+        - **Returns**:
+            - A list of lists containing the embeddings of the documents.
+        """
         return [self._embed(text) for text in texts]
 
     def embed_query(self, text: str) -> list[float]:
-        """Embed query text."""
+        """
+        Embeds a single query text.
+
+        - **Parameters**:
+            - `text`: The query text to embed.
+
+        - **Returns**:
+            - A list of floating-point numbers representing the embedding of the query text.
+        """
         return self._embed(text)
 
+
 if __name__ == "__main__":
-    # se = SentenceEmbedding(
-    #     pretrained_model_name_or_path="ignore/BAAI--bge-m3", cache_dir="ignore"
-    # )
     se = SimpleEmbedding()
     print(se.embed_query("hello world"))
-    print(se.embed_query("hello world"))
-    print(se.embed_query("hello world"))
-    print(se.embed_query("hello world"))

pycityagent/llm/llm.py

@@ -1,14 +1,9 @@
 """UrbanLLM: 智能能力类及其定义"""
 
+import asyncio
 import json
 import logging
-
-from openai import APIConnectionError, AsyncOpenAI, OpenAI, OpenAIError
-from zhipuai import ZhipuAI
-
-logging.getLogger("zhipuai").setLevel(logging.WARNING)
-
-import asyncio
+import time
 import os
 from http import HTTPStatus
 from io import BytesIO
@@ -17,21 +12,37 @@ from typing import Any, Optional, Union
 import dashscope
 import requests
 from dashscope import ImageSynthesis
+from openai import APIConnectionError, AsyncOpenAI, OpenAI, OpenAIError
 from PIL import Image
+from zhipuai import ZhipuAI
 
 from .llmconfig import *
 from .utils import *
 
+logging.getLogger("zhipuai").setLevel(logging.WARNING)
 os.environ["GRPC_VERBOSITY"] = "ERROR"
 
+__all__ = [
+    "LLM",
+]
+
 
 class LLM:
     """
-    大语言模型对象
-    The LLM Object used by Agent(Soul)
+    Main class for the Large Language Model (LLM) object used by Agent(Soul).
+
+    - **Description**:
+        - This class manages configurations and interactions with different large language model APIs.
+        - It initializes clients based on the specified request type and handles token usage and consumption reporting.
     """
 
     def __init__(self, config: LLMConfig) -> None:
+        """
+        Initializes the LLM instance.
+
+        - **Parameters**:
+            - `config`: An instance of `LLMConfig` containing configuration settings for the LLM.
+        """
         self.config = config
         if config.text["request_type"] not in ["openai", "deepseek", "qwen", "zhipuai"]:
             raise ValueError("Invalid request type for text request")
@@ -40,12 +51,14 @@ class LLM:
         self.request_number = 0
         self.semaphore = None
         self._current_client_index = 0
+        self._log_list = []
 
         api_keys = self.config.text["api_key"]
         if not isinstance(api_keys, list):
             api_keys = [api_keys]
 
         self._aclients = []
+        self._client_usage = []
 
         for api_key in api_keys:
             if self.config.text["request_type"] == "openai":
@@ -69,11 +82,31 @@ class LLM:
                     f"Unsupported `request_type` {self.config.text['request_type']}!"
                 )
             self._aclients.append(client)
+            self._client_usage.append({
+                "prompt_tokens": 0,
+                "completion_tokens": 0,
+                "request_number": 0
+            })
+
+    def get_log_list(self):
+        return self._log_list
+    
+    def clear_log_list(self):
+        self._log_list = []
 
     def set_semaphore(self, number_of_coroutine: int):
+        """
+        Sets the semaphore for controlling concurrent coroutines.
+
+        - **Parameters**:
+            - `number_of_coroutine`: The maximum number of concurrent coroutines allowed.
+        """
         self.semaphore = asyncio.Semaphore(number_of_coroutine)
 
     def clear_semaphore(self):
+        """
+        Clears the semaphore setting.
+        """
         self.semaphore = None
 
     def clear_used(self):
@@ -81,48 +114,80 @@ class LLM:
         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
+        for usage in self._client_usage:
+            usage["prompt_tokens"] = 0
+            usage["completion_tokens"] = 0 
+            usage["request_number"] = 0
+
+    def get_consumption(self):
+        consumption = {}
+        for i, usage in enumerate(self._client_usage):
+            consumption[f"api-key-{i+1}"] = {
+                "total_tokens": usage["prompt_tokens"] + usage["completion_tokens"],
+                "request_number": usage["request_number"]
+            }
+        return consumption
 
     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
+        Displays token usage and optionally calculates the estimated cost based on provided prices.
+
+        - **Parameters**:
+            - `input_price`: Price per million prompt tokens. Default is None.
+            - `output_price`: Price per million completion tokens. Default is None.
+
+        - **Returns**:
+            - A dictionary summarizing the token usage and, if applicable, the estimated cost.
         """
-        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,
+        total_stats = {
+            "total": 0,
+            "prompt": 0,
+            "completion": 0,
+            "requests": 0
         }
+        
+        for i, usage in enumerate(self._client_usage):
+            prompt_tokens = usage["prompt_tokens"]
+            completion_tokens = usage["completion_tokens"]
+            requests = usage["request_number"]
+            total_tokens = prompt_tokens + completion_tokens
+            
+            total_stats["total"] += total_tokens
+            total_stats["prompt"] += prompt_tokens
+            total_stats["completion"] += completion_tokens
+            total_stats["requests"] += requests
+            
+            rate = prompt_tokens / completion_tokens if completion_tokens != 0 else "nan"
+            tokens_per_request = total_tokens / requests if requests != 0 else "nan"
+            
+            print(f"\nAPI Key #{i+1}:")
+            print(f"Request Number: {requests}")
+            print("Token Usage:")
+            print(f"    - Total tokens: {total_tokens}")
+            print(f"    - Prompt tokens: {prompt_tokens}")
+            print(f"    - Completion tokens: {completion_tokens}")
+            print(f"    - Token per request: {tokens_per_request}")
+            print(f"    - Prompt:Completion ratio: {rate}:1")
+            
+            if input_price is not None and output_price is not None:
+                consumption = (prompt_tokens / 1000000 * input_price + 
+                             completion_tokens / 1000000 * output_price)
+                print(f"    - Cost Estimation: {consumption}")
+        
+        return total_stats
 
     def _get_next_client(self):
-        """获取下一个要使用的客户端"""
+        """
+        Retrieves the next client to be used for making requests.
+
+        - **Description**:
+            - This method cycles through the available clients in a round-robin fashion.
+
+        - **Returns**:
+            - The next client instance to be used for making requests.
+        """
         client = self._aclients[self._current_client_index]
         self._current_client_index = (self._current_client_index + 1) % len(
             self._aclients
@@ -143,8 +208,30 @@ Token Usage:
         tool_choice: Optional[dict[str, Any]] = None,
     ):
         """
-        异步版文本请求
+        Sends an asynchronous text request to the configured LLM API.
+
+        - **Description**:
+            - Attempts to send a text request up to `retries` times with exponential backoff on failure.
+            - Handles different request types and manages token usage statistics.
+
+        - **Parameters**:
+            - `dialog`: Messages to send as part of the chat completion request.
+            - `temperature`: Controls randomness in the model's output. Default is 1.
+            - `max_tokens`: Maximum number of tokens to generate in the response. Default is None.
+            - `top_p`: Limits the next token selection to a subset of tokens with a cumulative probability above this value. Default is None.
+            - `frequency_penalty`: Penalizes new tokens based on their existing frequency in the text so far. Default is None.
+            - `presence_penalty`: Penalizes new tokens based on whether they appear in the text so far. Default is None.
+            - `timeout`: Request timeout in seconds. Default is 300 seconds.
+            - `retries`: Number of retry attempts in case of failure. Default is 3.
+            - `tools`: List of dictionaries describing the tools that can be called by the model. Default is None.
+            - `tool_choice`: Dictionary specifying how the model should choose from the provided tools. Default is None.
+
+        - **Returns**:
+            - A string containing the message content or a dictionary with tool call arguments if tools are used.
+            - Raises exceptions if the request fails after all retry attempts.
         """
+        start_time = time.time()
+        log = {"request_time": start_time}
         if (
             self.config.text["request_type"] == "openai"
             or self.config.text["request_type"] == "deepseek"
@@ -153,57 +240,35 @@ Token Usage:
             for attempt in range(retries):
                 try:
                     client = self._get_next_client()
-                    if self.semaphore != None:
-                        async with self.semaphore:
-                            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.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 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
+                    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:
-                        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.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 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
+                        return response.choices[0].message.content
                 except APIConnectionError as e:
                     print("API connection error:", e)
                     if attempt < retries - 1:
@@ -248,9 +313,13 @@ Token Usage:
                     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
+                    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
@@ -273,15 +342,14 @@ Token Usage:
         self, img_path: Union[str, list[str]], prompt: Optional[str] = None
     ) -> str:
         """
-        图像理解
-        Image understanding
+        Analyzes and understands images using external APIs.
 
-        Args:
-        - img_path (Union[str, list[str]]): 目标图像的路径, 既可以是一个路径也可以是包含多张图片路径的list. The path of selected Image
-        - prompt (str): 理解提示词 - 例如理解方向. The understanding prompts
+        - **Args**:
+            img_path (Union[str, list[str]]): Path or list of paths to the images for analysis.
+            prompt (Optional[str]): Guidance text for understanding the images.
 
-        Returns:
-        - (str): the understanding content
+        - **Returns**:
+            str: The content derived from understanding the images.
         """
         ppt = "如何理解这幅图像？"
         if prompt != None:
@@ -354,16 +422,15 @@ Token Usage:
 
     async def img_generate(self, prompt: str, size: str = "512*512", quantity: int = 1):
         """
-        图像生成
-        Image generation
+        Generates images based on a given prompt.
 
-        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
+        - **Args**:
+            prompt (str): Prompt for generating images.
+            size (str): Size of the generated images, default is '512*512'.
+            quantity (int): Number of images to generate, default is 1.
 
-        Returns:
-        - (list[PIL.Image.Image]): 生成的图像列表. The list of generated Images.
+        - **Returns**:
+            list[PIL.Image.Image]: List of generated PIL Image objects.
         """
         rsp = ImageSynthesis.call(
             model=self.config.image_g["model"],

pycityagent/llm/llmconfig.py

@@ -1,3 +1,8 @@
+__all__ = [
+    "LLMConfig",
+]
+
+
 class LLMConfig:
     """
     大语言模型相关配置