Update app.py

app.py

@@ -1,19 +1,15 @@
-import os
-import gradio as gr
-from langchain.agents import AgentExecutor, Tool
-from langchain.prompts import PromptTemplate
+from langgraph.graph import StateGraph, MessagesState, START, END
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import HumanMessage
 from langchain_huggingface import HuggingFaceEndpoint
-from dotenv import load_dotenv  # Optional for .env files
-
-# Load environment variables from .env file
-load_dotenv()
+from langchain_core.tools import tool
+from typing import Literal
+import os
 
-# Get Hugging Face token from environment variable
-HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("Hugging Face API token not found. Set HF_TOKEN in your environment variables.")
+# Hugging Face Endpoint setup
+HF_TOKEN = os.getenv("HF_TOKEN")  # Ensure your API key is in the environment
 
-# Define the Hugging Face model endpoint
 llm = HuggingFaceEndpoint(
     repo_id="mistralai/Mistral-7B-Instruct-v0.3",
     huggingfacehub_api_token=HF_TOKEN.strip(),
@@ -21,85 +17,72 @@ llm = HuggingFaceEndpoint(
     max_new_tokens=150,
 )
 
-# Define specialized agents
-
-def flight_booking(input):
-    return f"Searching flights for: {input}"
-
-flight_tool = Tool(
-    name="Flight Booking Tool",
-    func=flight_booking,
-    description="Finds and books airline flights"
-)
-
-def hotel_booking(input):
-    return f"Searching hotels for: {input}"
-
-hotel_tool = Tool(
-    name="Hotel Booking Tool",
-    func=hotel_booking,
-    description="Searches and books accommodations"
-)
-
-def transportation_booking(input):
-    return f"Arranging transportation for: {input}"
-
-transport_tool = Tool(
-    name="Transportation Booking Tool",
-    func=transportation_booking,
-    description="Handles rental cars, shuttles, or trains"
-)
-
-def activity_booking(input):
-    return f"Finding activities for: {input}"
-
-activity_tool = Tool(
-    name="Activity Booking Tool",
-    func=activity_booking,
-    description="Books activities, tours, and events"
-)
-
-tools = [flight_tool, hotel_tool, transport_tool, activity_tool]
-
-# Define a prompt for the planner agent
-planner_prompt = PromptTemplate(
-    template="You are a travel planner. Your task is to coordinate tools to plan a trip based on user input: {input}",
-    input_variables=["input"]
+# Define tools for travel planning
+@tool
+def search_destination(query: str):
+    """Search for travel destinations based on preferences."""
+    if "beach" in query.lower():
+        return "Destinations: Bali, Maldives, Goa"
+    elif "adventure" in query.lower():
+        return "Destinations: Patagonia, Kilimanjaro, Swiss Alps"
+    return "Popular Destinations: Paris, New York, Tokyo"
+
+@tool
+def fetch_flights(destination: str):
+    """Fetch flight options to the destination."""
+    return f"Flights available to {destination}: Option 1 - $500, Option 2 - $700."
+
+@tool
+def fetch_hotels(destination: str):
+    """Fetch hotel recommendations for the destination."""
+    return f"Hotels in {destination}: Hotel A ($200/night), Hotel B ($150/night)."
+
+tools = [search_destination, fetch_flights, fetch_hotels]
+
+# Bind tools to the Hugging Face model
+llm = llm.bind_tools(tools)
+
+# Define the function to determine the next step
+def should_continue(state: MessagesState) -> Literal["tools", END]:
+    messages = state['messages']
+    last_message = messages[-1]
+    if last_message.tool_calls:
+        return "tools"
+    return END
+
+# Define the function to call the LLM
+def call_model(state: MessagesState):
+    messages = state['messages']
+    response = llm.invoke(messages)
+    return {"messages": [response]}
+
+# Create the graph
+workflow = StateGraph(MessagesState)
+
+# Define nodes
+workflow.add_node("agent", call_model)
+workflow.add_node("tools", ToolNode(tools))
+
+# Define edges
+workflow.add_edge(START, "agent")
+workflow.add_conditional_edges("agent", should_continue)
+workflow.add_edge("tools", "agent")
+
+# Initialize memory for persistence
+checkpointer = MemorySaver()
+
+# Compile the graph
+app = workflow.compile(checkpointer=checkpointer)
+
+# Use the graph
+initial_state = {"messages": [HumanMessage(content="I want a beach vacation.")]}
+
+final_state = app.invoke(initial_state)
+print(final_state["messages"][-1].content)
+
+# Continue the conversation
+next_state = app.invoke(
+    {"messages": [HumanMessage(content="Tell me about flights to Bali.")]},
+    config={"configurable": {"thread_id": 123}}  # Reuse thread_id to maintain context
 )
-
-# Create the agent
-planner_agent = AgentExecutor.from_agent_and_tools(
-    llm=llm,
-    tools=tools,
-    verbose=True
-)
-
-# Define the Gradio function
-def plan_trip(user_input):
-    """
-    Processes user input through the planner agent and returns the output.
-    """
-    try:
-        output = planner_agent.run(user_input)
-        return output
-    except Exception as e:
-        return f"Error: {e}"
-
-# Create the Gradio interface
-with gr.Blocks() as travel_planner_app:
-    gr.Markdown("## AI Travel Planner")
-    gr.Markdown("Enter your trip details below, and the AI will help you plan your travel itinerary!")
-    
-    with gr.Row():
-        user_input = gr.Textbox(
-            label="Trip Details",
-            placeholder="E.g., Plan a trip to Paris from May 10 to May 15 for a family of 4."
-        )
-    submit_button = gr.Button("Plan Trip")
-    output = gr.Textbox(label="Travel Plan")
-    
-    submit_button.click(plan_trip, inputs=user_input, outputs=output)
-
-# Launch the Gradio app
-if __name__ == "__main__":
-    travel_planner_app.launch()
+print(next_state["messages"][-1].content)