Add TEI rerank provider and batch chunking

.dockerignore

@@ -0,0 +1,19 @@
+.git
+.venv
+.coverage
+.pytest_cache
+__pycache__
+*.py[cod]
+*.egg-info
+test_data
+data
+
+node_modules
+web-studio/node_modules
+web-studio/dist
+openviking/web_studio/dist
+
+.mypy_cache
+.ruff_cache
+.pytest_cache
+.DS_Store

docs/en/guides/01-configuration.md

@@ -809,7 +809,7 @@ The `PermissionDeniedError` message names the exact key to add for the blocked h
 
 ### rerank
 
-Reranking model for search result refinement. Supports VikingDB (Volcengine), Cohere, and OpenAI-compatible APIs.
+Reranking model for search result refinement. Supports VikingDB (Volcengine), Cohere, OpenAI-compatible APIs, LiteLLM, and Hugging Face Text Embeddings Inference (TEI).
 
 **Volcengine (VikingDB):**
 
@@ -840,25 +840,45 @@ Reranking model for search result refinement. Supports VikingDB (Volcengine), Co
 }
 ```
 
+**Hugging Face Text Embeddings Inference (TEI):**
+
+```json
+{
+  "rerank": {
+    "provider": "tei",
+    "api_base": "http://localhost:8080",
+    "api_key": "optional-tei-api-key",
+    "model": "BAAI/bge-reranker-v2-m3",
+    "batch_size": 32,
+    "threshold": 0.05
+  }
+}
+```
+
+For TEI, `api_base` may be either the server base URL (`http://localhost:8080`) or the full rerank endpoint (`http://localhost:8080/rerank`). `api_key` is optional and is sent as a Bearer token when configured. TEI is auto-detected when only `api_base` is set; if your TEI deployment also uses `api_key`, set `"provider": "tei"` explicitly so it is not treated as an OpenAI-compatible rerank endpoint.
+
 **Parameters**
 
 | Parameter | Type | Description |
 |-----------|------|-------------|
-| `provider` | str | `"vikingdb"`, `"cohere"`, or `"openai"`. Auto-detected if omitted. |
+| `provider` | str | `"vikingdb"`, `"cohere"`, `"openai"`, `"litellm"`, or `"tei"`. Auto-detected if omitted. |
 | `ak` | str | VikingDB Access Key (vikingdb provider only) |
 | `sk` | str | VikingDB Secret Key (vikingdb provider only) |
 | `model_name` | str | Model name (vikingdb provider only, default: `doubao-seed-rerank`) |
-| `api_key` | str | API key (for `openai` or `cohere` providers) |
-| `api_base` | str | Endpoint URL (for `openai` provider) |
-| `model` | str | Model name (for `openai` providers) |
+| `api_key` | str | API key (for `openai` or `cohere` providers, optional for `tei`) |
+| `api_base` | str | Endpoint URL (for `openai` provider) or TEI base/rerank URL (for `tei`) |
+| `model` | str | Model name (for `openai` and `litellm`; optional label for TEI usage tracking) |
 | `timeout` | float | HTTP request timeout in seconds for OpenAI-compatible providers. Increase for slow or cold-starting local rerank servers. Default: `30.0` |
+| `batch_size` | int | Maximum number of documents sent in a single rerank provider call. TEI deployments commonly cap this at `32`; larger candidate sets are chunked. Default: `32` |
 | `threshold` | float | Score threshold between `0.0` and `1.0`; results below this are filtered out. Default: `0.1` |
-| `extra_headers` | object | Custom HTTP headers (for OpenAI-compatible providers, optional) |
+| `extra_headers` | object | Custom HTTP headers (for OpenAI-compatible or TEI providers, optional) |
 
 **Supported providers:**
 - `vikingdb`: Volcengine VikingDB Rerank API (uses AK/SK)
 - `cohere`: Cohere Rerank API
 - `openai`: OpenAI-compatible Rerank API
+- `litellm`: LiteLLM rerank API
+- `tei`: Hugging Face Text Embeddings Inference rerank API
 
 If rerank is not configured, search uses vector similarity only.
 

openviking/models/rerank/__init__.py

@@ -8,12 +8,14 @@
 - cohere: Cohere Rerank v3.5 API
 - litellm: LiteLLM rerank (supports multiple providers)
 - openai: OpenAI-compatible rerank API
+- tei: Hugging Face Text Embeddings Inference rerank API
 """
 
 from openviking.models.rerank.base import RerankBase
 from openviking.models.rerank.cohere_rerank import CohereRerankClient
 from openviking.models.rerank.litellm_rerank import LiteLLMRerankClient
 from openviking.models.rerank.openai_rerank import OpenAIRerankClient
+from openviking.models.rerank.tei_rerank import TEIRerankClient
 from openviking.models.rerank.volcengine_rerank import RerankClient
 
 __all__ = [
@@ -22,4 +24,5 @@
     "CohereRerankClient",
     "LiteLLMRerankClient",
     "OpenAIRerankClient",
+    "TEIRerankClient",
 ]

openviking/models/rerank/tei_rerank.py

@@ -0,0 +1,178 @@
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: AGPL-3.0
+"""
+Text Embeddings Inference rerank API client.
+
+Hugging Face Text Embeddings Inference (TEI) exposes rerank models through a
+provider-specific `/rerank` endpoint. Its request/response shape differs from
+OpenAI-compatible rerank APIs, so it needs a dedicated adapter.
+"""
+
+import time
+from typing import Dict, List, Optional
+
+import requests
+
+from openviking.models.rerank.base import RerankBase
+from openviking_cli.utils import get_logger
+
+logger = get_logger(__name__)
+
+
+class TEIRerankClient(RerankBase):
+    """
+    TEI rerank API client.
+
+    TEI accepts `texts` and returns a list of `{index, score}` items:
+    https://huggingface.co/docs/text-embeddings-inference
+    """
+
+    def __init__(
+        self,
+        api_base: str,
+        api_key: Optional[str] = None,
+        model_name: Optional[str] = None,
+        extra_headers: Optional[Dict[str, str]] = None,
+        batch_size: int = 32,
+    ) -> None:
+        """
+        Initialize TEI rerank client.
+
+        Args:
+            api_base: TEI base URL (`http://host:port`) or full rerank endpoint.
+            api_key: Optional Bearer token for TEI deployments that enforce auth.
+            model_name: Optional model name used for usage tracking.
+            extra_headers: Optional extra headers for API requests.
+            batch_size: Maximum number of documents to send per TEI request.
+        """
+        super().__init__()
+        self.api_base = api_base
+        self.api_key = api_key
+        self.model_name = model_name
+        self.extra_headers = extra_headers or {}
+        self.batch_size = max(1, int(batch_size))
+        self.provider = "tei"
+
+    @property
+    def rerank_url(self) -> str:
+        """Return the full TEI rerank URL while accepting base or endpoint config."""
+        base = self.api_base.rstrip("/")
+        if base.endswith("/rerank"):
+            return base
+        return f"{base}/rerank"
+
+    def rerank_batch(self, query: str, documents: List[str]) -> Optional[List[float]]:
+        """
+        Batch rerank documents against a query.
+
+        Args:
+            query: Query text
+            documents: List of document texts to rank
+
+        Returns:
+            List of rerank scores in the same order as input documents, or None
+            when rerank fails and the caller should fall back.
+        """
+        if not documents:
+            return []
+
+        scores = [0.0] * len(documents)
+        for start in range(0, len(documents), self.batch_size):
+            chunk = documents[start : start + self.batch_size]
+            chunk_scores = self._rerank_chunk(query, chunk)
+            if chunk_scores is None:
+                return None
+            scores[start : start + len(chunk_scores)] = chunk_scores
+
+        logger.debug(
+            "[TEIRerankClient] Reranked %s documents in %s request(s)",
+            len(documents),
+            (len(documents) + self.batch_size - 1) // self.batch_size,
+        )
+        return scores
+
+    def _rerank_chunk(self, query: str, documents: List[str]) -> Optional[List[float]]:
+        """Rerank one TEI-sized chunk and return scores in chunk-local order."""
+
+        req_body = {
+            "query": query,
+            "texts": documents,
+            "raw_scores": False,
+        }
+
+        try:
+            headers = {"Content-Type": "application/json"}
+            if self.api_key:
+                headers["Authorization"] = f"Bearer {self.api_key}"
+            if self.extra_headers:
+                headers.update(self.extra_headers)
+
+            started = time.monotonic()
+            response = requests.post(
+                url=self.rerank_url,
+                headers=headers,
+                json=req_body,
+                timeout=30,
+            )
+            response.raise_for_status()
+            result = response.json()
+
+            self._extract_and_update_token_usage(
+                {"results": result} if isinstance(result, list) else result,
+                query,
+                documents,
+                duration_seconds=time.monotonic() - started,
+            )
+
+            results = self._extract_results(result)
+            if not results:
+                logger.warning(f"[TEIRerankClient] Unexpected response format: {result}")
+                return None
+
+            scores = [0.0] * len(documents)
+            for item in results:
+                idx = item.get("index")
+                if idx is None or not (0 <= idx < len(documents)):
+                    logger.warning(
+                        "[TEIRerankClient] Out-of-bounds or missing index in result: %s", item
+                    )
+                    return None
+                scores[idx] = float(item.get("score", item.get("relevance_score", 0.0)))
+
+            return scores
+
+        except Exception as e:
+            logger.error(f"[TEIRerankClient] Rerank failed: {e}")
+            return None
+
+    @staticmethod
+    def _extract_results(result) -> Optional[List[dict]]:
+        """Extract TEI rerank rows from supported response shapes."""
+        if isinstance(result, list):
+            return result
+        if isinstance(result, dict):
+            rows = result.get("results")
+            if isinstance(rows, list):
+                return rows
+        return None
+
+    @classmethod
+    def from_config(cls, config) -> Optional["TEIRerankClient"]:
+        """
+        Create TEIRerankClient from RerankConfig.
+
+        Args:
+            config: RerankConfig instance with provider='tei'
+
+        Returns:
+            TEIRerankClient instance or None if config is not available
+        """
+        if not config or not config.is_available():
+            return None
+        return cls(
+            api_base=config.api_base,
+            api_key=config.api_key,
+            model_name=config.model,
+            extra_headers=config.extra_headers,
+            batch_size=config.batch_size,
+        )

openviking/models/rerank/volcengine_rerank.py

@@ -219,6 +219,11 @@ def from_config(cls, config) -> Optional["RerankClient"]:
 
             return OpenAIRerankClient.from_config(config)
 
+        if provider == "tei":
+            from openviking.models.rerank.tei_rerank import TEIRerankClient
+
+            return TEIRerankClient.from_config(config)
+
         return cls(
             ak=config.ak,
             sk=config.sk,

openviking/server/routers/admin.py

@@ -66,6 +66,25 @@ class MigrateLegacyDataRequest(BaseModel):
     action: str = "migrate"
 
 
+class _LegacyCleanupUserIdentifier:
+    """Minimal user identity for deleting pre-validation legacy accounts."""
+
+    def __init__(self, account_id: str, user_id: str):
+        self._account_id = account_id
+        self._user_id = user_id
+
+    @property
+    def account_id(self) -> str:
+        return self._account_id
+
+    @property
+    def user_id(self) -> str:
+        return self._user_id
+
+    def user_space_name(self) -> str:
+        return self._user_id
+
+
 def _get_api_key_manager(request: Request):
     """Get APIKeyManager from app state."""
     return get_api_key_manager_or_raise(request)
@@ -133,6 +152,17 @@ def _validate_register_user_role(ctx: RequestContext, role: str) -> Role:
     return resolved_role
 
 
+def _cleanup_user_identifier(account_id: str) -> UserIdentifier | _LegacyCleanupUserIdentifier:
+    """Build a cleanup identity, allowing already-existing legacy account ids."""
+    try:
+        return UserIdentifier(account_id, "system")
+    except ValueError:
+        logger.warning(
+            "Using legacy cleanup identity for non-conforming account_id=%s", account_id
+        )
+        return _LegacyCleanupUserIdentifier(account_id, "system")
+
+
 async def _run_legacy_migration_task(
     task_id: str,
     migration: LegacyDataMigration,
@@ -264,7 +294,7 @@ async def delete_account(
 
     # Build a ROOT-level context scoped to the target account for cleanup
     cleanup_ctx = RequestContext(
-        user=UserIdentifier(account_id, "system"),
+        user=_cleanup_user_identifier(account_id),
         role=Role.ROOT,
     )