""" Inference API routes for model loading and text generation. """ import sys import time import uuid from pathlib import Path from fastapi import APIRouter, HTTPException, Request from fastapi.responses import StreamingResponse, JSONResponse from typing import Optional import json import logging import asyncio import threading # Add backend directory to path backend_path = Path(__file__).parent.parent.parent if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) # Import backend functions try: from core.inference import get_inference_backend from utils.models import ModelConfig from utils.inference import load_inference_config except ImportError: parent_backend = backend_path.parent / "backend" if str(parent_backend) not in sys.path: sys.path.insert(0, str(parent_backend)) from core.inference import get_inference_backend from utils.models import ModelConfig from utils.inference import load_inference_config from models.inference import ( LoadRequest, UnloadRequest, GenerateRequest, LoadResponse, UnloadResponse, InferenceStatusResponse, ChatCompletionRequest, ChatCompletionChunk, ChatCompletion, ChunkChoice, ChoiceDelta, CompletionChoice, CompletionMessage, ) router = APIRouter() logger = logging.getLogger(__name__) # Configure logger if not logger.handlers: handler = logging.StreamHandler() handler.setLevel(logging.INFO) formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') handler.setFormatter(formatter) logger.addHandler(handler) logger.setLevel(logging.INFO) @router.post("/load", response_model=LoadResponse) async def load_model(request: LoadRequest): """ Load a model for inference. The model_path should be a clean identifier from GET /models/list. Returns inference configuration parameters (temperature, top_p, top_k, min_p) from the model's YAML config, falling back to default.yaml for missing values. """ try: backend = get_inference_backend() # Create config using clean factory method # is_lora is auto-detected from adapter_config.json on disk/HF config = ModelConfig.from_identifier( model_id=request.model_path, hf_token=request.hf_token, ) if not config: raise HTTPException( status_code=400, detail=f"Invalid model identifier: {request.model_path}" ) # Load the model success = backend.load_model( config=config, max_seq_length=request.max_seq_length, load_in_4bit=request.load_in_4bit, hf_token=request.hf_token, ) if not success: raise HTTPException( status_code=500, detail=f"Failed to load model: {config.display_name}" ) logger.info(f"Loaded model: {config.identifier}") # Load inference configuration parameters inference_config = load_inference_config(config.identifier) return LoadResponse( status="loaded", model=config.identifier, display_name=config.display_name, is_vision=config.is_vision, is_lora=config.is_lora, inference=inference_config, ) except HTTPException: raise except Exception as e: logger.error(f"Error loading model: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to load model: {str(e)}" ) @router.post("/unload", response_model=UnloadResponse) async def unload_model(request: UnloadRequest): """ Unload a model from memory. """ try: backend = get_inference_backend() backend.unload_model(request.model_path) logger.info(f"Unloaded model: {request.model_path}") return UnloadResponse(status="unloaded", model=request.model_path) except Exception as e: logger.error(f"Error unloading model: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to unload model: {str(e)}" ) @router.post("/generate/stream") async def generate_stream(request: GenerateRequest): """ Generate a chat response with Server-Sent Events (SSE) streaming. For vision models, provide image_base64 with the base64-encoded image. """ backend = get_inference_backend() if not backend.active_model_name: raise HTTPException( status_code=400, detail="No model loaded. Call POST /inference/load first." ) # Decode image if provided (for vision models) image = None if request.image_base64: try: import base64 from PIL import Image from io import BytesIO # Check if current model supports vision model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_vision"): raise HTTPException( status_code=400, detail="Image provided but current model is text-only. Load a vision model." ) image_data = base64.b64decode(request.image_base64) image = Image.open(BytesIO(image_data)) image = backend.resize_image(image) except HTTPException: raise except Exception as e: raise HTTPException( status_code=400, detail=f"Failed to decode image: {str(e)}" ) async def stream(): try: for chunk in backend.generate_chat_response( messages=request.messages, system_prompt=request.system_prompt, image=image, temperature=request.temperature, top_p=request.top_p, top_k=request.top_k, max_new_tokens=request.max_new_tokens, repetition_penalty=request.repetition_penalty, ): yield f"data: {json.dumps({'content': chunk})}\n\n" yield "data: [DONE]\n\n" except Exception as e: backend.reset_generation_state() logger.error(f"Error during generation: {e}", exc_info=True) yield f"data: {json.dumps({'error': str(e)})}\n\n" return StreamingResponse( stream(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", } ) @router.get("/status", response_model=InferenceStatusResponse) async def get_status(): """ Get current inference backend status. """ try: backend = get_inference_backend() is_vision = False if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) is_vision = model_info.get("is_vision", False) return InferenceStatusResponse( active_model=backend.active_model_name, is_vision=is_vision, loading=list(getattr(backend, 'loading_models', set())), loaded=list(backend.models.keys()), ) except Exception as e: logger.error(f"Error getting status: {e}", exc_info=True) raise HTTPException( status_code=500, detail=f"Failed to get status: {str(e)}" ) # ===================================================================== # OpenAI-Compatible Chat Completions (/chat/completions) # ===================================================================== def _extract_content_parts( messages: list, ) -> tuple[str, list[dict], "Optional[str]"]: """ Parse OpenAI-format messages into components the inference backend expects. Handles both plain-string ``content`` and multimodal content-part arrays (``[{type: "text", ...}, {type: "image_url", ...}]``). Returns: system_prompt: The system message text (or a default). chat_messages: Non-system messages with content flattened to strings. image_base64: Base64 data of the *first* image found, or ``None``. """ system_prompt = "You are a helpful AI assistant." chat_messages: list[dict] = [] first_image_b64: Optional[str] = None for msg in messages: # ── System messages → extract as system_prompt ──────── if msg.role == "system": if isinstance(msg.content, str): system_prompt = msg.content elif isinstance(msg.content, list): # Unlikely but handle: join text parts system_prompt = "\n".join( p.text for p in msg.content if p.type == "text" ) continue # ── User / assistant messages ───────────────────────── if isinstance(msg.content, str): # Plain string content — pass through chat_messages.append({"role": msg.role, "content": msg.content}) elif isinstance(msg.content, list): # Multimodal content parts text_parts: list[str] = [] for part in msg.content: if part.type == "text": text_parts.append(part.text) elif part.type == "image_url" and first_image_b64 is None: url = part.image_url.url if url.startswith("data:"): # data:image/png;base64, → extract first_image_b64 = url.split(",", 1)[1] if "," in url else None else: logger.warning( f"Remote image URLs not yet supported: {url[:80]}..." ) combined_text = "\n".join(text_parts) if text_parts else "" chat_messages.append({"role": msg.role, "content": combined_text}) return system_prompt, chat_messages, first_image_b64 @router.post("/chat/completions") async def openai_chat_completions(payload: ChatCompletionRequest, request: Request): """ OpenAI-compatible chat completions endpoint. Supports multimodal messages: ``content`` may be a plain string or a list of content parts (``text`` / ``image_url``). Streaming (default): returns SSE chunks matching OpenAI's format. Non-streaming: returns a single ChatCompletion JSON object. """ backend = get_inference_backend() if not backend.active_model_name: raise HTTPException( status_code=400, detail="No model loaded. Call POST /inference/load first.", ) # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( payload.messages ) # If no non-system messages were provided, error out if not chat_messages: raise HTTPException( status_code=400, detail="At least one non-system message is required.", ) # ── Decode image (from content parts OR legacy field) ───── # Content-part images take priority; fall back to legacy field image_b64 = extracted_image_b64 or payload.image_base64 image = None if image_b64: try: import base64 from PIL import Image from io import BytesIO model_info = backend.models.get(backend.active_model_name, {}) if not model_info.get("is_vision"): raise HTTPException( status_code=400, detail="Image provided but current model is text-only. Load a vision model.", ) image_data = base64.b64decode(image_b64) image = Image.open(BytesIO(image_data)) image = backend.resize_image(image) except HTTPException: raise except Exception as e: raise HTTPException(status_code=400, detail=f"Failed to decode image: {e}") # ── Shared generation kwargs ────────────────────────────── gen_kwargs = dict( messages=chat_messages, system_prompt=system_prompt, image=image, temperature=payload.temperature, top_p=payload.top_p, top_k=payload.top_k, min_p=payload.min_p, max_new_tokens=payload.max_tokens or 512, repetition_penalty=payload.repetition_penalty, ) # ── Choose generation path (adapter-controlled or standard) ── cancel_event = threading.Event() if payload.use_adapter is not None: # Compare mode: toggle adapter state atomically with generation def generate(): return backend.generate_with_adapter_control( use_adapter=payload.use_adapter, cancel_event=cancel_event, **gen_kwargs, ) else: # Standard path: no adapter toggling def generate(): return backend.generate_chat_response(cancel_event=cancel_event, **gen_kwargs) model_name = backend.active_model_name or payload.model completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" created = int(time.time()) # ── Streaming response ──────────────────────────────────────── if payload.stream: async def stream_chunks(): try: # First chunk: send the role first_chunk = ChatCompletionChunk( id=completion_id, created=created, model=model_name, choices=[ChunkChoice( delta=ChoiceDelta(role="assistant"), finish_reason=None, )], ) yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" # Content chunks — generate_chat_response yields cumulative # text, so we diff to get incremental deltas. prev_text = "" for cumulative in generate(): if await request.is_disconnected(): cancel_event.set() backend.reset_generation_state() return new_text = cumulative[len(prev_text):] prev_text = cumulative if not new_text: continue chunk = ChatCompletionChunk( id=completion_id, created=created, model=model_name, choices=[ChunkChoice( delta=ChoiceDelta(content=new_text), finish_reason=None, )], ) yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" # Final chunk: finish_reason = stop final_chunk = ChatCompletionChunk( id=completion_id, created=created, model=model_name, choices=[ChunkChoice( delta=ChoiceDelta(), finish_reason="stop", )], ) yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" yield "data: [DONE]\n\n" except asyncio.CancelledError: cancel_event.set() backend.reset_generation_state() raise except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI streaming: {e}", exc_info=True) error_chunk = { "error": {"message": str(e), "type": "server_error"}, } yield f"data: {json.dumps(error_chunk)}\n\n" return StreamingResponse( stream_chunks(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) # ── Non-streaming response ──────────────────────────────────── else: try: full_text = "" for token in generate(): full_text = token # generate_stream yields cumulative text response = ChatCompletion( id=completion_id, created=created, model=model_name, choices=[CompletionChoice( message=CompletionMessage(content=full_text), finish_reason="stop", )], ) return JSONResponse(content=response.model_dump()) except Exception as e: backend.reset_generation_state() logger.error(f"Error during OpenAI completion: {e}", exc_info=True) raise HTTPException(status_code=500, detail=str(e))