Feature/merged default chroma (#150)

* Change default db to chroma, remove old db setup

* fix test error

* steps towards working code embedding

* Adding embedding dbs for now

* refresh code embeddings

* update code embedding db

* update embedding db

* finalize db update

.gitignore

@@ -60,4 +60,5 @@ local_tasks/
 scripts/setup.sh
 setup.sh
 .direnv
-docs/_build
+docs/_build
+.chroma

automata/agent/providers.py

@@ -165,7 +165,7 @@ def _build_initial_messages(
         assert "user_input_instructions" in instruction_formatter
 
         messages_config = load_config(
-            ConfigCategory.INSTRUCTION.value, self.config.instruction_version.value
+            ConfigCategory.INSTRUCTION.to_path(), self.config.instruction_version.to_path()
         )
         initial_messages = messages_config["initial_messages"]
 

automata/cli/options.py

@@ -20,7 +20,7 @@ def common_options(command: click.Command, *args, **kwargs) -> click.Command:
         ),
         click.option(
             "--index-file",
-            default="index.scip",
+            default="automata.scip",
             help="Which index file to use for the embedding modifications.",
         ),
         click.option(
@@ -30,7 +30,7 @@ def common_options(command: click.Command, *args, **kwargs) -> click.Command:
         ),
         click.option(
             "--doc-embedding-file",
-            default="symbol_doc_embedding_l3.json",
+            default="symbol_doc_embedding_l2.json",
             help="Which embedding file to save to.",
         ),
     ]

automata/cli/scripts/run_agent_config_validation.py

@@ -91,7 +91,9 @@ def test_yaml_compatibility(file_path) -> None:
 
 if __name__ == "__main__":
     # Find all .yaml files in the specified directory
-    yaml_files = glob.glob(os.path.join(get_config_fpath(), ConfigCategory.AGENT.value, "*.yaml"))
+    yaml_files = glob.glob(
+        os.path.join(get_config_fpath(), ConfigCategory.AGENT.to_path(), "*.yaml")
+    )
 
     # Run validation and compatibility tests on each YAML file
     for yaml_file in yaml_files:

automata/cli/scripts/run_code_embedding.py

@@ -1,10 +1,7 @@
 import logging
-import os
 
 from tqdm import tqdm
 
-from automata.config.base import ConfigCategory
-from automata.core.utils import get_config_fpath
 from automata.llm.providers.openai import OpenAIEmbeddingProvider
 from automata.memory_store.symbol_code_embedding import SymbolCodeEmbeddingHandler
 from automata.singletons.dependency_factory import dependency_factory
@@ -22,20 +19,10 @@ def main(*args, **kwargs) -> str:
 
     py_module_loader.initialize()
 
-    scip_fpath = os.path.join(
-        get_config_fpath(), ConfigCategory.SYMBOL.value, kwargs.get("index-file", "index.scip")
-    )
-    code_embedding_fpath = os.path.join(
-        get_config_fpath(),
-        ConfigCategory.SYMBOL.value,
-        kwargs.get("code-embedding-file", "symbol_code_embedding.json"),
-    )
     embedding_provider = OpenAIEmbeddingProvider()
 
     dependency_factory.set_overrides(
         **{
-            "symbol_graph_scip_fpath": scip_fpath,
-            "code_embedding_fpath": code_embedding_fpath,
             "embedding_provider": embedding_provider,
             "disable_synchronization": True,
         }
@@ -55,8 +42,8 @@ def main(*args, **kwargs) -> str:
     for symbol in tqdm(filtered_symbols):
         try:
             symbol_code_embedding_handler.process_embedding(symbol)
-            symbol_code_embedding_handler.embedding_db.save()
         except Exception as e:
             logger.error(f"Failed to update embedding for {symbol.dotpath}: {e}")
 
+    symbol_code_embedding_handler.embedding_db.save()
     return "Success"

automata/cli/scripts/run_doc_embedding.py

@@ -1,21 +1,17 @@
 import logging
-import os
 
 from tqdm import tqdm
 
-from automata.config.base import ConfigCategory
 from automata.context_providers.symbol_synchronization import (
     SymbolProviderSynchronizationContext,
 )
-from automata.core.utils import get_config_fpath
 from automata.llm.providers.openai import OpenAIEmbeddingProvider
 from automata.memory_store.symbol_code_embedding import SymbolCodeEmbeddingHandler
 from automata.memory_store.symbol_doc_embedding import SymbolDocEmbeddingHandler
 from automata.singletons.dependency_factory import dependency_factory
 from automata.singletons.py_module_loader import py_module_loader, pyast_module_loader
 from automata.symbol.graph import SymbolGraph
 from automata.symbol.symbol_utils import get_rankable_symbols
-from automata.symbol_embedding.vector_databases import JSONSymbolEmbeddingVectorDatabase
 
 logger = logging.getLogger(__name__)
 
@@ -24,38 +20,15 @@ def initialize_providers(embedding_level, **kwargs):
     py_module_loader.initialize()
     pyast_module_loader.initialize()
 
-    scip_fpath = os.path.join(
-        get_config_fpath(), ConfigCategory.SYMBOL.value, kwargs.get("index-file", "index.scip")
-    )
-    code_embedding_fpath = os.path.join(
-        get_config_fpath(),
-        ConfigCategory.SYMBOL.value,
-        kwargs.get("code-embedding-file", "symbol_code_embedding.json"),
-    )
-    doc_embedding_fpath = os.path.join(
-        get_config_fpath(),
-        ConfigCategory.SYMBOL.value,
-        kwargs.get("doc-embedding-file", f"symbol_doc_embedding_l{embedding_level}.json"),
-    )
-
     embedding_provider = OpenAIEmbeddingProvider()
 
     overrides = {
-        "symbol_graph_scip_fpath": scip_fpath,
-        "code_embedding_fpath": code_embedding_fpath,
-        "doc_embedding_fpath": doc_embedding_fpath,
         "embedding_provider": embedding_provider,
         "disable_synchronization": True,
     }
 
     if embedding_level == 3:
-        doc_embedding_fpath_l2 = os.path.join(
-            get_config_fpath(),
-            ConfigCategory.SYMBOL.value,
-            kwargs.get("doc-embedding-file", "symbol_doc_embedding_l2.json"),
-        )
-        doc_embedding_db_l2 = JSONSymbolEmbeddingVectorDatabase(doc_embedding_fpath_l2)
-        overrides["py_retriever_doc_embedding_db"] = doc_embedding_db_l2
+        raise NotImplementedError("Embedding level 3 is not supported at this moment.")
 
     dependency_factory.set_overrides(**overrides)
 
@@ -91,6 +64,6 @@ def main(*args, **kwargs) -> str:
     for symbol in tqdm(filtered_symbols):
         logger.info(f"Caching embedding for {symbol}")
         symbol_doc_embedding_handler.process_embedding(symbol)
-        symbol_doc_embedding_handler.embedding_db.save()
+    symbol_doc_embedding_handler.embedding_db.save()
 
     return "Success"

automata/cli/scripts/run_doc_post_process.py

@@ -16,7 +16,7 @@ def main(*args, **kwargs) -> str:
     doc_writer = PyDocWriter(get_root_fpath())
 
     embedding_path = os.path.join(
-        get_config_fpath(), ConfigCategory.SYMBOL.value, "symbol_doc_embedding_l2.json"
+        get_config_fpath(), ConfigCategory.SYMBOL.to_path(), "symbol_doc_embedding_l2.json"
     )
 
     embedding_db = JSONSymbolEmbeddingVectorDatabase(embedding_path)

automata/config/base.py

@@ -6,11 +6,30 @@
 import yaml
 from pydantic import BaseModel, PrivateAttr
 
-from automata.core.utils import convert_kebab_to_snake
+from automata.core.utils import convert_kebab_to_snake_case
 from automata.tools.base import Tool
 
 
-class ConfigCategory(Enum):
+class PathEnum(Enum):
+
+    """An abstract class for enums that represent paths"""
+
+    def to_path(self) -> str:
+        return convert_kebab_to_snake_case(self.value)
+
+
+class EmbeddingDataCategory(PathEnum):
+    """
+    A class to represent the different categories of configuration options
+    Corresponds folders in automata/configs/*
+    """
+
+    CODE_EMBEDDING = "code-embedding"
+    DOC_EMBEDDING = "doc-embedding-l2"
+    INDICES = "indices"
+
+
+class ConfigCategory(PathEnum):
     """
     A class to represent the different categories of configuration options
     Corresponds folders in automata/configs/*
@@ -19,19 +38,19 @@ class ConfigCategory(Enum):
     AGENT = "agent"
     PROMPT = "prompt"
     SYMBOL = "symbol"
-    INSTRUCTION = "instruction_configs"
+    INSTRUCTION = "instruction-configs"
 
 
-class InstructionConfigVersion(Enum):
+class InstructionConfigVersion(PathEnum):
     """
     InstructionConfigVersion: Enum of instruction versions.
     Corresponds files in automata/configs/instruction_configs/*.yaml
     """
 
-    AGENT_INTRODUCTION = "agent_introduction"
+    AGENT_INTRODUCTION = "agent-introduction"
 
 
-class AgentConfigName(Enum):
+class AgentConfigName(PathEnum):
     """
     AgentConfigName: Enum of agent config names.
     Corresponds files in automata/config/agent/*.yaml
@@ -45,7 +64,7 @@ class AgentConfigName(Enum):
     AUTOMATA_MAIN = "automata-main"
 
 
-class LLMProvider(Enum):
+class LLMProvider(PathEnum):
     OPENAI = "openai"
 
 
@@ -84,12 +103,11 @@ def get_llm_provider() -> LLMProvider:
     def _load_automata_yaml_config(cls, config_name: AgentConfigName) -> Dict:
         file_dir_path = os.path.dirname(os.path.abspath(__file__))
         # convert kebab to snake case to support file naming convention
-        config_file_name = convert_kebab_to_snake(config_name.value)
         config_abs_path = os.path.join(
             file_dir_path,
-            ConfigCategory.AGENT.value,
-            cls.get_llm_provider().value,
-            f"{config_file_name}.yaml",
+            ConfigCategory.AGENT.to_path(),
+            cls.get_llm_provider().to_path(),
+            f"{config_name.to_path()}.yaml",
         )
 
         if not os.path.isfile(config_abs_path):

automata/config/openai_agent.py

@@ -77,7 +77,12 @@ def load(cls, config_name: AgentConfigName) -> "OpenAIAutomataAgentConfig":
             return OpenAIAutomataAgentConfig()
 
         loaded_yaml = cls._load_automata_yaml_config(config_name)
-        return OpenAIAutomataAgentConfig(**loaded_yaml)
+        casted_config = OpenAIAutomataAgentConfig(**loaded_yaml)
+        # FIXME - Why is this cast necessary to ensure correct versioning?
+        casted_config.instruction_version = InstructionConfigVersion(
+            casted_config.instruction_version
+        )
+        return casted_config
 
     @staticmethod
     def get_llm_provider() -> LLMProvider:

automata/core/base/database/vector.py

@@ -247,7 +247,7 @@ def load(self) -> None:
 
     def save(self) -> None:
         """As Chroma is a live database, no specific save action is required."""
-        pass
+        self.client.persist()
 
     def clear(self) -> None:
         """Clears all entries in the collection, Use with care!"""

automata/core/utils.py

@@ -31,6 +31,11 @@ def get_root_fpath() -> str:
     return os.path.join(get_root_py_fpath(), "..")
 
 
+def get_embedding_data_fpath() -> str:
+    """Get the path to the root of the Automata config directory."""
+    return os.path.join(get_root_fpath(), "embedding_data")
+
+
 def get_config_fpath() -> str:
     """Get the path to the root of the Automata config directory."""
     return os.path.join(get_root_py_fpath(), "config")
@@ -61,7 +66,7 @@ def format_text(format_variables: Dict[str, str], input_text: str) -> str:
     return input_text
 
 
-def convert_kebab_to_snake(s: str) -> str:
+def convert_kebab_to_snake_case(s: str) -> str:
     """Convert a kebab-case string to snake_case."""
     return s.replace("-", "_")
 

automata/memory_store/symbol_code_embedding.py

@@ -1,9 +1,9 @@
 import logging
 
+from automata.core.base.database.vector import VectorDatabaseProvider
 from automata.symbol.base import Symbol
 from automata.symbol_embedding.builders import SymbolCodeEmbeddingBuilder
 from automata.symbol_embedding.handler import SymbolEmbeddingHandler
-from automata.symbol_embedding.vector_databases import JSONSymbolEmbeddingVectorDatabase
 
 logger = logging.getLogger(__name__)
 
@@ -13,18 +13,14 @@ class SymbolCodeEmbeddingHandler(SymbolEmbeddingHandler):
 
     def __init__(
         self,
-        embedding_db: JSONSymbolEmbeddingVectorDatabase,
+        embedding_db: VectorDatabaseProvider,
         embedding_builder: SymbolCodeEmbeddingBuilder,
     ) -> None:
         super().__init__(embedding_db, embedding_builder)
 
     def process_embedding(self, symbol: Symbol) -> None:
         """Process the embedding for a `Symbol` by updating if the source code has changed."""
         source_code = self.embedding_builder.fetch_embedding_source_code(symbol)
-
-        if not source_code:
-            raise ValueError(f"Symbol {symbol} has no source code")
-
         if self.embedding_db.contains(symbol.dotpath):
             self.update_existing_embedding(source_code, symbol)
         else:
@@ -37,13 +33,12 @@ def update_existing_embedding(self, source_code: str, symbol: Symbol) -> None:
         of the existing embedding. If there are differences, update the embedding.
         """
         existing_embedding = self.embedding_db.get(symbol.dotpath)
+
         if existing_embedding.document != source_code:
-            logger.debug("Building a new embedding for %s", symbol)
             self.embedding_db.discard(symbol.dotpath)
             symbol_embedding = self.embedding_builder.build(source_code, symbol)
             self.embedding_db.add(symbol_embedding)
         elif existing_embedding.symbol != symbol:
-            logger.debug("Updating the embedding for %s", symbol)
             self.embedding_db.discard(symbol.dotpath)
             existing_embedding.symbol = symbol
             self.embedding_db.add(existing_embedding)

automata/memory_store/symbol_doc_embedding.py

@@ -1,9 +1,9 @@
 import logging
 
+from automata.core.base.database.vector import VectorDatabaseProvider
 from automata.symbol.base import Symbol, SymbolDescriptor
 from automata.symbol_embedding.builders import SymbolDocEmbeddingBuilder
 from automata.symbol_embedding.handler import SymbolEmbeddingHandler
-from automata.symbol_embedding.vector_databases import JSONSymbolEmbeddingVectorDatabase
 
 logger = logging.getLogger(__name__)
 
@@ -13,7 +13,7 @@ class SymbolDocEmbeddingHandler(SymbolEmbeddingHandler):
 
     def __init__(
         self,
-        embedding_db: JSONSymbolEmbeddingVectorDatabase,
+        embedding_db: VectorDatabaseProvider,
         embedding_builder: SymbolDocEmbeddingBuilder,
     ) -> None:
         super().__init__(embedding_db, embedding_builder)
@@ -35,10 +35,10 @@ def process_embedding(self, symbol: Symbol) -> None:
             return
         if symbol.symbol_kind_by_suffix() == SymbolDescriptor.PyKind.Class:
             symbol_embedding = self.embedding_builder.build(source_code, symbol)
-        else:
-            if not isinstance(self.embedding_builder, SymbolDocEmbeddingBuilder):
-                raise ValueError("SymbolDocEmbeddingHandler requires a SymbolDocEmbeddingBuilder")
+        elif isinstance(self.embedding_builder, SymbolDocEmbeddingBuilder):
             symbol_embedding = self.embedding_builder.build_non_class(source_code, symbol)
+        else:
+            raise ValueError("SymbolDocEmbeddingHandler requires a SymbolDocEmbeddingBuilder")
 
         self.embedding_db.add(symbol_embedding)