code: # %% [markdown]
  # # Code documentation Q&A bot example with LangChain
  #
  # This Q&A bot will allow you to query your own documentation easily using questions. We'll also demonstrate the use of LangChain and LanceDB using the OpenAI API.
  #
  # In this example we'll use Pandas 2.0 documentation, but, this could be replaced for your own docs as well
  
  import os
  import openai
  import argparse
  import lancedb
  import re
  import pickle
  import requests
  import zipfile
  from pathlib import Path
  
  from langchain.document_loaders import BSHTMLLoader
  from langchain.embeddings import OpenAIEmbeddings
  from langchain.text_splitter import RecursiveCharacterTextSplitter
  from langchain.vectorstores import LanceDB
  from langchain.llms import OpenAI
  from langchain.chains import RetrievalQA
  
  
  def get_document_title(document):
      m = str(document.metadata["source"])
      title = re.findall("pandas.documentation(.*).html", m)
      if title[0] is not None:
          return title[0]
      return ""
  
  
  def arg_parse():
      default_query = "What are the major differences in pandas 2.0?"
  
      parser = argparse.ArgumentParser(description="Code Documentation QA Bot")
      parser.add_argument(
          "--query", type=str, default=default_query, help="query to search"
      )
      parser.add_argument("--openai-key", type=str, help="OpenAI API Key")
      args = parser.parse_args()
  
      if not args.openai_key:
          if "OPENAI_API_KEY" not in os.environ:
              raise ValueError(
                  "OPENAI_API_KEY environment variable not set. Please set it or pass --openai_key"
              )
      else:
          openai.api_key = args.openai_key
  
      return args
  
  
  if __name__ == "__main__":
      args = arg_parse()
  
      docs_path = Path("docs.pkl")
      docs = []
  
      pandas_docs = requests.get(
          "https://eto-public.s3.us-west-2.amazonaws.com/datasets/pandas_docs/pandas.documentation.zip"
      )
      with open("/tmp/pandas.documentation.zip", "wb") as f:
          f.write(pandas_docs.content)
  
      file = zipfile.ZipFile("/tmp/pandas.documentation.zip")
      file.extractall(path="/tmp/pandas_docs")
  
      if not docs_path.exists():
          for p in Path("/tmp/pandas_docs/pandas.documentation").rglob("*.html"):
              print(p)
              if p.is_dir():
                  continue
              loader = BSHTMLLoader(p, open_encoding="utf8")
              raw_document = loader.load()
  
              m = {}
              m["title"] = get_document_title(raw_document[0])
              m["version"] = "2.0rc0"
              raw_document[0].metadata = raw_document[0].metadata | m
              raw_document[0].metadata["source"] = str(raw_document[0].metadata["source"])
              docs = docs + raw_document
  
          with docs_path.open("wb") as fh:
              pickle.dump(docs, fh)
      else:
          with docs_path.open("rb") as fh:
              docs = pickle.load(fh)
      print("Loaded {} documents".format(len(docs)))
      text_splitter = RecursiveCharacterTextSplitter(
          chunk_size=1000,
          chunk_overlap=200,
      )
      documents = text_splitter.split_documents(docs)
      embeddings = OpenAIEmbeddings()
  
      db = lancedb.connect("/tmp/lancedb")
      table = db.create_table(
          "pandas_docs",
          data=[
              {
                  "vector": embeddings.embed_query("Hello World"),
                  "text": "Hello World",
                  "id": "1",
              }
          ],
          mode="overwrite",
      )
      docsearch = LanceDB.from_documents(documents, embeddings, connection=table)
  
      qa = RetrievalQA.from_chain_type(
          llm=OpenAI(), chain_type="stuff", retriever=docsearch.as_retriever()
      )
  
      result = qa.run(args.query)
      print(result)
  
apis: lancedb.connect
extract_api: [((859, 905), 're.findall', 're.findall', (['"""pandas.documentation(.*).html"""', 'm'], {}), "('pandas.documentation(.*).html', m)\n", (869, 905), False, 'import re\n'), ((1074, 1138), 'argparse.ArgumentParser', 'argparse.ArgumentParser', ([], {'description': '"""Code Documentation QA Bot"""'}), "(description='Code Documentation QA Bot')\n", (1097, 1138), False, 'import argparse\n'), ((1704, 1720), 'pathlib.Path', 'Path', (['"""docs.pkl"""'], {}), "('docs.pkl')\n", (1708, 1720), False, 'from pathlib import Path\n'), ((1754, 1871), 'requests.get', 'requests.get', (['"""https://eto-public.s3.us-west-2.amazonaws.com/datasets/pandas_docs/pandas.documentation.zip"""'], {}), "(\n    'https://eto-public.s3.us-west-2.amazonaws.com/datasets/pandas_docs/pandas.documentation.zip'\n    )\n", (1766, 1871), False, 'import requests\n'), ((1984, 2032), 'zipfile.ZipFile', 'zipfile.ZipFile', (['"""/tmp/pandas.documentation.zip"""'], {}), "('/tmp/pandas.documentation.zip')\n", (1999, 2032), False, 'import zipfile\n'), ((2909, 2975), 'langchain.text_splitter.RecursiveCharacterTextSplitter', 'RecursiveCharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(200)'}), '(chunk_size=1000, chunk_overlap=200)\n', (2939, 2975), False, 'from langchain.text_splitter import RecursiveCharacterTextSplitter\n'), ((3068, 3086), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (3084, 3086), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((3097, 3128), 'lancedb.connect', 'lancedb.connect', (['"""/tmp/lancedb"""'], {}), "('/tmp/lancedb')\n", (3112, 3128), False, 'import lancedb\n'), ((3414, 3477), 'langchain.vectorstores.LanceDB.from_documents', 'LanceDB.from_documents', (['documents', 'embeddings'], {'connection': 'table'}), '(documents, embeddings, connection=table)\n', (3436, 3477), False, 'from langchain.vectorstores import LanceDB\n'), ((2284, 2321), 'langchain.document_loaders.BSHTMLLoader', 'BSHTMLLoader', (['p'], {'open_encoding': '"""utf8"""'}), "(p, open_encoding='utf8')\n", (2296, 2321), False, 'from langchain.document_loaders import BSHTMLLoader\n'), ((2730, 2751), 'pickle.dump', 'pickle.dump', (['docs', 'fh'], {}), '(docs, fh)\n', (2741, 2751), False, 'import pickle\n'), ((2822, 2837), 'pickle.load', 'pickle.load', (['fh'], {}), '(fh)\n', (2833, 2837), False, 'import pickle\n'), ((3529, 3537), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (3535, 3537), False, 'from langchain.llms import OpenAI\n'), ((2127, 2172), 'pathlib.Path', 'Path', (['"""/tmp/pandas_docs/pandas.documentation"""'], {}), "('/tmp/pandas_docs/pandas.documentation')\n", (2131, 2172), False, 'from pathlib import Path\n')]