HaiLin/ai-ticket-platform
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add secure agent and MCP extensions

Browse Source
This commit is contained in:
hailin
2026-07-27 19:02:37 +08:00
parent 07c9ea1cc4
commit d93cb75c10
29 changed files with 3186 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
backend/apps/modelhub/api.py
View File

@@ -158,7 +158,7 @@ def create_deployment(request, payload: ModelDeploymentIn):
raise HttpError(400, "Temperature must be between 0 and 2")
if not 1 <= payload.max_output_tokens <= 32768:
raise HttpError(400, "Max output tokens are out of range")
if payload.capability not in {"text_chat", "vision_chat"}:
if payload.capability not in {"text_chat", "vision_chat", "embedding", "rerank"}:
raise HttpError(400, "Unknown model capability")
deployment = ModelDeployment.objects.create(**payload.dict())
deployment = ModelDeployment.objects.select_related("provider").get(id=deployment.id)
@@ -177,7 +177,12 @@ def update_deployment(request, deployment_id: int, payload: ModelDeploymentUpdat
require_permission(request, "model.manage")
deployment = ModelDeployment.objects.select_related("provider").get(id=deployment_id)
values = payload.dict(exclude_unset=True)
if "capability" in values and values["capability"] not in {"text_chat", "vision_chat"}:
if "capability" in values and values["capability"] not in {
"text_chat",
"vision_chat",
"embedding",
"rerank",
}:
raise HttpError(400, "Unknown model capability")
for field, value in values.items():
setattr(deployment, field, value)
@@ -196,6 +201,8 @@ def update_deployment(request, deployment_id: int, payload: ModelDeploymentUpdat
def test_deployment(request, deployment_id: int, payload: ModelTestIn):
require_permission(request, "model.manage")
deployment = ModelDeployment.objects.select_related("provider").get(id=deployment_id)
if deployment.capability == "embedding":
raise HttpError(400, "Embedding 模型请通过知识检索验证")
invocation = create_model_test(deployment, payload.prompt, request.user)
run_model_test.delay(invocation.id)
invocation.refresh_from_db()

39
backend/apps/modelhub/services.py
View File

@@ -1,3 +1,9 @@
"""模型调用执行器。
统一处理超时预算、可重试错误、凭据脱敏、调用快照和 Token/耗时记录。
"""
from collections.abc import Callable
from dataclasses import dataclass
from time import monotonic, sleep
from typing import Any
@@ -38,12 +44,18 @@ def _request_completion(
*,
temperature: float,
max_output_tokens: int,
timeout_seconds: float | None = None,
) -> CompletionResult:
provider = deployment.provider
request_timeout = min(max(provider.default_timeout_seconds, 1), 300)
if timeout_seconds is not None:
request_timeout = min(request_timeout, max(timeout_seconds, 1))
client = OpenAI(
api_key=decrypt_secret(provider.api_key_encrypted),
base_url=provider.base_url,
timeout=min(max(provider.default_timeout_seconds, 1), 300),
timeout=request_timeout,
# Retry here so the platform can expose it instead of waiting on opaque SDK retries.
max_retries=0,
)
response = client.chat.completions.create(
model=deployment.model_name,
@@ -84,7 +96,11 @@ def execute_invocation(
*,
temperature: float | None = None,
max_output_tokens: int | None = None,
on_retry: Callable[[int, str], None] | None = None,
total_timeout_seconds: float | None = None,
) -> ModelInvocation:
"""在总时间预算内调用模型，分类错误、限制重试并记录脱敏终态。"""
deployment = invocation.deployment
if deployment is None or not deployment.enabled or not deployment.provider.enabled:
invocation.status = "failed"
@@ -111,6 +127,12 @@ def execute_invocation(
result: CompletionResult | None = None
last_error: Exception | None = None
for attempt in range(2):
elapsed = monotonic() - started
remaining = total_timeout_seconds - elapsed if total_timeout_seconds is not None else None
if remaining is not None and remaining <= 0:
invocation.error_code = "timeout"
invocation.error_message = "模型请求超过 Agent 时间预算"
break
try:
result = _request_completion(
deployment,
@@ -119,6 +141,7 @@ def execute_invocation(
if temperature is not None
else deployment.default_temperature,
max_output_tokens=max_output_tokens or deployment.max_output_tokens,
timeout_seconds=remaining,
)
invocation.retry_count = attempt
break
@@ -130,6 +153,18 @@ def execute_invocation(
invocation.error_code = code
invocation.error_message = message
break
if total_timeout_seconds is not None:
remaining_after_error = total_timeout_seconds - (monotonic() - started)
if remaining_after_error <= 1:
invocation.error_code = code
invocation.error_message = message
break
invocation.retry_count = attempt + 1
invocation.error_code = code
invocation.error_message = message
invocation.save(update_fields=["retry_count", "error_code", "error_message"])
if on_retry:
on_retry(attempt + 1, code)
sleep(1)
invocation.duration_ms = max(int((monotonic() - started) * 1000), 0)
@@ -156,6 +191,8 @@ def create_model_test(
prompt: str,
user=None,
) -> ModelInvocation:
"""创建可审计的异步模型连通测试，不绕过正式模型执行器。"""
return ModelInvocation.objects.create(
deployment=deployment,
invocation_type="model_test",

39
backend/apps/modelhub/tests/test_modelhub_api.py
View File

@@ -89,6 +89,45 @@ def test_model_test_runs_through_recorded_invocation(client, db, monkeypatch):
assert response.json()["output_tokens"] == 4
def test_embedding_deployment_can_be_configured_but_uses_knowledge_search_for_test(client, db):
login_admin(client)
provider_id = create_provider(client).json()["id"]
deployment = client.post(
"/api/modelhub/deployments",
data={
"provider_id": provider_id,
"display_name": "知识向量模型",
"model_name": "embedding-model",
"capability": "embedding",
},
content_type="application/json",
)
assert deployment.status_code == 200
assert deployment.json()["capability"] == "embedding"
response = client.post(
f"/api/modelhub/deployments/{deployment.json()['id']}/test",
data={"prompt": "测试连接"},
content_type="application/json",
)
assert response.status_code == 400
assert "知识检索" in response.json()["detail"]
rerank = client.post(
"/api/modelhub/deployments",
data={
"provider_id": provider_id,
"display_name": "知识重排模型",
"model_name": "rerank-model",
"capability": "rerank",
},
content_type="application/json",
)
assert rerank.status_code == 200
assert rerank.json()["capability"] == "rerank"
def test_non_ai_admin_cannot_manage_modelhub(client, db):
user = get_user_model().objects.create_user(username="requester", password="password")
client.force_login(user)

33
backend/apps/modelhub/tests/test_services.py
View File

@@ -27,6 +27,7 @@ def create_invocation():
def test_transient_error_retries_once(db, monkeypatch):
invocation = create_invocation()
calls = []
retries = []
def request(*args, **kwargs):
calls.append(1)
@@ -37,12 +38,17 @@ def test_transient_error_retries_once(db, monkeypatch):
monkeypatch.setattr("apps.modelhub.services._request_completion", request)
monkeypatch.setattr("apps.modelhub.services.sleep", lambda *_: None)
execute_invocation(invocation, invocation.input_snapshot["messages"])
execute_invocation(
invocation,
invocation.input_snapshot["messages"],
on_retry=lambda retry_number, code: retries.append((retry_number, code)),
)
invocation.refresh_from_db()
assert len(calls) == 2
assert invocation.status == "succeeded"
assert invocation.retry_count == 1
assert retries == [(1, "timeout")]
def test_non_transient_error_fails_without_exposing_raw_error(db, monkeypatch):
@@ -61,6 +67,31 @@ def test_non_transient_error_fails_without_exposing_raw_error(db, monkeypatch):
assert "secret provider detail" not in invocation.error_message
def test_agent_time_budget_does_not_repeat_a_full_timeout(db, monkeypatch):
invocation = create_invocation()
calls = []
clock = iter([0.0, 0.0, 19.5, 20.0])
def request(*args, **kwargs):
calls.append(kwargs["timeout_seconds"])
raise APITimeoutError(request=None)
monkeypatch.setattr("apps.modelhub.services._request_completion", request)
monkeypatch.setattr("apps.modelhub.services.monotonic", lambda: next(clock))
execute_invocation(
invocation,
invocation.input_snapshot["messages"],
total_timeout_seconds=20,
)
invocation.refresh_from_db()
assert calls == [20]
assert invocation.status == "failed"
assert invocation.error_code == "timeout"
assert invocation.retry_count == 0
def provider_error(error_type, status_code):
request = httpx.Request("POST", "https://models.example.com/v1/chat/completions")
response = httpx.Response(status_code, request=request)