video-render/video_render/llm.py

from __future__ import annotations

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

from google import genai
from google.genai import types as genai_types
import requests

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

logger = logging.getLogger(__name__)

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
        self.client = genai.Client()

    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
            ],
        }

        try:
            response = self._call_gemini(payload)
        except Exception as exc:
            logger.error("Gemini API request falhou: %s", exc)
            raise RuntimeError("Gemini API request falhou") from exc

        raw_text = self._extract_response_text(response)

        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

    def _call_gemini(self, payload: Dict[str, Any]) -> Any:
        contents = [
            {
                "role": "user",
                "parts": [
                    {"text": self.prompt_template},
                    {"text": json.dumps(payload, ensure_ascii=False)},
                ],
            }
        ]

        request_kwargs: Dict[str, Any] = {
            "model": self.settings.gemini.model,
            "contents": contents,
        }

        config = self._build_generation_config()
        if config is not None:
            request_kwargs["config"] = config

        return self.client.models.generate_content(**request_kwargs)

    def _build_generation_config(self) -> Optional[genai_types.GenerateContentConfig]:
        config_kwargs: Dict[str, Any] = {}
        if self.settings.gemini.temperature is not None:
            config_kwargs["temperature"] = self.settings.gemini.temperature
        if self.settings.gemini.top_p is not None:
            config_kwargs["top_p"] = self.settings.gemini.top_p
        if self.settings.gemini.top_k is not None:
            config_kwargs["top_k"] = self.settings.gemini.top_k

        if not config_kwargs:
            return None

        return genai_types.GenerateContentConfig(**config_kwargs)

    @staticmethod
    def _extract_response_text(response: Any) -> str:
        text = getattr(response, "text", None)
        if text:
            return str(text).strip()

        candidates = getattr(response, "candidates", None) or []
        for candidate in candidates:
            content = getattr(candidate, "content", None)
            if not content:
                continue
            parts = getattr(content, "parts", None) or []
            for part in parts:
                part_text = getattr(part, "text", None)
                if part_text:
                    return str(part_text).strip()

        raise RuntimeError("Resposta do Gemini sem texto")

    @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,
            "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",
        }

        response = requests.post(
            url=OPENROUTER_ENDPOINT,
            data=json.dumps(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)