Repository: microsoft/autogen
Pull request:  New Retriever API
I communicated these changes over discord
<!-- Thank you for your contribution! Please review https://microsoft.github.io/autogen/docs/Contribute before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?
This PR creates a new Retriever API that makes plugging different retrievers simpler. It also removes hardcoding of particular vectorDBs and makes a centralized module that deal with optional vectordb dependencies
<!-- Please give a short summary of the change and the problem this solves. -->
Retrieval functions were hardcoded. This PR introduces a new `Retriever` abstract interface that simply has 3 functions.
`Retriever.init_db()` - To initialize connection/client or nothing if in serverless setting
`Retriever.ingest_data()` - To upsert data in the db
`Retriever.query()` - To run queries
Any DB provider can simply implement these functions and register as a supported retriever for Autogen.
**Note:** This PR will not break the existing examples/notebooks.
This is still WIP but I tested the current chromadb implementation and it works well with the existing examples.
## Related issue number

<!-- For example: "Closes #1234" -->
Closes https://github.com/microsoft/autogen/issues/586 https://github.com/microsoft/autogen/issues/416 and a few others
## Checks

- [ ] I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation to build and test documentation locally.
- [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR.
- [ ] I've made sure all auto checks have passed.


Problem statement:
Math Chat using LanceDb Integration
I want to use MathChat with LanceDB. 
Is there any way to use MathChat with LanceDB? Any Plans to integrate AutoGen with LanceDB? 
I am open to doing an integration. Please Let me know whatever way it works.

Hints from issue discussion:
Hi @PrashantDixit-dev , would you mind explaining a little bit more about how you would like to use lancedb in Math Chat? I know some are asking for lancedb for RAG, but Math Chat with lancedb is new to me, and I would like to learn more. Thanks.

Commit message: update

Patch to review:
diff --git a/autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py b/autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py
index 619befdefbd5..2ba4b12ad235 100644
--- a/autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py
+++ b/autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py
@@ -1,7 +1,7 @@
 from typing import Callable, Dict, List, Optional
 
 from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
-from autogen.retrieve_utils import get_files_from_dir, split_files_to_chunks, TEXT_FORMATS
+from autogen.agentchat.contrib.retrieve_utils import get_files_from_dir, split_files_to_chunks, TEXT_FORMATS
 import logging
 
 logger = logging.getLogger(__name__)
diff --git a/autogen/agentchat/contrib/retrieve_user_proxy_agent.py b/autogen/agentchat/contrib/retrieve_user_proxy_agent.py
index 5edc13f2408a..a29f79ff3003 100644
--- a/autogen/agentchat/contrib/retrieve_user_proxy_agent.py
+++ b/autogen/agentchat/contrib/retrieve_user_proxy_agent.py
@@ -1,14 +1,11 @@
 import re
 
-try:
-    import chromadb
-except ImportError:
-    raise ImportError("Please install dependencies first. `pip install pyautogen[retrievechat]`")
 from autogen.agentchat.agent import Agent
 from autogen.agentchat import UserProxyAgent
-from autogen.retrieve_utils import create_vector_db_from_dir, query_vector_db, TEXT_FORMATS
+from autogen.agentchat.contrib.retriever.retrieve_utils import create_vector_db_from_dir, query_vector_db, TEXT_FORMATS
 from autogen.token_count_utils import count_token
 from autogen.code_utils import extract_code
+from autogen.agentchat.contrib.retriever import get_retriever
 
 from typing import Callable, Dict, Optional, Union, List, Tuple, Any
 from IPython import get_ipython
@@ -95,10 +92,9 @@ def __init__(
                 To use default config, set to None. Otherwise, set to a dictionary with the following keys:
                 - task (Optional, str): the task of the retrieve chat. Possible values are "code", "qa" and "default". System
                     prompt will be different for different tasks. The default value is `default`, which supports both code and qa.
-                - client (Optional, chromadb.Client): the chromadb client. If key not provided, a default client `chromadb.Client()`
-                    will be used. If you want to use other vector db, extend this class and override the `retrieve_docs` function.
-                - docs_path (Optional, Union[str, List[str]]): the path to the docs directory. It can also be the path to a single file,
-                    the url to a single file or a list of directories, files and urls. Default is None, which works only if the collection is already created.
+                - client (Optional, Any): the vectordb client/connection. If key not provided, the Retreiver class should handle it.
+                - docs_path (Optional, str): the path to the docs directory. It can also be the path to a single file,
+                    or the url to a single file. Default is None, which works only if the collection is already created.
                 - collection_name (Optional, str): the name of the collection.
                     If key not provided, a default name `autogen-docs` will be used.
                 - model (Optional, str): the model to use for the retrieve chat.
@@ -122,9 +118,9 @@ def __init__(
                 - customized_answer_prefix (Optional, str): the customized answer prefix for the retrieve chat. Default is "".
                     If not "" and the customized_answer_prefix is not in the answer, `Update Context` will be triggered.
                 - update_context (Optional, bool): if False, will not apply `Update Context` for interactive retrieval. Default is True.
-                - get_or_create (Optional, bool): if True, will create/return a collection for the retrieve chat. This is the same as that used in chromadb.
-                    Default is False. Will raise ValueError if the collection already exists and get_or_create is False. Will be set to True if docs_path is None.
-                - custom_token_count_function (Optional, Callable): a custom function to count the number of tokens in a string.
+                - get_or_create (Optional, bool): if True, will create/recreate a collection for the retrieve chat.
+                    This is the same as that used in retriever. Default is False. Will be set to False if docs_path is None.
+                - custom_token_count_function(Optional, Callable): a custom function to count the number of tokens in a string.
                     The function should take (text:str, model:str) as input and return the token_count(int). the retrieve_config["model"] will be passed in the function.
                     Default is autogen.token_count_utils.count_token that uses tiktoken, which may not be accurate for non-OpenAI models.
                 - custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings.
@@ -135,7 +131,7 @@ def __init__(
             **kwargs (dict): other kwargs in [UserProxyAgent](../user_proxy_agent#__init__).
 
         Example of overriding retrieve_docs:
-        If you have set up a customized vector db, and it's not compatible with chromadb, you can easily plug in it with below code.
+        If you have set up a customized vector db, and it's not compatible with retriever, you can easily plug in it with below code.
         ```python
         class MyRetrieveUserProxyAgent(RetrieveUserProxyAgent):
             def query_vector_db(
@@ -167,8 +163,9 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =
         )
 
         self._retrieve_config = {} if retrieve_config is None else retrieve_config
+        self._retriever_type = self._retrieve_config.get("retriever_type")
         self._task = self._retrieve_config.get("task", "default")
-        self._client = self._retrieve_config.get("client", chromadb.Client())
+        self._client = self._retrieve_config.get("client", None)
         self._docs_path = self._retrieve_config.get("docs_path", None)
         self._collection_name = self._retrieve_config.get("collection_name", "autogen-docs")
         self._model = self._retrieve_config.get("model", "gpt-4")
@@ -355,13 +352,9 @@ def _generate_retrieve_user_reply(
 
     def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
         """Retrieve docs based on the given problem and assign the results to the class property `_results`.
-        In case you want to customize the retrieval process, such as using a different vector db whose APIs are not
-        compatible with chromadb or filter results with metadata, you can override this function. Just keep the current
-        parameters and add your own parameters with default values, and keep the results in below type.
 
         Type of the results: Dict[str, List[List[Any]]], should have keys "ids" and "documents", "ids" for the ids of
-        the retrieved docs and "documents" for the contents of the retrieved docs. Any other keys are optional. Refer
-        to `chromadb.api.types.QueryResult` as an example.
+        the retrieved docs and "documents" for the contents of the retrieved docs. Any other keys are optional.
             ids: List[string]
             documents: List[List[string]]
 
@@ -372,31 +365,27 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =
         """
         if not self._collection or not self._get_or_create:
             print("Trying to create collection.")
-            self._client = create_vector_db_from_dir(
-                dir_path=self._docs_path,
+            retriever_class = get_retriever(self._retriever_type)
+            self.retriever = retriever_class(
+                name=self._collection_name,
+                embedding_model_name=self._embedding_model,
+                embedding_function=self._embedding_function,
                 max_tokens=self._chunk_token_size,
-                client=self._client,
-                collection_name=self._collection_name,
                 chunk_mode=self._chunk_mode,
                 must_break_at_empty_line=self._must_break_at_empty_line,
-                embedding_model=self._embedding_model,
-                get_or_create=self._get_or_create,
-                embedding_function=self._embedding_function,
                 custom_text_split_function=self.custom_text_split_function,
+                use_existing=not self._get_or_create,
+                client=self._client,
                 custom_text_types=self._custom_text_types,
                 recursive=self._recursive,
             )
             self._collection = True
-            self._get_or_create = True
-
-        results = query_vector_db(
-            query_texts=[problem],
-            n_results=n_results,
-            search_string=search_string,
-            client=self._client,
-            collection_name=self._collection_name,
-            embedding_model=self._embedding_model,
-            embedding_function=self._embedding_function,
+            self._get_or_create = False
+        self.retriever.ingest_data(self._docs_path)
+        results = self.retriever.query(
+            texts=[problem],
+            top_k=n_results,
+            filter=search_string,
         )
         self._search_string = search_string
         self._results = results
diff --git a/autogen/agentchat/contrib/retriever/__init__.py b/autogen/agentchat/contrib/retriever/__init__.py
new file mode 100644
index 000000000000..9bfb09598a64
--- /dev/null
+++ b/autogen/agentchat/contrib/retriever/__init__.py
@@ -0,0 +1,17 @@
+from typing import Optional
+from .chromadb import ChromaDB
+from .lancedb import LanceDB
+
+AVILABLE_RETRIEVERS = ["lanchedb", "chromadb"]
+DEFAULT_RETRIEVER = "lancedb"
+
+
+def get_retriever(type: Optional[str] = None):
+    """Return a retriever instance."""
+    type = type or DEFAULT_RETRIEVER
+    if type == "chromadb":
+        return ChromaDB
+    elif type == "lancedb":
+        return LanceDB
+    else:
+        raise ValueError(f"Unknown retriever type {type}")
diff --git a/autogen/agentchat/contrib/retriever/base.py b/autogen/agentchat/contrib/retriever/base.py
new file mode 100644
index 000000000000..4c5dfd69f442
--- /dev/null
+++ b/autogen/agentchat/contrib/retriever/base.py
@@ -0,0 +1,67 @@
+from abc import ABC, abstractmethod
+from typing import List, Union, Callable, Any
+
+
+class Retriever(ABC):
+    def __init__(
+        self,
+        path="./db",
+        name="vectorstore",
+        embedding_model_name="all-MiniLM-L6-v2",
+        embedding_function=None,
+        max_tokens: int = 4000,
+        chunk_mode: str = "multi_lines",
+        must_break_at_empty_line: bool = True,
+        custom_text_split_function: Callable = None,
+        use_existing=False,
+        client=None,
+    ):
+        """
+        Args:
+            path: path to the folder where the database is stored
+            name: name of the database
+            embedding_model_name: name of the embedding model to use
+            embedding_function: function to use to embed the text
+            max_tokens: maximum number of tokens to embed
+            chunk_mode: mode to chunk the text. Can be "multi_lines" or "single_line"
+            must_break_at_empty_line: whether to break the text at empty lines when chunking
+            custom_text_split_function: custom function to split the text into chunks
+        """
+        self.path = path
+        self.name = name
+        self.embedding_model_name = embedding_model_name
+        self.embedding_function = embedding_function
+        self.max_tokens = max_tokens
+        self.chunk_mode = chunk_mode
+        self.must_break_at_empty_line = must_break_at_empty_line
+        self.custom_text_split_function = custom_text_split_function
+        self.use_existing = use_existing
+        self.client = client
+
+        self.init_db()
+
+    @abstractmethod
+    def ingest_data(self, data_dir):
+        """
+        Create a vector database from a directory of files.
+        Args:
+            data_dir: path to the directory containing the text files
+        """
+        pass
+
+    @abstractmethod
+    def query(self, texts: List[str], top_k: int = 10, filter: Any = None):
+        """
+        Query the database.
+        Args:
+            query: query string or list of query strings
+            top_k: number of results to return
+        """
+        pass
+
+    @abstractmethod
+    def init_db(self):
+        """
+        Initialize the database.
+        """
+        pass
diff --git a/autogen/agentchat/contrib/retriever/chromadb.py b/autogen/agentchat/contrib/retriever/chromadb.py
new file mode 100644
index 000000000000..05b0951970f7
--- /dev/null
+++ b/autogen/agentchat/contrib/retriever/chromadb.py
@@ -0,0 +1,79 @@
+from typing import List
+from .base import Retriever
+from .retrieve_utils import split_text_to_chunks, extract_text_from_pdf, split_files_to_chunks, get_files_from_dir
+
+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
+except ImportError:
+    raise ImportError("Please install chromadb: pip install chromadb")
+
+
+class ChromaDB(Retriever):
+    def init_db(self):
+        if self.client is None:
+            self.client = chromadb.PersistentClient(path=self.path)
+        embedding_function = (
+            ef.SentenceTransformerEmbeddingFunction(self.embedding_model_name)
+            if self.embedding_function is None
+            else self.embedding_function
+        )
+        self.collection = self.client.create_collection(
+            self.name,
+            get_or_create=self.use_existing,
+            embedding_function=embedding_function,
+            # https://github.com/nmslib/hnswlib#supported-distances
+            # https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184
+            # https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md
+            metadata={"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32},  # ip, l2, cosine
+        )
+
+    def ingest_data(self, data_dir):
+        """
+        Create a vector database from a directory of files.
+        Args:
+            data_dir: path to the directory containing the text files
+        """
+        if self.client is None:
+            self.init_db()
+        if self.custom_text_split_function is not None:
+            chunks = split_files_to_chunks(
+                get_files_from_dir(data_dir), custom_text_split_function=self.custom_text_split_function
+            )
+        else:
+            chunks = split_files_to_chunks(
+                get_files_from_dir(data_dir), self.max_tokens, self.chunk_mode, self.must_break_at_empty_line
+            )
+        print(f"Found {len(chunks)} chunks.")  #
+        # Upsert in batch of 40000 or less if the total number of chunks is less than 40000
+        for i in range(0, len(chunks), min(40000, len(chunks))):
+            end_idx = i + min(40000, len(chunks) - i)
+            self.collection.upsert(
+                documents=chunks[i:end_idx],
+                ids=[f"doc_{j}" for j in range(i, end_idx)],  # unique for each doc
+            )
+
+    def query(self, texts: List[str], top_k: int = 10, filter: str = None):
+        if self.client is None:
+            self.init_db()
+        # the collection's embedding function is always the default one, but we want to use the one we used to create the
+        # collection. So we compute the embeddings ourselves and pass it to the query function.
+        embedding_function = (
+            ef.SentenceTransformerEmbeddingFunction(self.embedding_model_name)
+            if self.embedding_function is None
+            else self.embedding_function
+        )
+        query_embeddings = embedding_function(texts)
+        # Query/search n most similar results. You can also .get by id
+        results = self.collection.query(
+            query_embeddings=query_embeddings,
+            n_results=top_k,
+            where_document={"$contains": filter} if filter else None,  # optional filter
+        )
+        return results
diff --git a/autogen/agentchat/contrib/retriever/lancedb.py b/autogen/agentchat/contrib/retriever/lancedb.py
new file mode 100644
index 000000000000..cf7af621f1b5
--- /dev/null
+++ b/autogen/agentchat/contrib/retriever/lancedb.py
@@ -0,0 +1,96 @@
+from typing import Callable, List
+from collections import defaultdict
+from .base import Retriever
+from .retrieve_utils import split_text_to_chunks, extract_text_from_pdf, split_files_to_chunks, get_files_from_dir
+
+try:
+    import lancedb
+    from lancedb.embeddings import get_registry, EmbeddingFunction, with_embeddings
+    from lancedb.pydantic import LanceModel, Vector
+    import pyarrow as pa
+except ImportError:
+    raise ImportError("Please install lancedb: pip install lancedb")
+
+
+class LanceDB(Retriever):
+    db = None
+
+    def init_db(self):
+        if self.db is None:
+            self.db = lancedb.connect(self.path)
+        self.embedding_function = (
+            get_registry().get("sentence-transformers").create(name=self.embedding_model_name, show_progress_bar=True)
+            if self.embedding_function is None
+            else self.embedding_function
+        )
+        if self.use_existing and self.name in self.db.table_names():
+            self.table = self.db.open_table(self.name)
+            # logger.info(f"Reusing existing table {self.name}")
+        else:
+            # logger.info(f"Creating new table {self.name}")
+            schema = self._get_schema(self.embedding_function)
+            self.table = self.db.create_table(self.name, schema=schema, mode="overwrite")
+
+    def ingest_data(self, data_dir):
+        """
+        Create a vector database from a directory of files.
+        Args:
+            data_dir: path to the directory containing the text files
+        """
+        if self.db is None:
+            self.init_db()
+        if self.custom_text_split_function is not None:
+            chunks = split_files_to_chunks(
+                get_files_from_dir(data_dir), custom_text_split_function=self.custom_text_split_function
+            )
+        else:
+            chunks = split_files_to_chunks(
+                get_files_from_dir(data_dir), self.max_tokens, self.chunk_mode, self.must_break_at_empty_line
+            )
+        print(f"Found {len(chunks)} chunks.")  #
+        data = [{"documents": docs, "ids": idx} for idx, docs in enumerate(chunks)]
+        if isinstance(self.embedding_function, EmbeddingFunction):  # this means we are using embedding API
+            self.table.add(data)
+        elif isinstance(self.embedding_function, Callable):
+            pa_table = pa.Table.from_pylist(data)
+            data = with_embeddings(self.embedding_function, pa_table)
+            self.table.add(data)
+
+    def query(self, texts: List[str], top_k: int = 10, filter: str = None):
+        if self.db is None:
+            self.init_db()
+        texts = [texts] if isinstance(texts, str) else texts
+        results = defaultdict(list)
+        for text in texts:
+            query = self.embedding_function(text) if isinstance(self.embedding_function, Callable) else text
+            print("query: ", query)
+            result = self.table.search(query)
+            if filter is not None:
+                result = result.where(f"documents LIKE '%{filter}%'")
+            result = result.limit(top_k).to_arrow().to_pydict()
+            for k, v in result.items():
+                results[k].append(v)
+
+        return results
+
+    def _get_schema(self, embedding_function):
+        if isinstance(embedding_function, EmbeddingFunction):
+
+            class Schema(LanceModel):
+                vector: Vector(embedding_function.ndims()) = embedding_function.VectorField()
+                documents: str = embedding_function.SourceField()
+                ids: str
+
+            return Schema
+        elif isinstance(embedding_function, Callable):
+            dim = embedding_function("test").shape[0]  # TODO: check this
+            schema = pa.schema(
+                [
+                    pa.field("Vector", pa.list_(pa.float32(), dim)),
+                    pa.field("documents", pa.string()),
+                    pa.field("ids", pa.string()),
+                ]
+            )
+            return schema
+        else:
+            raise ValueError("embedding_function should be a callable or an EmbeddingFunction instance")
diff --git a/autogen/agentchat/contrib/retriever/retrieve_utils.py b/autogen/agentchat/contrib/retriever/retrieve_utils.py
new file mode 100644
index 000000000000..b02ab66d6a31
--- /dev/null
+++ b/autogen/agentchat/contrib/retriever/retrieve_utils.py
@@ -0,0 +1,211 @@
+from typing import List, Union, Callable
+import os
+import requests
+from urllib.parse import urlparse
+import glob
+import logging
+from autogen.token_count_utils import count_token
+
+try:
+    from unstructured.partition.auto import partition
+
+    HAS_UNSTRUCTURED = True
+except ImportError:
+    HAS_UNSTRUCTURED = False
+
+logger = logging.getLogger(__name__)
+TEXT_FORMATS = [
+    "txt",
+    "json",
+    "csv",
+    "tsv",
+    "md",
+    "html",
+    "htm",
+    "rtf",
+    "rst",
+    "jsonl",
+    "log",
+    "xml",
+    "yaml",
+    "yml",
+    "pdf",
+]
+UNSTRUCTURED_FORMATS = ["docx", "doc", "odt", "pptx", "ppt", "xlsx", "eml", "msg", "epub"]
+if HAS_UNSTRUCTURED:
+    TEXT_FORMATS += UNSTRUCTURED_FORMATS
+    TEXT_FORMATS = list(set(TEXT_FORMATS))
+VALID_CHUNK_MODES = frozenset({"one_line", "multi_lines"})
+
+
+def split_text_to_chunks(
+    text: str,
+    max_tokens: int = 4000,
+    chunk_mode: str = "multi_lines",
+    must_break_at_empty_line: bool = True,
+    overlap: int = 10,
+):
+    """Split a long text into chunks of max_tokens."""
+    if chunk_mode not in VALID_CHUNK_MODES:
+        raise AssertionError
+    if chunk_mode == "one_line":
+        must_break_at_empty_line = False
+    chunks = []
+    lines = text.split("\n")
+    lines_tokens = [count_token(line) for line in lines]
+    sum_tokens = sum(lines_tokens)
+    while sum_tokens > max_tokens:
+        if chunk_mode == "one_line":
+            estimated_line_cut = 2
+        else:
+            estimated_line_cut = int(max_tokens / sum_tokens * len(lines)) + 1
+        cnt = 0
+        prev = ""
+        for cnt in reversed(range(estimated_line_cut)):
+            if must_break_at_empty_line and lines[cnt].strip() != "":
+                continue
+            if sum(lines_tokens[:cnt]) <= max_tokens:
+                prev = "\n".join(lines[:cnt])
+                break
+        if cnt == 0:
+            logger.warning(
+                f"max_tokens is too small to fit a single line of text. Breaking this line:\n\t{lines[0][:100]} ..."
+            )
+            if not must_break_at_empty_line:
+                split_len = int(max_tokens / lines_tokens[0] * 0.9 * len(lines[0]))
+                prev = lines[0][:split_len]
+                lines[0] = lines[0][split_len:]
+                lines_tokens[0] = count_token(lines[0])
+            else:
+                logger.warning("Failed to split docs with must_break_at_empty_line being True, set to False.")
+                must_break_at_empty_line = False
+        chunks.append(prev) if len(prev) > 10 else None  # don't add chunks less than 10 characters
+        lines = lines[cnt:]
+        lines_tokens = lines_tokens[cnt:]
+        sum_tokens = sum(lines_tokens)
+    text_to_chunk = "\n".join(lines)
+    chunks.append(text_to_chunk) if len(text_to_chunk) > 10 else None  # don't add chunks less than 10 characters
+    return chunks
+
+
+def extract_text_from_pdf(file: str) -> str:
+    """Extract text from PDF files"""
+    import pypdf  # optional dependency
+
+    text = ""
+    with open(file, "rb") as f:
+        reader = pypdf.PdfReader(f)
+        if reader.is_encrypted:  # Check if the PDF is encrypted
+            try:
+                reader.decrypt("")
+            except pypdf.errors.FileNotDecryptedError as e:
+                logger.warning(f"Could not decrypt PDF {file}, {e}")
+                return text  # Return empty text if PDF could not be decrypted
+
+        for page_num in range(len(reader.pages)):
+            page = reader.pages[page_num]
+            text += page.extract_text()
+
+    if not text.strip():  # Debugging line to check if text is empty
+        logger.warning(f"Could not decrypt PDF {file}")
+
+    return text
+
+
+def split_files_to_chunks(
+    files: list,
+    max_tokens: int = 4000,
+    chunk_mode: str = "multi_lines",
+    must_break_at_empty_line: bool = True,
+    custom_text_split_function: Callable = None,
+):
+    """Split a list of files into chunks of max_tokens."""
+
+    chunks = []
+
+    for file in files:
+        _, file_extension = os.path.splitext(file)
+        file_extension = file_extension.lower()
+
+        if HAS_UNSTRUCTURED and file_extension[1:] in UNSTRUCTURED_FORMATS:
+            text = partition(file)
+            text = "\n".join([t.text for t in text]) if len(text) > 0 else ""
+        elif file_extension == ".pdf":
+            text = extract_text_from_pdf(file)
+        else:  # For non-PDF text-based files
+            with open(file, "r", encoding="utf-8", errors="ignore") as f:
+                text = f.read()
+
+        if not text.strip():  # Debugging line to check if text is empty after reading
+            logger.warning(f"No text available in file: {file}")
+            continue  # Skip to the next file if no text is available
+
+        if custom_text_split_function is not None:
+            chunks += custom_text_split_function(text)
+        else:
+            chunks += split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)
+
+    return chunks
+
+
+def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMATS, recursive: bool = True):
+    """Return a list of all the files in a given directory."""
+    if len(types) == 0:
+        raise ValueError("types cannot be empty.")
+    types = [t[1:].lower() if t.startswith(".") else t.lower() for t in set(types)]
+    types += [t.upper() for t in types]
+
+    files = []
+    # If the path is a list of files or urls, process and return them
+    if isinstance(dir_path, list):
+        for item in dir_path:
+            if os.path.isfile(item):
+                files.append(item)
+            elif is_url(item):
+                files.append(get_file_from_url(item))
+            else:
+                logger.warning(f"File {item} does not exist. Skipping.")
+        return files
+
+    # If the path is a file, return it
+    if os.path.isfile(dir_path):
+        return [dir_path]
+
+    # If the path is a url, download it and return the downloaded file
+    if is_url(dir_path):
+        return [get_file_from_url(dir_path)]
+
+    if os.path.exists(dir_path):
+        for type in types:
+            if recursive:
+                files += glob.glob(os.path.join(dir_path, f"**/*.{type}"), recursive=True)
+            else:
+                files += glob.glob(os.path.join(dir_path, f"*.{type}"), recursive=False)
+    else:
+        logger.error(f"Directory {dir_path} does not exist.")
+        raise ValueError(f"Directory {dir_path} does not exist.")
+    return files
+
+
+def get_file_from_url(url: str, save_path: str = None):
+    """Download a file from a URL."""
+    if save_path is None:
+        os.makedirs("/tmp/chromadb", exist_ok=True)
+        save_path = os.path.join("/tmp/chromadb", os.path.basename(url))
+    else:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    with requests.get(url, stream=True) as r:
+        r.raise_for_status()
+        with open(save_path, "wb") as f:
+            for chunk in r.iter_content(chunk_size=8192):
+                f.write(chunk)
+    return save_path
+
+
+def is_url(string: str):
+    """Return True if the string is a valid URL."""
+    try:
+        result = urlparse(string)
+        return all([result.scheme, result.netloc])
+    except ValueError:
+        return False
diff --git a/autogen/retrieve_utils.py b/autogen/retrieve_utils.py
index ba856e3b238e..6e92a938c270 100644
--- a/autogen/retrieve_utils.py
+++ b/autogen/retrieve_utils.py
@@ -1,374 +1,4 @@
-from typing import List, Union, Callable
-import os
-import requests
-from urllib.parse import urlparse
-import glob
-import chromadb
+from . import logger
+from .agentchat.contrib.retriever.retrieve_utils import *
 
-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:
-    from unstructured.partition.auto import partition
-
-    HAS_UNSTRUCTURED = True
-except ImportError:
-    HAS_UNSTRUCTURED = False
-
-logger = logging.getLogger(__name__)
-TEXT_FORMATS = [
-    "txt",
-    "json",
-    "csv",
-    "tsv",
-    "md",
-    "html",
-    "htm",
-    "rtf",
-    "rst",
-    "jsonl",
-    "log",
-    "xml",
-    "yaml",
-    "yml",
-    "pdf",
-]
-UNSTRUCTURED_FORMATS = [
-    "doc",
-    "docx",
-    "epub",
-    "msg",
-    "odt",
-    "org",
-    "pdf",
-    "ppt",
-    "pptx",
-    "rtf",
-    "rst",
-    "xlsx",
-]
-if HAS_UNSTRUCTURED:
-    TEXT_FORMATS += UNSTRUCTURED_FORMATS
-    TEXT_FORMATS = list(set(TEXT_FORMATS))
-VALID_CHUNK_MODES = frozenset({"one_line", "multi_lines"})
-
-
-def split_text_to_chunks(
-    text: str,
-    max_tokens: int = 4000,
-    chunk_mode: str = "multi_lines",
-    must_break_at_empty_line: bool = True,
-    overlap: int = 10,
-):
-    """Split a long text into chunks of max_tokens."""
-    if chunk_mode not in VALID_CHUNK_MODES:
-        raise AssertionError
-    if chunk_mode == "one_line":
-        must_break_at_empty_line = False
-    chunks = []
-    lines = text.split("\n")
-    lines_tokens = [count_token(line) for line in lines]
-    sum_tokens = sum(lines_tokens)
-    while sum_tokens > max_tokens:
-        if chunk_mode == "one_line":
-            estimated_line_cut = 2
-        else:
-            estimated_line_cut = int(max_tokens / sum_tokens * len(lines)) + 1
-        cnt = 0
-        prev = ""
-        for cnt in reversed(range(estimated_line_cut)):
-            if must_break_at_empty_line and lines[cnt].strip() != "":
-                continue
-            if sum(lines_tokens[:cnt]) <= max_tokens:
-                prev = "\n".join(lines[:cnt])
-                break
-        if cnt == 0:
-            logger.warning(
-                f"max_tokens is too small to fit a single line of text. Breaking this line:\n\t{lines[0][:100]} ..."
-            )
-            if not must_break_at_empty_line:
-                split_len = int(max_tokens / lines_tokens[0] * 0.9 * len(lines[0]))
-                prev = lines[0][:split_len]
-                lines[0] = lines[0][split_len:]
-                lines_tokens[0] = count_token(lines[0])
-            else:
-                logger.warning("Failed to split docs with must_break_at_empty_line being True, set to False.")
-                must_break_at_empty_line = False
-        chunks.append(prev) if len(prev) > 10 else None  # don't add chunks less than 10 characters
-        lines = lines[cnt:]
-        lines_tokens = lines_tokens[cnt:]
-        sum_tokens = sum(lines_tokens)
-    text_to_chunk = "\n".join(lines)
-    chunks.append(text_to_chunk) if len(text_to_chunk) > 10 else None  # don't add chunks less than 10 characters
-    return chunks
-
-
-def extract_text_from_pdf(file: str) -> str:
-    """Extract text from PDF files"""
-    text = ""
-    with open(file, "rb") as f:
-        reader = pypdf.PdfReader(f)
-        if reader.is_encrypted:  # Check if the PDF is encrypted
-            try:
-                reader.decrypt("")
-            except pypdf.errors.FileNotDecryptedError as e:
-                logger.warning(f"Could not decrypt PDF {file}, {e}")
-                return text  # Return empty text if PDF could not be decrypted
-
-        for page_num in range(len(reader.pages)):
-            page = reader.pages[page_num]
-            text += page.extract_text()
-
-    if not text.strip():  # Debugging line to check if text is empty
-        logger.warning(f"Could not decrypt PDF {file}")
-
-    return text
-
-
-def split_files_to_chunks(
-    files: list,
-    max_tokens: int = 4000,
-    chunk_mode: str = "multi_lines",
-    must_break_at_empty_line: bool = True,
-    custom_text_split_function: Callable = None,
-):
-    """Split a list of files into chunks of max_tokens."""
-
-    chunks = []
-
-    for file in files:
-        _, file_extension = os.path.splitext(file)
-        file_extension = file_extension.lower()
-
-        if HAS_UNSTRUCTURED and file_extension[1:] in UNSTRUCTURED_FORMATS:
-            text = partition(file)
-            text = "\n".join([t.text for t in text]) if len(text) > 0 else ""
-        elif file_extension == ".pdf":
-            text = extract_text_from_pdf(file)
-        else:  # For non-PDF text-based files
-            with open(file, "r", encoding="utf-8", errors="ignore") as f:
-                text = f.read()
-
-        if not text.strip():  # Debugging line to check if text is empty after reading
-            logger.warning(f"No text available in file: {file}")
-            continue  # Skip to the next file if no text is available
-
-        if custom_text_split_function is not None:
-            chunks += custom_text_split_function(text)
-        else:
-            chunks += split_text_to_chunks(text, max_tokens, chunk_mode, must_break_at_empty_line)
-
-    return chunks
-
-
-def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMATS, recursive: bool = True):
-    """Return a list of all the files in a given directory, a url, a file path or a list of them."""
-    if len(types) == 0:
-        raise ValueError("types cannot be empty.")
-    types = [t[1:].lower() if t.startswith(".") else t.lower() for t in set(types)]
-    types += [t.upper() for t in types]
-
-    files = []
-    # If the path is a list of files or urls, process and return them
-    if isinstance(dir_path, list):
-        for item in dir_path:
-            if os.path.isfile(item):
-                files.append(item)
-            elif is_url(item):
-                files.append(get_file_from_url(item))
-            elif os.path.exists(item):
-                try:
-                    files.extend(get_files_from_dir(item, types, recursive))
-                except ValueError:
-                    logger.warning(f"Directory {item} does not exist. Skipping.")
-            else:
-                logger.warning(f"File {item} does not exist. Skipping.")
-        return files
-
-    # If the path is a file, return it
-    if os.path.isfile(dir_path):
-        return [dir_path]
-
-    # If the path is a url, download it and return the downloaded file
-    if is_url(dir_path):
-        return [get_file_from_url(dir_path)]
-
-    if os.path.exists(dir_path):
-        for type in types:
-            if recursive:
-                files += glob.glob(os.path.join(dir_path, f"**/*.{type}"), recursive=True)
-            else:
-                files += glob.glob(os.path.join(dir_path, f"*.{type}"), recursive=False)
-    else:
-        logger.error(f"Directory {dir_path} does not exist.")
-        raise ValueError(f"Directory {dir_path} does not exist.")
-    return files
-
-
-def get_file_from_url(url: str, save_path: str = None):
-    """Download a file from a URL."""
-    if save_path is None:
-        os.makedirs("/tmp/chromadb", exist_ok=True)
-        save_path = os.path.join("/tmp/chromadb", os.path.basename(url))
-    else:
-        os.makedirs(os.path.dirname(save_path), exist_ok=True)
-    with requests.get(url, stream=True) as r:
-        r.raise_for_status()
-        with open(save_path, "wb") as f:
-            for chunk in r.iter_content(chunk_size=8192):
-                f.write(chunk)
-    return save_path
-
-
-def is_url(string: str):
-    """Return True if the string is a valid URL."""
-    try:
-        result = urlparse(string)
-        return all([result.scheme, result.netloc])
-    except ValueError:
-        return False
-
-
-def create_vector_db_from_dir(
-    dir_path: Union[str, List[str]],
-    max_tokens: int = 4000,
-    client: API = None,
-    db_path: str = "/tmp/chromadb.db",
-    collection_name: str = "all-my-documents",
-    get_or_create: bool = False,
-    chunk_mode: str = "multi_lines",
-    must_break_at_empty_line: bool = True,
-    embedding_model: str = "all-MiniLM-L6-v2",
-    embedding_function: Callable = None,
-    custom_text_split_function: Callable = None,
-    custom_text_types: List[str] = TEXT_FORMATS,
-    recursive: bool = True,
-) -> 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.
-
-    Args:
-        dir_path (Union[str, List[str]]): the path to the directory, file, url or a list of them.
-        max_tokens (Optional, int): the maximum number of tokens per chunk. Default is 4000.
-        client (Optional, API): the chromadb client. Default is None.
-        db_path (Optional, str): the path to the chromadb. Default is "/tmp/chromadb.db".
-        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
-        get_or_create (Optional, bool): Whether to get or create the collection. Default is False. If True, the collection
-            will be returned if it already exists. Will raise ValueError if the collection already exists and get_or_create is False.
-        chunk_mode (Optional, str): the chunk mode. Default is "multi_lines".
-        must_break_at_empty_line (Optional, bool): Whether to break at empty line. Default is True.
-        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
-            embedding_function is not None.
-        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`.
-        custom_text_split_function (Optional, Callable): a custom function to split a string into a list of strings.
-            Default is None, will use the default function in `autogen.retrieve_utils.split_text_to_chunks`.
-        custom_text_types (Optional, List[str]): a list of file types to be processed. Default is TEXT_FORMATS.
-        recursive (Optional, bool): whether to search documents recursively in the dir_path. Default is True.
-
-    Returns:
-        API: the chromadb client.
-    """
-    if client is None:
-        client = chromadb.PersistentClient(path=db_path)
-    try:
-        embedding_function = (
-            ef.SentenceTransformerEmbeddingFunction(embedding_model)
-            if embedding_function is None
-            else embedding_function
-        )
-        collection = client.create_collection(
-            collection_name,
-            get_or_create=get_or_create,
-            embedding_function=embedding_function,
-            # https://github.com/nmslib/hnswlib#supported-distances
-            # https://github.com/chroma-core/chroma/blob/566bc80f6c8ee29f7d99b6322654f32183c368c4/chromadb/segment/impl/vector/local_hnsw.py#L184
-            # https://github.com/nmslib/hnswlib/blob/master/ALGO_PARAMS.md
-            metadata={"hnsw:space": "ip", "hnsw:construction_ef": 30, "hnsw:M": 32},  # ip, l2, cosine
-        )
-
-        if custom_text_split_function is not None:
-            chunks = split_files_to_chunks(
-                get_files_from_dir(dir_path, custom_text_types, recursive),
-                custom_text_split_function=custom_text_split_function,
-            )
-        else:
-            chunks = split_files_to_chunks(
-                get_files_from_dir(dir_path, custom_text_types, recursive),
-                max_tokens,
-                chunk_mode,
-                must_break_at_empty_line,
-            )
-        logger.info(f"Found {len(chunks)} chunks.")
-        # Upsert in batch of 40000 or less if the total number of chunks is less than 40000
-        for i in range(0, len(chunks), min(40000, len(chunks))):
-            end_idx = i + min(40000, len(chunks) - i)
-            collection.upsert(
-                documents=chunks[i:end_idx],
-                ids=[f"doc_{j}" for j in range(i, end_idx)],  # unique for each doc
-            )
-    except ValueError as e:
-        logger.warning(f"{e}")
-    return client
-
-
-def query_vector_db(
-    query_texts: List[str],
-    n_results: int = 10,
-    client: API = None,
-    db_path: str = "/tmp/chromadb.db",
-    collection_name: str = "all-my-documents",
-    search_string: str = "",
-    embedding_model: str = "all-MiniLM-L6-v2",
-    embedding_function: Callable = None,
-) -> QueryResult:
-    """Query a vector db. We support chromadb compatible APIs, it's not required if you prepared your own vector db
-        and query function.
-
-    Args:
-        query_texts (List[str]): the query texts.
-        n_results (Optional, int): the number of results to return. Default is 10.
-        client (Optional, API): the chromadb compatible client. Default is None, a chromadb client will be used.
-        db_path (Optional, str): the path to the vector db. Default is "/tmp/chromadb.db".
-        collection_name (Optional, str): the name of the collection. Default is "all-my-documents".
-        search_string (Optional, str): the search string. Default is "".
-        embedding_model (Optional, str): the embedding model to use. Default is "all-MiniLM-L6-v2". Will be ignored if
-            embedding_function is not None.
-        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:
-        QueryResult: the query result. The format is:
-            class QueryResult(TypedDict):
-                ids: List[IDs]
-                embeddings: Optional[List[List[Embedding]]]
-                documents: Optional[List[List[Document]]]
-                metadatas: Optional[List[List[Metadata]]]
-                distances: Optional[List[List[float]]]
-    """
-    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
-    # collection. So we compute the embeddings ourselves and pass it to the query function.
-    collection = client.get_collection(collection_name)
-    embedding_function = (
-        ef.SentenceTransformerEmbeddingFunction(embedding_model) if embedding_function is None else embedding_function
-    )
-    query_embeddings = embedding_function(query_texts)
-    # Query/search n most similar results. You can also .get by id
-    results = collection.query(
-        query_embeddings=query_embeddings,
-        n_results=n_results,
-        where_document={"$contains": search_string} if search_string else None,  # optional filter
-    )
-    return results
+logger.warning("This module is deprecated. Please use autogen.agentchat.contrib.retriever.retrieve_utils instead.")
diff --git a/setup.py b/setup.py
index 4dc5d4e9dd64..76f1e9134e47 100644
--- a/setup.py
+++ b/setup.py
@@ -49,8 +49,8 @@
         ],
         "blendsearch": ["flaml[blendsearch]"],
         "mathchat": ["sympy", "pydantic==1.10.9", "wolframalpha"],
-        "retrievechat": ["chromadb", "sentence_transformers", "pypdf", "ipython"],
-        "teachable": ["chromadb"],
+        "retrievechat": ["chromadb", "lancedb", "sentence_transformers", "pypdf", "ipython"],
+        "teachable": ["chromadb", "lancedb"],
         "lmm": ["replicate", "pillow"],
     },
     classifiers=[
diff --git a/test/agentchat/contrib/retrievers/test_chromadb.py b/test/agentchat/contrib/retrievers/test_chromadb.py
new file mode 100644
index 000000000000..f73250517253
--- /dev/null
+++ b/test/agentchat/contrib/retrievers/test_chromadb.py
@@ -0,0 +1,35 @@
+import os
+import pytest
+from autogen.agentchat.contrib.retriever.retrieve_utils import (
+    split_text_to_chunks,
+    extract_text_from_pdf,
+    split_files_to_chunks,
+    get_files_from_dir,
+    is_url,
+)
+
+try:
+    from autogen.agentchat.contrib.retriever.chromadb import ChromaDB
+    import chromadb
+except ImportError:
+    skip = True
+else:
+    skip = False
+
+test_dir = os.path.join(os.path.dirname(__file__), "test_files")
+
+
+@pytest.mark.skipif(skip, reason="chromadb is not installed")
+def test_chromadb():
+    db_path = "/tmp/test_retrieve_utils_chromadb.db"
+    client = chromadb.PersistentClient(path=db_path)
+    if os.path.exists(db_path):
+        vectorstore = ChromaDB(path=db_path, use_existing=True)
+    else:
+        vectorstore = ChromaDB(path=db_path)
+    vectorstore.ingest_data(test_dir)
+
+    assert client.get_collection("vectorstore")
+
+    results = vectorstore.query(["autogen"])
+    assert isinstance(results, dict) and any("autogen" in res[0].lower() for res in results.get("documents", []))
diff --git a/test/agentchat/contrib/retrievers/test_lancedb.py b/test/agentchat/contrib/retrievers/test_lancedb.py
new file mode 100644
index 000000000000..5eb82eab041e
--- /dev/null
+++ b/test/agentchat/contrib/retrievers/test_lancedb.py
@@ -0,0 +1,35 @@
+import os
+import pytest
+from autogen.agentchat.contrib.retriever.retrieve_utils import (
+    split_text_to_chunks,
+    extract_text_from_pdf,
+    split_files_to_chunks,
+    get_files_from_dir,
+    is_url,
+)
+
+try:
+    from autogen.agentchat.contrib.retriever.lancedb import LanceDB
+    import lancedb
+except ImportError:
+    skip = True
+else:
+    skip = False
+
+test_dir = os.path.join(os.path.dirname(__file__), "test_files")
+
+
+@pytest.mark.skipif(skip, reason="lancedb is not installed")
+def test_lancedb():
+    db_path = "/tmp/test_lancedb_store"
+    db = lancedb.connect(db_path)
+    if os.path.exists(db_path):
+        vectorstore = LanceDB(path=db_path, use_existing=True)
+    else:
+        vectorstore = LanceDB(path=db_path)
+    vectorstore.ingest_data(test_dir)
+
+    assert "vectorstore" in db.table_names()
+
+    results = vectorstore.query(["autogen"])
+    assert isinstance(results, dict) and any("autogen" in res[0].lower() for res in results.get("documents", []))
diff --git a/test/test_retrieve_utils.py b/test/test_retrieve_utils.py
index 013228e47ce6..ad6ad3df9e63 100644
--- a/test/test_retrieve_utils.py
+++ b/test/test_retrieve_utils.py
@@ -1,25 +1,26 @@
 """
 Unit test for retrieve_utils.py
 """
+import os
+import sys
+import pytest
+
 try:
-    import chromadb
-    from autogen.retrieve_utils import (
+    from autogen.agentchat.contrib.retriever.retrieve_utils import (
         split_text_to_chunks,
         extract_text_from_pdf,
         split_files_to_chunks,
         get_files_from_dir,
         is_url,
-        create_vector_db_from_dir,
-        query_vector_db,
     )
+    from autogen.agentchat.contrib.retriever import DEFAULT_RETRIEVER, get_retriever
     from autogen.token_count_utils import count_token
+
+    Retriever = get_retriever(DEFAULT_RETRIEVER)
 except ImportError:
     skip = True
 else:
     skip = False
-import os
-import sys
-import pytest
 
 try:
     from unstructured.partition.auto import partition
@@ -93,129 +94,30 @@ def test_is_url(self):
         assert is_url("https://www.example.com")
         assert not is_url("not_a_url")
 
-    def test_create_vector_db_from_dir(self):
-        db_path = "/tmp/test_retrieve_utils_chromadb.db"
-        if os.path.exists(db_path):
-            client = chromadb.PersistentClient(path=db_path)
-        else:
-            client = chromadb.PersistentClient(path=db_path)
-            create_vector_db_from_dir(test_dir, client=client)
-
-        assert client.get_collection("all-my-documents")
-
-    def test_query_vector_db(self):
-        db_path = "/tmp/test_retrieve_utils_chromadb.db"
-        if os.path.exists(db_path):
-            client = chromadb.PersistentClient(path=db_path)
-        else:  # If the database does not exist, create it first
-            client = chromadb.PersistentClient(path=db_path)
-            create_vector_db_from_dir(test_dir, client=client)
-
-        results = query_vector_db(["autogen"], client=client)
-        assert isinstance(results, dict) and any("autogen" in res[0].lower() for res in results.get("documents", []))
-
-    def test_custom_vector_db(self):
-        try:
-            import lancedb
-        except ImportError:
-            return
-        from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
-
-        db_path = "/tmp/lancedb"
-
-        def create_lancedb():
-            db = lancedb.connect(db_path)
-            data = [
-                {"vector": [1.1, 1.2], "id": 1, "documents": "This is a test document spark"},
-                {"vector": [0.2, 1.8], "id": 2, "documents": "This is another test document"},
-                {"vector": [0.1, 0.3], "id": 3, "documents": "This is a third test document spark"},
-                {"vector": [0.5, 0.7], "id": 4, "documents": "This is a fourth test document"},
-                {"vector": [2.1, 1.3], "id": 5, "documents": "This is a fifth test document spark"},
-                {"vector": [5.1, 8.3], "id": 6, "documents": "This is a sixth test document"},
-            ]
-            try:
-                db.create_table("my_table", data)
-            except OSError:
-                pass
-
-        class MyRetrieveUserProxyAgent(RetrieveUserProxyAgent):
-            def query_vector_db(
-                self,
-                query_texts,
-                n_results=10,
-                search_string="",
-            ):
-                if query_texts:
-                    vector = [0.1, 0.3]
-                db = lancedb.connect(db_path)
-                table = db.open_table("my_table")
-                query = table.search(vector).where(f"documents LIKE '%{search_string}%'").limit(n_results).to_df()
-                return {"ids": [query["id"].tolist()], "documents": [query["documents"].tolist()]}
-
-            def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
-                results = self.query_vector_db(
-                    query_texts=[problem],
-                    n_results=n_results,
-                    search_string=search_string,
-                )
-
-                self._results = results
-                print("doc_ids: ", results["ids"])
-
-        ragragproxyagent = MyRetrieveUserProxyAgent(
-            name="ragproxyagent",
-            human_input_mode="NEVER",
-            max_consecutive_auto_reply=2,
-            retrieve_config={
-                "task": "qa",
-                "chunk_token_size": 2000,
-                "client": "__",
-                "embedding_model": "all-mpnet-base-v2",
-            },
-        )
-
-        create_lancedb()
-        ragragproxyagent.retrieve_docs("This is a test document spark", n_results=10, search_string="spark")
-        assert ragragproxyagent._results["ids"] == [[3, 1, 5]]
-
     def test_custom_text_split_function(self):
         def custom_text_split_function(text):
             return [text[: len(text) // 2], text[len(text) // 2 :]]
 
-        db_path = "/tmp/test_retrieve_utils_chromadb.db"
-        client = chromadb.PersistentClient(path=db_path)
-        create_vector_db_from_dir(
-            os.path.join(test_dir, "example.txt"),
-            client=client,
-            collection_name="mytestcollection",
+        db_path = "/tmp/test_retrieve_utils"
+        retriever = Retriever(
+            path=db_path,
+            name="mytestcollection",
             custom_text_split_function=custom_text_split_function,
-            get_or_create=True,
+            use_existing=False,
             recursive=False,
         )
-        results = query_vector_db(["autogen"], client=client, collection_name="mytestcollection", n_results=1)
+        retriever.ingest_data(os.path.join(test_dir, "example.txt"))
+        results = retriever.query(["autogen"], top_k=1)
         assert (
             "AutoGen is an advanced tool designed to assist developers in harnessing the capabilities"
             in results.get("documents")[0][0]
         )
 
     def test_retrieve_utils(self):
-        client = chromadb.PersistentClient(path="/tmp/chromadb")
-        create_vector_db_from_dir(
-            dir_path="./website/docs",
-            client=client,
-            collection_name="autogen-docs",
-            custom_text_types=["txt", "md", "rtf", "rst"],
-            get_or_create=True,
-        )
-        results = query_vector_db(
-            query_texts=[
-                "How can I use AutoGen UserProxyAgent and AssistantAgent to do code generation?",
-            ],
-            n_results=4,
-            client=client,
-            collection_name="autogen-docs",
-            search_string="AutoGen",
-        )
+        retriever = Retriever(path="/tmp/chromadb", name="autogen-docs", use_existing=False)
+        retriever.ingest_data("./website/docs")
+        results = retriever.query(["autogen"], top_k=4, filter="AutoGen")
+
         print(results["ids"][0])
         assert len(results["ids"][0]) == 4
 


Code review comments:
diff_hunk: @@ -122,9 +118,9 @@ def __init__(
                 - customized_answer_prefix (Optional, str): the customized answer prefix for the retrieve chat. Default is "".
                     If not "" and the customized_answer_prefix is not in the answer, `Update Context` will be triggered.
                 - update_context (Optional, bool): if False, will not apply `Update Context` for interactive retrieval. Default is True.
-                - get_or_create (Optional, bool): if True, will create/return a collection for the retrieve chat. This is the same as that used in chromadb.
-                    Default is False. Will raise ValueError if the collection already exists and get_or_create is False. Will be set to True if docs_path is None.
-                - custom_token_count_function (Optional, Callable): a custom function to count the number of tokens in a string.
+                - get_or_create (Optional, bool): if True, will create/recreate a collection for the retrieve chat.
path: autogen/agentchat/contrib/retrieve_user_proxy_agent.py
text: @user1:
I suggest we use a new parameter such as `db_action` and send a warning if `get_or_create` is not None, before deprecation, the two are the same.
`db_action` has 3 values: 
- get: Create a collection if it doesn't exist, get the existing collection if it exists.
- recreate: Create a collection if it doesn't exist, remove the existing one and recreate a new one if it exists.
- create: Create a collection if it doesn't exist, raise an error if it exists.

@author:
should we use `overwrite` instead of `recreate`?

@user1:
Both are OK to me. I suggested `recreate` because for chromadb, we need to actually remove the old one and create a new one. The two collections will have the same name but different id.

@author:
@user1 I've actually changed the logic of creating and querying retrievers and I think this change probably won't be required.

@user1:
I think your new design still misses `Create a collection if it doesn't exist, raise an error if it exists.`

diff_hunk: @@ -0,0 +1,17 @@
+from typing import Optional
+from .chromadb import ChromaDB
+from .lancedb import LanceDB
path: autogen/agentchat/contrib/retriever/__init__.py
text: @user1:
These two imports may need to move into the function so that any db can work without other dbs.

@author:
sounds good. I'll update

diff_hunk: @@ -372,31 +365,27 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =
         """
         if not self._collection or not self._get_or_create:
             print("Trying to create collection.")
-            self._client = create_vector_db_from_dir(
-                dir_path=self._docs_path,
+            retriever_class = get_retriever(self._retriever_type)
+            self.retriever = retriever_class(
+                name=self._collection_name,
+                embedding_model_name=self._embedding_model,
+                embedding_function=self._embedding_function,
                 max_tokens=self._chunk_token_size,
-                client=self._client,
-                collection_name=self._collection_name,
                 chunk_mode=self._chunk_mode,
                 must_break_at_empty_line=self._must_break_at_empty_line,
-                embedding_model=self._embedding_model,
-                get_or_create=self._get_or_create,
-                embedding_function=self._embedding_function,
                 custom_text_split_function=self.custom_text_split_function,
+                use_existing=not self._get_or_create,
+                client=self._client,
                 custom_text_types=self._custom_text_types,
                 recursive=self._recursive,
             )
             self._collection = True
-            self._get_or_create = True
-
-        results = query_vector_db(
-            query_texts=[problem],
-            n_results=n_results,
-            search_string=search_string,
-            client=self._client,
-            collection_name=self._collection_name,
-            embedding_model=self._embedding_model,
-            embedding_function=self._embedding_function,
+            self._get_or_create = False
+        self.retriever.ingest_data(self._docs_path)
path: autogen/agentchat/contrib/retrieve_user_proxy_agent.py
text: @user1:
Do we need to `ingest_data` at every single call of `retrieve_docs`?

@author:
Ohh not really. I think that what's the code was doing before, no? I'll just check if the index already exists and only create it if it doesn't?

@user1:
In the old code, `create_vector_db_from_dir` happens only under the condition `if not self._collection or not self._get_or_create:`, `ingest_data` is a part of `create_vector_db_from_dir`, looks like better to be under the same condition. WDYT?

@user1:
Btw, how do you `check if the index already exists`? I think it would be interesting if we could check if there are new contents in the `docs_path` and automatically add the new contents into the same collection. This could be addressed in a new PR.

@author:
@user1 I've made the change, can you please take a look again?

diff_hunk: @@ -49,8 +49,8 @@
         ],
         "blendsearch": ["flaml[blendsearch]"],
         "mathchat": ["sympy", "pydantic==1.10.9", "wolframalpha"],
-        "retrievechat": ["chromadb", "sentence_transformers", "pypdf", "ipython"],
-        "teachable": ["chromadb"],
+        "retrievechat": ["chromadb", "lancedb", "sentence_transformers", "pypdf", "ipython"],
+        "teachable": ["chromadb", "lancedb"],
path: setup.py
text: @user1:
It would be better to have only one or even no db here. We can discuss about it later.

@author:
OKay we can do that

diff_hunk: @@ -1,7 +1,7 @@
 from typing import Callable, Dict, List, Optional
 
 from autogen.agentchat.contrib.retrieve_user_proxy_agent import RetrieveUserProxyAgent
-from autogen.retrieve_utils import get_files_from_dir, split_files_to_chunks, TEXT_FORMATS
+from autogen.agentchat.contrib.retrieve_utils import get_files_from_dir, split_files_to_chunks, TEXT_FORMATS
path: autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py
text: @user1:
qdrant is another vector db like lancedb and chromadb, would you like to remove the `QdrantRetrieveUserProxyAgent` and keep only `RetrieveUserProxyAgent` ?

@user3:
Please add @user2 as a reviewer when the change involves qdrant in any PR.