composiohq__composio/python/composio/core/provider/base.py

"""
BaseProvider module

Defines the barebones provider metaclass that needs to be subclassed for every provider.
"""

from __future__ import annotations

import typing as t

import typing_extensions as te

if t.TYPE_CHECKING:
    from composio.core.models.tools import Modifiers, ToolExecutionResponse

TTool = t.TypeVar("TTool")
TToolCollection = t.TypeVar("TToolCollection")


class ExecuteToolFn(t.Protocol):
    def __call__(
        self,
        slug: str,
        arguments: t.Dict,
        *,
        modifiers: t.Optional[Modifiers] = None,
        user_id: t.Optional[str] = None,
    ) -> ToolExecutionResponse:
        """
        Execute a wrapped tool by slug, passing an arbitrary input dict.
        This function is used by the providers to execute tools for the helper methods.
        Returns a dict with the following keys:
            - data: The data returned by the tool.
            - error: The error returned by the tool.
            - successful: Whether the tool was successful.
        """
        ...


class ToolCallSession(t.Protocol):
    """Execution contract implemented by ToolRouterSession."""

    def execute(
        self,
        tool_slug: str,
        *,
        arguments: t.Optional[t.Dict[str, t.Any]] = None,
    ) -> t.Any: ...


ToolCallExecutionTarget: t.TypeAlias = t.Union[str, ToolCallSession]


class SchemaConfig(te.TypedDict):
    skip_defaults: te.NotRequired[bool]


class BaseProviderConfig(te.TypedDict):
    schema_config: te.NotRequired[SchemaConfig]


class BaseProvider(t.Generic[TTool, TToolCollection]):
    """
    BaseProvider class

    All providers should inherit from this class and implement `wrap_tools` so that
    they can be used with the core Composio class.
    """

    name: str
    """Name of the provider"""

    __schema_skip_defaults__ = False

    execute_tool: ExecuteToolFn
    """
    The function to execute a tool for the provider's helper methods.
    This is automatically injected by the core SDK.
    """

    def __init__(self, **kwargs: t.Unpack[BaseProviderConfig]) -> None:
        self.skip_default = kwargs.get("schema_config", {}).get(
            "skip_defaults", self.__schema_skip_defaults__
        )

    def set_execute_tool_fn(self, execute_tool_fn: ExecuteToolFn) -> None:
        self.execute_tool = execute_tool_fn

    def resolve_tool_call_execution_target(
        self,
        *,
        user_id: t.Optional[str],
        session: t.Optional[ToolCallSession],
    ) -> ToolCallExecutionTarget:
        """Resolve exactly one direct user or Tool Router session target."""
        if (user_id is None) == (session is None):
            raise ValueError("Provide exactly one of user_id or session")
        if session is not None:
            return session
        return t.cast(str, user_id)

    def execute_tool_for_target(
        self,
        *,
        target: ToolCallExecutionTarget,
        slug: str,
        arguments: t.Dict[str, t.Any],
        modifiers: t.Optional[Modifiers] = None,
    ) -> ToolExecutionResponse:
        """Execute normalized arguments through the matching SDK boundary."""
        if isinstance(target, str):
            return self.execute_tool(
                slug=slug,
                arguments=arguments,
                modifiers=modifiers,
                user_id=target,
            )

        if modifiers is not None:
            raise ValueError(
                "Direct execution modifiers cannot be used with a Tool Router session"
            )

        result = target.execute(tool_slug=slug, arguments=arguments)
        return {
            "data": t.cast(t.Dict, result.data),
            "error": result.error,
            "successful": result.error is None,
        }