Update app.py

app.py

@@ -1,250 +1,108 @@
 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
+import requests
+from typing import Dict, List
+from langchain_core.messages import HumanMessage
 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 ---
+# Weather and location tools
 @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}
+    """Get the latitude and longitude of a location."""
+    return {"lat": 51.1, "lng": -0.1}  # London coordinates as dummy response
 
 @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
+    """Get the weather at a location."""
+    return {"temperature": "21°C", "description": "Sunny"}  # Dummy response
+        
+
+def stream_from_agent(message: str, history: List[Dict[str, str]]) -> gr.ChatMessage:
+    """Process messages through the LangChain agent with visible reasoning."""
+    
+    # Initialize the agent
+    llm = ChatOpenAI(temperature=0, model="gpt-4")
+    memory = MemorySaver()
+    tools = [get_lat_lng, get_weather]
+    agent_executor = create_react_agent(llm, tools, checkpointer=memory)
+    
+    # Add message to history
+    past_messages = [HumanMessage(content=message)]
+    for h in history:
+        if h["role"] == "user":
+            past_messages.append(HumanMessage(content=h["content"]))
+    
+    messages_to_display = []
+    final_response = None
+    
+    for chunk in agent_executor.stream(
+        {"messages": past_messages},
+        config={"configurable": {"thread_id": "abc123"}}
+    ):
+        # Handle agent's actions and tool usage
+        if chunk.get("agent"):
+            for msg in chunk["agent"]["messages"]:
+                if msg.content:
+                    final_response = msg.content
+                
+                # Handle tool calls
+                for tool_call in msg.tool_calls:
+                    tool_message = ChatMessage(
+                        content=f"Parameters: {tool_call['args']}",
+                        metadata={
+                            "title": f"🛠️ Using {tool_call['name']}",
+                            "id": tool_call["id"],
+                            "status": "pending",
+                        }
+                    )
+                    messages_to_display.append(tool_message)
+                    yield messages_to_display
+                    tool_message.metadata["status"] = "done"
+                    
+        # Handle tool responses
+        if chunk.get("tools"):
+            for tool_response in chunk["tools"]["messages"]:
+                # Find the corresponding tool message
+                for msg in messages_to_display:
+                    if msg.metadata.get("id") == tool_response.tool_call_id:
+                        msg.content += f"\nResult: {tool_response.content}"
+                        yield messages_to_display
+    
+    # Add the final response as a regular message
+    if final_response:
+        messages_to_display.append(ChatMessage(content=final_response))
+        yield messages_to_display
+
+# Create the Gradio interface
 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.",
+    type="messages",
+    title="🌤️ Weather Assistant",
+    description="Ask about the weather anywhere! Watch as I gather the information step by step.",
     examples=[
-        ["What's the weather like in Tokyo?"],
-        ["Is it sunny in Paris right now?"],
-        ["Should I bring an umbrella in New York today?"]
+        "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
+    example_icons=["https://cdn3.iconfinder.com/data/icons/landmark-outline/432/japan_tower_tokyo_landmark_travel_architecture_tourism_view-256.png", 
+                   "https://cdn2.iconfinder.com/data/icons/city-building-1/200/ArcdeTriomphe-256.png",
+                   "https://cdn2.iconfinder.com/data/icons/city-icons-for-offscreen-magazine/80/new-york-256.png"
+                  ],
+    save_history=True,
+    editable=True
+
 )
 
-# --- 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)
+    # Load environment variables
+    try:
+        from dotenv import load_dotenv
+        load_dotenv()
+    except ImportError:
+        pass
+
+    demo.launch(debug=True)