code: from typing import List, Optional, Tuple, Dict, Callable, Any, Union
  from functools import reduce
  import os
  import os
  from pathlib import Path
  import re
  from .utils import maybe_is_text, maybe_is_truncated
  from .qaprompts import (
      summary_prompt,
      qa_prompt,
      search_prompt,
      citation_prompt,
      make_chain,
  )
  from dataclasses import dataclass
  from .readers import read_doc
  from langchain.vectorstores import FAISS
  from langchain.embeddings.openai import OpenAIEmbeddings
  from langchain.chat_models import ChatOpenAI
  from langchain.llms.base import LLM
  from langchain.chains import LLMChain
  from langchain.callbacks import get_openai_callback
  from langchain.cache import SQLiteCache
  import langchain
  from datetime import datetime
  
  CACHE_PATH = Path.home() / ".paperqa" / "llm_cache.db"
  os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
  langchain.llm_cache = SQLiteCache(CACHE_PATH)
  
  
  @dataclass
  class Answer:
      """A class to hold the answer to a question."""
  
      question: str
      answer: str = ""
      context: str = ""
      contexts: List[Any] = None
      references: str = ""
      formatted_answer: str = ""
      passages: Dict[str, str] = None
      tokens: int = 0
      question_embedding: List[float] = None
      from_embed: bool = False
      def __post_init__(self):
          """Initialize the answer."""
          if self.contexts is None:
              self.contexts = []
          if self.passages is None:
              self.passages = {}
  
      def __str__(self) -> str:
          """Return the answer as a string."""
          return self.formatted_answer
  
  
  class Docs:
      """A collection of documents to be used for answering questions."""
  
      def __init__(
          self,
          chunk_size_limit: int = 3000,
          llm: Optional[Union[LLM, str]] = None,
          summary_llm: Optional[Union[LLM, str]] = None,
          name: str = "default",
          index_path: Optional[Path] = None,
      ) -> None:
          """Initialize the collection of documents.
  
  
  
          Args:
              chunk_size_limit: The maximum number of characters to use for a single chunk of text.
              llm: The language model to use for answering questions. Default - OpenAI chat-gpt-turbo
              summary_llm: The language model to use for summarizing documents. If None, llm is used.
              name: The name of the collection.
              index_path: The path to the index file IF pickled. If None, defaults to using name in $HOME/.paperqa/name
          """
          self.docs = dict()  # self.docs[path] = dict(texts=texts, metadata=metadata, key=key)
          self.chunk_size_limit = chunk_size_limit
          self.keys = set()
          self._faiss_index = None
          self.update_llm(llm, summary_llm)
          if index_path is None:
              index_path = Path.home() / ".paperqa" / name
          self.index_path = index_path
          self.name = name
  
      def update_llm(
          self,
          llm: Optional[Union[LLM, str]] = None,
          summary_llm: Optional[Union[LLM, str]] = None,
      ) -> None:
          """Update the LLM for answering questions."""
          if llm is None:
              llm = "gpt-3.5-turbo"
          if type(llm) is str:
              llm = ChatOpenAI(temperature=0.1, model=llm)
          if type(summary_llm) is str:
              summary_llm = ChatOpenAI(temperature=0.1, model=summary_llm)
          self.llm = llm
          if summary_llm is None:
              summary_llm = llm
          self.summary_llm = summary_llm
          self.summary_chain = make_chain(prompt=summary_prompt, llm=summary_llm)
          self.qa_chain = make_chain(prompt=qa_prompt, llm=llm)
          self.search_chain = make_chain(prompt=search_prompt, llm=summary_llm)
          self.cite_chain = make_chain(prompt=citation_prompt, llm=summary_llm)
  
      def add(
          self,
          path: str,
          citation: Optional[str] = None,
          key: Optional[str] = None,
          disable_check: bool = False,
          chunk_chars: Optional[int] = 3000,
      ) -> None:
          """Add a document to the collection."""
  
          # first check to see if we already have this document
          # this way we don't make api call to create citation on file we already have
          if path in self.docs:
              raise ValueError(f"Document {path} already in collection.")
  
          if citation is None:
              # peak first chunk
              texts, _ = read_doc(path, "", "", chunk_chars=chunk_chars)
              with get_openai_callback() as cb:
                  citation = self.cite_chain.run(texts[0])
              if len(citation) < 3 or "Unknown" in citation or "insufficient" in citation:
                  citation = f"Unknown, {os.path.basename(path)}, {datetime.now().year}"
  
          if key is None:
              # get first name and year from citation
              try:
                  author = re.search(r"([A-Z][a-z]+)", citation).group(1)
              except AttributeError:
                  # panicking - no word??
                  raise ValueError(
                      f"Could not parse key from citation {citation}. Consider just passing key explicitly - e.g. docs.py (path, citation, key='mykey')"
                  )
              try:
                  year = re.search(r"(\d{4})", citation).group(1)
              except AttributeError:
                  year = ""
              key = f"{author}{year}"
          suffix = ""
          while key + suffix in self.keys:
              # move suffix to next letter
              if suffix == "":
                  suffix = "a"
              else:
                  suffix = chr(ord(suffix) + 1)
          key += suffix
          self.keys.add(key)
  
          texts, metadata = read_doc(path, citation, key, chunk_chars=chunk_chars)
  
          # loose check to see if document was loaded
          if len("".join(texts)) < 10 or (
              not disable_check and not maybe_is_text("".join(texts))
          ):
              raise ValueError(
                  f"This does not look like a text document: {path}. Path disable_check to ignore this error."
              )
  
          self.docs[path] = dict(texts=texts, metadata=metadata, key=key)
          if self._faiss_index is not None:
              self._faiss_index.add_texts(texts, metadatas=metadata)
  
      def add_from_embeddings(
              self,
              path: str,
              texts,
              text_embeddings: List[float],
              metadatas
  
      ) -> None:
  
          """Add a document to the collection."""
  
          # first check to see if we already have this document
          # this way we don't make api call to create citation on file we already have
          if path in self.docs:
              raise ValueError(f"Document {path} already in collection.")
  
          key = metadatas[0]['dockey']
  
          suffix = ""
          while key + suffix in self.keys:
              # move suffix to next letter
              if suffix == "":
                  suffix = "a"
              else:
                  suffix = chr(ord(suffix) + 1)
          key += suffix
          self.keys.add(key)
  
          if key != metadatas[0]['dockey']:
              for j in range(len(metadatas)):
                  metadatas[j]['dockey'] = key
          
          self.docs[path] = dict(texts=texts, metadata=metadatas, key=key)
  
          if self._faiss_index is not None:
              self._faiss_index.add_embeddings([*zip(texts, text_embeddings)],
                                               metadatas=metadatas)
          else:
              """Instantiate FAISS"""
              self._faiss_index = FAISS.from_embeddings([*zip(texts, text_embeddings)],
                                                        metadatas=metadatas,
                                                        embedding=OpenAIEmbeddings())
  
      def clear(self) -> None:
          """Clear the collection of documents."""
          self.docs = dict()
          self.keys = set()
          self._faiss_index = None
          # delete index file
          pkl = self.index_path / "index.pkl"
          if pkl.exists():
              pkl.unlink()
          fs = self.index_path / "index.faiss"
          if fs.exists():
              fs.unlink()
  
      @property
      def doc_previews(self) -> List[Tuple[int, str, str]]:
          """Return a list of tuples of (key, citation) for each document."""
          return [
              (
                  len(doc["texts"]),
                  doc["metadata"][0]["dockey"],
                  doc["metadata"][0]["citation"],
              )
              for doc in self.docs.values()
          ]
  
      # to pickle, we have to save the index as a file
      def __getstate__(self):
          if self._faiss_index is None and len(self.docs) > 0:
              self._build_faiss_index()
          state = self.__dict__.copy()
          if self._faiss_index is not None:
              state["_faiss_index"].save_local(self.index_path)
          del state["_faiss_index"]
          # remove LLMs (they can have callbacks, which can't be pickled)
          del state["summary_chain"]
          del state["qa_chain"]
          del state["cite_chain"]
          del state["search_chain"]
          return state
  
      def __setstate__(self, state):
          self.__dict__.update(state)
          try:
              self._faiss_index = FAISS.load_local(self.index_path, OpenAIEmbeddings())
          except:
              # they use some special exception type, but I don't want to import it
              self._faiss_index = None
          self.update_llm("gpt-3.5-turbo")
  
      def _build_faiss_index(self):
          if self._faiss_index is None:
              texts = reduce(
                  lambda x, y: x + y, [doc["texts"] for doc in self.docs.values()], []
              )
              metadatas = reduce(
                  lambda x, y: x + y, [doc["metadata"] for doc in self.docs.values()], []
              )
              self._faiss_index = FAISS.from_texts(
                  texts, OpenAIEmbeddings(), metadatas=metadatas
              )
  
      def get_evidence(
          self,
          answer: Answer,
          k: int = 3,
          max_sources: int = 5,
          marginal_relevance: bool = True,
          key_filter: Optional[List[str]] = None,
      ) -> str:
          if self._faiss_index is None:
              self._build_faiss_index()
          _k = k
          if key_filter is not None:
              _k = k * 10  # heuristic
          # want to work through indices but less k
          if marginal_relevance:
              if not answer.from_embed:
                  docs = self._faiss_index.max_marginal_relevance_search(
                      answer.question, k=_k, fetch_k=5 * _k
                  )
              else:
                  docs = self._faiss_index.max_marginal_relevance_search_by_vector(
                      answer.question_embedding, k=_k, fetch_k=5 * _k
                  )
          else:
              docs = self._faiss_index.similarity_search(
                  answer.question, k=_k, fetch_k=5 * _k
              )
          for doc in docs:
              if key_filter is not None and doc.metadata["dockey"] not in key_filter:
                  continue
              c = (
                  doc.metadata["key"],
                  doc.metadata["citation"],
                  self.summary_chain.run(
                      question=answer.question,
                      context_str=doc.page_content,
                      citation=doc.metadata["citation"],
                  ),
                  doc.page_content,
              )
              if "Not applicable" not in c[2]:
                  answer.contexts.append(c)
                  yield answer
              if len(answer.contexts) == max_sources:
                  break
          context_str = "\n\n".join(
              [f"{k}: {s}" for k, c, s, t in answer.contexts if "Not applicable" not in s]
          )
          valid_keys = [k for k, c, s, t in answer.contexts if "Not applicable" not in s]
          if len(valid_keys) > 0:
              context_str += "\n\nValid keys: " + ", ".join(valid_keys)
          answer.context = context_str
          yield answer
  
      def generate_search_query(self, query: str) -> List[str]:
          """Generate a list of search strings that can be used to find
          relevant papers.
  
          Args:
              query (str): The query to generate search strings for.
          """
  
          search_query = self.search_chain.run(question=query)
          queries = [s for s in search_query.split("\n") if len(s) > 3]
          # remove 2., 3. from queries
          queries = [re.sub(r"^\d+\.\s*", "", q) for q in queries]
          return queries
  
      def query_gen(
          self,
          query: str,
          k: int = 10,
          max_sources: int = 5,
          length_prompt: str = "about 100 words",
          marginal_relevance: bool = True,
      ):
          yield from self._query(
              query,
              k=k,
              max_sources=max_sources,
              length_prompt=length_prompt,
              marginal_relevance=marginal_relevance,
          )
  
      def query(
          self,
          query: str,
          k: int = 10,
          max_sources: int = 5,
          length_prompt: str = "about 100 words",
          marginal_relevance: bool = True,
          embedding: Optional[List[float]] = None
      ):
          for answer in self._query(
              query,
              k=k,
              max_sources=max_sources,
              length_prompt=length_prompt,
              marginal_relevance=marginal_relevance,
              embedding=embedding
          ):
              pass
          return answer
  
      def _query(
          self,
          query: str,
          k: int,
          max_sources: int,
          length_prompt: str,
          marginal_relevance: bool,
          embedding
      ):
          if k < max_sources:
              raise ValueError("k should be greater than max_sources")
          tokens = 0
          answer = Answer(query, question_embedding=embedding)
          if embedding is not None:
              answer.question_embedding = embedding
              answer.from_embed = True
  
          with get_openai_callback() as cb:
              for answer in self.get_evidence(
                  answer,
                  k=k,
                  max_sources=max_sources,
                  marginal_relevance=marginal_relevance,
              ):
                  yield answer
              tokens += cb.total_tokens
          context_str, citations = answer.context, answer.contexts
          bib = dict()
          passages = dict()
          if len(context_str) < 10:
              answer_text = (
                  "I cannot answer this question due to insufficient information."
              )
          else:
              with get_openai_callback() as cb:
                  answer_text = self.qa_chain.run(
                      question=query, context_str=context_str, length=length_prompt
                  )
                  tokens += cb.total_tokens
          # it still happens lol
          if "(Foo2012)" in answer_text:
              answer_text = answer_text.replace("(Foo2012)", "")
          for key, citation, summary, text in citations:
              # do check for whole key (so we don't catch Callahan2019a with Callahan2019)
              skey = key.split(" ")[0]
              if skey + " " in answer_text or skey + ")" or skey + "," in answer_text:
                  bib[skey] = citation
                  passages[key] = text
          bib_str = "\n\n".join(
              [f"{i+1}. ({k}): {c}" for i, (k, c) in enumerate(bib.items())]
          )
          formatted_answer = f"Question: {query}\n\n{answer_text}\n"
          if len(bib) > 0:
              formatted_answer += f"\nReferences\n\n{bib_str}\n"
          formatted_answer += f"\nTokens Used: {tokens} Cost: ${tokens/1000 * 0.002:.2f}"
          answer.answer = answer_text
          answer.formatted_answer = formatted_answer
          answer.references = bib_str
          answer.passages = passages
          answer.tokens = tokens
          yield answer
  
apis: langchain.embeddings.openai.OpenAIEmbeddings langchain.cache.SQLiteCache langchain.chat_models.ChatOpenAI langchain.callbacks.get_openai_callback
extract_api: [((906, 929), 'langchain.cache.SQLiteCache', 'SQLiteCache', (['CACHE_PATH'], {}), '(CACHE_PATH)\n', (917, 929), False, 'from langchain.cache import SQLiteCache\n'), ((839, 866), 'os.path.dirname', 'os.path.dirname', (['CACHE_PATH'], {}), '(CACHE_PATH)\n', (854, 866), False, 'import os\n'), ((784, 795), 'pathlib.Path.home', 'Path.home', ([], {}), '()\n', (793, 795), False, 'from pathlib import Path\n'), ((3260, 3298), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.1)', 'model': 'llm'}), '(temperature=0.1, model=llm)\n', (3270, 3298), False, 'from langchain.chat_models import ChatOpenAI\n'), ((3364, 3410), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'temperature': '(0.1)', 'model': 'summary_llm'}), '(temperature=0.1, model=summary_llm)\n', (3374, 3410), False, 'from langchain.chat_models import ChatOpenAI\n'), ((12572, 12601), 're.sub', 're.sub', (['"""^\\\\d+\\\\.\\\\s*"""', '""""""', 'q'], {}), "('^\\\\d+\\\\.\\\\s*', '', q)\n", (12578, 12601), False, 'import re\n'), ((14124, 14145), 'langchain.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (14143, 14145), False, 'from langchain.callbacks import get_openai_callback\n'), ((4525, 4546), 'langchain.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (4544, 4546), False, 'from langchain.callbacks import get_openai_callback\n'), ((9363, 9381), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (9379, 9381), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((9981, 9999), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (9997, 9999), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((14738, 14759), 'langchain.callbacks.get_openai_callback', 'get_openai_callback', ([], {}), '()\n', (14757, 14759), False, 'from langchain.callbacks import get_openai_callback\n'), ((2842, 2853), 'pathlib.Path.home', 'Path.home', ([], {}), '()\n', (2851, 2853), False, 'from pathlib import Path\n'), ((7812, 7830), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (7828, 7830), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((4742, 4764), 'os.path.basename', 'os.path.basename', (['path'], {}), '(path)\n', (4758, 4764), False, 'import os\n'), ((4914, 4950), 're.search', 're.search', (['"""([A-Z][a-z]+)"""', 'citation'], {}), "('([A-Z][a-z]+)', citation)\n", (4923, 4950), False, 'import re\n'), ((5286, 5317), 're.search', 're.search', (['"""(\\\\d{4})"""', 'citation'], {}), "('(\\\\d{4})', citation)\n", (5295, 5317), False, 'import re\n'), ((4768, 4782), 'datetime.datetime.now', 'datetime.now', ([], {}), '()\n', (4780, 4782), False, 'from datetime import datetime\n')]