[Blocking Issue] Add tests dependencies for qdrant and fix chromadb errors (#435)

* Add tests dependencies for qdrant

* Update chromadb API

* Update chromadb API version

* Fix typehint

* Add py 3.9 condition

* Fix client creation error
This commit is contained in:
Li Jiang 2023-10-27 10:52:38 +08:00 committed by GitHub
parent 50ac547637
commit 27e619e46e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 20 additions and 7 deletions

View File

@ -56,6 +56,10 @@ jobs:
- name: Install packages for Teachable when needed
run: |
pip install -e .[teachable]
- name: Install packages for RetrieveChat with QDrant when needed
if: matrix.python-version == '3.9'
run: |
pip install qdrant_client[fastembed]
- name: Coverage
if: matrix.python-version == '3.9'
env:

View File

@ -18,10 +18,10 @@ except ImportError as e:
class QdrantRetrieveUserProxyAgent(RetrieveUserProxyAgent):
def __init__(
self,
name="RetrieveChatAgent",
human_input_mode: str | None = "ALWAYS",
is_termination_msg: Callable[[Dict], bool] | None = None,
retrieve_config: Dict | None = None,
name="RetrieveChatAgent", # default set to RetrieveChatAgent
human_input_mode: Optional[str] = "ALWAYS",
is_termination_msg: Optional[Callable[[Dict], bool]] = None,
retrieve_config: Optional[Dict] = None, # config for the retrieve agent
**kwargs,
):
"""

View File

@ -361,7 +361,7 @@ class RetrieveUserProxyAgent(UserProxyAgent):
"""
if not self._collection or self._get_or_create:
print("Trying to create collection.")
create_vector_db_from_dir(
self._client = create_vector_db_from_dir(
dir_path=self._docs_path,
max_tokens=self._chunk_token_size,
client=self._client,

View File

@ -5,7 +5,11 @@ from urllib.parse import urlparse
import glob
import tiktoken
import chromadb
from chromadb.api import API
if chromadb.__version__ < "0.4.15":
from chromadb.api import API
else:
from chromadb.api import ClientAPI as API
from chromadb.api.types import QueryResult
import chromadb.utils.embedding_functions as ef
import logging
@ -287,7 +291,7 @@ def create_vector_db_from_dir(
embedding_model: str = "all-MiniLM-L6-v2",
embedding_function: Callable = None,
custom_text_split_function: Callable = None,
):
) -> API:
"""Create a vector db from all the files in a given directory, the directory can also be a single file or a url to
a single file. We support chromadb compatible APIs to create the vector db, this function is not required if
you prepared your own vector db.
@ -307,6 +311,9 @@ def create_vector_db_from_dir(
embedding_function (Optional, Callable): the embedding function to use. Default is None, SentenceTransformer with
the given `embedding_model` will be used. If you want to use OpenAI, Cohere, HuggingFace or other embedding
functions, you can pass it here, follow the examples in `https://docs.trychroma.com/embeddings`.
Returns:
API: the chromadb client.
"""
if client is None:
client = chromadb.PersistentClient(path=db_path)
@ -344,6 +351,7 @@ def create_vector_db_from_dir(
)
except ValueError as e:
logger.warning(f"{e}")
return client
def query_vector_db(

View File

@ -53,6 +53,7 @@ setuptools.setup(
"sympy",
"tiktoken",
"wolframalpha",
"qdrant_client[fastembed]",
],
"blendsearch": ["flaml[blendsearch]"],
"mathchat": ["sympy", "pydantic==1.10.9", "wolframalpha"],