airbyte.agents.connectors

  1# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
  2"""Airbyte Agents connectors, and the single-action `execute` interface.
  3
  4> ## ⚠️ Experimental Interface
  5>
  6> **The Airbyte Agents Python interfaces are experimental.** Class names, method signatures,
  7> and result models may change or be removed without notice between minor versions of
  8> PyAirbyte. Pin an exact PyAirbyte version if you depend on them.
  9"""
 10
 11from __future__ import annotations
 12
 13from enum import Enum
 14from typing import TYPE_CHECKING, Any
 15
 16from airbyte._direct_connectors import api_util as _api_util
 17from airbyte._direct_connectors.actions import (
 18    UNSUPPORTED_ACTIONS,
 19    AgentAction,
 20    AgentReadAction,
 21    AgentWriteAction,
 22    _build_params,
 23)
 24from airbyte._direct_connectors.models import AgentConnectorDetails, AgentExecuteResult
 25from airbyte.exceptions import PyAirbyteInputError
 26
 27
 28if TYPE_CHECKING:
 29    from collections.abc import Iterator
 30
 31    from airbyte.cloud._credentials import _AirbyteCredentials
 32
 33
 34class AgentConnector:
 35    """A connector in an Airbyte Agents workspace.
 36
 37    Get one from `AgentWorkspace.get_connector()` rather than constructing it directly.
 38
 39    ```python
 40    from airbyte import agents
 41
 42    workspace = agents.AgentWorkspace.from_env()
 43    connector = workspace.get_connector("GitHub")  # by ID or name (case insensitive)
 44    result = connector.list_entities("issues", api_args={"repository": "airbytehq/PyAirbyte"})
 45    for entity in result.entities:
 46        print(entity["title"])
 47    ```
 48    """
 49
 50    def __init__(
 51        self,
 52        connector_id: str,
 53        *,
 54        credentials: _AirbyteCredentials,
 55        name: str | None = None,
 56        workspace_id: str | None = None,
 57    ) -> None:
 58        """Initialize an `AgentConnector`. Prefer `AgentWorkspace.get_connector()`."""
 59        self.connector_id = connector_id
 60        """The connector ID."""
 61
 62        self.workspace_id = workspace_id
 63        """The workspace ID."""
 64
 65        self._credentials = credentials
 66        self._name = name
 67        self._details: AgentConnectorDetails | None = None
 68
 69    @property
 70    def name(self) -> str | None:
 71        """The connector name, fetched from the Agents API if not already known."""
 72        if self._name is None:
 73            self._name = self.inspect().name
 74        return self._name
 75
 76    def inspect(self, *, force_refresh: bool = False) -> AgentConnectorDetails:
 77        """Return connector metadata from the Agents API `inspect` endpoint.
 78
 79        The result is cached; pass `force_refresh=True` to fetch it again.
 80        """
 81        if self._details is None or force_refresh:
 82            self._details = AgentConnectorDetails.model_validate(
 83                _api_util.inspect_agent_connector(
 84                    connector_id=self.connector_id,
 85                    credentials=self._credentials,
 86                    organization_id=self._credentials.organization_id,
 87                )
 88            )
 89            self._name = self._details.name or self._name
 90        return self._details
 91
 92    def execute(  # noqa: PLR0913  # Explicit args are the point of this public API.
 93        self,
 94        entity_type: str,
 95        action: AgentAction | str,
 96        api_args: dict[str, Any] | None = None,
 97        *,
 98        select_fields: list[str] | None = None,
 99        exclude_fields: list[str] | None = None,
100        page_size: int | None = None,
101        cursor: str | None = None,
102        workspace_id: str | None = None,
103        skip_truncation: bool = True,
104        intent: str | None = None,
105    ) -> AgentExecuteResult:
106        """Execute a single action against one entity type on this connector.
107
108        `entity_type` and `action` are connector-specific, for example `issues` and `list`.
109        Use `inspect()` to see what a connector supports.
110
111        `api_args` holds connector-specific arguments passed through to the connector, for
112        example `{"repository": "airbytehq/PyAirbyte"}`. All other arguments are interpreted
113        by PyAirbyte or by the Agents API itself:
114
115        - `select_fields` and `exclude_fields` prune fields from returned entities.
116        - `page_size` and `cursor` are merged into `api_args` as pagination arguments. Context
117          Store `search` uses `limit` and `cursor`, `sql_select` uses the top-level `cursor`, and
118          direct connector actions use their own pagination arguments in `api_args`. Pass the
119          result cursor according to the action type.
120        - `workspace_id` selects the workspace an action runs against. It defaults to the
121          connector's workspace and is currently sent for `sql_select`, the only action the
122          Agents API scopes by workspace; direct connector actions are scoped by the connector.
123        - `skip_truncation` disables the Agents API's default truncation of large payloads.
124        - `intent` is a free-text description of why the action is being run, which some
125          connectors use to refine results.
126
127        The `download` action is rejected, because it returns a binary stream and PyAirbyte
128        does not yet support streaming responses.
129        """
130        if action in UNSUPPORTED_ACTIONS:
131            raise PyAirbyteInputError(
132                message=f"The {action!r} action is not supported by PyAirbyte.",
133                guidance=(
134                    "This action returns a binary stream instead of JSON, and PyAirbyte does "
135                    "not yet support streaming responses."
136                ),
137                context={"entity_type": entity_type, "action": action},
138            )
139
140        if action not in {*AgentReadAction, *AgentWriteAction}:
141            action_names = ", ".join(
142                member.value for member in (*AgentReadAction, *AgentWriteAction)
143            )
144            raise PyAirbyteInputError(
145                message=f"The {action!r} action is not a valid action name for `execute`.",
146                guidance=f"Use one of: {action_names}.",
147                context={"entity_type": entity_type, "action": action},
148            )
149
150        action_value = action.value if isinstance(action, Enum) else action
151        params = _build_params(api_args=api_args, page_size=page_size, cursor=cursor)
152        if action_value == AgentReadAction.SQL_SELECT.value:
153            resolved_workspace_id = workspace_id or self.workspace_id
154            if (
155                resolved_workspace_id is not None
156                and params.get("workspace_id") is None
157                and params.get("workspace_name") is None
158            ):
159                params.pop("workspace_name", None)
160                params["workspace_id"] = resolved_workspace_id
161        request_body: dict[str, Any] = {
162            "entity": entity_type,
163            "action": action_value,
164            "params": params,
165            "skip_truncation": skip_truncation,
166        }
167        if select_fields is not None:
168            request_body["select_fields"] = select_fields
169        if exclude_fields is not None:
170            request_body["exclude_fields"] = exclude_fields
171        if intent is not None:
172            request_body["intent"] = intent
173
174        return AgentExecuteResult.model_validate(
175            _api_util.execute_agent_connector_action(
176                connector_id=self.connector_id,
177                request_body=request_body,
178                credentials=self._credentials,
179                organization_id=self._credentials.organization_id,
180            )
181        )
182
183    def list_entities(
184        self,
185        entity_type: str,
186        api_args: dict[str, Any] | None = None,
187        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
188    ) -> AgentExecuteResult:
189        """Run the `list` action, which returns a page of entities of `entity_type`."""
190        return self.execute(entity_type, "list", api_args, **kwargs)
191
192    def iter_entities(
193        self,
194        entity_type: str,
195        api_args: dict[str, Any] | None = None,
196        *,
197        limit: int | None = None,
198        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `list_entities()`.
199    ) -> Iterator[dict[str, Any]]:
200        """Yield entities of `entity_type`, following the connector's pagination cursor.
201
202        This is the pagination-free way to read entities: each page is fetched lazily as
203        the caller iterates, so no cursor bookkeeping is needed.
204
205        ```python
206        for issue in connector.iter_entities("issues", {"repository": "airbytehq/PyAirbyte"}):
207            print(issue["title"])
208        ```
209
210        `limit` caps how many entities are yielded in total, which matters for entity types
211        with no natural end. Pass `page_size` to control how many are fetched per request.
212
213        Iteration stops early if the connector reports another page without advancing its
214        cursor, rather than requesting the same page forever.
215
216        Use `list_entities()` instead when a single page is enough, or when the result's
217        `status`, `warning`, or `execution_metadata` are needed.
218        """
219        cursor: str | None = kwargs.pop("cursor", None)
220        yield from _api_util.iter_paged_entities(
221            lambda page_cursor: self.list_entities(
222                entity_type,
223                api_args,
224                cursor=page_cursor,
225                **kwargs,
226            ),
227            limit=limit,
228            cursor=cursor,
229        )
230
231    def search_entities(
232        self,
233        entity_type: str,
234        api_args: dict[str, Any] | None = None,
235        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
236    ) -> AgentExecuteResult:
237        """Run the `search` action, which returns matching entities of `entity_type`."""
238        return self.execute(entity_type, "search", api_args, **kwargs)
239
240    def get_entity(
241        self,
242        entity_type: str,
243        api_args: dict[str, Any] | None = None,
244        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
245    ) -> AgentExecuteResult:
246        """Run the `get` action, which returns a single entity of `entity_type`."""
247        return self.execute(entity_type, "get", api_args, **kwargs)
248
249    def create_entity(
250        self,
251        entity_type: str,
252        api_args: dict[str, Any] | None = None,
253        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
254    ) -> AgentExecuteResult:
255        """Run the `create` action, which creates an entity of `entity_type`."""
256        return self.execute(entity_type, "create", api_args, **kwargs)
257
258    def update_entity(
259        self,
260        entity_type: str,
261        api_args: dict[str, Any] | None = None,
262        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
263    ) -> AgentExecuteResult:
264        """Run the `update` action, which updates an entity of `entity_type`."""
265        return self.execute(entity_type, "update", api_args, **kwargs)
266
267    def delete_entity(
268        self,
269        entity_type: str,
270        api_args: dict[str, Any] | None = None,
271        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
272    ) -> AgentExecuteResult:
273        """Run the `delete` action, which deletes an entity of `entity_type`."""
274        return self.execute(entity_type, "delete", api_args, **kwargs)