openai: compose healthchecks + dependency ordering, registries, LiteLLM routing
docker-compose.yml gains healthchecks and depends_on/condition chains for the litellm/langfuse/postgres tier so dependants wait for a genuinely ready service instead of a started container. Also plumbs AGAP_MCP_TOKEN into the adolf and adolf-llm containers, sourced from openai/.env (gitignored), for the kb#180 bearer auth on the agap MCP server; shared-mcp.json consumes it via bearerTokenEnvVar so the Kimi backbone authenticates too. agent-registry.yaml / agent_registry.py: the version-controlled source of truth for agent identities and trust classes -- the same ids the agap-mcp token map resolves to (`adolf`, `claude-coder`; note `claude-code-cli` is the runtime entry, not an agent identity). model-registry.yaml, litellm-config.yaml, auto-router-routes.json and provision_litellm_keys.py: model tiering, virtual-key provisioning and auto-router routes. tei-reranker/ is the local reranker service backing Hindsight recall. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
85
openai/tei-reranker/server.py
Normal file
85
openai/tei-reranker/server.py
Normal file
@@ -0,0 +1,85 @@
|
||||
"""Minimal TEI-compatible cross-encoder rerank server (GPU).
|
||||
|
||||
Why this exists: HuggingFace's official Text-Embeddings-Inference GPU images
|
||||
require CUDA compute capability >= 7.5 (Turing+). This box has a GTX 1070
|
||||
(Pascal, 6.1), so the stock TEI image won't run. Plain CUDA torch DOES support
|
||||
Pascal (that's why ollama works here), so we serve the same
|
||||
`jina-reranker-v2-base-multilingual` cross-encoder via torch and expose only the
|
||||
two endpoints Hindsight's `tei` reranker provider calls:
|
||||
GET /info -> JSON (init/health probe)
|
||||
POST /rerank -> {"query": str, "texts": [str], ...}
|
||||
-> bare list [{"index": i, "score": f}, ...] sorted desc
|
||||
See hindsight_api/engine/cross_encoder.py::RemoteTEICrossEncoder for the client.
|
||||
"""
|
||||
|
||||
import os
|
||||
import torch
|
||||
from fastapi import FastAPI
|
||||
from pydantic import BaseModel
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
|
||||
MODEL_ID = os.environ.get("RERANKER_MODEL", "jinaai/jina-reranker-v2-base-multilingual")
|
||||
DEVICE = os.environ.get("RERANKER_DEVICE", "cuda")
|
||||
MAX_LENGTH = int(os.environ.get("RERANKER_MAX_LENGTH", "1024"))
|
||||
# fp16 on GPU halves the ~1.1GB fp32 footprint; Pascal supports fp16 storage.
|
||||
DTYPE = torch.float16 if DEVICE == "cuda" else torch.float32
|
||||
|
||||
app = FastAPI(title="tei-reranker")
|
||||
_model = None
|
||||
|
||||
|
||||
def _load():
|
||||
global _model
|
||||
if _model is not None:
|
||||
return
|
||||
m = AutoModelForSequenceClassification.from_pretrained(
|
||||
MODEL_ID, torch_dtype=DTYPE, trust_remote_code=True
|
||||
)
|
||||
m.to(DEVICE)
|
||||
m.eval()
|
||||
_model = m
|
||||
|
||||
|
||||
@app.on_event("startup")
|
||||
def startup():
|
||||
_load()
|
||||
|
||||
|
||||
class RerankRequest(BaseModel):
|
||||
query: str
|
||||
texts: list[str]
|
||||
return_text: bool = False
|
||||
truncate: bool | None = None
|
||||
raw_scores: bool | None = None
|
||||
|
||||
|
||||
@app.get("/info")
|
||||
def info():
|
||||
# Hindsight only needs a 200 JSON here to consider the server initialized.
|
||||
return {
|
||||
"model_id": MODEL_ID,
|
||||
"model_dtype": str(DTYPE).replace("torch.", ""),
|
||||
"model_type": {"reranker": {}},
|
||||
"max_input_length": MAX_LENGTH,
|
||||
"device": DEVICE,
|
||||
}
|
||||
|
||||
|
||||
@app.get("/health")
|
||||
def health():
|
||||
return {"status": "ok" if _model is not None else "loading"}
|
||||
|
||||
|
||||
@app.post("/rerank")
|
||||
def rerank(req: RerankRequest):
|
||||
if not req.texts:
|
||||
return []
|
||||
pairs = [[req.query, t] for t in req.texts]
|
||||
with torch.no_grad():
|
||||
# jina-reranker-v2 exposes compute_score (batches + moves to device).
|
||||
scores = _model.compute_score(pairs, max_length=MAX_LENGTH)
|
||||
if not isinstance(scores, list):
|
||||
scores = [scores]
|
||||
results = [{"index": i, "score": float(s)} for i, s in enumerate(scores)]
|
||||
results.sort(key=lambda r: r["score"], reverse=True)
|
||||
return results
|
||||
Reference in New Issue
Block a user