code: from dotenv import load_dotenv
  from typing import List
  import lancedb
  import openai
  import os
  
  from gpt_pdf_bot.shared import embed_text
  
  
  load_dotenv()
  
  
  class ChatBot:
      def __init__(self, table_name: str):
          self.db = lancedb.connect(uri=".lancedb")
          self.table = self.db[table_name]
  
      def run(self):
          while True:
              query = input("\nAsk me a question: ")
  
              if query == "exit":
                  break
  
              print(f"Searching for {query}", end="\n\n")
              context = self.retrieve_context(query=query)
              print(f"Found {len(context)} results", end="\n\n")
  
              texts = self.retrieve_text_from_context(context=context)
              annotations = self.create_annotations_from_context(context=context)
  
              prompt = self.generate_prompt(query=query, context=texts)
  
              response = self.ask_gpt(prompt, annotations=annotations)
              print(response, end="\n\n")
  
      def retrieve_context(self, query: str, limit: int = 5) -> List[str]:
          emb = embed_text([query])[0]
  
          # TODO: return and pass metadata so we can display the source and page number
          context = self.table.search(emb).limit(limit).to_df()
          return context
  
      def retrieve_text_from_context(self, context: List[str]) -> List[str]:
          """Extracts only the text from the context, so it can be used in the prompt"""
          # curly braces will mess up prompt templating, so remove them
          texts = context["text"].str.replace("{", "").str.replace("}", "").tolist()
          return texts
  
      def create_annotations_from_context(self, context: List[str]) -> List[str]:
          """Extracts the document metadata from the context,
          so it can be displayed with the LLM's response"""
          metadata = context["metadata"].tolist()
  
          annotations = []
          for meta in metadata:
              source = meta["source"]
              page_num = meta["page_num"]
              annotations.append(f"- file name: {source}, page number: {page_num}")
          return annotations
  
      def generate_prompt(self, query: str, context: List[str]) -> str:
          # TODO: this is a naive implementation that doesn't handle the case where the context is too long
          prepared_context = "\n\n---\n\n".join(context)
          prompt_tmpl = f"""Answer the question based on the context below.
              {context}
  
              Question: {query}
              Answer:"""
          return prompt_tmpl.format(query=query, context=prepared_context)
  
      def prepare_messages_for_chat(self, text: str) -> list:
          messages = [
              {"role": "user", "content": text},
          ]
          return messages
  
      def ask_gpt(self, prompt: str, annotations: List[str]) -> str:
          res = openai.ChatCompletion.create(
              model="gpt-3.5-turbo",
              messages=self.prepare_messages_for_chat(prompt),
              temperature=0.7,
              api_key=os.environ["OPENAI_API_KEY"],
          )
          response = res["choices"][0]["message"]["content"].strip()
          annotations = "\n".join(annotations).strip()
  
          answer = f"Answer: \n\n{response}\n\n---\n\n"
          answer += f"Here are the sources I used:\n{annotations}"
          return answer
  
apis: lancedb.connect
extract_api: [((139, 152), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (150, 152), False, 'from dotenv import load_dotenv\n'), ((229, 260), 'lancedb.connect', 'lancedb.connect', ([], {'uri': '""".lancedb"""'}), "(uri='.lancedb')\n", (244, 260), False, 'import lancedb\n'), ((1044, 1063), 'gpt_pdf_bot.shared.embed_text', 'embed_text', (['[query]'], {}), '([query])\n', (1054, 1063), False, 'from gpt_pdf_bot.shared import embed_text\n')]