🔧 Fix vectorstore typing

README.md

@@ -6,6 +6,7 @@
 🤖 Megabots provides State-of-the-art, production ready LLM apps made mega-easy, so you don't have to build them from scratch 🤯 Create a bot, now 🫵
 
 👉 Join us on Discord: https://discord.gg/zkqDWk5S7P
+✈️ Work is managed in this project: https://github.com/users/momegas/projects/5/views/2
 
 The Megabots library can be used to create bots that:
 

megabots/__init__.py

@@ -13,7 +13,7 @@ from langchain.prompts import PromptTemplate
 from langchain.chains.question_answering import load_qa_chain
 from langchain.chains.conversational_retrieval.prompts import QA_PROMPT
 from langchain.document_loaders import DirectoryLoader
-from megabots.vectorstores import vectorstore
+from megabots.vectorstores import VectorStore, vectorstore
 
 load_dotenv()
 
@@ -26,8 +26,7 @@ class Bot:
         prompt_variables: list[str] | None = None,
         index: str | None = None,
         sources: bool | None = False,
-        # TODO: Fix this typing
-        vectorstore: Any | None = None,
+        vectorstore: VectorStore | None = None,
         memory: str | None = None,
         verbose: bool = False,
         temperature: int = 0,
@@ -84,7 +83,7 @@ class Bot:
             )
         self.loader = DirectoryLoader(index, recursive=True)
 
-    def load_or_create_index(self, index: str, vectorstore=None):
+    def load_or_create_index(self, index: str, vectorstore: VectorStore | None = None):
         # Load an existing index from disk or create a new one if not available
         if vectorstore is not None:
             self.search_index = vectorstore.client.from_documents(

megabots/vectorstores.py

@@ -10,6 +10,14 @@ class MilvusVectorStore:
         self.client = Milvus
 
 
+class ChromaVectorStore:
+    pass
+
+
+# Generic type variable for all vectorstores
+VectorStore = type("VectorStore", (MilvusVectorStore, ChromaVectorStore), {})
+
+
 SUPPORTED_VECTORSTORES = {
     "milvus": {
         "impl": MilvusVectorStore,
@@ -18,7 +26,7 @@ SUPPORTED_VECTORSTORES = {
 }
 
 
-def vectorstore(name: str) -> MilvusVectorStore:
+def vectorstore(name: str) -> VectorStore:
     """Return a vectorstore object."""
 
     if name is None: