app.py

import gradio as gr
import random
from dotenv import load_dotenv
import os

from typing import TypedDict, Annotated
from langgraph.graph.message import add_messages
from langchain_core.messages import AnyMessage, HumanMessage, AIMessage
from langgraph.prebuilt import ToolNode
from langgraph.graph import START, StateGraph
from langgraph.prebuilt import tools_condition
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langchain_community.embeddings import HuggingFaceEmbeddings

from langgraph.prebuilt import create_react_agent
from langgraph.store.memory import InMemoryStore



# Import our custom tools from their modules
from tools import weather_info_tool, hub_stats_tool, duckduckgo_search_tool
from retriever import load_guest_dataset

# Load environment variables from .env file
load_dotenv()

HUGGINGFACEHUB_API_TOKEN=os.getenv("HUGGINGFACE_TOKEN")

# Generate the chat interface, including the tools
llm = HuggingFaceEndpoint(
    repo_id="Qwen/Qwen2.5-Coder-32B-Instruct",
    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
)

chat = ChatHuggingFace(llm=llm, verbose=True)

# Initialize memory store with HuggingFace embeddings
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
store = InMemoryStore(
    index={
        "dims": 384,  # Dimension for the MiniLM model
        "embed": embeddings,
    }
) 



# Load the guest dataset and initialize the guest info tool
guest_info_tool = load_guest_dataset()




tools = [guest_info_tool, weather_info_tool, hub_stats_tool, duckduckgo_search_tool]
chat_with_tools = chat.bind_tools(tools)

# Create Alfred with all the tools
agent = create_react_agent(
    "openai:gpt-4o-mini",
    tools=tools,
    store=store,
)

# Generate the AgentState and Agent graph
class AgentState(TypedDict):
    messages: Annotated[list[AnyMessage], add_messages]

def assistant(state: AgentState):
    # Prepare messages for the agent

    messages = [{"role": "user", "content": msg.content} for msg in state["messages"]]
    
    # Invoke the agent with the prepared messages
    response = agent.invoke({"messages": messages})
    print(response)
    
    # Ensure the response is a list of message dictionaries
    response_messages = [
        {"role": "assistant", "content": msg.content} for msg in response["messages"]
    ]
    print(response_messages)

    # Extract the response content from the last message
    response_content = response_messages[-1]["content"]
    print(response_content)
    
    return {
        "messages": [response_content],
    }


## The graph
builder = StateGraph(AgentState)

# Define nodes: these do the work
builder.add_node("assistant", assistant)
builder.add_node("tools", ToolNode(tools))

# Define the graph

builder.add_edge(START, "assistant")
builder.add_conditional_edges(
    "assistant",
    # If the latest message requires a tool, route to tools
    # Otherwise, provide a direct response
    tools_condition,
)

builder.add_edge("tools", "assistant")
alfred = builder.compile()



def GradioUI(chain):
    with gr.Blocks() as demo:
        chatbot = gr.Chatbot()
        msg = gr.Textbox()
        clear = gr.Button("Clear")

        def user(user_message, history):
            return "", history + [[user_message, None]]

        def bot(history):
            messages = [HumanMessage(content=history[-1][0])]
            response = chain.invoke({"messages": messages})
            bot_message = response["messages"][-1].content
            history[-1][1] = bot_message
            return history

        msg.submit(user, [msg, chatbot], [msg, chatbot]).then(
            bot, chatbot, chatbot
        )
        clear.click(lambda: None, None, chatbot, queue=False)

    return demo


if __name__ == "__main__":
    GradioUI(alfred).launch()