app.py

import os
import time
import uuid
import torch
import logging
from typing import Dict, Any, Optional
from fastapi import FastAPI, Depends, Security, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse

import config
from schemas import (
    ChatCompletionRequest,
    ModelLoadRequest,
    TokenizeRequest,
    TokenizeResponse,
    ResetRequest,
    ResetResponse,
    BenchmarkQueryRequest,
    BenchmarkQueryResponse,
    ModelListResponse,
    ModelInfo,
    BenchmarkMetrics
)
from model_registry import PLAIN_MODELS, PAIRED_RIF_MODELS, validate_model_id
import model_service
import tokenization
import prompt_builder
import telemetry
import security
import errors

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("plain_baseline_api")

app = FastAPI(
    title="Kalpanā Plain Model Baselines API",
    description=(
        "## Plain LLM Baselines — No RIF Memory\n\n"
        "This API serves **standard open-source language models** (Llama 3, Qwen) "
        "**without** the Kalpanā Resonant Interference Field (RIF) memory layer.\n\n"
        "It exists as the **control group** for the "
        "[Kalpanā Multi-Model Benchmark](https://huggingface.co/spaces/MaduRox/Kalpana-Multi-Model-Benchmark), "
        "enabling fair side-by-side comparisons between:\n\n"
        "| System | Memory | Context Handling |\n"
        "|---|---|---|\n"
        "| **This API (Plain)** | Standard KV-Cache — O(N) | Full conversation re-sent every turn |\n"
        "| **[Kalpanā RIF API](https://huggingface.co/spaces/MaduRox/Kalpana-API-Public)** | "
        "RIF Holographic State — O(1), ~8 MB fixed | Knowledge Pack holds all context |\n\n"
        "### Available Models\n"
        "| Plain Model | Paired RIF Model |\n"
        "|---|---|\n"
        "| `plain-llama-3-8b` | `kalpana-llama-3-8b-rif` |\n"
        "| `plain-llama-3.2-3b` | `kalpana-llama-3.2-3b-rif` |\n"
        "| `plain-qwen-0.5b` | `kalpana-qwen-0.5b-rif` |\n\n"
        "### How It Works\n"
        "Each model is served with identical quantization (INT4 NF4) and hardware. "
        "The **only** difference is whether RIF memory is active. "
        "This isolates the impact of RIF on accuracy, latency, memory, and cost.\n\n"
        "### Related Spaces\n"
        "- **RIF API:** [Kalpana-API-Public](https://huggingface.co/spaces/MaduRox/Kalpana-API-Public)\n"
        "- **Benchmark:** [Kalpana-Multi-Model-Benchmark](https://huggingface.co/spaces/MaduRox/Kalpana-Multi-Model-Benchmark)\n"
        "- **Chat App:** [Kalpana-Chat](https://huggingface.co/spaces/MaduRox/Kalpana-Chat)"
    ),
    version="1.0.0",
    contact={"name": "Vijñāna AI", "email": "support@vijnanaai.com", "url": "https://huggingface.co/MaduRox"},
    openapi_tags=[
        {"name": "Chat", "description": "Standard chat completions — OpenAI-compatible `/v1/chat/completions` format. No RIF memory."},
        {"name": "Benchmark", "description": "Dedicated benchmark endpoint with truncation control and detailed metrics."},
        {"name": "Models", "description": "List available models, pre-load into GPU, and inspect configuration."},
        {"name": "System", "description": "Health check, diagnostics, and state reset."},
    ]
)

# CORS setup
origins = [origin.strip() for origin in config.ALLOWED_ORIGINS.split(",") if origin.strip()]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins if origins else ["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Exception handlers
errors.setup_exception_handlers(app)

# Env-based testing mock toggle
MOCK_INFERENCE = os.getenv("MOCK_INFERENCE", "false").lower() == "true"

@app.get("/health", tags=["System"],
         summary="Health Check",
         description="Returns API status, loaded model, available models, and hardware info.")
async def health():
    device = "cuda" if torch.cuda.is_available() else "cpu"
    active = model_service.active_model_id
    
    return {
        "status": "ok",
        "service": "plain-multi-model-baseline",
        "version": "1.0.0",
        "device": device,
        "inference_profile": config.INFERENCE_PROFILE,
        "active_model": active,
        "model_loaded": active is not None,
        "available_models": list(PLAIN_MODELS.keys()),
        "hardware": telemetry.get_hardware_info()
    }

@app.get("/v1/models", response_model=ModelListResponse, tags=["Models"],
         summary="List Available Models",
         description="Returns all plain baseline models. Each model has a `paired_rif_model` field showing its RIF-enabled counterpart for benchmark comparison.")
async def list_models(api_key: None = Depends(security.verify_api_key)):
    data = []
    for model_id in PLAIN_MODELS:
        data.append(
            ModelInfo(
                id=model_id,
                paired_rif_model=PAIRED_RIF_MODELS[model_id]
            )
        )
    return ModelListResponse(data=data)

@app.get("/v1/config", tags=["Models"],
         summary="Get Full Configuration",
         description="Returns detailed configuration for all models including HuggingFace IDs, quantization, runtime, context limits, and model hashes.")
async def get_config(api_key: None = Depends(security.verify_api_key)):
    models_meta = {}
    for m_id, hf_id in PLAIN_MODELS.items():
        ctx_limit = model_service.get_context_limit(m_id)
        # Check active model metadata if loaded
        is_active = (m_id == model_service.active_model_id)
        revision = model_service.loaded_model_revision if is_active else "main"
        m_hash = model_service.loaded_model_hash if is_active else "not-loaded"
        
        models_meta[m_id] = {
            "hf_model_id": hf_id,
            "paired_rif_model": PAIRED_RIF_MODELS[m_id],
            "quantization": "nf4" if config.INFERENCE_PROFILE == "gpu-transformers" else "Q4_K_M",
            "runtime": "transformers" if config.INFERENCE_PROFILE == "gpu-transformers" else "llama.cpp",
            "context_limit": ctx_limit,
            "model_revision": revision,
            "model_hash": m_hash
        }
        
    return {
        "service_type": "plain_baseline",
        "service_version": "1.0.0",
        "inference_profile": config.INFERENCE_PROFILE,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "active_model": model_service.active_model_id,
        "models": models_meta
    }

@app.post("/v1/models/load", tags=["Models"],
          summary="Pre-warm a Model",
          description="Pre-loads a model into GPU memory so the first inference request is instant. Models are swapped on demand — only one model is active at a time.")
async def load_model(request: ModelLoadRequest, api_key: None = Depends(security.verify_api_key)):
    validate_model_id(request.model)
    
    start_time = time.time()
    
    if not MOCK_INFERENCE:
        model_service.load_model_and_tokenizer(request.model)
        load_time = model_service.model_load_time_ms
    else:
        model_service.active_model_id = request.model
        model_service.loaded_model_revision = "main"
        model_service.loaded_model_hash = "mock-hash"
        load_time = 1500.0 # simulated
        time.sleep(1.5)
        
    return {
        "status": "success",
        "model": request.model,
        "paired_rif_model": PAIRED_RIF_MODELS[request.model],
        "loaded": True,
        "load_time_ms": load_time,
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "runtime": "transformers" if config.INFERENCE_PROFILE == "gpu-transformers" else "llama.cpp"
    }

@app.post("/v1/tokenize", response_model=TokenizeResponse, tags=["Models"],
          summary="Count Tokens",
          description="Counts the number of tokens a text would produce with the specified model's tokenizer. Useful for pre-flight budget checks before running benchmark queries.")
async def tokenize(request: TokenizeRequest, api_key: None = Depends(security.verify_api_key)):
    validate_model_id(request.model)
    
    if MOCK_INFERENCE:
        # Mock counting by dividing characters by 4 roughly
        count = len(request.text) // 4
    else:
        count = tokenization.count_tokens(request.model, request.text)
        
    return TokenizeResponse(model=request.model, token_count=count)

@app.post("/v1/reset", response_model=ResetResponse, tags=["System"],
          summary="Reset Server State",
          description="Resets server state between benchmark runs. Optionally unloads the active model from GPU memory to ensure clean cold-start measurements.")
async def reset(request: ResetRequest, api_key: None = Depends(security.verify_api_key)):
    if request.unload_model:
        model_service.unload_all_models()
        model_remained = False
    else:
        model_remained = model_service.active_model_id is not None
        
    return ResetResponse(
        run_id=request.run_id,
        model_remained_loaded=model_remained
    )

def execute_inference(
    model_id: str,
    prompt: str,
    max_tokens: int,
    temperature: float,
    seed: Optional[int]
) -> str:
    """
    Synchronous helper to run inference on loaded GGUF or Transformer model.
    """
    if MOCK_INFERENCE:
        time.sleep(0.5) # Simulate generation latency
        return f"Mock baseline response for question. (Model: {model_id}, temp: {temperature}, seed: {seed})"
        
    # Standard inference execution
    if config.INFERENCE_PROFILE == "gpu-transformers":
        # Ensure correct model is loaded
        model, tokenizer = model_service.load_model_and_tokenizer(model_id)
        
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        if seed is not None:
            torch.manual_seed(seed)
            
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=max_tokens,
                temperature=temperature,
                do_sample=temperature > 0.0,
                pad_token_id=tokenizer.pad_token_id
            )
            
        input_len = inputs.input_ids.shape[1]
        decoded = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
        return decoded
        
    elif config.INFERENCE_PROFILE == "cpu-gguf":
        model, tokenizer = model_service.load_model_and_tokenizer(model_id)
        
        # model is llama_cpp.Llama
        res = model.create_completion(
            prompt=prompt,
            max_tokens=max_tokens,
            temperature=temperature,
            seed=seed if seed is not None else -1
        )
        return res["choices"][0]["text"]
        
    else:
        raise ValueError(f"Unknown INFERENCE_PROFILE: {config.INFERENCE_PROFILE}")

@app.post("/v1/chat/completions", tags=["Chat"],
          summary="Chat Completions",
          description=(
              "Generate an AI response — compatible with OpenAI's `/v1/chat/completions` format.\n\n"
              "**This is the plain baseline** — no RIF memory, no Knowledge Packs. "
              "The full conversation history must be sent in `messages` on every request. "
              "Context grows linearly (O(N)) with conversation length.\n\n"
              "Compare this against the RIF-enabled version at "
              "[Kalpana-API-Public](https://huggingface.co/spaces/MaduRox/Kalpana-API-Public) "
              "to see the difference RIF memory makes."
          ))
async def chat_completions(request: ChatCompletionRequest, api_key: None = Depends(security.verify_api_key)):
    validate_model_id(request.model)
    
    start_time = time.time()
    
    # 1. Check if model is loaded, record load status
    switched = (model_service.active_model_id is not None and model_service.active_model_id != request.model)
    cold_start = (model_service.active_model_id is None)
    
    rss_before = telemetry.get_system_ram_mb()
    
    # Pre-load or warm up
    if not MOCK_INFERENCE:
        model, tokenizer = model_service.load_model_and_tokenizer(request.model)
        load_time = model_service.model_load_time_ms
    else:
        model_service.active_model_id = request.model
        load_time = 0.0
        tokenizer = None
        
    # 2. Format chat prompt using standard template
    messages_dicts = [{"role": m.role, "content": m.content} for m in request.messages]
    
    if MOCK_INFERENCE:
        prompt_str = f"Mocked chat prompt for message history of length {len(messages_dicts)}"
        prompt_tokens = len(prompt_str) // 4
    else:
        prompt_str = tokenizer.apply_chat_template(messages_dicts, tokenize=False, add_generation_prompt=True)
        prompt_tokens = len(tokenizer.encode(prompt_str, add_special_tokens=False))
        
    # 3. Generate response text
    t_gen_start = time.time()
    answer = execute_inference(
        model_id=request.model,
        prompt=prompt_str,
        max_tokens=request.max_tokens,
        temperature=request.temperature,
        seed=request.seed
    )
    generation_time = (time.time() - t_gen_start) * 1000.0
    
    if MOCK_INFERENCE:
        completion_tokens = len(answer) // 4
    else:
        completion_tokens = len(tokenizer.encode(answer, add_special_tokens=False))
        
    total_tokens = prompt_tokens + completion_tokens
    end_to_end = (time.time() - start_time) * 1000.0
    
    tokens_per_second = (completion_tokens / (generation_time / 1000.0)) if generation_time > 0 else 0.0
    rss_after = telemetry.get_system_ram_mb()
    
    # Format standard OpenAI Response structure plus custom benchmark metrics
    request_uuid = f"chatcmpl-{uuid.uuid4()}"
    
    metrics = BenchmarkMetrics(
        request_id=request_uuid,
        run_id=request.run_id,
        question_id=request.question_id,
        model_id=request.model,
        paired_rif_model_id=PAIRED_RIF_MODELS[request.model],
        cold_start=cold_start,
        model_switched=switched,
        model_load_time_ms=load_time,
        prompt_build_time_ms=0.0,
        tokenization_time_ms=0.0,
        prompt_tokens=prompt_tokens,
        completion_tokens=completion_tokens,
        total_tokens=total_tokens,
        generation_time_ms=generation_time,
        end_to_end_time_ms=end_to_end,
        tokens_per_second=tokens_per_second,
        rss_before_mb=rss_before,
        rss_after_mb=rss_after,
        peak_rss_mb=max(rss_before, rss_after),
        gpu_memory_allocated_mb=telemetry.get_gpu_vram_mb(),
        context_truncated=False,
        tokens_discarded=0,
        truncation_policy="none",
        runtime="transformers" if config.INFERENCE_PROFILE == "gpu-transformers" else "llama.cpp",
        device="cuda" if torch.cuda.is_available() else "cpu"
    )
    
    response_content = {
        "id": request_uuid,
        "object": "chat.completion",
        "created": int(time.time()),
        "model": request.model,
        "choices": [
            {
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": answer
                },
                "finish_reason": "stop"
            }
        ],
        "usage": {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens
        },
        "benchmark_metrics": metrics.model_dump()
    }
    
    if request.return_final_prompt:
        response_content["final_prompt"] = prompt_str
        
    return JSONResponse(content=response_content)

@app.post("/v1/benchmark/query", response_model=BenchmarkQueryResponse, tags=["Benchmark"],
          summary="Run a Benchmark Query",
          description=(
              "Executes a single benchmark query with full control over truncation, token budgets, and seeding.\n\n"
              "The benchmark orchestrator uses this endpoint to run identical questions against the plain model "
              "and its paired RIF model, then compares accuracy, latency, and resource usage.\n\n"
              "Returns detailed `timings` and `resources` for precise performance measurement."
          ))
async def benchmark_query(request: BenchmarkQueryRequest, api_key: None = Depends(security.verify_api_key)):
    validate_model_id(request.model)
    
    start_time = time.time()
    
    # 1. Warm model up / load if needed
    switched = (model_service.active_model_id is not None and model_service.active_model_id != request.model)
    cold_start = (model_service.active_model_id is None)
    
    rss_before = telemetry.get_system_ram_mb()
    
    if not MOCK_INFERENCE:
        model, tokenizer = model_service.load_model_and_tokenizer(request.model)
        load_time = model_service.model_load_time_ms
        context_limit = model_service.get_context_limit(request.model)
    else:
        model_service.active_model_id = request.model
        load_time = 0.0
        tokenizer = None
        context_limit = 32768
        
    # 2. Reconstruct prompt string with deterministic truncation
    prompt_build_start = time.time()
    
    if MOCK_INFERENCE:
        source_tokens = len(request.source_text) // 4
        # Simulated truncation
        context_truncated = source_tokens > (request.prompt_budget_tokens - 500)
        tokens_discarded = max(0, source_tokens - (request.prompt_budget_tokens - 500))
        prompt_tokens = min(source_tokens + 500, request.prompt_budget_tokens)
        
        final_prompt_str = f"Mocked prompt with truncated source content. System: {request.system_prompt}. Question: {request.question}"
        messages_dicts = []
    else:
        # Build standard chat prompt dynamically applying the truncation policy on token lists
        final_prompt_str, messages_dicts, source_tokens, tokens_discarded, context_truncated = prompt_builder.build_benchmark_prompt(
            tokenizer=tokenizer,
            source_text=request.source_text,
            question=request.question,
            system_prompt=request.system_prompt,
            prompt_budget=request.prompt_budget_tokens,
            max_output_tokens=request.max_output_tokens,
            truncation_policy=request.truncation_policy,
            context_limit=context_limit
        )
        prompt_tokens = len(tokenizer.encode(final_prompt_str, add_special_tokens=False))
        
    prompt_build_time = (time.time() - prompt_build_start) * 1000.0
    
    # 3. Generate response
    t_gen_start = time.time()
    answer = execute_inference(
        model_id=request.model,
        prompt=final_prompt_str,
        max_tokens=request.max_output_tokens,
        temperature=request.temperature,
        seed=request.seed
    )
    generation_time = (time.time() - t_gen_start) * 1000.0
    
    if MOCK_INFERENCE:
        completion_tokens = len(answer) // 4
    else:
        completion_tokens = len(tokenizer.encode(answer, add_special_tokens=False))
        
    end_to_end = (time.time() - start_time) * 1000.0
    rss_after = telemetry.get_system_ram_mb()
    
    timings = {
        "model_load_time_ms": load_time,
        "prompt_build_time_ms": prompt_build_time,
        "generation_time_ms": generation_time,
        "end_to_end_time_ms": end_to_end
    }
    
    resources = {
        "rss_before_mb": rss_before,
        "rss_after_mb": rss_after,
        "peak_rss_mb": max(rss_before, rss_after),
        "gpu_memory_allocated_mb": telemetry.get_gpu_vram_mb(),
        "device": "cuda" if torch.cuda.is_available() else "cpu",
        "runtime": "transformers" if config.INFERENCE_PROFILE == "gpu-transformers" else "llama.cpp"
    }
    
    response = BenchmarkQueryResponse(
        run_id=request.run_id,
        question_id=request.question_id,
        system="plain",
        model=request.model,
        paired_rif_model=PAIRED_RIF_MODELS[request.model],
        answer=answer,
        source_tokens=source_tokens,
        prompt_tokens=prompt_tokens,
        tokens_discarded=tokens_discarded,
        context_truncated=context_truncated,
        truncation_policy=request.truncation_policy,
        timings=timings,
        resources=resources
    )
    
    if request.return_final_prompt:
        response.final_prompt = final_prompt_str
        
    return response

from fastapi.responses import RedirectResponse

@app.get("/", include_in_schema=False)
def root_redirect():
    return RedirectResponse(url="/docs")