cgoncalves/agentic-rag

1import os2from typing import TypedDict, Annotated3from dotenv import load_dotenv4from langgraph.graph.message import add_messages5from langchain_core.messages import AnyMessage, HumanMessage, SystemMessage6from langgraph.prebuilt import ToolNode7from langgraph.graph import START, StateGraph8from langgraph.prebuilt import tools_condition9from langgraph.checkpoint.memory import MemorySaver10from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace11from langfuse.callback import CallbackHandler12from retriever import guest_info_tool13from tools import search_tool, weather_info_tool, hub_stats_tool14 15import gradio as gr16 17load_dotenv()18 19# Initialize Langfuse20langfuse_handler = CallbackHandler(21    secret_key=os.getenv("LANGFUSE_SECRET_KEY"),22    public_key=os.getenv("LANGFUSE_PUBLIC_KEY"),23    host=os.getenv("LANGFUSE_HOST"),24)25 26# Generate the chat interface, including the tools27llm = HuggingFaceEndpoint(28    repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",29    huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN")30)31chat = ChatHuggingFace(llm=llm, verbose=True)32tools = [guest_info_tool, search_tool, weather_info_tool, hub_stats_tool]33chat_with_tools = chat.bind_tools(tools)34 35 36# Define the AgentState structure37class AgentState(TypedDict):38    messages: Annotated[list[AnyMessage], add_messages]39 40 41def assistant(state: AgentState):42    # Define the system prompt (this is not part of the conversation history)43    system_message = SystemMessage(content="""You are Alfred, a helpful and sophisticated assistant.44 45Your capabilities:46- Answer questions using your knowledge47- Search the web for recent or factual information using the DuckDuckGoSearchResults tool48- Retrieve information about guests using the guest_info_tool49- Use weather_info_tool to give information about the weather50 51Guidelines:52- Be concise, polite and helpful53- When you don't know something, use the appropriate tool rather than guessing54- For guest information requests, always use the guest_info_tool first55- For factual or current information, use the search tool56- Present information in a clear, organized manner57 58Always think carefully about which tool is most appropriate for the user's request.59""")60    # Call the agent with the system prompt and conversation history (state messages)61    assistant_response = chat_with_tools.invoke([system_message] + state.get("messages"))62    # Return the updated conversation state including the new assistant response63    return {"messages": state.get("messages") + [assistant_response]}64 65# Build the graph66builder = StateGraph(AgentState)67builder.add_node("assistant", assistant)68builder.add_node("tools", ToolNode(tools))69builder.add_edge(START, "assistant")70builder.add_conditional_edges("assistant", tools_condition)71builder.add_edge("tools", "assistant")72checkpointer = MemorySaver()73alfred = builder.compile(checkpointer=checkpointer)74config = {"configurable": {"thread_id": "1"}, "callbacks": [langfuse_handler]}75 76# Gradio chat function77def chat_fn(message, history):78    state = {"messages": [HumanMessage(content=message)]}79    result = alfred.invoke(state, config)80    return result["messages"][-1].content81 82# Launch Gradio interface83interface = gr.ChatInterface(fn=chat_fn,type="messages")84interface.launch()85