NarratoAI/app/services/tavily_search.py

"""Tavily-powered web search helpers for plot analysis."""

from __future__ import annotations

import os
from typing import Any

import requests
from loguru import logger


TAVILY_API_BASE_URL = "https://api.tavily.com"
DEFAULT_SEARCH_DEPTH = "basic"
DEFAULT_MAX_RESULTS = 5
DEFAULT_TIMEOUT = 20


class TavilySearchError(RuntimeError):
    """Raised when Tavily search cannot be completed."""


def _trim_text(value: Any, max_chars: int) -> str:
    text = str(value or "").strip()
    if len(text) <= max_chars:
        return text
    return f"{text[:max_chars].rstrip()}..."


def search_short_drama(
    short_name: str,
    api_key: str | None = None,
    *,
    search_depth: str = DEFAULT_SEARCH_DEPTH,
    max_results: int = DEFAULT_MAX_RESULTS,
    timeout: int = DEFAULT_TIMEOUT,
) -> dict[str, Any]:
    """Search web context for a short drama name with Tavily."""
    return search_story_context(
        short_name,
        api_key,
        search_keywords="短剧 剧情 介绍 人物 结局",
        empty_name_message="短剧名称不能为空",
        search_depth=search_depth,
        max_results=max_results,
        timeout=timeout,
    )


def search_story_context(
    title: str,
    api_key: str | None = None,
    *,
    search_keywords: str = "剧情 介绍 人物 结局",
    empty_name_message: str = "作品名称不能为空",
    search_depth: str = DEFAULT_SEARCH_DEPTH,
    max_results: int = DEFAULT_MAX_RESULTS,
    timeout: int = DEFAULT_TIMEOUT,
) -> dict[str, Any]:
    """Search web context for a story title with Tavily."""
    title = str(title or "").strip()
    if not title:
        raise TavilySearchError(empty_name_message)

    api_key = (api_key or os.getenv("TAVILY_API_KEY") or "").strip()
    if not api_key:
        raise TavilySearchError("Tavily API Key 未配置")

    query = f"{title} {search_keywords}".strip()
    payload = {
        "query": query,
        "search_depth": search_depth or DEFAULT_SEARCH_DEPTH,
        "topic": "general",
        "max_results": max(1, min(int(max_results or DEFAULT_MAX_RESULTS), 10)),
        "include_answer": True,
        "include_raw_content": False,
        "include_images": False,
    }

    try:
        response = requests.post(
            f"{TAVILY_API_BASE_URL}/search",
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            json=payload,
            timeout=timeout,
        )
    except requests.RequestException as exc:
        raise TavilySearchError(f"Tavily 请求失败: {exc}") from exc

    if response.status_code >= 400:
        message = _trim_text(response.text, 500)
        raise TavilySearchError(f"Tavily 请求失败: HTTP {response.status_code} {message}")

    try:
        data = response.json()
    except ValueError as exc:
        raise TavilySearchError("Tavily 返回内容不是有效 JSON") from exc

    logger.info(
        "Tavily 剧情检索完成: query={}, results={}",
        query,
        len(data.get("results") or []),
    )
    return data

def format_search_context(search_data: dict[str, Any], *, max_chars: int = 6000) -> str:
    """Format Tavily response into compact LLM context."""
    if not search_data:
        return ""

    lines = [
        "# Tavily 联网检索结果",
        f"检索 query: {search_data.get('query', '')}",
    ]

    answer = _trim_text(search_data.get("answer"), 1200)
    if answer:
        lines.extend(["", "## 综合回答", answer])

    results = search_data.get("results") or []
    if results:
        lines.extend(["", "## 搜索来源"])
    for index, result in enumerate(results, start=1):
        title = _trim_text(result.get("title"), 120)
        url = _trim_text(result.get("url"), 240)
        content = _trim_text(result.get("content") or result.get("raw_content"), 700)
        lines.extend(
            [
                f"{index}. 标题: {title}",
                f"   来源: {url}",
                f"   摘要: {content}",
            ]
        )

    return _trim_text("\n".join(lines).strip(), max_chars)