video-render/video_render/llm.py

from __future__ import annotations

import json
import logging
from pathlib import Path
from typing import Dict, List

import requests

from video_render.config import BASE_DIR, Settings
from video_render.transcription import TranscriptionResult

logger = logging.getLogger(__name__)

GEMINI_ENDPOINT_TEMPLATE = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
OPENROUTER_ENDPOINT = "https://openrouter.ai/api/v1/chat/completions"


class GeminiHighlighter:
    def __init__(self, settings: Settings) -> None:
        if not settings.gemini.api_key:
            raise RuntimeError("GEMINI_API_KEY nao foi definido")

        prompt_path = Path(settings.gemini.prompt_path)

        if not prompt_path.is_absolute():
            prompt_path = BASE_DIR / prompt_path

        if not prompt_path.exists():
            raise FileNotFoundError(f"Prompt do Gemini nao encontrado: {prompt_path}")

        self.prompt_template = prompt_path.read_text(encoding="utf-8")
        self.settings = settings

    def generate_highlights(self, transcription: TranscriptionResult) -> List[Dict]:
        payload = {
            "transcript": transcription.full_text,
            "segments": [
                {
                    "start": segment.start,
                    "end": segment.end,
                    "text": segment.text,
                }
                for segment in transcription.segments
            ],
        }

        body = {
            "contents": [
                {
                    "role": "user",
                    "parts": [
                        {"text": self.prompt_template},
                        {"text": json.dumps(payload, ensure_ascii=False)},
                    ],
                }
            ]
        }

        if self.settings.gemini.temperature is not None:
            body["generationConfig"] = {
                "temperature": self.settings.gemini.temperature,
            }
            if self.settings.gemini.top_p is not None:
                body["generationConfig"]["topP"] = self.settings.gemini.top_p
            if self.settings.gemini.top_k is not None:
                body["generationConfig"]["topK"] = self.settings.gemini.top_k

        url = GEMINI_ENDPOINT_TEMPLATE.format(model=self.settings.gemini.model)
        params = {"key": self.settings.gemini.api_key}

        response = requests.post(url, params=params, json=body, timeout=120)
        response.raise_for_status()
        data = response.json()

        candidates = data.get("candidates") or []
        if not candidates:
            raise RuntimeError("Gemini nao retornou candidatos")

        text_parts = candidates[0].get("content", {}).get("parts", [])
        if not text_parts:
            raise RuntimeError("Resposta do Gemini sem conteudo")

        raw_text = text_parts[0].get("text")
        if not raw_text:
            raise RuntimeError("Resposta do Gemini sem texto")

        parsed = self._extract_json(raw_text)
        highlights = parsed.get("highlights")
        if not isinstance(highlights, list):
            raise ValueError("Resposta do Gemini invalida: campo 'highlights' ausente")
        return highlights

    @staticmethod
    def _extract_json(response_text: str) -> Dict:
        try:
            return json.loads(response_text)
        except json.JSONDecodeError:
            start = response_text.find("{")
            end = response_text.rfind("}")
            if start == -1 or end == -1:
                raise
            subset = response_text[start : end + 1]
            return json.loads(subset)


class OpenRouterCopywriter:
    def __init__(self, settings: Settings) -> None:
        if not settings.openrouter.api_key:
            raise RuntimeError("OPENROUTER_API_KEY nao foi definido")
        self.settings = settings

    def generate_titles(self, highlights: List[Dict]) -> List[str]:
        if not highlights:
            return []

        prompt = (
            "Voce e um copywriter especializado em titulos curtos e virais para reels.\n"
            "Recebera uma lista de trechos destacados de um video com resumo e tempo.\n"
            "Produza um titulo envolvente (ate 60 caracteres) para cada item.\n"
            "Responda apenas em JSON com a seguinte estrutura:\n"
            '{"titles": ["titulo 1", "titulo 2"]}\n'
            "Titulos devem ser em portugues, usar verbos fortes e refletir o resumo."
        )

        user_payload = {
            "highlights": [
                {
                    "start": item.get("start"),
                    "end": item.get("end"),
                    "summary": item.get("summary"),
                }
                for item in highlights
            ]
        }

        body = {
            "model": self.settings.openrouter.model,
            "temperature": self.settings.openrouter.temperature,
            "max_tokens": self.settings.openrouter.max_output_tokens,
            "messages": [
                {"role": "system", "content": prompt},
                {
                    "role": "user",
                    "content": json.dumps(user_payload, ensure_ascii=False),
                },
            ],
        }

        headers = {
            "Authorization": f"Bearer {self.settings.openrouter.api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://localhost",
            "X-Title": "video-render-pipeline",
        }

        response = requests.post(
            OPENROUTER_ENDPOINT, json=body, headers=headers, timeout=120
        )
        response.raise_for_status()
        data = response.json()

        choices = data.get("choices") or []
        if not choices:
            raise RuntimeError("OpenRouter nao retornou escolhas")

        message = choices[0].get("message", {}).get("content")
        if not message:
            raise RuntimeError("Resposta do OpenRouter sem conteudo")

        parsed = self._extract_json(message)
        titles = parsed.get("titles")
        if not isinstance(titles, list):
            raise ValueError("Resposta do OpenRouter invalida: campo 'titles'")
        return [str(title) for title in titles]

    @staticmethod
    def _extract_json(response_text: str) -> Dict:
        try:
            return json.loads(response_text)
        except json.JSONDecodeError:
            start = response_text.find("{")
            end = response_text.rfind("}")
            if start == -1 or end == -1:
                raise
            subset = response_text[start : end + 1]
            return json.loads(subset)