main.py

"""
Devrayog AI - Main API Server
FastAPI-based API for Devrayog AI models
"""

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
from typing import Optional
import time
import os

app = FastAPI(
    title="Devrayog AI API",
    description="India's First Solo-Built AI - API Server",
    version="0.1.0"
)

# CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Request/Response models
class ChatRequest(BaseModel):
    message: str
    model: str = "devrayog-agni"
    max_tokens: int = 512
    temperature: float = 0.7
    system_prompt: Optional[str] = None

class ChatResponse(BaseModel):
    response: str
    model: str
    tokens_used: int
    time_taken: float

class ModelInfo(BaseModel):
    name: str
    description: str
    parameters: str
    status: str

# Routes
@app.get("/")
async def root():
    return {
        "name": "Devrayog AI",
        "version": "0.1.0",
        "status": "running",
        "message": "Welcome to Devrayog AI - India's First Solo-Built AI",
        "endpoints": {
            "chat": "/api/v1/chat",
            "models": "/api/v1/models",
            "health": "/health"
        }
    }

@app.get("/health")
async def health():
    return {"status": "healthy", "timestamp": time.time()}

@app.get("/api/v1/models")
async def list_models():
    return {
        "models": [
            {
                "id": "devrayog-agni",
                "name": "Devrayog Agni",
                "description": "Fast and sharp - optimized for quick responses",
                "parameters": "3.8B",
                "status": "coming_soon"
            },
            {
                "id": "devrayog-vayu",
                "name": "Devrayog Vayu",
                "description": "Creative and flowing - optimized for storytelling",
                "parameters": "9B",
                "status": "coming_soon"
            },
            {
                "id": "devrayog-akash",
                "name": "Devrayog Akash",
                "description": "Most powerful - optimized for complex reasoning",
                "parameters": "8B",
                "status": "coming_soon"
            }
        ]
    }

@app.post("/api/v1/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    start_time = time.time()
    
    # For now, return a placeholder response
    # Once models are deployed, this will call the actual model
    response = f"[Devrayog AI - {request.model}] Thank you for your message: '{request.message}'. This is a placeholder response. The actual model will be deployed soon on Hugging Face."
    
    time_taken = time.time() - start_time
    
    return ChatResponse(
        response=response,
        model=request.model,
        tokens_used=len(response.split()),
        time_taken=round(time_taken, 3)
    )

if __name__ == "__main__":
    import uvicorn
    port = int(os.environ.get("PORT", 8000))
    uvicorn.run(app, host="0.0.0.0", port=port)