Upload 2 files
Browse files- app.py +94 -0
- requirements.txt +9 -0
app.py
ADDED
@@ -0,0 +1,94 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import os
|
3 |
+
from typing import Annotated
|
4 |
+
from typing_extensions import TypedDict
|
5 |
+
from langchain_community.utilities import ArxivAPIWrapper, WikipediaAPIWrapper
|
6 |
+
from langchain_community.tools import ArxivQueryRun, WikipediaQueryRun
|
7 |
+
from langgraph.graph.message import add_messages
|
8 |
+
from langgraph.graph import StateGraph, START, END
|
9 |
+
from langchain_groq import ChatGroq
|
10 |
+
from langgraph.prebuilt import ToolNode, tools_condition
|
11 |
+
|
12 |
+
# Initialize tools
|
13 |
+
arxiv_wrapper = ArxivAPIWrapper(top_k_results=1, doc_content_chars_max=300)
|
14 |
+
arxiv_tool = ArxivQueryRun(api_wrapper=arxiv_wrapper)
|
15 |
+
wiki_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=300)
|
16 |
+
wiki_tool = WikipediaQueryRun(api_wrapper=wiki_wrapper)
|
17 |
+
tools = [wiki_tool, arxiv_tool]
|
18 |
+
|
19 |
+
# Define State
|
20 |
+
class State(TypedDict):
|
21 |
+
messages: Annotated[list, add_messages]
|
22 |
+
|
23 |
+
# Initialize LLM
|
24 |
+
@st.cache_resource
|
25 |
+
def initialize_llm():
|
26 |
+
groq_api_key = os.getenv("GROQ_API_KEY")
|
27 |
+
if not groq_api_key:
|
28 |
+
st.error("Please set the GROQ_API_KEY environment variable.")
|
29 |
+
st.stop()
|
30 |
+
return ChatGroq(groq_api_key=groq_api_key, model_name="Gemma2-9b-It")
|
31 |
+
|
32 |
+
llm = initialize_llm()
|
33 |
+
llm_with_tools = llm.bind_tools(tools=tools)
|
34 |
+
|
35 |
+
# Define chatbot function
|
36 |
+
def chatbot(state: State):
|
37 |
+
return {"messages": [llm_with_tools.invoke(state["messages"])]}
|
38 |
+
|
39 |
+
# Build graph
|
40 |
+
@st.cache_resource
|
41 |
+
def build_graph():
|
42 |
+
graph_builder = StateGraph(State)
|
43 |
+
graph_builder.add_node("chatbot", chatbot)
|
44 |
+
tool_node = ToolNode(tools=tools)
|
45 |
+
graph_builder.add_node("tools", tool_node)
|
46 |
+
graph_builder.add_conditional_edges("chatbot", tools_condition)
|
47 |
+
graph_builder.add_edge("tools", "chatbot")
|
48 |
+
graph_builder.add_edge(START, "chatbot")
|
49 |
+
return graph_builder.compile()
|
50 |
+
|
51 |
+
graph = build_graph()
|
52 |
+
|
53 |
+
# Streamlit UI
|
54 |
+
st.title("WIKXIV AI: Wikipedia and ArXiv Chatbot")
|
55 |
+
|
56 |
+
# Initialize chat history
|
57 |
+
if "messages" not in st.session_state:
|
58 |
+
st.session_state.messages = []
|
59 |
+
|
60 |
+
# Display chat messages from history on app rerun
|
61 |
+
for message in st.session_state.messages:
|
62 |
+
with st.chat_message(message["role"]):
|
63 |
+
st.markdown(message["content"])
|
64 |
+
|
65 |
+
# React to user input
|
66 |
+
if prompt := st.chat_input("What is your question?"):
|
67 |
+
# Display user message in chat message container
|
68 |
+
st.chat_message("user").markdown(prompt)
|
69 |
+
# Add user message to chat history
|
70 |
+
st.session_state.messages.append({"role": "user", "content": prompt})
|
71 |
+
|
72 |
+
# Get bot response
|
73 |
+
events = graph.stream(
|
74 |
+
{"messages": [("user", prompt)]},
|
75 |
+
stream_mode="values"
|
76 |
+
)
|
77 |
+
|
78 |
+
# Display assistant response in chat message container
|
79 |
+
with st.chat_message("assistant"):
|
80 |
+
message_placeholder = st.empty()
|
81 |
+
full_response = ""
|
82 |
+
for event in events:
|
83 |
+
message = event["messages"][-1]
|
84 |
+
full_response += message.content
|
85 |
+
message_placeholder.markdown(full_response + "▌")
|
86 |
+
message_placeholder.markdown(full_response)
|
87 |
+
|
88 |
+
# Add assistant response to chat history
|
89 |
+
st.session_state.messages.append({"role": "assistant", "content": full_response})
|
90 |
+
|
91 |
+
# Display info about the app
|
92 |
+
st.sidebar.title("About")
|
93 |
+
st.sidebar.info("This is a Streamlit app made by theaimart\nthat uses LangGraph to create a chatbot with access to Wikipedia and ArXiv tools.")
|
94 |
+
# st.sidebar.title("made with love by theaimart")
|
requirements.txt
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
langgraph
|
2 |
+
langsmith
|
3 |
+
langchain
|
4 |
+
langchain_groq
|
5 |
+
langchain_community
|
6 |
+
arxiv
|
7 |
+
wikipedia-api
|
8 |
+
streamlit
|
9 |
+
wikipedia
|