THEAIMART commited on
Commit
82b0e20
·
verified ·
1 Parent(s): eb78aec

Upload 2 files

Browse files
Files changed (2) hide show
  1. app.py +94 -0
  2. 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