acp-sdk 0.10.0__py3-none-any.whl → 0.11.0__py3-none-any.whl

acp_sdk/server/app.py

@@ -5,19 +5,24 @@ from contextlib import asynccontextmanager
 from datetime import timedelta
 from enum import Enum
 
-from fastapi import Depends, FastAPI, HTTPException, status
+import httpx
+import obstore.store
+from fastapi import Depends, FastAPI, HTTPException, Request, status
 from fastapi.applications import AppType, Lifespan
 from fastapi.encoders import jsonable_encoder
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
+from obstore.exceptions import NotFoundError
 
 from acp_sdk.models import (
-    Agent as AgentModel,
-)
-from acp_sdk.models import (
+    ACPError,
     AgentName,
     AgentReadResponse,
     AgentsListResponse,
+    AwaitResume,
+    PingResponse,
+    ResourceId,
+    ResourceUrl,
     Run,
     RunCancelResponse,
     RunCreateRequest,
@@ -28,11 +33,15 @@ from acp_sdk.models import (
     RunReadResponse,
     RunResumeRequest,
     RunResumeResponse,
+    RunStatus,
+    Session,
+    SessionId,
+    SessionReadResponse,
+)
+from acp_sdk.models import (
+    AgentManifest as AgentModel,
 )
-from acp_sdk.models.errors import ACPError
-from acp_sdk.models.models import AwaitResume, RunStatus
-from acp_sdk.models.schemas import PingResponse
-from acp_sdk.server.agent import Agent
+from acp_sdk.server.agent import AgentManifest
 from acp_sdk.server.errors import (
     RequestValidationError,
     StarletteHTTPException,
@@ -42,9 +51,9 @@ from acp_sdk.server.errors import (
     validation_exception_handler,
 )
 from acp_sdk.server.executor import CancelData, Executor, RunData
-from acp_sdk.server.session import Session
 from acp_sdk.server.store import MemoryStore, Store
 from acp_sdk.server.utils import stream_sse, wait_util_stop
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class Headers(str, Enum):
@@ -52,23 +61,34 @@ class Headers(str, Enum):
 
 
 def create_app(
-    *agents: Agent,
+    *agents: AgentManifest,
     store: Store | None = None,
+    resource_store: ResourceStore | None = None,
+    resource_loader: ResourceLoader | None = None,
+    forward_resources: bool = True,
     lifespan: Lifespan[AppType] | None = None,
     dependencies: list[Depends] | None = None,
 ) -> FastAPI:
+    if not forward_resources and (
+        resource_store is None
+        or isinstance(resource_store._store, (obstore.store.MemoryStore, obstore.store.LocalStore))
+    ):
+        raise ValueError("Resource forwarding must be enabled when resource store does not support HTTP URLs")
+
     executor: ThreadPoolExecutor
+    client = httpx.AsyncClient()
 
     @asynccontextmanager
     async def internal_lifespan(app: FastAPI) -> AsyncGenerator[None]:
         nonlocal executor
-        with ThreadPoolExecutor() as exec:
-            executor = exec
-            if not lifespan:
-                yield None
-            else:
-                async with lifespan(app) as state:
-                    yield state
+        async with client:
+            with ThreadPoolExecutor() as exec:
+                executor = exec
+                if not lifespan:
+                    yield None
+                else:
+                    async with lifespan(app) as state:
+                        yield state
 
     app = FastAPI(
         lifespan=internal_lifespan,
@@ -83,7 +103,7 @@ def create_app(
         allow_credentials=True,
     )
 
-    agents: dict[AgentName, Agent] = {agent.name: agent for agent in agents}
+    agents: dict[AgentName, AgentManifest] = {agent.name: agent for agent in agents}
 
     store = store or MemoryStore(limit=1000, ttl=timedelta(hours=1))
     run_store = store.as_store(model=RunData, prefix="run_")
@@ -91,6 +111,9 @@ def create_app(
     run_resume_store = store.as_store(model=AwaitResume, prefix="run_resume_")
     session_store = store.as_store(model=Session, prefix="session_")
 
+    resource_loader = resource_loader or ResourceLoader(client=client)
+    resource_store = resource_store or ResourceStore(store=obstore.store.MemoryStore())
+
     app.exception_handler(ACPError)(acp_error_handler)
     app.exception_handler(StarletteHTTPException)(http_exception_handler)
     app.exception_handler(RequestValidationError)(validation_exception_handler)
@@ -107,7 +130,7 @@ def create_app(
             run_data.run.status = RunStatus.CANCELLING
         return run_data
 
-    def find_agent(agent_name: AgentName) -> Agent:
+    def find_agent(agent_name: AgentName) -> AgentManifest:
         agent = agents.get(agent_name, None)
         if not agent:
             raise HTTPException(status_code=404, detail=f"Agent {agent_name} not found")
@@ -132,36 +155,54 @@ def create_app(
         return PingResponse()
 
     @app.post("/runs")
-    async def create_run(request: RunCreateRequest) -> RunCreateResponse:
+    async def create_run(request: RunCreateRequest, req: Request) -> RunCreateResponse:
         agent = find_agent(request.agent_name)
 
-        session = (
-            (await session_store.get(request.session_id)) or Session(id=request.session_id)
+        if request.session_id and request.session and request.session_id != request.session.id:
+            raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail="Session ID mismatch")
+
+        session = request.session or (
+            (
+                await session_store.get(request.session_id)
+                or Session(id=request.session_id, loader=resource_loader, store=resource_store)
+            )
             if request.session_id
-            else Session()
+            else Session(loader=resource_loader, store=resource_store)
         )
+
         nonlocal executor
         run_data = RunData(
-            run=Run(agent_name=agent.name, session_id=session.id),
-            input=request.input,
+            run=Run(
+                agent_name=agent.name,
+                session_id=session.id,
+            )
         )
         await run_store.set(run_data.key, run_data)
-
-        session.append(run_data.run.run_id)
         await session_store.set(session.id, session)
 
         headers = {Headers.RUN_ID: str(run_data.run.run_id)}
         ready = asyncio.Event()
 
+        async def create_resource_url(id: ResourceId) -> ResourceUrl:
+            if forward_resources:
+                return ResourceUrl(url=str(req.url_for("get_resource", resource_id=id)))
+            else:
+                return await resource_store.url(id)
+
         Executor(
             agent=agent,
             run_data=run_data,
-            history=await session.history(run_store),
+            session=session,
+            session_store=session_store,
             run_store=run_store,
             cancel_store=run_cancel_store,
             resume_store=run_resume_store,
             executor=executor,
-        ).execute(wait=ready)
+            request=req,
+            resource_store=resource_store,
+            resource_loader=resource_loader,
+            create_resource_url=create_resource_url,
+        ).execute(request.input, wait=ready)
 
         match request.mode:
             case RunMode.STREAM:
@@ -242,4 +283,21 @@ def create_app(
         run_data.run.status = RunStatus.CANCELLING
         return JSONResponse(status_code=status.HTTP_202_ACCEPTED, content=jsonable_encoder(run_data.run))
 
+    @app.get("/sessions/{session_id}")
+    async def read_session(session_id: SessionId) -> SessionReadResponse:
+        session = await session_store.get(session_id)
+        if not session:
+            raise HTTPException(status_code=404, detail=f"Session {session_id} not found")
+        return session
+
+    if forward_resources:
+
+        @app.get("/resources/{resource_id}", name="get_resource")
+        async def read_resource(resource_id: ResourceId) -> StreamingResponse:
+            try:
+                result = await resource_store.load(resource_id)
+                return StreamingResponse(result)
+            except NotFoundError:
+                raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Resource {resource_id} not found")
+
     return app

acp_sdk/server/context.py

@@ -1,22 +1,30 @@
 from concurrent.futures import ThreadPoolExecutor
 
 import janus
+from fastapi import Request
 
-from acp_sdk.models import SessionId
+from acp_sdk.models import Session
 from acp_sdk.server.types import RunYield, RunYieldResume
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class Context:
     def __init__(
         self,
         *,
-        session_id: SessionId | None = None,
+        session: Session,
+        store: ResourceStore,
+        loader: ResourceLoader,
         executor: ThreadPoolExecutor,
+        request: Request,
         yield_queue: janus.Queue[RunYield],
         yield_resume_queue: janus.Queue[RunYieldResume],
     ) -> None:
-        self.session_id = session_id
+        self.session = session
+        self.storage = store
+        self.loader = loader
         self.executor = executor
+        self.request = request
         self._yield_queue = yield_queue
         self._yield_resume_queue = yield_resume_queue
 

acp_sdk/server/executor.py

@@ -1,10 +1,14 @@
 import asyncio
+import inspect
 import logging
-from collections.abc import AsyncIterator
+import uuid
+from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Generator
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
-from typing import Self
+from typing import Callable, Self
 
+import janus
+from fastapi import Request
 from pydantic import BaseModel, ValidationError
 
 from acp_sdk.instrumentation import get_tracer
@@ -22,6 +26,8 @@ from acp_sdk.models import (
     MessageCreatedEvent,
     MessagePart,
     MessagePartEvent,
+    ResourceId,
+    ResourceUrl,
     Run,
     RunAwaitingEvent,
     RunCancelledEvent,
@@ -30,15 +36,18 @@ from acp_sdk.models import (
     RunFailedEvent,
     RunInProgressEvent,
     RunStatus,
+    Session,
 )
-from acp_sdk.server.agent import Agent
+from acp_sdk.server.agent import AgentManifest
+from acp_sdk.server.context import Context
 from acp_sdk.server.logging import logger
 from acp_sdk.server.store import Store
+from acp_sdk.server.types import RunYield, RunYieldResume
+from acp_sdk.shared import ResourceLoader, ResourceStore
 
 
 class RunData(BaseModel):
     run: Run
-    input: list[Message]
     events: list[Event] = []
 
     @property
@@ -62,27 +71,38 @@ class Executor:
     def __init__(
         self,
         *,
-        agent: Agent,
+        agent: AgentManifest,
         run_data: RunData,
-        history: list[Message],
+        session: Session,
         executor: ThreadPoolExecutor,
+        request: Request,
         run_store: Store[RunData],
         cancel_store: Store[CancelData],
         resume_store: Store[AwaitResume],
+        session_store: Store[Session],
+        resource_store: ResourceStore,
+        resource_loader: ResourceLoader,
+        create_resource_url: Callable[[ResourceId], Awaitable[ResourceUrl]],
     ) -> None:
         self.agent = agent
-        self.history = history
+        self.session = session
         self.run_data = run_data
         self.executor = executor
+        self.request = request
 
         self.run_store = run_store
         self.cancel_store = cancel_store
         self.resume_store = resume_store
+        self.session_store = session_store
+        self.resource_store = resource_store
+        self.resource_loader = resource_loader
+
+        self.create_resource_url = create_resource_url
 
         self.logger = logging.LoggerAdapter(logger, {"run_id": str(run_data.run.run_id)})
 
-    def execute(self, *, wait: asyncio.Event) -> None:
-        self.task = asyncio.create_task(self._execute(self.run_data, executor=self.executor, wait=wait))
+    def execute(self, input: list[Message], *, wait: asyncio.Event) -> None:
+        self.task = asyncio.create_task(self._execute(input=input, executor=self.executor, wait=wait))
         self.watcher = asyncio.create_task(self._watch_for_cancellation())
 
     async def _push(self) -> None:
@@ -108,7 +128,16 @@ class Executor:
             except Exception:
                 logger.warning("Cancellation watcher failed, restarting")
 
-    async def _execute(self, run_data: RunData, *, executor: ThreadPoolExecutor, wait: asyncio.Event) -> None:
+    async def _record_session(self, history: list[Message]) -> None:
+        for message in history:
+            id = uuid.uuid4()
+            url = await self.create_resource_url(id)
+            await self.resource_store.store(id, message.model_dump_json().encode())
+            self.session.history.append(url)
+        await self.session_store.set(self.session.id, self.session)
+
+    async def _execute(self, input: list[Message], *, executor: ThreadPoolExecutor, wait: asyncio.Event) -> None:
+        run_data = self.run_data
         with get_tracer().start_as_current_span("run"):
             in_message = False
 
@@ -118,15 +147,22 @@ class Executor:
                     message = run_data.run.output[-1]
                     message.completed_at = datetime.now(timezone.utc)
                     await self._emit(MessageCompletedEvent(message=message))
+                    session_history.append(message)
                     in_message = False
 
+            session_history = input.copy()
             try:
                 await wait.wait()
 
                 await self._emit(RunCreatedEvent(run=run_data.run))
 
-                generator = self.agent.execute(
-                    input=self.history + run_data.input, session_id=run_data.run.session_id, executor=executor
+                generator = self._execute_agent(
+                    input=input,
+                    session=self.session,
+                    storage=self.resource_store,
+                    loader=self.resource_loader,
+                    executor=executor,
+                    request=self.request,
                 )
                 self.logger.info("Run started")
 
@@ -141,18 +177,21 @@ class Executor:
                         if isinstance(next, str):
                             next = MessagePart(content=next)
                         if not in_message:
-                            run_data.run.output.append(Message(parts=[], completed_at=None))
+                            run_data.run.output.append(
+                                Message(role=f"agent/{self.agent.name}", parts=[], completed_at=None)
+                            )
                             in_message = True
                             await self._emit(MessageCreatedEvent(message=run_data.run.output[-1]))
                         run_data.run.output[-1].parts.append(next)
                         await self._emit(MessagePartEvent(part=next))
                     elif isinstance(next, Message):
                         await flush_message()
-                        run_data.run.output.append(next)
+                        run_data.run.output.append(next.model_copy(update={"role": f"agent/{self.agent.name}"}))
                         await self._emit(MessageCreatedEvent(message=next))
                         for part in next.parts:
                             await self._emit(MessagePartEvent(part=part))
                         await self._emit(MessageCompletedEvent(message=next))
+                        session_history.append(next)
                     elif isinstance(next, AwaitRequest):
                         run_data.run.await_request = next
                         run_data.run.status = RunStatus.AWAITING
@@ -180,6 +219,10 @@ class Executor:
                 await flush_message()
                 run_data.run.status = RunStatus.COMPLETED
                 run_data.run.finished_at = datetime.now(timezone.utc)
+                try:
+                    await self._record_session(session_history)
+                except Exception as e:
+                    self.logger.warning(f"Failed to record session: {e}")
                 await self._emit(RunCompletedEvent(run=run_data.run))
                 self.logger.info("Run completed")
             except asyncio.CancelledError:
@@ -196,3 +239,85 @@ class Executor:
                 run_data.run.finished_at = datetime.now(timezone.utc)
                 await self._emit(RunFailedEvent(run=run_data.run))
                 self.logger.exception("Run failed")
+
+    async def _execute_agent(
+        self,
+        input: list[Message],
+        session: Session,
+        storage: ResourceStore,
+        loader: ResourceLoader,
+        executor: ThreadPoolExecutor,
+        request: Request,
+    ) -> AsyncGenerator[RunYield, RunYieldResume]:
+        yield_queue: janus.Queue[RunYield] = janus.Queue()
+        yield_resume_queue: janus.Queue[RunYieldResume] = janus.Queue()
+
+        context = Context(
+            session=session,
+            store=storage,
+            loader=loader,
+            executor=executor,
+            request=request,
+            yield_queue=yield_queue,
+            yield_resume_queue=yield_resume_queue,
+        )
+
+        if inspect.isasyncgenfunction(self.agent.run):
+            run = asyncio.create_task(self._run_async_gen(input, context))
+        elif inspect.iscoroutinefunction(self.agent.run):
+            run = asyncio.create_task(self._run_coro(input, context))
+        elif inspect.isgeneratorfunction(self.agent.run):
+            run = asyncio.get_running_loop().run_in_executor(executor, self._run_gen, input, context)
+        else:
+            run = asyncio.get_running_loop().run_in_executor(executor, self._run_func, input, context)
+
+        try:
+            while not run.done() or yield_queue.async_q.qsize() > 0:
+                value = yield await yield_queue.async_q.get()
+                if isinstance(value, Exception):
+                    raise value
+                await yield_resume_queue.async_q.put(value)
+        except janus.AsyncQueueShutDown:
+            pass
+
+    async def _run_async_gen(self, input: list[Message], context: Context) -> None:
+        try:
+            gen: AsyncGenerator[RunYield, RunYieldResume] = self.agent.run(input, context)
+            value = None
+            while True:
+                value = await context.yield_async(await gen.asend(value))
+        except StopAsyncIteration:
+            pass
+        except Exception as e:
+            await context.yield_async(e)
+        finally:
+            context.shutdown()
+
+    async def _run_coro(self, input: list[Message], context: Context) -> None:
+        try:
+            await context.yield_async(await self.agent.run(input, context))
+        except Exception as e:
+            await context.yield_async(e)
+        finally:
+            context.shutdown()
+
+    def _run_gen(self, input: list[Message], context: Context) -> None:
+        try:
+            gen: Generator[RunYield, RunYieldResume] = self.agent.run(input, context)
+            value = None
+            while True:
+                value = context.yield_sync(gen.send(value))
+        except StopIteration:
+            pass
+        except Exception as e:
+            context.yield_sync(e)
+        finally:
+            context.shutdown()
+
+    def _run_func(self, input: list[Message], context: Context) -> None:
+        try:
+            context.yield_sync(self.agent.run(input, context))
+        except Exception as e:
+            context.yield_sync(e)
+        finally:
+            context.shutdown()

acp_sdk/server/server.py

@@ -1,5 +1,6 @@
 import asyncio
 import os
+import re
 from collections.abc import AsyncGenerator, Awaitable
 from contextlib import asynccontextmanager
 from typing import Any, Callable
@@ -10,7 +11,7 @@ import uvicorn.config
 from fastapi import FastAPI
 
 from acp_sdk.models import Metadata
-from acp_sdk.server.agent import Agent
+from acp_sdk.server.agent import AgentManifest
 from acp_sdk.server.agent import agent as agent_decorator
 from acp_sdk.server.app import create_app
 from acp_sdk.server.logging import configure_logger as configure_logger_func
@@ -18,11 +19,12 @@ from acp_sdk.server.logging import logger
 from acp_sdk.server.store import Store
 from acp_sdk.server.telemetry import configure_telemetry as configure_telemetry_func
 from acp_sdk.server.utils import async_request_with_retry
+from acp_sdk.shared.resources import ResourceLoader, ResourceStore
 
 
 class Server:
     def __init__(self) -> None:
-        self.agents: list[Agent] = []
+        self.agents: list[AgentManifest] = []
         self.server: uvicorn.Server | None = None
 
     def agent(
@@ -41,7 +43,7 @@ class Server:
 
         return decorator
 
-    def register(self, *agents: Agent) -> None:
+    def register(self, *agents: AgentManifest) -> None:
         self.agents.extend(agents)
 
     @asynccontextmanager
@@ -55,6 +57,8 @@ class Server:
         configure_telemetry: bool = False,
         self_registration: bool = True,
         store: Store | None = None,
+        resource_store: ResourceStore | None = None,
+        resource_loader: ResourceLoader | None = None,
         host: str = "127.0.0.1",
         port: int = 8000,
         uds: str | None = None,
@@ -117,7 +121,13 @@ class Server:
 
         import uvicorn
 
-        app = create_app(*self.agents, lifespan=self.lifespan, store=store)
+        app = create_app(
+            *self.agents,
+            lifespan=self.lifespan,
+            store=store,
+            resource_loader=resource_loader,
+            resource_store=resource_store,
+        )
 
         if configure_logger:
             configure_logger_func()
@@ -184,6 +194,8 @@ class Server:
         configure_telemetry: bool = False,
         self_registration: bool = True,
         store: Store | None = None,
+        resource_store: ResourceStore | None = None,
+        resource_loader: ResourceLoader | None = None,
         host: str = "127.0.0.1",
         port: int = 8000,
         uds: str | None = None,
@@ -242,6 +254,8 @@ class Server:
                 configure_telemetry=configure_telemetry,
                 self_registration=self_registration,
                 store=store,
+                resource_store=resource_store,
+                resource_loader=resource_loader,
                 host=host,
                 port=port,
                 uds=uds,
@@ -313,22 +327,33 @@ class Server:
             return
 
         url = os.getenv("PLATFORM_URL", "http://127.0.0.1:8333")
+        host = re.sub(r"localhost|127\.0\.0\.1", "host.docker.internal", self.server.config.host)
         request_data = {
-            "location": f"http://{self.server.config.host}:{self.server.config.port}",
+            "location": f"http://{host}:{self.server.config.port}",
         }
+        await async_request_with_retry(lambda client, data=request_data: client.get(f"{url}/api/v1/providers"))
         try:
             await async_request_with_retry(
-                lambda client, data=request_data: client.post(f"{url}/api/v1/providers/register/unmanaged", json=data)
+                lambda client, data=request_data: client.post(
+                    f"{url}/api/v1/providers", json=data, params={"auto_remove": True}
+                )
             )
             logger.info("Agent registered to the beeai server.")
 
             # check missing env keyes
             envs_request = await async_request_with_retry(lambda client: client.get(f"{url}/api/v1/variables"))
             envs = envs_request.get("env")
+            os.environ["LLM_MODEL"] = "dummy"
+            os.environ["LLM_API_KEY"] = "dummy"
+            os.environ["LLM_API_BASE"] = f"{url.rstrip('/')}/api/v1/llm"
+
             for agent in self.agents:
                 # register all available envs
                 missing_keyes = []
                 for env in agent.metadata.model_dump().get("env", []):
+                    # Those envs are set to use LLM gateway from platform server
+                    if env["name"] in {"LLM_MODEL", "LLM_API_KEY", "LLM_API_BASE"}:
+                        continue
                     server_env = envs.get(env.get("name"))
                     if server_env:
                         logger.debug(f"Env variable {env['name']} = '{server_env}' added dynamically")

acp_sdk/server/store/memory_store.py

@@ -5,25 +5,25 @@ from typing import Generic
 
 from cachetools import TTLCache
 
-from acp_sdk.server.store.store import Store, T
+from acp_sdk.server.store.store import Store, StoreModel, T
 from acp_sdk.server.store.utils import Stringable
 
 
 class MemoryStore(Store[T], Generic[T]):
     def __init__(self, *, limit: int, ttl: int | None = None) -> None:
         super().__init__()
-        self._cache: TTLCache[str, T] = TTLCache(maxsize=limit, ttl=ttl, timer=datetime.now)
+        self._cache: TTLCache[str, str] = TTLCache(maxsize=limit, ttl=ttl, timer=datetime.now)
         self._event = asyncio.Event()
 
     async def get(self, key: Stringable) -> T | None:
         value = self._cache.get(str(key))
-        return value.model_copy(deep=True) if value else value
+        return StoreModel.model_validate_json(value) if value else value
 
     async def set(self, key: Stringable, value: T | None) -> None:
         if value is None:
             del self._cache[str(key)]
         else:
-            self._cache[str(key)] = value.model_copy(deep=True)
+            self._cache[str(key)] = value.model_dump_json()
         self._event.set()
 
     async def watch(self, key: Stringable, *, ready: asyncio.Event | None = None) -> AsyncIterator[T | None]:

acp_sdk/server/utils.py

@@ -9,7 +9,7 @@ from pydantic import BaseModel
 from acp_sdk.models import RunStatus
 from acp_sdk.server.executor import RunData
 from acp_sdk.server.logging import logger
-from acp_sdk.server.store.store import Store
+from acp_sdk.server.store import Store
 
 
 def encode_sse(model: BaseModel) -> str:

acp_sdk/shared/__init__.py

@@ -0,0 +1,2 @@
+from acp_sdk.shared.resources import ResourceLoader as ResourceLoader
+from acp_sdk.shared.resources import ResourceStore as ResourceStore