feat: Split some main dependencies into optional dependencies

pilot/vector_store/connector.py

@@ -1,9 +1,15 @@
 from pilot.vector_store.chroma_store import ChromaStore
 
-from pilot.vector_store.milvus_store import MilvusStore
 from pilot.vector_store.weaviate_store import WeaviateStore
 
-connector = {"Chroma": ChromaStore, "Milvus": MilvusStore, "Weaviate": WeaviateStore}
+connector = {"Chroma": ChromaStore, "Weaviate": WeaviateStore}
+
+try:
+    from pilot.vector_store.milvus_store import MilvusStore
+
+    connector["Milvus"] = MilvusStore
+except:
+    pass
 
 
 class VectorStoreConnector:

requirements.txt

@@ -5,8 +5,8 @@ async-timeout==4.0.2
 attrs==22.2.0
 cchardet==2.1.7
 chardet==5.1.0
-contourpy==1.0.7
-cycler==0.11.0
+# contourpy==1.0.7
+# cycler==0.11.0
 filelock==3.9.0
 fonttools==4.38.0
 frozenlist==1.3.3
@@ -14,20 +14,20 @@ huggingface-hub==0.14.1
 importlib-resources==5.12.0
 
 sqlparse==0.4.4
-kiwisolver==1.4.4
-matplotlib==3.7.1
+# kiwisolver==1.4.4
+# matplotlib==3.7.1
 multidict==6.0.4
 packaging==23.0
 psutil==5.9.4
-pycocotools==2.0.6
-pyparsing==3.0.9
+# pycocotools==2.0.6
+# pyparsing==3.0.9
 python-dateutil==2.8.2
 pyyaml==6.0
 tokenizers==0.13.2
 tqdm==4.64.1
 transformers>=4.31.0
 transformers_stream_generator
-timm==0.6.13
+# timm==0.6.13
 spacy==3.5.3
 webdataset==0.2.48
 yarl==1.8.2
@@ -40,18 +40,17 @@ peft
 pycocoevalcap
 cpm_kernels
 umap-learn
-notebook
+# notebook
 gradio==3.23
 gradio-client==0.0.8
-wandb
-llama-index==0.5.27
+# wandb
+# llama-index==0.5.27
 
 # TODO move bitsandbytes to optional
 bitsandbytes
 accelerate>=0.20.3
 
 unstructured==0.6.3
-grpcio==1.47.5
 gpt4all==0.3.0
 diskcache==5.6.1
 
@@ -61,7 +60,7 @@ gTTS==2.3.1
 langchain
 nltk
 python-dotenv==1.0.0
-pymilvus==2.2.1
+
 vcrpy
 chromadb==0.3.22
 markdown2
@@ -74,18 +73,7 @@ bardapi==0.1.29
 
 # database
 
+# TODO moved to optional dependencies
 pymysql
 duckdb
-duckdb-engine
-pymssql
-
-# Testing dependencies
-pytest
-asynctest
-pytest-asyncio
-pytest-benchmark
-pytest-cov
-pytest-integration
-pytest-mock
-pytest-recording
-pytesseract==0.3.10
+duckdb-engine

requirements/test-requirements.txt

@@ -0,0 +1,10 @@
+# Testing dependencies
+pytest
+asynctest
+pytest-asyncio
+pytest-benchmark
+pytest-cov
+pytest-integration
+pytest-mock
+pytest-recording
+pytesseract==0.3.10

setup.py

@@ -117,10 +117,30 @@ def llama_cpp_python_cuda_requires():
 
 
 def llama_cpp_requires():
+    """
+    pip install "db-gpt[llama_cpp]"
+    """
     setup_spec.extras["llama_cpp"] = ["llama-cpp-python"]
     llama_cpp_python_cuda_requires()
 
 
+def all_vector_store_requires():
+    """
+    pip install "db-gpt[vstore]"
+    """
+    setup_spec.extras["vstore"] = [
+        "grpcio==1.47.5",  # maybe delete it
+        "pymilvus==2.2.1",
+    ]
+
+
+def all_datasource_requires():
+    """
+    pip install "db-gpt[datasource]"
+    """
+    setup_spec.extras["datasource"] = ["pymssql", "pymysql"]
+
+
 def all_requires():
     requires = set()
     for _, pkgs in setup_spec.extras.items():
@@ -130,11 +150,15 @@ def all_requires():
 
 
 llama_cpp_requires()
+all_vector_store_requires()
+all_datasource_requires()
+
+# must be last
 all_requires()
 
 setuptools.setup(
     name="db-gpt",
-    packages=find_packages(),
+    packages=find_packages(exclude=("tests", "*.tests", "*.tests.*", "examples")),
     version="0.3.5",
     author="csunny",
     author_email="cfqcsunny@gmail.com",