Search agent improvements

.pre-commit-config.yaml

@@ -14,3 +14,4 @@ repos:
           # pre-commit's default_language_version, see
           # https://pre-commit.com/#top_level-default_language_version
           language_version: python3.11
+          args: ["--line-length", "100"]

streamlit_agent/search_and_chat.py

@@ -1,7 +1,7 @@
-from langchain.agents import initialize_agent, load_tools
-from langchain.agents import AgentType
+from langchain.agents import initialize_agent, AgentType
 from langchain.callbacks import StreamlitCallbackHandler
 from langchain.chat_models import ChatOpenAI
+from langchain.tools import DuckDuckGoSearchRun
 import streamlit as st
 
 st.set_page_config(page_title="LangChain: Chat with search", page_icon="🦜")
@@ -9,9 +9,7 @@ st.title("🦜 LangChain: Chat with search")
 
 openai_api_key = st.sidebar.text_input("OpenAI API Key", type="password")
 if "messages" not in st.session_state:
-    st.session_state["messages"] = [
-        {"role": "assistant", "content": "How can I help you?"}
-    ]
+    st.session_state["messages"] = [{"role": "assistant", "content": "How can I help you?"}]
 
 for msg in st.session_state.messages:
     st.chat_message(msg["role"]).write(msg["content"])
@@ -24,14 +22,15 @@ if prompt := st.chat_input(placeholder="Who won the Women's U.S. Open in 2018?")
         st.info("Please add your OpenAI API key to continue.")
         st.stop()
 
-    llm = ChatOpenAI(
-        model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, streaming=True
-    )
+    llm = ChatOpenAI(model_name="gpt-3.5-turbo", openai_api_key=openai_api_key, streaming=True)
     search_agent = initialize_agent(
-        load_tools(["ddg-search"]), llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION
+        tools=[DuckDuckGoSearchRun(name="Search")],
+        llm=llm,
+        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
+        handle_parsing_errors=True,
     )
     with st.chat_message("assistant"):
-        st_cb = StreamlitCallbackHandler(st.container())
+        st_cb = StreamlitCallbackHandler(st.container(), expand_new_thoughts=False)
         response = search_agent.run(st.session_state.messages, callbacks=[st_cb])
         st.session_state.messages.append({"role": "assistant", "content": response})
         st.write(response)