Agent/agent/tools/builtin/content/platforms/aigc_channel.py

"""
AIGC-Channel 平台实现（9 个中文平台）

后端：aigc-channel.aiddit.com
平台：xhs / gzh / sph / github / toutiao / douyin / bili / zhihu / weibo
"""

import json
from typing import Any, Dict, List, Optional

import httpx

from agent.tools.models import ToolResult
from agent.tools.utils.image import build_image_grid, encode_base64, load_images
from agent.tools.builtin.content.registry import (
    PlatformDef, ParamSpec, register_platform,
)

BASE_URL = "http://aigc-channel.aiddit.com/aigc/channel"
DEFAULT_TIMEOUT = 60.0


# ── 平台注册 ──

_XHS_SEARCH_PARAMS = {
    "sort_type": ParamSpec(
        values=["综合排序", "最新发布", "最多点赞"],
        default="综合排序",
    ),
    "publish_time": ParamSpec(
        values=["不限", "近1天", "近7天", "近30天"],
        default="不限",
    ),
    "content_type": ParamSpec(
        values=["不限", "图文", "视频", "文章"],
        default="不限",
    ),
    "filter_note_range": ParamSpec(
        values=["不限", "1分钟以内", "1-5分钟", "5分钟以上"],
        default="不限",
        note="仅视频内容生效",
    ),
}

_COMMON_CONTENT_TYPE = {
    "content_type": ParamSpec(
        values=["视频", "图文"],
        default="",
        note="留空不限",
    ),
}

# 9 个中文平台定义
_AIGC_PLATFORMS = [
    PlatformDef(id="xhs",     name="小红书",   aliases=["RED", "xiaohongshu"], search_params=_XHS_SEARCH_PARAMS, supports_suggest=True),
    PlatformDef(id="gzh",     name="公众号",   aliases=["微信公众号", "wechat"], search_params=_COMMON_CONTENT_TYPE),
    PlatformDef(id="sph",     name="视频号",   aliases=["微信视频号"], search_params=_COMMON_CONTENT_TYPE),
    PlatformDef(id="github",  name="GitHub",   aliases=["gh"], search_params=_COMMON_CONTENT_TYPE),
    PlatformDef(id="toutiao", name="头条",     aliases=["今日头条", "toutiao"], search_params=_COMMON_CONTENT_TYPE, supports_suggest=True),
    PlatformDef(id="douyin",  name="抖音",     aliases=["TikTok"], search_params=_COMMON_CONTENT_TYPE, supports_suggest=True),
    PlatformDef(id="bili",    name="B站",      aliases=["哔哩哔哩", "bilibili"], search_params=_COMMON_CONTENT_TYPE, supports_suggest=True),
    PlatformDef(id="zhihu",   name="知乎",     aliases=[], search_params=_COMMON_CONTENT_TYPE, supports_suggest=True),
    PlatformDef(id="weibo",   name="微博",     aliases=["sina"], search_params=_COMMON_CONTENT_TYPE),
]

# suggest API 额外支持 wx（微信搜一搜），但它不是搜索平台
_SUGGEST_ONLY_CHANNELS = {"wx": "微信"}


# ── 搜索实现 ──

async def search(
    platform_id: str,
    keyword: str,
    max_count: int = 20,
    cursor: str = "",
    extras: Optional[Dict[str, Any]] = None,
) -> ToolResult:
    """AIGC-Channel 统一搜索"""
    extras = extras or {}

    if platform_id == "xhs":
        payload = {
            "type": platform_id,
            "keyword": keyword,
            "cursor": cursor,
            "content_type": extras.get("content_type", "不限"),
            "sort_type": extras.get("sort_type", "综合排序"),
            "publish_time": extras.get("publish_time", "不限"),
            "filter_note_range": extras.get("filter_note_range", "不限"),
        }
    else:
        payload = {
            "type": platform_id,
            "keyword": keyword,
            "cursor": cursor or "0",
            "max_count": max_count,
            "content_type": extras.get("content_type", ""),
        }

    try:
        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
            response = await client.post(
                f"{BASE_URL}/data",
                json=payload,
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            data = response.json()
    except httpx.HTTPStatusError as e:
        return ToolResult(title="搜索失败", output="", error=f"HTTP {e.response.status_code}: {e.response.text}")
    except Exception as e:
        return ToolResult(title="搜索失败", output="", error=str(e))

    posts = data.get("data", [])

    # 构建概览摘要
    summary_list = []
    for idx, post in enumerate(posts, 1):
        body = post.get("body_text", "") or ""
        title = post.get("title") or body[:20] or ""
        summary_list.append({
            "index": idx,
            "title": title,
            "body_text": body[:100] + ("..." if len(body) > 100 else ""),
            "like_count": post.get("like_count"),
            "comment_count": post.get("comment_count"),
            "channel": post.get("channel"),
            "link": post.get("link"),
            "content_type": post.get("content_type"),
        })

    # 封面拼图
    images = []
    try:
        collage_obj = await _build_collage(posts)
        if collage_obj:
            images.append(collage_obj)
    except Exception as e:
        import logging
        logging.getLogger(__name__).warning("Error generating collage: %s", e)

    return ToolResult(
        title=f"搜索: {keyword} ({platform_id})",
        output=json.dumps({"data": summary_list}, ensure_ascii=False, indent=2),
        long_term_memory=f"Searched '{keyword}' on {platform_id}, {len(posts)} results. Use content_detail to view full details.",
        images=images,
        metadata={"posts": posts},  # 完整数据传给上层缓存
    )


# ── 详情实现（从缓存获取，不需要额外 HTTP） ──

async def detail(post: Dict[str, Any], extras: Optional[Dict[str, Any]] = None) -> ToolResult:
    """返回单条帖子的完整内容"""
    title = post.get("title") or post.get("body_text", "")[:30] or "无标题"

    images = []
    for img_url in post.get("images", []):
        if img_url:
            images.append({"type": "url", "url": img_url})

    return ToolResult(
        title=f"详情: {title}",
        output=json.dumps(post, ensure_ascii=False, indent=2),
        long_term_memory=f"Viewed detail: {title}",
        images=images,
    )


# ── 建议词实现 ──

async def suggest(channel: str, keyword: str) -> ToolResult:
    """获取搜索建议词"""
    try:
        async with httpx.AsyncClient(timeout=DEFAULT_TIMEOUT) as client:
            response = await client.post(
                f"{BASE_URL}/suggest",
                json={"type": channel, "keyword": keyword},
                headers={"Content-Type": "application/json"},
            )
            response.raise_for_status()
            data = response.json()
    except Exception as e:
        return ToolResult(title="建议词获取失败", output="", error=str(e))

    suggestion_count = sum(len(item.get("list", [])) for item in data.get("data", []))
    return ToolResult(
        title=f"建议词: {keyword} ({channel})",
        output=json.dumps(data, ensure_ascii=False, indent=2),
        long_term_memory=f"Got {suggestion_count} suggestions for '{keyword}' on {channel}",
    )


# ── 拼图辅助 ──

async def _build_collage(posts: List[Dict[str, Any]]) -> Optional[str]:
    """封面图网格拼图"""
    urls, titles = [], []
    for post in posts:
        imgs = post.get("images", [])
        if imgs and imgs[0]:
            urls.append(imgs[0])
            titles.append(post.get("title", "") or "")

    if not urls:
        return None

    loaded = await load_images(urls)
    valid_images, valid_labels = [], []
    for (_, img), title in zip(loaded, titles):
        if img is not None:
            valid_images.append(img)
            valid_labels.append(title)

    if not valid_images:
        return None

    grid = build_image_grid(images=valid_images, labels=valid_labels)
    import io
    buf = io.BytesIO()
    grid.save(buf, format="PNG")
    img_bytes = buf.getvalue()
    
    # 尝试上传到 CDN，替换冗长的 base64
    try:
        from agent.tools.builtin.file.image_cdn import _upload_bytes_to_oss
        import hashlib
        
        md5_hash = hashlib.md5(img_bytes).hexdigest()[:12]
        filename = f"collage_search_{md5_hash}.png"
        cdn_url = await _upload_bytes_to_oss(img_bytes, filename)
        return {"type": "url", "url": cdn_url}
    except Exception as e:
        import logging
        logging.getLogger(__name__).warning("Failed to upload collage to CDN: %s", e)
        # 降级：还是用 base64 但可能会超长
        b64, _ = encode_base64(grid, format="PNG")
        return {"type": "base64", "media_type": "image/png", "data": b64}


# ── 注册所有 AIGC 平台 ──

def _register_all():
    for p in _AIGC_PLATFORMS:
        p.search_impl = search
        p.detail_impl = detail
        if p.supports_suggest:
            p.suggest_impl = suggest
            p.suggest_channels = [p.id]
        register_platform(p)

    # wx 只有 suggest，没有搜索
    # suggest 调用时 channel 传 "wx"，但不注册为独立平台

_register_all()