"""Routes liées à Ollama : statut de connexion, modèles, téléchargement.""" from __future__ import annotations import asyncio import json import httpx from fastapi import APIRouter, HTTPException, Query from fastapi.responses import StreamingResponse from pydantic import BaseModel from .. import agent_config from ..config import settings from ..ollama_client import ollama router = APIRouter(prefix="/api", tags=["ollama"]) # État en mémoire des préchargements. Loki est lancé avec un worker unique ; ce # suivi permet au navigateur d'interroger une route courte pendant que le long # chargement Ollama continue sans maintenir la requête HTTP initiale ouverte. _warm_states: dict[str, dict[str, str]] = {} _warm_tasks: dict[str, asyncio.Task[None]] = {} @router.get("/status") async def status() -> dict: """État de la connexion Ollama (point vert/rouge de la barre supérieure).""" try: version = await ollama.ping() return { "connected": True, "host": ollama.host, "version": version.get("version"), "default_model": settings.default_model, } except (httpx.HTTPError, OSError) as exc: return { "connected": False, "host": ollama.host, "error": str(exc), "default_model": settings.default_model, } @router.get("/models") async def list_models() -> dict: """Modèles installés localement, formatés pour le sélecteur de l'UI.""" try: raw = await ollama.list_models_cached() except (httpx.HTTPError, OSError): return {"models": []} models = [] for m in raw: details = m.get("details", {}) size_go = round(m.get("size", 0) / 1_000_000_000, 1) models.append( { "name": m.get("name"), "size_go": size_go, "parameter_size": details.get("parameter_size"), "quantization": details.get("quantization_level"), "family": details.get("family"), } ) return {"models": models, "default": settings.default_model} class WarmRequest(BaseModel): name: str keep_alive: str = "30m" async def _placement_of(name: str) -> dict: """Où le modèle vient-il d'être chargé ? (GPU / CPU / mixte), via /api/ps.""" try: for m in await ollama.ps(): if m.get("name") == name or m.get("model") == name: size = m.get("size", 0) or 0 vram = m.get("size_vram", 0) or 0 if size <= 0: return {} pct = int(vram / size * 100) where = "gpu" if pct >= 99 else "cpu" if pct <= 1 else "mixte" return {"processor": where, "gpu_percent": str(pct)} except (httpx.HTTPError, OSError): pass return {} async def _warm_in_background(name: str, keep_alive: str, options: dict | None) -> None: try: await ollama.warm(name, keep_alive, options) state = {"state": "loaded"} state.update(await _placement_of(name)) _warm_states[name] = state except Exception as exc: _warm_states[name] = { "state": "error", "error": f"préchargement impossible : {str(exc)[:500]}", } def start_model_warm(name: str, keep_alive: str) -> None: """Démarre au plus une tâche de préchargement par modèle. Les options de génération du modèle (num_ctx, num_batch, num_gpu…) sont envoyées au préchargement pour qu'Ollama charge exactement le runner que le chat utilisera — évite un rechargement complet au premier message. """ current = _warm_tasks.get(name) if current and not current.done(): return cfg = agent_config.get_config(name) options = agent_config.ollama_options(cfg) _warm_states[name] = {"state": "loading"} task = asyncio.create_task(_warm_in_background(name, keep_alive, options)) _warm_tasks[name] = task def forget(done: asyncio.Task[None]) -> None: if _warm_tasks.get(name) is done: _warm_tasks.pop(name, None) task.add_done_callback(forget) @router.post("/models/warm", status_code=202) async def warm_model(req: WarmRequest) -> dict: """Démarre le préchargement sans exposer sa durée au reverse proxy.""" name = req.name.strip() if not name: raise HTTPException(400, "nom de modèle vide") start_model_warm(name, req.keep_alive) return {"warming": name, "state": _warm_states[name]["state"]} @router.get("/models/warm/status") async def warm_status(name: str = Query(min_length=1)) -> dict: """État court du préchargement : idle, loading, loaded ou error.""" return _warm_states.get(name.strip(), {"state": "idle"}) @router.get("/models/loaded") async def loaded_models() -> dict: """Modèles actuellement chargés en mémoire + placement GPU/CPU (/api/ps).""" try: loaded = await ollama.ps() except (httpx.HTTPError, OSError): return {"loaded": []} result = [] for m in loaded: size = m.get("size", 0) or 0 vram = m.get("size_vram", 0) or 0 result.append({ "name": m.get("name") or m.get("model"), "on_gpu": bool(size and vram >= size * 0.99), "gpu_percent": int(vram / size * 100) if size else 0, }) return {"loaded": result} class PullRequest(BaseModel): name: str class DeleteModelRequest(BaseModel): name: str @router.delete("/models") async def delete_model(req: DeleteModelRequest) -> dict: """Supprime explicitement un modèle de l'instance Ollama.""" if not req.name.strip(): raise HTTPException(400, "nom de modèle vide") try: await ollama.delete_model(req.name.strip()) except httpx.HTTPStatusError as exc: detail = exc.response.text[:500] or str(exc) raise HTTPException(502, f"Ollama : {detail}") from exc except (httpx.HTTPError, OSError) as exc: raise HTTPException(502, f"Ollama injoignable : {exc}") from exc return {"deleted": req.name.strip()} @router.post("/models/pull") async def pull_model(req: PullRequest) -> StreamingResponse: """Télécharge un modèle en streamant la progression (SSE).""" async def event_stream(): try: async for chunk in ollama.pull_model(req.name): yield f"data: {json.dumps(chunk)}\n\n" except (httpx.HTTPError, OSError) as exc: yield f"data: {json.dumps({'error': str(exc)})}\n\n" yield "data: [DONE]\n\n" return StreamingResponse(event_stream(), media_type="text/event-stream")