Skip to content

Commit

Permalink
[Blocking Issue] Add tests dependencies for qdrant and fix chromadb e…
Browse files Browse the repository at this point in the history
…rrors (#435)

* Add tests dependencies for qdrant

* Update chromadb API

* Update chromadb API version

* Fix typehint

* Add py 3.9 condition

* Fix client creation error
  • Loading branch information
thinkall authored Oct 27, 2023
1 parent 50ac547 commit 27e619e
Show file tree
Hide file tree
Showing 5 changed files with 20 additions and 7 deletions.
4 changes: 4 additions & 0 deletions .github/workflows/openai.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
8 changes: 4 additions & 4 deletions autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,10 +18,10 @@
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,
):
"""
Expand Down
2 changes: 1 addition & 1 deletion autogen/agentchat/contrib/retrieve_user_proxy_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -361,7 +361,7 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =
"""
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,
Expand Down
12 changes: 10 additions & 2 deletions autogen/retrieve_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,11 @@
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
Expand Down Expand Up @@ -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.
Expand All @@ -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)
Expand Down Expand Up @@ -344,6 +351,7 @@ def create_vector_db_from_dir(
)
except ValueError as e:
logger.warning(f"{e}")
return client


def query_vector_db(
Expand Down
1 change: 1 addition & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
"sympy",
"tiktoken",
"wolframalpha",
"qdrant_client[fastembed]",
],
"blendsearch": ["flaml[blendsearch]"],
"mathchat": ["sympy", "pydantic==1.10.9", "wolframalpha"],
Expand Down

0 comments on commit 27e619e

Please sign in to comment.