app.py

"""
Hugging Face Spaces (Docker) — Contrastive 질의 임베딩 API.

모델은 HF Hub에서 로드 (환경변수 HF_MODEL_ID, 예: <org>/<model-id>).
POST로 텍스트 받아 768-dim 벡터 반환. 백엔드 RAG_CONTRASTIVE_EMBEDDING_URL 연동용.

환경변수:
- HF_MODEL_ID: SentenceTransformer 모델 ID (필수)
- HF_TOKEN: 비공개 리포 접근용 토큰 (선택)
"""
import os
from typing import List, Union

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel

MODEL_ID = (os.environ.get("HF_MODEL_ID") or "").strip()
HF_TOKEN = (os.environ.get("HF_TOKEN") or "").strip() or None

model = None

app = FastAPI(title="Contrastive Embedding API")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)


class EmbedRequest(BaseModel):
    """백엔드 contrastive_retriever가 보내는 형식."""
    inputs: str


def load_model() -> None:
    global model
    from sentence_transformers import SentenceTransformer

    print(f"[Embed] 모델 로딩: {MODEL_ID}")
    model = SentenceTransformer(MODEL_ID, token=HF_TOKEN)
    print("[Embed] 로드 완료.")


@app.on_event("startup")
def startup() -> None:
    if MODEL_ID:
        load_model()
    else:
        print("[Embed] HF_MODEL_ID 미설정. 설정 후 재시작.")


@app.post("/")
def embed(body: EmbedRequest) -> Union[List[List[float]], dict]:
    """
    질의 문자열을 768-dim 벡터로 변환.
    응답: [[float,...]] (백엔드가 기대하는 형식) 또는 {"embedding": [float,...]}
    """
    global model
    if not model:
        return {"embedding": [], "error": "HF_MODEL_ID not set"}

    text = (body.inputs or "").strip()
    if not text:
        return {"embedding": [], "error": "empty inputs"}

    vec = model.encode([text], normalize_embeddings=True)
    # 백엔드: list of list 또는 {"embedding": ...} 둘 다 처리
    return vec.tolist()


@app.get("/")
def root():
    """
    Hugging Face Spaces / 컨테이너 헬스체크용 루트 엔드포인트.

    - GET /            → 200 OK
    - GET /?logs=...   → 200 OK (path는 / 그대로, query만 다름)
    - UptimeRobot HTTP(S) 모니터의 HEAD / 도 FastAPI가 자동 처리
    """
    return {
        "status": "ok",
        "service": "contrastive-embedding",
        "model_loaded": model is not None,
    }


@app.get("/health")
def health():
    """단순 헬스 체크 엔드포인트."""
    return {"status": "ok", "model_loaded": model is not None}