code: """
  Unit test for retrieve_utils.py
  """
  try:
      import chromadb
      from autogen.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.token_count_utils import count_token
  except ImportError:
      skip = True
  else:
      skip = False
  import os
  import sys
  import pytest
  
  try:
      from unstructured.partition.auto import partition
  
      HAS_UNSTRUCTURED = True
  except ImportError:
      HAS_UNSTRUCTURED = False
  
  test_dir = os.path.join(os.path.dirname(__file__), "test_files")
  expected_text = """AutoGen is an advanced tool designed to assist developers in harnessing the capabilities
  of Large Language Models (LLMs) for various applications. The primary purpose of AutoGen is to automate and
  simplify the process of building applications that leverage the power of LLMs, allowing for seamless
  integration, testing, and deployment."""
  
  
  @pytest.mark.skipif(skip, reason="dependency is not installed")
  class TestRetrieveUtils:
      def test_split_text_to_chunks(self):
          long_text = "A" * 10000
          chunks = split_text_to_chunks(long_text, max_tokens=1000)
          assert all(count_token(chunk) <= 1000 for chunk in chunks)
  
      def test_split_text_to_chunks_raises_on_invalid_chunk_mode(self):
          with pytest.raises(AssertionError):
              split_text_to_chunks("A" * 10000, chunk_mode="bogus_chunk_mode")
  
      def test_extract_text_from_pdf(self):
          pdf_file_path = os.path.join(test_dir, "example.pdf")
          assert "".join(expected_text.split()) == "".join(extract_text_from_pdf(pdf_file_path).strip().split())
  
      def test_split_files_to_chunks(self):
          pdf_file_path = os.path.join(test_dir, "example.pdf")
          txt_file_path = os.path.join(test_dir, "example.txt")
          chunks = split_files_to_chunks([pdf_file_path, txt_file_path])
          assert all(
              isinstance(chunk, str) and "AutoGen is an advanced tool designed to assist developers" in chunk.strip()
              for chunk in chunks
          )
  
      def test_get_files_from_dir(self):
          files = get_files_from_dir(test_dir, recursive=False)
          assert all(os.path.isfile(file) for file in files)
          pdf_file_path = os.path.join(test_dir, "example.pdf")
          txt_file_path = os.path.join(test_dir, "example.txt")
          files = get_files_from_dir([pdf_file_path, txt_file_path])
          assert all(os.path.isfile(file) for file in files)
          files = get_files_from_dir(
              [
                  pdf_file_path,
                  txt_file_path,
                  os.path.join(test_dir, "..", "..", "website/docs"),
                  "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
              ],
              recursive=True,
          )
          assert all(os.path.isfile(file) for file in files)
          files = get_files_from_dir(
              [
                  pdf_file_path,
                  txt_file_path,
                  os.path.join(test_dir, "..", "..", "website/docs"),
                  "https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
              ],
              recursive=True,
              types=["pdf", "txt"],
          )
          assert all(os.path.isfile(file) for file in files)
          assert len(files) == 3
  
      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",
              custom_text_split_function=custom_text_split_function,
              get_or_create=True,
              recursive=False,
          )
          results = query_vector_db(["autogen"], client=client, collection_name="mytestcollection", n_results=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",
          )
          print(results["ids"][0])
          assert len(results["ids"][0]) == 4
  
      @pytest.mark.skipif(
          not HAS_UNSTRUCTURED,
          reason="do not run if unstructured is not installed",
      )
      def test_unstructured(self):
          pdf_file_path = os.path.join(test_dir, "example.pdf")
          txt_file_path = os.path.join(test_dir, "example.txt")
          word_file_path = os.path.join(test_dir, "example.docx")
          chunks = split_files_to_chunks([pdf_file_path, txt_file_path, word_file_path])
          assert all(
              isinstance(chunk, str) and "AutoGen is an advanced tool designed to assist developers" in chunk.strip()
              for chunk in chunks
          )
  
  
  if __name__ == "__main__":
      pytest.main()
  
      db_path = "/tmp/test_retrieve_utils_chromadb.db"
      if os.path.exists(db_path):
          os.remove(db_path)  # Delete the database file after tests are finished
  
apis: lancedb.connect
extract_api: [((1021, 1083), 'pytest.mark.skipif', 'pytest.mark.skipif', (['skip'], {'reason': '"""dependency is not installed"""'}), "(skip, reason='dependency is not installed')\n", (1039, 1083), False, 'import pytest\n'), ((619, 644), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (634, 644), False, 'import os\n'), ((8695, 8794), 'pytest.mark.skipif', 'pytest.mark.skipif', (['(not HAS_UNSTRUCTURED)'], {'reason': '"""do not run if unstructured is not installed"""'}), "(not HAS_UNSTRUCTURED, reason=\n    'do not run if unstructured is not installed')\n", (8713, 8794), False, 'import pytest\n'), ((9332, 9345), 'pytest.main', 'pytest.main', ([], {}), '()\n', (9343, 9345), False, 'import pytest\n'), ((9407, 9430), 'os.path.exists', 'os.path.exists', (['db_path'], {}), '(db_path)\n', (9421, 9430), False, 'import os\n'), ((1199, 1247), 'autogen.retrieve_utils.split_text_to_chunks', 'split_text_to_chunks', (['long_text'], {'max_tokens': '(1000)'}), '(long_text, max_tokens=1000)\n', (1219, 1247), False, 'from autogen.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\n'), ((1574, 1611), 'os.path.join', 'os.path.join', (['test_dir', '"""example.pdf"""'], {}), "(test_dir, 'example.pdf')\n", (1586, 1611), False, 'import os\n'), ((1790, 1827), 'os.path.join', 'os.path.join', (['test_dir', '"""example.pdf"""'], {}), "(test_dir, 'example.pdf')\n", (1802, 1827), False, 'import os\n'), ((1852, 1889), 'os.path.join', 'os.path.join', (['test_dir', '"""example.txt"""'], {}), "(test_dir, 'example.txt')\n", (1864, 1889), False, 'import os\n'), ((1907, 1960), 'autogen.retrieve_utils.split_files_to_chunks', 'split_files_to_chunks', (['[pdf_file_path, txt_file_path]'], {}), '([pdf_file_path, txt_file_path])\n', (1928, 1960), False, 'from autogen.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\n'), ((2195, 2240), 'autogen.retrieve_utils.get_files_from_dir', 'get_files_from_dir', (['test_dir'], {'recursive': '(False)'}), '(test_dir, recursive=False)\n', (2213, 2240), False, 'from autogen.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\n'), ((2324, 2361), 'os.path.join', 'os.path.join', (['test_dir', '"""example.pdf"""'], {}), "(test_dir, 'example.pdf')\n", (2336, 2361), False, 'import os\n'), ((2386, 2423), 'os.path.join', 'os.path.join', (['test_dir', '"""example.txt"""'], {}), "(test_dir, 'example.txt')\n", (2398, 2423), False, 'import os\n'), ((2440, 2490), 'autogen.retrieve_utils.get_files_from_dir', 'get_files_from_dir', (['[pdf_file_path, txt_file_path]'], {}), '([pdf_file_path, txt_file_path])\n', (2458, 2490), False, 'from autogen.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\n'), ((3414, 3447), 'autogen.retrieve_utils.is_url', 'is_url', (['"""https://www.example.com"""'], {}), "('https://www.example.com')\n", (3420, 3447), False, 'from autogen.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\n'), ((3602, 3625), 'os.path.exists', 'os.path.exists', (['db_path'], {}), '(db_path)\n', (3616, 3625), False, 'import os\n'), ((3989, 4012), 'os.path.exists', 'os.path.exists', (['db_path'], {}), '(db_path)\n', (4003, 4012), False, 'import os\n'), ((4283, 4326), 'autogen.retrieve_utils.query_vector_db', 'query_vector_db', (["['autogen']"], {'client': 'client'}), "(['autogen'], client=client)\n", (4298, 4326), False, 'from autogen.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\n'), ((7323, 7362), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'db_path'}), '(path=db_path)\n', (7348, 7362), False, 'import chromadb\n'), ((7680, 7777), 'autogen.retrieve_utils.query_vector_db', 'query_vector_db', (["['autogen']"], {'client': 'client', 'collection_name': '"""mytestcollection"""', 'n_results': '(1)'}), "(['autogen'], client=client, collection_name=\n    'mytestcollection', n_results=1)\n", (7695, 7777), False, 'from autogen.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\n'), ((8002, 8049), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': '"""/tmp/chromadb"""'}), "(path='/tmp/chromadb')\n", (8027, 8049), False, 'import chromadb\n'), ((8058, 8232), 'autogen.retrieve_utils.create_vector_db_from_dir', '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)'}), "(dir_path='./website/docs', client=client,\n    collection_name='autogen-docs', custom_text_types=['txt', 'md', 'rtf',\n    'rst'], get_or_create=True)\n", (8083, 8232), False, 'from autogen.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\n'), ((8314, 8524), 'autogen.retrieve_utils.query_vector_db', 'query_vector_db', ([], {'query_texts': "['How can I use AutoGen UserProxyAgent and AssistantAgent to do code generation?'\n    ]", 'n_results': '(4)', 'client': 'client', 'collection_name': '"""autogen-docs"""', 'search_string': '"""AutoGen"""'}), "(query_texts=[\n    'How can I use AutoGen UserProxyAgent and AssistantAgent to do code generation?'\n    ], n_results=4, client=client, collection_name='autogen-docs',\n    search_string='AutoGen')\n", (8329, 8524), False, 'from autogen.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\n'), ((8870, 8907), 'os.path.join', 'os.path.join', (['test_dir', '"""example.pdf"""'], {}), "(test_dir, 'example.pdf')\n", (8882, 8907), False, 'import os\n'), ((8932, 8969), 'os.path.join', 'os.path.join', (['test_dir', '"""example.txt"""'], {}), "(test_dir, 'example.txt')\n", (8944, 8969), False, 'import os\n'), ((8995, 9033), 'os.path.join', 'os.path.join', (['test_dir', '"""example.docx"""'], {}), "(test_dir, 'example.docx')\n", (9007, 9033), False, 'import os\n'), ((9051, 9120), 'autogen.retrieve_utils.split_files_to_chunks', 'split_files_to_chunks', (['[pdf_file_path, txt_file_path, word_file_path]'], {}), '([pdf_file_path, txt_file_path, word_file_path])\n', (9072, 9120), False, 'from autogen.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\n'), ((9440, 9458), 'os.remove', 'os.remove', (['db_path'], {}), '(db_path)\n', (9449, 9458), False, 'import os\n'), ((1399, 1428), 'pytest.raises', 'pytest.raises', (['AssertionError'], {}), '(AssertionError)\n', (1412, 1428), False, 'import pytest\n'), ((1442, 1506), 'autogen.retrieve_utils.split_text_to_chunks', 'split_text_to_chunks', (["('A' * 10000)"], {'chunk_mode': '"""bogus_chunk_mode"""'}), "('A' * 10000, chunk_mode='bogus_chunk_mode')\n", (1462, 1506), False, 'from autogen.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\n'), ((3467, 3486), 'autogen.retrieve_utils.is_url', 'is_url', (['"""not_a_url"""'], {}), "('not_a_url')\n", (3473, 3486), False, 'from autogen.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\n'), ((3648, 3687), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'db_path'}), '(path=db_path)\n', (3673, 3687), False, 'import chromadb\n'), ((3723, 3762), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'db_path'}), '(path=db_path)\n', (3748, 3762), False, 'import chromadb\n'), ((3775, 3825), 'autogen.retrieve_utils.create_vector_db_from_dir', 'create_vector_db_from_dir', (['test_dir'], {'client': 'client'}), '(test_dir, client=client)\n', (3800, 3825), False, 'from autogen.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\n'), ((4035, 4074), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'db_path'}), '(path=db_path)\n', (4060, 4074), False, 'import chromadb\n'), ((4161, 4200), 'chromadb.PersistentClient', 'chromadb.PersistentClient', ([], {'path': 'db_path'}), '(path=db_path)\n', (4186, 4200), False, 'import chromadb\n'), ((4213, 4263), 'autogen.retrieve_utils.create_vector_db_from_dir', 'create_vector_db_from_dir', (['test_dir'], {'client': 'client'}), '(test_dir, client=client)\n', (4238, 4263), False, 'from autogen.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\n'), ((4747, 4771), 'lancedb.connect', 'lancedb.connect', (['db_path'], {}), '(db_path)\n', (4762, 4771), False, 'import lancedb\n'), ((7410, 7447), 'os.path.join', 'os.path.join', (['test_dir', '"""example.txt"""'], {}), "(test_dir, 'example.txt')\n", (7422, 7447), False, 'import os\n'), ((2260, 2280), 'os.path.isfile', 'os.path.isfile', (['file'], {}), '(file)\n', (2274, 2280), False, 'import os\n'), ((2510, 2530), 'os.path.isfile', 'os.path.isfile', (['file'], {}), '(file)\n', (2524, 2530), False, 'import os\n'), ((2678, 2728), 'os.path.join', 'os.path.join', (['test_dir', '""".."""', '""".."""', '"""website/docs"""'], {}), "(test_dir, '..', '..', 'website/docs')\n", (2690, 2728), False, 'import os\n'), ((2888, 2908), 'os.path.isfile', 'os.path.isfile', (['file'], {}), '(file)\n', (2902, 2908), False, 'import os\n'), ((3056, 3106), 'os.path.join', 'os.path.join', (['test_dir', '""".."""', '""".."""', '"""website/docs"""'], {}), "(test_dir, '..', '..', 'website/docs')\n", (3068, 3106), False, 'import os\n'), ((3300, 3320), 'os.path.isfile', 'os.path.isfile', (['file'], {}), '(file)\n', (3314, 3320), False, 'import os\n'), ((5827, 5851), 'lancedb.connect', 'lancedb.connect', (['db_path'], {}), '(db_path)\n', (5842, 5851), False, 'import lancedb\n'), ((1267, 1285), 'autogen.token_count_utils.count_token', 'count_token', (['chunk'], {}), '(chunk)\n', (1278, 1285), False, 'from autogen.token_count_utils import count_token\n'), ((1669, 1705), 'autogen.retrieve_utils.extract_text_from_pdf', 'extract_text_from_pdf', (['pdf_file_path'], {}), '(pdf_file_path)\n', (1690, 1705), False, 'from autogen.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\n')]