spider_ai_seo/spiders/ai_seo/kimi.py

# coding=utf-8
import asyncio
from functools import partial, wraps

import pyperclip
from playwright.async_api import Browser

from abs_spider import AbstractAiSeoSpider
from domain.ai_seo import AiAnswer, AiSearchResult
from utils import create_logger
from glom import glom, Coalesce

from utils.ai_seo_api_utils import AiSeoApis

logger = create_logger(__name__)


class KimiSpider(AbstractAiSeoSpider):

    async def do_check_session(self) -> bool:
        self.completed_event = asyncio.Event()
        await self.browser_page.goto(self.get_home_url(), timeout=300000)
        await asyncio.sleep(3)
        user_name_element = self.browser_page.locator("//span[@class='user-name']")
        if await user_name_element.is_visible() and not await user_name_element.text_content() == '登录':
            return True
        return False

    def __init__(self, browser: Browser, prompt: str, keyword: str, think: bool = False):
        super().__init__(browser, prompt, keyword, think)
        self.__listen_response = self.handle_listen_response_error(self.__listen_response)

    def get_home_url(self) -> str:
        return 'https://www.kimi.ai'

    def get_platform_id(self) -> int:
        return 4

    def get_platform_name(self) -> str:
        return 'Kimi'

    async def _do_spider(self) -> AiAnswer:
        self.completed_event = asyncio.Event()
        await self.browser_page.goto('https://www.kimi.ai', timeout=600000)
        self.ai_answer = AiAnswer(self.get_platform_id(), self.get_platform_name(), self.prompt, self.keyword)
        await asyncio.sleep(3)
        confirm_btn = self.browser_page.locator('//button[text()="知道了"]')
        if await confirm_btn.is_visible():
            await confirm_btn.click()
        # if self.think:
        #     think_btn = self.browser_page.locator("span:text('长思考 (k1.5)')").locator('..')
        #     if await think_btn.is_visible():
        #         clazz = (await think_btn.get_attribute('class')).split(' ')
        #         if 'open' not in clazz:
        #             await think_btn.click()
        #             await asyncio.sleep(2)
        chat_input_element = self.browser_page.locator("//div[@class='chat-input']")
        await chat_input_element.click()
        # 输入提问词
        await self.browser_page.keyboard.type(self.prompt)
        await asyncio.sleep(2)
        await self.browser_page.keyboard.press('Enter')
        # 监听请求
        # self.browser_page.on('response', partial(self.__listen_response))
        # await self.completed_event.wait()
        await asyncio.sleep(2)

        # 等待复制按钮可见
        copy_btn_xpath = "//div[@class='segment-assistant-actions-content']/div[@class='simple-button size-small'][1]"
        await self.browser_page.wait_for_selector(copy_btn_xpath, timeout=600000)
        copy_btn = self.browser_page.locator(copy_btn_xpath)
        await copy_btn.click()
        # 读取剪贴板
        self.ai_answer.answer = pyperclip.paste()
        logger.debug(f"ai回复: {self.ai_answer.answer}")

        # 报错检查
        if self.fail_status:
            await AiSeoApis.update_spider_session(self.session_info['id'], 2)
            raise self.fail_exception

        # 关闭侧边栏
        sidebar_element = self.browser_page.locator("//div[@class='expand-btn']")
        if await sidebar_element.is_visible():
            await sidebar_element.click()
        # 获取回答元素
        answer_element = self.browser_page.locator("//div[@class='segment-container']").nth(-1)
        box = await answer_element.bounding_box()
        logger.debug(f'answer_element: {box}')
        view_port_height = box['height'] + 500
        # 调整视口大小
        await self.browser_page.set_viewport_size({
            'width': 1920,
            'height': int(view_port_height)
        })
        # 打开搜索结果
        search_list_content_element = self.browser_page.locator("//div[contains(@class, 'side-console-container')]")
        search_list_element = self.browser_page.locator("//div[@class='search-plus']")
        if await search_list_element.is_visible() and not await search_list_content_element.is_visible():
            await search_list_element.click()

        # 获取搜索结果
        search_list = []
        search_elements = await self.browser_page.locator("//div[@class='sites']/a[@class='site']").all()
        for search_element in search_elements:
            result = AiSearchResult()
            result.url = await search_element.get_attribute('href')

            children = await search_element.locator("xpath=./child::*").all()
            result.title = await children[1].inner_text()
            result.body = await children[2].inner_text()

            # 获取信源信息元素
            result_host_elements = await children[0].locator("xpath=./child::*").all()
            try:
                result.host_name = await result_host_elements[1].inner_text()
            except Exception:
                result.host_name = ''
            if len(result_host_elements) >= 3:
                try:
                    result.publish_time = await result_host_elements[2].inner_text()
                    result.publish_time = result.publish_time.replace('/', '-')
                except Exception:
                    result.publish_time = 0
                    logger.error(f"{result.title}获取发布时间失败")
            search_list.append(result)
        self.ai_answer.search_result = search_list


        # 截图
        screenshot_path = self._get_screenshot_path()
        self.ai_answer.screenshot_file = screenshot_path
        await self.browser_page.screenshot(path=screenshot_path)
        return self.ai_answer

    async def __listen_response(self, response):
        if '/segment/scroll' in response.url:
            json_data = await response.json()
            if json_data['items']:
                logger.debug(json_data)
                detail = json_data['items'][-1]
                if 'error' in detail:
                    logger.error(f"kimi回复错误: {detail['error']['detail']}")
                    self.fail_status = True
                    self.fail_exception = Exception(detail['error']['detail'])
                    self.completed_event.set()
                    return
                content = detail['content']
                if self.think:
                    self.ai_answer.search_result = self.get_search_list_enable_think(detail)
                else:
                    self.ai_answer.search_result = self.get_search_list_disable_think(detail)
                self.ai_answer.answer = content
                logger.debug(f"ai回复: {content}")
                self.completed_event.set()

    def handle_listen_response_error(self, func):
        """
        装饰器 用于处理请求回调中的异常
        :param func:
        :return:
        """

        @wraps(func)
        async def wrapper(*args, **kwargs):
            try:
                return await func(*args, **kwargs)
            except Exception as e:
                logger.error(f"{self.get_platform_name()}响应异常: {e}", exc_info=True)
                # 标记失败状态 记录异常
                self.fail_status = True
                self.fail_exception = e
                self.completed_event.set()

        return wrapper

    def get_search_list_disable_think(self, detail):
        """
        未开启深度思考时 获取搜索结果
        :param detail:
        :return:
        """
        answer_search_list = []
        search_result_list = detail.get('search_plus', [])
        for search_result in search_result_list:
            event = search_result.get('event', '')
            msg = search_result.get('msg', {})
            msg_type = msg.get('type', '')
            if event == 'search_plus' and msg_type == 'get_res':
                answer_search_list.append(
                    AiSearchResult(msg['title'], msg['url'], msg['site_name'], msg['snippet'], msg['date']))
                logger.debug(f"ai参考资料: {msg['title']}({msg['url']})")
        return answer_search_list

    def get_search_list_enable_think(self, detail):
        """
        开启深度思考时 获取搜索结果
        :param detail:
        :return:
        """
        answer_search_list = []
        keys = 'contents.zones.0.sections.0.k1.search_results'
        search_result_list = glom(detail, keys, default=[])
        for search_result in search_result_list:
                answer_search_list.append(
                    AiSearchResult(search_result['title'], search_result['url'], search_result['site_name'], search_result['snippet'], search_result['date']))
                logger.debug(f"ai参考资料: {search_result['title']}({search_result['url']})")
        return answer_search_list