Merge branch 'harrison/combine-docs-parse' into harrison/map-rerank

docs/modules/chains/combine_docs_examples/qa_with_sources.ipynb

@@ -56,23 +56,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "docsearch = FAISS.from_texts(texts, embeddings, metadatas=[{\"source\": i} for i in range(len(texts))])"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 4,
-   "id": "5286f58f",
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "docsearch = FAISS.from_texts(texts, embeddings, metadatas=[{\"source\": i} for i in range(len(texts))])\n",
+    "\n",
     "query = \"What did the president say about Justice Breyer\"\n",
     "docs = docsearch.similarity_search(query)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "005a47e9",
    "metadata": {},
    "outputs": [],
@@ -93,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "fc1a5ed6",
    "metadata": {},
    "outputs": [],
@@ -103,24 +95,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 6,
    "id": "7d766417",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'output_text': ' The president thanked Justice Breyer for his service.\\nSOURCES: 30-pl'}"
+       "{'answer': ' The president thanked Justice Breyer for his service.',\n",
+       " 'sources': '30-pl'}"
       ]
      },
-     "execution_count": 7,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "query = \"What did the president say about Justice Breyer\"\n",
-    "chain({\"input_documents\": docs, \"question\": query}, return_only_outputs=True)"
+    "chain.combine_and_parse(**{\"docs\": docs, \"question\": query})"
    ]
   },
   {
@@ -135,7 +128,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 7,
    "id": "921db0a4",
    "metadata": {},
    "outputs": [],
@@ -145,24 +138,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 8,
    "id": "e417926a",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'output_text': ' The president thanked Justice Breyer for his service.\\nSOURCES: 30-pl'}"
+       "{'answer': ' The president thanked Justice Breyer for his service.',\n",
+       " 'sources': '30-pl'}"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 8,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "query = \"What did the president say about Justice Breyer\"\n",
-    "chain({\"input_documents\": docs, \"question\": query}, return_only_outputs=True)"
+    "chain.combine_and_parse(**{\"docs\": docs, \"question\": query})"
    ]
   },
   {
@@ -177,7 +171,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "15af265f",
    "metadata": {},
    "outputs": [],
@@ -187,7 +181,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "21b136e5",
    "metadata": {},
    "outputs": [
@@ -201,7 +195,7 @@
        " 'output_text': ' The president thanked Justice Breyer for his service.\\nSOURCES: 30-pl'}"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -222,7 +216,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "id": "904835c8",
    "metadata": {},
    "outputs": [],
@@ -232,24 +226,25 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 12,
    "id": "f60875c6",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'output_text': \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country and that he was a retiring Justice of the United States Supreme Court. He also thanked him for his service and praised his career as a top litigator in private practice, a former federal public defender, and a family of public school educators and police officers. He noted Justice Breyer's reputation as a consensus builder and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of securing the border and fixing the immigration system in order to advance liberty and justice, and mentioned the new technology, joint patrols, dedicated immigration judges, and commitments to support partners in South and Central America that have been put in place. He also expressed his commitment to the LGBTQ+ community, noting the need for the bipartisan Equality Act and the importance of protecting transgender Americans from state laws targeting them. He also highlighted his commitment to bipartisanship, noting the 80 bipartisan bills he signed into law last year, and his plans to strengthen the Violence Against Women Act. Additionally, he announced that the Justice Department will name a chief prosecutor for pandemic fraud and his plan to lower the deficit by more than one trillion dollars in a\"}"
+       "{'answer': \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country and that he was a retiring Justice of the United States Supreme Court. He also thanked him for his service and praised his career as a top litigator in private practice, a former federal public defender, and a family of public school educators and police officers. He noted Justice Breyer's reputation as a consensus builder and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of securing the border and fixing the immigration system in order to advance liberty and justice, and mentioned the new technology, joint patrols, dedicated immigration judges, and commitments to support partners in South and Central America that have been put in place. He also expressed his commitment to the LGBTQ+ community, noting the need for the bipartisan Equality Act and the importance of protecting transgender Americans from state laws targeting them. He also highlighted his commitment to bipartisanship, noting the 80 bipartisan bills he signed into law last year, and his plans to strengthen the Violence Against Women Act. Additionally, he announced that the Justice Department will name a chief prosecutor for pandemic fraud and his plan to lower the deficit by more than one trillion dollars in a\",\n",
+       " 'sources': ''}"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 12,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
     "query = \"What did the president say about Justice Breyer\"\n",
-    "chain({\"input_documents\": docs, \"question\": query}, return_only_outputs=True)"
+    "chain.combine_and_parse(**{\"docs\": docs, \"question\": query})"
    ]
   },
   {
@@ -264,7 +259,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 13,
    "id": "3396a773",
    "metadata": {},
    "outputs": [],
@@ -274,21 +269,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 14,
    "id": "be5739ef",
    "metadata": {},
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "{'refine_steps': ['\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country and that he was a retiring Justice of the United States Supreme Court.',\n",
-       "  \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his career as a top litigator in private practice, a former federal public defender, and his family of public school educators and police officers. He also noted Justice Breyer's consensus-building skills and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of advancing liberty and justice by securing the border and fixing the immigration system, noting the new technology and joint patrols with Mexico and Guatemala to catch more human traffickers, as well as the dedicated immigration judges and commitments to support partners in South and Central America to host more refugees and secure their own borders. \\nSource: 31\",\n",
-       "  \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his career as a top litigator in private practice, a former federal public defender, and his family of public school educators and police officers. He also noted Justice Breyer's consensus-building skills and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of advancing liberty and justice by securing the border and fixing the immigration system, noting the new technology and joint patrols with Mexico and Guatemala to catch more human traffickers, as well as the dedicated immigration judges and commitments to support partners in South and Central America to host more refugees and secure their own borders. Additionally, he mentioned the need for the bipartisan Equality Act to be passed and signed into law, and the importance of strengthening the Violence Against Women Act. He also offered a Unity Agenda for the Nation, which includes beating the opioid epidemic. \\nSource: 31, 33\",\n",
-       "  \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his career as a top litigator in private practice, a former federal public defender, and his family of public school educators and police officers. He also noted Justice Breyer's consensus-building skills and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of advancing liberty and justice by securing the border and fixing the immigration system, noting the new technology and joint patrols with Mexico and Guatemala to catch more human traffickers, as well as the dedicated immigration judges and commitments to support partners in South and Central America to host more refugees and secure their own borders. Additionally, he mentioned the need for the bipartisan Equality Act to be passed and signed into law, and the importance of strengthening the Violence Against Women Act. He also offered a Unity Agenda for the Nation, which includes beating the opioid epidemic, and announced that the Justice Department will name a chief prosecutor for pandemic fraud. Source: 31, 33, 20\"],\n",
-       " 'output_text': \"\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his career as a top litigator in private practice, a former federal public defender, and his family of public school educators and police officers. He also noted Justice Breyer's consensus-building skills and the broad range of support he has received from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. He also highlighted the importance of advancing liberty and justice by securing the border and fixing the immigration system, noting the new technology and joint patrols with Mexico and Guatemala to catch more human traffickers, as well as the dedicated immigration judges and commitments to support partners in South and Central America to host more refugees and secure their own borders. Additionally, he mentioned the need for the bipartisan Equality Act to be passed and signed into law, and the importance of strengthening the Violence Against Women Act. He also offered a Unity Agenda for the Nation, which includes beating the opioid epidemic, and announced that the Justice Department will name a chief prosecutor for pandemic fraud. Source: 31, 33, 20\"}"
+       "{'refine_steps': ['\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country and that he was a veteran, Constitutional scholar, and retiring Justice of the United States Supreme Court.',\n",
+       "  '\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his experience as a veteran, Constitutional scholar, former top litigator in private practice, former federal public defender, and retiring Justice of the United States Supreme Court. He also noted that Justice Breyer came from a family of public school educators and police officers, and that he was a consensus builder who had received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. The president also highlighted his commitment to advancing liberty and justice by securing the border and fixing the immigration system, including installing new technology, setting up joint patrols with Mexico and Guatemala, putting in place dedicated immigration judges, and securing commitments and supporting partners in South and Central America. Source: 31',\n",
+       "  '\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his experience as a veteran, Constitutional scholar, former top litigator in private practice, former federal public defender, and retiring Justice of the United States Supreme Court. He also noted that Justice Breyer came from a family of public school educators and police officers, and that he was a consensus builder who had received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. The president also highlighted his commitment to advancing liberty and justice by securing the border and fixing the immigration system, including installing new technology, setting up joint patrols with Mexico and Guatemala, putting in place dedicated immigration judges, and securing commitments and supporting partners in South and Central America. He also expressed his commitment to the LGBTQ+ community by offering a Unity Agenda for the Nation, which includes beating the opioid epidemic, and passing the bipartisan Equality Act. Source: 31, 33',\n",
+       "  '\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his experience as a veteran, Constitutional scholar, former top litigator in private practice, former federal public defender, and retiring Justice of the United States Supreme Court. He also noted that Justice Breyer came from a family of public school educators and police officers, and that he was a consensus builder who had received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. The president also highlighted his commitment to advancing liberty and justice by securing the border and fixing the immigration system, including installing new technology, setting up joint patrols with Mexico and Guatemala, putting in place dedicated immigration judges, and securing commitments and supporting partners in South and Central America. He also expressed his commitment to the LGBTQ+ community by offering a Unity Agenda for the Nation, which includes beating the opioid epidemic, and passing the bipartisan Equality Act. Additionally, the president mentioned his plan to lower costs to give families a fair shot, lower the deficit, and go after criminals who stole billions in relief money meant for small businesses and millions of Americans. He also announced that the Justice Department will name a chief prosecutor for pandemic fraud. Source: 20, 31, 33'],\n",
+       " 'output_text': '\\n\\nThe president said that he was honoring Justice Breyer for his dedication to serving the country, his experience as a veteran, Constitutional scholar, former top litigator in private practice, former federal public defender, and retiring Justice of the United States Supreme Court. He also noted that Justice Breyer came from a family of public school educators and police officers, and that he was a consensus builder who had received a broad range of support from the Fraternal Order of Police to former judges appointed by Democrats and Republicans. The president also highlighted his commitment to advancing liberty and justice by securing the border and fixing the immigration system, including installing new technology, setting up joint patrols with Mexico and Guatemala, putting in place dedicated immigration judges, and securing commitments and supporting partners in South and Central America. He also expressed his commitment to the LGBTQ+ community by offering a Unity Agenda for the Nation, which includes beating the opioid epidemic, and passing the bipartisan Equality Act. Additionally, the president mentioned his plan to lower costs to give families a fair shot, lower the deficit, and go after criminals who stole billions in relief money meant for small businesses and millions of Americans. He also announced that the Justice Department will name a chief prosecutor for pandemic fraud. Source: 20, 31, 33'}"
       ]
      },
-     "execution_count": 15,
+     "execution_count": 14,
      "metadata": {},
      "output_type": "execute_result"
     }

docs/modules/chains/combine_docs_examples/vector_db_qa_with_sources.ipynb

@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 2,
    "id": "17d1306e",
    "metadata": {},
    "outputs": [],
@@ -41,7 +41,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 3,
    "id": "0e745d99",
    "metadata": {},
    "outputs": [],
@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 4,
    "id": "f42d79dc",
    "metadata": {},
    "outputs": [],
@@ -63,7 +63,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 5,
    "id": "8aa571ae",
    "metadata": {},
    "outputs": [],
@@ -73,7 +73,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 6,
    "id": "aa859d4c",
    "metadata": {},
    "outputs": [],
@@ -85,7 +85,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 7,
    "id": "8ba36fa7",
    "metadata": {},
    "outputs": [
@@ -96,7 +96,7 @@
        " 'sources': '30-pl'}"
       ]
      },
-     "execution_count": 9,
+     "execution_count": 7,
      "metadata": {},
      "output_type": "execute_result"
     }
@@ -104,11 +104,19 @@
    "source": [
     "chain({\"question\": \"What did the president say about Justice Breyer\"}, return_only_outputs=True)"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9e8ded9a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.9.0 64-bit ('llm-env')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -122,7 +130,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.0"
+   "version": "3.10.9"
   },
   "vscode": {
    "interpreter": {

langchain/chains/combine_documents/base.py

@@ -1,12 +1,13 @@
 """Base interface for chains combining documents."""
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 from pydantic import BaseModel
 
 from langchain.chains.base import Chain
 from langchain.docstore.document import Document
+from langchain.prompts.base import BaseOutputParser
 
 
 class BaseCombineDocumentsChain(Chain, BaseModel, ABC):
@@ -42,6 +43,21 @@ class BaseCombineDocumentsChain(Chain, BaseModel, ABC):
     def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
         """Combine documents into a single string."""
 
+    @property
+    @abstractmethod
+    def output_parser(self) -> Optional[BaseOutputParser]:
+        """Output parser to use for results of combine_docs."""
+
+    def combine_and_parse(
+        self, docs: List[Document], **kwargs: Any
+    ) -> Union[str, List[str], Dict[str, str]]:
+        """Combine documents and parse the result."""
+        result, _ = self.combine_docs(docs, **kwargs)
+        if self.output_parser is not None:
+            return self.output_parser.parse(result)
+        else:
+            return result
+
     def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
         docs = inputs[self.input_key]
         # Other keys are assumed to be needed for LLM prediction

langchain/chains/combine_documents/map_reduce.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel, Extra, root_validator
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.docstore.document import Document
+from langchain.prompts.base import BaseOutputParser
 
 
 def _split_list_of_docs(
@@ -113,6 +114,11 @@ class MapReduceDocumentsChain(BaseCombineDocumentsChain, BaseModel):
         else:
             return self.combine_document_chain
 
+    @property
+    def output_parser(self) -> Optional[BaseOutputParser]:
+        """Output parser to use for results of combine_docs."""
+        return self.combine_document_chain.output_parser
+
     def combine_docs(
         self, docs: List[Document], token_max: int = 3000, **kwargs: Any
     ) -> Tuple[str, dict]:

langchain/chains/combine_documents/map_rerank.py

@@ -0,0 +1,70 @@
+"""Combining documents by mapping a chain over them first, then reranking results."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Dict, List, Optional
+
+from pydantic import BaseModel, Extra, root_validator
+
+from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
+from langchain.chains.llm import LLMChain
+from langchain.docstore.document import Document
+
+
+
+class MapRerankDocumentsChain(BaseCombineDocumentsChain, BaseModel):
+    """Combining documents by mapping a chain over them, then reranking results."""
+
+    llm_chain: LLMChain
+    """Chain to apply to each document individually."""
+    document_variable_name: str
+    """The variable name in the llm_chain to put the documents in.
+    If only one variable in the llm_chain, this need not be provided."""
+    rank_key: str
+    """Key in output of llm_chain to rank on."""
+    metadata_keys: Optional[List[str]] = None
+
+    class Config:
+        """Configuration for this pydantic object."""
+
+        extra = Extra.forbid
+        arbitrary_types_allowed = True
+
+    @root_validator(pre=True)
+    def get_default_document_variable_name(cls, values: Dict) -> Dict:
+        """Get default document variable name, if not provided."""
+        if "document_variable_name" not in values:
+            llm_chain_variables = values["llm_chain"].prompt.input_variables
+            if len(llm_chain_variables) == 1:
+                values["document_variable_name"] = llm_chain_variables[0]
+            else:
+                raise ValueError(
+                    "document_variable_name must be provided if there are "
+                    "multiple llm_chain input_variables"
+                )
+        else:
+            llm_chain_variables = values["llm_chain"].prompt.input_variables
+            if values["document_variable_name"] not in llm_chain_variables:
+                raise ValueError(
+                    f"document_variable_name {values['document_variable_name']} was "
+                    f"not found in llm_chain input_variables: {llm_chain_variables}"
+                )
+        return values
+
+    def combine_docs(
+        self, docs: List[Document], **kwargs: Any
+    ) -> str:
+        """Combine documents in a map rerank manner.
+
+        Combine by mapping first chain over all documents, then reranking the results.
+        """
+        results = self.llm_chain.apply_and_parse(
+            # FYI - this is parallelized and so it is fast.
+            [{**{self.document_variable_name: d.page_content}, **kwargs} for d in docs]
+        )
+        sorted_res = sorted(zip(results, docs), key=lambda x: -x[0][self.rank_key])
+        output, document = sorted_res[0]
+        if self.metadata_keys is not None:
+            for key in self.metadata_keys:
+                output[key] = document.metadata[key]
+        return output

langchain/chains/combine_documents/refine.py

@@ -2,14 +2,14 @@
 
 from __future__ import annotations
 
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 
 from pydantic import BaseModel, Extra, Field, root_validator
 
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.docstore.document import Document
-from langchain.prompts.base import BasePromptTemplate
+from langchain.prompts.base import BaseOutputParser, BasePromptTemplate
 from langchain.prompts.prompt import PromptTemplate
 
 
@@ -74,6 +74,11 @@ class RefineDocumentsChain(BaseCombineDocumentsChain, BaseModel):
                 )
         return values
 
+    @property
+    def output_parser(self) -> Optional[BaseOutputParser]:
+        """Output parser to use for results of combine_docs."""
+        return self.refine_llm_chain.prompt.output_parser
+
     def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
         """Combine by mapping first chain over all, then stuffing into final chain."""
         base_info = {"page_content": docs[0].page_content}

langchain/chains/combine_documents/stuff.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel, Extra, Field, root_validator
 from langchain.chains.combine_documents.base import BaseCombineDocumentsChain
 from langchain.chains.llm import LLMChain
 from langchain.docstore.document import Document
-from langchain.prompts.base import BasePromptTemplate
+from langchain.prompts.base import BaseOutputParser, BasePromptTemplate
 from langchain.prompts.prompt import PromptTemplate
 
 
@@ -78,6 +78,11 @@ class StuffDocumentsChain(BaseCombineDocumentsChain, BaseModel):
         prompt = self.llm_chain.prompt.format(**inputs)
         return self.llm_chain.llm.get_num_tokens(prompt)
 
+    @property
+    def output_parser(self) -> Optional[BaseOutputParser]:
+        """Output parser to use for results of combine_docs."""
+        return self.llm_chain.prompt.output_parser
+
     def combine_docs(self, docs: List[Document], **kwargs: Any) -> Tuple[str, dict]:
         """Stuff all documents into one prompt and pass to LLM."""
         inputs = self._get_inputs(docs, **kwargs)

langchain/chains/qa_with_sources/base.py

@@ -19,7 +19,7 @@ from langchain.chains.qa_with_sources.map_reduce_prompt import (
 )
 from langchain.docstore.document import Document
 from langchain.llms.base import BaseLLM
-from langchain.prompts.base import BasePromptTemplate
+from langchain.prompts.base import BasePromptTemplate, RegexParser
 
 
 class BaseQAWithSourcesChain(Chain, BaseModel, ABC):
@@ -29,8 +29,6 @@ class BaseQAWithSourcesChain(Chain, BaseModel, ABC):
     """Chain to use to combine documents."""
     question_key: str = "question"  #: :meta private:
     input_docs_key: str = "docs"  #: :meta private:
-    answer_key: str = "answer"  #: :meta private:
-    sources_answer_key: str = "sources"  #: :meta private:
 
     @classmethod
     def from_llm(
@@ -79,7 +77,13 @@ class BaseQAWithSourcesChain(Chain, BaseModel, ABC):
 
         :meta private:
         """
-        return [self.answer_key, self.sources_answer_key]
+        output_parser = self.combine_document_chain.output_parser
+        if not isinstance(output_parser, RegexParser):
+            raise ValueError(
+                "Output parser of combine_document_chain should be a RegexParser,"
+                f" got {output_parser}"
+            )
+        return output_parser.output_keys
 
     @root_validator(pre=True)
     def validate_question_chain(cls, values: Dict) -> Dict:
@@ -94,9 +98,14 @@ class BaseQAWithSourcesChain(Chain, BaseModel, ABC):
         return values
 
     @root_validator()
-    def validate_combine_chain_can_be_constructed(cls, values: Dict) -> Dict:
-        """Validate that the combine chain can be constructed."""
-        # Try to construct the combine documents chains.
+    def validate_combine_chain_output(cls, values: Dict) -> Dict:
+        """Validate that the combine chain outputs a dictionary."""
+        combine_docs_chain = values["combine_document_chain"]
+        if not isinstance(combine_docs_chain.output_parser, RegexParser):
+            raise ValueError(
+                "Output parser of combine_document_chain should be a RegexParser,"
+                f" got {combine_docs_chain.output_parser}"
+            )
 
         return values
 
@@ -106,12 +115,8 @@ class BaseQAWithSourcesChain(Chain, BaseModel, ABC):
 
     def _call(self, inputs: Dict[str, Any]) -> Dict[str, str]:
         docs = self._get_docs(inputs)
-        answer, _ = self.combine_document_chain.combine_docs(docs, **inputs)
-        if "\nSOURCES: " in answer:
-            answer, sources = answer.split("\nSOURCES: ")
-        else:
-            sources = ""
-        return {self.answer_key: answer, self.sources_answer_key: sources}
+        answer = self.combine_document_chain.combine_and_parse(docs, **inputs)
+        return answer
 
 
 class QAWithSourcesChain(BaseQAWithSourcesChain, BaseModel):

langchain/chains/qa_with_sources/map_reduce_prompt.py

@@ -1,5 +1,12 @@
 # flake8: noqa
 from langchain.prompts import PromptTemplate
+from langchain.prompts.base import RegexParser
+
+output_parser = RegexParser(
+    regex=r"(.*?)\nSOURCES: (.*)",
+    output_keys=["answer", "sources"],
+    default_output_key="answer",
+)
 
 question_prompt_template = """Use the following portion of a long document to see if any of the text is relevant to answer the question. 
 Return any relevant text verbatim.
@@ -46,7 +53,9 @@ QUESTION: {question}
 =========
 FINAL ANSWER:"""
 COMBINE_PROMPT = PromptTemplate(
-    template=combine_prompt_template, input_variables=["summaries", "question"]
+    template=combine_prompt_template,
+    input_variables=["summaries", "question"],
+    output_parser=output_parser,
 )
 
 EXAMPLE_PROMPT = PromptTemplate(

langchain/chains/qa_with_sources/refine_prompts.py

@@ -1,5 +1,12 @@
 # flake8: noqa
 from langchain.prompts import PromptTemplate
+from langchain.prompts.base import RegexParser
+
+output_parser = RegexParser(
+    regex=r"(.*?)\nSOURCES: (.*)",
+    output_keys=["answer", "sources"],
+    default_output_key="answer",
+)
 
 DEFAULT_REFINE_PROMPT_TMPL = (
     "The original question is as follows: {question}\n"
@@ -17,6 +24,7 @@ DEFAULT_REFINE_PROMPT_TMPL = (
 DEFAULT_REFINE_PROMPT = PromptTemplate(
     input_variables=["question", "existing_answer", "context_str"],
     template=DEFAULT_REFINE_PROMPT_TMPL,
+    output_parser=output_parser,
 )
 
 

langchain/chains/qa_with_sources/stuff_prompt.py

@@ -1,5 +1,12 @@
 # flake8: noqa
 from langchain.prompts import PromptTemplate
+from langchain.prompts.base import RegexParser
+
+output_parser = RegexParser(
+    regex=r"(.*?)\nSOURCES: (.*)",
+    output_keys=["answer", "sources"],
+    default_output_key="answer",
+)
 
 template = """Given the following extracted parts of a long document and a question, create a final answer with references ("SOURCES"). 
 If you don't know the answer, just say that you don't know. Don't try to make up an answer.
@@ -36,7 +43,11 @@ QUESTION: {question}
 {summaries}
 =========
 FINAL ANSWER:"""
-PROMPT = PromptTemplate(template=template, input_variables=["summaries", "question"])
+PROMPT = PromptTemplate(
+    template=template,
+    input_variables=["summaries", "question"],
+    output_parser=output_parser,
+)
 
 EXAMPLE_PROMPT = PromptTemplate(
     template="Content: {page_content}\nSource: {source}",

langchain/chains/question_answering/map_rerank_prompt.py

@@ -0,0 +1,42 @@
+from langchain.prompts import PromptTemplate
+from langchain.prompts.base import BaseOutputParser
+import re
+
+template = """Use the following document to answer the given question. In addition to providing an answer, please also give your answer a score from 0-100 in terms of how good it is (higher is better). 
+
+What decides the score? A good score is factually accurate, and FULLY answers the question in a way the user would find helpful. If the document does not contain the answer, the score should be 0. You should only give a score of 100 if you are absolutely positive this is the best answer. Keep in mind that you will also be answering this question with other documents, so one of them could have a better answer.
+
+Use the following format:
+
+Document:
+---------------
+Document text here
+---------------
+Question: Question here
+Answer: Answer here
+Score: Score (between 0 and 100) here
+
+Begin!
+
+Document:
+---------------
+{context}
+---------------
+Question: {question}
+Answer:"""
+
+
+class ScoreOutputParser(BaseOutputParser):
+
+    def parse(self, text: str):
+        regex = r"(.*?)\nScore: (.*)"
+        match = re.search(regex, text)
+        if match:
+            question = match.group(1)
+            answer = match.group(2)
+            return {"answer": question, "score": int(answer)}
+        else:
+            raise ValueError(f"Could not parse output: {text}")
+
+
+prompt = PromptTemplate(template=template, input_variables=['context', 'question'], output_parser=ScoreOutputParser())