app.py

import os
import sys
import json
import gradio as gr
import asyncio
import subprocess
import threading
import requests
import time
import random
import string
from typing import Dict, List, Optional
import uuid

from hf_utils import HuggingFaceSpaceHelper

# Initialize helpers
hf_helper = HuggingFaceSpaceHelper()

# Install required packages for HF Spaces if needed
if hf_helper.is_in_space:
    hf_helper.install_dependencies([
        "pymongo", "python-dotenv", "gradio", "requests"
    ])

# Import after potentially installing dependencies
try:
    from db_helper import MongoDBHelper
    db = MongoDBHelper(hf_helper.get_mongodb_uri())
except Exception as e:
    print(f"Warning: MongoDB connection failed: {e}")
    print("API key management will not work!")
    db = None

# Function to start the API in the background
def start_api_server():
    """Start the API server in a separate process"""
    api_process = subprocess.Popen(
        [sys.executable, "pyscout_api.py"],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )
    return api_process

# Determine the API base URL based on environment
# In Docker, use the environment variable if set
API_BASE_URL = os.getenv("API_BASE_URL", f"http://{hf_helper.get_hostname()}:8000")

# Function to check API health
def check_api_health():
    """Check if the API is running and healthy"""
    try:
        response = requests.get(f"{API_BASE_URL}/health", timeout=5)
        if response.status_code == 200:
            return response.json()
        return {"status": "error", "code": response.status_code}
    except Exception as e:
        return {"status": "error", "message": str(e)}

# Utility functions for the Gradio UI
def generate_api_key(email: str, name: str, organization: str = ""):
    """Generate a new API key for a user"""
    if not email or not name:
        return "Error: Email and name are required"
    
    if not db:
        # Generate a dummy key for demonstration if MongoDB is not connected
        dummy_key = f"PyScoutAI-demo-{uuid.uuid4().hex[:16]}"
        return f"Generated demo API key (MongoDB not connected):\n\n{dummy_key}\n\nThis key won't be validated in actual requests."
    
    try:
        # Create a user ID from email
        user_id = email.strip().lower()
        
        # Generate the API key
        api_key = db.generate_api_key(user_id, name)
        
        return f"API key generated successfully:\n\n{api_key}\n\nStore this key safely. It won't be displayed again."
        
    except Exception as e:
        return f"Error generating API key: {str(e)}"

def list_user_api_keys(email: str):
    """List all API keys for a user"""
    if not email:
        return "Error: Email is required"
    
    if not db:
        return "Error: MongoDB not connected. Cannot list API keys."
    
    try:
        # Get user ID from email
        user_id = email.strip().lower()
        
        # Get all API keys for the user
        keys = db.get_user_api_keys(user_id)
        
        if not keys:
            return f"No API keys found for {email}"
        
        result = f"Found {len(keys)} API key(s) for {email}:\n\n"
        for i, key in enumerate(keys):
            status = "Active" if key.get("is_active", False) else "Revoked"
            last_used = key.get("last_used", "Never")
            if isinstance(last_used, str):
                last_used_str = last_used
            else:
                last_used_str = last_used.strftime("%Y-%m-%d %H:%M:%S") if last_used else "Never"
                
            result += f"{i+1}. {key.get('name', 'Unnamed')} - {key.get('key')}\n"
            result += f"   Status: {status}, Created: {key.get('created_at').strftime('%Y-%m-%d')}, "
            result += f"Last used: {last_used_str}\n\n"
        
        return result
        
    except Exception as e:
        return f"Error listing API keys: {str(e)}"

def revoke_api_key(api_key: str):
    """Revoke an API key"""
    if not api_key:
        return "Error: API key is required"
    
    if not db:
        return "Error: MongoDB not connected. Cannot revoke API key."
    
    try:
        if not api_key.startswith("PyScoutAI-"):
            return "Error: Invalid API key format. Keys should start with 'PyScoutAI-'."
            
        success = db.revoke_api_key(api_key)
        
        if success:
            return f"API key {api_key} revoked successfully"
        else:
            return f"API key {api_key} not found or already revoked"
            
    except Exception as e:
        return f"Error revoking API key: {str(e)}"

def test_api(api_key: str, prompt: str, model: str = "meta-llama/Llama-3.3-70B-Instruct-Turbo", temperature: float = 0.7):
    """Test the API with a simple chat completion request"""
    if not api_key:
        return "Error: API key is required"
    
    if not prompt:
        return "Error: Prompt is required"
    
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {api_key}"
    }
    
    data = {
        "model": model,
        "messages": [
            {"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        "temperature": temperature
    }
    
    try:
        start_time = time.time()
        response = requests.post(
            f"{API_BASE_URL}/v1/chat/completions",
            headers=headers,
            json=data,
            timeout=60
        )
        
        elapsed = time.time() - start_time
        
        if response.status_code == 200:
            result = response.json()
            content = result["choices"][0]["message"]["content"]
            tokens = result.get("usage", {}).get("total_tokens", "unknown")
            
            return f"Response (in {elapsed:.2f}s, {tokens} tokens):\n\n{content}"
        else:
            error_detail = "Unknown error"
            try:
                error_detail = response.json().get("detail", "Unknown error")
            except:
                error_detail = response.text
                
            return f"API Error (status {response.status_code}):\n{error_detail}"
            
    except requests.RequestException as e:
        return f"Request error: {str(e)}"
    except Exception as e:
        return f"Unexpected error: {str(e)}"

def list_models(api_key: str):
    """List available models from the API"""
    if not api_key:
        return "Error: API key is required"
    
    headers = {
        "Authorization": f"Bearer {api_key}"
    }
    
    try:
        response = requests.get(
            f"{API_BASE_URL}/v1/models",
            headers=headers,
            timeout=10
        )
        
        if response.status_code == 200:
            models = response.json()
            result = "Available models:\n\n"
            for i, model in enumerate(models.get("data", [])):
                result += f"{i+1}. {model.get('id')}\n"
            return result
        else:
            error_detail = "Unknown error"
            try:
                error_detail = response.json().get("detail", "Unknown error")
            except:
                error_detail = response.text
                
            return f"API Error (status {response.status_code}):\n{error_detail}"
            
    except Exception as e:
        return f"Error listing models: {str(e)}"

def create_ui():
    """Create the Gradio UI"""
    with gr.Blocks(title="PyScoutAI API Manager") as app:
        gr.Markdown("# PyScoutAI API Manager")
        gr.Markdown("Manage API keys and test the PyScoutAI API")
        
        # API Status
        with gr.Row():
            check_api_btn = gr.Button("Check API Status")
            api_status = gr.JSON(label="API Status")
            check_api_btn.click(check_api_health, outputs=[api_status])
        
        with gr.Tabs():
            # API Key Management Tab
            with gr.TabItem("Manage API Keys"):
                with gr.Tab("Generate API Key"):
                    email_input = gr.Textbox(label="Email")
                    name_input = gr.Textbox(label="Name")
                    org_input = gr.Textbox(label="Organization (optional)")
                    gen_key_btn = gr.Button("Generate API Key")
                    key_output = gr.Textbox(label="Generated Key", lines=5)
                    
                    gen_key_btn.click(
                        generate_api_key,
                        inputs=[email_input, name_input, org_input],
                        outputs=[key_output]
                    )
                
                with gr.Tab("List User Keys"):
                    email_list_input = gr.Textbox(label="Email")
                    list_keys_btn = gr.Button("List API Keys")
                    keys_output = gr.Textbox(label="User API Keys", lines=10)
                    
                    list_keys_btn.click(
                        list_user_api_keys,
                        inputs=[email_list_input],
                        outputs=[keys_output]
                    )
                
                with gr.Tab("Revoke API Key"):
                    key_revoke_input = gr.Textbox(label="API Key to Revoke")
                    revoke_btn = gr.Button("Revoke API Key")
                    revoke_output = gr.Textbox(label="Result")
                    
                    revoke_btn.click(
                        revoke_api_key,
                        inputs=[key_revoke_input],
                        outputs=[revoke_output]
                    )
            
            # API Testing Tab
            with gr.TabItem("Test API"):
                with gr.Row():
                    api_key_input = gr.Textbox(label="API Key")
                    model_input = gr.Dropdown(
                        choices=[
                            "meta-llama/Llama-3.3-70B-Instruct-Turbo",
                            "meta-llama/Meta-Llama-3.1-70B-Instruct",
                            "mistralai/Mistral-Small-24B-Instruct-2501",
                            "deepseek-ai/DeepSeek-V3"
                        ],
                        label="Model",
                        value="meta-llama/Llama-3.3-70B-Instruct-Turbo"
                    )
                    temperature_input = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.7,
                        step=0.1,
                        label="Temperature"
                    )
                
                prompt_input = gr.Textbox(label="Prompt", lines=3)
                test_btn = gr.Button("Send Request")
                list_models_btn = gr.Button("List Available Models")
                api_output = gr.Textbox(label="Response", lines=15)
                
                test_btn.click(
                    test_api,
                    inputs=[api_key_input, prompt_input, model_input, temperature_input],
                    outputs=[api_output]
                )
                
                list_models_btn.click(
                    list_models,
                    inputs=[api_key_input],
                    outputs=[api_output]
                )
    
    return app

def main():
    """Main entry point for the app"""
    # Check if API server is running
    api_health = check_api_health()
    
    if "status" in api_health and api_health["status"] == "error":
        print("API server doesn't seem to be running. Starting it...")
        # Start the API server if it's not already running
        api_process = start_api_server()