gradio_app.py

"""
Gradio Chat UI for RAG Agent
"""

import os
import sys
import gradio as gr
import vertexai
from vertexai.agent_engines._agent_engines import AgentEngine
from google.cloud import aiplatform
from dotenv import load_dotenv

# Add the project root to the path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# Load environment variables from rag_agent/.env
env_path = os.path.join(os.path.dirname(__file__), "rag_agent", ".env")
load_dotenv(env_path)

PROJECT_ID = os.environ.get("GOOGLE_CLOUD_PROJECT")
LOCATION = os.environ.get("GOOGLE_CLOUD_LOCATION", "us-central1")

# Initialize Vertex AI
vertexai.init(project=PROJECT_ID, location=LOCATION)


def list_available_agents():
    """List all available agents from Agent Engine."""
    try:
        client = aiplatform.gapic.AgentEnginesServiceClient(
            client_options={"api_endpoint": f"{LOCATION}-aiplatform.googleapis.com"}
        )
        parent = f"projects/{PROJECT_ID}/locations/{LOCATION}"
        
        agents = []
        for agent in client.list_agent_engines(parent=parent):
            agent_info = {
                "name": agent.name,
                "display_name": agent.display_name,
            }
            agents.append(agent_info)
        
        return agents
    except Exception as e:
        print(f"Error listing agents: {e}")
        return []


def get_agent_names():
    """Get list of agent display names for dropdown."""
    agents = list_available_agents()
    if not agents:
        return ["No agents found"]
    return [agent["display_name"] for agent in agents]


def get_agent_by_display_name(display_name):
    """Get agent resource name by display name."""
    agents = list_available_agents()
    for agent in agents:
        if agent["display_name"] == display_name:
            return agent["name"]
    return None


def chat_with_agent(message, history, agent_name):
    """Send message to selected agent and get response."""
    if not message:
        return history
    
    if agent_name == "No agents found":
        history.append((message, "Error: No agents available. Please deploy an agent first."))
        return history
    
    try:
        # Get the agent resource name
        agent_resource_name = get_agent_by_display_name(agent_name)
        if not agent_resource_name:
            history.append((message, f"Error: Could not find agent '{agent_name}'"))
            return history
        
        # Create AgentEngine instance
        agent_engine = AgentEngine(name=agent_resource_name)
        
        # Send query to agent
        response = agent_engine.query(query=message)
        
        # Extract response text
        if hasattr(response, 'text'):
            response_text = response.text
        elif isinstance(response, dict) and 'text' in response:
            response_text = response['text']
        else:
            response_text = str(response)
        
        history.append((message, response_text))
        
    except Exception as e:
        error_msg = f"Error communicating with agent: {str(e)}"
        history.append((message, error_msg))
    
    return history


def refresh_agents():
    """Refresh the list of available agents."""
    return gr.Dropdown(choices=get_agent_names(), value=get_agent_names()[0])


# Create Gradio interface
with gr.Blocks(title="RAG Agent Chat", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🤖 RAG Agent Chat Interface")
    gr.Markdown("Select an agent and start chatting!")
    
    with gr.Row():
        with gr.Column(scale=3):
            agent_dropdown = gr.Dropdown(
                choices=get_agent_names(),
                value=get_agent_names()[0] if get_agent_names() else "No agents found",
                label="Select Agent",
                interactive=True
            )
        with gr.Column(scale=1):
            refresh_btn = gr.Button("🔄 Refresh Agents", size="sm")
    
    chatbot = gr.Chatbot(
        label="Chat History",
        height=500,
        show_copy_button=True
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Your Message",
            placeholder="Type your message here...",
            scale=4,
            lines=2
        )
        submit_btn = gr.Button("Send", variant="primary", scale=1)
    
    clear_btn = gr.Button("Clear Chat")
    
    gr.Markdown("""
    ### 📝 Agent Capabilities
    - **Query Documents**: Ask questions and retrieve information from document corpora
    - **List Corpora**: See all available document collections
    - **Create Corpus**: Create new document collections
    - **Add Data**: Add documents to existing corpora
    - **Get Corpus Info**: View detailed corpus information
    - **Delete Document/Corpus**: Remove documents or entire corpora
    """)
    
    # Event handlers
    def submit_message(message, history, agent_name):
        history = chat_with_agent(message, history, agent_name)
        return "", history
    
    submit_btn.click(
        submit_message,
        inputs=[msg, chatbot, agent_dropdown],
        outputs=[msg, chatbot]
    )
    
    msg.submit(
        submit_message,
        inputs=[msg, chatbot, agent_dropdown],
        outputs=[msg, chatbot]
    )
    
    clear_btn.click(lambda: [], outputs=chatbot)
    
    refresh_btn.click(
        refresh_agents,
        outputs=agent_dropdown
    )


if __name__ == "__main__":
    # Check if required environment variables are set
    if not PROJECT_ID:
        print("⚠️  Warning: GOOGLE_CLOUD_PROJECT environment variable not set")
        print("Please set it in your .env file or environment")
    
    print(f"🚀 Starting Gradio app...")
    print(f"📍 Project: {PROJECT_ID}")
    print(f"📍 Location: {LOCATION}")
    print(f"🔍 Found {len(get_agent_names())} agent(s)")
    
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )