142 lines
5.1 KiB
Python
142 lines
5.1 KiB
Python
import json
|
||
from datetime import timedelta
|
||
|
||
from celery import shared_task
|
||
from django.conf import settings
|
||
from django.db import transaction
|
||
from django.utils import timezone
|
||
from jsonschema import Draft202012Validator, ValidationError
|
||
|
||
from apps.modelhub.models import ModelInvocation
|
||
from apps.modelhub.services import execute_invocation
|
||
from apps.modelhub.structured_output import (
|
||
STRUCTURED_OUTPUT_TEST_EXPECTED,
|
||
STRUCTURED_OUTPUT_TEST_SCHEMA,
|
||
)
|
||
|
||
|
||
@shared_task(name="apps.modelhub.tasks.run_model_test")
|
||
def run_model_test(invocation_id: int) -> None:
|
||
invocation = ModelInvocation.objects.select_related("deployment__provider").get(
|
||
id=invocation_id
|
||
)
|
||
messages = invocation.input_snapshot.get("messages", [])
|
||
output_schema = invocation.input_snapshot.get("output_schema")
|
||
execute_invocation(
|
||
invocation,
|
||
messages,
|
||
output_schema=output_schema,
|
||
schema_name=invocation.input_snapshot.get("schema_name"),
|
||
)
|
||
if invocation.invocation_type != "structured_output_test":
|
||
return
|
||
|
||
deployment = invocation.deployment
|
||
if deployment is None:
|
||
return
|
||
message = invocation.error_message
|
||
passed = invocation.status == "succeeded"
|
||
if passed:
|
||
try:
|
||
output = json.loads(invocation.output_text)
|
||
Draft202012Validator(STRUCTURED_OUTPUT_TEST_SCHEMA).validate(output)
|
||
if output != STRUCTURED_OUTPUT_TEST_EXPECTED:
|
||
raise ValidationError("测试值与预期不一致")
|
||
except (json.JSONDecodeError, ValidationError) as error:
|
||
passed = False
|
||
invocation.status = "failed"
|
||
invocation.error_code = "structured_output_invalid"
|
||
invocation.error_message = f"结构化输出测试未通过：{error}"
|
||
invocation.save(update_fields=["status", "error_code", "error_message"])
|
||
message = invocation.error_message
|
||
|
||
deployment.structured_output_test_status = "passed" if passed else "failed"
|
||
deployment.structured_output_verified_at = timezone.now() if passed else None
|
||
deployment.structured_output_test_message = (
|
||
"结构化输出测试已通过" if passed else (message or "结构化输出测试失败")
|
||
)
|
||
deployment.save(
|
||
update_fields=[
|
||
"structured_output_test_status",
|
||
"structured_output_verified_at",
|
||
"structured_output_test_message",
|
||
"updated_at",
|
||
]
|
||
)
|
||
from apps.audit.models import OperationAuditLog
|
||
|
||
OperationAuditLog.objects.create(
|
||
user=invocation.created_by,
|
||
module="modelhub",
|
||
action="structured_output_test_result",
|
||
target_type="model_deployment",
|
||
target_id=str(deployment.id),
|
||
request_snapshot={
|
||
"invocation_id": str(invocation.public_id),
|
||
"status": deployment.structured_output_test_status,
|
||
"error_code": invocation.error_code,
|
||
},
|
||
result="success" if passed else "failed",
|
||
)
|
||
|
||
|
||
def _retained_input_metadata(snapshot: dict) -> dict:
|
||
structured = snapshot.get("structured_output")
|
||
retained: dict = {}
|
||
if isinstance(structured, dict):
|
||
retained["structured_output"] = {
|
||
key: structured.get(key)
|
||
for key in (
|
||
"mode",
|
||
"schema_name",
|
||
"schema_hash",
|
||
"response_format_sent",
|
||
)
|
||
if key in structured
|
||
}
|
||
if snapshot.get("schema_name"):
|
||
retained["schema_name"] = snapshot["schema_name"]
|
||
retained["diagnostics_scrubbed"] = True
|
||
return retained
|
||
|
||
|
||
@shared_task(name="apps.modelhub.tasks.scrub_expired_model_invocation_diagnostics")
|
||
def scrub_expired_model_invocation_diagnostics() -> int:
|
||
"""擦除过期诊断正文，保留调用元数据、Schema 证据、Token 与终态。"""
|
||
|
||
cutoff = timezone.now() - timedelta(
|
||
days=settings.MODEL_INVOCATION_DIAGNOSTIC_RETENTION_DAYS
|
||
)
|
||
from apps.agents.models import AgentRunAttempt
|
||
|
||
scrubbed_count = 0
|
||
while True:
|
||
invocation_ids = list(
|
||
ModelInvocation.objects.filter(
|
||
created_at__lt=cutoff,
|
||
diagnostics_scrubbed_at__isnull=True,
|
||
).values_list("id", flat=True)[:1000]
|
||
)
|
||
if not invocation_ids:
|
||
break
|
||
now = timezone.now()
|
||
with transaction.atomic():
|
||
invocations = list(
|
||
ModelInvocation.objects.select_for_update().filter(id__in=invocation_ids)
|
||
)
|
||
for invocation in invocations:
|
||
invocation.input_snapshot = _retained_input_metadata(
|
||
invocation.input_snapshot or {}
|
||
)
|
||
invocation.output_text = ""
|
||
invocation.diagnostics_scrubbed_at = now
|
||
ModelInvocation.objects.bulk_update(
|
||
invocations,
|
||
["input_snapshot", "output_text", "diagnostics_scrubbed_at"],
|
||
)
|
||
AgentRunAttempt.objects.filter(
|
||
model_invocation_id__in=invocation_ids
|
||
).update(raw_output="")
|
||
scrubbed_count += len(invocations)
|
||
return scrubbed_count
|