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())
  )