Writing a custom framework adapter

from __future__ import annotations

import uuid
from typing import Any, AsyncGenerator

from ag_ui.core import BaseEvent
from ag_ui.core.types import RunAgentInput

from idun_agent_engine.agent.base import BaseAgent
from idun_agent_schema.engine.agent_framework import AgentFramework
from idun_agent_schema.engine.capabilities import (
    AgentCapabilities,
    CapabilityFlags,
    InputDescriptor,
    OutputDescriptor,
)
from idun_agent_schema.engine.sessions import HistoryCapabilities


class MyAdapter(BaseAgent):
    def __init__(self) -> None:
        super().__init__()  # critical: wires run_event_observers
        self._id = str(uuid.uuid4())
        self._instance: Any = None

    @property
    def id(self) -> str:
        return self._id

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

    @property
    def agent_instance(self) -> Any:
        if self._instance is None:
            raise RuntimeError("Agent not initialized")
        return self._instance

    @property
    def copilotkit_agent_instance(self) -> Any:
        # If your framework has no CopilotKit wrapper, return the same object.
        # Browser clients calling /agent/run still work via the engine's
        # AG-UI translation layer.
        return self.agent_instance

    @property
    def infos(self) -> dict[str, Any]:
        return {"framework": "MyFramework", "version": "0.1.0"}

    async def initialize(
        self,
        config: dict,
        observability: list | None = None,
    ) -> None:
        # 1. Validate config against your Pydantic model
        # 2. Instantiate your framework's agent
        # 3. Wire observability callbacks if your framework supports them
        from my_framework import build_agent

        self._instance = build_agent(config)

    async def invoke(self, message: Any) -> Any:
        # Deprecated /agent/invoke path. Optional if you do not need
        # the legacy shim; raise NotImplementedError to skip.
        return await self._instance.run(message)

    async def stream(self, message: Any) -> AsyncGenerator[BaseEvent, None]:
        # Deprecated /agent/stream path. Optional; raise NotImplementedError
        # to skip the legacy shim.
        if False:  # pragma: no cover
            yield  # type: ignore[misc]
        raise NotImplementedError("stream() not implemented; use run()")

    async def run(
        self, input_data: RunAgentInput
    ) -> AsyncGenerator[BaseEvent, None]:
        # Canonical AG-UI entry point. The /agent/run route iterates this
        # generator and forwards each event to the SSE response. Translate
        # your framework's native events into ag_ui.core event types here.
        async for event in self._instance.run_ag_ui(input_data):
            yield event

    def discover_capabilities(self) -> AgentCapabilities:
        return AgentCapabilities(
            framework=AgentFramework.CUSTOM,
            capabilities=CapabilityFlags(streaming=True, history=False, thread_id=True),
            input=InputDescriptor(mode="chat"),
            output=OutputDescriptor(mode="text"),
        )

    def history_capabilities(self) -> HistoryCapabilities:
        return HistoryCapabilities(can_list=False, can_get=False)