airbyte.agents

113class AgentConnector:
114    """A connector in an Airbyte Agents workspace.
115
116    Get one from `AgentWorkspace.get_connector()` rather than constructing it directly.
117
118    ```python
119    from airbyte import agents
120
121    workspace = agents.AgentWorkspace.from_env()
122    connector = workspace.get_connector("GitHub")  # by ID or name (case insensitive)
123    result = connector.list_entities("issues", api_args={"repository": "airbytehq/PyAirbyte"})
124    for entity in result.entities:
125        print(entity["title"])
126    ```
127    """
128
129    def __init__(
130        self,
131        connector_id: str,
132        *,
133        credentials: _AirbyteCredentials,
134        name: str | None = None,
135    ) -> None:
136        """Initialize an `AgentConnector`. Prefer `AgentWorkspace.get_connector()`."""
137        self.connector_id = connector_id
138        """The connector ID."""
139
140        self._credentials = credentials
141        self._name = name
142        self._details: AgentConnectorDetails | None = None
143
144    @property
145    def name(self) -> str | None:
146        """The connector name, fetched from the Agents API if not already known."""
147        if self._name is None:
148            self._name = self.describe().name
149        return self._name
150
151    def describe(self, *, force_refresh: bool = False) -> AgentConnectorDetails:
152        """Return connector metadata from the Agents API `inspect` endpoint.
153
154        The result is cached; pass `force_refresh=True` to fetch it again.
155        """
156        if self._details is None or force_refresh:
157            self._details = AgentConnectorDetails.model_validate(
158                _api_util.inspect_agent_connector(
159                    connector_id=self.connector_id,
160                    credentials=self._credentials,
161                    organization_id=self._credentials.organization_id,
162                )
163            )
164            self._name = self._details.name or self._name
165        return self._details
166
167    def execute(  # noqa: PLR0913  # Explicit args are the point of this public API.
168        self,
169        entity_type: str,
170        action: str,
171        api_args: dict[str, Any] | None = None,
172        *,
173        select_fields: list[str] | None = None,
174        exclude_fields: list[str] | None = None,
175        page_size: int | None = None,
176        cursor: str | None = None,
177        skip_truncation: bool = True,
178        intent: str | None = None,
179    ) -> AgentExecuteResult:
180        """Execute a single action against one entity type on this connector.
181
182        `entity_type` and `action` are connector-specific, for example `issues` and `list`.
183        Use `describe()` to see what a connector supports.
184
185        `api_args` holds connector-specific arguments passed through to the connector, for
186        example `{"repository": "airbytehq/PyAirbyte"}`. All other arguments are interpreted
187        by PyAirbyte or by the Agents API itself:
188
189        - `select_fields` and `exclude_fields` prune fields from returned entities.
190        - `page_size` and `cursor` are merged into `api_args` as pagination arguments, where
191          the connector receives `page_size` as its own `limit`. Pass the `end_cursor` of a
192          previous result as `cursor` to fetch the next page.
193        - `skip_truncation` disables the Agents API's default truncation of large payloads.
194        - `intent` is a free-text description of why the action is being run, which some
195          connectors use to refine results.
196
197        The `download` action is rejected, because it returns a binary stream and PyAirbyte
198        does not yet support streaming responses.
199        """
200        if action in UNSUPPORTED_ACTIONS:
201            raise PyAirbyteInputError(
202                message=f"The {action!r} action is not supported by PyAirbyte.",
203                guidance=(
204                    "This action returns a binary stream instead of JSON, and PyAirbyte does "
205                    "not yet support streaming responses."
206                ),
207                context={"entity_type": entity_type, "action": action},
208            )
209
210        request_body: dict[str, Any] = {
211            "entity": entity_type,
212            "action": action,
213            "params": _build_params(api_args=api_args, page_size=page_size, cursor=cursor),
214            "skip_truncation": skip_truncation,
215        }
216        if select_fields is not None:
217            request_body["select_fields"] = select_fields
218        if exclude_fields is not None:
219            request_body["exclude_fields"] = exclude_fields
220        if intent is not None:
221            request_body["intent"] = intent
222
223        return AgentExecuteResult.model_validate(
224            _api_util.execute_agent_connector_action(
225                connector_id=self.connector_id,
226                request_body=request_body,
227                credentials=self._credentials,
228                organization_id=self._credentials.organization_id,
229            )
230        )
231
232    def list_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 `list` action, which returns a page of entities of `entity_type`."""
239        return self.execute(entity_type, "list", api_args, **kwargs)
240
241    def iter_entities(
242        self,
243        entity_type: str,
244        api_args: dict[str, Any] | None = None,
245        *,
246        limit: int | None = None,
247        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `list_entities()`.
248    ) -> Iterator[dict[str, Any]]:
249        """Yield entities of `entity_type`, following the connector's pagination cursor.
250
251        This is the pagination-free way to read entities: each page is fetched lazily as
252        the caller iterates, so no cursor bookkeeping is needed.
253
254        ```python
255        for issue in connector.iter_entities("issues", {"repository": "airbytehq/PyAirbyte"}):
256            print(issue["title"])
257        ```
258
259        `limit` caps how many entities are yielded in total, which matters for entity types
260        with no natural end. Pass `page_size` to control how many are fetched per request.
261
262        Iteration stops early if the connector reports another page without advancing its
263        cursor, rather than requesting the same page forever.
264
265        Use `list_entities()` instead when a single page is enough, or when the result's
266        `status`, `warning`, or `execution_metadata` are needed.
267        """
268        cursor: str | None = kwargs.pop("cursor", None)
269        seen_cursors: set[str] = set()
270        yielded = 0
271
272        while True:
273            result = self.list_entities(entity_type, api_args, cursor=cursor, **kwargs)
274            for agent_entity in result.entities:
275                yield agent_entity
276                yielded += 1
277                if limit is not None and yielded >= limit:
278                    return
279
280            cursor = result.end_cursor
281            if not result.has_next_page or cursor is None or cursor in seen_cursors:
282                return
283            seen_cursors.add(cursor)
284
285    def search_entities(
286        self,
287        entity_type: str,
288        api_args: dict[str, Any] | None = None,
289        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
290    ) -> AgentExecuteResult:
291        """Run the `search` action, which returns matching entities of `entity_type`."""
292        return self.execute(entity_type, "search", api_args, **kwargs)
293
294    def get_entity(
295        self,
296        entity_type: str,
297        api_args: dict[str, Any] | None = None,
298        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
299    ) -> AgentExecuteResult:
300        """Run the `get` action, which returns a single entity of `entity_type`."""
301        return self.execute(entity_type, "get", api_args, **kwargs)
302
303    def create_entity(
304        self,
305        entity_type: str,
306        api_args: dict[str, Any] | None = None,
307        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
308    ) -> AgentExecuteResult:
309        """Run the `create` action, which creates an entity of `entity_type`."""
310        return self.execute(entity_type, "create", api_args, **kwargs)
311
312    def update_entity(
313        self,
314        entity_type: str,
315        api_args: dict[str, Any] | None = None,
316        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
317    ) -> AgentExecuteResult:
318        """Run the `update` action, which updates an entity of `entity_type`."""
319        return self.execute(entity_type, "update", api_args, **kwargs)
320
321    def delete_entity(
322        self,
323        entity_type: str,
324        api_args: dict[str, Any] | None = None,
325        **kwargs: Any,  # noqa: ANN401  # Forwarded verbatim to `execute()`.
326    ) -> AgentExecuteResult:
327        """Run the `delete` action, which deletes an entity of `entity_type`."""
328        return self.execute(entity_type, "delete", api_args, **kwargs)

 72class AgentConnectorDetails(BaseModel):
 73    """Connector metadata returned by the Agents API `inspect` endpoint."""
 74
 75    model_config = ConfigDict(extra="allow")
 76
 77    connector_id: str
 78    """The connector ID."""
 79
 80    name: str | None = None
 81    """The connector name."""
 82
 83    workspace_id: str | None = None
 84    """The ID of the workspace the connector belongs to."""
 85
 86    organization_id: str | None = None
 87    """The ID of the organization the connector belongs to."""
 88
 89    source_definition_id: str | None = None
 90    """The ID of the underlying Airbyte source definition."""
 91
 92    source_definition_name: str | None = None
 93    """The name of the underlying Airbyte source definition, for example `GitHub`."""
 94
 95    context_store_readiness: AgentContextStoreReadiness | None = None
 96    """Context Store readiness information, when reported."""
 97
 98    warnings: list[Any] = Field(default_factory=list)
 99    """Warnings reported by the Agents API, for example degraded capabilities."""
100
101    @property
102    def context_store_entities(self) -> list[str]:
103        """The entity names this connector can cache in the Context Store.
104
105        Note that this lists Context Store-supported entities specifically. The Agents API
106        does not publish an exhaustive list of executable entity and action pairs, so an
107        entity may be executable via `AgentConnector.execute()` without appearing here.
108        """
109        if self.context_store_readiness is None:
110            return []
111        return [
112            entity.entity
113            for entity in self.context_store_readiness.supported_context_store_entities
114        ]

36class AgentConnectorInfo(BaseModel):
37    """Summary information about a connector, as returned by the Agents API."""
38
39    model_config = ConfigDict(extra="allow")
40
41    id: str
42    """The connector ID."""
43
44    name: str | None = None
45    """The connector name, for example `GitHub - <workspace_id>`."""

129class AgentConnectorMetadata(BaseModel):
130    """Connector-reported metadata about a single action's result, including pagination."""
131
132    model_config = ConfigDict(extra="allow")
133
134    has_next_page: bool | None = None
135    """Whether more entities are available after this page, when the connector reports it."""
136
137    end_cursor: str | None = None
138    """The cursor to pass as `cursor` to fetch the next page, when one is available."""

48class AgentContextStoreEntity(BaseModel):
49    """An entity that a connector supports caching in the Airbyte Context Store."""
50
51    model_config = ConfigDict(extra="allow")
52
53    entity: str
54    """The entity name, for example `issues`."""
55
56    suggested: bool | None = None
57    """Whether Airbyte suggests caching this entity."""

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

141class AgentExecuteResult(BaseModel):
142    """The result of executing a single action against an Airbyte Agents connector."""
143
144    model_config = ConfigDict(extra="allow")
145
146    status: str
147    """The execution status reported by the Agents API, for example `success`."""
148
149    result: Any = None
150    """The action's payload. Entity-returning actions put a list of entities here."""
151
152    connector_metadata: AgentConnectorMetadata = Field(default_factory=AgentConnectorMetadata)
153    """Connector-reported metadata about the result, including pagination cursors."""
154
155    execution_metadata: AgentExecutionMetadata = Field(default_factory=AgentExecutionMetadata)
156    """Metadata describing how the action was executed."""
157
158    warning: dict[str, Any] | None = None
159    """A warning reported alongside an otherwise successful result."""
160
161    @property
162    def entities(self) -> list[dict[str, Any]]:
163        """The result as a list of entities.
164
165        Raises `PyAirbyteInputError` if the action did not return a list of entities. Use
166        `result` for actions whose payload is not a list of entities.
167        """
168        if not isinstance(self.result, list):
169            raise PyAirbyteInputError(
170                message="This action did not return a list of entities.",
171                guidance="Use the `result` attribute to read non-entity result payloads.",
172                context={"result_type": type(self.result).__name__},
173            )
174
175        invalid_types = sorted(
176            {type(entity).__name__ for entity in self.result if not isinstance(entity, dict)}
177        )
178        if invalid_types:
179            raise PyAirbyteInputError(
180                message="This action returned a list that is not a list of entities.",
181                guidance="Use the `result` attribute to read non-entity result payloads.",
182                context={"unexpected_item_types": invalid_types},
183            )
184        return self.result
185
186    @property
187    def has_next_page(self) -> bool:
188        """Whether the connector reported more entities after this page."""
189        return bool(self.connector_metadata.has_next_page)
190
191    @property
192    def end_cursor(self) -> str | None:
193        """The cursor for the next page, or `None` when there is no next page."""
194        return self.connector_metadata.end_cursor

117class AgentExecutionMetadata(BaseModel):
118    """Metadata describing how an Agents connector action was executed."""
119
120    model_config = ConfigDict(extra="allow")
121
122    connector_instance_id: str | None = None
123    """The connector instance that served the request."""
124
125    execution_time_ms: int | None = None
126    """The server-side execution time, in milliseconds."""

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