builtin.py

"""Google ADK 内置集成函数 / Google ADK Built-in Integration Functions

提供快速创建 Google ADK 兼容模型和工具的便捷函数。
Provides convenient functions for quickly creating Google ADK-compatible models and tools.
"""

from typing import Any, Callable, List, Optional, Union

from typing_extensions import Unpack

from agentrun.integration.builtin import (
    knowledgebase_toolset as _knowledgebase_toolset,
)
from agentrun.integration.builtin import model as _model
from agentrun.integration.builtin import ModelArgs
from agentrun.integration.builtin import sandbox_toolset as _sandbox_toolset
from agentrun.integration.builtin import skill_tools as _skill_tools
from agentrun.integration.builtin import tool_resource as _tool_resource
from agentrun.integration.builtin import toolset as _toolset
from agentrun.integration.utils.tool import Tool
from agentrun.model import ModelProxy, ModelService
from agentrun.sandbox import TemplateType
from agentrun.tool.tool import Tool as ToolResourceType
from agentrun.toolset import ToolSet
from agentrun.utils.config import Config


def model(
    name: Union[str, ModelProxy, ModelService],
    **kwargs: Unpack[ModelArgs],
):
    """获取 AgentRun 模型并转换为 LangChain ``BaseChatModel``。 / Google ADK Built-in Integration Functions"""

    m = _model(input=name, **kwargs)  # type: ignore
    return m.to_google_adk()


def toolset(
    name: Union[str, ToolSet],
    *,
    prefix: Optional[str] = None,
    modify_tool_name: Optional[Callable[[Tool], Tool]] = None,
    filter_tools_by_name: Optional[Callable[[str], bool]] = None,
    config: Optional[Config] = None,
) -> List[Any]:
    """将内置工具集封装为 LangChain ``StructuredTool`` 列表。 / Google ADK Built-in Integration Functions"""

    ts = _toolset(input=name, config=config)
    return ts.to_google_adk(
        prefix=prefix,
        modify_tool_name=modify_tool_name,
        filter_tools_by_name=filter_tools_by_name,
    )


def tool_resource(
    name: Union[str, ToolResourceType],
    *,
    prefix: Optional[str] = None,
    modify_tool_name: Optional[Callable[[Tool], Tool]] = None,
    filter_tools_by_name: Optional[Callable[[str], bool]] = None,
    config: Optional[Config] = None,
) -> List[Any]:
    """将 ToolResource 封装为 Google ADK 工具列表。 / Google ADK Built-in ToolResource Integration"""

    ts = _tool_resource(input=name, config=config)
    return ts.to_google_adk(
        prefix=prefix,
        modify_tool_name=modify_tool_name,
        filter_tools_by_name=filter_tools_by_name,
    )


def sandbox_toolset(
    template_name: str,
    *,
    template_type: TemplateType = TemplateType.CODE_INTERPRETER,
    config: Optional[Config] = None,
    sandbox_idle_timeout_seconds: int = 600,
    prefix: Optional[str] = None,
) -> List[Any]:
    """将沙箱模板封装为 LangChain ``StructuredTool`` 列表。 / Google ADK Built-in Integration Functions"""

    return _sandbox_toolset(
        template_name=template_name,
        template_type=template_type,
        config=config,
        sandbox_idle_timeout_seconds=sandbox_idle_timeout_seconds,
    ).to_google_adk(prefix=prefix)


def knowledgebase_toolset(
    knowledge_base_names: List[str],
    *,
    prefix: Optional[str] = None,
    modify_tool_name: Optional[Callable[[Tool], Tool]] = None,
    filter_tools_by_name: Optional[Callable[[str], bool]] = None,
    config: Optional[Config] = None,
) -> List[Any]:
    """将知识库检索封装为 Google ADK 工具列表。 / Google ADK Built-in Integration Functions"""

    return _knowledgebase_toolset(
        knowledge_base_names=knowledge_base_names,
        config=config,
    ).to_google_adk(
        prefix=prefix,
        modify_tool_name=modify_tool_name,
        filter_tools_by_name=filter_tools_by_name,
    )


def skill_tools(
    name: Optional[Union[str, List[str]]] = None,
    *,
    skills_dir: str = ".skills",
    prefix: Optional[str] = None,
    config: Optional[Config] = None,
) -> List[Any]:
    """将 Skill 封装为 Google ADK 工具列表。 / Google ADK Built-in Skill Integration"""

    ts = _skill_tools(name=name, skills_dir=skills_dir, config=config)
    return ts.to_google_adk(
        prefix=prefix,
    )