refactor

langchain-streamlit-demo/app.py

@@ -409,6 +409,13 @@ if st.session_state.llm:
             if st.session_state.ls_tracer:
                 callbacks.append(st.session_state.ls_tracer)
 
+            config: Dict[str, Any] = dict(
+                callbacks=callbacks,
+                tags=["Streamlit Chat"],
+            )
+            if st.session_state.provider == "Anthropic":
+                config["max_concurrency"] = 5
+
             use_document_chat = all(
                 [
                     document_chat,
@@ -420,46 +427,18 @@ if st.session_state.llm:
             try:
                 full_response: Union[str, None]
                 if use_document_chat:
-                    if document_chat_chain_type == "Summarization":
-                        st.session_state.doc_chain = get_rag_summarization_chain(
-                            prompt,
-                            st.session_state.retriever,
-                            st.session_state.llm,
-                        )
+                    if document_chat_chain_type in ("Summarization", "Q&A Generation"):
+                        if document_chat_chain_type == "Summarization":
+                            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.invoke(
                             prompt,
-                            dict(
-                                callbacks=callbacks,
-                                tags=["Streamlit Chat"],
-                            ),
-                        )
-
-                        st.markdown(full_response)
-                    elif document_chat_chain_type == "Q&A Generation":
-                        config: Dict[str, Any] = dict(
-                            callbacks=callbacks,
-                            tags=["Streamlit Chat"],
-                        )
-                        if st.session_state.provider == "Anthropic":
-                            config["max_concurrency"] = 5
-                        raw_results = st.session_state.doc_chain.invoke(prompt, config)
-                        results = raw_results.QuestionAnswerPairs
-
-                        def _to_str(idx, qap):
-                            question_piece = f"{idx}. **Q:** {qap.question}"
-                            whitespace = " " * (len(str(idx)) + 2)
-                            answer_piece = f"{whitespace}**A:** {qap.answer}"
-                            return f"{question_piece}\n\n{answer_piece}"
-
-                        full_response = "\n\n".join(
-                            [
-                                _to_str(idx, qap)
-                                for idx, qap in enumerate(results, start=1)
-                            ],
+                            config,
                         )
 
-                        st.markdown(full_response)
-
                     else:
                         st_handler = StreamlitCallbackHandler(st.container())
                         callbacks.append(st_handler)
@@ -470,7 +449,9 @@ if st.session_state.llm:
                             return_only_outputs=True,
                         )[st.session_state.doc_chain.output_key]
                         st_handler._complete_current_thought()
-                        st.markdown(full_response)
+
+                    st.markdown(full_response)
+
                 else:
                     message_placeholder = st.empty()
                     stream_handler = StreamHandler(message_placeholder)

langchain-streamlit-demo/qagen.py

@@ -15,10 +15,24 @@ class QuestionAnswerPair(BaseModel):
     question: str = Field(..., description="The question that will be answered.")
     answer: str = Field(..., description="The answer to the question that was asked.")
 
+    def to_str(self, idx: int) -> str:
+        question_piece = f"{idx}. **Q:** {self.question}"
+        whitespace = " " * (len(str(idx)) + 2)
+        answer_piece = f"{whitespace}**A:** {self.answer}"
+        return f"{question_piece}\n\n{answer_piece}"
+
 
 class QuestionAnswerPairList(BaseModel):
     QuestionAnswerPairs: List[QuestionAnswerPair]
 
+    def to_str(self) -> str:
+        return "\n\n".join(
+            [
+                qap.to_str(idx)
+                for idx, qap in enumerate(self.QuestionAnswerPairs, start=1)
+            ],
+        )
+
 
 PYDANTIC_PARSER: PydanticOutputParser = PydanticOutputParser(
     pydantic_object=QuestionAnswerPairList,
@@ -74,4 +88,5 @@ def get_rag_qa_gen_chain(
         | CHAT_PROMPT
         | llm
         | OutputFixingParser.from_llm(llm=llm, parser=PYDANTIC_PARSER)
+        | (lambda parsed_output: combine_qa_pair_lists(parsed_output).to_str())
     )