get_rag_qa_gen_chain

langchain-streamlit-demo/app.py

@@ -354,7 +354,8 @@ if st.session_state.llm:
             st.session_state.doc_chain = "summarization"
         elif document_chat_chain_type == "Q&A Generation":
             st.session_state.doc_chain = get_qa_gen_chain(st.session_state.llm)
-
+            # from qagen import get_rag_qa_gen_chain
+            # st.session_state.doc_chain = get_rag_qa_gen_chain(st.session_state.retriever, st.session_state.llm)
         else:
             st.session_state.doc_chain = RetrievalQA.from_chain_type(
                 llm=st.session_state.llm,
@@ -432,6 +433,7 @@ if st.session_state.llm:
                             ],
                             config,
                         )
+                        # raw_results = st.session_state.doc_chain.invoke(prompt, config)
                         results = combine_qa_pair_lists(raw_results).QuestionAnswerPairs
 
                         def _to_str(idx, qap):

langchain-streamlit-demo/qagen.py

@@ -6,7 +6,8 @@ from langchain.prompts.chat import (
     ChatPromptTemplate,
 )
 from langchain.schema.language_model import BaseLanguageModel
-from langchain.schema.runnable import RunnableSequence
+from langchain.schema.retriever import BaseRetriever
+from langchain.schema.runnable import RunnablePassthrough, RunnableSequence
 from pydantic import BaseModel, Field
 
 
@@ -67,3 +68,16 @@ def get_qa_gen_chain(llm: BaseLanguageModel) -> RunnableSequence:
     return (
         CHAT_PROMPT | llm | OutputFixingParser.from_llm(llm=llm, parser=PYDANTIC_PARSER)
     )
+
+
+def get_rag_qa_gen_chain(
+    retriever: BaseRetriever,
+    llm: BaseLanguageModel,
+    input_key: str = "prompt",
+) -> RunnableSequence:
+    return (
+        {"context": retriever, input_key: RunnablePassthrough()}
+        | CHAT_PROMPT
+        | llm
+        | OutputFixingParser.from_llm(llm=llm, parser=PYDANTIC_PARSER)
+    )