diff --git a/.gitignore b/.gitignore index 922116d12..aa5edd74a 100644 --- a/.gitignore +++ b/.gitignore @@ -1,7 +1,7 @@ ### Python template # Byte-compiled / optimized / DLL files -__pycache__/ +__pycache__ *.py[cod] *$py.class diff --git a/config/config2.example.yaml b/config/config2.example.yaml index 3a5cc3585..c5454ec32 100644 --- a/config/config2.example.yaml +++ b/config/config2.example.yaml @@ -4,6 +4,7 @@ llm: api_key: "YOUR_API_KEY" model: "gpt-4-turbo-preview" # or gpt-3.5-turbo-1106 / gpt-4-1106-preview proxy: "YOUR_PROXY" # for LLM API requests + # timeout: 600 # Optional. If set to 0, default value is 300. pricing_plan: "" # Optional. If invalid, it will be automatically filled in with the value of the `model`. # Azure-exclusive pricing plan mappings: # - gpt-3.5-turbo 4k: "gpt-3.5-turbo-1106" diff --git a/metagpt/actions/action_node.py b/metagpt/actions/action_node.py index 59559ad04..27dde5a8c 100644 --- a/metagpt/actions/action_node.py +++ b/metagpt/actions/action_node.py @@ -17,6 +17,7 @@ from pydantic import BaseModel, Field, create_model, model_validator from tenacity import retry, stop_after_attempt, wait_random_exponential from metagpt.actions.action_outcls_registry import register_action_outcls +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.llm import BaseLLM from metagpt.logs import logger from metagpt.provider.postprocess.llm_output_postprocess import llm_output_postprocess @@ -416,7 +417,7 @@ class ActionNode: images: Optional[Union[str, list[str]]] = None, system_msgs: Optional[list[str]] = None, schema="markdown", # compatible to original format - timeout=3, + timeout=USE_CONFIG_TIMEOUT, ) -> (str, BaseModel): """Use ActionOutput to wrap the output of aask""" content = await self.llm.aask(prompt, system_msgs, images=images, timeout=timeout) @@ -448,7 +449,9 @@ class ActionNode: def set_context(self, context): self.set_recursive("context", context) - async def simple_fill(self, schema, mode, images: Optional[Union[str, list[str]]] = None, timeout=3, exclude=None): + async def simple_fill( + self, schema, mode, images: Optional[Union[str, list[str]]] = None, timeout=USE_CONFIG_TIMEOUT, exclude=None + ): prompt = self.compile(context=self.context, schema=schema, mode=mode, exclude=exclude) if schema != "raw": @@ -473,7 +476,7 @@ class ActionNode: mode="auto", strgy="simple", images: Optional[Union[str, list[str]]] = None, - timeout=3, + timeout=USE_CONFIG_TIMEOUT, exclude=[], ): """Fill the node(s) with mode. diff --git a/metagpt/configs/llm_config.py b/metagpt/configs/llm_config.py index fa9bc0b1b..af8f56372 100644 --- a/metagpt/configs/llm_config.py +++ b/metagpt/configs/llm_config.py @@ -10,6 +10,7 @@ from typing import Optional from pydantic import field_validator +from metagpt.const import LLM_API_TIMEOUT from metagpt.utils.yaml_model import YamlModel @@ -74,7 +75,7 @@ class LLMConfig(YamlModel): stream: bool = False logprobs: Optional[bool] = None # https://cookbook.openai.com/examples/using_logprobs top_logprobs: Optional[int] = None - timeout: int = 60 + timeout: int = 600 # For Network proxy: Optional[str] = None @@ -88,3 +89,8 @@ class LLMConfig(YamlModel): if v in ["", None, "YOUR_API_KEY"]: raise ValueError("Please set your API key in config2.yaml") return v + + @field_validator("timeout") + @classmethod + def check_timeout(cls, v): + return v or LLM_API_TIMEOUT diff --git a/metagpt/const.py b/metagpt/const.py index 6dbbfe0c1..e4cebfd96 100644 --- a/metagpt/const.py +++ b/metagpt/const.py @@ -123,7 +123,6 @@ BASE64_FORMAT = "base64" # REDIS REDIS_KEY = "REDIS_KEY" -LLM_API_TIMEOUT = 300 # Message id IGNORED_MESSAGE_ID = "0" @@ -132,3 +131,7 @@ IGNORED_MESSAGE_ID = "0" GENERALIZATION = "Generalize" COMPOSITION = "Composite" AGGREGATION = "Aggregate" + +# Timeout +USE_CONFIG_TIMEOUT = 0 # Using llm.timeout configuration. +LLM_API_TIMEOUT = 300 diff --git a/metagpt/environment/__init__.py b/metagpt/environment/__init__.py index a844827d8..28981f2f8 100644 --- a/metagpt/environment/__init__.py +++ b/metagpt/environment/__init__.py @@ -4,10 +4,9 @@ from metagpt.environment.base_env import Environment from metagpt.environment.android_env.android_env import AndroidEnv -from metagpt.environment.minecraft_env.minecraft_env import MinecraftExtEnv from metagpt.environment.werewolf_env.werewolf_env import WerewolfEnv from metagpt.environment.stanford_town_env.stanford_town_env import StanfordTownEnv from metagpt.environment.software_env.software_env import SoftwareEnv -__all__ = ["AndroidEnv", "MinecraftExtEnv", "WerewolfEnv", "StanfordTownEnv", "SoftwareEnv", "Environment"] +__all__ = ["AndroidEnv", "WerewolfEnv", "StanfordTownEnv", "SoftwareEnv", "Environment"] diff --git a/metagpt/learn/skill_loader.py b/metagpt/learn/skill_loader.py index bcf28bb87..e98f73cf9 100644 --- a/metagpt/learn/skill_loader.py +++ b/metagpt/learn/skill_loader.py @@ -9,11 +9,11 @@ from pathlib import Path from typing import Dict, List, Optional -import aiofiles import yaml from pydantic import BaseModel, Field from metagpt.context import Context +from metagpt.utils.common import aread class Example(BaseModel): @@ -68,8 +68,7 @@ class SkillsDeclaration(BaseModel): async def load(skill_yaml_file_name: Path = None) -> "SkillsDeclaration": if not skill_yaml_file_name: skill_yaml_file_name = Path(__file__).parent.parent.parent / "docs/.well-known/skills.yaml" - async with aiofiles.open(str(skill_yaml_file_name), mode="r") as reader: - data = await reader.read(-1) + data = await aread(filename=skill_yaml_file_name) skill_data = yaml.safe_load(data) return SkillsDeclaration(**skill_data) diff --git a/metagpt/provider/anthropic_api.py b/metagpt/provider/anthropic_api.py index 872f9b2c7..1aeacbe83 100644 --- a/metagpt/provider/anthropic_api.py +++ b/metagpt/provider/anthropic_api.py @@ -5,6 +5,7 @@ from anthropic import AsyncAnthropic from anthropic.types import Message, Usage from metagpt.configs.llm_config import LLMConfig, LLMType +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream from metagpt.provider.base_llm import BaseLLM from metagpt.provider.llm_provider_registry import register_provider @@ -41,15 +42,15 @@ class AnthropicLLM(BaseLLM): def get_choice_text(self, resp: Message) -> str: return resp.content[0].text - async def _achat_completion(self, messages: list[dict], timeout: int = 3) -> Message: + async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message: resp: Message = await self.aclient.messages.create(**self._const_kwargs(messages)) self._update_costs(resp.usage, self.model) return resp - async def acompletion(self, messages: list[dict], timeout: int = 3) -> Message: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> Message: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: stream = await self.aclient.messages.create(**self._const_kwargs(messages, stream=True)) collected_content = [] usage = Usage(input_tokens=0, output_tokens=0) diff --git a/metagpt/provider/base_llm.py b/metagpt/provider/base_llm.py index 601980d5e..db2757ec3 100644 --- a/metagpt/provider/base_llm.py +++ b/metagpt/provider/base_llm.py @@ -23,6 +23,7 @@ from tenacity import ( ) from metagpt.configs.llm_config import LLMConfig +from metagpt.const import LLM_API_TIMEOUT, USE_CONFIG_TIMEOUT from metagpt.logs import logger from metagpt.schema import Message from metagpt.utils.common import log_and_reraise @@ -130,7 +131,7 @@ class BaseLLM(ABC): system_msgs: Optional[list[str]] = None, format_msgs: Optional[list[dict[str, str]]] = None, images: Optional[Union[str, list[str]]] = None, - timeout=3, + timeout=USE_CONFIG_TIMEOUT, stream=True, ) -> str: if system_msgs: @@ -146,31 +147,31 @@ class BaseLLM(ABC): else: message.extend(msg) logger.debug(message) - rsp = await self.acompletion_text(message, stream=stream, timeout=timeout) + rsp = await self.acompletion_text(message, stream=stream, timeout=self.get_timeout(timeout)) return rsp def _extract_assistant_rsp(self, context): return "\n".join([i["content"] for i in context if i["role"] == "assistant"]) - async def aask_batch(self, msgs: list, timeout=3) -> str: + async def aask_batch(self, msgs: list, timeout=USE_CONFIG_TIMEOUT) -> str: """Sequential questioning""" context = [] for msg in msgs: umsg = self._user_msg(msg) context.append(umsg) - rsp_text = await self.acompletion_text(context, timeout=timeout) + rsp_text = await self.acompletion_text(context, timeout=self.get_timeout(timeout)) context.append(self._assistant_msg(rsp_text)) return self._extract_assistant_rsp(context) - async def aask_code(self, messages: Union[str, Message, list[dict]], timeout=3, **kwargs) -> dict: + async def aask_code(self, messages: Union[str, Message, list[dict]], timeout=USE_CONFIG_TIMEOUT, **kwargs) -> dict: raise NotImplementedError @abstractmethod - async def _achat_completion(self, messages: list[dict], timeout=3): + async def _achat_completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): """_achat_completion implemented by inherited class""" @abstractmethod - async def acompletion(self, messages: list[dict], timeout=3): + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): """Asynchronous version of completion All GPTAPIs are required to provide the standard OpenAI completion interface [ @@ -181,7 +182,7 @@ class BaseLLM(ABC): """ @abstractmethod - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: """_achat_completion_stream implemented by inherited class""" @retry( @@ -191,11 +192,13 @@ class BaseLLM(ABC): retry=retry_if_exception_type(ConnectionError), retry_error_callback=log_and_reraise, ) - async def acompletion_text(self, messages: list[dict], stream: bool = False, timeout: int = 3) -> str: + async def acompletion_text( + self, messages: list[dict], stream: bool = False, timeout: int = USE_CONFIG_TIMEOUT + ) -> str: """Asynchronous version of completion. Return str. Support stream-print""" if stream: - return await self._achat_completion_stream(messages, timeout=timeout) - resp = await self._achat_completion(messages, timeout=timeout) + return await self._achat_completion_stream(messages, timeout=self.get_timeout(timeout)) + resp = await self._achat_completion(messages, timeout=self.get_timeout(timeout)) return self.get_choice_text(resp) def get_choice_text(self, rsp: dict) -> str: @@ -258,3 +261,6 @@ class BaseLLM(ABC): """Set model and return self. For example, `with_model("gpt-3.5-turbo")`.""" self.config.model = model return self + + def get_timeout(self, timeout: int) -> int: + return timeout or self.config.timeout or LLM_API_TIMEOUT diff --git a/metagpt/provider/dashscope_api.py b/metagpt/provider/dashscope_api.py index 21f3ef351..82224e893 100644 --- a/metagpt/provider/dashscope_api.py +++ b/metagpt/provider/dashscope_api.py @@ -25,6 +25,7 @@ from dashscope.common.error import ( UnsupportedApiProtocol, ) +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream from metagpt.provider.base_llm import BaseLLM, LLMConfig from metagpt.provider.llm_provider_registry import LLMType, register_provider @@ -202,16 +203,16 @@ class DashScopeLLM(BaseLLM): self._update_costs(dict(resp.usage)) return resp.output - async def _achat_completion(self, messages: list[dict], timeout: int = 3) -> GenerationOutput: + async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> GenerationOutput: resp: GenerationResponse = await self.aclient.acall(**self._const_kwargs(messages, stream=False)) self._check_response(resp) self._update_costs(dict(resp.usage)) return resp.output - async def acompletion(self, messages: list[dict], timeout=3) -> GenerationOutput: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> GenerationOutput: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: resp = await self.aclient.acall(**self._const_kwargs(messages, stream=True)) collected_content = [] usage = {} diff --git a/metagpt/provider/general_api_base.py b/metagpt/provider/general_api_base.py index 1b9149396..8e5da8f16 100644 --- a/metagpt/provider/general_api_base.py +++ b/metagpt/provider/general_api_base.py @@ -573,7 +573,7 @@ class APIRequestor: total=request_timeout[1], ) else: - timeout = aiohttp.ClientTimeout(total=request_timeout if request_timeout else TIMEOUT_SECS) + timeout = aiohttp.ClientTimeout(total=request_timeout or TIMEOUT_SECS) if files: # TODO: Use `aiohttp.MultipartWriter` to create the multipart form data here. diff --git a/metagpt/provider/google_gemini_api.py b/metagpt/provider/google_gemini_api.py index 7370747a5..7a35f0a9d 100644 --- a/metagpt/provider/google_gemini_api.py +++ b/metagpt/provider/google_gemini_api.py @@ -2,6 +2,7 @@ # -*- coding: utf-8 -*- # @Desc : Google Gemini LLM from https://ai.google.dev/tutorials/python_quickstart +import os from typing import Optional, Union import google.generativeai as genai @@ -15,7 +16,8 @@ from google.generativeai.types.generation_types import ( ) from metagpt.configs.llm_config import LLMConfig, LLMType -from metagpt.logs import log_llm_stream +from metagpt.const import USE_CONFIG_TIMEOUT +from metagpt.logs import log_llm_stream, logger from metagpt.provider.base_llm import BaseLLM from metagpt.provider.llm_provider_registry import register_provider from metagpt.schema import Message @@ -52,6 +54,10 @@ class GeminiLLM(BaseLLM): self.llm = GeminiGenerativeModel(model_name=self.model) def __init_gemini(self, config: LLMConfig): + if config.proxy: + logger.info(f"Use proxy: {config.proxy}") + os.environ["HTTP_PROXY"] = config.proxy + os.environ["HTTP_PROXYS"] = config.proxy genai.configure(api_key=config.api_key) def _user_msg(self, msg: str, images: Optional[Union[str, list[str]]] = None) -> dict[str, str]: @@ -118,16 +124,18 @@ class GeminiLLM(BaseLLM): self._update_costs(usage) return resp - async def _achat_completion(self, messages: list[dict], timeout: int = 3) -> "AsyncGenerateContentResponse": + async def _achat_completion( + self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT + ) -> "AsyncGenerateContentResponse": resp: AsyncGenerateContentResponse = await self.llm.generate_content_async(**self._const_kwargs(messages)) usage = await self.aget_usage(messages, resp.text) self._update_costs(usage) return resp - async def acompletion(self, messages: list[dict], timeout=3) -> dict: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> dict: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: resp: AsyncGenerateContentResponse = await self.llm.generate_content_async( **self._const_kwargs(messages, stream=True) ) diff --git a/metagpt/provider/human_provider.py b/metagpt/provider/human_provider.py index e5f37c5b9..f205ecd1f 100644 --- a/metagpt/provider/human_provider.py +++ b/metagpt/provider/human_provider.py @@ -6,6 +6,7 @@ Author: garylin2099 from typing import Optional from metagpt.configs.llm_config import LLMConfig +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import logger from metagpt.provider.base_llm import BaseLLM @@ -18,7 +19,7 @@ class HumanProvider(BaseLLM): def __init__(self, config: LLMConfig): pass - def ask(self, msg: str, timeout=3) -> str: + def ask(self, msg: str, timeout=USE_CONFIG_TIMEOUT) -> str: logger.info("It's your turn, please type in your response. You may also refer to the context below") rsp = input(msg) if rsp in ["exit", "quit"]: @@ -31,20 +32,20 @@ class HumanProvider(BaseLLM): system_msgs: Optional[list[str]] = None, format_msgs: Optional[list[dict[str, str]]] = None, generator: bool = False, - timeout=3, + timeout=USE_CONFIG_TIMEOUT, ) -> str: - return self.ask(msg, timeout=timeout) + return self.ask(msg, timeout=self.get_timeout(timeout)) - async def _achat_completion(self, messages: list[dict], timeout=3): + async def _achat_completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): pass - async def acompletion(self, messages: list[dict], timeout=3): + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): """dummy implementation of abstract method in base""" return [] - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: pass - async def acompletion_text(self, messages: list[dict], stream=False, timeout=3) -> str: + async def acompletion_text(self, messages: list[dict], stream=False, timeout=USE_CONFIG_TIMEOUT) -> str: """dummy implementation of abstract method in base""" return "" diff --git a/metagpt/provider/ollama_api.py b/metagpt/provider/ollama_api.py index f65d7e411..2913eb1dd 100644 --- a/metagpt/provider/ollama_api.py +++ b/metagpt/provider/ollama_api.py @@ -5,7 +5,7 @@ import json from metagpt.configs.llm_config import LLMConfig, LLMType -from metagpt.const import LLM_API_TIMEOUT +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream from metagpt.provider.base_llm import BaseLLM from metagpt.provider.general_api_requestor import GeneralAPIRequestor @@ -50,28 +50,28 @@ class OllamaLLM(BaseLLM): chunk = chunk.decode(encoding) return json.loads(chunk) - async def _achat_completion(self, messages: list[dict], timeout: int = 3) -> dict: + async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> dict: resp, _, _ = await self.client.arequest( method=self.http_method, url=self.suffix_url, params=self._const_kwargs(messages), - request_timeout=LLM_API_TIMEOUT, + request_timeout=self.get_timeout(timeout), ) resp = self._decode_and_load(resp) usage = self.get_usage(resp) self._update_costs(usage) return resp - async def acompletion(self, messages: list[dict], timeout=3) -> dict: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> dict: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: stream_resp, _, _ = await self.client.arequest( method=self.http_method, url=self.suffix_url, stream=True, params=self._const_kwargs(messages, stream=True), - request_timeout=LLM_API_TIMEOUT, + request_timeout=self.get_timeout(timeout), ) collected_content = [] diff --git a/metagpt/provider/openai_api.py b/metagpt/provider/openai_api.py index 2fb64dc85..dbfed72df 100644 --- a/metagpt/provider/openai_api.py +++ b/metagpt/provider/openai_api.py @@ -25,6 +25,7 @@ from tenacity import ( ) from metagpt.configs.llm_config import LLMConfig, LLMType +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream, logger from metagpt.provider.base_llm import BaseLLM from metagpt.provider.constant import GENERAL_FUNCTION_SCHEMA @@ -74,9 +75,9 @@ class OpenAILLM(BaseLLM): return params - async def _achat_completion_stream(self, messages: list[dict], timeout=3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> str: response: AsyncStream[ChatCompletionChunk] = await self.aclient.chat.completions.create( - **self._cons_kwargs(messages, timeout=timeout), stream=True + **self._cons_kwargs(messages, timeout=self.get_timeout(timeout)), stream=True ) usage = None collected_messages = [] @@ -104,7 +105,7 @@ class OpenAILLM(BaseLLM): self._update_costs(usage) return full_reply_content - def _cons_kwargs(self, messages: list[dict], timeout=3, **extra_kwargs) -> dict: + def _cons_kwargs(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT, **extra_kwargs) -> dict: kwargs = { "messages": messages, "max_tokens": self._get_max_tokens(messages), @@ -112,20 +113,20 @@ class OpenAILLM(BaseLLM): # "stop": None, # default it's None and gpt4-v can't have this one "temperature": self.config.temperature, "model": self.model, - "timeout": max(self.config.timeout, timeout), + "timeout": self.get_timeout(timeout), } if extra_kwargs: kwargs.update(extra_kwargs) return kwargs - async def _achat_completion(self, messages: list[dict], timeout=3) -> ChatCompletion: - kwargs = self._cons_kwargs(messages, timeout=timeout) + async def _achat_completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> ChatCompletion: + kwargs = self._cons_kwargs(messages, timeout=self.get_timeout(timeout)) rsp: ChatCompletion = await self.aclient.chat.completions.create(**kwargs) self._update_costs(rsp.usage) return rsp - async def acompletion(self, messages: list[dict], timeout=3) -> ChatCompletion: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> ChatCompletion: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) @retry( wait=wait_random_exponential(min=1, max=60), @@ -134,24 +135,24 @@ class OpenAILLM(BaseLLM): retry=retry_if_exception_type(APIConnectionError), retry_error_callback=log_and_reraise, ) - async def acompletion_text(self, messages: list[dict], stream=False, timeout=3) -> str: + async def acompletion_text(self, messages: list[dict], stream=False, timeout=USE_CONFIG_TIMEOUT) -> str: """when streaming, print each token in place.""" if stream: return await self._achat_completion_stream(messages, timeout=timeout) - rsp = await self._achat_completion(messages, timeout=timeout) + rsp = await self._achat_completion(messages, timeout=self.get_timeout(timeout)) return self.get_choice_text(rsp) async def _achat_completion_function( - self, messages: list[dict], timeout: int = 3, **chat_configs + self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT, **chat_configs ) -> ChatCompletion: messages = self.format_msg(messages) - kwargs = self._cons_kwargs(messages=messages, timeout=timeout, **chat_configs) + kwargs = self._cons_kwargs(messages=messages, timeout=self.get_timeout(timeout), **chat_configs) rsp: ChatCompletion = await self.aclient.chat.completions.create(**kwargs) self._update_costs(rsp.usage) return rsp - async def aask_code(self, messages: list[dict], timeout: int = 3, **kwargs) -> dict: + async def aask_code(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT, **kwargs) -> dict: """Use function of tools to ask a code. Note: Keep kwargs consistent with https://platform.openai.com/docs/api-reference/chat/create diff --git a/metagpt/provider/qianfan_api.py b/metagpt/provider/qianfan_api.py index 50916fa3e..3d78c8bfc 100644 --- a/metagpt/provider/qianfan_api.py +++ b/metagpt/provider/qianfan_api.py @@ -9,6 +9,7 @@ from qianfan import ChatCompletion from qianfan.resources.typing import JsonBody from metagpt.configs.llm_config import LLMConfig, LLMType +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream from metagpt.provider.base_llm import BaseLLM from metagpt.provider.llm_provider_registry import register_provider @@ -107,15 +108,15 @@ class QianFanLLM(BaseLLM): self._update_costs(resp.body.get("usage", {})) return resp.body - async def _achat_completion(self, messages: list[dict], timeout: int = 3) -> JsonBody: + async def _achat_completion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> JsonBody: resp = await self.aclient.ado(**self._const_kwargs(messages=messages, stream=False)) self._update_costs(resp.body.get("usage", {})) return resp.body - async def acompletion(self, messages: list[dict], timeout: int = 3) -> JsonBody: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> JsonBody: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: resp = await self.aclient.ado(**self._const_kwargs(messages=messages, stream=True)) collected_content = [] usage = {} diff --git a/metagpt/provider/spark_api.py b/metagpt/provider/spark_api.py index 882c6ce85..594267259 100644 --- a/metagpt/provider/spark_api.py +++ b/metagpt/provider/spark_api.py @@ -17,6 +17,7 @@ from wsgiref.handlers import format_date_time import websocket # 使用websocket_client from metagpt.configs.llm_config import LLMConfig, LLMType +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import logger from metagpt.provider.base_llm import BaseLLM from metagpt.provider.llm_provider_registry import register_provider @@ -31,19 +32,19 @@ class SparkLLM(BaseLLM): def get_choice_text(self, rsp: dict) -> str: return rsp["payload"]["choices"]["text"][-1]["content"] - async def _achat_completion_stream(self, messages: list[dict], timeout: int = 3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout: int = USE_CONFIG_TIMEOUT) -> str: pass - async def acompletion_text(self, messages: list[dict], stream=False, timeout: int = 3) -> str: + async def acompletion_text(self, messages: list[dict], stream=False, timeout: int = USE_CONFIG_TIMEOUT) -> str: # 不支持 # logger.warning("当前方法无法支持异步运行。当你使用acompletion时,并不能并行访问。") w = GetMessageFromWeb(messages, self.config) return w.run() - async def _achat_completion(self, messages: list[dict], timeout=3): + async def _achat_completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): pass - async def acompletion(self, messages: list[dict], timeout=3): + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT): # 不支持异步 w = GetMessageFromWeb(messages, self.config) return w.run() diff --git a/metagpt/provider/zhipuai_api.py b/metagpt/provider/zhipuai_api.py index 14ad1a36b..2db441991 100644 --- a/metagpt/provider/zhipuai_api.py +++ b/metagpt/provider/zhipuai_api.py @@ -8,6 +8,7 @@ from typing import Optional from zhipuai.types.chat.chat_completion import Completion from metagpt.configs.llm_config import LLMConfig, LLMType +from metagpt.const import USE_CONFIG_TIMEOUT from metagpt.logs import log_llm_stream from metagpt.provider.base_llm import BaseLLM from metagpt.provider.llm_provider_registry import register_provider @@ -45,22 +46,22 @@ class ZhiPuAILLM(BaseLLM): kwargs = {"model": self.model, "messages": messages, "stream": stream, "temperature": 0.3} return kwargs - def completion(self, messages: list[dict], timeout=3) -> dict: + def completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> dict: resp: Completion = self.llm.chat.completions.create(**self._const_kwargs(messages)) usage = resp.usage.model_dump() self._update_costs(usage) return resp.model_dump() - async def _achat_completion(self, messages: list[dict], timeout=3) -> dict: + async def _achat_completion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> dict: resp = await self.llm.acreate(**self._const_kwargs(messages)) usage = resp.get("usage", {}) self._update_costs(usage) return resp - async def acompletion(self, messages: list[dict], timeout=3) -> dict: - return await self._achat_completion(messages, timeout=timeout) + async def acompletion(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> dict: + return await self._achat_completion(messages, timeout=self.get_timeout(timeout)) - async def _achat_completion_stream(self, messages: list[dict], timeout=3) -> str: + async def _achat_completion_stream(self, messages: list[dict], timeout=USE_CONFIG_TIMEOUT) -> str: response = await self.llm.acreate_stream(**self._const_kwargs(messages, stream=True)) collected_content = [] usage = {} diff --git a/metagpt/roles/__init__.py b/metagpt/roles/__init__.py index f033a5dfa..08a0406b3 100644 --- a/metagpt/roles/__init__.py +++ b/metagpt/roles/__init__.py @@ -14,7 +14,6 @@ from metagpt.roles.engineer import Engineer from metagpt.roles.qa_engineer import QaEngineer from metagpt.roles.searcher import Searcher from metagpt.roles.sales import Sales -from metagpt.roles.customer_service import CustomerService __all__ = [ @@ -26,5 +25,4 @@ __all__ = [ "QaEngineer", "Searcher", "Sales", - "CustomerService", ] diff --git a/metagpt/utils/common.py b/metagpt/utils/common.py index ee0f74989..e443c3466 100644 --- a/metagpt/utils/common.py +++ b/metagpt/utils/common.py @@ -18,6 +18,7 @@ import csv import importlib import inspect import json +import mimetypes import os import platform import re @@ -29,6 +30,7 @@ from typing import Any, Callable, List, Literal, Tuple, Union from urllib.parse import quote, unquote import aiofiles +import chardet import loguru import requests from PIL import Image @@ -663,14 +665,21 @@ def role_raise_decorator(func): @handle_exception -async def aread(filename: str | Path, encoding=None) -> str: +async def aread(filename: str | Path, encoding="utf-8") -> str: """Read file asynchronously.""" - async with aiofiles.open(str(filename), mode="r", encoding=encoding) as reader: - content = await reader.read() + try: + async with aiofiles.open(str(filename), mode="r", encoding=encoding) as reader: + content = await reader.read() + except UnicodeDecodeError: + async with aiofiles.open(str(filename), mode="rb") as reader: + raw = await reader.read() + result = chardet.detect(raw) + detected_encoding = result["encoding"] + content = raw.decode(detected_encoding) return content -async def awrite(filename: str | Path, data: str, encoding=None): +async def awrite(filename: str | Path, data: str, encoding="utf-8"): """Write file asynchronously.""" pathname = Path(filename) pathname.parent.mkdir(parents=True, exist_ok=True) @@ -811,3 +820,21 @@ See FAQ 5.8 """ ) raise retry_state.outcome.exception() + + +def get_markdown_codeblock_type(filename: str) -> str: + """Return the markdown code-block type corresponding to the file extension.""" + mime_type, _ = mimetypes.guess_type(filename) + mappings = { + "text/x-shellscript": "bash", + "text/x-c++src": "cpp", + "text/css": "css", + "text/html": "html", + "text/x-java": "java", + "application/javascript": "javascript", + "application/json": "json", + "text/x-python": "python", + "text/x-ruby": "ruby", + "application/sql": "sql", + } + return mappings.get(mime_type, "text") diff --git a/metagpt/utils/dependency_file.py b/metagpt/utils/dependency_file.py index d3add1171..0a375051c 100644 --- a/metagpt/utils/dependency_file.py +++ b/metagpt/utils/dependency_file.py @@ -13,9 +13,7 @@ import re from pathlib import Path from typing import Set -import aiofiles - -from metagpt.utils.common import aread +from metagpt.utils.common import aread, awrite from metagpt.utils.exceptions import handle_exception @@ -45,8 +43,7 @@ class DependencyFile: async def save(self): """Save dependencies to the file asynchronously.""" data = json.dumps(self._dependencies) - async with aiofiles.open(str(self._filename), mode="w") as writer: - await writer.write(data) + await awrite(filename=self._filename, data=data) async def update(self, filename: Path | str, dependencies: Set[Path | str], persist=True): """Update dependencies for a file asynchronously. diff --git a/metagpt/utils/file_repository.py b/metagpt/utils/file_repository.py index d2a06963a..d19f2b705 100644 --- a/metagpt/utils/file_repository.py +++ b/metagpt/utils/file_repository.py @@ -14,11 +14,9 @@ from datetime import datetime from pathlib import Path from typing import Dict, List, Set -import aiofiles - from metagpt.logs import logger from metagpt.schema import Document -from metagpt.utils.common import aread +from metagpt.utils.common import aread, awrite from metagpt.utils.json_to_markdown import json_to_markdown @@ -55,8 +53,7 @@ class FileRepository: pathname = self.workdir / filename pathname.parent.mkdir(parents=True, exist_ok=True) content = content if content else "" # avoid `argument must be str, not None` to make it continue - async with aiofiles.open(str(pathname), mode="w") as writer: - await writer.write(content) + await awrite(filename=str(pathname), data=content) logger.info(f"save to: {str(pathname)}") if dependencies is not None: diff --git a/metagpt/utils/mermaid.py b/metagpt/utils/mermaid.py index ae3c5118f..e1d140e84 100644 --- a/metagpt/utils/mermaid.py +++ b/metagpt/utils/mermaid.py @@ -9,11 +9,9 @@ import asyncio import os from pathlib import Path -import aiofiles - from metagpt.config2 import config from metagpt.logs import logger -from metagpt.utils.common import check_cmd_exists +from metagpt.utils.common import awrite, check_cmd_exists async def mermaid_to_file(engine, mermaid_code, output_file_without_suffix, width=2048, height=2048) -> int: @@ -30,9 +28,7 @@ async def mermaid_to_file(engine, mermaid_code, output_file_without_suffix, widt if dir_name and not os.path.exists(dir_name): os.makedirs(dir_name) tmp = Path(f"{output_file_without_suffix}.mmd") - async with aiofiles.open(tmp, "w", encoding="utf-8") as f: - await f.write(mermaid_code) - # tmp.write_text(mermaid_code, encoding="utf-8") + await awrite(filename=tmp, data=mermaid_code) if engine == "nodejs": if check_cmd_exists(config.mermaid.path) != 0: diff --git a/metagpt/utils/repair_llm_raw_output.py b/metagpt/utils/repair_llm_raw_output.py index b8756e8c6..17e095c5f 100644 --- a/metagpt/utils/repair_llm_raw_output.py +++ b/metagpt/utils/repair_llm_raw_output.py @@ -340,7 +340,9 @@ def extract_state_value_from_output(content: str) -> str: content (str): llm's output from `Role._think` """ content = content.strip() # deal the output cases like " 0", "0\n" and so on. - pattern = r"([0-9])" # TODO find the number using a more proper method not just extract from content using pattern + pattern = ( + r"(? 0 else "-1" diff --git a/metagpt/utils/repo_to_markdown.py b/metagpt/utils/repo_to_markdown.py new file mode 100644 index 000000000..76dfe1b82 --- /dev/null +++ b/metagpt/utils/repo_to_markdown.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +This file provides functionality to convert a local repository into a markdown representation. +""" +from __future__ import annotations + +import mimetypes +from pathlib import Path + +from gitignore_parser import parse_gitignore + +from metagpt.logs import logger +from metagpt.utils.common import aread, awrite, get_markdown_codeblock_type, list_files +from metagpt.utils.tree import tree + + +async def repo_to_markdown(repo_path: str | Path, output: str | Path = None, gitignore: str | Path = None) -> str: + """ + Convert a local repository into a markdown representation. + + This function takes a path to a local repository and generates a markdown representation of the repository structure, + including directory trees and file listings. + + Args: + repo_path (str | Path): The path to the local repository. + output (str | Path, optional): The path to save the generated markdown file. Defaults to None. + gitignore (str | Path, optional): The path to the .gitignore file. Defaults to None. + + Returns: + str: The markdown representation of the repository. + """ + repo_path = Path(repo_path) + gitignore = Path(gitignore or Path(__file__).parent / "../../.gitignore").resolve() + + markdown = await _write_dir_tree(repo_path=repo_path, gitignore=gitignore) + + gitignore_rules = parse_gitignore(full_path=str(gitignore)) + markdown += await _write_files(repo_path=repo_path, gitignore_rules=gitignore_rules) + + if output: + await awrite(filename=str(output), data=markdown, encoding="utf-8") + return markdown + + +async def _write_dir_tree(repo_path: Path, gitignore: Path) -> str: + try: + content = tree(repo_path, gitignore, run_command=True) + except Exception as e: + logger.info(f"{e}, using safe mode.") + content = tree(repo_path, gitignore, run_command=False) + + doc = f"## Directory Tree\n```text\n{content}\n```\n---\n\n" + return doc + + +async def _write_files(repo_path, gitignore_rules) -> str: + filenames = list_files(repo_path) + markdown = "" + for filename in filenames: + if gitignore_rules(str(filename)): + continue + markdown += await _write_file(filename=filename, repo_path=repo_path) + return markdown + + +async def _write_file(filename: Path, repo_path: Path) -> str: + relative_path = filename.relative_to(repo_path) + markdown = f"## {relative_path}\n" + + mime_type, _ = mimetypes.guess_type(filename.name) + if "text/" not in mime_type: + logger.info(f"Ignore content: {filename}") + markdown += "\n---\n\n" + return markdown + content = await aread(filename, encoding="utf-8") + content = content.replace("```", "\\`\\`\\`").replace("---", "\\-\\-\\-") + code_block_type = get_markdown_codeblock_type(filename.name) + markdown += f"```{code_block_type}\n{content}\n```\n---\n\n" + return markdown diff --git a/metagpt/utils/tree.py b/metagpt/utils/tree.py new file mode 100644 index 000000000..bd7922290 --- /dev/null +++ b/metagpt/utils/tree.py @@ -0,0 +1,140 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +""" +@Time : 2024/3/11 +@Author : mashenquan +@File : tree.py +@Desc : Implement the same functionality as the `tree` command. + Example: + >>> print_tree(".") + utils + +-- serialize.py + +-- project_repo.py + +-- tree.py + +-- mmdc_playwright.py + +-- cost_manager.py + +-- __pycache__ + | +-- __init__.cpython-39.pyc + | +-- redis.cpython-39.pyc + | +-- singleton.cpython-39.pyc + | +-- embedding.cpython-39.pyc + | +-- make_sk_kernel.cpython-39.pyc + | +-- file_repository.cpython-39.pyc + +-- file.py + +-- save_code.py + +-- common.py + +-- redis.py +""" +from __future__ import annotations + +import subprocess +from pathlib import Path +from typing import Callable, Dict, List + +from gitignore_parser import parse_gitignore + + +def tree(root: str | Path, gitignore: str | Path = None, run_command: bool = False) -> str: + """ + Recursively traverses the directory structure and prints it out in a tree-like format. + + Args: + root (str or Path): The root directory from which to start traversing. + gitignore (str or Path): The filename of gitignore file. + run_command (bool): Whether to execute `tree` command. Execute the `tree` command and return the result if True, + otherwise execute python code instead. + + Returns: + str: A string representation of the directory tree. + + Example: + >>> tree(".") + utils + +-- serialize.py + +-- project_repo.py + +-- tree.py + +-- mmdc_playwright.py + +-- __pycache__ + | +-- __init__.cpython-39.pyc + | +-- redis.cpython-39.pyc + | +-- singleton.cpython-39.pyc + +-- parse_docstring.py + + >>> tree(".", gitignore="../../.gitignore") + utils + +-- serialize.py + +-- project_repo.py + +-- tree.py + +-- mmdc_playwright.py + +-- parse_docstring.py + + >>> tree(".", gitignore="../../.gitignore", run_command=True) + utils + ├── serialize.py + ├── project_repo.py + ├── tree.py + ├── mmdc_playwright.py + └── parse_docstring.py + + + """ + root = Path(root).resolve() + if run_command: + return _execute_tree(root, gitignore) + + git_ignore_rules = parse_gitignore(gitignore) if gitignore else None + dir_ = {root.name: _list_children(root=root, git_ignore_rules=git_ignore_rules)} + v = _print_tree(dir_) + return "\n".join(v) + + +def _list_children(root: Path, git_ignore_rules: Callable) -> Dict[str, Dict]: + dir_ = {} + for i in root.iterdir(): + if git_ignore_rules and git_ignore_rules(str(i)): + continue + try: + if i.is_file(): + dir_[i.name] = {} + else: + dir_[i.name] = _list_children(root=i, git_ignore_rules=git_ignore_rules) + except (FileNotFoundError, PermissionError, OSError): + dir_[i.name] = {} + return dir_ + + +def _print_tree(dir_: Dict[str:Dict]) -> List[str]: + ret = [] + for name, children in dir_.items(): + ret.append(name) + if not children: + continue + lines = _print_tree(children) + for j, v in enumerate(lines): + if v[0] not in ["+", " ", "|"]: + ret = _add_line(ret) + row = f"+-- {v}" + else: + row = f" {v}" + ret.append(row) + return ret + + +def _add_line(rows: List[str]) -> List[str]: + for i in range(len(rows) - 1, -1, -1): + v = rows[i] + if v[0] != " ": + return rows + rows[i] = "|" + v[1:] + return rows + + +def _execute_tree(root: Path, gitignore: str | Path) -> str: + args = ["--gitfile", str(gitignore)] if gitignore else [] + try: + result = subprocess.run(["tree"] + args + [str(root)], capture_output=True, text=True, check=True) + if result.returncode != 0: + raise ValueError(f"tree exits with code {result.returncode}") + return result.stdout + except subprocess.CalledProcessError as e: + raise e diff --git a/requirements.txt b/requirements.txt index 83565278b..2e6cce034 100644 --- a/requirements.txt +++ b/requirements.txt @@ -10,14 +10,6 @@ typer==0.9.0 # godot==0.1.1 # google_api_python_client==2.93.0 # Used by search_engine.py lancedb==0.4.0 -llama-index-core==0.10.15 -llama-index-embeddings-azure-openai==0.1.6 -llama-index-embeddings-openai==0.1.5 -llama-index-llms-azure-openai==0.1.4 -llama-index-readers-file==0.1.4 -llama-index-retrievers-bm25==0.1.3 -llama-index-vector-stores-faiss==0.1.1 -chromadb==0.4.23 loguru==0.6.0 meilisearch==0.21.0 numpy==1.24.3 @@ -34,7 +26,7 @@ PyYAML==6.0.1 # sentence_transformers==2.2.2 setuptools==65.6.3 tenacity==8.2.3 -tiktoken==0.5.2 +tiktoken==0.6.0 tqdm==4.66.2 #unstructured[local-inference] # selenium>4 diff --git a/setup.py b/setup.py index 7a14c6182..f834b4c44 100644 --- a/setup.py +++ b/setup.py @@ -28,6 +28,16 @@ extras_require = { "search-google": ["google-api-python-client==2.94.0"], "search-ddg": ["duckduckgo-search~=4.1.1"], "ocr": ["paddlepaddle==2.4.2", "paddleocr>=2.0.1", "tabulate==0.9.0"], + "rag": [ + "llama-index-core==0.10.15", + "llama-index-embeddings-azure-openai==0.1.6", + "llama-index-embeddings-openai==0.1.5", + "llama-index-llms-azure-openai==0.1.4", + "llama-index-readers-file==0.1.4", + "llama-index-retrievers-bm25==0.1.3", + "llama-index-vector-stores-faiss==0.1.1", + "chromadb==0.4.23", + ], } extras_require["test"] = [ @@ -42,7 +52,6 @@ extras_require["test"] = [ "connexion[uvicorn]~=3.0.5", "azure-cognitiveservices-speech~=1.31.0", "aioboto3~=11.3.0", - "chromadb==0.4.23", "gradio==3.0.0", "grpcio-status==1.48.2", "pylint==3.0.3", @@ -57,7 +66,7 @@ extras_require["dev"] = (["pylint~=3.0.3", "black~=23.3.0", "isort~=5.12.0", "pr setup( name="metagpt", - version="0.7.4", + version="0.7.6", description="The Multi-Agent Framework", long_description=long_description, long_description_content_type="text/markdown", diff --git a/tests/metagpt/roles/test_tutorial_assistant.py b/tests/metagpt/roles/test_tutorial_assistant.py index c12c2b26e..732f346fd 100644 --- a/tests/metagpt/roles/test_tutorial_assistant.py +++ b/tests/metagpt/roles/test_tutorial_assistant.py @@ -6,11 +6,11 @@ @File : test_tutorial_assistant.py """ -import aiofiles import pytest from metagpt.const import TUTORIAL_PATH from metagpt.roles.tutorial_assistant import TutorialAssistant +from metagpt.utils.common import aread @pytest.mark.asyncio @@ -20,9 +20,8 @@ async def test_tutorial_assistant(language: str, topic: str, context): msg = await role.run(topic) assert TUTORIAL_PATH.exists() filename = msg.content - async with aiofiles.open(filename, mode="r", encoding="utf-8") as reader: - content = await reader.read() - assert "pip" in content + content = await aread(filename=filename) + assert "pip" in content if __name__ == "__main__": diff --git a/tests/metagpt/utils/test_common.py b/tests/metagpt/utils/test_common.py index b365f424f..75e8ef4ad 100644 --- a/tests/metagpt/utils/test_common.py +++ b/tests/metagpt/utils/test_common.py @@ -13,7 +13,6 @@ import uuid from pathlib import Path from typing import Any, Set -import aiofiles import pytest from pydantic import BaseModel @@ -125,9 +124,7 @@ class TestGetProjectRoot: async def test_parse_data_exception(self, filename, want): pathname = Path(__file__).parent.parent.parent / "data/output_parser" / filename assert pathname.exists() - async with aiofiles.open(str(pathname), mode="r") as reader: - data = await reader.read() - + data = await aread(filename=pathname) result = OutputParser.parse_data(data=data) assert want in result @@ -198,12 +195,25 @@ class TestGetProjectRoot: @pytest.mark.asyncio async def test_read_write(self): - pathname = Path(__file__).parent / uuid.uuid4().hex / "test.tmp" + pathname = Path(__file__).parent / f"../../../workspace/unittest/{uuid.uuid4().hex}" / "test.tmp" await awrite(pathname, "ABC") data = await aread(pathname) assert data == "ABC" pathname.unlink(missing_ok=True) + @pytest.mark.asyncio + async def test_read_write_error_charset(self): + pathname = Path(__file__).parent / f"../../../workspace/unittest/{uuid.uuid4().hex}" / "test.txt" + content = "中国abc123\u27f6" + await awrite(filename=pathname, data=content) + data = await aread(filename=pathname) + assert data == content + + content = "GB18030 是中国国家标准局发布的新一代中文字符集标准,是 GBK 的升级版,支持更广泛的字符范围。" + await awrite(filename=pathname, data=content, encoding="gb2312") + data = await aread(filename=pathname, encoding="utf-8") + assert data == content + if __name__ == "__main__": pytest.main([__file__, "-s"]) diff --git a/tests/metagpt/utils/test_git_repository.py b/tests/metagpt/utils/test_git_repository.py index ea28b8f0b..480a22e24 100644 --- a/tests/metagpt/utils/test_git_repository.py +++ b/tests/metagpt/utils/test_git_repository.py @@ -10,15 +10,14 @@ import shutil from pathlib import Path -import aiofiles import pytest +from metagpt.utils.common import awrite from metagpt.utils.git_repository import GitRepository async def mock_file(filename, content=""): - async with aiofiles.open(str(filename), mode="w") as file: - await file.write(content) + await awrite(filename=filename, data=content) async def mock_repo(local_path) -> (GitRepository, Path): diff --git a/tests/metagpt/utils/test_repo_to_markdown.py b/tests/metagpt/utils/test_repo_to_markdown.py new file mode 100644 index 000000000..914c50dd7 --- /dev/null +++ b/tests/metagpt/utils/test_repo_to_markdown.py @@ -0,0 +1,25 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +import uuid +from pathlib import Path + +import pytest + +from metagpt.utils.repo_to_markdown import repo_to_markdown + + +@pytest.mark.parametrize( + ["repo_path", "output"], + [(Path(__file__).parent.parent, Path(__file__).parent.parent.parent / f"workspace/unittest/{uuid.uuid4().hex}.md")], +) +@pytest.mark.asyncio +async def test_repo_to_markdown(repo_path: Path, output: Path): + markdown = await repo_to_markdown(repo_path=repo_path, output=output) + assert output.exists() + assert markdown + + output.unlink(missing_ok=True) + + +if __name__ == "__main__": + pytest.main([__file__, "-s"]) diff --git a/tests/metagpt/utils/test_s3.py b/tests/metagpt/utils/test_s3.py index 34c21612c..ef13c2325 100644 --- a/tests/metagpt/utils/test_s3.py +++ b/tests/metagpt/utils/test_s3.py @@ -9,7 +9,6 @@ import uuid from pathlib import Path import aioboto3 -import aiofiles import pytest from metagpt.config2 import Config @@ -46,7 +45,7 @@ async def test_s3(mocker): conn = S3(s3) object_name = "unittest.bak" await conn.upload_file(bucket=s3.bucket, local_path=__file__, object_name=object_name) - pathname = (Path(__file__).parent / uuid.uuid4().hex).with_suffix(".bak") + pathname = (Path(__file__).parent / "../../../workspace/unittest" / uuid.uuid4().hex).with_suffix(".bak") pathname.unlink(missing_ok=True) await conn.download_file(bucket=s3.bucket, object_name=object_name, local_path=str(pathname)) assert pathname.exists() @@ -54,8 +53,7 @@ async def test_s3(mocker): assert url bin_data = await conn.get_object(bucket=s3.bucket, object_name=object_name) assert bin_data - async with aiofiles.open(__file__, mode="r", encoding="utf-8") as reader: - data = await reader.read() + data = await aread(filename=__file__) res = await conn.cache(data, ".bak", "script") assert "http" in res @@ -69,8 +67,6 @@ async def test_s3(mocker): except Exception: pass - await reader.close() - if __name__ == "__main__": pytest.main([__file__, "-s"]) diff --git a/tests/metagpt/utils/test_tree.py b/tests/metagpt/utils/test_tree.py new file mode 100644 index 000000000..03a2a5606 --- /dev/null +++ b/tests/metagpt/utils/test_tree.py @@ -0,0 +1,64 @@ +from pathlib import Path +from typing import List + +import pytest + +from metagpt.utils.tree import _print_tree, tree + + +@pytest.mark.parametrize( + ("root", "rules"), + [ + (str(Path(__file__).parent / "../.."), None), + (str(Path(__file__).parent / "../.."), str(Path(__file__).parent / "../../../.gitignore")), + ], +) +def test_tree(root: str, rules: str): + v = tree(root=root, gitignore=rules) + assert v + + +@pytest.mark.parametrize( + ("root", "rules"), + [ + (str(Path(__file__).parent / "../.."), None), + (str(Path(__file__).parent / "../.."), str(Path(__file__).parent / "../../../.gitignore")), + ], +) +def test_tree_command(root: str, rules: str): + v = tree(root=root, gitignore=rules, run_command=True) + assert v + + +@pytest.mark.parametrize( + ("tree", "want"), + [ + ({"a": {"b": {}, "c": {}}}, ["a", "+-- b", "+-- c"]), + ({"a": {"b": {}, "c": {"d": {}}}}, ["a", "+-- b", "+-- c", " +-- d"]), + ( + {"a": {"b": {"e": {"f": {}, "g": {}}}, "c": {"d": {}}}}, + ["a", "+-- b", "| +-- e", "| +-- f", "| +-- g", "+-- c", " +-- d"], + ), + ( + {"h": {"a": {"b": {"e": {"f": {}, "g": {}}}, "c": {"d": {}}}, "i": {}}}, + [ + "h", + "+-- a", + "| +-- b", + "| | +-- e", + "| | +-- f", + "| | +-- g", + "| +-- c", + "| +-- d", + "+-- i", + ], + ), + ], +) +def test__print_tree(tree: dict, want: List[str]): + v = _print_tree(tree) + assert v == want + + +if __name__ == "__main__": + pytest.main([__file__, "-s"])