airbyte.agents

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