refactor

langchain-streamlit-demo/app.py

@@ -7,7 +7,6 @@ import anthropic
 import langsmith.utils
 import openai
 import streamlit as st
-from langchain.callbacks import StreamlitCallbackHandler
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.callbacks.tracers.langchain import LangChainTracer, wait_for_all_tracers
 from langchain.callbacks.tracers.run_collector import RunCollectorCallbackHandler
@@ -360,40 +359,17 @@ for msg in STMEMORY.messages:
 
 # --- Current Chat ---
 if st.session_state.llm:
- # --- Document Chat ---
- if st.session_state.retriever:
- if document_chat_chain_type == "Summarization":
- st.session_state.doc_chain = "summarization"
- elif document_chat_chain_type == "Q&A Generation":
- 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,
- chain_type=document_chat_chain_type,
- retriever=st.session_state.retriever,
- memory=MEMORY,
- )
-
- else:
- # --- Regular Chat ---
- chat_prompt = ChatPromptTemplate.from_messages(
- [
- (
- "system",
- system_prompt + "\nIt's currently {time}.",
- ),
- MessagesPlaceholder(variable_name="chat_history"),
- ("human", "{query}"),
- ],
- ).partial(time=lambda: str(datetime.now()))
- st.session_state.chain = LLMChain(
- prompt=chat_prompt,
- llm=st.session_state.llm,
- memory=MEMORY,
- )
+ # --- Regular Chat ---
+ chat_prompt = ChatPromptTemplate.from_messages(
+ [
+ (
+ "system",
+ system_prompt + "\nIt's currently {time}.",
+ ),
+ MessagesPlaceholder(variable_name="chat_history"),
+ ("human", "{query}"),
+ ],
+ ).partial(time=lambda: str(datetime.now()))
 
  # --- Chat Input ---
  prompt = st.chat_input(placeholder="Ask me a question!")
@@ -419,57 +395,60 @@ if st.session_state.llm:
  use_document_chat = all(
  [
  document_chat,
- st.session_state.doc_chain,
  st.session_state.retriever,
  ],
  )
 
+ full_response = None
+
  try:
- full_response: Union[str, None]
- if use_document_chat:
- 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(
+ if not use_document_chat:
+ message_placeholder = st.empty()
+ stream_handler = StreamHandler(message_placeholder)
+ callbacks.append(stream_handler)
+ st.session_state.chain = LLMChain(
+ prompt=chat_prompt,
+ llm=st.session_state.llm,
+ memory=MEMORY,
+ ) | (lambda output: output["text"])
+ config = {"callbacks": callbacks, "tags": ["Streamlit Chat"]}
+ full_response = st.session_state.chain.invoke(prompt, config)
+ message_placeholder.markdown(full_response)
+
+ else:
+
+ def get_rag_runnable():
+ if document_chat_chain_type == "Q&A Generation":
+ return get_rag_qa_gen_chain(
+ st.session_state.retriever,
+ st.session_state.llm,
+ )
+ elif document_chat_chain_type == "Summarization":
+ return get_rag_summarization_chain(
  prompt,
  st.session_state.retriever,
  st.session_state.llm,
  )
- full_response = st.session_state.doc_chain.invoke(
- prompt,
- config,
- )
-
- else:
- st_handler = StreamlitCallbackHandler(st.container())
- callbacks.append(st_handler)
- full_response = st.session_state.doc_chain(
- {"query": prompt},
- callbacks=callbacks,
- tags=["Streamlit Chat"],
- return_only_outputs=True,
- )[st.session_state.doc_chain.output_key]
- st_handler._complete_current_thought()
-
+ else:
+ return RetrievalQA.from_chain_type(
+ llm=st.session_state.llm,
+ chain_type=document_chat_chain_type,
+ retriever=st.session_state.retriever,
+ memory=MEMORY,
+ output_key="output_text",
+ ) | (lambda output: output["output_text"])
+
+ st.session_state.doc_chain = get_rag_runnable()
+
+ full_response = st.session_state.doc_chain.invoke(prompt, config)
  st.markdown(full_response)
 
- else:
- message_placeholder = st.empty()
- stream_handler = StreamHandler(message_placeholder)
- callbacks.append(stream_handler)
- full_response = st.session_state.chain(
- {"query": prompt},
- callbacks=callbacks,
- tags=["Streamlit Chat"],
- return_only_outputs=True,
- )[st.session_state.chain.output_key]
- message_placeholder.markdown(full_response)
  except (openai.error.AuthenticationError, anthropic.AuthenticationError):
  st.error(
  f"Please enter a valid {st.session_state.provider} API key.",
  icon="❌",
  )
- full_response = None
- if full_response:
+ if full_response is not None:
  # --- Tracing ---
  if st.session_state.client:
  st.session_state.run = RUN_COLLECTOR.traced_runs[0]

langchain-streamlit-demo/qagen.py

@@ -1,4 +1,3 @@
-from functools import reduce
 from typing import List
 
 from langchain.output_parsers import PydanticOutputParser, OutputFixingParser
@@ -59,25 +58,6 @@ CHAT_PROMPT = ChatPromptTemplate.from_messages(
 ).partial(format_instructions=PYDANTIC_PARSER.get_format_instructions)
 
 
-def combine_qa_pair_lists(
- qa_pair_lists: List[QuestionAnswerPairList],
-) -> QuestionAnswerPairList:
- def reducer(
- accumulator: QuestionAnswerPairList,
- current: QuestionAnswerPairList,
- ) -> QuestionAnswerPairList:
- return QuestionAnswerPairList(
- QuestionAnswerPairs=accumulator.QuestionAnswerPairs
- + current.QuestionAnswerPairs,
- )
-
- return reduce(
- reducer,
- qa_pair_lists,
- QuestionAnswerPairList(QuestionAnswerPairs=[]),
- )
-
-
 def get_rag_qa_gen_chain(
  retriever: BaseRetriever,
  llm: BaseLanguageModel,
@@ -88,5 +68,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())
+ | (lambda parsed_output: parsed_output.to_str())
  )