main.py

import streamlit as st
import requests
import json
import uuid
from typing import Dict, Any
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

# Configuration
API_BASE_URL = os.getenv("API_BASE_URL", "https://pixerse-pixerse.hf.space")  # Default to the env value
CHAT_ENDPOINT = f"{API_BASE_URL}/api/mcp/chat"

def call_chat_api(query: str, user_id: str) -> Dict[str, Any]:
    """Call the FastAPI chat endpoint"""
    payload = {
        "query": query,
        "user_id": user_id
    }
    
    try:
        response = requests.post(
            CHAT_ENDPOINT,
            json=payload,
            headers={"Content-Type": "application/json"},
            timeout=30
        )
        
        if response.status_code == 200:
            return {
                "success": True,
                "data": response.json()
            }
        else:
            return {
                "success": False,
                "error": f"API Error: {response.status_code} - {response.text}"
            }
    except requests.exceptions.RequestException as e:
        return {
            "success": False,
            "error": f"Connection Error: {str(e)}"
        }

def main():
    st.set_page_config(
        page_title="PiXerse Chatbot",
        page_icon="🤖",
        layout="wide"
    )
    
    st.title("🤖 PiXerse AI Chatbot")
    st.markdown("---")
    
    # Initialize session state
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    if "user_id" not in st.session_state:
        st.session_state.user_id = str(uuid.uuid4())
    
    # Sidebar for configuration
    with st.sidebar:
        st.header("⚙️ Configuration")
        
        # Display user ID
        st.text_input(
            "User ID", 
            value=st.session_state.user_id, 
            disabled=True,
            help="Your unique session identifier"
        )
        
        # API Status check
        st.subheader("🔍 API Status")
        if st.button("Check API Status"):
            try:
                health_response = requests.get(f"{API_BASE_URL}/docs", timeout=5)
                if health_response.status_code == 200:
                    st.success("✅ API is running")
                else:
                    st.error("❌ API is not responding properly")
            except:
                st.error("❌ Cannot connect to API")
        
        # Clear chat button
        if st.button("🗑️ Clear Chat History"):
            st.session_state.messages = []
            st.rerun()
    
    # Main chat interface
    st.subheader("💬 Chat Interface")
    
    # Display chat messages
    chat_container = st.container()
    with chat_container:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
                
                # Show additional info for assistant messages
                if message["role"] == "assistant" and "metadata" in message:
                    with st.expander("📊 Response Details"):
                        metadata = message["metadata"]
                        col_meta1, col_meta2 = st.columns(2)
                        
                        with col_meta1:
                            st.metric("Token Usage", metadata.get("token_usage", "N/A"))
                            
                        with col_meta2:
                            st.metric("Tools Used", len(metadata.get("tools_used", [])))
                        
                        if metadata.get("tools_used"):
                            st.write("**Tools Used:**")
                            for tool in metadata["tools_used"]:
                                st.code(tool, language="text")
                        
                        if metadata.get("tools_response"):
                            st.write("**Tool Responses:**")
                            for i, response in enumerate(metadata["tools_response"]):
                                st.json(response)

    # Chat input (outside columns to avoid Streamlit API error)
    if prompt := st.chat_input("Type your message here..."):
        # Add user message to chat
        st.session_state.messages.append({
            "role": "user", 
            "content": prompt
        })
        
        # Rerun to display new message immediately
        st.rerun()

    # Process last message if it's from user and no assistant response yet
    if (st.session_state.messages and 
        st.session_state.messages[-1]["role"] == "user" and
        (len(st.session_state.messages) == 1 or 
         st.session_state.messages[-2]["role"] == "assistant")):
        
        user_message = st.session_state.messages[-1]["content"]
        
        # Show processing message
        with st.chat_message("assistant"):
            with st.spinner("🤔 Thinking..."):
                response = call_chat_api(user_message, st.session_state.user_id)
            
            if response["success"]:
                data = response["data"]
                assistant_message = data["response"]
                
                # Display response
                st.markdown(assistant_message)
                
                # Add assistant message to chat with metadata
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": assistant_message,
                    "metadata": {
                        "token_usage": data.get("token_usage", 0),
                        "tools_used": data.get("tools_used", []),
                        "tools_response": data.get("tools_response", [])
                    }
                })
                
            else:
                error_message = f"❌ Error: {response['error']}"
                st.error(error_message)
                
                # Add error message to chat
                st.session_state.messages.append({
                    "role": "assistant",
                    "content": error_message
                })
            
            # Rerun to update the display
            st.rerun()

if __name__ == "__main__":
    # Run Streamlit UI
    main()