app.py

import os
import gradio as gr
from gradio import ChatMessage
import requests # Although not used in dummy tools, keep if future tools need it
from typing import Dict, List, Generator, Sequence
from langchain_core.messages import HumanMessage, BaseMessage
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import MemorySaver
from langgraph.prebuilt import create_react_agent
import logging

# Configure logging for better debugging on Spaces
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# --- Tool Definitions ---
@tool
def get_lat_lng(location_description: str) -> dict[str, float]:
    """Get the latitude and longitude of a location description (e.g., 'Paris', 'Tokyo, Japan')."""
    # In a real app, you'd call a geocoding API here.
    logger.info(f"Tool 'get_lat_lng' called with location: {location_description}")
    # Dummy response for demonstration
    if "london" in location_description.lower():
        return {"lat": 51.5074, "lng": -0.1278}
    elif "tokyo" in location_description.lower():
         return {"lat": 35.6895, "lng": 139.6917}
    elif "paris" in location_description.lower():
         return {"lat": 48.8566, "lng": 2.3522}
    elif "new york" in location_description.lower():
         return {"lat": 40.7128, "lng": -74.0060}
    else:
        # Default dummy response
        return {"lat": 51.1, "lng": -0.1}

@tool
def get_weather(lat: float, lng: float) -> dict[str, str]:
    """Get the current weather conditions at a specific latitude and longitude."""
    # In a real app, you'd call a weather API (e.g., OpenWeatherMap) here.
    logger.info(f"Tool 'get_weather' called with lat: {lat}, lng: {lng}")
    # Dummy response based on latitude for variety
    if 40 < lat < 50: # Approx Paris/New York
        return {"temperature": "18°C", "description": "Cloudy"}
    elif lat > 50: # Approx London
        return {"temperature": "15°C", "description": "Rainy"}
    else: # Approx Tokyo / Default
        return {"temperature": "25°C", "description": "Sunny"}

# --- Agent and Streaming Logic ---
def initialize_agent():
    """Initializes the LangChain agent."""
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        logger.error("OPENAI_API_KEY environment variable not set.")
        # Option 1: Raise an error to stop the app
        # raise ValueError("OpenAI API Key not found. Please set it in the Space secrets.")
        # Option 2: Return None and handle it in the stream function
        return None

    try:
        llm = ChatOpenAI(temperature=0, model="gpt-4", openai_api_key=api_key)
        # Note: MemorySaver() is in-memory. State will be lost on space restarts/sleeps.
        # For persistent memory across sessions/restarts, you'd need a persistent checkpointer (e.g., using Redis, SQL).
        memory = MemorySaver()
        tools = [get_lat_lng, get_weather]
        agent_executor = create_react_agent(llm, tools, checkpointer=memory)
        logger.info("LangChain agent initialized successfully.")
        return agent_executor
    except Exception as e:
        logger.error(f"Failed to initialize LangChain agent: {e}", exc_info=True)
        return None

# Initialize agent once when the script starts
agent_executor = initialize_agent()

# Define the streaming function for Gradio ChatInterface
def stream_from_agent(message: str, history: List[List[str]]) -> Generator[Sequence[ChatMessage], None, None]:
    """
    Processes user messages through the LangChain agent, yielding intermediate steps.

    Args:
        message: The user's input message.
        history: The conversation history provided by Gradio (list of [user, assistant] pairs).

    Yields:
        A sequence of Gradio ChatMessage objects representing the agent's thoughts and actions.
    """
    global agent_executor # Use the globally initialized agent

    if agent_executor is None:
        error_msg = "Agent initialization failed. Please check the logs and ensure the OPENAI_API_KEY secret is set correctly."
        yield [ChatMessage(role="assistant", content=error_msg)]
        return

    logger.info(f"Received message: {message}")
    logger.info(f"History: {history}")

    # Convert Gradio history to LangChain message format
    # Note: create_react_agent expects a list of BaseMessages under the "messages" key.
    # It typically works best with a single HumanMessage as input per turn for the ReAct loop.
    # We will use the memory checkpointer to handle history persistence within the agent's context.
    langchain_message = HumanMessage(content=message)

    messages_to_display: List[ChatMessage] = []
    final_response_content = ""

    try:
        # Note: Using a fixed thread_id means all users share the same memory state if MemorySaver is used.
        # For isolated user sessions, you'd need a mechanism to generate/retrieve unique thread_ids per user/session.
        # This often requires integrating with Gradio's state or session management.
        # For simplicity here, we use a fixed ID as in the original code.
        thread_id = "shared_weather_thread_123"
        config = {"configurable": {"thread_id": thread_id}}

        # Stream the agent's execution steps
        for chunk in agent_executor.stream({"messages": [langchain_message]}, config=config):
            logger.debug(f"Agent chunk received: {chunk}") # Use debug level for verbose chunk logging

            # Check for Agent Actions (Tool Calls)
            if agent_action := chunk.get("agent"):
                 # Often the agent's rationale or decision to use a tool might be here
                 # Depending on the specific agent type, you might parse agent_action differently
                 if agent_action.get("messages"):
                      for msg in agent_action["messages"]:
                           if hasattr(msg, 'tool_calls') and msg.tool_calls:
                                for tool_call in msg.tool_calls:
                                    # Display the tool call intention
                                    tool_msg = ChatMessage(
                                        role="assistant", # Show tool usage as assistant action
                                        content=f"Parameters: `{tool_call['args']}`",
                                        metadata={
                                            "title": f"🛠️ Calling Tool: `{tool_call['name']}`",
                                            "tool_call_id": tool_call["id"], # Store ID to match response
                                        }
                                    )
                                    messages_to_display.append(tool_msg)
                                    yield messages_to_display
                           # Capture potential intermediate reasoning if available (depends on agent/LLM)
                           elif hasattr(msg, 'content') and isinstance(msg.content, str) and msg.content:
                               # Avoid displaying the *final* answer prematurely if it appears mid-stream
                               # The final answer is usually in the last chunk's 'agent' message list
                               pass # We'll handle the final answer specifically later


            # Check for Tool Execution Results
            if tool_chunk := chunk.get("tools"):
                 if tool_chunk.get("messages"):
                    for tool_response in tool_chunk["messages"]:
                        # Find the corresponding tool call message to update it
                        found = False
                        for i, msg in enumerate(messages_to_display):
                            if msg.metadata and msg.metadata.get("tool_call_id") == tool_response.tool_call_id:
                                # Update the existing tool message with the result
                                updated_content = msg.content + f"\nResult: `{tool_response.content}`"
                                messages_to_display[i] = ChatMessage(
                                    role=msg.role,
                                    content=updated_content,
                                    metadata=msg.metadata # Keep original metadata
                                )
                                found = True
                                break
                        if found:
                            yield messages_to_display
                        else:
                             # If matching call not found (shouldn't happen often), display separately
                             tool_result_msg = ChatMessage(
                                 role="tool", # Or keep as assistant? 'tool' role might not render well by default
                                 content=f"Tool Result (`{tool_response.tool_call_id}`): `{tool_response.content}`"
                             )
                             messages_to_display.append(tool_result_msg)
                             yield messages_to_display


            # Check for the Final Agent Response
            # The final answer is typically the last message in the 'agent' chunk's list
            if agent_final := chunk.get("agent"):
                if agent_final.get("messages"):
                    last_message = agent_final["messages"][-1]
                    # Ensure it's the final response (often not a tool call)
                    if hasattr(last_message, 'content') and not (hasattr(last_message, 'tool_calls') and last_message.tool_calls):
                        final_response_content = last_message.content


        # After the loop, ensure the final response is added if it hasn't been implicitly handled
        if final_response_content:
            # Check if the last displayed message is already the final response
            is_already_displayed = False
            if messages_to_display:
                last_displayed = messages_to_display[-1]
                # Simple check: if last displayed message has no tool metadata and content matches
                if not (last_displayed.metadata and "tool_call_id" in last_displayed.metadata) and last_displayed.content == final_response_content:
                     is_already_displayed = True

            if not is_already_displayed:
                final_msg = ChatMessage(role="assistant", content=final_response_content)
                messages_to_display.append(final_msg)
                yield messages_to_display
        elif not messages_to_display:
            # Handle cases where the agent might not produce a final response (e.g., errors)
             yield [ChatMessage(role="assistant", content="Sorry, I couldn't process that request.")]


    except Exception as e:
        logger.error(f"Error during agent stream: {e}", exc_info=True)
        error_message = f"An error occurred: {e}"
        yield [ChatMessage(role="assistant", content=error_message)]


# --- Gradio Interface Definition ---
# Use gr.ChatInterface with type="messages" for full ChatMessage object support
demo = gr.ChatInterface(
    fn=stream_from_agent,
    chatbot=gr.Chatbot(
        bubble_full_width=False,
        show_copy_button=True,
        render=False # Render manually for better control if needed, but False is fine here
        ),
    input_components=[gr.Textbox(label="Ask the weather assistant")], # Customize input textbox
    # `type="messages"` passes message/history using gr.ChatMessage objects (needed for metadata)
    # However, ChatInterface's standard history format is List[List[str]].
    # Let's stick to the standard fn signature for ChatInterface if possible
    # and convert history inside the function if needed.
    # Reverting fn signature slightly based on typical ChatInterface usage.
    # If type="messages" is used, fn signature might expect different types.
    # Sticking to standard List[List[str]] history for compatibility.
    # Let's adjust the stream_from_agent function signature slightly if needed.
    # **Correction**: `gr.ChatInterface` *does* handle the `List[List[str]]` history format even when yielding `ChatMessage`.
    # The function signature `(message: str, history: List[List[str]])` is correct.

    title="🌤️ Weather Assistant with LangGraph ReAct Agent",
    description="Ask about the weather anywhere! Watch the agent think step-by-step as it uses tools.",
    examples=[
        ["What's the weather like in Tokyo?"],
        ["Is it sunny in Paris right now?"],
        ["Should I bring an umbrella in New York today?"]
    ],
    cache_examples=False, # Disable caching for dynamic examples if needed
    theme="soft", # Optional: Apply a theme
    retry_btn=None, # Disable retry button if stream handles errors
    undo_btn="Delete Previous", # Customize undo button text
    clear_btn="Clear Conversation", # Customize clear button text
)

# --- Launch the App ---
if __name__ == "__main__":
    # Launch the Gradio app
    # share=False is default and recommended for Spaces
    # debug=True can be helpful during development but disable for production
    # server_name="0.0.0.0" allows access within the Space's network
    demo.launch(server_name="0.0.0.0", server_port=7860)