api/ai.py

from __future__ import annotations

from fastapi import APIRouter, Header, HTTPException, Request
from fastapi.concurrency import run_in_threadpool
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, ConfigDict, Field

from api.image_inputs import parse_image_edit_request, read_image_sources
from api.support import require_identity, resolve_image_base_url
from services.auth_service import auth_service
from services.content_filter import check_request, request_text
from services.log_service import LoggedCall
from services.protocol import (
    anthropic_v1_messages,
    openai_v1_chat_complete,
    openai_v1_image_edit,
    openai_v1_image_generations,
    openai_v1_models,
    openai_v1_response,
)
from utils.helper import has_response_image_generation_tool, is_image_chat_request, parse_image_count


class ImageGenerationRequest(BaseModel):
    prompt: str = Field(..., min_length=1)
    model: str = "gpt-image-2"
    n: int = Field(default=1, ge=1, le=4)
    size: str | None = None
    response_format: str = "b64_json"
    account_pool_strategy: str = "own_first"
    history_disabled: bool = True
    stream: bool | None = None


class ChatCompletionRequest(BaseModel):
    model_config = ConfigDict(extra="allow")
    model: str | None = None
    prompt: str | None = None
    n: int | None = None
    stream: bool | None = None
    modalities: list[str] | None = None
    messages: list[dict[str, object]] | None = None


class ResponseCreateRequest(BaseModel):
    model_config = ConfigDict(extra="allow")
    model: str | None = None
    input: object | None = None
    tools: list[dict[str, object]] | None = None
    tool_choice: object | None = None
    stream: bool | None = None


class AnthropicMessageRequest(BaseModel):
    model_config = ConfigDict(extra="allow")
    model: str | None = None
    messages: list[dict[str, object]] | None = None
    system: object | None = None
    stream: bool | None = None


async def filter_or_log(call: LoggedCall, text: str) -> None:
    try:
        await run_in_threadpool(check_request, text)
    except HTTPException as exc:
        call.log("调用失败", status="failed", error=str(exc.detail))
        raise


def _to_positive_image_count(value: object, default: int = 1) -> int:
    try:
        return max(1, int(value or default))
    except (TypeError, ValueError):
        return default


def _count_data_images(value: object) -> int:
    if not isinstance(value, list):
        return 0
    return sum(
        1
        for item in value
        if isinstance(item, dict) and (str(item.get("b64_json") or "").strip() or str(item.get("url") or "").strip())
    )


def _count_image_generation_items(value: object) -> int:
    if not isinstance(value, list):
        return 0
    count = 0
    for item in value:
        if not isinstance(item, dict):
            continue
        if str(item.get("type") or "") == "image_generation_call" and str(item.get("result") or "").strip():
            count += 1
    return count


def _count_data_url_images(text: object) -> int:
    return str(text or "").count("data:image/")


def _count_successful_images(value: object) -> int:
    if not isinstance(value, dict):
        return 0

    count = _count_data_images(value.get("data"))
    if count:
        return count

    output = value.get("output")
    count = _count_image_generation_items(output)
    if count:
        return count

    item = value.get("item")
    if isinstance(item, dict):
        count = _count_image_generation_items([item])
        if count:
            return count

    choices = value.get("choices")
    if isinstance(choices, list):
        return sum(
            _count_data_url_images(
                ((choice.get("message") if isinstance(choice, dict) else {}) or {}).get("content")
                if isinstance(choice, dict)
                else ""
            )
            + _count_data_url_images(
                ((choice.get("delta") if isinstance(choice, dict) else {}) or {}).get("content")
                if isinstance(choice, dict)
                else ""
            )
            for choice in choices
        )

    response = value.get("response")
    if isinstance(response, dict):
        return _count_successful_images(response)

    return 0


def _attach_account_pool(payload: dict[str, object], identity: dict[str, object]) -> None:
    has_personal_pool = identity.get("role") == "admin" or bool(identity.get("account_pool_enabled"))
    payload["account_owner_id"] = str(identity.get("id") or "").strip()
    if not has_personal_pool:
        payload["account_pool_strategy"] = "admin_only"
        payload["account_allow_shared"] = False
        return
    payload["account_pool_strategy"] = str(payload.get("account_pool_strategy") or "own_first")
    payload["account_allow_shared"] = bool(payload.get("account_allow_shared", True))


def _split_grants(grants: list[dict[str, object]], amount: int) -> tuple[list[dict[str, object]], list[dict[str, object]]]:
    remaining = max(0, int(amount or 0))
    used: list[dict[str, object]] = []
    refunded: list[dict[str, object]] = []
    for grant in grants:
        if not isinstance(grant, dict):
            continue
        grant_amount = max(0, int(grant.get("amount") or 0))
        grant_id = str(grant.get("id") or "").strip()
        if not grant_id or grant_amount <= 0:
            continue
        take = min(grant_amount, remaining)
        if take > 0:
            used.append({"id": grant_id, "amount": take})
        if grant_amount - take > 0:
            refunded.append({"id": grant_id, "amount": grant_amount - take})
        remaining -= take
    return used, refunded


class CreditReservation:
    def __init__(self, user_id: str = "", amount: int = 0):
        self.user_id = user_id
        self.amount = max(0, int(amount or 0))
        self.daily_amount = 0
        self.paid_amount = 0
        self.package_reservations: list[dict[str, object]] = []
        self.grant_reservations: list[dict[str, object]] = []
        self.settled = False

    @classmethod
    def reserve(cls, identity: dict[str, object], amount: int) -> "CreditReservation":
        if str(identity.get("role") or "").strip().lower() != "normal":
            return cls()
        reservation = cls(str(identity.get("id") or "").strip(), amount)
        if not reservation.user_id or reservation.amount <= 0:
            return reservation
        details = auth_service.reserve_credits(reservation.user_id, reservation.amount).get("_credit_reservation", {})
        if not isinstance(details, dict):
            details = {}
        reservation.daily_amount = max(0, int(details.get("daily") or 0))
        reservation.paid_amount = max(0, int(details.get("permanent_paid") if "permanent_paid" in details else details.get("paid") or 0))
        reservation.package_reservations = details.get("packages") if isinstance(details.get("packages"), list) else []
        reservation.grant_reservations = details.get("grants") if isinstance(details.get("grants"), list) else []
        return reservation

    def settle(self, success_count: int) -> None:
        if self.settled:
            return
        if not self.user_id or self.amount <= 0:
            self.settled = True
            return

        consume_count = min(self.amount, max(0, int(success_count or 0)))
        daily_to_consume = min(self.daily_amount, consume_count)
        remaining = consume_count - daily_to_consume
        packages_to_consume, packages_to_refund = _split_grants(self.package_reservations, remaining)
        remaining -= sum(max(0, int(item.get("amount") or 0)) for item in packages_to_consume)
        grants_to_consume, grants_to_refund = _split_grants(self.grant_reservations, remaining)
        remaining -= sum(max(0, int(item.get("amount") or 0)) for item in grants_to_consume)
        paid_to_consume = min(self.paid_amount, remaining)

        daily_to_refund = self.daily_amount - daily_to_consume
        paid_to_refund = self.paid_amount - paid_to_consume

        if consume_count > 0:
            auth_service.consume_reserved_credits(
                self.user_id,
                consume_count,
                daily_amount=daily_to_consume,
                paid_amount=paid_to_consume,
                package_reservations=packages_to_consume,
                grant_reservations=grants_to_consume,
            )
        refund_count = self.amount - consume_count
        if refund_count > 0:
            auth_service.refund_reserved_credits(
                self.user_id,
                refund_count,
                daily_amount=daily_to_refund,
                paid_amount=paid_to_refund,
                package_reservations=packages_to_refund,
                grant_reservations=grants_to_refund,
            )
        self.settled = True


async def run_metered_image_call(
    call: LoggedCall,
    handler,
    payload: dict[str, object],
    identity: dict[str, object],
    *,
    requested_amount: int | None = None,
):
    amount = _to_positive_image_count(requested_amount if requested_amount is not None else payload.get("n"), 1)
    try:
        reservation = CreditReservation.reserve(identity, amount)
    except ValueError as exc:
        raise HTTPException(status_code=402, detail={"error": str(exc)}) from exc

    stream_success_count = 0

    def on_stream_item(item: object) -> None:
        nonlocal stream_success_count
        stream_success_count += _count_successful_images(item)

    def on_stream_complete() -> None:
        reservation.settle(stream_success_count)

    def on_stream_error(_exc: BaseException) -> None:
        reservation.settle(stream_success_count)

    try:
        result = await call.run(
            handler,
            payload,
            on_stream_item=on_stream_item,
            on_stream_complete=on_stream_complete,
            on_stream_error=on_stream_error,
        )
    except Exception:
        reservation.settle(0)
        raise
    if isinstance(result, StreamingResponse):
        return result
    reservation.settle(_count_successful_images(result))
    return result


def create_router() -> APIRouter:
    router = APIRouter()

    @router.get("/v1/models")
    async def list_models(authorization: str | None = Header(default=None)):
        require_identity(authorization)
        try:
            return await run_in_threadpool(openai_v1_models.list_models)
        except Exception as exc:
            raise HTTPException(status_code=502, detail={"error": str(exc)}) from exc

    @router.post("/v1/images/generations")
    async def generate_images(
            body: ImageGenerationRequest,
            request: Request,
            authorization: str | None = Header(default=None),
    ):
        identity = require_identity(authorization)
        payload = body.model_dump(mode="python")
        payload["base_url"] = resolve_image_base_url(request)
        _attach_account_pool(payload, identity)
        call = LoggedCall(identity, "/v1/images/generations", body.model, "文生图", request_text=body.prompt)
        await filter_or_log(call, body.prompt)
        return await run_metered_image_call(call, openai_v1_image_generations.handle, payload, identity)

    @router.post("/v1/images/edits")
    async def edit_images(
            request: Request,
            authorization: str | None = Header(default=None),
    ):
        identity = require_identity(authorization)
        payload, image_sources = await parse_image_edit_request(request)
        prompt = str(payload["prompt"])
        model = str(payload["model"])
        call = LoggedCall(identity, "/v1/images/edits", model, "图生图", request_text=prompt)
        await filter_or_log(call, prompt)
        payload["images"] = await read_image_sources(image_sources)
        payload["base_url"] = resolve_image_base_url(request)
        _attach_account_pool(payload, identity)
        return await run_metered_image_call(call, openai_v1_image_edit.handle, payload, identity)

    @router.post("/v1/chat/completions")
    async def create_chat_completion(
        body: ChatCompletionRequest,
        request: Request,
        authorization: str | None = Header(default=None),
    ):
        identity = require_identity(authorization)
        payload = body.model_dump(mode="python")
        model = str(payload.get("model") or "auto")
        request_preview = request_text(payload.get("prompt"), payload.get("messages"))
        is_image_request = is_image_chat_request(payload)
        if is_image_request:
            payload["base_url"] = resolve_image_base_url(request)
            _attach_account_pool(payload, identity)
        call = LoggedCall(identity, "/v1/chat/completions", model, "图片生成" if is_image_request else "文本生成", request_text=request_preview)
        await filter_or_log(call, request_preview)
        if is_image_request:
            return await run_metered_image_call(
                call,
                openai_v1_chat_complete.handle,
                payload,
                identity,
                requested_amount=parse_image_count(payload.get("n")),
            )
        return await call.run(openai_v1_chat_complete.handle, payload)

    @router.post("/v1/responses")
    async def create_response(
        body: ResponseCreateRequest,
        request: Request,
        authorization: str | None = Header(default=None),
    ):
        identity = require_identity(authorization)
        payload = body.model_dump(mode="python")
        model = str(payload.get("model") or "auto")
        request_preview = request_text(payload.get("input"), payload.get("instructions"))
        is_image_request = has_response_image_generation_tool(payload)
        if is_image_request:
            payload["base_url"] = resolve_image_base_url(request)
            _attach_account_pool(payload, identity)
        call = LoggedCall(identity, "/v1/responses", model, "Responses 图片生成" if is_image_request else "Responses", request_text=request_preview)
        await filter_or_log(call, request_preview)
        if is_image_request:
            return await run_metered_image_call(
                call,
                openai_v1_response.handle,
                payload,
                identity,
                requested_amount=1,
            )
        return await call.run(openai_v1_response.handle, payload)

    @router.post("/v1/messages")
    async def create_message(
            body: AnthropicMessageRequest,
            authorization: str | None = Header(default=None),
            x_api_key: str | None = Header(default=None, alias="x-api-key"),
            anthropic_version: str | None = Header(default=None, alias="anthropic-version"),
    ):
        identity = require_identity(authorization or (f"Bearer {x_api_key}" if x_api_key else None))
        payload = body.model_dump(mode="python")
        model = str(payload.get("model") or "auto")
        request_preview = request_text(payload.get("system"), payload.get("messages"), payload.get("tools"))
        call = LoggedCall(identity, "/v1/messages", model, "Messages", request_text=request_preview)
        await filter_or_log(call, request_preview)
        return await call.run(anthropic_v1_messages.handle, payload, sse="anthropic")

    return router