add context and config2

metagpt/utils/cost_manager.py

@@ -80,3 +80,20 @@ class CostManager(BaseModel):
     def get_costs(self) -> Costs:
         """Get all costs"""
         return Costs(self.total_prompt_tokens, self.total_completion_tokens, self.total_cost, self.total_budget)
+
+
+class TokenCostManager(CostManager):
+    """open llm model is self-host, it's free and without cost"""
+
+    def update_cost(self, prompt_tokens, completion_tokens, model):
+        """
+        Update the total cost, prompt tokens, and completion tokens.
+
+        Args:
+        prompt_tokens (int): The number of tokens used in the prompt.
+        completion_tokens (int): The number of tokens used in the completion.
+        model (str): The model used for the API call.
+        """
+        self.total_prompt_tokens += prompt_tokens
+        self.total_completion_tokens += completion_tokens
+        logger.info(f"prompt_tokens: {prompt_tokens}, completion_tokens: {completion_tokens}")

metagpt/utils/embedding.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2024/1/4 20:58
+@Author  : alexanderwu
+@File    : embedding.py
+"""
+from langchain_community.embeddings import OpenAIEmbeddings
+
+from metagpt.config2 import config
+
+
+def get_embedding():
+    llm = config.get_openai_llm()
+    embedding = OpenAIEmbeddings(openai_api_key=llm.api_key, openai_api_base=llm.base_url)
+    return embedding

metagpt/utils/redis.py

@@ -12,26 +12,25 @@ from datetime import timedelta
 
 import aioredis  # https://aioredis.readthedocs.io/en/latest/getting-started/
 
-from metagpt.config import CONFIG
+from metagpt.configs.redis_config import RedisConfig
 from metagpt.logs import logger
 
 
 class Redis:
-    def __init__(self):
+    def __init__(self, config: RedisConfig = None):
+        self.config = config
         self._client = None
 
     async def _connect(self, force=False):
         if self._client and not force:
             return True
-        if not self.is_configured:
-            return False
 
         try:
             self._client = await aioredis.from_url(
-                f"redis://{CONFIG.REDIS_HOST}:{CONFIG.REDIS_PORT}",
-                username=CONFIG.REDIS_USER,
-                password=CONFIG.REDIS_PASSWORD,
-                db=CONFIG.REDIS_DB,
+                self.config.to_url(),
+                username=self.config.username,
+                password=self.config.password,
+                db=self.config.db,
             )
             return True
         except Exception as e:
@@ -62,18 +61,3 @@ class Redis:
             return
         await self._client.close()
         self._client = None
-
-    @property
-    def is_valid(self) -> bool:
-        return self._client is not None
-
-    @property
-    def is_configured(self) -> bool:
-        return bool(
-            CONFIG.REDIS_HOST
-            and CONFIG.REDIS_HOST != "YOUR_REDIS_HOST"
-            and CONFIG.REDIS_PORT
-            and CONFIG.REDIS_PORT != "YOUR_REDIS_PORT"
-            and CONFIG.REDIS_DB is not None
-            and CONFIG.REDIS_PASSWORD is not None
-        )

metagpt/utils/s3.py

@@ -8,7 +8,7 @@ from typing import Optional
 import aioboto3
 import aiofiles
 
-from metagpt.config import CONFIG
+from metagpt.config2 import S3Config
 from metagpt.const import BASE64_FORMAT
 from metagpt.logs import logger
 
@@ -16,13 +16,14 @@ from metagpt.logs import logger
 class S3:
     """A class for interacting with Amazon S3 storage."""
 
-    def __init__(self):
+    def __init__(self, config: S3Config):
         self.session = aioboto3.Session()
+        self.config = config
         self.auth_config = {
             "service_name": "s3",
-            "aws_access_key_id": CONFIG.S3_ACCESS_KEY,
-            "aws_secret_access_key": CONFIG.S3_SECRET_KEY,
-            "endpoint_url": CONFIG.S3_ENDPOINT_URL,
+            "aws_access_key_id": config.access_key,
+            "aws_secret_access_key": config.secret_key,
+            "endpoint_url": config.endpoint,
         }
 
     async def upload_file(
@@ -139,8 +140,8 @@ class S3:
                 data = base64.b64decode(data) if format == BASE64_FORMAT else data.encode(encoding="utf-8")
                 await file.write(data)
 
-            bucket = CONFIG.S3_BUCKET
-            object_pathname = CONFIG.S3_BUCKET or "system"
+            bucket = self.config.bucket
+            object_pathname = self.config.bucket or "system"
             object_pathname += f"/{object_name}"
             object_pathname = os.path.normpath(object_pathname)
             await self.upload_file(bucket=bucket, local_path=str(pathname), object_name=object_pathname)
@@ -151,20 +152,3 @@ class S3:
             logger.exception(f"{e}, stack:{traceback.format_exc()}")
             pathname.unlink(missing_ok=True)
             return None
-
-    @property
-    def is_valid(self):
-        return self.is_configured
-
-    @property
-    def is_configured(self) -> bool:
-        return bool(
-            CONFIG.S3_ACCESS_KEY
-            and CONFIG.S3_ACCESS_KEY != "YOUR_S3_ACCESS_KEY"
-            and CONFIG.S3_SECRET_KEY
-            and CONFIG.S3_SECRET_KEY != "YOUR_S3_SECRET_KEY"
-            and CONFIG.S3_ENDPOINT_URL
-            and CONFIG.S3_ENDPOINT_URL != "YOUR_S3_ENDPOINT_URL"
-            and CONFIG.S3_BUCKET
-            and CONFIG.S3_BUCKET != "YOUR_S3_BUCKET"
-        )

metagpt/utils/yaml_model.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+"""
+@Time    : 2024/1/4 10:18
+@Author  : alexanderwu
+@File    : YamlModel.py
+"""
+from pathlib import Path
+from typing import Dict, Optional
+
+import yaml
+from pydantic import BaseModel, model_validator
+
+
+class YamlModel(BaseModel):
+    extra_fields: Optional[Dict[str, str]] = None
+
+    @classmethod
+    def read_yaml(cls, file_path: Path) -> Dict:
+        with open(file_path, "r") as file:
+            return yaml.safe_load(file)
+
+    @classmethod
+    def model_validate_yaml(cls, file_path: Path) -> "YamlModel":
+        return cls(**cls.read_yaml(file_path))
+
+    def model_dump_yaml(self, file_path: Path) -> None:
+        with open(file_path, "w") as file:
+            yaml.dump(self.model_dump(), file)
+
+
+class YamlModelWithoutDefault(YamlModel):
+    @model_validator(mode="before")
+    @classmethod
+    def check_not_default_config(cls, values):
+        if any(["YOUR" in v for v in values]):
+            raise ValueError("Please set your S3 config in config.yaml")
+        return values