model_service_api.py

import aiohttp
import logging
from typing import List, Dict, Optional
from fastapi import FastAPI, HTTPException, Query, Path
from fastapi.responses import JSONResponse
from pydantic import BaseModel
from datetime import datetime, timedelta

# Setup logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__name__)

# FastAPI app
app = FastAPI(
    title="Free Models API",
    description="API to fetch and filter free AI models from OpenRouter",
    version="1.0.0"
)

# Pydantic models for response structure
class Architecture(BaseModel):
    modality: str
    input_modalities: List[str]
    output_modalities: List[str]
    tokenizer: str
    instruct_type: Optional[str]

class Pricing(BaseModel):
    prompt: str
    completion: str
    request: str
    image: str
    audio: str
    web_search: str
    internal_reasoning: str

class TopProvider(BaseModel):
    context_length: int
    max_completion_tokens: int
    is_moderated: bool

class Model(BaseModel):
    id: str
    canonical_slug: str
    hugging_face_id: Optional[str]
    name: str
    created: int
    description: str
    context_length: int
    architecture: Architecture
    pricing: Pricing
    top_provider: TopProvider
    per_request_limits: Optional[Dict]
    supported_parameters: List[str]

class FreeModelsResponse(BaseModel):
    success: bool
    message: str
    count: int
    models: List[Model]
    filtered_at: str

# Cache for storing API responses
cache = {
    "data": None,
    "timestamp": None,
    "ttl_minutes": 15  # Cache for 15 minutes
}

# Supported categories
SUPPORTED_CATEGORIES = [
    "chat", "reasoning", "vision", "coding", "lightweight", 
    "roleplay", "marketing", "programming", "image", "text"
]

class OpenRouterClient:
    """Client for interacting with OpenRouter API"""
    
    def __init__(self, api_key: Optional[str] = None):
        self.base_url = "https://openrouter.ai/api/v1"
        self.api_key = api_key
        self.session = None
    
    async def __aenter__(self):
        timeout = aiohttp.ClientTimeout(total=30)
        connector = aiohttp.TCPConnector(limit=10)
        
        headers = {
            "User-Agent": "Free-Models-API/1.0.0",
            "Content-Type": "application/json"
        }
        
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"
        
        self.session = aiohttp.ClientSession(
            timeout=timeout,
            connector=connector,
            headers=headers
        )
        return self
    
    async def __aexit__(self, exc_type, exc_val, exc_tb):
        if self.session:
            await self.session.close()
    
    async def get_models(self, category: Optional[str] = None) -> Dict:
        """Fetch all available models from OpenRouter"""
        url = f"{self.base_url}/models"
        params = {}
        
        if category:
            params["category"] = category
        
        try:
            logger.info(f"Fetching models from OpenRouter API{f' with category: {category}' if category else ''}...")
            async with self.session.get(url, params=params) as response:
                if response.status != 200:
                    logger.error(f"OpenRouter API returned status {response.status}")
                    raise Exception(f"API request failed with status {response.status}")
                
                data = await response.json()
                logger.info(f"Successfully fetched {len(data.get('data', []))} models")
                return data
        
        except Exception as e:
            logger.error(f"Error fetching models: {str(e)}")
            raise

def is_cache_valid() -> bool:
    """Check if cached data is still valid"""
    if not cache["timestamp"] or not cache["data"]:
        return False
    
    expiry_time = cache["timestamp"] + timedelta(minutes=cache["ttl_minutes"])
    return datetime.now() < expiry_time

def filter_free_models(models_data: Dict) -> List[Dict]:
    """Filter models where both prompt and completion pricing are '0'"""
    if not models_data.get("data"):
        return []
    
    free_models = []
    
    for model in models_data["data"]:
        pricing = model.get("pricing", {})
        prompt_price = pricing.get("prompt", "")
        completion_price = pricing.get("completion", "")
        
        # Check if both prompt and completion are "0" (free)
        if prompt_price == "0" and completion_price == "0":
            free_models.append(model)
            logger.debug(f"Found free model: {model.get('name', 'Unknown')}")
    
    logger.info(f"Found {len(free_models)} free models out of {len(models_data['data'])} total models")
    return free_models

async def get_free_models_data(category: Optional[str] = None) -> List[Dict]:
    """Get free models data with caching"""
    
    # Use cached data if valid
    if is_cache_valid() and not category:  # Don't use cache for category-specific requests
        logger.info("Using cached data")
        return filter_free_models(cache["data"])
    
    # Fetch fresh data
    async with OpenRouterClient() as client:
        try:
            models_data = await client.get_models(category=category)
            
            # Update cache only for general requests (no category)
            if not category:
                cache["data"] = models_data
                cache["timestamp"] = datetime.now()
                logger.info("Updated cache with fresh data")
            
            return filter_free_models(models_data)
        
        except Exception as e:
            # If we have cached data, use it as fallback
            if cache["data"] and not category:
                logger.warning(f"API request failed, using cached data: {str(e)}")
                return filter_free_models(cache["data"])
            else:
                raise HTTPException(
                    status_code=503,
                    detail=f"Failed to fetch models and no cached data available: {str(e)}"
                )

@app.get("/", response_model=Dict)
async def root():
    """Root endpoint with API information"""
    return {
        "message": "Free Models API - Filter free AI models from OpenRouter",
        "version": "1.0.0",
        "endpoints": {
            "free_models": "/api/free-models",
            "free_models_by_category": "/api/free-models?category=programming",
            "category_route": "/api/category/{category}/free-models",
            "free_model_names": "/api/free-models/names",
            "specific_model": "/api/free-models/{model_id}",
            "supported_categories": "/api/categories",
            "health": "/health"
        },
        "description": "This API fetches models from OpenRouter and filters for free models (prompt=0, completion=0)"
    }

@app.get("/api/categories")
async def get_supported_categories():
    """Get list of supported categories"""
    return {
        "success": True,
        "message": "List of supported categories",
        "categories": SUPPORTED_CATEGORIES,
        "usage": "Use these categories with /api/category/{category}/free-models"
    }

@app.get("/api/category/{category}/free-models", response_model=FreeModelsResponse)
async def get_free_models_by_category(
    category: str = Path(..., description="Category to filter by (e.g., chat, coding, vision)")
):
    """Get free models filtered by specific category"""
    
    # Validate category
    if category.lower() not in [cat.lower() for cat in SUPPORTED_CATEGORIES]:
        raise HTTPException(
            status_code=400,
            detail=f"Unsupported category '{category}'. Supported categories: {', '.join(SUPPORTED_CATEGORIES)}"
        )
    
    try:
        free_models = await get_free_models_data(category=category.lower())
        
        return FreeModelsResponse(
            success=True,
            message=f"Successfully retrieved {len(free_models)} free models in category '{category}'",
            count=len(free_models),
            models=free_models,
            filtered_at=datetime.now().isoformat()
        )
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in get_free_models_by_category: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/free-models", response_model=FreeModelsResponse)
async def get_free_models(
    category: Optional[str] = Query(None, description="Filter by category (e.g., 'programming', 'chat')")
):
    """Get all free AI models from OpenRouter"""
    try:
        free_models = await get_free_models_data(category=category)
        
        return FreeModelsResponse(
            success=True,
            message=f"Successfully retrieved {len(free_models)} free models" + (f" in category '{category}'" if category else ""),
            count=len(free_models),
            models=free_models,
            filtered_at=datetime.now().isoformat()
        )
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in get_free_models: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/free-models/names/{name}", response_model=Dict)
async def get_free_model_names(
    name: bool = False,
    category: Optional[str] = Query(None, description="Filter by category")
):
    """Get just the names and IDs of free models"""
    try:
        free_models = await get_free_models_data(category=category)
        
        if name:
            return {
                "success": True,
                "message": f"Retrieved {len(free_models)} free model names",
                "count": len(free_models),
                "models": [{"id": model["id"]} for model in free_models],
                "filtered_at": datetime.now().isoformat()
            }

        model_names = [
            {
                "id": model["id"],
                "name": model["name"],
                "input_modalities": model.get("architecture", {}).get("input_modalities"),
                "output_modalities": model.get("architecture", {}).get("output_modalities"),
                "per_request_limits": model.get("per_request_limits", {}),
                "description": model.get("description", "")[:100] + "..." if len(model.get("description", "")) > 100 else model.get("description", "")
            }
            for model in free_models
        ]
        
        return {
            "success": True,
            "message": f"Retrieved {len(model_names)} free model names",
            "count": len(model_names),
            "models": model_names,
            "filtered_at": datetime.now().isoformat()
        }
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in get_free_model_names: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/api/free-models/{model_id}")
async def get_free_model_by_id(model_id: str):
    """Get specific free model by ID"""
    try:
        free_models = await get_free_models_data()
        
        for model in free_models:
            if model["id"] == model_id:
                return {
                    "success": True,
                    "message": f"Found model {model_id}",
                    "model": model
                }
        
        raise HTTPException(
            status_code=404,
            detail=f"Free model with ID '{model_id}' not found"
        )
    
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error in get_free_model_by_id: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health_check():
    """Health check endpoint"""
    try:
        # Test a simple API call to check connectivity
        async with OpenRouterClient() as client:
            await client.get_models()
        
        return {
            "status": "healthy",
            "timestamp": datetime.now().isoformat(),
            "cache_status": "valid" if is_cache_valid() else "expired",
            "openrouter_api": "accessible"
        }
    
    except Exception as e:
        return JSONResponse(
            status_code=503,
            content={
                "status": "unhealthy",
                "timestamp": datetime.now().isoformat(),
                "error": str(e),
                "cache_status": "valid" if is_cache_valid() else "expired",
                "openrouter_api": "inaccessible"
            }
        )

@app.get("/ping")
async def ping():
    """Ping endpoint to check API status"""
    return {"status": "ok"}