Update wanderlust.py

wanderlust.py

@@ -1,134 +1,391 @@
-import gradio as gr
-from langchain_groq import ChatGroq
-from langchain_tavily import TavilySearch
-import os
-from dotenv import load_dotenv
-
-# Load environment variables
-load_dotenv()
-
-# Set API keys
-TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
-GROQ_API_KEY = os.getenv("GROQ_API_KEY")
-
-# Initialize models
-llm = ChatGroq(api_key=GROQ_API_KEY, model="llama3-8b-8192")  # Correct Groq model
-search_tool = TavilySearch(
-    max_results=5,
-    topic="general"
-)
-
-# State stored as a dictionary
-state = {
-    "destination": None,
-    "budget": None,
-    "activities": None,
-    "duration": None,
-    "accommodation": None,
-    "step": -1  # -1 to wait for START
-}
-
-# Define questions
-questions = [
-    "Where would you like to travel?",
-    "What is your budget for this trip?",
-    "What kind of activities do you prefer? (e.g., adventure, relaxation, sightseeing)",
-    "How many days do you plan to stay?",
-    "Do you prefer hotels, hostels, or Airbnbs?"
-]
-
-
-# Function to handle user responses and ask the next question
-def ask_question(user_response, chat_history):
-    # Check if user wants to start the process
-    if state["step"] == -1:
-        if user_response.strip().lower() == "start":
-            state["step"] = 0
-            initial_message = {"role": "assistant", "content": questions[0]}
-            chat_history.append(initial_message)
-            return chat_history, ""
-        else:
-            chat_history.append({"role": "assistant", "content": "❗ Type `START` to begin the travel itinerary process."})
-            return chat_history, ""
-
-    # Handle user responses
-    if user_response:
-        if state["step"] == 0:
-            state["destination"] = user_response
-        elif state["step"] == 1:
-            state["budget"] = user_response
-        elif state["step"] == 2:
-            state["activities"] = user_response
-        elif state["step"] == 3:
-            state["duration"] = user_response
-        elif state["step"] == 4:
-            state["accommodation"] = user_response
-
-        state["step"] += 1
-
-    # Check if all questions are answered
-    if state["step"] < len(questions):
-        next_question = questions[state["step"]]
-        chat_history.append({"role": "assistant", "content": next_question})
-        return chat_history, ""
-    else:
-        itinerary = generate_itinerary()
-        chat_history.append({"role": "assistant", "content": itinerary})
-        state["step"] = -1  # Reset for next session
-        return chat_history, ""
-
-
-# Generate the travel itinerary based on user inputs
-def generate_itinerary():
-    # Search for travel recommendations
-    search_query = f"Top things to do in {state['destination']} for {state['duration']} days with a budget of {state['budget']} and preference for {state['activities']}"
-    search_results = search_tool.run(search_query)
-
-    # Generate itinerary using search data and user preferences
-    itinerary_prompt = f"""
-    Create a detailed {state['duration']}-day travel itinerary for {state['destination']}.
-    Consider the following preferences:
-    - Budget: {state['budget']}
-    - Activities: {state['activities']}
-    - Accommodation type: {state['accommodation']}
-    
-    Use the following recommendations:
-    {search_results}
-    """
-    itinerary = llm.predict(itinerary_prompt)
-
-    return f"✅ Your travel itinerary is ready!\n\n{itinerary}"
-
-
-# Restart conversation (Start Over button)
-def start_over():
-    state.update({
-        "destination": None,
-        "budget": None,
-        "activities": None,
-        "duration": None,
-        "accommodation": None,
-        "step": -1  # Reset for next session
-    })
-    initial_message = {"role": "assistant", "content": "✅ Restarted! Type `START` to begin again."}
-    return [initial_message], ""
-
-
-# Gradio Interface
-with gr.Blocks() as app:
-    gr.Markdown("# 🌍 AI-Powered Travel Itinerary Generator")
-
-    # Chat interface with updated `type="messages"`
-    chatbot = gr.Chatbot(label="Travel Bot", type="messages")
-    user_input = gr.Textbox(label="Your Answer", placeholder="Type `START` to begin...")
-    submit_btn = gr.Button("Send")
-
-    # Handle conversation flow
-    submit_btn.click(ask_question, inputs=[user_input, chatbot], outputs=[chatbot, user_input])
-
-    # Start over button
-    start_btn = gr.Button("Start Over")
-    start_btn.click(start_over, inputs=[], outputs=[chatbot, user_input])
-
-# Run the Gradio app
-app.launch()
+# --- Imports ---
+import gradio as gr
+from langchain_groq import ChatGroq
+from langchain_community.tools.tavily_search import TavilySearchResults # Updated import
+from langchain_core.messages import HumanMessage, AIMessage, BaseMessage
+import os
+from dotenv import load_dotenv
+from typing import TypedDict, List, Optional, Dict, Sequence
+from langgraph.graph import StateGraph, END
+
+# --- Environment Setup ---
+load_dotenv()
+
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+
+if not TAVILY_API_KEY or not GROQ_API_KEY:
+    raise ValueError("API keys for Tavily and Groq must be set in .env file")
+
+# --- Constants ---
+ORDERED_FIELDS = [
+    "destination",
+    "budget",
+    "activities",
+    "duration",
+    "accommodation",
+]
+
+QUESTIONS_MAP = {
+    "destination": "🌍 Where would you like to travel?",
+    "budget": "💰 What is your budget for this trip?",
+    "activities": "🤸 What kind of activities do you prefer? (e.g., adventure, relaxation, sightseeing)",
+    "duration": "⏳ How many days do you plan to stay?",
+    "accommodation": "🏨 Do you prefer hotels, hostels, or Airbnbs?",
+}
+
+INITIAL_MESSAGE = "👋 Welcome! Type `START` to begin planning your travel itinerary."
+START_CONFIRMATION = "🚀 Great! Let's plan your trip. " + QUESTIONS_MAP[ORDERED_FIELDS[0]]
+RESTART_MESSAGE = "✅ Restarted! Type `START` to begin again."
+INVALID_START_MESSAGE = "❗ Please type `START` to begin the travel itinerary process."
+ITINERARY_READY_MESSAGE = "✅ Your travel itinerary is ready!"
+
+# --- LangChain / LangGraph Components ---
+
+# Initialize models and tools
+llm = ChatGroq(api_key=GROQ_API_KEY, model="llama3-8b-8192")
+# Note: TavilySearchResults often works better when integrated as a Langchain Tool
+# but for direct use like this, the previous langchain_tavily.TavilySearch is also fine.
+# Let's use the more standard Tool wrapper approach for better compatibility.
+search_tool = TavilySearchResults(max_results=5, tavily_api_key=TAVILY_API_KEY)
+
+
+# Define the state for our graph
+class AgentState(TypedDict):
+    messages: Sequence[BaseMessage]  # Conversation history
+    user_info: Dict[str, str]      # Collected user preferences
+    missing_fields: List[str]      # Fields still needed
+    search_results: Optional[str]  # Results from Tavily search
+    itinerary: Optional[str]       # Final generated itinerary
+
+# --- Node Functions ---
+
+def process_user_input(state: AgentState) -> Dict:
+    """Processes the latest user message to update user_info."""
+    last_message = state['messages'][-1]
+    if not isinstance(last_message, HumanMessage):
+        # Should not happen in normal flow, but good safeguard
+        return {}
+
+    user_response = last_message.content.strip()
+    missing_fields = state.get('missing_fields', [])
+    user_info = state.get('user_info', {})
+
+    if not missing_fields:
+        # All info gathered, nothing to process here for now
+        # This node is primarily for capturing answers to questions
+        return {}
+
+    # Assume the user is answering the question related to the *first* missing field
+    field_to_update = missing_fields[0]
+    user_info[field_to_update] = user_response
+
+    # Remove the field we just got
+    updated_missing_fields = missing_fields[1:]
+
+    print(f"--- Processed Input: Field '{field_to_update}' = '{user_response}' ---")
+    print(f"--- Remaining Fields: {updated_missing_fields} ---")
+
+    return {"user_info": user_info, "missing_fields": updated_missing_fields}
+
+def ask_next_question(state: AgentState) -> Dict:
+    """Adds the next question to the messages list."""
+    missing_fields = state.get('missing_fields', [])
+    if not missing_fields:
+        # Should not be called if no fields are missing, but handle defensively
+        return {"messages": state['messages'] + [AIMessage(content="Something went wrong, no more questions to ask.")]} # type: ignore
+
+    next_field = missing_fields[0]
+    question = QUESTIONS_MAP[next_field]
+
+    print(f"--- Asking Question for: {next_field} ---")
+    # Append the new question
+    updated_messages = state['messages'] + [AIMessage(content=question)] # type: ignore
+    return {"messages": updated_messages}
+
+
+def run_search(state: AgentState) -> Dict:
+    """Runs Tavily search based on collected user info, handling descriptive terms."""
+    user_info = state['user_info']
+    print("--- Running Search ---")
+
+    # Construct a more descriptive query for the search engine
+    query_parts = [f"Travel itinerary ideas for {user_info.get('destination', 'anywhere')}"]
+    if user_info.get('duration'):
+         # Try to clarify duration if it's non-numeric, otherwise use as is
+         duration_desc = user_info.get('duration')
+         query_parts.append(f"for about {duration_desc} days.")
+
+    if user_info.get('budget'):
+        # Frame budget as a description
+        query_parts.append(f"User's budget is described as: '{user_info.get('budget')}'.")
+
+    if user_info.get('activities'):
+        # Frame activities as a description
+        query_parts.append(f"User is looking for activities like: '{user_info.get('activities')}'.")
+
+    if user_info.get('accommodation'):
+        # Frame accommodation as a description
+        query_parts.append(f"User's accommodation preference is: '{user_info.get('accommodation')}'.")
+
+    search_query = " ".join(query_parts)
+    print(f"--- Refined Search Query: {search_query} ---")
+
+    try:
+        results = search_tool.invoke(search_query)
+        # Handle potential result formats (string, list of docs, dict)
+        if isinstance(results, list):
+             search_results_str = "\n\n".join([getattr(doc, 'page_content', str(doc)) for doc in results]) # Safer access
+        elif isinstance(results, dict) and 'answer' in results:
+             search_results_str = results['answer']
+        elif isinstance(results, dict) and 'result' in results: # Another common format
+             search_results_str = results['result']
+        else:
+             search_results_str = str(results)
+
+        print(f"--- Search Results Found ---")
+        return {"search_results": search_results_str}
+    except Exception as e:
+        print(f"--- Search Failed: {e} ---")
+        # Provide a more informative error message if possible
+        error_details = str(e)
+        return {"search_results": f"Search failed or timed out. Details: {error_details}"}
+
+
+def generate_itinerary(state: AgentState) -> Dict:
+    """Generates the final itinerary using the LLM, interpreting flexible user inputs."""
+    user_info = state['user_info']
+    search_results = state.get('search_results', "No search results available.") # Provide default
+    print("--- Generating Itinerary ---")
+
+    # --- Enhanced Prompt ---
+    itinerary_prompt = f"""
+    You are an expert travel planner. Create a detailed and engaging travel itinerary based on the following user preferences:
+
+    **User Preferences:**
+    - Destination: {user_info.get('destination', 'Not specified')}
+    - Duration: {user_info.get('duration', 'Not specified')} days
+    - Budget Description: '{user_info.get('budget', 'Not specified')}'
+    - Preferred Activities Description: '{user_info.get('activities', 'Not specified')}'
+    - Preferred Accommodation Description: '{user_info.get('accommodation', 'Not specified')}'
+
+    **Your Task:**
+    1.  **Interpret Preferences:** Carefully interpret the user's descriptions for budget, activities, and accommodation.
+        * If the budget is descriptive (e.g., 'moderate', 'budget-friendly', 'a bit flexible', 'around $X'), tailor suggestions to match that level for the specific destination. Avoid extreme high-cost or only free options unless explicitly requested. 'Moderate' usually implies a balance of value, comfort, and experiences.
+        * If activities are described generally (e.g., 'mix of famous and offbeat', 'relaxing', 'cultural immersion', 'adventure'), create an itinerary that reflects this. A 'mix' should include popular landmarks and hidden gems. 'Relaxing' should include downtime.
+        * Interpret accommodation descriptions (e.g., 'mid-range', 'cheap but clean', 'boutique hotel') based on typical offerings at the destination.
+    2.  **Use Search Results:** Incorporate relevant and specific suggestions from the search results below, but *only* if they align with the interpreted user preferences. Do not blindly copy search results.
+    3.  **Create a Coherent Plan:** Structure the itinerary logically, often day-by-day. Include suggestions for specific activities, potential dining spots (matching budget), and estimated timings where appropriate.
+    4.  **Engaging Tone:** Present the itinerary in an exciting and appealing way.
+
+    **Supporting Search Results:**
+    ```
+    {search_results}
+    ```
+
+    **Generate the Itinerary:**
+    """
+    # --- End of Enhanced Prompt ---
+
+    try:
+        response = llm.invoke(itinerary_prompt)
+        # Ensure content extraction handles potential variations
+        itinerary_content = getattr(response, 'content', str(response))
+
+        print("--- Itinerary Generated ---")
+
+        final_message = AIMessage(content=f"{ITINERARY_READY_MESSAGE}\n\n{itinerary_content}")
+        # Ensure messages list exists and append correctly
+        updated_messages = list(state.get('messages', [])) + [final_message]
+
+        return {"itinerary": itinerary_content, "messages": updated_messages}
+    except Exception as e:
+        print(f"--- Itinerary Generation Failed: {e} ---")
+        error_details = str(e)
+        error_message = AIMessage(content=f"Sorry, I encountered an error while generating the itinerary. Details: {error_details}")
+        # Ensure messages list exists and append correctly
+        updated_messages = list(state.get('messages', [])) + [error_message]
+        return {"itinerary": None, "messages": updated_messages}
+
+
+
+def should_ask_question_or_search(state: AgentState) -> str:
+    """Determines the next step based on whether all info is collected."""
+    missing_fields = state.get('missing_fields', [])
+    if not missing_fields:
+        print("--- Condition: All info gathered, proceed to search ---")
+        return "run_search"
+    else:
+        print("--- Condition: More info needed, ask next question ---")
+        return "ask_next_question"
+
+# --- Build the Graph ---
+
+graph_builder = StateGraph(AgentState)
+
+# Define nodes
+graph_builder.add_node("process_user_input", process_user_input)
+graph_builder.add_node("ask_next_question", ask_next_question)
+graph_builder.add_node("run_search", run_search) # Uses the updated function
+graph_builder.add_node("generate_itinerary", generate_itinerary)
+
+# Define edges
+graph_builder.set_entry_point("process_user_input")
+graph_builder.add_conditional_edges(
+    "process_user_input",
+    should_ask_question_or_search,
+    {"ask_next_question": "ask_next_question", "run_search": "run_search"}
+)
+graph_builder.add_edge("ask_next_question", END)
+graph_builder.add_edge("run_search", "generate_itinerary")
+graph_builder.add_edge("generate_itinerary", END)
+
+# Compile the graph
+travel_agent_app = graph_builder.compile()
+
+# --- Gradio Interface ---
+
+# Function to handle the conversation logic with LangGraph state
+def handle_user_message(user_input: str, history: List[List[str | None]], current_state_dict: Optional[dict]) -> tuple:
+    user_input_cleaned = user_input.strip().lower()
+
+    # Initialize state if it doesn't exist (first interaction)
+    if current_state_dict is None:
+        current_state_dict = {
+            "messages": [AIMessage(content=INITIAL_MESSAGE)],
+            "user_info": {},
+            "missing_fields": [], # Will be populated if user types START
+            "search_results": None,
+            "itinerary": None,
+        }
+
+    # Handle START command
+    if user_input_cleaned == "start" and not current_state_dict.get("missing_fields"): # Only start if not already started
+        print("--- Received START ---")
+        current_state_dict["missing_fields"] = list(ORDERED_FIELDS) # Initialize missing fields
+        current_state_dict["user_info"] = {} # Reset user info
+        current_state_dict["messages"] = [AIMessage(content=START_CONFIRMATION)] # type: ignore
+        # No graph execution needed yet, just update state and return the first question
+        history.append([None, START_CONFIRMATION]) # Gradio format needs None for user message here
+
+    # Handle case where user types something other than START initially
+    elif not current_state_dict.get("missing_fields") and user_input_cleaned != "start":
+         print("--- Waiting for START ---")
+         current_state_dict["messages"] = current_state_dict["messages"] + [ # type: ignore
+              HumanMessage(content=user_input),
+              AIMessage(content=INVALID_START_MESSAGE)
+         ]
+         history.append([user_input, INVALID_START_MESSAGE])
+
+
+    # Handle user responses after START
+    elif current_state_dict.get("missing_fields") or current_state_dict.get("itinerary"): # Process if questions pending or itinerary just generated
+        print(f"--- User Input: {user_input} ---")
+         # Prevent processing if itinerary was just generated and user typed something else
+        if current_state_dict.get("itinerary") and not current_state_dict.get("missing_fields"):
+             print("--- Itinerary already generated, waiting for START OVER ---")
+             # Optionally add a message like "Type START OVER to begin again."
+             history.append([user_input, "Itinerary generated. Please click 'Start Over' to plan a new trip."])
+             # Keep state as is, just update history
+             return history, current_state_dict, ""
+
+
+        # Add user message to state's messages list
+        current_messages = current_state_dict.get("messages", [])
+        current_messages.append(HumanMessage(content=user_input))
+        current_state_dict["messages"] = current_messages
+
+        # Invoke the graph
+        print("--- Invoking Graph ---")
+        # Ensure state keys match AgentState before invoking
+        graph_input_state = AgentState(
+             messages=current_state_dict.get("messages", []),
+             user_info=current_state_dict.get("user_info", {}),
+             missing_fields=current_state_dict.get("missing_fields", []),
+             search_results=current_state_dict.get("search_results"),
+             itinerary=current_state_dict.get("itinerary"),
+        )
+        # Use stream or invoke. Invoke is simpler for this request/response cycle.
+        final_state = travel_agent_app.invoke(graph_input_state)
+        print("--- Graph Execution Complete ---")
+
+        # Update the state dictionary from the graph's final state
+        current_state_dict.update(final_state)
+
+        # Update Gradio history
+        # The graph adds the AI response(s) to state['messages']
+        # Get the *last* AI message added by the graph
+        ai_response = final_state['messages'][-1].content if final_state['messages'] and isinstance(final_state['messages'][-1], AIMessage) else "Error: No response."
+        history.append([user_input, ai_response])
+
+
+    # Handle unexpected state (fallback)
+    else:
+        print("--- Unexpected State - Resetting ---")
+        history.append([user_input, "Something went wrong. Please type START to begin."])
+        current_state_dict = { # Reset state
+            "messages": [AIMessage(content=INITIAL_MESSAGE)],
+            "user_info": {},
+            "missing_fields": [],
+            "search_results": None,
+            "itinerary": None,
+        }
+
+
+    # Return updated history, the persistent state dictionary, and clear the input box
+    return history, current_state_dict, ""
+
+# Function to reset the state (Start Over button)
+def start_over() -> tuple:
+    print("--- Starting Over ---")
+    initial_state = {
+        "messages": [AIMessage(content=RESTART_MESSAGE)],
+        "user_info": {},
+        "missing_fields": [],
+        "search_results": None,
+        "itinerary": None,
+    }
+    # Gradio history format: List of [user_msg, assistant_msg] pairs
+    initial_history = [[None, RESTART_MESSAGE]]
+    return initial_history, initial_state, "" # Return history, state, clear input
+
+
+# --- Gradio UI Definition ---
+with gr.Blocks(theme=gr.themes.Soft()) as app:
+    gr.Markdown("# 🌍 AI-Powered Travel Itinerary Generator (LangGraph Version)")
+    gr.Markdown("Type `START` to begin planning your trip. Answer the questions, and I'll generate a personalized itinerary for you!")
+
+    # Store the LangGraph state between interactions
+    agent_state = gr.State(value=None) # Initialize state as None
+
+    # Chat interface
+    chatbot = gr.Chatbot(
+        label="Travel Bot",
+        bubble_full_width=False,
+        value=[[None, INITIAL_MESSAGE]] # Initial message
+        )
+    user_input = gr.Textbox(label="Your Message", placeholder="Type here...", scale=3)
+    submit_btn = gr.Button("Send", scale=1)
+    start_over_btn = gr.Button("Start Over", scale=1)
+
+    # Button and Textbox actions
+    submit_btn.click(
+        fn=handle_user_message,
+        inputs=[user_input, chatbot, agent_state],
+        outputs=[chatbot, agent_state, user_input] # Update chatbot, state, and clear input
+    )
+    user_input.submit( # Allow Enter key submission
+         fn=handle_user_message,
+        inputs=[user_input, chatbot, agent_state],
+        outputs=[chatbot, agent_state, user_input]
+    )
+    start_over_btn.click(
+        fn=start_over,
+        inputs=[],
+        outputs=[chatbot, agent_state, user_input] # Reset chatbot, state, and clear input
+    )
+
+
+# --- Run the App ---
+if __name__ == "__main__":
+    app.launch(debug=True)# Debug=True provides more logs