add embedding service

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.12-slim
+FROM python:3.11-slim
 
 LABEL maintainer="All API Collection"
 LABEL description="All API Collection - Document extraction, conversion, and database query API"
@@ -24,6 +24,8 @@ RUN pip install --no-cache-dir "youtube-transcript-api>=1.2.4"
 
 COPY --chown=appuser:appuser . .
 
+RUN mkdir -p /app/models && python3 -c "from huggingface_hub import snapshot_download; snapshot_download(repo_id='BAAI/bge-small-en-v1.5', local_dir='/app/models/bge-384'); snapshot_download(repo_id='BAAI/bge-base-en-v1.5', local_dir='/app/models/bge-768'); snapshot_download(repo_id='BAAI/bge-large-en-v1.5', local_dir='/app/models/bge-1024')" && chown -R appuser:appuser /app/models
+
 RUN mkdir -p /app/logs && \
     chown -R appuser:appuser /app/logs
 

app/api/deps.py

@@ -6,6 +6,7 @@ from app.core.security import require_api_key
 from app.services.auth_service import AuthService
 from app.services.converter_service import ConverterService
 from app.services.database_service import DatabaseService
+from app.services.embeddings_service import EmbeddingService
 from app.services.extraction_service import ExtractionService
 from app.services.ocr_service import OCRService
 
@@ -30,5 +31,10 @@ def get_database_service() -> DatabaseService:
     return DatabaseService()
 
 
+def get_embeddings_service() -> EmbeddingService:
+    from app.api.server import _embedding_service
+    return _embedding_service
+
+
 def require_auth(token: str = Depends(require_api_key)) -> str:
     return token

app/api/server.py

@@ -10,11 +10,14 @@ from fastapi.middleware.gzip import GZipMiddleware
 from app.config import get_settings
 from app.core.database import pool_manager
 from app.core.logger import get_logger
+from app.services.embeddings_service import EmbeddingService
 from app.api.v1.router import api_v1_router
 
 _logger = get_logger(__name__)
 _settings = get_settings()
 
+_embedding_service: EmbeddingService = EmbeddingService()
+
 
 async def _self_ping():
     import httpx
@@ -34,6 +37,11 @@ async def _self_ping():
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    _logger.info("Initializing embedding service (loading all models)...")
+    loop = asyncio.get_running_loop()
+    await loop.run_in_executor(None, _embedding_service.load_all_models)
+    _logger.info("Embedding service initialized with dims: %s", _embedding_service.loaded_dimensions)
+
     asyncio.create_task(_self_ping())
     yield
     _logger.info("Shutting down database connection pools...")
@@ -51,6 +59,7 @@ def create_application() -> FastAPI:
             {"name": "Convert", "description": "Single-file and single-URL conversion"},
             {"name": "Batch", "description": "Bulk conversion of files and URLs"},
             {"name": "System", "description": "Health, info, and supported formats"},
+            {"name": "Embeddings", "description": "Text embedding generation using transformer models"},
             {"name": "Verify", "description": "Phone number and identity verification"},
         ],
         lifespan=lifespan,

app/api/v1/embeddings.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import asyncio
+import concurrent.futures
+import os
+
+from fastapi import APIRouter, Depends, HTTPException
+
+from app.api.deps import require_auth, get_embeddings_service
+from app.core.logger import get_logger
+from app.models.schemas import EmbeddingRequest, EmbeddingResponse
+from app.services.embeddings_service import EmbeddingService
+
+router = APIRouter()
+_logger = get_logger(__name__)
+_MAX_WORKERS = min(32, (os.cpu_count() or 1) + 4)
+_thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=_MAX_WORKERS)
+
+
+@router.post(
+    "/embeddings",
+    response_model=EmbeddingResponse,
+    summary="Generate text embeddings",
+)
+async def create_embeddings(
+    body: EmbeddingRequest,
+    token: str = Depends(require_auth),
+    embedding_service: EmbeddingService = Depends(get_embeddings_service),
+) -> EmbeddingResponse:
+    _logger.info("Embedding request: dim=%s, content_len=%s", body.dimension, len(body.content))
+
+    if not embedding_service.is_loaded(body.dimension):
+        _logger.error("Model dim=%s not loaded. Loaded: %s", body.dimension, embedding_service.loaded_dimensions)
+        raise HTTPException(
+            status_code=503,
+            detail={
+                "success": False,
+                "message": f"Model for dimension {body.dimension} not loaded. Loaded: {embedding_service.loaded_dimensions}",
+            },
+        )
+
+    try:
+        loop = asyncio.get_running_loop()
+        embeddings = await loop.run_in_executor(
+            _thread_pool,
+            embedding_service.generate_embedding,
+            body.content,
+            body.dimension,
+        )
+
+        _logger.info("Embedding success: dim=%s, vector_len=%s", body.dimension, len(embeddings))
+        return EmbeddingResponse(
+            success=True,
+            embeddings=embeddings,
+            dimension=body.dimension,
+        )
+    except Exception as exc:
+        _logger.error("Embedding error: %s", exc)
+        raise HTTPException(
+            status_code=500,
+            detail={"success": False, "message": str(exc)},
+        )

app/api/v1/router.py

@@ -2,7 +2,7 @@ from __future__ import annotations
 
 from fastapi import APIRouter
 
-from app.api.v1 import batch, convert, database, system
+from app.api.v1 import batch, convert, database, embeddings, system
 from app.api.verify import router as verify_router
 
 api_v1_router = APIRouter()
@@ -10,4 +10,5 @@ api_v1_router.include_router(convert.router, tags=["Convert"])
 api_v1_router.include_router(batch.router, tags=["Batch"])
 api_v1_router.include_router(system.router, tags=["System"])
 api_v1_router.include_router(database.router, tags=["Database"])
+api_v1_router.include_router(embeddings.router, tags=["Embeddings"])
 api_v1_router.include_router(verify_router, prefix="/verify", tags=["Verify"])

app/config.py

@@ -22,7 +22,7 @@ class Settings(BaseSettings):
     enable_colors: bool = True
 
     api_key: str = "changeme"
-    max_upload_bytes: int = 100 * 1024 * 1024
+    max_upload_bytes: int = 15 * 1024 * 1024
     max_batch_files: int = 10
     max_batch_urls: int = 20
 

app/models/__init__.py

@@ -7,6 +7,8 @@ from app.models.schemas import (
     BatchUrlRequest,
     ConversionMetadata,
     ConversionResponse,
+    EmbeddingRequest,
+    EmbeddingResponse,
     HealthResponse,
     InfoResponse,
     SpacyLabelsResponse,
@@ -19,6 +21,8 @@ __all__ = [
     "ConversionResult",
     "ConversionMetadata",
     "ConversionResponse",
+    "EmbeddingRequest",
+    "EmbeddingResponse",
     "UrlRequest",
     "BatchUrlRequest",
     "BatchFileResult",

app/models/schemas.py

@@ -222,3 +222,14 @@ class DatabaseQueryResponse(BaseModel):
     execution_time_ms: float
     results: Optional[List[StatementResultSchema]] = None
     error: Optional[DatabaseQueryError] = None
+
+
+class EmbeddingRequest(BaseModel):
+    content: str = Field(..., min_length=1, description="Text to embed")
+    dimension: int = Field(default=384, ge=384, le=1024, description="Target embedding dimension (384, 768, or 1024)")
+
+
+class EmbeddingResponse(BaseModel):
+    success: bool
+    embeddings: List[float]
+    dimension: int

app/services/__init__.py

@@ -2,12 +2,14 @@ from __future__ import annotations
 
 from app.services.auth_service import AuthService
 from app.services.converter_service import ConverterService
+from app.services.embeddings_service import EmbeddingService
 from app.services.extraction_service import ExtractionService
 from app.services.ocr_service import OCRService
 
 __all__ = [
     "AuthService",
     "ConverterService",
+    "EmbeddingService",
     "ExtractionService",
     "OCRService",
 ]

app/services/embeddings_service.py

@@ -0,0 +1,73 @@
+from __future__ import annotations
+
+import logging
+import os
+from typing import Dict, List, Optional
+
+import numpy as np
+from sentence_transformers import SentenceTransformer
+
+_logger = logging.getLogger(__name__)
+
+_MODEL_MAP: Dict[int, str] = {
+    384: "BAAI/bge-small-en-v1.5",
+    768: "BAAI/bge-base-en-v1.5",
+    1024: "BAAI/bge-large-en-v1.5",
+}
+
+
+class EmbeddingService:
+    def __init__(self, models_dir: Optional[str] = None) -> None:
+        self._models: Dict[int, SentenceTransformer] = {}
+        self._models_dir = models_dir or os.path.join(os.getcwd(), "models")
+        self._device = "cuda"
+        try:
+            import torch
+            if not torch.cuda.is_available():
+                self._device = "cpu"
+        except ImportError:
+            self._device = "cpu"
+
+        self._loaded_dimensions: List[int] = []
+
+    def load_model(self, dimension: int) -> None:
+        if dimension in self._models:
+            return
+        if dimension not in _MODEL_MAP:
+            raise ValueError(f"Unsupported dimension {dimension}. Supported: {list(_MODEL_MAP.keys())}")
+
+        model_name = _MODEL_MAP[dimension]
+        local_path = os.path.join(self._models_dir, f"bge-{dimension}")
+
+        _logger.info("Loading embedding model dim=%s from %s", dimension, local_path if os.path.isdir(local_path) else model_name)
+        model = SentenceTransformer(
+            local_path if os.path.isdir(local_path) else model_name,
+            device=self._device,
+        )
+        model.eval()
+        self._models[dimension] = model
+        self._loaded_dimensions.append(dimension)
+        _logger.info("Loaded embedding model dim=%s (device=%s)", dimension, self._device)
+
+    def load_all_models(self) -> None:
+        for dim in _MODEL_MAP:
+            self.load_model(dim)
+
+    def generate_embedding(self, text: str, dimension: int) -> List[float]:
+        if dimension not in self._models:
+            raise ValueError(f"Model for dimension {dimension} not loaded")
+        model = self._models[dimension]
+        result: np.ndarray = model.encode(
+            text,
+            normalize_embeddings=True,
+            convert_to_numpy=True,
+            show_progress_bar=False,
+        )
+        return result.tolist()
+
+    @property
+    def loaded_dimensions(self) -> List[int]:
+        return list(self._loaded_dimensions)
+
+    def is_loaded(self, dimension: int) -> bool:
+        return dimension in self._models

requirements.txt

@@ -11,6 +11,7 @@ onnxruntime>=1.18.0
 pillow>=10.0.0
 pypdfium2>=4.30.0
 pandas>=2.0.0
+sentence-transformers==3.3.1
 spacy>=3.7.0
 phonenumbers>=8.13.0