Updated to new callback syntax

* Fixed callback bugs

* Bumbped version

* Switched to new rmrkl version

Author: whitead

paperqa/__init__.py

@@ -1,3 +1,3 @@
-from .docs import Docs, maybe_is_text, Answer
-from .version import __version__
 from .agent import run_agent
+from .docs import Answer, Docs, maybe_is_text
+from .version import __version__

paperqa/agent.py

@@ -1,12 +1,12 @@
-from langchain.tools import BaseTool
-from .docs import Answer, Docs
-from langchain.agents import initialize_agent
-from langchain.chat_models import ChatOpenAI
+from langchain.agents import AgentType, initialize_agent
 from langchain.chains import LLMChain
-from langchain.agents import AgentType
-from .qaprompts import select_paper_prompt, make_chain
+from langchain.chat_models import ChatOpenAI
+from langchain.tools import BaseTool
 from rmrkl import ChatZeroShotAgent, RetryAgentExecutor
 
+from .docs import Answer, Docs
+from .qaprompts import make_chain, select_paper_prompt
+
 
 def status(answer: Answer, docs: Docs):
     return f" Status: Current Papers: {len(docs.doc_previews())} Current Evidence: {len(answer.contexts)} Current Cost: ${answer.cost:.2f}"
@@ -91,13 +91,12 @@ def __init__(self, docs, answer):
         self.answer = answer
 
     def _run(self, query: str) -> str:
-        self.answer = self.docs.query(
-            query, answer=self.answer
-        )
+        self.answer = self.docs.query(query, answer=self.answer)
         if "cannot answer" in self.answer.answer:
             self.answer = Answer(self.answer.question)
-            return "Failed to answer question. Deleting evidence. Consider rephrasing question or evidence statement." + status(
-                self.answer, self.docs
+            return (
+                "Failed to answer question. Deleting evidence. Consider rephrasing question or evidence statement."
+                + status(self.answer, self.docs)
             )
         return self.answer.answer + status(self.answer, self.docs)
 
@@ -108,7 +107,9 @@ def _arun(self, query: str) -> str:
 
 class Search(BaseTool):
     name = "Paper Search"
-    description = "Search for papers to add to cur. Input should be a string of keywords."
+    description = (
+        "Search for papers to add to cur. Input should be a string of keywords."
+    )
     docs: Docs = None
     answer: Answer = None
 
@@ -143,7 +144,6 @@ def _arun(self, query: str) -> str:
 
 
 def make_tools(docs, answer):
-
     tools = []
 
     tools.append(Search(docs, answer))

paperqa/contrib/zotero.py

@@ -1,17 +1,17 @@
 # This file gets PDF files from the user's Zotero library
-import os
-from typing import Union, Optional
-from pathlib import Path
 import logging
+import os
 from collections import namedtuple
+from pathlib import Path
+from typing import Optional, Union
 
 try:
     from pyzotero import zotero
 except ImportError:
     raise ImportError("Please install pyzotero: `pip install pyzotero`")
 from ..paths import CACHE_PATH
-from ..utils import count_pdf_pages
 from ..types import StrPath
+from ..utils import count_pdf_pages
 
 ZoteroPaper = namedtuple(
     "ZoteroPaper", ["key", "title", "pdf", "num_pages", "zotero_key", "details"]

paperqa/docs.py

@@ -1,33 +1,30 @@
-from typing import List, Optional, Tuple, Union, Callable
-from functools import reduce
+import asyncio
 import os
+import re
 import sys
-import asyncio
+from datetime import datetime
+from functools import reduce
 from pathlib import Path
-import re
-from .paths import CACHE_PATH
-from .utils import maybe_is_text, md5sum
-from .qaprompts import (
-    summary_prompt,
-    qa_prompt,
-    search_prompt,
-    citation_prompt,
-    select_paper_prompt,
-    make_chain,
-)
-from .types import Answer, Context
-from .readers import read_doc
-from langchain.vectorstores import FAISS
+from typing import Callable, List, Optional, Tuple, Union
+
+import langchain
+from langchain.cache import SQLiteCache
+from langchain.callbacks import OpenAICallbackHandler, get_openai_callback
+from langchain.callbacks.base import AsyncCallbackHandler
+from langchain.callbacks.manager import AsyncCallbackManager
+from langchain.chat_models import ChatOpenAI
 from langchain.docstore.document import Document
-from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.embeddings.base import Embeddings
-from langchain.chat_models import ChatOpenAI
+from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.llms.base import LLM
-from langchain.callbacks import get_openai_callback, OpenAICallbackHandler
-from langchain.callbacks.base import AsyncCallbackHandler, AsyncCallbackManager
-from langchain.cache import SQLiteCache
-import langchain
-from datetime import datetime
+from langchain.vectorstores import FAISS
+
+from .paths import CACHE_PATH
+from .qaprompts import (citation_prompt, make_chain, qa_prompt, search_prompt,
+                        select_paper_prompt, summary_prompt)
+from .readers import read_doc
+from .types import Answer, Context
+from .utils import maybe_is_text, md5sum
 
 os.makedirs(os.path.dirname(CACHE_PATH), exist_ok=True)
 langchain.llm_cache = SQLiteCache(CACHE_PATH)
@@ -44,7 +41,6 @@ def __init__(
         name: str = "default",
         index_path: Optional[Path] = None,
         embeddings: Optional[Embeddings] = None,
-        get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x : []
     ) -> None:
         """Initialize the collection of documents.
 
@@ -55,7 +51,6 @@ def __init__(
             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
             embeddings: The embeddings to use for indexing documents. Default - OpenAI embeddings
-            get_callbacks: A function that allows callbacks to built per stage of the pipeline.
         """
         self.docs = dict()
         self.chunk_size_limit = chunk_size_limit
@@ -70,7 +65,6 @@ def __init__(
         if embeddings is None:
             embeddings = OpenAIEmbeddings()
         self.embeddings = embeddings
-        self.get_callbacks = get_callbacks
 
     def update_llm(
         self,
@@ -96,7 +90,6 @@ def add(
         key: Optional[str] = None,
         disable_check: bool = False,
         chunk_chars: Optional[int] = 3000,
-        overwrite: bool = False,
     ) -> None:
         """Add a document to the collection."""
 
@@ -110,8 +103,9 @@ def add(
             cite_chain = make_chain(prompt=citation_prompt, llm=self.summary_llm)
             # peak first chunk
             texts, _ = read_doc(path, "", "", chunk_chars=chunk_chars)
-            with get_openai_callback():
-                citation = cite_chain.run(texts[0])
+            if len(texts) == 0:
+                raise ValueError(f"Could not read document {path}. Is it empty?")
+            citation = 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}"
 
@@ -207,12 +201,10 @@ def __getstate__(self):
             state["_faiss_index"].save_local(self.index_path)
         del state["_faiss_index"]
         del state["_doc_index"]
-        del state["get_callbacks"]
         return state
 
     def __setstate__(self, state):
         self.__dict__.update(state)
-        self.get_callbacks = lambda x: []
         try:
             self._faiss_index = FAISS.load_local(self.index_path, self.embeddings)
         except:
@@ -240,8 +232,9 @@ def get_evidence(
         k: int = 3,
         max_sources: int = 5,
         marginal_relevance: bool = True,
-        key_filter: Optional[List[str]] = None    
-        ) -> Answer:
+        key_filter: Optional[List[str]] = None,
+        get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x: [],
+    ) -> Answer:
         # special case for jupyter notebooks
         if "get_ipython" in globals() or "google.colab" in sys.modules:
             import nest_asyncio
@@ -258,7 +251,8 @@ def get_evidence(
                 k=k,
                 max_sources=max_sources,
                 marginal_relevance=marginal_relevance,
-                key_filter=key_filter
+                key_filter=key_filter,
+                get_callbacks=get_callbacks,
             )
         )
 
@@ -269,6 +263,7 @@ async def aget_evidence(
         max_sources: int = 5,
         marginal_relevance: bool = True,
         key_filter: Optional[List[str]] = None,
+        get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x: [],
     ) -> Answer:
         if len(self.docs) == 0:
             return answer
@@ -293,29 +288,32 @@ async def process(doc):
             # check if it is already in answer (possible in agent setting)
             if doc.metadata["key"] in [c.key for c in answer.contexts]:
                 return None, None
-            cb = OpenAICallbackHandler()
-            manager = AsyncCallbackManager([cb] + self.get_callbacks('evidence:' + doc.metadata['key']))
-            summary_chain = make_chain(summary_prompt, self.summary_llm, manager)
+            callbacks = [OpenAICallbackHandler()] + get_callbacks(
+                "evidence:" + doc.metadata["key"]
+            )
+            summary_chain = make_chain(summary_prompt, self.summary_llm)
             c = Context(
                 key=doc.metadata["key"],
                 citation=doc.metadata["citation"],
                 context=await summary_chain.arun(
                     question=answer.question,
                     context_str=doc.page_content,
                     citation=doc.metadata["citation"],
+                    callbacks=callbacks,
                 ),
                 text=doc.page_content,
             )
             if "Not applicable" not in c.context:
                 return c, cb
             return None, None
 
-        results = await asyncio.gather(*[process(doc) for doc in docs])
+        with get_openai_callback() as cb:
+            results = await asyncio.gather(*[process(doc) for doc in docs])
         # filter out failures
         results = [r for r in results if r[0] is not None]
         answer.tokens += sum([cb.total_tokens for _, cb in results])
         answer.cost += sum([cb.total_cost for _, cb in results])
-        contexts = [c for c,_ in results if c is not None]
+        contexts = [c for c, _ in results if c is not None]
         if len(contexts) == 0:
             return answer
         contexts = sorted(contexts, key=lambda x: len(x.context), reverse=True)
@@ -365,6 +363,7 @@ def query(
         marginal_relevance: bool = True,
         answer: Optional[Answer] = None,
         key_filter: Optional[bool] = None,
+        get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x: [],
     ) -> Answer:
         # special case for jupyter notebooks
         if "get_ipython" in globals() or "google.colab" in sys.modules:
@@ -385,6 +384,7 @@ def query(
                 marginal_relevance=marginal_relevance,
                 answer=answer,
                 key_filter=key_filter,
+                get_callbacks=get_callbacks,
             )
         )
 
@@ -397,6 +397,7 @@ async def aquery(
         marginal_relevance: bool = True,
         answer: Optional[Answer] = None,
         key_filter: Optional[bool] = None,
+        get_callbacks: Callable[[str], AsyncCallbackHandler] = lambda x: [],
     ) -> Answer:
         if k < max_sources:
             raise ValueError("k should be greater than max_sources")
@@ -414,6 +415,7 @@ async def aquery(
                 max_sources=max_sources,
                 marginal_relevance=marginal_relevance,
                 key_filter=keys if key_filter else None,
+                get_callbacks=get_callbacks,
             )
         context_str, contexts = answer.context, answer.contexts
         bib = dict()
@@ -424,11 +426,14 @@ async def aquery(
             )
         else:
             cb = OpenAICallbackHandler()
-            manager = AsyncCallbackManager([cb] + self.get_callbacks('answer'))
-            qa_chain = make_chain(qa_prompt, self.llm, manager)
+            callbacks = [OpenAICallbackHandler()] + get_callbacks("answer")
+            qa_chain = make_chain(qa_prompt, self.llm)
             answer_text = await qa_chain.arun(
-                    question=query, context_str=context_str, length=length_prompt
-                )
+                question=query,
+                context_str=context_str,
+                length=length_prompt,
+                callbacks=callbacks,
+            )
             answer.tokens += cb.total_tokens
             answer.cost += cb.total_cost
         # it still happens lol

paperqa/qaprompts.py

@@ -1,9 +1,13 @@
-import langchain.prompts as prompts
 from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+import langchain.prompts as prompts
+from langchain.callbacks.manager import AsyncCallbackManagerForChainRun
 from langchain.chains import LLMChain
 from langchain.chat_models import ChatOpenAI
-from langchain.schema import SystemMessage
-from langchain.prompts.chat import HumanMessagePromptTemplate, ChatPromptTemplate
+from langchain.prompts.chat import (ChatPromptTemplate,
+                                    HumanMessagePromptTemplate)
+from langchain.schema import LLMResult, SystemMessage
 
 summary_prompt = prompts.PromptTemplate(
     input_variables=["question", "context_str", "citation"],
@@ -73,11 +77,23 @@ def _get_datetime():
     partial_variables={"date": _get_datetime},
 )
 
-def make_chain(prompt, llm, callback_manager=None):
-    if callback_manager is not None:
-        # need to clone to attach
-        llm = llm.copy()
-        llm.callback_manager = callback_manager
+
+class FallbackLLMChain(LLMChain):
+    """Chain that falls back to synchronous generation if the async generation fails."""
+
+    async def agenerate(
+        self,
+        input_list: List[Dict[str, Any]],
+        run_manager: Optional[AsyncCallbackManagerForChainRun] = None,
+    ) -> LLMResult:
+        """Generate LLM result from inputs."""
+        try:
+            return await super().agenerate(input_list, run_manager=run_manager)
+        except NotImplementedError as e:
+            return self.generate(input_list, run_manager=run_manager)
+
+
+def make_chain(prompt, llm):
     if type(llm) == ChatOpenAI:
         system_message_prompt = SystemMessage(
             content="You are a scholarly researcher that answers in an unbiased, scholarly tone. "
@@ -87,4 +103,4 @@ def make_chain(prompt, llm, callback_manager=None):
         prompt = ChatPromptTemplate.from_messages(
             [system_message_prompt, human_message_prompt]
         )
-    return LLMChain(prompt=prompt, llm=llm)
+    return FallbackLLMChain(prompt=prompt, llm=llm)

paperqa/readers.py

@@ -1,15 +1,16 @@
-import os
-from .paths import OCR_CACHE_PATH
-from .version import __version__
-from html2text import html2text
-from pathlib import Path
 import json
 import logging
+import os
 from hashlib import md5
+from pathlib import Path
 
-from langchain.text_splitter import TokenTextSplitter
+from html2text import html2text
 from langchain.cache import SQLiteCache
 from langchain.schema import Generation
+from langchain.text_splitter import TokenTextSplitter
+
+from .paths import OCR_CACHE_PATH
+from .version import __version__
 
 OCR_CACHE = None
 
@@ -69,7 +70,6 @@ def parse_pdf(path, citation, key, chunk_chars=2000, overlap=50):
 
 
 def parse_txt(path, citation, key, chunk_chars=2000, overlap=50, html=False):
-
     try:
         with open(path) as f:
             doc = f.read()