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"])
+ )