airbyte.agents

152class AgentConnector:
153    """A connector in an Airbyte Agents workspace.
154
155    Get one from `AgentWorkspace.get_connector()` rather than constructing it directly.
156
157    ```python
158    from airbyte import agents
159
160    workspace = agents.AgentWorkspace.from_env()
161    connector = workspace.get_connector("GitHub")  # by ID or name (case insensitive)
162    result = connector.list_entities("issues", api_args={"repository": "airbytehq/PyAirbyte"})
163    for entity in result.entities:
164        print(entity["title"])
165    ```
166    """
167
168    def __init__(
169        self,
170        connector_id: str,
171        *,
172        credentials: _AirbyteCredentials,
173        name: str | None = None,
174        workspace_id: str | None = None,
175    ) -> None:
176        """Initialize an `AgentConnector`. Prefer `AgentWorkspace.get_connector()`."""
177        self.connector_id = connector_id
178        """The connector ID."""
179
180        self.workspace_id = workspace_id
181        """The workspace ID."""
182
183        self._credentials = credentials
184        self._name = name
185        self._details: AgentConnectorDetails | None = None
186
187    @property
188    def name(self) -> str | None:
189        """The connector name, fetched from the Agents API if not already known."""
190        if self._name is None:
191            self._name = self.inspect().name
192        return self._name
193
194    def inspect(self, *, force_refresh: bool = False) -> AgentConnectorDetails:
195        """Return connector metadata from the Agents API `inspect` endpoint.
196
197        The result is cached; pass `force_refresh=True` to fetch it again.
198        """
199        if self._details is None or force_refresh:
200            self._details = AgentConnectorDetails.model_validate(
201                _api_util.inspect_agent_connector(
202                    connector_id=self.connector_id,
203                    credentials=self._credentials,
204                    organization_id=self._credentials.organization_id,
205                )
206            )
207            self._name = self._details.name or self._name
208        return self._details
209
210    def execute(  # noqa: PLR0913  # Explicit args are the point of this public API.
211        self,
212        entity_type: str,
213        action: AgentAction | str,
214        api_args: dict[str, Any] | None = None,
215        *,
216        select_fields: list[str] | None = None,
217        exclude_fields: list[str] | None = None,
218        page_size: int | None = None,
219        cursor: str | None = None,
220        workspace_id: str | None = None,
221        skip_truncation: bool = True,
222        intent: str | None = None,
223    ) -> AgentExecuteResult:
224        """Execute a single action against one entity type on this connector.
225
226        `entity_type` and `action` are connector-specific, for example `issues` and `list`.
227        Use `inspect()` to see what a connector supports.
228
229        `api_args` holds connector-specific arguments passed through to the connector, for
230        example `{"repository": "airbytehq/PyAirbyte"}`. All other arguments are interpreted
231        by PyAirbyte or by the Agents API itself:
232
233        - `select_fields` and `exclude_fields` prune fields from returned entities.
234        - `page_size` and `cursor` are merged into `api_args` as pagination arguments. Context
235          Store `search` uses `limit` and `cursor`, `sql_select` uses the top-level `cursor`, and
236          direct connector actions use their own pagination arguments in `api_args`. Pass the
237          result cursor according to the action type.
238        - `workspace_id` selects the workspace an action runs against. It defaults to the
239          connector's workspace and is currently sent for `sql_select`, the only action the
240          Agents API scopes by workspace; direct connector actions are scoped by the connector.
241        - `skip_truncation` disables the Agents API's default truncation of large payloads.
242        - `intent` is a free-text description of why the action is being run, which some
243          connectors use to refine results.
244
245        The `download` action is rejected, because it returns a binary stream and PyAirbyte
246        does not yet support streaming responses.
247        """
248        if action in UNSUPPORTED_ACTIONS:
249            raise PyAirbyteInputError(
250                message=f"The {action!r} action is not supported by PyAirbyte.",
251                guidance=(
252                    "This action returns a binary stream instead of JSON, and PyAirbyte does "
253                    "not yet support streaming responses."
254                ),
255                context={"entity_type": entity_type, "action": action},
256            )
257
258        if action not in {*AgentReadAction, *AgentWriteAction}:
259            action_names = ", ".join(
260                member.value for member in (*AgentReadAction, *AgentWriteAction)
261            )
262            raise PyAirbyteInputError(
263                message=f"The {action!r} action is not a valid action name for `execute`.",
264                guidance=f"Use one of: {action_names}.",
265                context={"entity_type": entity_type, "action": action},
266            )
267
268        action_value = action.value if isinstance(action, Enum) else action
269        params = _build_params(api_args=api_args, page_size=page_size, cursor=cursor)
270        if action_value == AgentReadAction.SQL_SELECT.value:
271            resolved_workspace_id = workspace_id or self.workspace_id
272            if (
273                resolved_workspace_id is not None
274                and params.get("workspace_id") is None
275                and params.get("workspace_name") is None
276            ):
277                params.pop("workspace_name", None)
278                params["workspace_id"] = resolved_workspace_id
279        request_body: dict[str, Any] = {
280            "entity": entity_type,
281            "action": action_value,
282            "params": params,
283            "skip_truncation": skip_truncation,
284        }
285        if select_fields is not None:
286            request_body["select_fields"] = select_fields
287        if exclude_fields is not None:
288            request_body["exclude_fields"] = exclude_fields
289        if intent is not None:
290            request_body["intent"] = intent
291
292        return AgentExecuteResult.model_validate(
293            _api_util.execute_agent_connector_action(
294                connector_id=self.connector_id,
295                request_body=request_body,
296                credentials=self._credentials,
297                organization_id=self._credentials.organization_id,
298            )
299        )
300
301    def list_entities(
302        self,
303        entity_type: str,
304        api_args: dict[str, Any] | None = None,
305        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
306    ) -> AgentExecuteResult:
307        """Run the `list` action, which returns a page of entities of `entity_type`."""
308        return self.execute(entity_type, "list", api_args, **kwargs)
309
310    def iter_entities(
311        self,
312        entity_type: str,
313        api_args: dict[str, Any] | None = None,
314        *,
315        limit: int | None = None,
316        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `list_entities()`.
317    ) -> Iterator[dict[str, Any]]:
318        """Yield entities of `entity_type`, following the connector's pagination cursor.
319
320        This is the pagination-free way to read entities: each page is fetched lazily as
321        the caller iterates, so no cursor bookkeeping is needed.
322
323        ```python
324        for issue in connector.iter_entities("issues", {"repository": "airbytehq/PyAirbyte"}):
325            print(issue["title"])
326        ```
327
328        `limit` caps how many entities are yielded in total, which matters for entity types
329        with no natural end. Pass `page_size` to control how many are fetched per request.
330
331        Iteration stops early if the connector reports another page without advancing its
332        cursor, rather than requesting the same page forever.
333
334        Use `list_entities()` instead when a single page is enough, or when the result's
335        `status`, `warning`, or `execution_metadata` are needed.
336        """
337        cursor: str | None = kwargs.pop("cursor", None)
338        seen_cursors: set[str] = set()
339        yielded = 0
340
341        while True:
342            result = self.list_entities(entity_type, api_args, cursor=cursor, **kwargs)
343            for agent_entity in result.entities:
344                yield agent_entity
345                yielded += 1
346                if limit is not None and yielded >= limit:
347                    return
348
349            cursor = result.end_cursor
350            if not result.has_next_page or cursor is None or cursor in seen_cursors:
351                return
352            seen_cursors.add(cursor)
353
354    def search_entities(
355        self,
356        entity_type: str,
357        api_args: dict[str, Any] | None = None,
358        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
359    ) -> AgentExecuteResult:
360        """Run the `search` action, which returns matching entities of `entity_type`."""
361        return self.execute(entity_type, "search", api_args, **kwargs)
362
363    def get_entity(
364        self,
365        entity_type: str,
366        api_args: dict[str, Any] | None = None,
367        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
368    ) -> AgentExecuteResult:
369        """Run the `get` action, which returns a single entity of `entity_type`."""
370        return self.execute(entity_type, "get", api_args, **kwargs)
371
372    def create_entity(
373        self,
374        entity_type: str,
375        api_args: dict[str, Any] | None = None,
376        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
377    ) -> AgentExecuteResult:
378        """Run the `create` action, which creates an entity of `entity_type`."""
379        return self.execute(entity_type, "create", api_args, **kwargs)
380
381    def update_entity(
382        self,
383        entity_type: str,
384        api_args: dict[str, Any] | None = None,
385        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
386    ) -> AgentExecuteResult:
387        """Run the `update` action, which updates an entity of `entity_type`."""
388        return self.execute(entity_type, "update", api_args, **kwargs)
389
390    def delete_entity(
391        self,
392        entity_type: str,
393        api_args: dict[str, Any] | None = None,
394        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
395    ) -> AgentExecuteResult:
396        """Run the `delete` action, which deletes an entity of `entity_type`."""
397        return self.execute(entity_type, "delete", api_args, **kwargs)

150class AgentConnectorDetails(BaseModel):
151    """Connector metadata returned by the Agents API `inspect` endpoint."""
152
153    model_config = ConfigDict(extra="allow", populate_by_name=True)
154
155    connector_id: str
156    """The connector ID."""
157
158    name: str | None = None
159    """The connector name."""
160
161    workspace_id: str | None = None
162    """The ID of the workspace the connector belongs to."""
163
164    organization_id: str | None = None
165    """The ID of the organization the connector belongs to."""
166
167    source_definition_id: str | None = None
168    """The ID of the underlying Airbyte source definition."""
169
170    integration_name: str | None = Field(default=None, alias="source_definition_name")
171    """Name of the underlying integration, for example `GitHub` or `Snowflake`."""
172
173    docs_skill_id: str | None = None
174    """Skill ID to pass to `AgentWorkspace.get_skill(...).read_docs()` (MCP:
175    `read_agent_skill_docs`) for this connector's usage docs."""
176
177    context_store_readiness: AgentContextStoreReadiness | None = None
178    """Context Store readiness information, when reported."""
179
180    warnings: list[Any] = Field(default_factory=list)
181    """Warnings reported by the Agents API, for example degraded capabilities."""
182
183    @property
184    def context_store_entities(self) -> list[str]:
185        """The entity names this connector can cache in the Context Store.
186
187        Note that this lists Context Store-supported entities specifically. The Agents API
188        does not publish an exhaustive list of executable entity and action pairs, so an
189        entity may be executable via `AgentConnector.execute()` without appearing here.
190        """
191        if self.context_store_readiness is None:
192            return []
193        return [
194            entity.entity
195            for entity in self.context_store_readiness.supported_context_store_entities
196        ]

42class AgentConnectorInfo(BaseModel):
43    """Summary information about a connector, as returned by the Agents API."""
44
45    model_config = ConfigDict(extra="allow")
46
47    id: str
48    """The connector ID."""
49
50    name: str | None = None
51    """The connector name, for example `GitHub - <workspace_id>`."""

211class AgentConnectorMetadata(BaseModel):
212    """Connector-reported metadata about a single action's result, including pagination."""
213
214    model_config = ConfigDict(extra="allow")
215
216    has_next_page: bool | None = None
217    """Whether more entities are available after this page, when the connector reports it."""
218
219    end_cursor: str | None = None
220    """The cursor for the next page, when one is available. Pass it as `cursor` for Context Store
221    `search`, or as the connector's own cursor argument in `api_args` for direct connector
222    actions."""

54class AgentContextStoreEntity(BaseModel):
55    """An entity that a connector supports caching in the Airbyte Context Store."""
56
57    model_config = ConfigDict(extra="allow")
58
59    entity: str
60    """The entity name, for example `issues`."""
61
62    suggested: bool | None = None
63    """Whether Airbyte suggests caching this entity."""

66class AgentContextStoreReadiness(BaseModel):
67    """Context Store readiness information for a connector."""
68
69    model_config = ConfigDict(extra="allow")
70
71    supported_context_store_entities: list[AgentContextStoreEntity] = Field(default_factory=list)
72    """The entities this connector can cache in the Context Store."""
73
74    configured_cache_entities: list[dict[str, Any]] = Field(default_factory=list)
75    """The entities currently configured for caching, with their sync status."""

225class AgentExecuteResult(BaseModel):
226    """The result of executing a single action against an Airbyte Agents connector."""
227
228    model_config = ConfigDict(extra="allow")
229
230    status: str
231    """The execution status reported by the Agents API, for example `success`."""
232
233    result: Any = None
234    """The action's payload. Entity-returning actions put a list of entities here."""
235
236    connector_metadata: AgentConnectorMetadata = Field(default_factory=AgentConnectorMetadata)
237    """Connector-reported metadata about the result, including pagination cursors."""
238
239    execution_metadata: AgentExecutionMetadata = Field(default_factory=AgentExecutionMetadata)
240    """Metadata describing how the action was executed."""
241
242    warning: dict[str, Any] | None = None
243    """A warning reported alongside an otherwise successful result."""
244
245    @field_validator("connector_metadata", "execution_metadata", mode="before")
246    @classmethod
247    def _none_to_empty(cls, value: object) -> object:
248        return {} if value is None else value
249
250    @property
251    def entities(self) -> list[dict[str, Any]]:
252        """The result as a list of entities.
253
254        Raises `PyAirbyteInputError` if the action did not return a list of entities. Use
255        `result` for actions whose payload is not a list of entities.
256        """
257        if not isinstance(self.result, list):
258            raise PyAirbyteInputError(
259                message="This action did not return a list of entities.",
260                guidance="Use the `result` attribute to read non-entity result payloads.",
261                context={"result_type": type(self.result).__name__},
262            )
263
264        invalid_types = sorted(
265            {type(entity).__name__ for entity in self.result if not isinstance(entity, dict)}
266        )
267        if invalid_types:
268            raise PyAirbyteInputError(
269                message="This action returned a list that is not a list of entities.",
270                guidance="Use the `result` attribute to read non-entity result payloads.",
271                context={"unexpected_item_types": invalid_types},
272            )
273        return self.result
274
275    @property
276    def has_next_page(self) -> bool:
277        """Whether the connector reported more entities after this page."""
278        return bool(self.connector_metadata.has_next_page)
279
280    @property
281    def end_cursor(self) -> str | None:
282        """The cursor for the next page, or `None` when there is no next page."""
283        return self.connector_metadata.end_cursor

199class AgentExecutionMetadata(BaseModel):
200    """Metadata describing how an Agents connector action was executed."""
201
202    model_config = ConfigDict(extra="allow")
203
204    connector_instance_id: str | None = None
205    """The connector instance that served the request."""
206
207    execution_time_ms: int | None = None
208    """The server-side execution time, in milliseconds."""

 77class AgentOrganization:
 78    """An organization on the Airbyte Agents platform.
 79
 80    ```python
 81    from airbyte import agents
 82
 83    organization = agents.AgentOrganization.from_env()
 84    workspace = organization.get_workspace("my-workspace")  # by ID or name (case insensitive)
 85    ```
 86    """
 87
 88    def __init__(
 89        self,
 90        *,
 91        organization_id: str | None = None,
 92        client_id: str | SecretString | None = None,
 93        client_secret: str | SecretString | None = None,
 94        bearer_token: str | SecretString | None = None,
 95        public_api_root: str | None = None,
 96        config_api_root: str | None = None,
 97    ) -> None:
 98        """Initialize an `AgentOrganization`.
 99
100        Credentials fall back to the `AIRBYTE_CLOUD_*` environment variables when they are
101        not passed explicitly. API roots default to the `AIRBYTE_CLOUD_API_URL` /
102        `AIRBYTE_CLOUD_CONFIG_API_URL` environment variables (public Airbyte Cloud when unset);
103        custom roots require `AIRBYTE_AGENTS_API_URL`. The organization ID is optional: the
104        Agents API infers it when the credentials belong to exactly one organization.
105        """
106        self._credentials = _AirbyteCredentials.from_auth(
107            organization_id=organization_id,
108            client_id=client_id,
109            client_secret=client_secret,
110            bearer_token=bearer_token,
111            public_api_root=public_api_root,
112            config_api_root=config_api_root,
113            # Mirrors `CloudWorkspace.__init__`: any explicit credential disables env
114            # fallback, since an env bearer token plus explicit client creds is rejected
115            # as mutually exclusive auth.
116            env_vars=not (client_id or client_secret or bearer_token),
117        )
118
119        self.organization_id: str | None = self._credentials.organization_id
120        """The organization ID, when known."""
121
122    @classmethod
123    def from_env(cls, organization_id: str | None = None) -> AgentOrganization:
124        """Create an `AgentOrganization` from the `AIRBYTE_CLOUD_*` environment variables."""
125        return cls(organization_id=organization_id)
126
127    def list_workspaces(self) -> list[AgentWorkspace]:
128        """List the workspaces visible to these credentials in the Agents API."""
129        return [
130            self._workspace_from_info(info)
131            for info in (
132                AgentWorkspaceInfo.model_validate(record)
133                for record in _api_util.list_agent_workspaces(
134                    credentials=self._credentials,
135                    organization_id=self.organization_id,
136                )
137            )
138        ]
139
140    def get_workspace(
141        self,
142        id_or_name: str | None = None,
143        /,
144        *,
145        workspace_id: str | None = None,
146        name: str | None = None,
147    ) -> AgentWorkspace:
148        """Get a workspace in this organization, by ID or by name.
149
150        Pass a single positional value to look the workspace up by either its ID or its
151        name, or name the argument to be explicit.
152
153        Lookup by an explicit `workspace_id` does not call the Agents API. Every other form
154        lists the organization's workspaces and matches on ID first, then on an exact name,
155        ignoring case.
156        """
157        lookup = _resolve_workspace_lookup(
158            id_or_name,
159            workspace_id=workspace_id,
160            name=name,
161        )
162
163        if lookup.workspace_id and not lookup.name:
164            return self._workspace_from_info(AgentWorkspaceInfo(id=lookup.workspace_id))
165
166        workspaces = self.list_workspaces()
167        if lookup.workspace_id:
168            id_matches = [
169                workspace
170                for workspace in workspaces
171                if workspace.workspace_id == lookup.workspace_id
172            ]
173            if id_matches:
174                return id_matches[0]
175
176        name_lower = (lookup.name or "").lower()
177        matches = [
178            workspace
179            for workspace in workspaces
180            if workspace.name and workspace.name.lower() == name_lower
181        ]
182        if not matches:
183            raise AirbyteError(
184                message="No workspace found with the given ID or name.",
185                guidance="Use `list_workspaces()` to see the available workspaces.",
186                context={"name": lookup.name},
187            )
188        if len(matches) > 1:
189            raise AirbyteError(
190                message="Multiple workspaces matched the given name.",
191                guidance="Pass `workspace_id` instead of `name`.",
192                context={"name": lookup.name, "match_count": len(matches)},
193            )
194        return matches[0]
195
196    def as_cloud_organization(self) -> CloudOrganization:
197        """Return this organization as an `airbyte.cloud.CloudOrganization`.
198
199        Every Agents organization is also a Cloud organization, so this conversion needs no
200        API call. It requires a known organization ID, and raises `PyAirbyteInputError`
201        when the ID is unknown.
202        """
203        if not self.organization_id:
204            raise PyAirbyteInputError(
205                message="Organization ID is required to convert to a `CloudOrganization`.",
206                guidance=(
207                    "Provide `organization_id`, or set the `AIRBYTE_CLOUD_ORGANIZATION_ID` "
208                    "environment variable."
209                ),
210            )
211        return CloudOrganization(
212            organization_id=self.organization_id,
213            client_id=self._credentials.client_id,
214            client_secret=self._credentials.client_secret,
215            bearer_token=self._credentials.bearer_token,
216            public_api_root=self._credentials.public_api_root,
217            config_api_root=self._credentials.config_api_root,
218        )
219
220    @classmethod
221    def from_cloud_organization(
222        cls,
223        cloud_organization: CloudOrganization,
224    ) -> AgentOrganization:
225        """Return a Cloud organization as an `AgentOrganization`.
226
227        Whether the organization can actually execute connector actions depends on its
228        Airbyte Agents subscription, which is only knowable per workspace. Use
229        `AgentWorkspace.from_cloud_workspace()` for an authoritative eligibility check.
230
231        Raises `AirbyteAgentsUnavailableError` when the Cloud organization uses non-public Cloud
232        API roots unless `AIRBYTE_AGENTS_API_URL` is set; the Cloud API roots are carried over so
233        the token exchange stays on the same deployment.
234        """
235        credentials = cloud_organization._credentials  # noqa: SLF001  # Same-domain conversion.
236        _api_util.check_public_cloud_api_roots(credentials)
237        return cls(
238            organization_id=cloud_organization.organization_id,
239            client_id=credentials.client_id,
240            client_secret=credentials.client_secret,
241            bearer_token=credentials.bearer_token,
242            public_api_root=credentials.public_api_root,
243            config_api_root=credentials.config_api_root,
244        )
245
246    def _workspace_from_info(self, info: AgentWorkspaceInfo) -> AgentWorkspace:
247        """Build an `AgentWorkspace` from workspace info, reusing these credentials."""
248        return AgentWorkspace(
249            workspace_id=info.id,
250            organization_id=info.organization_id or self.organization_id,
251            name=info.name,
252            client_id=self._credentials.client_id,
253            client_secret=self._credentials.client_secret,
254            bearer_token=self._credentials.bearer_token,
255            public_api_root=self._credentials.public_api_root,
256            config_api_root=self._credentials.config_api_root,
257        )