Skip to content

Commit

Permalink
Support no chromadb
Browse files Browse the repository at this point in the history
  • Loading branch information
thinkall committed Nov 5, 2023
1 parent 0dd0fc5 commit 64e401d
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 17 deletions.
19 changes: 11 additions & 8 deletions autogen/agentchat/contrib/retrieve_user_proxy_agent.py
Original file line number Diff line number Diff line change
@@ -1,18 +1,12 @@
import re

try:
import chromadb
except ImportError:
raise ImportError("Please install dependencies first. `pip install pyautogen[retrievechat]`")
from typing import Callable, Dict, Optional, Union, List, Tuple, Any
from IPython import get_ipython
from autogen.agentchat.agent import Agent
from autogen.agentchat import UserProxyAgent
from autogen.retrieve_utils import create_vector_db_from_dir, query_vector_db
from autogen.token_count_utils import count_token
from autogen.code_utils import extract_code

from typing import Callable, Dict, Optional, Union, List, Tuple, Any
from IPython import get_ipython

try:
from termcolor import colored
except ImportError:
Expand All @@ -21,6 +15,15 @@ def colored(x, *args, **kwargs):
return x


try:
import chromadb
except ImportError:

class chromadb:
class Client:
pass


PROMPT_DEFAULT = """You're a retrieve augmented chatbot. You answer user's questions based on your own knowledge and the
context provided by the user. You should follow the following steps to answer a question:
Step 1, you estimate the user's intent based on the question and context. The intent can be a code generation task or
Expand Down
36 changes: 27 additions & 9 deletions autogen/retrieve_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,18 +3,30 @@
import requests
from urllib.parse import urlparse
import glob
import chromadb

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
import pypdf
from autogen.token_count_utils import count_token

try:
import chromadb

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

HAS_CHROMADB = True
except ImportError:
HAS_CHROMADB = False

class API:
pass

class QueryResult(dict):
pass


logger = logging.getLogger(__name__)
TEXT_FORMATS = [
"txt",
Expand Down Expand Up @@ -88,6 +100,8 @@ def split_text_to_chunks(

def extract_text_from_pdf(file: str) -> str:
"""Extract text from PDF files"""
import pypdf

text = ""
with open(file, "rb") as f:
reader = pypdf.PdfReader(f)
Expand Down Expand Up @@ -240,6 +254,8 @@ def create_vector_db_from_dir(
Returns:
API: the chromadb client.
"""
if not HAS_CHROMADB:
raise ImportError("Please install dependencies first. `pip install pyautogen[retrievechat]`")
if client is None:
client = chromadb.PersistentClient(path=db_path)
try:
Expand Down Expand Up @@ -314,6 +330,8 @@ class QueryResult(TypedDict):
metadatas: Optional[List[List[Metadata]]]
distances: Optional[List[List[float]]]
"""
if not HAS_CHROMADB:
raise ImportError("Please install dependencies first. `pip install pyautogen[retrievechat]`")
if client is None:
client = chromadb.PersistentClient(path=db_path)
# the collection's embedding function is always the default one, but we want to use the one we used to create the
Expand Down

0 comments on commit 64e401d

Please sign in to comment.