Kaynağa Gözat

Refactor: Use enum for config keys (#1376)

Boxuan Li 1 yıl önce
ebeveyn
işleme
831e934dab

+ 11 - 10
agenthub/monologue_agent/utils/memory.py

@@ -11,11 +11,12 @@ from openai._exceptions import APIConnectionError, RateLimitError, InternalServe
 
 from opendevin import config
 from opendevin.logger import opendevin_logger as logger
+from opendevin.schema.config import ConfigType
 from . import json
 
-num_retries = config.get('LLM_NUM_RETRIES')
-retry_min_wait = config.get('LLM_RETRY_MIN_WAIT')
-retry_max_wait = config.get('LLM_RETRY_MAX_WAIT')
+num_retries = config.get(ConfigType.LLM_NUM_RETRIES)
+retry_min_wait = config.get(ConfigType.LLM_RETRY_MIN_WAIT)
+retry_max_wait = config.get(ConfigType.LLM_RETRY_MAX_WAIT)
 
 # llama-index includes a retry decorator around openai.get_embeddings() function
 # it is initialized with hard-coded values and errors
@@ -46,7 +47,7 @@ def wrapper_get_embeddings(*args, **kwargs):
 
 llama_openai.get_embeddings = wrapper_get_embeddings
 
-embedding_strategy = config.get('LLM_EMBEDDING_MODEL')
+embedding_strategy = config.get(ConfigType.LLM_EMBEDDING_MODEL)
 
 # TODO: More embeddings: https://docs.llamaindex.ai/en/stable/examples/embeddings/OpenAI/
 # There's probably a more programmatic way to do this.
@@ -54,24 +55,24 @@ if embedding_strategy == 'llama2':
     from llama_index.embeddings.ollama import OllamaEmbedding
     embed_model = OllamaEmbedding(
         model_name='llama2',
-        base_url=config.get('LLM_BASE_URL', required=True),
+        base_url=config.get(ConfigType.LLM_BASE_URL, required=True),
         ollama_additional_kwargs={'mirostat': 0},
     )
 elif embedding_strategy == 'openai':
     from llama_index.embeddings.openai import OpenAIEmbedding
     embed_model = OpenAIEmbedding(
         model='text-embedding-ada-002',
-        api_key=config.get('LLM_API_KEY', required=True)
+        api_key=config.get(ConfigType.LLM_API_KEY, required=True)
     )
 elif embedding_strategy == 'azureopenai':
     # Need to instruct to set these env variables in documentation
     from llama_index.embeddings.azure_openai import AzureOpenAIEmbedding
     embed_model = AzureOpenAIEmbedding(
         model='text-embedding-ada-002',
-        deployment_name=config.get('LLM_EMBEDDING_DEPLOYMENT_NAME', required=True),
-        api_key=config.get('LLM_API_KEY', required=True),
-        azure_endpoint=config.get('LLM_BASE_URL', required=True),
-        api_version=config.get('LLM_API_VERSION', required=True),
+        deployment_name=config.get(ConfigType.LLM_EMBEDDING_DEPLOYMENT_NAME, required=True),
+        api_key=config.get(ConfigType.LLM_API_KEY, required=True),
+        azure_endpoint=config.get(ConfigType.LLM_BASE_URL, required=True),
+        api_version=config.get(ConfigType.LLM_API_VERSION, required=True),
     )
 elif (embedding_strategy is not None) and (embedding_strategy.lower() == 'none'):
     # TODO: this works but is not elegant enough. The incentive is when

+ 2 - 1
agenthub/monologue_agent/utils/prompts.py

@@ -14,6 +14,7 @@ from opendevin.observation import (
 )
 from opendevin.exceptions import LLMOutputError
 from opendevin import config
+from opendevin.schema.config import ConfigType
 
 ACTION_PROMPT = """
 You're a thoughtful robot. Your main task is this:
@@ -146,7 +147,7 @@ def get_request_action_prompt(
         'monologue': json.dumps(thoughts, indent=2),
         'background_commands': bg_commands_message,
         'hint': hint,
-        'WORKSPACE_MOUNT_PATH_IN_SANDBOX': config.get('WORKSPACE_MOUNT_PATH_IN_SANDBOX'),
+        'WORKSPACE_MOUNT_PATH_IN_SANDBOX': config.get(ConfigType.WORKSPACE_MOUNT_PATH_IN_SANDBOX),
     }
 
 

+ 2 - 1
opendevin/action/fileop.py

@@ -13,6 +13,7 @@ from opendevin.observation import (
 from opendevin.schema import ActionType
 from opendevin.sandbox import E2BBox
 from opendevin import config
+from opendevin.schema.config import ConfigType
 
 from .base import ExecutableAction
 
@@ -32,7 +33,7 @@ def resolve_path(file_path):
     path_in_workspace = abs_path_in_sandbox.relative_to(Path(SANDBOX_PATH_PREFIX))
 
     # Get path relative to host
-    path_in_host_workspace = Path(config.get('WORKSPACE_BASE')) / path_in_workspace
+    path_in_host_workspace = Path(config.get(ConfigType.WORKSPACE_BASE)) / path_in_workspace
 
     return path_in_host_workspace
 

+ 3 - 1
opendevin/config.py

@@ -159,10 +159,12 @@ def finalize_config():
 finalize_config()
 
 
-def get(key: str, required: bool = False):
+def get(key: ConfigType, required: bool = False):
     """
     Get a key from the environment variables or config.toml or default configs.
     """
+    if not isinstance(key, ConfigType):
+        raise ValueError(f"key '{key}' must be an instance of ConfigType Enum")
     value = config.get(key)
     if not value and required:
         raise KeyError(f"Please set '{key}' in `config.toml` or `.env`.")

+ 3 - 2
opendevin/controller/agent_controller.py

@@ -3,6 +3,7 @@ from typing import Callable, List, Type
 
 
 from opendevin import config
+from opendevin.schema.config import ConfigType
 from opendevin.action import (
     Action,
     AgentFinishAction,
@@ -25,8 +26,8 @@ from opendevin.action.tasks import TaskStateChangedAction
 from opendevin.schema import TaskState
 from opendevin.controller.action_manager import ActionManager
 
-MAX_ITERATIONS = config.get('MAX_ITERATIONS')
-MAX_CHARS = config.get('MAX_CHARS')
+MAX_ITERATIONS = config.get(ConfigType.MAX_ITERATIONS)
+MAX_CHARS = config.get(ConfigType.MAX_CHARS)
 
 
 class AgentController:

+ 8 - 7
opendevin/llm/llm.py

@@ -4,17 +4,18 @@ from litellm.exceptions import APIConnectionError, RateLimitError, ServiceUnavai
 from functools import partial
 
 from opendevin import config
+from opendevin.schema.config import ConfigType
 from opendevin.logger import llm_prompt_logger, llm_response_logger
 from opendevin.logger import opendevin_logger as logger
 
 
-DEFAULT_API_KEY = config.get('LLM_API_KEY')
-DEFAULT_BASE_URL = config.get('LLM_BASE_URL')
-DEFAULT_MODEL_NAME = config.get('LLM_MODEL')
-DEFAULT_API_VERSION = config.get('LLM_API_VERSION')
-LLM_NUM_RETRIES = config.get('LLM_NUM_RETRIES')
-LLM_RETRY_MIN_WAIT = config.get('LLM_RETRY_MIN_WAIT')
-LLM_RETRY_MAX_WAIT = config.get('LLM_RETRY_MAX_WAIT')
+DEFAULT_API_KEY = config.get(ConfigType.LLM_API_KEY)
+DEFAULT_BASE_URL = config.get(ConfigType.LLM_BASE_URL)
+DEFAULT_MODEL_NAME = config.get(ConfigType.LLM_MODEL)
+DEFAULT_API_VERSION = config.get(ConfigType.LLM_API_VERSION)
+LLM_NUM_RETRIES = config.get(ConfigType.LLM_NUM_RETRIES)
+LLM_RETRY_MIN_WAIT = config.get(ConfigType.LLM_RETRY_MIN_WAIT)
+LLM_RETRY_MAX_WAIT = config.get(ConfigType.LLM_RETRY_MAX_WAIT)
 
 
 class LLM:

+ 3 - 1
opendevin/logger.py

@@ -7,8 +7,10 @@ from opendevin import config
 from typing import Literal, Mapping
 from termcolor import colored
 
+from opendevin.schema.config import ConfigType
+
 DISABLE_COLOR_PRINTING = (
-    config.get('DISABLE_COLOR').lower() == 'true'
+    config.get(ConfigType.DISABLE_COLOR).lower() == 'true'
 )
 
 ColorType = Literal[

+ 3 - 3
opendevin/sandbox/docker/exec_box.py

@@ -27,9 +27,9 @@ SANDBOX_WORKSPACE_DIR = config.get(ConfigType.WORKSPACE_MOUNT_PATH_IN_SANDBOX)
 
 # FIXME: On some containers, the devin user doesn't have enough permission, e.g. to install packages
 # How do we make this more flexible?
-RUN_AS_DEVIN = config.get('RUN_AS_DEVIN').lower() != 'false'
+RUN_AS_DEVIN = config.get(ConfigType.RUN_AS_DEVIN).lower() != 'false'
 USER_ID = 1000
-if SANDBOX_USER_ID := config.get('SANDBOX_USER_ID'):
+if SANDBOX_USER_ID := config.get(ConfigType.SANDBOX_USER_ID):
     USER_ID = int(SANDBOX_USER_ID)
 elif hasattr(os, 'getuid'):
     USER_ID = os.getuid()
@@ -225,7 +225,7 @@ class DockerExecBox(Sandbox):
 
         try:
             # start the container
-            mount_dir = config.get('WORKSPACE_MOUNT_PATH')
+            mount_dir = config.get(ConfigType.WORKSPACE_MOUNT_PATH)
             self.container = self.docker_client.containers.run(
                 self.container_image,
                 command='tail -f /dev/null',

+ 6 - 5
opendevin/sandbox/docker/local_box.py

@@ -5,6 +5,7 @@ from opendevin.sandbox.sandbox import Sandbox
 from opendevin.sandbox.process import Process
 from opendevin.sandbox.docker.process import DockerProcess
 from opendevin import config
+from opendevin.schema.config import ConfigType
 
 # ===============================================================================
 #  ** WARNING **
@@ -33,7 +34,7 @@ class LocalBox(Sandbox):
         try:
             completed_process = subprocess.run(
                 cmd, shell=True, text=True, capture_output=True,
-                timeout=self.timeout, cwd=config.get('WORKSPACE_BASE')
+                timeout=self.timeout, cwd=config.get(ConfigType.WORKSPACE_BASE)
             )
             return completed_process.returncode, completed_process.stdout.strip()
         except subprocess.TimeoutExpired:
@@ -41,19 +42,19 @@ class LocalBox(Sandbox):
 
     def copy_to(self, host_src: str, sandbox_dest: str, recursive: bool = False):
         # mkdir -p sandbox_dest if it doesn't exist
-        res = subprocess.run(f'mkdir -p {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE'))
+        res = subprocess.run(f'mkdir -p {sandbox_dest}', shell=True, text=True, cwd=config.get(ConfigType.WORKSPACE_BASE))
         if res.returncode != 0:
             raise RuntimeError(f'Failed to create directory {sandbox_dest} in sandbox')
 
         if recursive:
             res = subprocess.run(
-                f'cp -r {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE')
+                f'cp -r {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get(ConfigType.WORKSPACE_BASE)
             )
             if res.returncode != 0:
                 raise RuntimeError(f'Failed to copy {host_src} to {sandbox_dest} in sandbox')
         else:
             res = subprocess.run(
-                f'cp {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get('WORKSPACE_BASE')
+                f'cp {host_src} {sandbox_dest}', shell=True, text=True, cwd=config.get(ConfigType.WORKSPACE_BASE)
             )
             if res.returncode != 0:
                 raise RuntimeError(f'Failed to copy {host_src} to {sandbox_dest} in sandbox')
@@ -61,7 +62,7 @@ class LocalBox(Sandbox):
     def execute_in_background(self, cmd: str) -> Process:
         process = subprocess.Popen(
             cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
-            text=True, cwd=config.get('WORKSPACE_BASE')
+            text=True, cwd=config.get(ConfigType.WORKSPACE_BASE)
         )
         bg_cmd = DockerProcess(
             id=self.cur_background_id, command=cmd, result=process, pid=process.pid

+ 2 - 2
opendevin/sandbox/docker/ssh_box.py

@@ -34,9 +34,9 @@ USE_HOST_NETWORK = config.get(ConfigType.USE_HOST_NETWORK)
 
 # FIXME: On some containers, the devin user doesn't have enough permission, e.g. to install packages
 # How do we make this more flexible?
-RUN_AS_DEVIN = config.get('RUN_AS_DEVIN').lower() != 'false'
+RUN_AS_DEVIN = config.get(ConfigType.RUN_AS_DEVIN).lower() != 'false'
 USER_ID = 1000
-if SANDBOX_USER_ID := config.get('SANDBOX_USER_ID'):
+if SANDBOX_USER_ID := config.get(ConfigType.SANDBOX_USER_ID):
     USER_ID = int(SANDBOX_USER_ID)
 elif hasattr(os, 'getuid'):
     USER_ID = os.getuid()

+ 2 - 1
opendevin/sandbox/e2b/sandbox.py

@@ -8,6 +8,7 @@ from e2b.sandbox.exception import (
 )
 
 from opendevin import config
+from opendevin.schema.config import ConfigType
 from opendevin.logger import opendevin_logger as logger
 from opendevin.sandbox.sandbox import Sandbox
 from opendevin.sandbox.e2b.process import E2BProcess
@@ -26,7 +27,7 @@ class E2BBox(Sandbox):
         timeout: int = 120,
     ):
         self.sandbox = E2BSandbox(
-            api_key=config.get('E2B_API_KEY'),
+            api_key=config.get(ConfigType.E2B_API_KEY),
             template=template,
             # It's possible to stream stdout and stderr from sandbox and from each process
             on_stderr=lambda x: logger.info(f'E2B sandbox stderr: {x}'),

+ 1 - 0
opendevin/schema/config.py

@@ -23,6 +23,7 @@ class ConfigType(str, Enum):
     AGENT = 'AGENT'
     E2B_API_KEY = 'E2B_API_KEY'
     SANDBOX_TYPE = 'SANDBOX_TYPE'
+    SANDBOX_USER_ID = 'SANDBOX_USER_ID'
     USE_HOST_NETWORK = 'USE_HOST_NETWORK'
     SSH_HOSTNAME = 'SSH_HOSTNAME'
     DISABLE_COLOR = 'DISABLE_COLOR'

+ 3 - 2
opendevin/server/listen.py

@@ -11,6 +11,7 @@ from fastapi.staticfiles import StaticFiles
 
 import agenthub  # noqa F401 (we import this to get the agents registered)
 from opendevin import config, files
+from opendevin.schema.config import ConfigType
 from opendevin.agent import Agent
 from opendevin.logger import opendevin_logger as logger
 from opendevin.server.agent import agent_manager
@@ -113,14 +114,14 @@ async def del_messages(
 
 @app.get('/api/refresh-files')
 def refresh_files():
-    structure = files.get_folder_structure(Path(str(config.get('WORKSPACE_BASE'))))
+    structure = files.get_folder_structure(Path(str(config.get(ConfigType.WORKSPACE_BASE))))
     return structure.to_dict()
 
 
 @app.get('/api/select-file')
 def select_file(file: str):
     try:
-        workspace_base = config.get('WORKSPACE_BASE')
+        workspace_base = config.get(ConfigType.WORKSPACE_BASE)
         file_path = Path(workspace_base, file)
         # The following will check if the file is within the workspace base and throw an exception if not
         file_path.resolve().relative_to(Path(workspace_base).resolve())

+ 10 - 7
tests/test_fileops.py

@@ -1,18 +1,21 @@
-from opendevin import config
-from opendevin.action import fileop
 from pathlib import Path
+
 import pytest
 
+from opendevin import config
+from opendevin.schema.config import ConfigType
+from opendevin.action import fileop
+
 
 def test_resolve_path():
-    assert fileop.resolve_path('test.txt') == Path(config.get('WORKSPACE_BASE')) / 'test.txt'
-    assert fileop.resolve_path('subdir/test.txt') == Path(config.get('WORKSPACE_BASE')) / 'subdir' / 'test.txt'
+    assert fileop.resolve_path('test.txt') == Path(config.get(ConfigType.WORKSPACE_BASE)) / 'test.txt'
+    assert fileop.resolve_path('subdir/test.txt') == Path(config.get(ConfigType.WORKSPACE_BASE)) / 'subdir' / 'test.txt'
     assert fileop.resolve_path(Path(fileop.SANDBOX_PATH_PREFIX) / 'test.txt') == \
-        Path(config.get('WORKSPACE_BASE')) / 'test.txt'
+        Path(config.get(ConfigType.WORKSPACE_BASE)) / 'test.txt'
     assert fileop.resolve_path(Path(fileop.SANDBOX_PATH_PREFIX) / 'subdir' / 'test.txt') == \
-        Path(config.get('WORKSPACE_BASE')) / 'subdir' / 'test.txt'
+        Path(config.get(ConfigType.WORKSPACE_BASE)) / 'subdir' / 'test.txt'
     assert fileop.resolve_path(Path(fileop.SANDBOX_PATH_PREFIX) / 'subdir' / '..' / 'test.txt') == \
-        Path(config.get('WORKSPACE_BASE')) / 'test.txt'
+        Path(config.get(ConfigType.WORKSPACE_BASE)) / 'test.txt'
     with pytest.raises(PermissionError):
         fileop.resolve_path(Path(fileop.SANDBOX_PATH_PREFIX) / '..' / 'test.txt')
     with pytest.raises(PermissionError):