Merge pull request emrgnt-cmplxty#58 from maks-ivanov/feature/automat…

…a-split-tool-management

Feature/automata split tool management

.gitignore

@@ -1,5 +1,6 @@
 **/sample_modules/**
 interactions.sqlite3
+dump/
 /.idea/
 .env
 local

README.md

@@ -29,7 +29,7 @@ Main Packages and Modules:
    A module which provides the ability to change the in-memory representation of a python module in the PythonIndexer and to
    write this out to disk
 
-6. automata.tools.documentation_tools.documentation_gpt:
+6. automata.tools.documentation.documentation_gpt:
    A simple chatbot that uses DocGPT to answer questions about documentation.
 
 7. automata.tools.oracle.codebase_oracle:

automata/configs/agent_configs/__init__.py

@@ -1 +0,0 @@
-from .automata_config_version import AutomataConfigVersion  # noqa F401

automata/configs/agent_configs/automata_master_v2.yaml

@@ -166,7 +166,7 @@ instruction_template: >
     Action:
     {
       "tool": "automata-writer-modify-module",
-      "input": "Update the file at path tools.documentation_tools.DocumentationGPT so that the method DocumentationGPT.generate_documentation now reads "INSERT_CODE", 
+      "input": "Update the file at path tools.documentation.DocumentationGPT so that the method DocumentationGPT.generate_documentation now reads "INSERT_CODE", 
     }
   '  
 

automata/configs/agent_configs/automata_writer_v2.yaml

@@ -63,7 +63,7 @@ instruction_template: >
     Action:
     {
       "tool": "python-writer-update-module", 
-      "input": "tools.python_tools.documentation_tools.documentation_gpt,DocumentationGPT,def hello_world(self):\n    return "Hello World"\n"
+      "input": "tools.python_tools.documentation.documentation_gpt,DocumentationGPT,def hello_world(self):\n    return "Hello World"\n"
     }
   User:
     Observation:

automata/configs/agent_configs/config_type.py

@@ -0,0 +1,71 @@
+import os
+from enum import Enum
+from typing import Any, Dict, List, Optional
+
+import yaml
+from pydantic import BaseModel
+
+
+class AutomataConfigVersion(Enum):
+    DEFAULT = "default"
+    TEST = "test"
+
+    AUTOMATA_MASTER_V1 = "automata_master_v1"
+    AUTOMATA_INDEXER_V1 = "automata_indexer_v1"
+    AUTOMATA_WRITER_V1 = "automata_writer_v1"
+
+    AUTOMATA_MASTER_V2 = "automata_master_v2"
+    AUTOMATA_INDEXER_V2 = "automata_indexer_v2"
+    AUTOMATA_WRITER_V2 = "automata_writer_v2"
+
+    AUTOMATA_MASTER_V3 = "automata_master_v3"
+
+    AUTOMATA_DOCSTRING_MANAGER_V1 = "automata_docstring_manager_v1"
+
+
+class AutomataAgentConfig(BaseModel):
+
+    """
+    Args:
+        config_version (AutomataConfigVersion): The config_version of the agent to use.
+        initial_payload (Dict[str, str]): Initial payload to send to the agent.
+        llm_toolkits (Dict[ToolkitType, Toolkit]): A dictionary of toolkits to use.
+        instructions (str): A string of instructions to execute.
+        instruction_template (str): A string of instructions to execute.
+        instruction_input_variables (List[str]): A list of input variables for the instruction template.
+        model (str): The model to use for the agent.
+        stream (bool): Whether to stream the results back to the master.
+        verbose (bool): Whether to print the results to stdout.
+        max_iters (int): The maximum number of iterations to run.
+        temperature (float): The temperature to use for the agent.
+        session_id (Optional[str]): The session ID to use for the agent.
+    """
+
+    class Config:
+        SUPPORTED_MDOELS = ["gpt-4", "gpt-3.5-turbo"]
+        arbitrary_types_allowed = True
+
+    config_version: str = "default"
+    initial_payload: Dict[str, str] = {}
+    llm_toolkits: Dict[
+        Any, Any
+    ] = {}  # Dict[ToolkitType, Toolkit], not specified due to circular import
+    instructions: str = ""
+    instruction_template: str = ""
+    instruction_input_variables: List[str] = []
+    model: str = "gpt-4"
+    stream: bool = False
+    verbose: bool = False
+    max_iters: int = 1_000_000
+    temperature: float = 0.7
+    session_id: Optional[str] = None
+
+    @classmethod
+    def load(cls, config_version: AutomataConfigVersion) -> "AutomataAgentConfig":
+        if config_version == AutomataConfigVersion.DEFAULT:
+            return AutomataAgentConfig()
+        file_dir_path = os.path.dirname(os.path.abspath(__file__))
+        config_abs_path = os.path.join(file_dir_path, f"{config_version.value}.yaml")
+        with open(config_abs_path, "r") as file:
+            loaded_yaml = yaml.safe_load(file)
+            return AutomataAgentConfig(**loaded_yaml)

automata/configs/agent_configs/test.yaml

@@ -0,0 +1,9 @@
+initial_payload: {}
+llm_toolkits: {}
+instructions: "Test instructions."
+model: "gpt-4"
+stream: false
+verbose: true
+max_iters: 100
+temperature: 0.8
+session_id: "test-session-id"

automata/core/__init__.py

@@ -1,2 +0,0 @@
-from .base.tool import Tool  # noqa F401
-from .base.tool_utils import Toolkit, ToolkitBuilder, ToolkitType, load_llm_toolkits  # noqa F401

automata/core/agents/automata_agent.py

@@ -6,15 +6,11 @@
     Example:
 
  
-        llm_toolkits = load_llm_toolkits(tools_list, **inputs)
+        llm_toolkits = build_llm_toolkits(tools_list, **inputs)
 
-        initial_payload = {
-            "overview": python_inexer.get_overview(),
-        }
         config_version = AutomataConfigVersion.AUTOMATA_MASTER_V3
         agent_config = AutomataAgentConfig.load(config_version)
         agent = (AutomataAgentBuilder(agent_config)
-            .with_initial_payload(initial_payload)
             .with_llm_toolkits(llm_toolkits)
             .with_instructions(instructions)
             .with_model(model)
@@ -31,64 +27,16 @@
 from typing import Dict, List, Optional, Tuple
 
 import openai
-import yaml
-from pydantic import BaseModel
 from termcolor import colored
 from transformers import GPT2Tokenizer
 
 from automata.config import *  # noqa F403
-from automata.configs.agent_configs import AutomataConfigVersion
-from automata.core import Toolkit, ToolkitType
-from automata.core.utils import format_config_path
+from automata.configs.agent_configs.config_type import AutomataAgentConfig
+from automata.core.base.tool import Toolkit, ToolkitType
 
 logger = logging.getLogger(__name__)
 
 
-class AutomataAgentConfig(BaseModel):
-
-    """
-    Args:
-        config_version (AutomataConfigVersion): The config_version of the agent to use.
-        initial_payload (Dict[str, str]): Initial payload to send to the agent.
-        llm_toolkits (Dict[ToolkitType, Toolkit]): A dictionary of toolkits to use.
-        instructions (str): A string of instructions to execute.
-        instruction_template (str): A string of instructions to execute.
-        instruction_input_variables (List[str]): A list of input variables for the instruction template.
-        model (str): The model to use for the agent.
-        stream (bool): Whether to stream the results back to the master.
-        verbose (bool): Whether to print the results to stdout.
-        max_iters (int): The maximum number of iterations to run.
-        temperature (float): The temperature to use for the agent.
-        session_id (Optional[str]): The session ID to use for the agent.
-    """
-
-    class Config:
-        AGENT_CONFIG_DIRECTORY = "agent_configs"
-        arbitrary_types_allowed = True
-
-    config_version: str = "default"
-    initial_payload: Dict[str, str] = {}
-    llm_toolkits: Dict[ToolkitType, Toolkit] = {}
-    instructions: str = ""
-    instruction_template: str = ""
-    instruction_input_variables: List[str] = []
-    model: str = "gpt-4"
-    stream: bool = False
-    verbose: bool = True
-    max_iters: int = 1_000_000
-    temperature: float = 0.7
-    session_id: Optional[str] = None
-
-    @classmethod
-    def load(cls, config_version: AutomataConfigVersion) -> "AutomataAgentConfig":
-        if config_version == AutomataConfigVersion.DEFAULT:
-            return AutomataAgentConfig()
-        config_path = format_config_path(cls.Config.AGENT_CONFIG_DIRECTORY, config_version.value)
-        with open(f"{config_path}.yaml", "r") as file:
-            loaded_yaml = yaml.safe_load(file)
-            return AutomataAgentConfig(**loaded_yaml)
-
-
 class AutomataAgentBuilder:
     def __init__(self, config: AutomataAgentConfig):
         self._instance = AutomataAgent(config)
@@ -107,25 +55,37 @@ def with_instructions(self, instructions: str):
 
     def with_model(self, model: str):
         self._instance.model = model
+        if model not in AutomataAgentConfig.Config.SUPPORTED_MDOELS:
+            raise ValueError(f"Model {model} not found in Supported OpenAI list of models.")
         return self
 
     def with_stream(self, stream: bool):
+        if not isinstance(stream, bool):
+            raise ValueError("Stream must be a boolean.")
         self._instance.stream = stream
         return self
 
     def with_verbose(self, verbose: bool):
+        if not isinstance(verbose, bool):
+            raise ValueError("Verbose must be a boolean.")
         self._instance.verbose = verbose
         return self
 
     def with_max_iters(self, max_iters: int):
+        if not isinstance(max_iters, int):
+            raise ValueError("Max iters must be an integer.")
         self._instance.max_iters = max_iters
         return self
 
     def with_temperature(self, temperature: float):
+        if not isinstance(temperature, float):
+            raise ValueError("Temperature iters must be a float.")
         self._instance.temperature = temperature
         return self
 
     def with_session_id(self, session_id: Optional[str]):
+        if session_id and not isinstance(session_id, str):
+            raise ValueError("Session Id must be a str.")
         self._instance.session_id = session_id
         return self