merge (#394)

.github/workflows/on-push-nbconvert-examples.yml

@@ -32,6 +32,6 @@ jobs:
           for notebook in $(find ./examples -name "*.ipynb"); do
             echo "Testing $notebook"
             cd "$(dirname "$notebook")"
-            poetry run python -m nbconvert --to notebook --execute "$(basename "$notebook")" --allow-errors
+            poetry run python -m nbconvert --to notebook --execute "$(basename "$notebook")"
             cd -
           done

automata/agent/__init__.py

@@ -1,3 +1,4 @@
+# sourcery skip: docstrings-for-modules
 from automata.agent.agent import Agent, AgentToolkitBuilder, AgentToolkitNames
 from automata.agent.openai_agent import (
     OpenAIAgentToolkitBuilder,

automata/agent/agent.py

@@ -1,3 +1,4 @@
+"""Defines the abstract base classes and enums for agent objects"""
 import logging
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -45,7 +46,7 @@ def __next__(self) -> LLMIterationResult:
         a new 'asisstant' message, and parsing the reply from the 'user'.
 
         Raises:
-            AgentStopIteration: If the agent has already completed its task
+            AgentStopIterationError: If the agent has already completed its task
             or exceeded the maximum number of iterations.
         """
         pass
@@ -105,6 +106,7 @@ class AgentToolkitNames(Enum):
     SYMBOL_SEARCH = "symbol-search"
     ADVANCED_CONTEXT_ORACLE = "advanced-context-oracle"
     DOCUMENT_ORACLE = "document-oracle"
+
     # Core tools
     PY_READER = "py-reader"
     PY_WRITER = "py-writer"

automata/agent/error.py

@@ -1,3 +1,4 @@
+"""Defines all agent-related exceptions."""
 from automata.core.base import AutomataError
 
 
@@ -7,7 +8,7 @@ class AgentMaxIterError(AutomataError):
     pass
 
 
-class AgentStopIteration(AutomataError):
+class AgentStopIterationError(AutomataError):
     """An exception raised when the agent iteration process terminates."""
 
     pass

automata/agent/openai_agent.py

@@ -1,3 +1,4 @@
+"""Defines the concrete OpenAIAutomataAgent class."""
 import logging
 import uuid
 from abc import ABC, abstractmethod
@@ -9,12 +10,12 @@
     AgentGeneralError,
     AgentMaxIterError,
     AgentResultError,
-    AgentStopIteration,
+    AgentStopIterationError,
 )
-from automata.config import ConfigCategory
-from automata.config.openai_config import OpenAIAutomataAgentConfig
+from automata.config import ConfigCategory, OpenAIAutomataAgentConfig
 from automata.core.utils import format_text, load_config
 from automata.llm import (
+    FunctionCall,
     LLMChatMessage,
     LLMConversation,
     LLMConversationDatabaseProvider,
@@ -25,7 +26,6 @@
     OpenAIFunction,
     OpenAITool,
 )
-from automata.llm.llm_base import FunctionCall
 from automata.tools import ToolExecution, ToolExecutor
 
 logger = logging.getLogger(__name__)
@@ -46,6 +46,7 @@ class OpenAIAutomataAgent(Agent):
     def __init__(
         self, instructions: str, config: OpenAIAutomataAgentConfig
     ) -> None:
+        # sourcery skip: docstrings-for-functions
         super().__init__(instructions)
         self.config = config
         self.iteration_count = 0
@@ -66,7 +67,7 @@ def __next__(self) -> LLMIterationResult:
         assistant and user messages.
 
         Raises:
-            AgentStopIteration: If the agent has already completed its task
+            AgentStopIterationError: If the agent has already completed its task
             or exceeded the maximum number of iterations.
 
         TODO:
@@ -76,7 +77,7 @@ def __next__(self) -> LLMIterationResult:
             self.completed
             or self.iteration_count >= self.config.max_iterations
         ):
-            raise AgentStopIteration
+            raise AgentStopIterationError
 
         logging.debug(f"\n{('-' * 120)}\nLatest Assistant Message -- \n")
         assistant_message = self.chat_provider.get_next_assistant_completion()
@@ -130,7 +131,7 @@ def run(self) -> str:
         or the max iterations are exceeded.
 
         The agent must be setup before running.
-        This implementation calls next() on self until a AgentStopIteration exception is raised,
+        This implementation calls next() on self until a AgentStopIterationError exception is raised,
         at which point it will break out of the loop and return the final result.
 
         Returns:
@@ -147,7 +148,7 @@ def run(self) -> str:
         while True:
             try:
                 next(self)
-            except AgentStopIteration:
+            except AgentStopIterationError:
                 break
 
         last_message = self._conversation.get_latest_message()
@@ -329,7 +330,8 @@ def _setup(self) -> None:
     def _get_termination_tool(self) -> OpenAITool:
         """Gets the tool responsible for terminating the OpenAI agent."""
 
-        def terminate(result: str):
+        def terminate(result: str) -> str:
+            """Terminates the agent run."""
             self.completed = True
             return result
 
@@ -358,6 +360,6 @@ def build_for_open_ai(self) -> List[OpenAITool]:
         pass
 
     @classmethod
-    def can_handle(cls, tool_manager: AgentToolkitNames):
+    def can_handle(cls, tool_manager: AgentToolkitNames) -> bool:
         """Checks if the ToolkitBuilder matches the expecte dtool_manager type"""
         return cls.TOOL_NAME == tool_manager

automata/cli/commands.py

@@ -187,7 +187,13 @@ def run_doc_post_process(ctx: click.Context, *args, **kwargs) -> None:
 )
 @click.pass_context
 def run_agent(ctx: click.Context, *args, **kwargs) -> None:
-    """Run the agent."""
+    """
+    Run the automata agent.
+
+    Ex:
+    poetry run automata run-agent --model="gpt-4" --toolkits="agent-search" --log-level=DEBUG --max-iterations=3 --instructions="Return the documentation for symbol search"
+
+    """
     from automata.cli.scripts.run_agent import main
 
     reconfigure_logging(kwargs.get("log-level", "DEBUG"))
@@ -204,10 +210,8 @@ def run_agent_eval(ctx: click.Context, *args, **kwargs) -> None:
     """
     Run the evaluation.
 
-    Here is an exmaple command -
+    Ex:
     poetry run automata run-agent-eval --evals-filepath=automata/config/eval/primary_agent_payload.json --model="gpt-4" --toolkits="document-oracle,py-reader" --log-level=DEBUG --max-iterations=3
-
-
     """
 
     from automata.cli.scripts.run_agent_eval import main

automata/cli/scripts/run_agent.py

@@ -86,7 +86,7 @@ def main(*args, **kwargs) -> str:
     agent_config_builder = (
         OpenAIAutomataAgentConfigBuilder.from_name(config_name)
         .with_tools(tools)
-        .with_model(kwargs.get("model", "gpt-4-0613"))
+        .with_model(kwargs.get("model", "gpt-4"))
     )
 
     max_iterations = kwargs.get("max_iterations", None)

automata/cli/scripts/run_doc_embedding.py

@@ -12,11 +12,9 @@
 from automata.context_providers.symbol_synchronization_context import (
     SymbolProviderSynchronizationContext,
 )
+from automata.experimental.memory_store import SymbolDocEmbeddingHandler
 from automata.llm import OpenAIEmbeddingProvider
-from automata.memory_store import (
-    SymbolCodeEmbeddingHandler,
-    SymbolDocEmbeddingHandler,
-)
+from automata.memory_store import SymbolCodeEmbeddingHandler
 from automata.singletons.dependency_factory import (
     DependencyFactory,
     dependency_factory,

automata/code_parsers/directory.py

@@ -1,3 +1,4 @@
+"""Defines concrete classes for manipulating directory structures."""
 import logging
 import os
 from typing import Dict, List, Optional
@@ -9,6 +10,7 @@ class Node:
     """Abstract base class for a node in the file tree"""
 
     def __init__(self, name: str, parent: Optional["Node"] = None) -> None:
+        # sourcery skip: docstrings-for-functions
         self.name = name
         self.parent = parent
 
@@ -17,13 +19,15 @@ class File(Node):
     """Represents a file in the tree"""
 
     def __init__(self, name: str, parent: Optional["Node"] = None) -> None:
+        # sourcery skip: docstrings-for-functions
         super().__init__(name, parent)
 
 
 class Directory(Node):
     """Represents a directory. Has children which can be directories or files"""
 
     def __init__(self, name: str, parent: Optional["Node"] = None) -> None:
+        # sourcery skip: docstrings-for-functions
         super().__init__(name, parent)
         self.children: Dict[str, Node] = {}
 
@@ -70,6 +74,7 @@ class DirectoryManager:
     """Handles operations related to directory structure."""
 
     def __init__(self, base_path: str) -> None:
+        # sourcery skip: docstrings-for-functions
         self.root = self._load_directory_structure(base_path)
 
     def _load_directory_structure(self, root_dir: str) -> "Directory":

automata/code_parsers/py/__init__.py

@@ -1,19 +1,8 @@
-from .context_processing.context_handler import (
-    PyContextHandler,
-    PyContextHandlerConfig,
-)
-from .context_processing.context_retriever import (
-    ContextComponent,
-    PyContextRetriever,
-)
-from .dotpath_map import DotPathMap
-from .py_reader import PyReader
+# sourcery skip: docstrings-for-packages
+from automata.code_parsers.py.dotpath_map import DotPathMap
+from automata.code_parsers.py.py_reader import PyReader
 
 __all__ = [
-    "PyContextRetriever",
-    "PyContextHandler",
-    "PyContextHandlerConfig",
-    "PyReader",
-    "ContextComponent",
     "DotPathMap",
+    "PyReader",
 ]

automata/code_parsers/py/dotpath_map.py

@@ -1,3 +1,4 @@
+"""Implements a map from module dotpaths to module filepaths"""
 import os.path
 from typing import Dict, Iterable, Tuple
 
@@ -23,6 +24,7 @@ class DotPathMap:
     DOT_SEP = "."
 
     def __init__(self, path: str, project_name: str) -> None:
+        # sourcery skip: docstrings-for-functions
         # TODO - Test that project_name works when path != local directory name
         self.prefix = project_name.replace(os.pathsep, DotPathMap.DOT_SEP)
         # Remove ending '.' in module fpath

automata/code_parsers/py/py_reader.py

@@ -1,3 +1,4 @@
+"""Defines the core logic for reading python files, e.g. `PyReader`"""
 from __future__ import annotations
 
 import copy

automata/code_writers/py/py_code_writer.py

@@ -1,3 +1,4 @@
+"""Module containing functions to build SymbolDocEmbedding objects."""
 import ast
 import logging
 import subprocess
@@ -12,17 +13,17 @@
 class PyCodeWriter:
     """A utility class for writing Python code along AST nodes"""
 
-    class ModuleNotFound(Exception):
+    class ModuleNotFoundError(Exception):
         """Raised when a module is not found in the module dictionary"""
 
         pass
 
-    class StatementNotFound(Exception):
+    class StatementNotFoundError(Exception):
         """Raised when a provided ast.Statement is not found in the module"""
 
         pass
 
-    class InvalidArguments(Exception):
+    class InvalidArgumentsError(Exception):
         """Raised when invalid arguments are passed to a method"""
 
         pass
@@ -43,11 +44,11 @@ def create_new_module(
         Create a new module object from source code, with option to write to disk.
 
         Raises:
-            PyCodeWriter.InvalidArguments: If the module already exists in the module dictionary.
-            PyCodeWriter.ModuleNotFound: If the module writeout fails.
+            PyCodeWriter.InvalidArgumentsError: If the module already exists in the module dictionary.
+            PyCodeWriter.ModuleNotFoundError: If the module writeout fails.
         """
         if module_dotpath in py_module_loader:
-            raise PyCodeWriter.InvalidArguments(
+            raise PyCodeWriter.InvalidArgumentsError(
                 "Module already exists in module dictionary."
             )
         py_module_loader.put_module(module_dotpath, module)
@@ -58,12 +59,12 @@ def write_module_to_disk(self, module_dotpath: str) -> None:
         """Write the modified module to a file at the specified output path
 
         Raises:
-            ModuleNotFound: If the module is not found in the module dictionary
+            ModuleNotFoundError: If the module is not found in the module dictionary
         """
         if not (
             module_ast := py_module_loader.fetch_ast_module(module_dotpath)
         ):
-            raise PyCodeWriter.ModuleNotFound(
+            raise PyCodeWriter.ModuleNotFoundError(
                 f"Module fpath found in module map for dotpath: {module_dotpath}"
             )
         source_code = ast.unparse(module_ast)
@@ -73,7 +74,7 @@ def write_module_to_disk(self, module_dotpath: str) -> None:
         )
 
         if not module_fpath:
-            raise PyCodeWriter.ModuleNotFound(
+            raise PyCodeWriter.ModuleNotFoundError(
                 f"Module fpath found in module map for dotpath: {module_dotpath}"
             )
         module_fpath = cast(str, module_fpath)
@@ -94,16 +95,14 @@ def upsert_to_module(
         """Upserts the nodes from a new_module into an existing module."""
 
         # For quick lookup, create a dictionary with key as the node name and value as the node.
-        module_dict = {
-            getattr(node, "name", None): node for node in module.body
-        }
+        nodes = {getattr(node, "name", None): node for node in module.body}
 
         for new_node in new_module.body:
             new_node_name = getattr(new_node, "name", None)
 
             # If the new node already exists in the module, remove the old node.
-            if new_node_name in module_dict:
-                module.body.remove(module_dict[new_node_name])
+            if new_node_name in nodes:
+                module.body.remove(nodes[new_node_name])
 
             # Add the new node (either as an insert or an update)
             module.body.append(new_node)
@@ -119,32 +118,30 @@ def delete_from_module(
         """
 
         # For quick lookup, create a dictionary with key as the node name and value as the node.
-        module_dict = {
-            getattr(node, "name", None): node for node in module.body
-        }
+        nodes = {getattr(node, "name", None): node for node in module.body}
 
         for deletion_node in deletion_module.body:
             deletion_node_name = getattr(deletion_node, "name", None)
 
             # If the node to be deleted is not found in the module, raise an exception.
-            if deletion_node_name not in module_dict:
-                raise PyCodeWriter.StatementNotFound(
+            if deletion_node_name not in nodes:
+                raise PyCodeWriter.StatementNotFoundError(
                     f"Node with name '{deletion_node_name}' not found in module"
                 )
 
             # Delete the node from the module body.
-            module.body.remove(module_dict[deletion_node_name])
+            module.body.remove(nodes[deletion_node_name])
 
     def delete_module(self, module_dotpath: str) -> None:
         """
         Create a new module object from source code, with option to write to disk.
 
         Raises:
-            PyCodeWriter.InvalidArguments: If the module already exists in the module dictionary.
-            PyCodeWriter.ModuleNotFound: If the module writeout fails.
+            PyCodeWriter.InvalidArgumentsError: If the module already exists in the module dictionary.
+            PyCodeWriter.ModuleNotFoundError: If the module writeout fails.
         """
         if module_dotpath not in py_module_loader:
-            raise PyCodeWriter.InvalidArguments(
+            raise PyCodeWriter.InvalidArgumentsError(
                 "Module does not exist in module dictionary."
             )
         py_module_loader.delete_module(module_dotpath)

automata/code_writers/py/py_doc_writer.py

@@ -1,3 +1,4 @@
+"""This module contains the PyDocWriter class, which is used to generate"""
 import logging
 import os
 import re