AgentCoreService/pqai_agent/toolkit/pq_video_searcher.py

from typing import Dict, List
import requests

from pqai_agent.toolkit.base import BaseToolkit
from pqai_agent.toolkit.function_tool import FunctionTool

class PQVideoSearcher(BaseToolkit):
    API_URL = "https://vlogapi.piaoquantv.com/longvideoapi/search/userandvideo/list"
    def search_pq_video(self, keywords: List[str]) -> List[Dict]:
        """
        Search for videos on the PQ Video (票圈视频) platform using keywords.
        使用关键词在票圈视频平台上搜索视频。使用时请按要求提供关键词列表。
        Args:
            keywords (List[str]): 关键词列表，至多5个词，且必须按相关性从高到低排列，每个词至多4个字。
        Returns:
            List[Dict]: List of video metadata dictionaries.
        """
        if not keywords:
            return []
        keyword = keywords[0]
        data = {
            'keyWord': keyword,
            'pageSize': "10",
            'pageNo': "1"
        }
        raw_data = '&'.join([f"{k}={v}" for k, v in data.items()])
        print(data)

        try:
            response = requests.post(self.API_URL, data=data)
            response.raise_for_status()  # Raise an error for bad responses
            result = response.json()
            data = result.get("data", {})
            videos = data.get("videos", {}).get("videos", [])
            ret = []
            for item in videos:
                if item.get("auditStatus") == 5 and item.get("transcodeStatus") == 3 and item.get("title"):
                    ret.append({
                        "id": item["id"],
                        "coverImgPath": item["coverImg"]["coverImgPath"],
                        "title": item["title"],
                        "totalTime": item["totalTime"],
                        "videoPath": item["videoPath"],
                        "playCountFormatStr": item["playCountFormatStr"]
                    })
            ret = ret[:4]
            return ret
        except Exception as e:
            return []

    def get_tools(self) -> List[FunctionTool]:
        return [FunctionTool(self.search_pq_video)]