Merge pull request llmware-ai#705 from llmware-ai/update-vector-db-imports-to-optional

consistent optional imports for all vector db

Author: doberst

examples/Embedding/embeddings_fast_start.py

@@ -2,13 +2,19 @@
 """
 #   This example shows the general recipe for creating an embedding.  This scenario uses ChromaDB in local
     file mode for no-install laptop deployment.
+
+    NOTE: you may need to install separately:  pip3 install chromadb.
+
 """
 
 
 import os
 from llmware.library import Library
 from llmware.retrieval import Query
 from llmware.setup import Setup
+from importlib import util
+if not util.find_spec("chromadb"):
+    print("\nto run this example with chromadb, you need to install the chromadb python sdk:  pip3 install chromadb")
 
 
 def embeddings_fast_start (library_name, vector_db="chromadb"):
@@ -56,7 +62,7 @@ def embeddings_fast_start (library_name, vector_db="chromadb"):
 
 if __name__ == "__main__":
 
-    #   set to 'chromadb' local file storage for  no-install fast start
+    #   set to 'chromadb' local file storage for no-install fast start
     db = "chromadb"
     embeddings_fast_start("embedding_test_1", vector_db=db)
 

fast_start/example-2-build_embeddings.py

@@ -14,6 +14,8 @@
     Note: we have updated the no-install vector db option to 'chromadb' from 'faiss' starting in
     llmware>=0.2.12, due to better support on Python 3.12
 
+    Note: you may need to install chromadb's python driver:  `pip3 install chromadb`
+
     -- This same basic recipe will work with any of the vector db and collection db by simply changing the name
 
 """
@@ -27,6 +29,10 @@
 from llmware.models import ModelCatalog
 from llmware.configs import LLMWareConfig
 
+from importlib import util
+if not util.find_spec("chromadb"):
+    print("\nto run this example with chromadb, you need to install the chromadb python sdk:  pip3 install chromadb")
+
 
 def setup_library(library_name):
 
@@ -138,7 +144,7 @@ def install_vector_embeddings(library, embedding_model_name):
     #  library = Library().load_library("example1_library")
 
     #  alternatively, to use this example as self-contained, then create a new library from scratch:
-    library = setup_library("example2_lib")
+    library = setup_library("example2_library")
 
     #   Step 2 - Select any embedding model in the LLMWare catalog
 

fast_start/example-5-rag-semantic-query.py

@@ -12,6 +12,8 @@
     3.  Select the most relevant results by document.
     4.  Loop through all of the documents - packaging the context and asking our questions to the LLM.
 
+    NOTE: to use chromadb, you may need to install the python sdk:  pip3 install chromadb.
+
 """
 
 
@@ -22,6 +24,9 @@
 from llmware.status import Status
 from llmware.prompts import Prompt
 from llmware.configs import LLMWareConfig
+from importlib import util
+if not util.find_spec("chromadb"):
+    print("\nto run this example with chromadb, you need to install the chromadb python sdk:  pip3 install chromadb")
 
 
 def semantic_rag (library_name, embedding_model_name, llm_model_name):
@@ -119,13 +124,14 @@ def semantic_rag (library_name, embedding_model_name, llm_model_name):
     #   --if you are using a Python version before 3.12, please feel free to substitute for "faiss"
     #   --for versions of Python >= 3.12, for the Fast Start examples (e.g., no install required), we
     #   recommend using chromadb or lancedb
+
     #   please double-check: `pip3 install chromadb` or pull the latest llmware version to get automatically
     #   -- if you have installed any other vector db, just change the name, e.g, "milvus" or "pg_vector"
 
     vector_db = "chromadb"
 
     # pick any name for the library
-    lib_name = "example5_library"
+    lib_name = "example_5_library"
 
     example_models = ["llmware/bling-1b-0.1", "llmware/bling-tiny-llama-v0", "llmware/dragon-yi-6b-gguf"]
 

llmware/embeddings.py

@@ -11,6 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
 # implied.  See the License for the specific language governing
 # permissions and limitations under the License.
+
 """The embeddings module implements the supported vector databases.
 
 The common abstraction for all supported vector databases is the EmbeddingHandler class, which supports
@@ -26,10 +27,15 @@
 import time
 import uuid
 import itertools
+from importlib import util
 
-from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection
 from pymongo import MongoClient
 
+try:
+    from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection
+except ImportError:
+    pass
+
 try:
     import faiss
 except ImportError:
@@ -384,6 +390,7 @@ def unset_text_index(self):
 
 
 class EmbeddingMilvus:
+
     """Implements the vector database Milvius.
 
     ``EmbeddingMivlus`` implements the interface to the ``Milvus`` vector store. It is used by the
@@ -408,6 +415,7 @@ class EmbeddingMilvus:
     embedding_milvus : EmbeddingMilvus
         A new ``EmbeddingMilvus`` object.
     """
+
     def __init__(self, library, model=None, model_name=None, embedding_dims=None):
 
         self.library = library
@@ -416,6 +424,10 @@ def __init__(self, library, model=None, model_name=None, embedding_dims=None):
         self.milvus_alias = "default"
 
         # Connect to milvus
+        # Instantiate client.
+        if not util.find_spec("pymilvus"):
+            raise DependencyNotInstalledException("pip3 install pymilvus")
+
         connections.connect(self.milvus_alias,
                             host=MilvusConfig.get_config("host"),
                             port=MilvusConfig.get_config("port"),
@@ -2400,6 +2412,9 @@ def __init__(self, library, model=None, model_name=None, embedding_dims=None):
         host = ChromaDBConfig.get_config('host')
 
         # Instantiate client.
+        if not util.find_spec("chromadb"):
+            raise DependencyNotInstalledException("pip3 install chromadb")
+
         if host is None and persistent_path is None:
             self.client = chromadb.EphemeralClient()
 

llmware/requirements.txt

@@ -3,7 +3,6 @@ datasets==2.15.0
 huggingface-hub==0.19.4  
 numpy>=1.23.2
 openai>=1.0
-pymilvus>=2.3.0
 pymongo>=4.7.0
 sentence-transformers==2.2.2
 tabulate==0.9.0
@@ -18,7 +17,6 @@ pgvector==0.2.4
 colorama==0.4.6
 einops==0.7.0
 librosa>=0.10.0
-chromadb>=0.4.22  
 
 requests~=2.31.0
 tqdm~=4.66.1

setup.py

@@ -58,7 +58,6 @@ def glob_fix(package_name, glob):
         'huggingface-hub==0.19.4',
         'numpy>=1.23.2',
         'openai>=1.0.0',
-        'pymilvus>=2.3.0',
         'pymongo>=4.7.0',
         'sentence-transformers==2.2.2',
         'tabulate==0.9.0',
@@ -72,11 +71,12 @@ def glob_fix(package_name, glob):
         'pgvector==0.2.4',
         'colorama==0.4.6',
         'einops==0.7.0',
-        'librosa>=0.10.0',
-        'chromadb>=0.4.22'
+        'librosa>=0.10.0'
     ],
 
     extras_require={
+        'milvus': ['pymilvus>=2.3.0'],
+        'chromadb': ['chromadb>=0.4.22'],
         'pinecone': ['pinecone-client==3.0.0'],
         'lancedb' :['lancedb==0.5.0'],
         'qdrant': ['qdrant-client==1.7.0'],