ai-video-fullstack/backend/services/brains/prompt_brain.py

"""Local prompt assistant, including prompt-only reusable tools."""

from __future__ import annotations

from uuid import uuid4

from models import AssistantConfig
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.frames.frames import OutputTransportMessageUrgentFrame, TTSSpeakFrame
from pipecat.processors.aggregators.llm_context import LLMContext
from pipecat.processors.frame_processor import FrameProcessor
from pipecat.services.llm_service import (
    FunctionCallParams,
    FunctionCallResultProperties,
)
from pipecat.utils.time import time_now_iso8601

from services.brains.base import BaseBrain, BrainRuntime, BrainSpec
from services.runtime_variables import (
    DynamicVariableStore,
)
from services.tool_executor import ToolExecutionError, ToolExecutor


class PromptBrain(BaseBrain):
    spec = BrainSpec(
        type="prompt",
        supported_runtime_modes=frozenset({"pipeline", "realtime"}),
        owns_context=True,
    )

    def __init__(self, cfg: AssistantConfig):
        self._cfg = cfg
        self._dynamic_enabled = True
        self._store = DynamicVariableStore.from_config(cfg)
        self._tools = ToolExecutor(self._store)
        self._runtime: BrainRuntime | None = None

    async def greeting(self, cfg: AssistantConfig) -> str:
        return self._store.render(cfg.greeting) if self._dynamic_enabled else cfg.greeting

    def system_prompt(self, cfg: AssistantConfig) -> str:
        return self._store.render(cfg.prompt) if self._dynamic_enabled else cfg.prompt

    def build_llm(self, cfg: AssistantConfig, context: LLMContext) -> FrameProcessor:
        from services.pipecat.service_factory import create_llm

        return create_llm(cfg)

    async def setup(self, cfg: AssistantConfig, runtime: BrainRuntime) -> None:
        self._runtime = runtime
        schemas: list[FunctionSchema] = []
        for tool in cfg.tools:
            if tool.type == "end_call":
                schema, handler = self._make_end_call_tool(tool, runtime)
            elif tool.type == "http":
                schema, handler = self._make_http_tool(tool, runtime)
            else:
                continue
            schemas.append(schema)
            runtime.llm.register_function(tool.function_name, handler)
        runtime.set_tools(schemas)

    def record_user_message(self, content: str) -> None:
        if not self._dynamic_enabled:
            return
        self._store.record("user", content)
        self._refresh_prompt()

    async def on_assistant_text_end(
        self,
        _turn_id: str,
        content: str,
        interrupted: bool,
    ) -> None:
        if content and not interrupted:
            self._store.record("agent", content, completed_agent_turn=True)
            self._refresh_prompt()

    def _refresh_prompt(self) -> None:
        if self._dynamic_enabled and self._runtime is not None:
            self._runtime.set_system_prompt(self._store.render(self._cfg.prompt))

    def _make_http_tool(self, tool, runtime: BrainRuntime):
        properties, required = self._tools.schema_parts(tool)
        self._tools.register_secrets(tool)

        async def call_http(params: FunctionCallParams) -> None:
            try:
                result = await self._tools.execute(tool, dict(params.arguments or {}))
                if result["updated_variables"]:
                    self._refresh_prompt()
                await params.result_callback(result)
            except (ToolExecutionError, ValueError) as exc:
                await params.result_callback(
                    {"status": "error", "message": f"HTTP 工具调用失败: {exc}"}
                )

        schema = FunctionSchema(
            name=tool.function_name,
            description=tool.description or f"调用 {tool.name}",
            properties=properties,
            required=required,
        )
        return schema, call_http

    @staticmethod
    def _make_end_call_tool(tool, runtime: BrainRuntime):
        config = (tool.definition or {}).get("config") or {}
        message_type = str(config.get("message_type") or "none")
        custom_message = str(config.get("custom_message") or "").strip()
        capture_reason = bool(config.get("capture_reason", True))

        async def end_call(params: FunctionCallParams) -> None:
            reason = str(params.arguments.get("reason") or "end_call_tool").strip()
            runtime.call_end.begin(reason)
            await params.result_callback(
                {"status": "success", "action": "ending_call"},
                properties=FunctionCallResultProperties(run_llm=False),
            )

            if message_type != "custom" or not custom_message:
                await runtime.call_end.finish()
                return

            turn_id = uuid4().hex
            timestamp = time_now_iso8601()
            for message in (
                {
                    "type": "assistant-text-start",
                    "turn_id": turn_id,
                    "timestamp": timestamp,
                },
                {
                    "type": "assistant-text-delta",
                    "turn_id": turn_id,
                    "delta": custom_message,
                },
                {
                    "type": "assistant-text-end",
                    "turn_id": turn_id,
                    "content": custom_message,
                    "interrupted": False,
                },
            ):
                await runtime.queue_frame(
                    OutputTransportMessageUrgentFrame(message=message)
                )
            runtime.call_end.arm_after_speech()
            await runtime.queue_frame(
                TTSSpeakFrame(custom_message, append_to_context=False)
            )

        properties = (
            {
                "reason": {
                    "type": "string",
                    "description": "结束本次通话的简短原因。",
                }
            }
            if capture_reason
            else {}
        )
        schema = FunctionSchema(
            name=tool.function_name,
            description=tool.description or "结束当前通话。",
            properties=properties,
            required=["reason"] if capture_reason else [],
        )
        return schema, end_call