Added streaming in the web ui

search_agent_ui.py

@@ -1,13 +1,15 @@
-import dotenv
+import datetime
 
+import dotenv
 import streamlit as st
 
-import web_rag as wr
-import web_crawler as wc
-
 from langchain_core.tracers.langchain import LangChainTracer
+from langchain.callbacks.base import BaseCallbackHandler
 from langsmith.client import Client
 
+import web_rag as wr
+import web_crawler as wc
+
 dotenv.load_dotenv()
 
 ls_tracer = LangChainTracer(
@@ -15,6 +17,14 @@ ls_tracer = LangChainTracer(
     client=Client()
 )
 
+class StreamHandler(BaseCallbackHandler):
+    def __init__(self, container, initial_text=""):
+        self.container = container
+        self.text = initial_text
+
+    def on_llm_new_token(self, token: str, **kwargs):
+        self.text += token
+        self.container.markdown(self.text)
 
 chat = wr.get_chat_llm(provider="cohere")
 
@@ -22,40 +32,64 @@ st.title("🔍 Simple Search Agent 💬")
 
 if "messages" not in st.session_state:
     st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
+    
+if "input_disabled" not in st.session_state:
+    st.session_state["input_disabled"] = False
 
 for message in st.session_state.messages:
     st.chat_message(message["role"]).write(message["content"])
+    if message["role"] == "assistant" and 'message_id' in message:
+        st.download_button(
+            label="Download",
+            data=message["content"],
+            file_name=f"{message['message_id']}.txt",
+            mime="text/plain"
+        )
 
-if prompt := st.chat_input():
+if prompt := st.chat_input("Enter you instructions...", disabled=st.session_state["input_disabled"] ):
+    
+    st.session_state["input_disabled"] = True
     
     st.chat_message("user").write(prompt)
     st.session_state.messages.append({"role": "user", "content": prompt})
-    
+
     message = "I first need to do some research"
     st.chat_message("assistant").write(message)
     st.session_state.messages.append({"role": "assistant", "content": message})
-    
+   
     with st.spinner("Optimizing search query"):
         optimize_search_query = wr.optimize_search_query(chat, query=prompt, callbacks=[ls_tracer])
-        
+            
     message = f"I'll search the web for: {optimize_search_query}"
     st.chat_message("assistant").write(message)
     st.session_state.messages.append({"role": "assistant", "content": message})
-    
-    
+       
     with st.spinner(f"Searching the web for: {optimize_search_query}"):
         sources = wc.get_sources(optimize_search_query, max_pages=20)
-        
+       
     with st.spinner(f"I'm now retrieveing the {len(sources)} webpages and documents I found (be patient)"):
         contents = wc.get_links_contents(sources)
 
-
     with st.spinner( f"Reading through the {len(contents)} sources I managed to retrieve"):
         vector_store = wc.vectorize(contents)
+    
+    message = f"Got {vector_store.index.ntotal} chunk of data"
+    st.chat_message("assistant").write(message)
+    st.session_state.messages.append({"role": "assistant", "content": message})
 
-    with st.spinner( "Ok I have now enough information to answer"):
-        response = wr.query_rag(chat, prompt, optimize_search_query, vector_store, callbacks=[ls_tracer])
-
-    st.chat_message("assistant").write(response)
-    st.session_state.messages.append({"role": "assistant", "content": response})
+    rag_prompt = wr.build_rag_prompt(prompt, optimize_search_query, vector_store, top_k=5, callbacks=[ls_tracer])  
+    with st.chat_message("assistant"):
+        st_cb = StreamHandler(st.empty())
+        result = chat.invoke(rag_prompt, stream=True, config={ "callbacks": [st_cb, ls_tracer]})
+        response = result.content.strip()
+        message_id = f"{prompt}{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}"
+        st.session_state.messages.append({"role": "assistant", "content": response})
+        if st.session_state.messages[-1]["role"] == "assistant":
+            st.download_button(
+                label="Download",
+                data=st.session_state.messages[-1]["content"],
+                file_name=f"{message_id}.txt",
+                mime="text/plain"
+            )
+    st.session_state["input_disabled"] = False
     

web_rag.py

@@ -120,6 +120,9 @@ def get_optimized_search_messages(query):
             Example:
                 Question: Write a short article about the solar system in the style of donald trump
                 Search query: solar system**
+            Exmaple:
+                Question: Write a short linkedin about how the "freakeconomics" book previsions didn't pan out
+                Search query: freakeconomics book predictions failed**
         """
     )
     human_message = HumanMessage(
@@ -209,9 +212,14 @@ def multi_query_rag(chat_llm, question, search_query, vectorstore, callbacks = [
     return response.content
 
 
-def query_rag(chat_llm, question, search_query, vectorstore, callbacks = []):
-    unique_docs = vectorstore.similarity_search(search_query, k=15, callbacks=callbacks, verbose=True)
+def build_rag_prompt(question, search_query, vectorstore, top_k = 10, callbacks = []):
+    unique_docs = vectorstore.similarity_search(
+        search_query, k=top_k, callbacks=callbacks, verbose=True)
     context = format_docs(unique_docs)
     prompt = get_rag_prompt_template().format(query=question, context=context)
+    return prompt    
+
+def query_rag(chat_llm, question, search_query, vectorstore, callbacks = []):
+    prompt = build_rag_prompt(question, search_query, vectorstore, callbacks)
     response = chat_llm.invoke(prompt, config={"callbacks": callbacks})
-    return response.content
+    return response.content