acp-sdk 0.0.6__py3-none-any.whl → 1.0.0rc1__py3-none-any.whl

acp_sdk/client/__init__.py

@@ -0,0 +1 @@
+from acp_sdk.client.client import Client

acp_sdk/client/client.py

@@ -0,0 +1,135 @@
+from types import TracebackType
+from typing import AsyncIterator
+
+import httpx
+from httpx_sse import EventSource, aconnect_sse
+
+from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
+
+from acp_sdk.models import (
+    Agent,
+    AgentName,
+    AgentReadResponse,
+    AgentsListResponse,
+    AwaitResume,
+    Message,
+    RunCancelResponse,
+    RunCreateRequest,
+    Run,
+    RunCreateResponse,
+    RunEvent,
+    RunId,
+    RunMode,
+    RunResumeRequest,
+    RunResumeResponse,
+)
+from pydantic import TypeAdapter
+
+
+class Client:
+    def __init__(
+        self, *, base_url: httpx.URL | str = "", client: httpx.AsyncClient | None = None
+    ):
+        self.base_url = base_url
+
+        self._client = self._init_client(client)
+
+    def _init_client(self, client: httpx.AsyncClient | None = None):
+        client = client or httpx.AsyncClient(base_url=self.base_url)
+        HTTPXClientInstrumentor.instrument_client(client)
+        return client
+
+    async def __aenter__(self):
+        await self._client.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None = None,
+        exc_value: BaseException | None = None,
+        traceback: TracebackType | None = None,
+    ):
+        await self._client.__aexit__(exc_type, exc_value, traceback)
+
+    async def agents(self) -> AsyncIterator[Agent]:
+        response = await self._client.get("/agents")
+        for agent in AgentsListResponse.model_validate(response.json()).agents:
+            yield agent
+
+    async def agent(self, *, name: AgentName) -> Agent:
+        response = await self._client.get(f"/agents/{name}")
+        return AgentReadResponse.model_validate(response.json())
+
+    async def run_sync(self, *, agent: AgentName, input: Message) -> Run:
+        response = await self._client.post(
+            "/runs",
+            json=RunCreateRequest(
+                agent_name=agent, input=input, mode=RunMode.SYNC
+            ).model_dump(),
+        )
+        return RunCreateResponse.model_validate(response.json())
+
+    async def run_async(self, *, agent: AgentName, input: Message) -> Run:
+        response = await self._client.post(
+            "/runs",
+            json=RunCreateRequest(
+                agent_name=agent, input=input, mode=RunMode.ASYNC
+            ).model_dump(),
+        )
+        return RunCreateResponse.model_validate(response.json())
+
+    async def run_stream(
+        self, *, agent: AgentName, input: Message
+    ) -> AsyncIterator[RunEvent]:
+        async with aconnect_sse(
+            self._client,
+            "POST",
+            "/runs",
+            json=RunCreateRequest(
+                agent_name=agent, input=input, mode=RunMode.STREAM
+            ).model_dump(),
+        ) as event_source:
+            async for event in self._validate_stream(event_source):
+                yield event
+
+    async def run_status(self, *, run_id: RunId) -> Run:
+        response = await self._client.get(f"/runs/{run_id}")
+        return Run.model_validate(response.json())
+
+    async def run_cancel(self, *, run_id: RunId) -> Run:
+        response = await self._client.post(f"/runs/{run_id}/cancel")
+        return RunCancelResponse.model_validate(response.json())
+
+    async def run_resume_sync(self, *, run_id: RunId, await_: AwaitResume) -> Run:
+        response = await self._client.post(
+            f"/runs/{run_id}",
+            json=RunResumeRequest(await_=await_, mode=RunMode.SYNC).model_dump(),
+        )
+        return RunResumeResponse.model_validate(response.json())
+
+    async def run_resume_async(self, *, run_id: RunId, await_: AwaitResume) -> Run:
+        response = await self._client.post(
+            f"/runs/{run_id}",
+            json=RunResumeRequest(await_=await_, mode=RunMode.ASYNC).model_dump(),
+        )
+        return RunResumeResponse.model_validate(response.json())
+
+    async def run_resume_stream(
+        self, *, run_id: RunId, await_: AwaitResume
+    ) -> AsyncIterator[RunEvent]:
+        async with aconnect_sse(
+            self._client,
+            "POST",
+            f"/runs/{run_id}",
+            json=RunResumeRequest(await_=await_, mode=RunMode.STREAM).model_dump(),
+        ) as event_source:
+            async for event in self._validate_stream(event_source):
+                yield event
+
+    async def _validate_stream(
+        self,
+        event_source: EventSource,
+    ) -> AsyncIterator[RunEvent]:
+        async for event in event_source.aiter_sse():
+            event = TypeAdapter(RunEvent).validate_json(event.data)
+            yield event

acp_sdk/models.py

@@ -0,0 +1,219 @@
+from enum import Enum
+from typing import Annotated, Literal, Union
+import uuid
+
+from pydantic import AnyUrl, BaseModel, ConfigDict, Field, RootModel
+
+
+class ACPError(BaseModel):
+    code: str
+    message: str
+
+
+class AnyModel(BaseModel):
+    model_config = ConfigDict(extra="allow")
+
+
+class MessagePartBase(BaseModel):
+    type: Literal["text", "image", "artifact"]
+
+
+class TextMessagePart(MessagePartBase):
+    type: Literal["text"] = "text"
+    content: str
+
+
+class ImageMessagePart(MessagePartBase):
+    type: Literal["image"] = "image"
+    content_url: AnyUrl
+
+
+class ArtifactMessagePart(MessagePartBase):
+    type: Literal["artifact"] = "artifact"
+    name: str
+    content_url: AnyUrl
+
+
+MessagePart = Union[TextMessagePart, ImageMessagePart, ArtifactMessagePart]
+
+
+class Message(RootModel):
+    root: list[MessagePart]
+
+    def __init__(self, *items: MessagePart):
+        super().__init__(root=list(items))
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+    def __add__(self, other: "Message") -> "Message":
+        if not isinstance(other, Message):
+            raise TypeError(f"Cannot concatenate Message with {type(other).__name__}")
+        return Message(*(self.root + other.root))
+
+    def __str__(self):
+        return "".join(
+            str(part) for part in self.root if isinstance(part, TextMessagePart)
+        )
+
+
+AgentName = str
+SessionId = str
+RunId = str
+
+
+class RunMode(str, Enum):
+    SYNC = "sync"
+    ASYNC = "async"
+    STREAM = "stream"
+
+
+class RunStatus(str, Enum):
+    CREATED = "created"
+    IN_PROGRESS = "in-progress"
+    AWAITING = "awaiting"
+    CANCELLING = "cancelling"
+    CANCELLED = "cancelled"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+    @property
+    def is_terminal(self) -> bool:
+        terminal_states = {RunStatus.COMPLETED, RunStatus.FAILED, RunStatus.CANCELLED}
+        return self in terminal_states
+
+
+class Await(BaseModel):
+    type: Literal["placeholder"] = "placeholder"
+
+
+class AwaitResume(BaseModel):
+    pass
+
+
+class Run(BaseModel):
+    run_id: RunId = str(uuid.uuid4())
+    agent_name: AgentName
+    session_id: SessionId | None = None
+    status: RunStatus = RunStatus.CREATED
+    await_: Await | None = Field(None, alias="await")
+    output: Message | None = None
+    error: ACPError | None = None
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    def model_dump_json(
+        self,
+        **kwargs,
+    ):
+        return super().model_dump_json(
+            by_alias=True,
+            **kwargs,
+        )
+
+
+class MessageEvent(BaseModel):
+    type: Literal["message"] = "message"
+    message: Message
+
+
+class AwaitEvent(BaseModel):
+    type: Literal["await"] = "await"
+    await_: Await | None = Field(alias="await")
+
+    model_config = ConfigDict(populate_by_name=True)
+
+    def model_dump_json(
+        self,
+        **kwargs,
+    ):
+        return super().model_dump_json(
+            by_alias=True,
+            **kwargs,
+        )
+
+
+class GenericEvent(BaseModel):
+    type: Literal["generic"] = "generic"
+    generic: AnyModel
+
+
+class CreatedEvent(BaseModel):
+    type: Literal["created"] = "created"
+    run: Run
+
+
+class InProgressEvent(BaseModel):
+    type: Literal["in-progress"] = "in-progress"
+    run: Run
+
+
+class FailedEvent(BaseModel):
+    type: Literal["failed"] = "failed"
+    run: Run
+
+
+class CancelledEvent(BaseModel):
+    type: Literal["cancelled"] = "cancelled"
+    run: Run
+
+
+class CompletedEvent(BaseModel):
+    type: Literal["completed"] = "completed"
+    run: Run
+
+
+RunEvent = Union[
+    CreatedEvent,
+    InProgressEvent,
+    MessageEvent,
+    AwaitEvent,
+    GenericEvent,
+    CancelledEvent,
+    FailedEvent,
+    CompletedEvent,
+]
+
+
+class RunCreateRequest(BaseModel):
+    agent_name: AgentName
+    session_id: SessionId | None = None
+    input: Message
+    mode: RunMode = RunMode.SYNC
+
+
+class RunCreateResponse(Run):
+    pass
+
+
+class RunResumeRequest(BaseModel):
+    await_: AwaitResume = Field(alias="await")
+    mode: RunMode
+
+
+class RunResumeResponse(Run):
+    pass
+
+
+class RunReadResponse(Run):
+    pass
+
+
+class RunCancelResponse(Run):
+    pass
+
+
+class Agent(BaseModel):
+    name: str
+    description: str | None = None
+
+
+class AgentsListResponse(BaseModel):
+    agents: list[Agent]
+
+
+class AgentReadResponse(Agent):
+    pass

acp_sdk/server/__init__.py

@@ -0,0 +1,2 @@
+from acp_sdk.server.server import create_app
+from acp_sdk.server.agent import Agent

acp_sdk/server/agent.py

@@ -0,0 +1,32 @@
+import abc
+from typing import AsyncGenerator
+
+from acp_sdk.models import (
+    AgentName,
+    Message,
+    Await,
+    AwaitResume,
+    SessionId,
+)
+from acp_sdk.server.context import Context
+
+
+class Agent(abc.ABC):
+    @property
+    def name(self) -> AgentName:
+        return self.__class__.__name__
+
+    @property
+    def description(self) -> str:
+        return ""
+
+    @abc.abstractmethod
+    def run(
+        self, input: Message, *, context: Context
+    ) -> AsyncGenerator[Message | Await, AwaitResume]:
+        pass
+
+    async def session(self, session_id: SessionId | None) -> SessionId | None:
+        if session_id:
+            raise NotImplementedError()
+        return None

acp_sdk/server/bundle.py

@@ -0,0 +1,133 @@
+import asyncio
+import logging
+
+from opentelemetry import trace
+from pydantic import ValidationError
+
+from acp_sdk.server.agent import Agent
+from acp_sdk.models import (
+    ACPError,
+    AnyModel,
+    Await,
+    AwaitEvent,
+    CancelledEvent,
+    CompletedEvent,
+    CreatedEvent,
+    FailedEvent,
+    GenericEvent,
+    InProgressEvent,
+    Message,
+    MessageEvent,
+    Run,
+    AwaitResume,
+    RunEvent,
+    RunStatus,
+)
+from acp_sdk.server.context import Context
+
+logger = logging.getLogger("uvicorn.error")
+
+
+class RunBundle:
+    def __init__(self, *, agent: Agent, run: Run, task: asyncio.Task | None = None):
+        self.agent = agent
+        self.run = run
+        self.task = task
+
+        self.stream_queue: asyncio.Queue[RunEvent] = asyncio.Queue()
+        self.composed_message = Message()
+
+        self.await_queue: asyncio.Queue[AwaitResume] = asyncio.Queue(maxsize=1)
+        self.await_or_terminate_event = asyncio.Event()
+
+    async def stream(self):
+        try:
+            while True:
+                event = await self.stream_queue.get()
+                yield event
+                self.stream_queue.task_done()
+        except asyncio.QueueShutDown:
+            pass
+
+    async def emit(self, event: RunEvent):
+        await self.stream_queue.put(event)
+
+    async def await_(self) -> AwaitResume:
+        self.stream_queue.shutdown()
+        self.await_queue.empty()
+        self.await_or_terminate_event.set()
+        self.await_or_terminate_event.clear()
+        resume = await self.await_queue.get()
+        self.await_queue.task_done()
+        return resume
+
+    async def resume(self, resume: AwaitResume):
+        self.stream_queue = asyncio.Queue()
+        await self.await_queue.put(resume)
+
+    async def join(self):
+        await self.await_or_terminate_event.wait()
+
+    async def execute(self, input: Message):
+        with trace.get_tracer(__name__).start_as_current_span("execute"):
+            run_logger = logging.LoggerAdapter(logger, {"run_id": self.run.run_id})
+
+            await self.emit(CreatedEvent(run=self.run))
+            try:
+                self.run.session_id = await self.agent.session(self.run.session_id)
+                run_logger.info("Session loaded")
+
+                generator = self.agent.run(
+                    input=input, context=Context(session_id=self.run.session_id)
+                )
+                run_logger.info("Run started")
+
+                self.run.status = RunStatus.IN_PROGRESS
+                await self.emit(InProgressEvent(run=self.run))
+
+                await_resume = None
+                while True:
+                    next = await generator.asend(await_resume)
+                    if isinstance(next, Message):
+                        self.composed_message += next
+                        await self.emit(MessageEvent(message=next))
+                    elif isinstance(next, Await):
+                        self.run.await_ = next
+                        self.run.status = RunStatus.AWAITING
+                        await self.emit(
+                            AwaitEvent.model_validate(
+                                {
+                                    "run_id": self.run.run_id,
+                                    "type": "await",
+                                    "await": next,
+                                }
+                            )
+                        )
+                        run_logger.info("Run awaited")
+                        await_resume = await self.await_()
+                        self.run.status = RunStatus.IN_PROGRESS
+                        await self.emit(InProgressEvent(run=self.run))
+                        run_logger.info("Run resumed")
+                    else:
+                        try:
+                            generic = AnyModel.model_validate(next)
+                            await self.emit(GenericEvent(generic=generic))
+                        except ValidationError:
+                            raise TypeError("Invalid yield")
+            except StopAsyncIteration:
+                self.run.output = self.composed_message
+                self.run.status = RunStatus.COMPLETED
+                await self.emit(CompletedEvent(run=self.run))
+                run_logger.info("Run completed")
+            except asyncio.CancelledError:
+                self.run.status = RunStatus.CANCELLED
+                await self.emit(CancelledEvent(run=self.run))
+                run_logger.info("Run cancelled")
+            except Exception as e:
+                self.run.error = ACPError(code="unspecified", message=str(e))
+                self.run.status = RunStatus.FAILED
+                await self.emit(FailedEvent(run=self.run))
+                run_logger.exception("Run failed")
+            finally:
+                self.await_or_terminate_event.set()
+                self.stream_queue.shutdown()

acp_sdk/server/context.py

@@ -0,0 +1,6 @@
+from acp_sdk.models import SessionId
+
+
+class Context:
+    def __init__(self, *, session_id: SessionId | None = None):
+        self.session_id = session_id

acp_sdk/server/server.py

@@ -0,0 +1,137 @@
+import asyncio
+
+from acp_sdk.server.telemetry import configure_telemetry
+from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
+
+from fastapi import FastAPI, HTTPException, status
+from fastapi.responses import JSONResponse, StreamingResponse
+
+from acp_sdk.server.agent import Agent
+from acp_sdk.models import (
+    AgentName,
+    Agent as AgentModel,
+    AgentsListResponse,
+    AgentReadResponse,
+    Run,
+    RunCancelResponse,
+    RunCreateRequest,
+    RunCreateResponse,
+    RunId,
+    RunMode,
+    RunReadResponse,
+    RunResumeRequest,
+    RunResumeResponse,
+    RunStatus,
+)
+from acp_sdk.server.bundle import RunBundle
+from acp_sdk.server.utils import stream_sse
+
+
+def create_app(*agents: Agent) -> FastAPI:
+    app = FastAPI(title="acp-agents")
+
+    configure_telemetry()
+    FastAPIInstrumentor.instrument_app(app)
+
+    agents: dict[AgentName, Agent] = {agent.name: agent for agent in agents}
+    runs: dict[RunId, RunBundle] = dict()
+
+    def find_run_bundle(run_id: RunId):
+        bundle = runs.get(run_id, None)
+        if not bundle:
+            raise HTTPException(status_code=404, detail=f"Run {run_id} not found")
+        return bundle
+
+    def find_agent(agent_name: AgentName):
+        agent = agents.get(agent_name, None)
+        if not agent:
+            raise HTTPException(status_code=404, detail=f"Agent {agent_name} not found")
+        return agent
+
+    @app.get("/agents")
+    async def list() -> AgentsListResponse:
+        return AgentsListResponse(
+            agents=[
+                AgentModel(name=agent.name, description=agent.description)
+                for agent in agents.values()
+            ]
+        )
+
+    @app.get("/agents/{name}")
+    async def read(name: AgentName) -> AgentReadResponse:
+        agent = find_agent(name)
+        return AgentModel(name=agent.name, description=agent.description)
+
+    @app.post("/runs")
+    async def create(request: RunCreateRequest) -> RunCreateResponse:
+        agent = find_agent(request.agent_name)
+        bundle = RunBundle(
+            agent=agent,
+            run=Run(
+                agent_name=agent.name,
+                session_id=request.session_id,
+            ),
+        )
+
+        bundle.task = asyncio.create_task(bundle.execute(request.input))
+        runs[bundle.run.run_id] = bundle
+
+        match request.mode:
+            case RunMode.STREAM:
+                return StreamingResponse(
+                    stream_sse(bundle),
+                    media_type="text/event-stream",
+                )
+            case RunMode.SYNC:
+                await bundle.join()
+                return bundle.run
+            case RunMode.ASYNC:
+                return JSONResponse(
+                    status_code=status.HTTP_202_ACCEPTED,
+                    content=bundle.run.model_dump(),
+                )
+            case _:
+                raise NotImplementedError()
+
+    @app.get("/runs/{run_id}")
+    async def read(run_id: RunId) -> RunReadResponse:
+        bundle = find_run_bundle(run_id)
+        return bundle.run
+
+    @app.post("/runs/{run_id}")
+    async def resume(run_id: RunId, request: RunResumeRequest) -> RunResumeResponse:
+        bundle = find_run_bundle(run_id)
+        bundle.stream_queue = asyncio.Queue()  # TODO improve
+        await bundle.await_queue.put(request.await_)
+        match request.mode:
+            case RunMode.STREAM:
+                return StreamingResponse(
+                    stream_sse(bundle),
+                    media_type="text/event-stream",
+                )
+            case RunMode.SYNC:
+                await bundle.join()
+                return bundle.run
+            case RunMode.ASYNC:
+                return JSONResponse(
+                    status_code=status.HTTP_202_ACCEPTED,
+                    content=bundle.run.model_dump(),
+                )
+            case _:
+                raise NotImplementedError()
+
+    @app.post("/runs/{run_id}/cancel")
+    async def cancel(run_id: RunId) -> RunCancelResponse:
+        bundle = find_run_bundle(run_id)
+        if bundle.run.status.is_terminal:
+            raise HTTPException(
+                status_code=403,
+                detail=f"Run with terminal status {bundle.run.status} can't be cancelled",
+            )
+        bundle.task.cancel()
+        bundle.run.status = RunStatus.CANCELLING
+        return JSONResponse(
+            status_code=status.HTTP_202_ACCEPTED, content=bundle.run.model_dump()
+        )
+
+    return app