yingqing/core/backend/apps/ai/providers/volcano.py

from dataclasses import dataclass
import base64
import json
import uuid
from io import BytesIO
from typing import Any, Iterator

import requests
from django.conf import settings

from .base import AIProviderResult


@dataclass
class VolcanoArkProvider:
    api_key: str | None = None
    base_url: str | None = None

    def __post_init__(self) -> None:
        self.api_key = self.api_key or settings.VOLCANO.get("ark_api_key")
        self.base_url = self.base_url or settings.VOLCANO.get("ark_base_url")

    def submit(self, payload: dict[str, Any]) -> AIProviderResult:
        # The exact endpoint is resolved by ModelConfig; this adapter keeps IO centralized.
        endpoint = payload.get("endpoint")
        if not endpoint:
            raise ValueError("Volcano request payload requires endpoint")

        response = requests.post(
            f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
            headers={"Authorization": f"Bearer {self.api_key}"},
            json=payload.get("body", {}),
            timeout=60,
        )
        response.raise_for_status()
        data = response.json()
        return AIProviderResult(
            provider_task_id=str(data.get("id") or data.get("task_id") or ""),
            status=str(data.get("status") or "submitted"),
            payload=data,
        )

    def chat_completion(self, *, model: str, messages: list[dict[str, str]], endpoint: str = "chat/completions") -> dict[str, Any]:
        if not self.api_key:
            raise ValueError("VOLCANO_ARK_API_KEY is not configured")

        response = requests.post(
            f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
            json={"model": model, "messages": messages},
            timeout=120,
        )
        response.raise_for_status()
        return response.json()

    def chat_completion_stream(
        self,
        *,
        model: str,
        messages: list[dict[str, Any]],
        endpoint: str = "chat/completions",
        temperature: float = 0.8,
        extra_body: dict[str, Any] | None = None,
    ) -> Iterator[dict[str, Any]]:
        """流式对话:逐块 yield {type:'delta'|'tool_call'|'done', ...}。
        OpenAI 兼容 SSE(火山 ARK / 各中转站同构),供脚本 agent 的 SSE 端实时转发。"""
        if not self.api_key:
            raise ValueError("api_key is not configured")
        body: dict[str, Any] = {"model": model, "messages": messages, "stream": True, "temperature": temperature}
        if extra_body:
            body.update(extra_body)
        with requests.post(
            f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "Accept": "text/event-stream",
            },
            json=body,
            stream=True,
            timeout=300,
        ) as response:
            response.raise_for_status()
            # SSE 响应常不带 charset,requests 会按 latin-1 解码 → 中文乱码。强制 UTF-8。
            response.encoding = "utf-8"
            for raw in response.iter_lines(decode_unicode=True):
                if not raw or not raw.startswith("data:"):
                    continue
                data = raw[5:].strip()
                if data == "[DONE]":
                    break
                try:
                    chunk = json.loads(data)
                except ValueError:
                    continue
                choices = chunk.get("choices") or []
                if not choices:
                    continue
                delta = choices[0].get("delta") or {}
                # 推理模型(豆包 seed-pro / 部分中转 o系/gemini)思考阶段只发 reasoning_content,
                # 不发 content。必须单独转发,否则整个思考期(可达几十秒~分钟)前端零输出 = 假死。
                reason = delta.get("reasoning_content")
                if reason:
                    yield {"type": "reasoning", "text": reason}
                piece = delta.get("content")
                if piece:
                    yield {"type": "delta", "text": piece}
                if delta.get("tool_calls"):
                    yield {"type": "tool_call", "tool_calls": delta["tool_calls"]}
        yield {"type": "done"}

    @staticmethod
    def extract_text(data: dict[str, Any]) -> str:
        choices = data.get("choices") or []
        if choices:
            message = choices[0].get("message") or {}
            content = message.get("content")
            if isinstance(content, str):
                return content
            if isinstance(content, list):
                return "\n".join(str(item.get("text", "")) for item in content if isinstance(item, dict))
        output = data.get("output")
        if isinstance(output, str):
            return output
        raise ValueError("Volcano response does not contain text content")

    def poll(self, provider_task_id: str) -> AIProviderResult:
        if not provider_task_id:
            raise ValueError("provider_task_id is required")

        return AIProviderResult(provider_task_id=provider_task_id, status="polling", payload={})

    def image_generation(
        self,
        *,
        model: str,
        prompt: str,
        endpoint: str = "images/generations",
        image: str | list[str] | None = None,
        size: str = "2K",
    ) -> dict[str, Any]:
        if not self.api_key:
            raise ValueError("VOLCANO_ARK_API_KEY is not configured")
        body: dict[str, Any] = {
            "model": model,
            "prompt": prompt,
            "response_format": "url",
            "watermark": False,
            "size": size,
            "sequential_image_generation": "disabled",
        }
        if image:
            body["image"] = image
        response = requests.post(
            f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
            json=body,
            timeout=180,
        )
        response.raise_for_status()
        return response.json()

    def create_video_task(
        self,
        *,
        model: str,
        endpoint: str,
        prompt: str,
        ratio: str = "9:16",
        duration: int = 15,
        resolution: str = "720p",
        reference_images: list[str] | None = None,
        generate_audio: bool = True,
    ) -> dict[str, Any]:
        if not self.api_key:
            raise ValueError("VOLCANO_ARK_API_KEY is not configured")
        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
        for image_url in reference_images or []:
            content.append({"type": "image_url", "image_url": {"url": image_url}, "role": "reference_image"})
        body = {
            "model": model,
            "content": content,
            "ratio": ratio,
            "duration": duration,
            "resolution": resolution,
            "watermark": False,
            # Seedance 直接出音效 + 人物声音(参考生视频);关掉则是哑片。默认开。
            "generate_audio": generate_audio,
        }
        response = requests.post(
            f"{self.base_url.rstrip('/')}/{endpoint.lstrip('/')}",
            headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
            json=body,
            timeout=120,
        )
        response.raise_for_status()
        return response.json()

    def poll_video_task(self, *, endpoint: str, provider_task_id: str) -> dict[str, Any]:
        if not self.api_key:
            raise ValueError("VOLCANO_ARK_API_KEY is not configured")
        response = requests.get(
            f"{self.base_url.rstrip('/')}/{endpoint.rstrip('/')}/{provider_task_id}",
            headers={"Authorization": f"Bearer {self.api_key}"},
            timeout=60,
        )
        response.raise_for_status()
        return response.json()

    @staticmethod
    def extract_first_media_url(data: dict[str, Any]) -> str:
        items = data.get("data") or []
        for item in items:
            if item.get("url"):
                return item["url"]
            if item.get("b64_json"):
                return item["b64_json"]
        content = data.get("content") or {}
        if content.get("video_url"):
            return content["video_url"]
        raise ValueError("Volcano response does not contain media url")

    @staticmethod
    def media_to_bytes(media: str) -> tuple[BytesIO, str]:
        if media.startswith("http://") or media.startswith("https://"):
            response = requests.get(media, timeout=180)
            response.raise_for_status()
            return BytesIO(response.content), response.headers.get("content-type", "application/octet-stream")
        if "," in media and media.startswith("data:"):
            header, raw = media.split(",", 1)
            content_type = header.split(";")[0].replace("data:", "") or "application/octet-stream"
            return BytesIO(base64.b64decode(raw)), content_type
        return BytesIO(base64.b64decode(media)), "image/png"


class TtsNotConfigured(ValueError):
    """语音合成凭证未配置(VOLC_TTS_APPID / VOLC_TTS_ACCESS_TOKEN)。"""


@dataclass
class VolcanoTtsProvider:
    """豆包语音合成大模型(openspeech V1 HTTP,非流式)。
    注意:与 ARK 不是同一套钥匙,要在 火山引擎控制台→语音技术→语音合成大模型 创建应用拿 APPID/Access Token。"""

    appid: str | None = None
    access_token: str | None = None
    cluster: str | None = None
    base_url: str | None = None

    def __post_init__(self) -> None:
        cfg = getattr(settings, "VOLC_TTS", {}) or {}
        self.appid = self.appid or cfg.get("appid")
        self.access_token = self.access_token or cfg.get("access_token")
        self.cluster = self.cluster or cfg.get("cluster") or "volcano_tts"
        self.base_url = self.base_url or cfg.get("base_url") or "https://openspeech.bytedance.com/api/v1/tts"

    @property
    def configured(self) -> bool:
        return bool(self.appid and self.access_token)

    def synthesize(self, *, text: str, voice_type: str, speed_ratio: float = 1.0, uid: str = "airshelf") -> tuple[bytes, int]:
        """合成一段语音。返回 (mp3 字节, 时长毫秒;接口没回时长则为 0)。"""
        if not self.configured:
            raise TtsNotConfigured(
                "语音合成未配置:请在后端环境变量设置 VOLC_TTS_APPID 和 VOLC_TTS_ACCESS_TOKEN"
                "(火山引擎控制台 → 语音技术 → 语音合成大模型 → 创建应用)"
            )
        body = {
            "app": {"appid": self.appid, "token": self.access_token, "cluster": self.cluster},
            "user": {"uid": uid},
            "audio": {"voice_type": voice_type, "encoding": "mp3", "speed_ratio": float(speed_ratio or 1.0)},
            "request": {"reqid": str(uuid.uuid4()), "text": text, "operation": "query"},
        }
        response = requests.post(
            self.base_url,
            headers={"Authorization": f"Bearer;{self.access_token}"},
            json=body,
            timeout=60,
        )
        response.raise_for_status()
        data = response.json()
        if data.get("code") != 3000 or not data.get("data"):
            raise ValueError(f"语音合成失败:{data.get('message') or data.get('code')}")
        audio = base64.b64decode(data["data"])
        duration_ms = 0
        try:
            duration_ms = int(float((data.get("addition") or {}).get("duration") or 0))
        except (TypeError, ValueError):
            duration_ms = 0
        return audio, duration_ms