rag summarization

langchain-streamlit-demo/app.py

@@ -27,7 +27,7 @@ from langsmith.client import Client
 from streamlit_feedback import streamlit_feedback
 
 from qagen import get_rag_qa_gen_chain
-from summarize import get_summarization_chain
+from summarize import get_rag_summarization_chain
 
 __version__ = "0.0.10"
 
@@ -421,14 +421,17 @@ if st.session_state.llm:
                 full_response: Union[str, None]
                 if use_document_chat:
                     if document_chat_chain_type == "Summarization":
-                        st.session_state.doc_chain = get_summarization_chain(
-                            st.session_state.llm,
+                        st.session_state.doc_chain = get_rag_summarization_chain(
                             prompt,
+                            st.session_state.retriever,
+                            st.session_state.llm,
                         )
-                        full_response = st.session_state.doc_chain.run(
-                            st.session_state.texts,
-                            callbacks=callbacks,
-                            tags=["Streamlit Chat"],
+                        full_response = st.session_state.doc_chain.invoke(
+                            prompt,
+                            dict(
+                                callbacks=callbacks,
+                                tags=["Streamlit Chat"],
+                            ),
                         )
 
                         st.markdown(full_response)

langchain-streamlit-demo/summarize.py

@@ -2,6 +2,8 @@ from langchain.chains.base import Chain
 from langchain.chains.summarize import load_summarize_chain
 from langchain.prompts import PromptTemplate
 from langchain.schema.language_model import BaseLanguageModel
+from langchain.schema.retriever import BaseRetriever
+from langchain.schema.runnable import RunnableSequence, RunnablePassthrough
 
 prompt_template = """Write a concise summary of the following text, based on the user input.
 User input: {query}
@@ -49,3 +51,16 @@ def get_summarization_chain(
         input_key="input_documents",
         output_key="output_text",
     )
+
+
+def get_rag_summarization_chain(
+    prompt: str,
+    retriever: BaseRetriever,
+    llm: BaseLanguageModel,
+    input_key: str = "prompt",
+) -> RunnableSequence:
+    return (
+        {"input_documents": retriever, input_key: RunnablePassthrough()}
+        | get_summarization_chain(llm, prompt)
+        | (lambda output: output["output_text"])
+    )