LLM_Search / src /streamlit_app.py
vatsal18's picture
Update src/streamlit_app.py
c21a97d verified
import streamlit as st
from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun, DuckDuckGoSearchRun
from langchain_community.utilities import WikipediaAPIWrapper, ArxivAPIWrapper
from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent,AgentType
from langchain.callbacks import StreamlitCallbackHandler
import os
from dotenv import load_dotenv
load_dotenv()
api_wiki=WikipediaAPIWrapper(top_k_results=3,doc_content_chars_max=250)
wiki=WikipediaQueryRun(api_wrapper=api_wiki)
api_wrapper_arxiv=ArxivAPIWrapper(top_k_results=1,doc_content_chars_max=250)
arxiv=ArxivQueryRun(api_wrapper=api_wrapper_arxiv)
search=DuckDuckGoSearchRun(name="Search")
st.title("๐Ÿ”Ž LangChain - Chat with search")
api_key = st.text_input("Enter your Open AI API key:", type="password")
if "messages" not in st.session_state:
st.session_state["messages"]=[
{"role":"assistant","content":"Hi,I'm a chatbot who can search the web. How can I help you?"}
]
for msg in st.session_state.messages:
st.chat_message(msg["role"]).write(msg['content'])
if prompt:=st.chat_input(placeholder="What is machine learning?"):
st.session_state.messages.append({"role":"user","content":prompt})
st.chat_message("user").write(prompt)
llm=ChatOpenAI(api_key=api_key, model="gpt-4o",streaming=True)
tools=[search,arxiv,wiki]
search_agent=initialize_agent(tools,llm,agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,handling_parsing_errors=True)
with st.chat_message("assistant"):
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)