File size: 2,032 Bytes
a511d8f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 | import streamlit as st
from langchain_groq import ChatGroq
from langchain_community.utilities import ArxivAPIWrapper,WikipediaAPIWrapper
from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun, DuckDuckGoSearchRun
from langchain.agents import initialize_agent, AgentType
from langchain.callbacks import StreamlitCallbackHandler
from dotenv import load_dotenv
import os
load_dotenv()
groq_api_key = os.getenv("GROQ_API_KEY")
arxiv_wrapper = ArxivAPIWrapper(top_k_results=3, doc_content_char_max = 200)
arxiv = ArxivQueryRun(api_wrapper=arxiv_wrapper)
wiki_wrapper = WikipediaAPIWrapper(top_k_results=3, doc_content_char_max = 200)
wiki = WikipediaQueryRun(api_wrapper=wiki_wrapper)
search = DuckDuckGoSearchRun(name="Search")
tools = [arxiv, wiki, search]
llm = ChatGroq(groq_api_key=groq_api_key, model="Llama3-8b-8192", streaming=True)
search_agent = initialize_agent(
tools=tools,
llm=llm,
agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
handle_parsing_errors=True
)
st.title("Search Engine with Agents")
st.write("This is a search engine that uses agents to search for information.")
st.write("You can search for information using Arxiv, Wikipedia, or DuckDuckGo.")
if "messages" not in st.session_state:
st.session_state.messages = [
{"role": "assistant", "content": "Hello! How can I assist you today?"}
]
for message in st.session_state.messages:
st.chat_message(message["role"]).markdown(message["content"])
prompt = st.chat_input("Enter your question:")
if prompt:
st.session_state.messages.append({"role": "user", "content": prompt})
st.chat_message("user").markdown(prompt)
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)
|