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