138 lines
4.4 KiB
Python
138 lines
4.4 KiB
Python
from django.contrib.auth import get_user_model
|
|
from django.core.management import call_command
|
|
|
|
from apps.modelhub.crypto import decrypt_secret
|
|
from apps.modelhub.models import ModelProvider
|
|
from apps.modelhub.services import CompletionResult
|
|
|
|
|
|
def login_admin(client):
|
|
call_command("bootstrap_governance")
|
|
user = get_user_model().objects.get(username="admin")
|
|
client.force_login(user)
|
|
return user
|
|
|
|
|
|
def create_provider(client):
|
|
return client.post(
|
|
"/api/modelhub/providers",
|
|
data={
|
|
"name": "企业模型渠道",
|
|
"base_url": "https://models.example.com/v1/",
|
|
"api_key": "sk-secret-1234",
|
|
"default_timeout_seconds": 30,
|
|
},
|
|
content_type="application/json",
|
|
)
|
|
|
|
|
|
def test_admin_manages_provider_without_exposing_secret(client, db):
|
|
login_admin(client)
|
|
|
|
response = create_provider(client)
|
|
|
|
assert response.status_code == 200
|
|
assert "api_key" not in response.json()
|
|
assert response.json()["credential_masked"] == "****1234"
|
|
provider = ModelProvider.objects.get()
|
|
assert provider.api_key_encrypted != "sk-secret-1234"
|
|
assert decrypt_secret(provider.api_key_encrypted) == "sk-secret-1234"
|
|
|
|
updated = client.patch(
|
|
f"/api/modelhub/providers/{provider.id}",
|
|
data={"name": "新渠道名称", "api_key": ""},
|
|
content_type="application/json",
|
|
)
|
|
assert updated.status_code == 200
|
|
provider.refresh_from_db()
|
|
assert decrypt_secret(provider.api_key_encrypted) == "sk-secret-1234"
|
|
|
|
replaced = client.patch(
|
|
f"/api/modelhub/providers/{provider.id}",
|
|
data={"api_key": "sk-new-secret-5678"},
|
|
content_type="application/json",
|
|
)
|
|
assert replaced.status_code == 200
|
|
provider.refresh_from_db()
|
|
assert decrypt_secret(provider.api_key_encrypted) == "sk-new-secret-5678"
|
|
assert replaced.json()["credential_masked"] == "****5678"
|
|
|
|
|
|
def test_model_test_runs_through_recorded_invocation(client, db, monkeypatch):
|
|
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": "support-model",
|
|
},
|
|
content_type="application/json",
|
|
)
|
|
assert deployment.status_code == 200
|
|
|
|
monkeypatch.setattr(
|
|
"apps.modelhub.services._request_completion",
|
|
lambda *args, **kwargs: CompletionResult("连接正常", 8, 4),
|
|
)
|
|
response = client.post(
|
|
f"/api/modelhub/deployments/{deployment.json()['id']}/test",
|
|
data={"prompt": "测试连接"},
|
|
content_type="application/json",
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["status"] == "succeeded"
|
|
assert response.json()["output_text"] == "连接正常"
|
|
assert response.json()["input_tokens"] == 8
|
|
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)
|
|
|
|
response = client.get("/api/modelhub/providers")
|
|
|
|
assert response.status_code == 403
|