pycityagent 2.0.0a6__py3-none-any.whl → 2.0.0a8__py3-none-any.whl

pycityagent/llm/llm.py

@@ -4,6 +4,7 @@ import json
 from openai import OpenAI, AsyncOpenAI, APIConnectionError, OpenAIError
 from zhipuai import ZhipuAI
 import logging
+
 logging.getLogger("zhipuai").setLevel(logging.WARNING)
 
 import asyncio
@@ -19,29 +20,38 @@ 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']:
+        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)
+        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):
+    def set_semaphore(self, number_of_coroutine: int):
         self.semaphore = asyncio.Semaphore(number_of_coroutine)
 
     def clear_semaphore(self):
@@ -56,19 +66,21 @@ class LLM:
         self.completion_tokens_used = 0
         self.request_number = 0
 
-    def show_consumption(self, input_price:Optional[float]=None, output_price:Optional[float]=None):
+    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
+            rate = self.prompt_tokens_used / self.completion_tokens_used
         else:
-            rate = 'nan'
+            rate = "nan"
         if self.request_number != 0:
-            TcA = total_token/self.request_number
+            TcA = total_token / self.request_number
         else:
-            TcA = 'nan'
+            TcA = "nan"
         out = f"""Request Number: {self.request_number}
 Token Usage:
     - Total tokens: {total_token}
@@ -77,23 +89,30 @@ Token Usage:
     - 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
+            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}
-
+        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]:
+        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
@@ -109,17 +128,17 @@ Token Usage:
         Returns:
         - (str): the response content
         """
-        if 'api_base' in self.config.text.keys():
-            api_base = self.config.text['api_base']
+        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':
+        if self.config.text["request_type"] == "openai":
             client = OpenAI(
-                api_key=self.config.text['api_key'], 
+                api_key=self.config.text["api_key"],
                 base_url=api_base,
             )
             response = client.chat.completions.create(
-                model=self.config.text['model'],
+                model=self.config.text["model"],
                 messages=dialog,
                 temperature=temperature,
                 max_tokens=max_tokens,
@@ -127,33 +146,33 @@ Token Usage:
                 frequency_penalty=frequency_penalty,
                 presence_penalty=presence_penalty,
                 tools=tools,
-                tool_choice=tool_choice
+                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.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':
+        elif self.config.text["request_type"] == "qwen":
             response = dashscope.Generation.call(
-                model=self.config.text['model'],
-                api_key=self.config.text['api_key'],
+                model=self.config.text["model"],
+                api_key=self.config.text["api_key"],
                 messages=dialog,
-                result_format='message'
+                result_format="message",
             )
-            if response.status_code == HTTPStatus.OK: # type: ignore
-                return response.output.choices[0]['message']['content'] # type: ignore
+            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':
+                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'],
+                api_key=self.config.text["api_key"],
                 base_url="https://api.deepseek.com/beta",
             )
             response = client.chat.completions.create(
-                model=self.config.text['model'],
+                model=self.config.text["model"],
                 messages=dialog,
                 temperature=temperature,
                 max_tokens=max_tokens,
@@ -162,63 +181,66 @@ Token Usage:
                 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.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'])
+        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'],
+                model=self.config.text["model"],
                 messages=dialog,
                 temperature=temperature,
                 top_p=top_p,
-                stream=False
+                stream=False,
             )
-            self.prompt_tokens_used += response.usage.prompt_tokens # type: ignore
-            self.completion_tokens_used += response.usage.completion_tokens # 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
-            return response.choices[0].message.content # type: ignore
+            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
-            ):
+        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':
+        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'],
+                                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
+                                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
+                                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
@@ -226,20 +248,20 @@ Token Usage:
                                 return response.choices[0].message.content
                     else:
                         response = await self._aclient.chat.completions.create(
-                            model=self.config.text['model'],
+                            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
+                            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
+                            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
@@ -248,77 +270,89 @@ Token Usage:
                 except APIConnectionError as e:
                     print("API connection error:", e)
                     if attempt < retries - 1:
-                        await asyncio.sleep(2 ** attempt)
+                        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
+                    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)
+                        await asyncio.sleep(2**attempt)
                     else:
                         raise e
-        elif self.config.text['request_type'] == 'zhipuai':
+        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'],
+                    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
+                        tool_choice=tool_choice,
                     )
                     task_id = response.id
-                    task_status = ''
+                    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
+                    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':
+                    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.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)
+                        return json.loads(
+                            result_response.choices[0]
+                            .message.tool_calls[0]
+                            .function.arguments
+                        )
                     else:
-                        return result_response.choices[0].message.content # type: ignore
+                        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)
+                        await asyncio.sleep(2**attempt)
                     else:
                         raise e
-        elif self.config.text['request_type'] == 'qwen':
+        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']}"}
+                headers = {
+                    "Content-Type": "application/json",
+                    "Authorization": f"{self.config.text['api_key']}",
+                }
                 payload = {
-                    'model': self.config.text['model'],
-                    'input': {
-                        'messages': dialog
-                    }
+                    "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']}")
+                    if "code" in response_json.keys():
+                        raise Exception(
+                            f"Error: {response_json['code']}, {response_json['message']}"
+                        )
                     else:
-                        return response_json['output']['text']
+                        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:
+    async def img_understand(
+        self, img_path: Union[str, list[str]], prompt: Optional[str] = None
+    ) -> str:
         """
         图像理解
         Image understanding
@@ -333,73 +367,73 @@ Token Usage:
         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']
+        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'], 
+                api_key=self.config.text["api_key"],
                 base_url=api_base,
             )
             content = []
-            content.append({'type': 'text', 'text': ppt})
+            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):
+                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}"
+                    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
-                }]
+                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':
+            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):
+                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})
+                    content.append({"image": "file://" + item})
+                content.append({"text": ppt})
 
-            dialog = [{
-                'role': 'user',
-                'content': content
-            }]
+            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
+                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")
+            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):
+    async def img_generate(self, prompt: str, size: str = "512*512", quantity: int = 1):
         """
         图像生成
         Image generation
@@ -413,11 +447,11 @@ Token Usage:
         - (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
+            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 = []
@@ -425,6 +459,8 @@ Token Usage:
                 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
+            print(
+                "Failed, status_code: %s, code: %s, message: %s"
+                % (rsp.status_code, rsp.code, rsp.message)
+            )
+            return None

pycityagent/llm/llmconfig.py

@@ -3,13 +3,11 @@ class LLMConfig:
     大语言模型相关配置
     The config of LLM
     """
-    def __init__(
-        self, 
-        config: dict
-    ) -> None:
+
+    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']
+        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

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

pycityagent/memory/memory.py

@@ -2,8 +2,7 @@ import asyncio
 import logging
 from copy import deepcopy
 from datetime import datetime
-from typing import (Any, Callable, Dict, List, Literal, Optional, Sequence,
-                    Tuple, Union)
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence, Tuple, Union
 
 import numpy as np
 from pyparsing import deque

pycityagent/memory/memory_base.py

@@ -6,8 +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, Iterable, List, Optional, Sequence, Tuple, Union
 
 from .const import *