rerank: add /v1/rerank cross-encoder endpoint + surface OCR in /v1/models

Reranker (e.g. BAAI/bge-reranker-v2-m3):
- new codai/api/rerank.py: /v1/rerank scores query↔document pairs with a native
  AutoModelForSequenceClassification cross-encoder (no new dep), cached +
  VRAM/thermal-managed via request_model like embeddings. Returns
  {index, relevance_score(sigmoid), document?} sorted desc, top_n honored.
- router: /v1/rerank added to _INFERENCE_PATHS so the front routes it by model
  capability (transformers → nvidia engine), like /v1/embeddings.
- app: include rerank_router.
- capabilities: add `reranking` + `ocr` fields; detect 'rerank'/'cross-encoder'
  names as reranking.

/v1/models discoverability:
- list_models() now appends enabled OCR engines (paddle/doctr/surya) as
  type='ocr' entries (a node subsystem invoked at /v1/ocr?engine=<id>), so
  clients can discover them.
- a reranker registered in models.json (embedding_models, capabilities
  ['reranking']) surfaces automatically via list_models, like bge-m3.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Mw2KQiswmD69T45fTfjKwW

codai/api/app.py

@@ -143,6 +143,7 @@ from codai.api.loras import router as loras_router
 from codai.api.spatial import router as spatial_router
 from codai.api.environments import router as environments_router
 from codai.api.ocr import router as ocr_router
+from codai.api.rerank import router as rerank_router
 from codai.admin.routes import router as admin_router
 
 # Import and add middleware
@@ -583,6 +584,7 @@ app.include_router(loras_router, tags=["LoRAs"])
 app.include_router(environments_router, tags=["Environments"])
 app.include_router(spatial_router, tags=["Spatial / 3D"])
 app.include_router(ocr_router, tags=["OCR"])
+app.include_router(rerank_router, tags=["Rerank"])
 app.include_router(admin_router, tags=["Admin"])
 
 

codai/api/rerank.py

+# CoderAI - OpenAI-compatible API server
+# Copyright (C) 2026 Stefy Lanza <stefy@nexlab.net>
+# GPLv3 - see the project LICENSE.
+
+"""Cross-encoder reranking endpoint (``/v1/rerank``).
+
+Scores query↔document relevance with a sequence-classification cross-encoder
+(e.g. ``BAAI/bge-reranker-v2-m3``). The model is loaded NATIVELY via transformers
+(``AutoModelForSequenceClassification`` — no extra dependency) and cached +
+VRAM-managed through :data:`multi_model_manager` like every other model, so it
+participates in eviction and thermal gating (``request_model`` waits on the
+thermal governor before serving).
+
+Response shape mirrors the common rerank APIs (Cohere/Jina): a list of
+``{index, relevance_score}`` sorted by score, optionally with the document text.
+``relevance_score`` is the sigmoid of the cross-encoder logit, in [0, 1]
+(monotonic, so ranking order is preserved).
+"""
+
+import asyncio
+from typing import List, Optional
+
+from fastapi import APIRouter, HTTPException, Request
+from pydantic import BaseModel, ConfigDict, Field
+
+from codai.models.manager import multi_model_manager
+
+router = APIRouter()
+
+global_args = None
+
+
+def set_global_args(args):
+    global global_args
+    global_args = args
+
+
+# Per-model-key asyncio locks so a burst of first-requests loads the model ONCE.
+_load_locks: dict = {}
+
+
+class RerankRequest(BaseModel):
+    model: str = Field(..., description="Reranker model id (e.g. 'bge-reranker-v2-m3').")
+    query: str = Field(..., description="The search query.")
+    documents: List[str] = Field(..., description="Documents to score against the query.")
+    top_n: Optional[int] = Field(None, description="Return only the top N results.")
+    return_documents: Optional[bool] = Field(
+        False, description="Include the document text in each result.")
+    max_length: Optional[int] = Field(
+        None, description="Max tokens per query+document pair (default: model n_ctx or 8192).")
+    model_config = ConfigDict(extra="allow")
+
+
+def _derive_device() -> str:
+    try:
+        import torch
+        if torch.cuda.is_available():
+            return "cuda:0"
+    except Exception:
+        pass
+    return "cpu"
+
+
+def _is_reranker(obj) -> bool:
+    return isinstance(obj, tuple) and len(obj) == 2 and obj[0] == "reranker"
+
+
+def _load_reranker(model_name: str, device: str, model_config: dict):
+    """Load a cross-encoder as ('reranker', (tokenizer, model, device))."""
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+    from codai.models.hf_loading import build_from_pretrained_kwargs
+
+    fp = build_from_pretrained_kwargs(model_config or {})
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSequenceClassification.from_pretrained(model_name, **fp)
+    if 'quantization_config' not in fp and 'device_map' not in fp:
+        model = model.to(device)
+    model.eval()
+    return ("reranker", (tokenizer, model, device))
+
+
+def _score(model_obj, query: str, docs: List[str], max_length: int) -> List[float]:
+    """Score each (query, doc) pair → sigmoid(logit) in [0, 1]."""
+    import torch
+
+    _tag, (tokenizer, model, device) = model_obj
+    pairs = [[query, d] for d in docs]
+    enc = tokenizer(pairs, padding=True, truncation=True,
+                    max_length=int(max_length), return_tensors='pt')
+    enc = {k: v.to(device) for k, v in enc.items()}
+    with torch.no_grad():
+        logits = model(**enc).logits.view(-1).float()
+        scores = torch.sigmoid(logits)
+    return scores.cpu().tolist()
+
+
+@router.post("/v1/rerank", summary="Rerank documents against a query (cross-encoder)")
+async def create_rerank(request: RerankRequest, http_request: Request = None):
+    if not request.documents:
+        raise HTTPException(status_code=400, detail="'documents' must be a non-empty list")
+    if not request.query:
+        raise HTTPException(status_code=400, detail="'query' is required")
+
+    # Resolve + reserve the model through the manager (thermal wait, caching, routing).
+    model_info = await asyncio.to_thread(
+        multi_model_manager.request_model, request.model, "rerank")
+    model_name = model_info.get('model_name')
+    if not model_name:
+        raise HTTPException(status_code=404,
+                            detail=model_info.get('error', f"Model '{request.model}' not found"))
+    model_key = model_info['model_key']
+    model_obj = model_info.get('model_object')
+
+    if not _is_reranker(model_obj):
+        lock = _load_locks.setdefault(model_key, asyncio.Lock())
+        async with lock:
+            model_obj = multi_model_manager.models.get(model_key)
+            if not _is_reranker(model_obj):
+                device = _derive_device()
+                _cfg = multi_model_manager.config.get(model_name) or {}
+                _snap = multi_model_manager.vram_before_load()
+                try:
+                    model_obj = await asyncio.get_event_loop().run_in_executor(
+                        None, _load_reranker, model_name, device, _cfg)
+                except Exception as e:
+                    raise HTTPException(status_code=500,
+                                        detail=f"Failed to load reranker '{model_name}': {e}")
+                multi_model_manager.add_model(model_key, model_obj)
+                multi_model_manager.current_model_key = model_key
+                try:
+                    multi_model_manager.record_vram_delta(model_key, _snap)
+                except Exception:
+                    pass
+
+    _cfg = multi_model_manager.config.get(model_name) or {}
+    max_len = int(request.max_length or _cfg.get('n_ctx') or 8192)
+    try:
+        scores = await asyncio.get_event_loop().run_in_executor(
+            None, _score, model_obj, request.query, request.documents, max_len)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Rerank failed: {e}")
+
+    results = [{"index": i, "relevance_score": float(s)} for i, s in enumerate(scores)]
+    results.sort(key=lambda r: r["relevance_score"], reverse=True)
+    if request.top_n:
+        results = results[:max(0, int(request.top_n))]
+    if request.return_documents:
+        for r in results:
+            r["document"] = request.documents[r["index"]]
+
+    total_tokens = len(request.query.split()) + sum(len(d.split()) for d in request.documents)
+    return {
+        "object": "rerank.result",
+        "model": request.model,
+        "results": results,
+        "usage": {"total_tokens": total_tokens},
+    }

codai/frontproxy/router.py

@@ -29,6 +29,7 @@ _INFERENCE_PATHS = {
     "/v1/chat/completions",
     "/v1/completions",
     "/v1/embeddings",
+    "/v1/rerank",
     "/v1/images/generations",
     "/v1/images/edits",
     "/v1/audio/speech",

codai/models/capabilities.py

@@ -73,6 +73,10 @@ class ModelCapabilities:
     model_3d_generation: bool = False   # text / image → 3D model (GLB)
     model_3d_to_image: bool = False     # 3D model → rendered 2D image / video
 
+    # Retrieval / document
+    reranking: bool = False             # cross-encoder query↔document reranking (bge-reranker, …)
+    ocr: bool = False                   # optical character recognition (paddle/doctr/surya)
+
     def to_list(self) -> List[str]:
         out = []
         for name, val in self.__dataclass_fields__.items():
@@ -96,6 +100,11 @@ def detect_model_capabilities(model_name: str) -> ModelCapabilities:
 
     n = model_name.lower()
 
+    # ── Reranking (cross-encoder) ────────────────────────────────────────────
+    if 'rerank' in n or 'cross-encoder' in n or 'cross_encoder' in n:
+        caps.reranking = True
+        return caps
+
     # ── 3D generation ────────────────────────────────────────────────────────
     if any(x in n for x in ['triposr', 'tsr', 'shap-e', 'shape-e', 'point-e', 'pointe',
                               'zero123', 'wonder3d', 'instant3d', 'one-2-3-45',

codai/models/manager.py

@@ -5129,6 +5129,26 @@ class MultiModelManager:
         except Exception:
             pass
 
+        # --- OCR engines (a node subsystem, not model-manager models) ---
+        # Surfaced so clients can discover them via /v1/models; invoked at /v1/ocr
+        # with ?engine=<id>. Listed once per node (collect_models dedups by id).
+        try:
+            from codai.admin.routes import config_manager
+            ocr = getattr(config_manager, "config", None)
+            ocr = getattr(ocr, "ocr", None) if ocr is not None else None
+            if ocr is not None and getattr(ocr, "enabled", False):
+                _ocr_engines = []
+                if getattr(ocr, "paddle_enabled", False):
+                    _ocr_engines.append("paddle")
+                if getattr(ocr, "doctr_enabled", False):
+                    _ocr_engines.append("doctr")
+                if getattr(ocr, "surya_enabled", False) and getattr(ocr, "surya_accept_license", False):
+                    _ocr_engines.append("surya")
+                for _eng in _ocr_engines:
+                    _add(_eng, "ocr", {"capabilities": ["ocr"], "backend": "ocr"})
+        except Exception:
+            pass
+
         # --- Fallback: runtime default_model ---
         if not models and self.default_model:
             model_id = self.default_model