feat: add secure agent and MCP extensions
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user