add sources in main chat

langchain-streamlit-demo/app.py

@@ -1,3 +1,4 @@
+import time
 from datetime import datetime
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -9,16 +10,16 @@ from defaults import default_values
 from langchain.agents.tools import tool
 from langchain.callbacks.base import BaseCallbackHandler
 from langchain.callbacks.manager import Callbacks
+
+# from langchain.callbacks.streamlit import StreamlitCallbackHandler
 from langchain.callbacks.tracers.langchain import LangChainTracer, wait_for_all_tracers
 from langchain.callbacks.tracers.run_collector import RunCollectorCallbackHandler
 from langchain.memory import ConversationBufferMemory, StreamlitChatMessageHistory
 from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
 from langchain.schema.document import Document
 from langchain.schema.retriever import BaseRetriever
-from langchain_community.callbacks import StreamlitCallbackHandler
 from langsmith.client import Client
 from llm_resources import (
-    get_agent,
     get_llm,
     get_runnable,
     get_texts_and_multiretriever,
@@ -44,6 +45,8 @@ def st_init_null(*variable_names) -> None:
 
 
 st_init_null(
+    "start_time",
+    "sources",
     "chain",
     "client",
     "doc_chain",
@@ -430,7 +433,7 @@ if st.session_state.llm:
     if question and question != "--" and not prompt:
         prompt = question
     if not uploaded_file:
-        st.error("Please upload a PDF to use the document chat feature.")
+        st.warning("Please upload a PDF to use the document chat feature.")
     elif prompt:
         feedback_update = None
         feedback = None
@@ -463,6 +466,7 @@ if st.session_state.llm:
             full_response: Union[str, None] = None
 
             message_placeholder = st.empty()
+            sources_placeholder = st.empty()
             default_tools = [
                 # DuckDuckGoSearchRun(),
                 # WikipediaQueryRun(api_wrapper=WikipediaAPIWrapper()),
@@ -474,10 +478,10 @@ if st.session_state.llm:
                 #     search_llm=get_llm(**get_llm_args_temp_zero),  # type: ignore
                 #     writer_llm=get_llm(**get_llm_args_temp_zero),  # type: ignore
                 # )
-                st_callback = StreamlitCallbackHandler(
-                    st.container(), expand_new_thoughts=False
-                )
-                callbacks.append(st_callback)
+                # st_callback = StreamlitCallbackHandler(
+                #     st.container(), expand_new_thoughts=False
+                # )
+                # callbacks.append(st_callback)
 
                 # @tool("web-research-assistant")
                 # def research_assistant_tool(question: str, callbacks: Callbacks = None):
@@ -526,10 +530,16 @@ if st.session_state.llm:
                         input_str,
                         config=get_config(callbacks),
                     )
+                    st.session_state.sources = response["source_documents"][:3]
                     with st.sidebar.expander("Sources"):
-                        for source in response["source_documents"][:3]:
+                        for source in st.session_state.sources:
                             st.markdown("-" * 50)
                             st.markdown(source.page_content)
+                    # with sources_placeholder:
+                    #     with st.expander("Sources", expanded=False):
+                    #         for source in response["source_documents"][:3]:
+                    #             st.markdown("-" * 50)
+                    #             st.markdown(source.page_content)
                     return response["output_text"]
 
                 # doc_chain_agent = get_doc_agent(
@@ -555,29 +565,31 @@ if st.session_state.llm:
 
                 TOOLS = TOOLS + [doc_chain_tool]
 
-                st.session_state.chain = get_agent(
-                    TOOLS,
-                    STMEMORY,
-                    st.session_state.llm,
-                    callbacks,
-                )
-                # else:
-                # st.session_state.chain = get_runnable(
-                #     True,  # use_document_chat,
-                #     document_chat_chain_type,
+                # st.session_state.chain = get_agent(
+                #     TOOLS,
+                #     STMEMORY,
                 #     st.session_state.llm,
-                #     st.session_state.retriever,
-                #     MEMORY,
-                #     chat_prompt,
-                #     prompt,
+                #     callbacks,
                 # )
+                # else:
+                st.session_state.chain = get_runnable(
+                    True,  # use_document_chat,
+                    document_chat_chain_type,
+                    st.session_state.llm,
+                    st.session_state.retriever,
+                    MEMORY,
+                    chat_prompt,
+                    prompt,
+                )
 
             # --- LLM call ---
             try:
-                full_response = st.session_state.chain.invoke(
-                    prompt,
-                    config=get_config(callbacks),
-                )
+                with st.spinner("Thinking..."):
+                    st.session_state.start_time = time.time()
+                    full_response = st.session_state.chain.invoke(
+                        prompt,
+                        config=get_config(callbacks),
+                    )
 
             except (openai.AuthenticationError, anthropic.AuthenticationError):
                 st.error(
@@ -587,7 +599,20 @@ if st.session_state.llm:
 
             # --- Display output ---
             if full_response is not None:
-                message_placeholder.markdown(full_response)
+                finish_time = time.time() - st.session_state.start_time
+                text_response, sources = (
+                    full_response["output_text"],
+                    full_response["source_documents"],
+                )
+                st.session_state.sources = sources[:3]
+                message_placeholder.markdown(
+                    text_response + "\n\n" + f"⏱️ {finish_time:.2f}s"
+                )
+                with sources_placeholder:
+                    with st.expander("Sources", expanded=False):
+                        for source in st.session_state.sources:
+                            st.markdown("-" * 50)
+                            st.markdown(source.page_content)
 
                 # --- Tracing ---
                 if st.session_state.client: