agentic-rag / app.py
cgoncalves's picture
Update space
f4a94cb
import os
from typing import TypedDict, Annotated
from dotenv import load_dotenv
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langgraph.checkpoint.memory import MemorySaver
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langfuse.callback import CallbackHandler
from retriever import guest_info_tool
from tools import search_tool, weather_info_tool, hub_stats_tool
import gradio as gr
load_dotenv()
# Initialize Langfuse
langfuse_handler = CallbackHandler(
secret_key=os.getenv("LANGFUSE_SECRET_KEY"),
public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),
host=os.getenv("LANGFUSE_HOST"),
)
# Generate the chat interface, including the tools
llm = HuggingFaceEndpoint(
repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN")
)
chat = ChatHuggingFace(llm=llm, verbose=True)
tools = [guest_info_tool, search_tool, weather_info_tool, hub_stats_tool]
chat_with_tools = chat.bind_tools(tools)
# Define the AgentState structure
class AgentState(TypedDict):
messages: Annotated[list[AnyMessage], add_messages]
def assistant(state: AgentState):
# Define the system prompt (this is not part of the conversation history)
system_message = SystemMessage(content="""You are Alfred, a helpful and sophisticated assistant.
Your capabilities:
- Answer questions using your knowledge
- Search the web for recent or factual information using the DuckDuckGoSearchResults tool
- Retrieve information about guests using the guest_info_tool
- Use weather_info_tool to give information about the weather
Guidelines:
- Be concise, polite and helpful
- When you don't know something, use the appropriate tool rather than guessing
- For guest information requests, always use the guest_info_tool first
- For factual or current information, use the search tool
- Present information in a clear, organized manner
Always think carefully about which tool is most appropriate for the user's request.
""")
# Call the agent with the system prompt and conversation history (state messages)
assistant_response = chat_with_tools.invoke([system_message] + state.get("messages"))
# Return the updated conversation state including the new assistant response
return {"messages": state.get("messages") + [assistant_response]}
# Build the graph
builder = StateGraph(AgentState)
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))
builder.add_edge(START, "assistant")
builder.add_conditional_edges("assistant", tools_condition)
builder.add_edge("tools", "assistant")
checkpointer = MemorySaver()
alfred = builder.compile(checkpointer=checkpointer)
config = {"configurable": {"thread_id": "1"}, "callbacks": [langfuse_handler]}
# Gradio chat function
def chat_fn(message, history):
state = {"messages": [HumanMessage(content=message)]}
result = alfred.invoke(state, config)
return result["messages"][-1].content
# Launch Gradio interface
interface = gr.ChatInterface(fn=chat_fn,type="messages")
interface.launch()