Add codebase embeddings and reranking endpoints

app/embeddings.py

@@ -0,0 +1,168 @@
+"""OpenAI-compatible codebase embedding request and response contracts."""
+from __future__ import annotations
+
+import base64
+import binascii
+import math
+import struct
+
+from app.site_routing import EMBEDDING_ENDPOINTS
+
+MODEL_ID = "bge-m3"
+PUBLIC_MODEL_ID = "baai/bge-m3"
+QWEN_MODEL_ID = "Qwen/Qwen3-Embedding-0.6B"
+PROFILES = frozenset(EMBEDDING_ENDPOINTS)
+DIMENSIONS = 1024
+MAX_INPUT_ITEMS = 2048
+
+
+class EmbeddingRequestError(ValueError):
+    def __init__(self, message: str, param: str):
+        super().__init__(message)
+        self.param = param
+
+
+class EmbeddingUpstreamError(Exception):
+    def __init__(self, code, message: str):
+        super().__init__(message)
+        self.upstream_code = code
+        self.message = message
+
+
+def model_details() -> list:
+    return [
+        {"id": MODEL_ID, "credits": None, "credits_by_profile": {}, "embedding": True},
+        {"id": QWEN_MODEL_ID, "credits": None, "credits_by_profile": {}, "embedding": True},
+    ]
+
+
+def is_embedding_model(model) -> bool:
+    return model in (MODEL_ID, QWEN_MODEL_ID)
+
+
+def codebase_model(model: str) -> str:
+    if model == MODEL_ID:
+        return "model2"
+    if model == QWEN_MODEL_ID:
+        return model
+    raise ValueError("不支持的 codebase embedding 模型")
+
+
+def _token_ids(values) -> list[int]:
+    if not values or any(type(value) is not int or not 0 <= value <= 2**31 - 1 for value in values):
+        raise EmbeddingRequestError("token ID 必须是非空整数数组", "input")
+    return values
+
+
+def _input(value) -> int:
+    if isinstance(value, str):
+        if not value:
+            raise EmbeddingRequestError("input 不能为空字符串", "input")
+        return 1
+    if not isinstance(value, list) or not value:
+        raise EmbeddingRequestError("input 必须是字符串、字符串数组或 token ID 数组", "input")
+    if all(isinstance(item, str) for item in value):
+        if len(value) > MAX_INPUT_ITEMS or any(not item for item in value):
+            raise EmbeddingRequestError(f"字符串数组须包含 1–{MAX_INPUT_ITEMS} 个非空字符串", "input")
+        return len(value)
+    if all(type(item) is int for item in value):
+        _token_ids(value)
+        return 1
+    if all(isinstance(item, list) for item in value):
+        if len(value) > MAX_INPUT_ITEMS:
+            raise EmbeddingRequestError(f"token 数组最多 {MAX_INPUT_ITEMS} 项", "input")
+        for item in value:
+            _token_ids(item)
+        return len(value)
+    raise EmbeddingRequestError("input 数组的元素类型必须一致", "input")
+
+
+def request_body(payload, model: str) -> tuple[dict, str, int]:
+    """Validate public input and pin the upstream representation to float vectors."""
+    if not isinstance(payload, dict):
+        raise EmbeddingRequestError("请求体必须是 JSON 对象", "")
+    encoding = payload.get("encoding_format", "float")
+    if encoding not in ("float", "base64"):
+        raise EmbeddingRequestError("encoding_format 必须是 float 或 base64", "encoding_format")
+    dimensions = payload.get("dimensions")
+    if dimensions is not None and (type(dimensions) is not int or dimensions != DIMENSIONS):
+        raise EmbeddingRequestError(f"当前向量模型的 dimensions 仅支持 {DIMENSIONS}", "dimensions")
+    user = payload.get("user")
+    if user is not None and (not isinstance(user, str) or len(user) > 512):
+        raise EmbeddingRequestError("user 必须是最多 512 字符的字符串", "user")
+    count = _input(payload.get("input"))
+    return {"model": model, "input": payload["input"], "encoding_format": "float"}, encoding, count
+
+
+def _numbers(value) -> list[float]:
+    if not isinstance(value, list) or len(value) != DIMENSIONS:
+        raise ValueError("embedding 响应的向量维度无效")
+    if any(type(number) not in (int, float) or not math.isfinite(number) for number in value):
+        raise ValueError("embedding 响应包含无效向量数值")
+    return value
+
+
+def _vector(value) -> list[float]:
+    if isinstance(value, list):
+        return _numbers(value)
+    if not isinstance(value, str):
+        raise ValueError("embedding 响应缺少向量")
+    try:
+        raw = base64.b64decode(value, validate=True)
+    except (ValueError, binascii.Error):
+        raise ValueError("embedding 响应的 base64 向量无效") from None
+    if len(raw) != DIMENSIONS * 4:
+        raise ValueError("embedding 响应的向量维度无效")
+    try:
+        return list(struct.unpack(f"<{DIMENSIONS}f", raw))
+    except struct.error:
+        raise ValueError("embedding 响应的 base64 向量无效") from None
+
+
+def _base64(values) -> str:
+    try:
+        return base64.b64encode(struct.pack(f"<{len(values)}f", *values)).decode("ascii")
+    except (struct.error, OverflowError):
+        raise ValueError("embedding 向量超出 float32 范围") from None
+
+
+def _usage(payload) -> dict:
+    source = payload.get("usage") if isinstance(payload.get("usage"), dict) else {}
+    result = {}
+    for key in ("prompt_tokens", "total_tokens", "credit"):
+        value = source.get(key)
+        if key == "prompt_tokens" and value is None:
+            value = source.get("input_tokens")
+        if type(value) in (int, float) and math.isfinite(value) and value >= 0:
+            result[key] = value
+    return result
+
+
+def parse_reply(payload, model: str, encoding: str, expected_count: int) -> dict:
+    """Validate every vector and restore OpenAI ordering and encoding."""
+    if not isinstance(payload, dict):
+        raise ValueError("embedding 响应必须是 JSON 对象")
+    if payload.get("code") not in (None, 0):
+        message = payload.get("msg") or payload.get("message") or "上游 embedding 接口错误"
+        raise EmbeddingUpstreamError(payload.get("code"), str(message)[:500])
+    items = payload.get("data")
+    if not isinstance(items, list) or len(items) != expected_count:
+        raise ValueError("embedding 响应的条目数量与请求不一致")
+    data = []
+    for position, item in enumerate(items):
+        if not isinstance(item, dict):
+            raise ValueError("embedding 响应包含无效条目")
+        index = item.get("index", position)
+        if type(index) is not int:
+            raise ValueError("embedding 响应的索引无效")
+        values = _vector(item.get("embedding"))
+        data.append((index, {"object": "embedding", "index": index,
+                             "embedding": _base64(values) if encoding == "base64" else values}))
+    if sorted(index for index, _item in data) != list(range(expected_count)):
+        raise ValueError("embedding 响应的索引不完整")
+    result = {"object": "list", "data": [item for _index, item in sorted(data, key=lambda row: row[0])],
+              "model": model}
+    usage = _usage(payload)
+    if usage:
+        result["usage"] = usage
+    return result

app/gateway_management.py

@@ -355,6 +355,11 @@ def admin_model_inventory(self):
         for media in self.gateway.media_model_details():
             row = facts.get(media["id"])
             facts[media["id"]] = {**media, **(row or {})}
+        for embedding in self.gateway.embedding_model_details():
+            row = facts.get(embedding["id"])
+            facts[embedding["id"]] = {**embedding, **(row or {})}
+        for reranker in self.gateway.rerank_model_details():
+            facts[reranker["id"]] = {**reranker, **(facts.get(reranker["id"]) or {})}
         result = []
         for source, row in facts.items():
             rule = model_policy.rule_for(self.CONFIG, source)
@@ -375,6 +380,8 @@ def admin_model_inventory(self):
     def admin_rate_probe_model(self, source):
         """Probe one rule's upstream model on every product that serves it, in the background."""
         upstream = model_policy.rule_for(self.CONFIG, source)["upstream_id"]
+        if self.gateway.embeddings.is_embedding_model(upstream) or self.gateway.rerank.is_rerank_model(upstream):
+            raise ValueError("向量与重排序模型不支持聊天倍率探测")
         return model_rates.start_probe(self.gateway, upstream, scope=source)
 
     def admin_rate_probe_all(self):
@@ -394,6 +401,8 @@ def admin_model_preview(self, source, rule):
         if pool is None:
             return {"candidates": [], "excluded": []}
         media = bool(self.gateway._media_all_kinds(upstream))
+        embedding = self.gateway.embeddings.is_embedding_model(upstream)
+        reranking = self.gateway.rerank.is_rerank_model(upstream)
         policy = model_policy.view(self.CONFIG)
         with pool._lock:
             for entry in pool.entries():
@@ -404,6 +413,8 @@ def admin_model_preview(self, source, rule):
                     reason = self.gateway._media_account_reason(entry, upstream, rule=rule, policy=policy)
                 elif not model_policy.credential_enabled(self.CONFIG, entry):
                     reason = "人工停用"
+                elif embedding or reranking:
+                    reason = self.gateway._codebase_account_reason(entry, upstream, rule=rule, policy=policy)
                 elif not rule.get("enabled", True):
                     reason = "模型已停用"
                 elif not model_policy.route_allowed(self.CONFIG, entry, source, rule=rule):

app/model_policy.py

@@ -3,7 +3,7 @@
 
 from fastapi import HTTPException
 
-from app import growth_tasks, perf_metrics, travel
+from app import growth_tasks, perf_metrics, travel, embeddings
 from app.model_capabilities import route_metadata
 from app.api_keys import api_principal
 _request_policy = ContextVar("gateway_model_policy", default=None)
@@ -113,13 +113,15 @@ def growth_auto_tasks(config, entry, policy=None):
     mapping = _policy(config, policy)["credentials"].get(entry.get("account_key"), {}).get("growth_auto") or {}
     return frozenset(code for code, enabled in mapping.items() if enabled is True)
 
-
 def default_rule(source):
+    if source in (embeddings.MODEL_ID, embeddings.PUBLIC_MODEL_ID):
+        return {"public_id": embeddings.PUBLIC_MODEL_ID, "upstream_id": embeddings.MODEL_ID,
+                "custom": False, "enabled": True, "keep_original": False,
+                "region": None, "profile": None, "credential_ids": []}
     return {"public_id": source, "upstream_id": source, "custom": False,
             "enabled": True, "keep_original": False,
             "region": None, "profile": None, "credential_ids": []}
 
-
 def rule_for(config, source, policy=None):
     return {**default_rule(source), **_policy(config, policy)["models"].get(source, {})}
 
@@ -128,6 +130,8 @@ def source_for(config, name):
     """Resolve a public model name or automatic alias to its policy source."""
     if name == "default-model":
         return "auto"
+    if name in (embeddings.MODEL_ID, embeddings.PUBLIC_MODEL_ID):
+        return embeddings.MODEL_ID
     data = snapshot(config)
     for source, rule in data.get("models", {}).items():
         if rule.get("public_id", source) == name:
@@ -137,7 +141,6 @@ def source_for(config, name):
     return name
 
 
-
 def _restriction_source(source):
     """Use the canonical automatic-model source for key restrictions."""
     return "auto" if source == "default-model" else source
@@ -177,7 +180,9 @@ def resolve(config, public_model):
         return state["upstream"]
     data = snapshot(config)
     source = next((key for key, value in data["models"].items()
-                   if value.get("public_id", key) == public_model), public_model)
+                   if value.get("public_id", key) == public_model),
+                  embeddings.MODEL_ID if public_model in (embeddings.MODEL_ID, embeddings.PUBLIC_MODEL_ID)
+                  else public_model)
     rule = {**default_rule(source), **data["models"].get(source, {})}
     if principal is not None and principal.models is not None and _restriction_source(source) not in principal.models:
         raise HTTPException(status_code=403, detail={"error": {
@@ -234,12 +239,17 @@ def public_name(fallback):
 
 
 def public_details(gateway, region=None):
-    """Publish chat and media routes using their own account eligibility checks."""
+    """Publish chat, media and embedding routes using their own eligibility checks."""
     config = gateway.CONFIG
     known = {item["id"]: item for item in gateway.current_model_details(region)}
     for media in gateway.media_model_details():
         row = known.get(media["id"])
         known[media["id"]] = {**media, **(row or {})}
+    for embedding in gateway.embedding_model_details():
+        row = known.get(embedding["id"])
+        known[embedding["id"]] = {**embedding, **(row or {})}
+    for reranker in gateway.rerank_model_details():
+        known[reranker["id"]] = {**reranker, **(known.get(reranker["id"]) or {})}
     policy = view(config)
     policies = policy.data["models"]
     owners = {rule.get("public_id", source): source for source, rule in policies.items()}
@@ -253,6 +263,8 @@ def public_details(gateway, region=None):
         if not rule["enabled"] or item is None:
             continue
         media = bool(item.get("media_kinds"))
+        embedding = bool(item.get("embedding"))
+        reranking = bool(item.get("rerank"))
         token = _request_policy.set({"source": source, "upstream": upstream,
                                      "public_model": rule["public_id"], "rule": rule})
         try:
@@ -272,13 +284,29 @@ def public_details(gateway, region=None):
                         if (profile in gateway.media_generation.MEDIA_PROFILES
                                 and route_allowed(config, entry, upstream, rule=rule, policy=policy)):
                             candidates = [entry]
+            elif embedding or reranking:
+                candidates = [entry for entry in entries if gateway._codebase_account_reason(
+                    entry, upstream, rule=rule, policy=policy, region=region) is None]
+                if pool is None:
+                    cm = config.get("cred")
+                    candidates = []
+                    if cm is not None:
+                        with cm._lock:
+                            headers = cm.get_headers()
+                        profile = gateway.profile_for_headers(headers)
+                        identity = gateway.account_key(profile, headers.get("X-User-Id"),
+                                                       headers.get("X-Enterprise-Id"))
+                        entry = {"profile": profile, "account_key": identity}
+                        if (profile in gateway.embeddings.PROFILES
+                                and route_allowed(config, entry, upstream, rule=rule, policy=policy)):
+                            candidates = [entry]
             else:
                 candidates = [entry for entry in entries if route_allowed(
                     config, entry, upstream, rule=rule, policy=policy)
                     and pool._eligible(entry, upstream, region=region)]
         finally:
             _request_policy.reset(token)
-        if not candidates and (pool is not None or media):
+        if not candidates and (pool is not None or media or embedding or reranking):
             continue
         prices, measured = {}, set()
         accounts = config.get("account_catalogs")
@@ -309,7 +337,7 @@ def public_details(gateway, region=None):
             best = min(prices.values(), default=None)
             row = {"id": name, "credits": best, "credits_by_profile": prices,
                    **route_metadata(gateway, upstream, entries=candidates, rule=rule, region=region)}
-            for key in ("media_kinds", "media_sources", "media_defaults"):
+            for key in ("media_kinds", "media_sources", "media_defaults", "embedding", "rerank"):
                 if key in item:
                     row[key] = item[key]
             if any(prices[profile] == best for profile in measured):

app/request_context.py

@@ -12,7 +12,8 @@
 # performs token estimation only and must not claim an upstream session.
 PATHS = {"/v1/chat/completions": "chat", "/v1/responses": "responses", "/v1/messages": "messages",
          "/v1/search": "search", "/v1/fetch": "fetch", "/v1/mcp": "mcp",
-         "/v1/images/generations": "images", "/v1/images/edits": "images", "/v1/videos": "videos"}
+         "/v1/images/generations": "images", "/v1/images/edits": "images", "/v1/videos": "videos",
+         "/v1/embeddings": "embeddings", "/v1/rerank": "rerank"}
 _SCOPE_KEY = "codebuddy.request_context"
 _current = ContextVar("request_context", default=None)