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, NamedTuple
 15
 16from airbyte.agents import _api_util
 17from airbyte.agents.models import AgentConnectorDetails, AgentExecuteResult
 18from airbyte.exceptions import PyAirbyteInputError
 19
 20
 21if TYPE_CHECKING:
 22    from collections.abc import Iterator
 23
 24    from airbyte.cloud._credentials import _AirbyteCredentials
 25
 26
 27UNSUPPORTED_ACTIONS: set[str] = {"download"}
 28"""Actions PyAirbyte rejects before sending them to the Agents API.
 29
 30`download` returns a binary stream rather than JSON, and PyAirbyte does not yet support
 31streaming responses, so it is rejected with actionable guidance instead of failing later
 32inside the transport layer.
 33"""
 34
 35
 36class AgentReadAction(str, Enum):
 37    """Connector actions that only read data.
 38
 39    The `search` action is the connector's native API search, parallel to `get` and `list`.
 40    The `sql_select` action runs one read-only SQL statement (or `SHOW TABLES`) on the query
 41    engine behind a destination connector. Pass `sql` and `sql_dialect` (and optionally
 42    `dry_run`) in `api_args`; `entity_type` is ignored for this action.
 43
 44    The `download` action is deliberately absent even though it reads: it returns a binary
 45    stream rather than JSON, which PyAirbyte does not yet support.
 46    """
 47
 48    LIST = "list"
 49    GET = "get"
 50    SEARCH = "search"
 51    SQL_SELECT = "sql_select"
 52
 53
 54class AgentWriteAction(str, Enum):
 55    """Connector actions that create, update, or delete data."""
 56
 57    CREATE = "create"
 58    UPDATE = "update"
 59    DELETE = "delete"
 60
 61
 62AgentAction = AgentReadAction | AgentWriteAction
 63"""Every connector action accepted by `AgentConnector.execute()`."""
 64
 65
 66_PAGINATION_ARGS: dict[str, str] = {"page_size": "limit", "cursor": "cursor"}
 67"""Pagination conveniences PyAirbyte merges into the connector's `params`.
 68
 69Maps the PyAirbyte argument name to the connector's own `params` key: the Agents API calls
 70page size `limit`, which PyAirbyte does not expose under that name because `limit` reads as
 71a cap on the whole result set rather than on one page.
 72"""
 73
 74
 75class _ConnectorLookup(NamedTuple):
 76    """What to look a connector up by, once the lookup arguments have been validated.
 77
 78    Both fields are set when the caller passed a positional value that could be either an
 79    ID or a name, in which case an ID match takes precedence over a name match.
 80    """
 81
 82    connector_id: str | None
 83    name: str | None
 84
 85
 86def _resolve_connector_lookup(
 87    id_or_name: str | None,
 88    /,
 89    *,
 90    id: str | None,  # noqa: A002  # Mirrors the public `id` alias it validates.
 91    connector_id: str | None,
 92    name: str | None,
 93) -> _ConnectorLookup:
 94    """Validate connector lookup arguments and return what to look the connector up by.
 95
 96    `id` and `connector_id` are synonyms, so exactly one of them, `name`, or the positional
 97    `id_or_name` is required. Conflicting synonym values are rejected, as is a blank value,
 98    which would otherwise be treated as an omitted argument.
 99    """
100    all_args = {
101        "id_or_name": id_or_name,
102        "id": id,
103        "connector_id": connector_id,
104        "name": name,
105    }
106
107    blank_args = sorted(
108        key for key, value in all_args.items() if value is not None and not value.strip()
109    )
110    if blank_args:
111        raise PyAirbyteInputError(
112            message="Connector lookup arguments cannot be blank.",
113            guidance="Omit the argument entirely, or pass a non-blank value.",
114            context={"blank_args": blank_args},
115        )
116
117    if id_or_name:
118        keyword_args = sorted(
119            key for key, value in all_args.items() if value and key != "id_or_name"
120        )
121        if keyword_args:
122            raise PyAirbyteInputError(
123                message="A positional connector lookup cannot be combined with keyword arguments.",
124                guidance="Pass the value positionally, or pass `id`, `connector_id`, or `name`.",
125                context={"keyword_args": keyword_args},
126            )
127        return _ConnectorLookup(connector_id=id_or_name, name=id_or_name)
128
129    provided = {
130        key: value for key, value in {"id": id, "connector_id": connector_id}.items() if value
131    }
132    if len(set(provided.values())) > 1:
133        raise PyAirbyteInputError(
134            message="`id` and `connector_id` were given conflicting values.",
135            guidance="These arguments are synonyms, so pass only one of them.",
136            context={"provided": sorted(provided)},
137        )
138
139    if bool(provided) == bool(name):
140        raise PyAirbyteInputError(
141            message="Exactly one connector lookup argument is required.",
142            guidance=(
143                "Pass a connector ID or name positionally, or as `id`, `connector_id`, "
144                "or `name`."
145            ),
146        )
147
148    return _ConnectorLookup(connector_id=next(iter(provided.values()), None), name=name)
149
150
151class AgentConnector:
152    """A connector in an Airbyte Agents workspace.
153
154    Get one from `AgentWorkspace.get_connector()` rather than constructing it directly.
155
156    ```python
157    from airbyte import agents
158
159    workspace = agents.AgentWorkspace.from_env()
160    connector = workspace.get_connector("GitHub")  # by ID or name (case insensitive)
161    result = connector.list_entities("issues", api_args={"repository": "airbytehq/PyAirbyte"})
162    for entity in result.entities:
163        print(entity["title"])
164    ```
165    """
166
167    def __init__(
168        self,
169        connector_id: str,
170        *,
171        credentials: _AirbyteCredentials,
172        name: str | None = None,
173        workspace_id: str | None = None,
174    ) -> None:
175        """Initialize an `AgentConnector`. Prefer `AgentWorkspace.get_connector()`."""
176        self.connector_id = connector_id
177        """The connector ID."""
178
179        self.workspace_id = workspace_id
180        """The workspace ID."""
181
182        self._credentials = credentials
183        self._name = name
184        self._details: AgentConnectorDetails | None = None
185
186    @property
187    def name(self) -> str | None:
188        """The connector name, fetched from the Agents API if not already known."""
189        if self._name is None:
190            self._name = self.inspect().name
191        return self._name
192
193    def inspect(self, *, force_refresh: bool = False) -> AgentConnectorDetails:
194        """Return connector metadata from the Agents API `inspect` endpoint.
195
196        The result is cached; pass `force_refresh=True` to fetch it again.
197        """
198        if self._details is None or force_refresh:
199            self._details = AgentConnectorDetails.model_validate(
200                _api_util.inspect_agent_connector(
201                    connector_id=self.connector_id,
202                    credentials=self._credentials,
203                    organization_id=self._credentials.organization_id,
204                )
205            )
206            self._name = self._details.name or self._name
207        return self._details
208
209    def execute(  # noqa: PLR0913  # Explicit args are the point of this public API.
210        self,
211        entity_type: str,
212        action: AgentAction | str,
213        api_args: dict[str, Any] | None = None,
214        *,
215        select_fields: list[str] | None = None,
216        exclude_fields: list[str] | None = None,
217        page_size: int | None = None,
218        cursor: str | None = None,
219        workspace_id: str | None = None,
220        skip_truncation: bool = True,
221        intent: str | None = None,
222    ) -> AgentExecuteResult:
223        """Execute a single action against one entity type on this connector.
224
225        `entity_type` and `action` are connector-specific, for example `issues` and `list`.
226        Use `inspect()` to see what a connector supports.
227
228        `api_args` holds connector-specific arguments passed through to the connector, for
229        example `{"repository": "airbytehq/PyAirbyte"}`. All other arguments are interpreted
230        by PyAirbyte or by the Agents API itself:
231
232        - `select_fields` and `exclude_fields` prune fields from returned entities.
233        - `page_size` and `cursor` are merged into `api_args` as pagination arguments. Context
234          Store `search` uses `limit` and `cursor`, `sql_select` uses the top-level `cursor`, and
235          direct connector actions use their own pagination arguments in `api_args`. Pass the
236          result cursor according to the action type.
237        - `workspace_id` selects the workspace an action runs against. It defaults to the
238          connector's workspace and is currently sent for `sql_select`, the only action the
239          Agents API scopes by workspace; direct connector actions are scoped by the connector.
240        - `skip_truncation` disables the Agents API's default truncation of large payloads.
241        - `intent` is a free-text description of why the action is being run, which some
242          connectors use to refine results.
243
244        The `download` action is rejected, because it returns a binary stream and PyAirbyte
245        does not yet support streaming responses.
246        """
247        if action in UNSUPPORTED_ACTIONS:
248            raise PyAirbyteInputError(
249                message=f"The {action!r} action is not supported by PyAirbyte.",
250                guidance=(
251                    "This action returns a binary stream instead of JSON, and PyAirbyte does "
252                    "not yet support streaming responses."
253                ),
254                context={"entity_type": entity_type, "action": action},
255            )
256
257        if action not in {*AgentReadAction, *AgentWriteAction}:
258            action_names = ", ".join(
259                member.value for member in (*AgentReadAction, *AgentWriteAction)
260            )
261            raise PyAirbyteInputError(
262                message=f"The {action!r} action is not a valid action name for `execute`.",
263                guidance=f"Use one of: {action_names}.",
264                context={"entity_type": entity_type, "action": action},
265            )
266
267        action_value = action.value if isinstance(action, Enum) else action
268        params = _build_params(api_args=api_args, page_size=page_size, cursor=cursor)
269        if action_value == AgentReadAction.SQL_SELECT.value:
270            resolved_workspace_id = workspace_id or self.workspace_id
271            if (
272                resolved_workspace_id is not None
273                and params.get("workspace_id") is None
274                and params.get("workspace_name") is None
275            ):
276                params.pop("workspace_name", None)
277                params["workspace_id"] = resolved_workspace_id
278        request_body: dict[str, Any] = {
279            "entity": entity_type,
280            "action": action_value,
281            "params": params,
282            "skip_truncation": skip_truncation,
283        }
284        if select_fields is not None:
285            request_body["select_fields"] = select_fields
286        if exclude_fields is not None:
287            request_body["exclude_fields"] = exclude_fields
288        if intent is not None:
289            request_body["intent"] = intent
290
291        return AgentExecuteResult.model_validate(
292            _api_util.execute_agent_connector_action(
293                connector_id=self.connector_id,
294                request_body=request_body,
295                credentials=self._credentials,
296                organization_id=self._credentials.organization_id,
297            )
298        )
299
300    def list_entities(
301        self,
302        entity_type: str,
303        api_args: dict[str, Any] | None = None,
304        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
305    ) -> AgentExecuteResult:
306        """Run the `list` action, which returns a page of entities of `entity_type`."""
307        return self.execute(entity_type, "list", api_args, **kwargs)
308
309    def iter_entities(
310        self,
311        entity_type: str,
312        api_args: dict[str, Any] | None = None,
313        *,
314        limit: int | None = None,
315        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `list_entities()`.
316    ) -> Iterator[dict[str, Any]]:
317        """Yield entities of `entity_type`, following the connector's pagination cursor.
318
319        This is the pagination-free way to read entities: each page is fetched lazily as
320        the caller iterates, so no cursor bookkeeping is needed.
321
322        ```python
323        for issue in connector.iter_entities("issues", {"repository": "airbytehq/PyAirbyte"}):
324            print(issue["title"])
325        ```
326
327        `limit` caps how many entities are yielded in total, which matters for entity types
328        with no natural end. Pass `page_size` to control how many are fetched per request.
329
330        Iteration stops early if the connector reports another page without advancing its
331        cursor, rather than requesting the same page forever.
332
333        Use `list_entities()` instead when a single page is enough, or when the result's
334        `status`, `warning`, or `execution_metadata` are needed.
335        """
336        cursor: str | None = kwargs.pop("cursor", None)
337        seen_cursors: set[str] = set()
338        yielded = 0
339
340        while True:
341            result = self.list_entities(entity_type, api_args, cursor=cursor, **kwargs)
342            for agent_entity in result.entities:
343                yield agent_entity
344                yielded += 1
345                if limit is not None and yielded >= limit:
346                    return
347
348            cursor = result.end_cursor
349            if not result.has_next_page or cursor is None or cursor in seen_cursors:
350                return
351            seen_cursors.add(cursor)
352
353    def search_entities(
354        self,
355        entity_type: str,
356        api_args: dict[str, Any] | None = None,
357        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
358    ) -> AgentExecuteResult:
359        """Run the `search` action, which returns matching entities of `entity_type`."""
360        return self.execute(entity_type, "search", api_args, **kwargs)
361
362    def get_entity(
363        self,
364        entity_type: str,
365        api_args: dict[str, Any] | None = None,
366        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
367    ) -> AgentExecuteResult:
368        """Run the `get` action, which returns a single entity of `entity_type`."""
369        return self.execute(entity_type, "get", api_args, **kwargs)
370
371    def create_entity(
372        self,
373        entity_type: str,
374        api_args: dict[str, Any] | None = None,
375        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
376    ) -> AgentExecuteResult:
377        """Run the `create` action, which creates an entity of `entity_type`."""
378        return self.execute(entity_type, "create", api_args, **kwargs)
379
380    def update_entity(
381        self,
382        entity_type: str,
383        api_args: dict[str, Any] | None = None,
384        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
385    ) -> AgentExecuteResult:
386        """Run the `update` action, which updates an entity of `entity_type`."""
387        return self.execute(entity_type, "update", api_args, **kwargs)
388
389    def delete_entity(
390        self,
391        entity_type: str,
392        api_args: dict[str, Any] | None = None,
393        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
394    ) -> AgentExecuteResult:
395        """Run the `delete` action, which deletes an entity of `entity_type`."""
396        return self.execute(entity_type, "delete", api_args, **kwargs)
397
398
399def _build_params(
400    *,
401    api_args: dict[str, Any] | None,
402    page_size: int | None,
403    cursor: str | None,
404) -> dict[str, Any]:
405    """Merge the pagination conveniences into the connector-specific `api_args`."""
406    params: dict[str, Any] = dict(api_args or {})
407    pagination: dict[str, Any] = {"page_size": page_size, "cursor": cursor}
408
409    conflicts = sorted(
410        name
411        for name, param_key in _PAGINATION_ARGS.items()
412        if pagination[name] is not None and param_key in params
413    )
414    if conflicts:
415        raise PyAirbyteInputError(
416            message="Pagination arguments were provided twice.",
417            guidance=(
418                "Pass each of `page_size` and `cursor` either as a keyword argument or "
419                "within `api_args`, but not both. Note that `page_size` is sent to the "
420                "connector as `limit`."
421            ),
422            context={"duplicated_args": conflicts},
423        )
424
425    params.update(
426        {
427            param_key: pagination[name]
428            for name, param_key in _PAGINATION_ARGS.items()
429            if pagination[name] is not None
430        }
431    )
432    return params