From 8b41b7a7f23d2614014781ae870bb3b581b533fb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=A2=A8=E8=8F=8A?= <277378677+maiphucgiang@users.noreply.github.com> Date: Thu, 29 Oct 2026 18:17:43 +0800 Subject: [PATCH] Add codebase embeddings and reranking endpoints - Route embeddings through the codebase service with BGE-M3 and Qwen aliases, bounded paid fallback, account-aware failover and usage handling. - Expose the fixed codebase reranker without unverified relevance scores, including model catalog, management and WebUI metadata. - Add preview simulations, English/Chinese client documentation and offline regression coverage. Tests: - uv run python -B -m pytest -q tests/test_embeddings.py tests/test_rerank.py tests/test_preview.py tests/test_observability.py tests/test_region_routing.py - cd web && vp check --- app/embeddings.py | 168 +++++++++++++++++ app/gateway_management.py | 11 ++ app/model_policy.py | 44 ++++- app/request_context.py | 3 +- app/rerank.py | 90 +++++++++ app/site_routing.py | 23 +++ converter.py | 353 ++++++++++++++++++++++++++++++++++- docs/clients.md | 24 +++ docs/clients.zh-CN.md | 24 +++ preview/upstream.py | 61 +++++- tests/test_embeddings.py | 258 +++++++++++++++++++++++++ tests/test_observability.py | 35 ++++ tests/test_preview.py | 28 +++ tests/test_region_routing.py | 3 +- tests/test_rerank.py | 159 ++++++++++++++++ web/src/api.ts | 2 + web/src/pages/Models.tsx | 2 + 17 files changed, 1274 insertions(+), 14 deletions(-) create mode 100644 app/embeddings.py create mode 100644 app/rerank.py create mode 100644 tests/test_embeddings.py create mode 100644 tests/test_rerank.py diff --git a/app/embeddings.py b/app/embeddings.py new file mode 100644 index 0000000..19da8ae --- /dev/null +++ b/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 diff --git a/app/gateway_management.py b/app/gateway_management.py index aa79393..fb14ddf 100644 --- a/app/gateway_management.py +++ b/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): diff --git a/app/model_policy.py b/app/model_policy.py index ad20cf8..2ce5e5b 100644 --- a/app/model_policy.py +++ b/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): diff --git a/app/request_context.py b/app/request_context.py index 1b36f7b..671a3ab 100644 --- a/app/request_context.py +++ b/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) diff --git a/app/rerank.py b/app/rerank.py new file mode 100644 index 0000000..f8a63a0 --- /dev/null +++ b/app/rerank.py @@ -0,0 +1,90 @@ +"""Public reranker contract for the fixed codebase reranking service.""" +from __future__ import annotations + +import math + +MODEL_ID = "reranker" +MAX_DOCUMENTS = 100 + + +class RerankRequestError(ValueError): + def __init__(self, message: str, param: str): + super().__init__(message) + self.param = param + + +class RerankUpstreamError(Exception): + def __init__(self, code, message: str): + super().__init__(message) + self.upstream_code = code + self.message = message + + +def model_details() -> dict: + return {"id": MODEL_ID, "credits": None, "credits_by_profile": {}, "rerank": True} + + +def is_rerank_model(model) -> bool: + return model == MODEL_ID + + +def request_body(payload) -> tuple[dict, list[str], int, bool]: + if not isinstance(payload, dict): + raise RerankRequestError("请求体必须是 JSON 对象", "") + query = payload.get("query") + if not isinstance(query, str) or not query.strip(): + raise RerankRequestError("query 必须是非空字符串", "query") + documents = payload.get("documents") + if (not isinstance(documents, list) or not 1 <= len(documents) <= MAX_DOCUMENTS + or any(not isinstance(document, str) for document in documents)): + raise RerankRequestError(f"documents 必须包含 1–{MAX_DOCUMENTS} 个字符串", "documents") + top_n = payload.get("top_n", len(documents)) + if type(top_n) is not int or not 1 <= top_n <= len(documents): + raise RerankRequestError("top_n 必须在 1 到 documents 数量之间", "top_n") + return_documents = payload.get("return_documents", False) + if type(return_documents) is not bool: + raise RerankRequestError("return_documents 必须是布尔值", "return_documents") + body = {"model": MODEL_ID, "query": query, "documents": documents, "top_n": top_n} + return body, documents, top_n, return_documents + + +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, documents, top_n, return_documents) -> dict: + if not isinstance(payload, dict): + raise ValueError("rerank 响应必须是 JSON 对象") + if payload.get("code") not in (None, 0): + message = payload.get("msg") or payload.get("message") or "上游 rerank 接口错误" + raise RerankUpstreamError(payload.get("code"), str(message)[:500]) + output = payload.get("output") + items = output.get("results") if isinstance(output, dict) else None + if not isinstance(items, list) or not items: + raise ValueError("rerank 响应缺少 results") + results = [] + seen = set() + for item in items: + if not isinstance(item, dict) or type(item.get("index")) is not int: + raise ValueError("rerank 响应包含无效索引") + index = item["index"] + if not 0 <= index < len(documents) or index in seen: + raise ValueError("rerank 响应包含重复或越界索引") + seen.add(index) + row = {"index": index} + if return_documents: + row["document"] = {"text": documents[index]} + results.append(row) + result = {"model": MODEL_ID, "results": results[:top_n]} + usage = _usage(payload) + if usage: + result["usage"] = usage + return result diff --git a/app/site_routing.py b/app/site_routing.py index e5ca3fe..e84dbe7 100644 --- a/app/site_routing.py +++ b/app/site_routing.py @@ -16,6 +16,8 @@ "intl-cli": INTERNATIONAL_ENDPOINT, "intl-work": "https://www.workbuddy.ai", } +# The paid bge-m3 endpoint is verified on WorkBuddy CN for both domestic products. +EMBEDDING_ENDPOINTS = {"cn-cli": "https://www.workbuddy.cn", "cn-work": "https://www.workbuddy.cn"} DOMAIN_PROFILES = { "www.codebuddy.cn": "cn-cli", "www.workbuddy.cn": "cn-work", @@ -154,5 +156,26 @@ def media_url_for_headers(headers: dict, path: str) -> str: return PROFILE_ENDPOINTS[profile_for_headers(headers)] + path +def codebase_url_for_headers(headers: dict, operation: str) -> str: + profile = profile_for_headers(headers) + if profile not in EMBEDDING_ENDPOINTS: + raise ValueError("Unsupported codebase profile") + return PROFILE_ENDPOINTS[profile] + "/codebase/openai/v1/" + operation + + +def embedding_url_for_headers(headers: dict) -> str: + return codebase_url_for_headers(headers, "embeddings") + + +def paid_embedding_url_for_headers(headers: dict) -> str: + profile = profile_for_headers(headers) + if profile not in EMBEDDING_ENDPOINTS: + raise ValueError("Unsupported embedding profile") + return EMBEDDING_ENDPOINTS[profile] + "/v2/embeddings" + + +def rerank_url_for_headers(headers: dict) -> str: + return codebase_url_for_headers(headers, "rerank") + def refresh_url_for_auth(auth: dict) -> str: return PROFILE_ENDPOINTS[profile_for_auth(auth)] + _REFRESH_PATH diff --git a/converter.py b/converter.py index 07b062c..58ec0dc 100644 --- a/converter.py +++ b/converter.py @@ -109,8 +109,9 @@ def apply_system_prompt_mode(body, profile, mode): from app.site_routing import (DOMESTIC, INTERNATIONAL, PROFILE_ENDPOINTS, site_for_auth, site_for_headers, profile_for_auth, profile_for_headers, profile_region, profile_product, profile_site, chat_url_for_headers, search_url_for_headers, fetch_url_for_headers, - refresh_url_for_auth, media_url_for_headers) -from app import dsml, mcp_server, server_tools, web_fetch, web_search, media_generation, media_store + refresh_url_for_auth, media_url_for_headers, embedding_url_for_headers, + paid_embedding_url_for_headers, rerank_url_for_headers) +from app import dsml, mcp_server, server_tools, web_fetch, web_search, media_generation, media_store, embeddings, rerank from app.inbound_limits import anthropic_path from app.client_profiles import (CLI_VERSION, CLI_USER_AGENT, credential_headers, credential_refresh_headers, catalog_cache_key, account_key, canonical_uid, enterprise_id_of, identity_email, @@ -4008,6 +4009,15 @@ def guard_model(name: str, *, region=None, resolved=False) -> None: if not resolved: name = model_policy.resolve(CONFIG, name) model_policy.check_resolved(CONFIG, name) + if embeddings.is_embedding_model(name): + public = model_policy.default_rule(name)["public_id"] + raise HTTPException(status_code=400, detail={"error": { + "message": f"{public} 是向量模型,请使用 /v1/embeddings", + "type": "invalid_request_error", "code": "embedding_model_not_chat", "param": "model"}}) + if rerank.is_rerank_model(name): + raise HTTPException(status_code=400, detail={"error": { + "message": "reranker 是重排序模型,请使用 /v1/rerank", + "type": "invalid_request_error", "code": "rerank_model_not_chat", "param": "model"}}) if not CONFIG.get("model_guard"): return if _model_profiles(name, region): @@ -4028,7 +4038,7 @@ def list_models(request: Request, authorization: Optional[str] = Header(default= data = [{"id": item["id"], "object": "model", "created": 1700000000, "owned_by": "codebuddy", "credits": item["credits"], "credits_by_profile": item["credits_by_profile"], **{key: item[key] for key in ("credits_estimated", "capabilities", "limits", "metadata_by_profile", - "media_kinds", "media_sources", "media_defaults") + "media_kinds", "media_sources", "media_defaults", "embedding", "rerank") if key in item}} for item in model_policy.public_details(sys.modules[__name__]) if principal is None or principal.models is None @@ -4036,6 +4046,343 @@ def list_models(request: Request, authorization: Optional[str] = Header(default= return {"object": "list", "data": data} + +EMBEDDING_TIMEOUT = 60.0 +EMBEDDING_MAX_INFLIGHT_PER_ACCOUNT = 16 +_EMBEDDING_ROUTE_REFUSED = re.compile(r"route config not found|model .{0,80}(?:not found|not supported|not served)", re.I) + + +def embedding_model_details() -> list: + return embeddings.model_details() + + +def rerank_model_details() -> list: + return [rerank.model_details()] + + +def _codebase_account_reason(entry, model, *, rule=None, policy=None, region=None): + pool = CONFIG.get("cred_pool") + profile = CredentialPool._entry_profile(entry) + if not model_policy.credential_enabled(CONFIG, entry, policy): + return "人工停用" + rule = rule or model_policy.active_rule(CONFIG, model, policy) + if not rule.get("enabled", True): + return "模型已停用" + if not model_policy.route_allowed(CONFIG, entry, model, rule=rule, policy=policy): + return "不在绑定范围" + if profile not in embeddings.PROFILES: + return "产品不支持 codebase 服务" + if region and not _in_region(profile, region): + return "不在区域范围" + if pool is not None: + if not pool._healthy(entry, policy): + return "认证熔断" + if not pool._model_healthy(entry, model): + return "模型额度冷却" + if not pool._model_servable(entry, model): + return "后端模型暂时不可用" + if not pool._has_credit(entry, profile): + return "额度不足或未知" + return None + + +def _embedding_cred_for(model, tried=(), *, kind="embedding"): + pool = CONFIG.get("cred_pool") + if pool is None: + cm = CONFIG.get("cred") + if cm is None or cm in {_cred_manager(item) for item in tried}: + raise HTTPException(status_code=503, detail={"error": { + "message": f"未找到可用的 {kind} 凭证", "type": "auth_error"}}) + with cm._lock: + headers = cm.get_headers() + profile = profile_for_headers(headers) + if profile not in embeddings.PROFILES: + raise HTTPException(status_code=503, detail={"error": { + "message": f"当前凭证产品不支持 {kind}", "type": "auth_error"}}) + identity = account_key(profile, headers.get("X-User-Id"), headers.get("X-Enterprise-Id")) + entry = {"profile": profile, "account_key": identity} + if not model_policy.route_allowed(CONFIG, entry, model): + raise HTTPException(status_code=503, detail={"error": { + "message": f"当前凭证不在 {kind} 模型路由范围内", "type": "service_unavailable", + "code": f"no_{kind}_account"}}) + _apply_session_headers(headers, profile, None, current_context()) + return cm, headers + + policy = model_policy.view(CONFIG) + excluded = {_cred_manager(item) for item in tried} + with pool._lock: + candidates = [entry for entry in pool._entries + if entry["cm"] not in excluded + and _codebase_account_reason(entry, model, policy=policy) is None] + if not candidates: + raise HTTPException(status_code=503, detail={"error": { + "message": f"没有符合 {kind} 路由条件的国内账号", "type": "service_unavailable", + "code": f"no_{kind}_account"}}) + busy = False + for entry in candidates: + cm = entry["cm"] + try: + with cm._lock: + headers = cm.get_headers() + generation = cm._generation + except Exception as error: + pool.cooldown(cm, reason=str(error), generation=cm._generation) + continue + profile = profile_for_headers(headers) + lease = pool._capacity.acquire(pool._capacity_key(entry), + EMBEDDING_MAX_INFLIGHT_PER_ACCOUNT, cm, generation) + if lease is None: + busy = True + continue + resources = request_resources.get() + if resources is not None: + resources.add(lease) + _apply_session_headers(headers, profile, None, current_context()) + return lease, headers + if busy: + raise HTTPException(status_code=503, headers={"Retry-After": "3"}, detail={"error": { + "message": f"{kind} 账号在途名额已满,请稍后重试", "type": "service_unavailable", + "code": "credential_concurrency_limit"}}) + raise HTTPException(status_code=503, detail={"error": { + "message": f"{kind} 账号暂时不可用,请稍后重试", "type": "service_unavailable"}}) + + +def _embedding_replayable(error) -> bool: + if isinstance(error, UpstreamHTTPError): + code = _upstream_diagnostic(error.raw)[0] + if error.status in (400, 404): + raw = error.raw.decode("utf-8", "replace") if error.raw else "" + return (_parse_not_servable(error.raw, error.status) is not None + or bool(_EMBEDDING_ROUTE_REFUSED.search(raw))) + return error.status in (401, 403, 429) and code not in NO_FAILOVER_UPSTREAM_CODES + return isinstance(error, REPLAYABLE_TRANSPORT) + + +async def _embedding_call(body, rid, *, model, public_model, encoding, expected_count, + use_codebase=True): + """Run codebase embeddings with a bounded paid fallback for bge-m3.""" + tried = [] + recovered = None + owned = True + cred = headers = None + while True: + if owned and cred is None: + cred, headers = await run_in_threadpool(_embedding_cred_for, model, tried) + profile = profile_for_headers(headers) + upstream_body = {**body, "model": embeddings.codebase_model(model)} if use_codebase else body + url = embedding_url_for_headers(headers) if use_codebase else paid_embedding_url_for_headers(headers) + observe_route(public_model=model_policy.public_name(public_model), + upstream_model=upstream_body["model"], + profile=profile, credential=account_key( + profile, headers.get("X-User-Id"), headers.get("X-Enterprise-Id"))) + t0 = time.time() + released = False + try: + async with _backend_stream(url, headers, upstream_body, timeout=EMBEDDING_TIMEOUT, + rid=rid, model_name="embeddings") as response: + if response.status_code != 200: + raw = await read_bounded_error(response) + code, message, error_body = _upstream_diagnostic(raw) + observe_attempt("upstream_error", status_code=response.status_code, + upstream_code=code, upstream_msg=message) + if response.status_code in (401, 403, 429) or (response.status_code == 404 and not use_codebase): + _note_cred_status(cred, response.status_code, model=model, raw=raw, + retry_after=parse_retry_after(response.headers.get("Retry-After"))) + raise UpstreamHTTPError( + response.status_code, raw, + retry_after=parse_retry_after(response.headers.get("Retry-After"))) + raw = await response.aread() + try: + result = embeddings.parse_reply( + json.loads(raw.decode("utf-8", "replace")), public_model, encoding, expected_count) + except embeddings.EmbeddingUpstreamError as error: + observe_failure("embedding_upstream_refused") + _media_note_restriction(cred, error.upstream_code, _upstream_diagnostic(raw)[2]) + raise HTTPException(status_code=502, detail={"error": { + "message": f"上游 embedding 接口返回错误 {error.upstream_code}:{error.message}", + "type": "upstream_error", "upstream_code": error.upstream_code}}) from None + except ValueError as error: + observe_failure("bad_embedding_response") + raise HTTPException(status_code=502, detail={"error": { + "message": str(error), "type": "upstream_error"}}) from None + except (httpx.HTTPError, UpstreamResponseError) as error: + status, raw = _upstream_failure(error, "embeddings", t0, rid) + if use_codebase and model == embeddings.MODEL_ID and ( + isinstance(error, httpx.ConnectError) + or (isinstance(error, UpstreamHTTPError) and error.status in (404, 501))): + release_credential(cred) + released = True + _log(f"[{rid}] ↻ codebase embedding 路由不可用,回退 /v2/embeddings") + return await _embedding_call(body, rid, model=model, public_model=public_model, + encoding=encoding, expected_count=expected_count, + use_codebase=False) + recovered = observe_failure_seq() + tried.append(cred) + surface = HTTPException( + status_code=status, + headers={"Retry-After": str(error.retry_after)} if getattr(error, "retry_after", None) else None, + detail=_safe_err_raw(raw, status)) + limit = _failover_limit() + if limit <= 0 or len(tried) > limit or not _embedding_replayable(error): + raise surface from None + _log(f"[{rid}] ↻ embedding 换号重放 {len(tried)}/{limit} | 上游 HTTP {status}") + release_credential(cred) + released = True + try: + cred, headers = await run_in_threadpool(_embedding_cred_for, model, tried) + except HTTPException: + raise surface from None + owned = True + continue + finally: + if owned and not released: + release_credential(cred) + if tried: + observe_recovery(recovered) + _note_cred_model_ok(cred, model) + _log(f"[{rid}] ✔ EMBEDDINGS | {profile} | {model} | {time.time() - t0:.1f}s") + return result + + +@app.post("/v1/embeddings") +async def create_embeddings(request: Request, + authorization: Optional[str] = Header(default=None), + x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key")): + _check_auth(request) + try: + payload = await request.json() + except Exception as error: + raise _bad_json(error) from None + if not isinstance(payload, dict): + raise HTTPException(status_code=400, detail={"error": { + "message": "请求体必须是 JSON 对象", "type": "invalid_request_error"}}) + requested_model = payload.get("model") + model = model_policy.resolve(CONFIG, requested_model) + model_policy.check_resolved(CONFIG, model) + if not embeddings.is_embedding_model(model): + raise HTTPException(status_code=400, detail={"error": { + "message": f"模型 {model} 不是 embedding 模型,请选择 /v1/models 中的向量模型", + "type": "invalid_request_error", "code": "model_not_embedding", "param": "model"}}) + try: + body, encoding, expected_count = embeddings.request_body(payload, model) + except embeddings.EmbeddingRequestError as error: + detail = {"message": str(error), "type": "invalid_request_error"} + if error.param: + detail["param"] = error.param + raise HTTPException(status_code=400, detail={"error": detail}) from None + rid = _request_id() + _log(f"[{rid}] ▶ EMBEDDINGS | {requested_model} | inputs={expected_count} | encoding={encoding}") + result = await _embedding_call(body, rid, model=model, public_model=requested_model, + encoding=encoding, expected_count=expected_count) + return JSONResponse(content=result) + + +async def _rerank_call(body, rid, *, model, public_model, documents, top_n, return_documents): + tried = [] + recovered = None + owned = True + cred = headers = None + while True: + if owned and cred is None: + cred, headers = await run_in_threadpool(_embedding_cred_for, model, tried, kind="rerank") + profile = profile_for_headers(headers) + url = rerank_url_for_headers(headers) + observe_route(public_model=model_policy.public_name(public_model), upstream_model=model, + profile=profile, credential=account_key( + profile, headers.get("X-User-Id"), headers.get("X-Enterprise-Id"))) + t0 = time.time() + released = False + try: + async with _backend_stream(url, headers, body, timeout=EMBEDDING_TIMEOUT, + rid=rid, model_name="rerank") as response: + if response.status_code != 200: + raw = await read_bounded_error(response) + code, message, error_body = _upstream_diagnostic(raw) + observe_attempt("upstream_error", status_code=response.status_code, + upstream_code=code, upstream_msg=message) + if response.status_code in (401, 403, 404, 429): + _note_cred_status(cred, response.status_code, model=model, raw=raw, + retry_after=parse_retry_after(response.headers.get("Retry-After"))) + raise UpstreamHTTPError( + response.status_code, raw, + retry_after=parse_retry_after(response.headers.get("Retry-After"))) + raw = await response.aread() + try: + result = rerank.parse_reply( + json.loads(raw.decode("utf-8", "replace")), documents, top_n, return_documents) + result["model"] = public_model + except rerank.RerankUpstreamError as error: + observe_failure("rerank_upstream_refused") + _media_note_restriction(cred, error.upstream_code, _upstream_diagnostic(raw)[2]) + raise HTTPException(status_code=502, detail={"error": { + "message": f"上游 rerank 接口返回错误 {error.upstream_code}:{error.message}", + "type": "upstream_error", "upstream_code": error.upstream_code}}) from None + except ValueError as error: + observe_failure("bad_rerank_response") + raise HTTPException(status_code=502, detail={"error": { + "message": str(error), "type": "upstream_error"}}) from None + except (httpx.HTTPError, UpstreamResponseError) as error: + status, raw = _upstream_failure(error, "rerank", t0, rid) + recovered = observe_failure_seq() + tried.append(cred) + surface = HTTPException( + status_code=status, + headers={"Retry-After": str(error.retry_after)} if getattr(error, "retry_after", None) else None, + detail=_safe_err_raw(raw, status)) + limit = _failover_limit() + if limit <= 0 or len(tried) > limit or not _embedding_replayable(error): + raise surface from None + _log(f"[{rid}] ↻ rerank 换号重放 {len(tried)}/{limit} | 上游 HTTP {status}") + release_credential(cred) + released = True + try: + cred, headers = await run_in_threadpool(_embedding_cred_for, model, tried, kind="rerank") + except HTTPException: + raise surface from None + owned = True + continue + finally: + if owned and not released: + release_credential(cred) + if tried: + observe_recovery(recovered) + _note_cred_model_ok(cred, model) + _log(f"[{rid}] ✔ RERANK | {profile} | {model} | {time.time() - t0:.1f}s") + return result + + +@app.post("/v1/rerank") +async def create_rerank(request: Request, + authorization: Optional[str] = Header(default=None), + x_api_key: Optional[str] = Header(default=None, alias="X-Api-Key")): + _check_auth(request) + try: + payload = await request.json() + except Exception as error: + raise _bad_json(error) from None + if not isinstance(payload, dict): + raise HTTPException(status_code=400, detail={"error": { + "message": "请求体必须是 JSON 对象", "type": "invalid_request_error"}}) + requested_model = payload.get("model") or rerank.MODEL_ID + model = model_policy.resolve(CONFIG, requested_model) + model_policy.check_resolved(CONFIG, model) + if not rerank.is_rerank_model(model): + raise HTTPException(status_code=400, detail={"error": { + "message": f"模型 {model} 不是重排序模型,请使用 reranker", + "type": "invalid_request_error", "code": "model_not_rerank", "param": "model"}}) + try: + body, documents, top_n, return_documents = rerank.request_body(payload) + except rerank.RerankRequestError as error: + detail = {"message": str(error), "type": "invalid_request_error"} + if error.param: + detail["param"] = error.param + raise HTTPException(status_code=400, detail={"error": detail}) from None + rid = _request_id() + _log(f"[{rid}] ▶ RERANK | {requested_model} | documents={len(documents)} | top_n={top_n}") + result = await _rerank_call(body, rid, model=model, public_model=requested_model, + documents=documents, top_n=top_n, return_documents=return_documents) + return JSONResponse(content=result) + @app.post("/v1/chat/completions") async def chat_completions(request: Request, authorization: Optional[str] = Header(default=None), diff --git a/docs/clients.md b/docs/clients.md index 3de14a4..4ba45f0 100644 --- a/docs/clients.md +++ b/docs/clients.md @@ -83,6 +83,30 @@ The generation endpoints are `POST /v1/chat/completions`, `POST /v1/responses` a Streaming policy is a server setting, not a client request field. The default `compatible` mode aggregates Responses and tool-bearing Chat/Messages as before; `realtime` sends all three protocols incrementally. Non-streaming requests remain validated JSON in either mode. See [Streaming modes](advanced.md#streaming-modes) for configuration, partial-output errors and the loss of automatic tool-argument regeneration. + +## Embeddings + +`POST /v1/embeddings` exposes `baai/bge-m3` and `Qwen/Qwen3-Embedding-0.6B` through `/codebase/openai/v1/embeddings`. It accepts text, text arrays and token-ID arrays, and returns 1024-dimensional `float` or `base64` vectors. `baai/bge-m3` maps to upstream `model2`; only a missing codebase route or connection failure falls back to `/v2/embeddings` as `bge-m3`. Qwen has no paid fallback. + +Codebase usage has been observed as zero, but that is not a billing guarantee. Routing uses enabled domestic accounts with positive balances, shared alias/binding rules and 16 in-flight requests per account; timeouts and 5xx responses do not switch endpoints. + +```bash +curl http://127.0.0.1:8787/v1/embeddings \ + -H "Authorization: Bearer YOUR_GATEWAY_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"baai/bge-m3","input":["first document","second document"],"encoding_format":"float"}' +``` + +## Reranking + +`POST /v1/rerank` exposes the fixed codebase reranker as `reranker`. Send `query`, 1–100 string `documents`, optional `top_n` and `return_documents`; results contain ordered indices and optional source documents, never relevance scores. The service truncates each document at roughly 600 characters. Its weights are not exposed: `bge-reranker-v2-m3` and `Qwen3-Reranker` are unverified candidates, not confirmed identities. + +```bash +curl http://127.0.0.1:8787/v1/rerank \ + -H "Authorization: Bearer YOUR_GATEWAY_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"reranker","query":"capital of France","documents":["Berlin","Paris"],"top_n":1}' +``` ## Web search and fetch Search and fetch endpoints backed by the upstream agenttool `/v1/search` and `/v1/webfetch`; no model is called and no tokens are billed. diff --git a/docs/clients.zh-CN.md b/docs/clients.zh-CN.md index f6f3d90..b6ba7c8 100644 --- a/docs/clients.zh-CN.md +++ b/docs/clients.zh-CN.md @@ -84,6 +84,30 @@ Cherry Studio、ZCode、LobeChat、NextChat、Open WebUI 或自研 SDK 客户端 流式策略是服务端设置,不是客户端请求字段。默认 `compatible` 保持 Responses 及带工具 Chat/Messages 的聚合行为;`realtime` 让三个协议都增量发送。两种模式的非流式请求均为经过校验的 JSON。配置、部分输出错误及不再自动重生成工具参数等限制见[流式模式](advanced.zh-CN.md#流式模式)。 +## 向量嵌入 + +`POST /v1/embeddings` 通过 `/codebase/openai/v1/embeddings` 提供 `baai/bge-m3` 和 `Qwen/Qwen3-Embedding-0.6B`。支持文本、文本数组和 token ID 数组,返回 1024 维 `float` 或 `base64` 向量。`baai/bge-m3` 对应上游 `model2`;仅在 codebase 路由不存在或连接失败时回退 `/v2/embeddings` 的 `bge-m3`,Qwen 无付费回退。 + +codebase 的 usage 实测可能恒为 0,但这不构成免费保证。仅选择已启用且有余额的国内账号,共用别名、启停和绑定规则;每账号允许 16 路并发,超时与 5xx 不切换端点。 + +```bash +curl http://127.0.0.1:8787/v1/embeddings \ + -H "Authorization: Bearer YOUR_GATEWAY_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"baai/bge-m3","input":["第一段文本","第二段文本"],"encoding_format":"float"}' +``` + +## 重排序 + +`POST /v1/rerank` 将固定的 codebase 重排序服务暴露为 `reranker`。请求包含 `query`、1–100 个字符串 `documents`,可选 `top_n` 和 `return_documents`;结果只含排序后的索引及可选原文,不含相关度分数。上游会把每篇文档截断到约 600 字符。真实权重未公开,`bge-reranker-v2-m3` 和 `Qwen3-Reranker` 仅为未证实候选。 + +```bash +curl http://127.0.0.1:8787/v1/rerank \ + -H "Authorization: Bearer YOUR_GATEWAY_API_KEY" \ + -H "Content-Type: application/json" \ + -d '{"model":"reranker","query":"法国首都","documents":["柏林","巴黎"],"top_n":1}' +``` + ## 联网搜索与抓取 搜索与抓取端点,上游为 agenttool 的 `/v1/search` 与 `/v1/webfetch`:不调用模型,不产生 token 计费。 diff --git a/preview/upstream.py b/preview/upstream.py index 11752a8..0c767ab 100644 --- a/preview/upstream.py +++ b/preview/upstream.py @@ -20,7 +20,7 @@ import httpx -from app import media_generation +from app import embeddings as embedding_contract, media_generation, rerank as rerank_contract from app.site_routing import DOMAIN_PROFILES, PROFILE_DOMAINS from preview.seed import (ACCOUNTS, CN_MODELS, COS_HOST, DAY, INTL_MODELS, MP4_BYTES, PNG_BYTES, SEED_VIDEO_TASK_ID, Account) @@ -163,6 +163,10 @@ def __init__(self, accounts=ACCOUNTS): "/v2/billing/meter/daily-checkin": self.checkin, "/billing/ide/trial": self.trial, "/v2/chat/completions": self.chat, + "/v2/embeddings": self.embeddings, + "/codebase/openai/v1/models": self.codebase_models, + "/codebase/openai/v1/embeddings": self.codebase_embeddings, + "/codebase/openai/v1/rerank": self.codebase_rerank, "/console/chat/completions": self.chat, "/agenttool/v1/search": self.search, "/agenttool/v1/webfetch": self.webfetch, @@ -447,6 +451,61 @@ def webfetch(self, request, body, state): "content_type": "text/markdown", "content_length": 64, "truncated": False, "status_code": 200}) + def codebase_models(self, request, body, state): + return Reply(payload={"object": "list", "data": [ + {"id": name, "object": "model"} for name in ( + "model1", "model2", embedding_contract.QWEN_MODEL_ID)]}) + + def codebase_embeddings(self, request, body, state): + if body.get("model") not in ("model1", "model2", embedding_contract.QWEN_MODEL_ID): + return fail(400, "route config not found") + raw = body.get("input") + items = [raw] if isinstance(raw, str) or (isinstance(raw, list) and all(type(v) is int for v in raw)) else raw + if not isinstance(items, list) or not items: + return fail(400, "invalid embedding input") + data = [] + for index, item in enumerate(items): + vector = [0.0] * embedding_contract.DIMENSIONS + vector[0], vector[1] = float(index), float(len(item) % 997) / 997 + data.append({"object": "embedding", "index": index, "embedding": vector}) + return Reply(payload={"object": "list", "data": data, "model": body["model"], + "usage": {"prompt_tokens": 0, "total_tokens": 0, "credit": 0}}) + + def codebase_rerank(self, request, body, state): + query = body.get("query") + documents = body.get("documents") + if not isinstance(query, str) or not query.strip(): + return fail(400, "invalid query", code=10001) + if (not isinstance(documents, list) or not 1 <= len(documents) <= rerank_contract.MAX_DOCUMENTS + or any(not isinstance(document, str) for document in documents)): + return fail(400, "documents exceeds maximum allowed (100)", code=10001) + top_n = body.get("top_n", len(documents)) + if type(top_n) is not int or not 1 <= top_n <= len(documents): + return fail(400, "invalid top_n", code=10001) + terms = set(query.lower().split()) + ranked = sorted(range(len(documents)), + key=lambda index: (-len(terms.intersection(documents[index].lower().split())), index)) + return Reply(payload={"output": {"results": [{"index": index} for index in ranked[:top_n]]}, + "request_id": "0" * 32, "usage": {"total_tokens": 0}}) + + def embeddings(self, request, body, state): + if body.get("model") != embedding_contract.MODEL_ID: + return fail(400, "route config not found") + raw = body.get("input") + items = [raw] if isinstance(raw, str) or (isinstance(raw, list) and all(type(v) is int for v in raw)) else raw + if not isinstance(items, list) or not items: + return fail(400, "invalid embedding input") + tokens = sum(max(1, len(item) // 4) if isinstance(item, str) else len(item) for item in items) + credit = round(tokens / 1000 * 0.01, 6) + self._charge(state, credit) + data = [] + for index, item in enumerate(items): + vector = [0.0] * embedding_contract.DIMENSIONS + vector[0], vector[1] = float(index), float(len(item) % 997) / 997 + data.append({"object": "embedding", "index": index, "embedding": vector}) + return Reply(payload={"object": "list", "data": data, "model": embedding_contract.MODEL_ID, + "usage": {"prompt_tokens": tokens, "total_tokens": tokens, "credit": credit}}) + # -- media generation ---------------------------------------------------------------- def media_file(self, request): diff --git a/tests/test_embeddings.py b/tests/test_embeddings.py new file mode 100644 index 0000000..60a6d9c --- /dev/null +++ b/tests/test_embeddings.py @@ -0,0 +1,258 @@ +#!/usr/bin/env python3 +"""Codebase embedding contracts, public aliases, routing and fallback.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import base64 +import json +import struct +import unittest +from unittest.mock import patch + +import httpx + +import converter +from app import embeddings, model_policy +from app.control_store import ControlStore +from app.gateway_management import Management +from tests import test_region_routing as fixtures + + +def vector(value=0.0): + return [float(value)] + [0.0] * (embeddings.DIMENSIONS - 1) + + +def embedding_payload(count=2, *, codebase=False): + usage = {"prompt_tokens": 0, "total_tokens": 0, "credit": 0} if codebase else { + "prompt_tokens": 7, "total_tokens": 7, "credit": 0.001} + return {"object": "list", "model": "model2" if codebase else "bge-m3", + "data": [{"object": "embedding", "index": index, "embedding": vector(index)} + for index in range(count)], "usage": usage} + + +class EmbeddingContractTests(unittest.TestCase): + def test_builtin_public_alias_and_codebase_names(self): + for source in (embeddings.MODEL_ID, embeddings.PUBLIC_MODEL_ID): + rule = model_policy.default_rule(source) + self.assertEqual(rule["public_id"], "baai/bge-m3") + self.assertEqual(rule["upstream_id"], "bge-m3") + self.assertEqual(model_policy.source_for({}, "baai/bge-m3"), "bge-m3") + self.assertEqual(embeddings.codebase_model("bge-m3"), "model2") + self.assertEqual(embeddings.codebase_model(embeddings.QWEN_MODEL_ID), embeddings.QWEN_MODEL_ID) + self.assertEqual(embeddings.is_embedding_model(embeddings.QWEN_MODEL_ID), True) + + def test_openai_input_shapes(self): + cases = (("hello", 1), (["hello", "world"], 2), ([1, 2, 3], 1), ([[1], [2, 3]], 2)) + for value, count in cases: + with self.subTest(value=value): + body, encoding, actual = embeddings.request_body( + {"input": value, "dimensions": 1024}, "bge-m3") + self.assertEqual(body, {"model": "bge-m3", "input": value, "encoding_format": "float"}) + self.assertEqual(encoding, "float") + self.assertEqual(actual, count) + for payload in ({}, {"input": ""}, {"input": []}, {"input": ["ok", 1]}, + {"input": [True]}, {"input": [[]]}, {"input": "x", "encoding_format": "hex"}, + {"input": "x", "dimensions": 512}, {"input": "x", "user": 1}): + with self.subTest(payload=payload), self.assertRaises(embeddings.EmbeddingRequestError): + embeddings.request_body(payload, "bge-m3") + + def test_reply_order_base64_and_usage(self): + payload = embedding_payload() + payload["data"].reverse() + result = embeddings.parse_reply(payload, "baai/bge-m3", "float", 2) + self.assertEqual([item["index"] for item in result["data"]], [0, 1]) + self.assertEqual(result["model"], "baai/bge-m3") + self.assertEqual(result["usage"], {"prompt_tokens": 7, "total_tokens": 7, "credit": 0.001}) + zero = embeddings.parse_reply(embedding_payload(codebase=True), "baai/bge-m3", "float", 2) + self.assertEqual(zero["usage"], {"prompt_tokens": 0, "total_tokens": 0, "credit": 0}) + + encoded = embeddings.parse_reply(payload, "baai/bge-m3", "base64", 2) + raw = base64.b64decode(encoded["data"][1]["embedding"], validate=True) + self.assertEqual(len(raw), embeddings.DIMENSIONS * 4) + self.assertEqual(struct.unpack("<2f", raw[:8]), (1.0, 0.0)) + + upstream_base64 = embedding_payload(1) + upstream_base64["data"][0]["embedding"] = base64.b64encode( + struct.pack(f"<{embeddings.DIMENSIONS}f", *vector(0.25))).decode() + decoded = embeddings.parse_reply(upstream_base64, "baai/bge-m3", "float", 1) + self.assertEqual(decoded["data"][0]["embedding"][0], 0.25) + + def test_bad_replies_are_rejected(self): + for mutate in ( + lambda body: body.update(data=[]), + lambda body: body["data"][0].update(index=9), + lambda body: body["data"][0].update(embedding=[0.0]), + lambda body: body["data"][0].update(embedding="not-base64"), + lambda body: body["data"][0].update(embedding=[float("nan")] * 1024)): + payload = embedding_payload() + mutate(payload) + with self.subTest(mutate=mutate), self.assertRaises(ValueError): + embeddings.parse_reply(payload, "baai/bge-m3", "float", 2) + with self.assertRaises(embeddings.EmbeddingUpstreamError) as caught: + embeddings.parse_reply({"code": 10040, "msg": "unavailable"}, "baai/bge-m3", "float", 1) + self.assertEqual(caught.exception.upstream_code, 10040) + + +class EmbeddingRoutingTests(fixtures.RegionRoutingTests): + def setUp(self): + super().setUp() + self.allowed_profiles = {"cn-cli", "cn-work"} + self.embedding_requests = [] + self.embedding_replies = [] + self.enterContext(patch.object(converter, "_log", side_effect=lambda *args, **kwargs: None)) + + @staticmethod + def input_items(raw): + if isinstance(raw, str) or (isinstance(raw, list) and all(type(value) is int for value in raw)): + return [raw] + return raw + + def handle_upstream(self, request): + if request.url.path not in ("/codebase/openai/v1/embeddings", "/v2/embeddings"): + return super().handle_upstream(request) + uid = request.headers.get("x-user-id") + self.assertIn(uid, self.allowed_profiles) + profile = self.account_profiles[uid] + self.assertIn(profile, embeddings.PROFILES) + codebase = request.url.path.startswith("/codebase/") + expected_host = httpx.URL(converter.PROFILE_ENDPOINTS[profile]).host if codebase else "www.workbuddy.cn" + self.assertEqual(request.url.host, expected_host) + self.assertEqual(request.headers["authorization"], "Bearer synthetic-access-" + uid) + body = json.loads(request.content) + self.embedding_requests.append((request, body)) + if self.embedding_replies: + status, payload = self.embedding_replies.pop(0) + else: + payload = embedding_payload(len(self.input_items(body["input"])), codebase=codebase) + status = 200 + return httpx.Response(status, json=payload) + + def post(self, value, **extra): + return self.client.post("/v1/embeddings", json={"model": "baai/bge-m3", "input": value, **extra}) + + def install_policy(self): + store = ControlStore(self.root / f"embedding-policy-{self.fixture_sequence}.sqlite3") + self.addCleanup(store.close) + self.enterContext(patch.dict(converter.CONFIG, {"control_store": store})) + converter.invalidate_model_table() + return store + + def test_public_catalog_and_codebase_batch_roundtrip(self): + models = self.client.get("/v1/models") + self.assertEqual(models.status_code, 200) + ids = {item["id"] for item in models.json()["data"]} + self.assertIn("baai/bge-m3", ids) + self.assertIn(embeddings.QWEN_MODEL_ID, ids) + self.assertNotIn("bge-m3", ids) + item = next(item for item in models.json()["data"] if item["id"] == "baai/bge-m3") + self.assertTrue(item["embedding"]) + + response = self.post(["first", "second"]) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["model"], "baai/bge-m3") + self.assertEqual(len(response.json()["data"]), 2) + self.assertEqual(response.json()["usage"], {"prompt_tokens": 0, "total_tokens": 0, "credit": 0}) + request, body = self.embedding_requests[-1] + self.assertEqual(body, {"model": "model2", "input": ["first", "second"], "encoding_format": "float"}) + self.assertEqual(request.url.path, "/codebase/openai/v1/embeddings") + + def test_qwen_uses_its_direct_codebase_name(self): + response = self.post("hello", model=embeddings.QWEN_MODEL_ID) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["model"], embeddings.QWEN_MODEL_ID) + request, body = self.embedding_requests[-1] + self.assertEqual(request.url.path, "/codebase/openai/v1/embeddings") + self.assertEqual(body["model"], embeddings.QWEN_MODEL_ID) + + def test_token_input_and_base64_response(self): + response = self.post([[1, 2], [3]], encoding_format="base64", dimensions=1024) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(len(response.json()["data"]), 2) + raw = base64.b64decode(response.json()["data"][0]["embedding"], validate=True) + self.assertEqual(len(raw), embeddings.DIMENSIONS * 4) + self.assertEqual(self.embedding_requests[-1][1]["input"], [[1, 2], [3]]) + self.assertEqual(self.embedding_requests[-1][1]["encoding_format"], "float") + + def test_old_name_wrong_endpoint_and_invalid_input_are_rejected(self): + old = self.post("hello", model="bge-m3") + self.assertEqual(old.status_code, 404) + self.assertEqual(old.json()["error"]["code"], "model_renamed") + chat = self.client.post("/v1/chat/completions", json={ + "model": "baai/bge-m3", "messages": [{"role": "user", "content": "hello"}]}) + self.assertEqual(chat.status_code, 400) + self.assertEqual(chat.json()["error"]["code"], "embedding_model_not_chat") + wrong = self.post("hello", model="shared-model") + self.assertEqual(wrong.status_code, 400) + self.assertEqual(wrong.json()["error"]["code"], "model_not_embedding") + invalid = self.post([]) + self.assertEqual(invalid.status_code, 400) + self.assertEqual(invalid.json()["error"]["param"], "input") + self.assertEqual(self.embedding_requests, []) + + def test_binding_disable_and_rate_probe_guard(self): + store = self.install_policy() + rule = {**model_policy.default_rule("bge-m3"), + "credential_ids": [self.entries["cn-work"]["account_key"]]} + store.update_model("bge-m3", rule, 0, known_models=["bge-m3"]) + response = self.post("hello") + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(self.embedding_requests[-1][0].headers["x-user-id"], "cn-work") + store.update_model("bge-m3", {**rule, "enabled": False}, 1, known_models=["bge-m3"]) + disabled = self.post("hello") + self.assertEqual(disabled.status_code, 404) + self.assertEqual(len(self.embedding_requests), 1) + with self.assertRaisesRegex(ValueError, "不支持聊天倍率探测"): + Management(converter).admin_rate_probe_model("bge-m3") + + def test_missing_codebase_route_falls_back_to_paid_bge(self): + self.embedding_replies = [(404, {"error": {"message": "route not found"}}), + (200, embedding_payload(1))] + response = self.post("hello") + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual([request.url.path for request, _body in self.embedding_requests], + ["/codebase/openai/v1/embeddings", "/v2/embeddings"]) + self.assertEqual(self.embedding_requests[0][1]["model"], "model2") + self.assertEqual(self.embedding_requests[1][1]["model"], "bge-m3") + self.assertEqual(response.json()["usage"]["total_tokens"], 7) + + def test_qwen_has_no_paid_fallback(self): + self.embedding_replies = [(404, {"error": {"message": "route not found"}}), + (200, embedding_payload(1))] + response = self.post("hello", model=embeddings.QWEN_MODEL_ID) + self.assertEqual(response.status_code, 404) + self.assertEqual(len(self.embedding_requests), 1) + + def test_authentication_failure_replays_on_another_account(self): + self.embedding_replies = [(401, {"error": {"message": "expired"}}), + (200, embedding_payload(1, codebase=True))] + with patch.dict(converter.CONFIG, {"failover_max": 1}): + response = self.post("hello") + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(len(self.embedding_requests), 2) + self.assertNotEqual(self.embedding_requests[0][0].headers["x-user-id"], + self.embedding_requests[1][0].headers["x-user-id"]) + + def test_server_and_business_failures_are_not_replayed(self): + self.embedding_replies = [(500, {"error": {"message": "unavailable"}}), (200, embedding_payload(1))] + response = self.post("hello") + self.assertEqual(response.status_code, 500) + self.assertEqual(len(self.embedding_requests), 1) + + self.embedding_replies = [(200, {"code": 10040, "msg": "unavailable"}), (200, embedding_payload(1))] + response = self.post("hello") + self.assertEqual(response.status_code, 502) + self.assertEqual(response.json()["error"]["upstream_code"], 10040) + self.assertEqual(len(self.embedding_requests), 2) + + def test_embedding_has_independent_sixteen_way_capacity(self): + with patch.dict(converter.CONFIG, {"max_inflight_per_account": 1}), \ + patch.object(self.pool._capacity, "acquire", wraps=self.pool._capacity.acquire) as acquire: + cred, _headers = converter._embedding_cred_for("bge-m3") + converter.release_credential(cred) + self.assertEqual(acquire.call_args.args[1], 16) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_observability.py b/tests/test_observability.py index e4d9208..06072b7 100644 --- a/tests/test_observability.py +++ b/tests/test_observability.py @@ -77,6 +77,41 @@ async def app(scope, receive, send): self.assertNotIn("private", json.dumps(record)) self.assertNotIn("sk-synthetic", json.dumps(record)) + async def test_embeddings_protocol_and_usage(self): + body = json.dumps({"object": "list", "model": "baai/bge-m3", "data": [], + "usage": {"prompt_tokens": 9, "total_tokens": 9, "credit": 0.01}}).encode() + messages = [{"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"application/json")]}, + {"type": "http.response.body", "body": body}] + + async def app(scope, receive, send): + for message in messages: + await send(message) + + await self.invoke(app, path="/v1/embeddings") + record = self.only_record() + self.assertEqual(record["protocol"], "embeddings") + self.assertEqual(record["public_model"], "baai/bge-m3") + self.assertEqual(record["input_tokens"], 9) + self.assertEqual(record["total_tokens"], 9) + self.assertEqual(record["credit"], 0.01) + + async def test_rerank_protocol(self): + body = json.dumps({"model": "reranker", "results": [], "usage": {"total_tokens": 0}}).encode() + messages = [{"type": "http.response.start", "status": 200, + "headers": [(b"content-type", b"application/json")]}, + {"type": "http.response.body", "body": body}] + + async def app(scope, receive, send): + for message in messages: + await send(message) + + await self.invoke(app, path="/v1/rerank") + record = self.only_record() + self.assertEqual(record["protocol"], "rerank") + self.assertEqual(record["public_model"], "reranker") + self.assertEqual(record["total_tokens"], 0) + async def test_non_responses_audit_omits_projection_fields(self): body = json.dumps({"model": "public-model", "choices": [{"message": {"content": "ok"}}]}).encode() messages = [{"type": "http.response.start", "status": 200, diff --git a/tests/test_preview.py b/tests/test_preview.py index 2933f82..48c88ad 100644 --- a/tests/test_preview.py +++ b/tests/test_preview.py @@ -159,6 +159,34 @@ def test_gzip_chat_body_is_decoded_before_dispatch(self): self.assertEqual(response.status_code, 200, response.text) self.assertIn("data: [DONE]", response.text) + + def test_embeddings_return_vectors_and_billed_usage(self): + token, _, _ = login("aurora") + response = httpx.post("https://www.workbuddy.cn/v2/embeddings", + headers={"Authorization": f"Bearer {token}"}, + json={"model": "bge-m3", "input": ["hello", "world"], "encoding_format": "float"}) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(len(response.json()["data"]), 2) + self.assertEqual(len(response.json()["data"][0]["embedding"]), 1024) + self.assertGreater(response.json()["usage"]["credit"], 0) + + def test_codebase_embeddings_and_rerank_have_no_simulated_charge(self): + token, uid, domain = login("aurora") + before = credits.fetch_credits(token, uid, domain)["credits"] + embedded = httpx.post("https://copilot.tencent.com/codebase/openai/v1/embeddings", + headers={"Authorization": f"Bearer {token}"}, + json={"model": "model2", "input": ["hello"], "encoding_format": "float"}) + self.assertEqual(embedded.status_code, 200, embedded.text) + self.assertEqual(embedded.json()["usage"]["total_tokens"], 0) + self.assertEqual(credits.fetch_credits(token, uid, domain)["credits"], before) + + ranked = httpx.post("https://copilot.tencent.com/codebase/openai/v1/rerank", + headers={"Authorization": f"Bearer {token}"}, + json={"query": "capital France", "documents": ["Berlin", "Paris capital France"], + "top_n": 1}) + self.assertEqual(ranked.status_code, 200, ranked.text) + self.assertEqual(ranked.json()["output"]["results"], [{"index": 1}]) + self.assertNotIn("relevance_score", ranked.text) def test_console_chat_uses_the_same_simulated_upstream(self): token, _, _ = login("aurora") response = httpx.post("https://copilot.tencent.com/console/chat/completions", diff --git a/tests/test_region_routing.py b/tests/test_region_routing.py index 7c0ae7d..e26a555 100644 --- a/tests/test_region_routing.py +++ b/tests/test_region_routing.py @@ -66,7 +66,8 @@ def success_sse(): def chat_model_ids(response): """Exclude media routes when asserting the legacy chat catalog.""" - return {item["id"] for item in response.json()["data"] if not item.get("media_kinds")} + return {item["id"] for item in response.json()["data"] + if not item.get("media_kinds") and not item.get("embedding") and not item.get("rerank")} class RegionRoutingTests(unittest.TestCase): diff --git a/tests/test_rerank.py b/tests/test_rerank.py new file mode 100644 index 0000000..83a05b7 --- /dev/null +++ b/tests/test_rerank.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +"""Codebase reranker contracts, routing and no-score response behavior.""" +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +import json +import unittest +from unittest.mock import patch + +import httpx + +import converter +from app import model_policy, rerank +from app.control_store import ControlStore +from app.gateway_management import Management +from tests import test_region_routing as fixtures + + +def rerank_payload(count=2): + indices = list(range(count - 1, -1, -1)) + return {"output": {"results": [{"index": index} for index in indices]}, + "request_id": "0" * 32, "usage": {"total_tokens": 0}} + + +class RerankContractTests(unittest.TestCase): + def test_request_and_no_score_reply(self): + body, documents, top_n, return_documents = rerank.request_body({ + "query": "capital of France", "documents": ["Berlin", "Paris", "London"], + "top_n": 2, "return_documents": True}) + self.assertEqual(body["model"], "reranker") + self.assertEqual(top_n, 2) + result = rerank.parse_reply(rerank_payload(3), documents, top_n, return_documents) + self.assertEqual(result["model"], "reranker") + self.assertEqual([item["index"] for item in result["results"]], [2, 1]) + self.assertEqual(result["results"][0]["document"], {"text": "London"}) + self.assertNotIn("relevance_score", json.dumps(result)) + self.assertEqual(result["usage"], {"total_tokens": 0}) + + def test_request_validation(self): + for payload in ({}, {"query": " ", "documents": ["x"]}, {"query": "x", "documents": []}, + {"query": "x", "documents": ["x"] * 101}, + {"query": "x", "documents": [1]}, {"query": "x", "documents": ["x"], "top_n": 0}, + {"query": "x", "documents": ["x"], "top_n": 2}, + {"query": "x", "documents": ["x"], "return_documents": 1}): + with self.subTest(payload=payload), self.assertRaises(rerank.RerankRequestError): + rerank.request_body(payload) + + def test_bad_replies_are_rejected(self): + for payload in ({}, {"output": {"results": []}}, + {"output": {"results": [{"index": 4}]}}, + {"output": {"results": [{"index": 0}, {"index": 0}]}}): + with self.subTest(payload=payload), self.assertRaises(ValueError): + rerank.parse_reply(payload, ["a", "b"], 2, False) + with self.assertRaises(rerank.RerankUpstreamError) as caught: + rerank.parse_reply({"code": 10001, "msg": "empty documents"}, ["a"], 1, False) + self.assertEqual(caught.exception.upstream_code, 10001) + + +class RerankRoutingTests(fixtures.RegionRoutingTests): + def setUp(self): + super().setUp() + self.allowed_profiles = {"cn-cli", "cn-work"} + self.rerank_requests = [] + self.rerank_replies = [] + self.enterContext(patch.object(converter, "_log", side_effect=lambda *args, **kwargs: None)) + + def handle_upstream(self, request): + if request.url.path != "/codebase/openai/v1/rerank": + return super().handle_upstream(request) + uid = request.headers.get("x-user-id") + self.assertIn(uid, self.allowed_profiles) + profile = self.account_profiles[uid] + self.assertIn(profile, ("cn-cli", "cn-work")) + self.assertEqual(request.url.host, httpx.URL(converter.PROFILE_ENDPOINTS[profile]).host) + self.assertEqual(request.headers["authorization"], "Bearer synthetic-access-" + uid) + body = json.loads(request.content) + self.rerank_requests.append((request, body)) + if self.rerank_replies: + status, payload = self.rerank_replies.pop(0) + else: + status, payload = 200, rerank_payload(len(body["documents"])) + return httpx.Response(status, json=payload) + + def post(self, documents, **extra): + return self.client.post("/v1/rerank", json={ + "model": "reranker", "query": "capital of France", "documents": documents, **extra}) + + def install_policy(self): + store = ControlStore(self.root / f"rerank-policy-{self.fixture_sequence}.sqlite3") + self.addCleanup(store.close) + self.enterContext(patch.dict(converter.CONFIG, {"control_store": store})) + converter.invalidate_model_table() + return store + + def test_public_catalog_and_roundtrip(self): + models = self.client.get("/v1/models") + item = next(item for item in models.json()["data"] if item["id"] == "reranker") + self.assertTrue(item["rerank"]) + response = self.post(["Berlin", "Paris", "London"], top_n=2, return_documents=True) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["model"], "reranker") + self.assertEqual([item["index"] for item in response.json()["results"]], [2, 1]) + self.assertEqual(response.json()["results"][0]["document"], {"text": "London"}) + self.assertNotIn("relevance_score", response.text) + request, body = self.rerank_requests[-1] + self.assertEqual(request.url.path, "/codebase/openai/v1/rerank") + self.assertEqual(body["model"], "reranker") + self.assertEqual(body["top_n"], 2) + + def test_default_model_and_validation(self): + response = self.client.post("/v1/rerank", json={"query": "x", "documents": ["doc"]}) + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["model"], "reranker") + for documents, extra in (([], {}), (["x"] * 101, {}), (["x"], {"top_n": 0})): + result = self.post(documents, **extra) + self.assertEqual(result.status_code, 400) + wrong = self.post(["x"], model="baai/bge-m3") + self.assertEqual(wrong.status_code, 400) + self.assertEqual(wrong.json()["error"]["code"], "model_not_rerank") + chat = self.client.post("/v1/chat/completions", json={ + "model": "reranker", "messages": [{"role": "user", "content": "hello"}]}) + self.assertEqual(chat.status_code, 400) + self.assertEqual(chat.json()["error"]["code"], "rerank_model_not_chat") + self.assertEqual(len(self.rerank_requests), 1) + + def test_alias_binding_disable_and_rate_probe_guard(self): + store = self.install_policy() + rule = {**model_policy.default_rule("reranker"), "public_id": "rerank-alias", + "credential_ids": [self.entries["cn-work"]["account_key"]]} + store.update_model("reranker", rule, 0, known_models=["reranker"]) + response = self.post(["a", "b"], model="rerank-alias") + self.assertEqual(response.status_code, 200, response.text) + self.assertEqual(response.json()["model"], "rerank-alias") + self.assertEqual(self.rerank_requests[-1][0].headers["x-user-id"], "cn-work") + self.assertEqual(self.rerank_requests[-1][1]["model"], "reranker") + store.update_model("reranker", {**rule, "enabled": False}, 1, known_models=["reranker"]) + disabled = self.post(["a"], model="rerank-alias") + self.assertEqual(disabled.status_code, 404) + self.assertEqual(len(self.rerank_requests), 1) + with self.assertRaisesRegex(ValueError, "不支持聊天倍率探测"): + Management(converter).admin_rate_probe_model("reranker") + + def test_server_and_business_failures_are_not_replayed(self): + self.rerank_replies = [(500, {"error": {"message": "unavailable"}}), (200, rerank_payload(1))] + response = self.post(["doc"]) + self.assertEqual(response.status_code, 500) + self.assertEqual(len(self.rerank_requests), 1) + + self.rerank_replies = [(200, {"code": 10001, "msg": "bad request"}), (200, rerank_payload(1))] + response = self.post(["doc"]) + self.assertEqual(response.status_code, 502) + self.assertEqual(response.json()["error"]["upstream_code"], 10001) + self.assertEqual(len(self.rerank_requests), 2) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/web/src/api.ts b/web/src/api.ts index 98a1060..3b0e401 100644 --- a/web/src/api.ts +++ b/web/src/api.ts @@ -267,6 +267,8 @@ export type ModelRule = { media_kinds?: string[]; media_sources?: string[]; media_defaults?: string[]; + embedding?: boolean; + rerank?: boolean; }; function rateOverride(value: unknown): Record | undefined { if (!value || typeof value !== "object" || Array.isArray(value)) return undefined; diff --git a/web/src/pages/Models.tsx b/web/src/pages/Models.tsx index 6b65e3a..d94a5d5 100644 --- a/web/src/pages/Models.tsx +++ b/web/src/pages/Models.tsx @@ -927,6 +927,8 @@ export function Models() { {mediaKindLabels[kind] ?? kind} ))} + {m.embedding && 向量} + {m.rerank && 重排序} {m.available === false && m.enabled && ( 暂无候选 )}