LQConstPlanWriterAgent/core/document_chat/skills/base.py

# -*- coding: utf-8 -*-
"""文档对话技能基类。"""

from abc import ABC, abstractmethod
from typing import Callable

from core.document_chat.schemas import DocumentChatSkillInput, DocumentChatSkillOutput


class BaseDocumentChatSkill(ABC):
    def __init__(self, name: str, function_name: str):
        self.name = name
        self.function_name = function_name

    @abstractmethod
    async def run(self, skill_input: DocumentChatSkillInput) -> DocumentChatSkillOutput:
        """执行技能并返回标准化输出。"""
        raise NotImplementedError

    async def run_stream(
        self,
        skill_input: DocumentChatSkillInput,
        on_chunk: Callable[[str], None],
    ) -> DocumentChatSkillOutput:
        """流式执行。每次生成一个 chunk 时调用 on_chunk，最终返回完整结果。

        默认实现调用非流式 run()，将整个 answer 一次性传给 on_chunk，
        子类可覆盖此方法实现真正的流式生成。
        """
        result = await self.run(skill_input)
        text = result.answer or result.proposed_content or ""
        if text:
            on_chunk(text)
        return result