ai-ticket-platform/backend/apps/modelhub/services.py

from dataclasses import dataclass
from time import monotonic, sleep
from typing import Any

from django.utils import timezone
from openai import (
    APIConnectionError,
    APIStatusError,
    APITimeoutError,
    AuthenticationError,
    BadRequestError,
    OpenAI,
    RateLimitError,
)

from apps.modelhub.crypto import CredentialEncryptionError, decrypt_secret
from apps.modelhub.models import ModelDeployment, ModelInvocation

TRANSIENT_ERRORS = (APIConnectionError, APITimeoutError, RateLimitError)


@dataclass(frozen=True)
class CompletionResult:
    content: str
    input_tokens: int | None
    output_tokens: int | None


class ModelGatewayError(RuntimeError):
    def __init__(self, code: str, message: str):
        super().__init__(message)
        self.code = code


def _request_completion(
    deployment: ModelDeployment,
    messages: list[dict[str, Any]],
    *,
    temperature: float,
    max_output_tokens: int,
) -> CompletionResult:
    provider = deployment.provider
    client = OpenAI(
        api_key=decrypt_secret(provider.api_key_encrypted),
        base_url=provider.base_url,
        timeout=min(max(provider.default_timeout_seconds, 1), 300),
    )
    response = client.chat.completions.create(
        model=deployment.model_name,
        messages=messages,
        temperature=temperature,
        max_tokens=max_output_tokens,
    )
    content = response.choices[0].message.content or ""
    usage = response.usage
    return CompletionResult(
        content=content,
        input_tokens=getattr(usage, "prompt_tokens", None) if usage else None,
        output_tokens=getattr(usage, "completion_tokens", None) if usage else None,
    )


def _error_details(error: Exception) -> tuple[str, str, bool]:
    if isinstance(error, AuthenticationError):
        return "authentication_failed", "模型渠道鉴权失败", False
    if isinstance(error, BadRequestError):
        return "bad_request", "模型请求参数不被渠道接受", False
    if isinstance(error, RateLimitError):
        return "rate_limited", "模型渠道当前限流", True
    if isinstance(error, APITimeoutError):
        return "timeout", "模型请求超时", True
    if isinstance(error, APIConnectionError):
        return "connection_failed", "无法连接模型渠道", True
    if isinstance(error, APIStatusError) and error.status_code >= 500:
        return "provider_unavailable", "模型渠道暂时不可用", True
    if isinstance(error, CredentialEncryptionError):
        return "credential_unavailable", str(error), False
    return "unexpected_error", "模型调用发生未预期错误", False


def execute_invocation(
    invocation: ModelInvocation,
    messages: list[dict[str, Any]],
    *,
    temperature: float | None = None,
    max_output_tokens: int | None = None,
) -> ModelInvocation:
    deployment = invocation.deployment
    if deployment is None or not deployment.enabled or not deployment.provider.enabled:
        invocation.status = "failed"
        invocation.error_code = "deployment_unavailable"
        invocation.error_message = "模型配置或渠道未启用"
        invocation.finished_at = timezone.now()
        invocation.save(update_fields=["status", "error_code", "error_message", "finished_at"])
        return invocation

    invocation.status = "running"
    invocation.started_at = timezone.now()
    invocation.model_snapshot = {
        "provider_id": deployment.provider_id,
        "provider_name": deployment.provider.name,
        "base_url": deployment.provider.base_url,
        "deployment_id": deployment.id,
        "model_name": deployment.model_name,
        "capability": deployment.capability,
        "temperature": temperature if temperature is not None else deployment.default_temperature,
        "max_output_tokens": max_output_tokens or deployment.max_output_tokens,
    }
    invocation.save(update_fields=["status", "started_at", "model_snapshot"])
    started = monotonic()
    result: CompletionResult | None = None
    last_error: Exception | None = None
    for attempt in range(2):
        try:
            result = _request_completion(
                deployment,
                messages,
                temperature=temperature
                if temperature is not None
                else deployment.default_temperature,
                max_output_tokens=max_output_tokens or deployment.max_output_tokens,
            )
            invocation.retry_count = attempt
            break
        except Exception as error:  # provider libraries expose several subclasses
            code, message, transient = _error_details(error)
            last_error = error
            invocation.retry_count = attempt
            if not transient or attempt == 1:
                invocation.error_code = code
                invocation.error_message = message
                break
            sleep(1)

    invocation.duration_ms = max(int((monotonic() - started) * 1000), 0)
    invocation.finished_at = timezone.now()
    if result is not None:
        invocation.status = "succeeded"
        invocation.output_text = result.content
        invocation.input_tokens = result.input_tokens
        invocation.output_tokens = result.output_tokens
        invocation.error_code = ""
        invocation.error_message = ""
    else:
        invocation.status = "failed"
        if not invocation.error_code:
            code, message, _ = _error_details(last_error or RuntimeError())
            invocation.error_code = code
            invocation.error_message = message
    invocation.save()
    return invocation


def create_model_test(
    deployment: ModelDeployment,
    prompt: str,
    user=None,
) -> ModelInvocation:
    return ModelInvocation.objects.create(
        deployment=deployment,
        invocation_type="model_test",
        input_snapshot={"messages": [{"role": "user", "content": prompt}]},
        created_by=user,
    )