Source code for acp_sdk.server.agent

import abc
import inspect
from collections.abc import AsyncGenerator, Coroutine, Generator
from typing import Callable

from acp_sdk.models import AgentName, Message, Metadata
from acp_sdk.server.context import Context
from acp_sdk.server.types import RunYield, RunYieldResume




class AgentManifest(abc.ABC):
    @property
    def name(self) -> AgentName:
        return self.__class__.__name__

    @property
    def description(self) -> str:
        return ""

    @property
    def metadata(self) -> Metadata:
        return Metadata()



    @abc.abstractmethod
    def run(
        self, input: list[Message], context: Context
    ) -> (
        AsyncGenerator[RunYield, RunYieldResume] | Generator[RunYield, RunYieldResume] | Coroutine[RunYield] | RunYield
    ):
        pass




Agent = AgentManifest




def agent(
    name: str | None = None,
    description: str | None = None,
    *,
    metadata: Metadata | None = None,
) -> Callable[[Callable], AgentManifest]:
    """Decorator to create an agent."""

    def decorator(fn: Callable) -> AgentManifest:
        signature = inspect.signature(fn)
        parameters = list(signature.parameters.values())

        if len(parameters) == 0:
            raise TypeError("The agent function must have at least 'input' argument")
        if len(parameters) > 2:
            raise TypeError("The agent function must have only 'input' and 'context' arguments")
        if len(parameters) == 2 and parameters[1].name != "context":
            raise TypeError("The second argument of the agent function must be 'context'")

        has_context_param = len(parameters) == 2

        class DecoratorAgentBase(AgentManifest):
            @property
            def name(self) -> str:
                return name or fn.__name__

            @property
            def description(self) -> str:
                return description or inspect.getdoc(fn) or ""

            @property
            def metadata(self) -> Metadata:
                return metadata or Metadata()

        agent: AgentManifest
        if inspect.isasyncgenfunction(fn):

            class AsyncGenDecoratorAgent(DecoratorAgentBase):
                async def run(self, input: list[Message], context: Context) -> AsyncGenerator[RunYield, RunYieldResume]:
                    try:
                        gen: AsyncGenerator[RunYield, RunYieldResume] = (
                            fn(input, context) if has_context_param else fn(input)
                        )
                        value = None
                        while True:
                            value = yield await gen.asend(value)
                    except StopAsyncIteration:
                        pass

            agent = AsyncGenDecoratorAgent()
        elif inspect.iscoroutinefunction(fn):

            class CoroDecoratorAgent(DecoratorAgentBase):
                async def run(self, input: list[Message], context: Context) -> Coroutine[RunYield]:
                    return await (fn(input, context) if has_context_param else fn(input))

            agent = CoroDecoratorAgent()
        elif inspect.isgeneratorfunction(fn):

            class GenDecoratorAgent(DecoratorAgentBase):
                def run(self, input: list[Message], context: Context) -> Generator[RunYield, RunYieldResume]:
                    yield from (fn(input, context) if has_context_param else fn(input))

            agent = GenDecoratorAgent()
        else:

            class FuncDecoratorAgent(DecoratorAgentBase):
                def run(self, input: list[Message], context: Context) -> RunYield:
                    return fn(input, context) if has_context_param else fn(input)

            agent = FuncDecoratorAgent()

        return agent

    return decorator