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")