airbyte.agents.organizations

  1# Copyright (c) 2026 Airbyte, Inc., all rights reserved.
  2"""Airbyte Agents organizations.
  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, NamedTuple
 14
 15from airbyte.agents import _api_util
 16from airbyte.agents.models import AgentWorkspaceInfo
 17from airbyte.agents.workspaces import AgentWorkspace
 18from airbyte.cloud._credentials import _AirbyteCredentials
 19from airbyte.cloud.organizations import CloudOrganization
 20from airbyte.exceptions import AirbyteError, PyAirbyteInputError
 21
 22
 23if TYPE_CHECKING:
 24    from airbyte.secrets.base import SecretString
 25
 26
 27class _WorkspaceLookup(NamedTuple):
 28    """What to look a workspace up by, once the lookup arguments have been validated.
 29
 30    Both fields are set when the caller passed a positional value that could be either an
 31    ID or a name, in which case an ID match takes precedence over a name match.
 32    """
 33
 34    workspace_id: str | None
 35    name: str | None
 36
 37
 38def _resolve_workspace_lookup(
 39    id_or_name: str | None,
 40    /,
 41    *,
 42    workspace_id: str | None,
 43    name: str | None,
 44) -> _WorkspaceLookup:
 45    """Validate workspace lookup arguments and return what to look the workspace up by.
 46
 47    Exactly one of `workspace_id`, `name`, or the positional `id_or_name` is required. A
 48    blank value is rejected rather than treated as an omitted argument.
 49    """
 50    all_args = {"id_or_name": id_or_name, "workspace_id": workspace_id, "name": name}
 51
 52    blank_args = sorted(
 53        key for key, value in all_args.items() if value is not None and not value.strip()
 54    )
 55    if blank_args:
 56        raise PyAirbyteInputError(
 57            message="Workspace lookup arguments cannot be blank.",
 58            guidance="Omit the argument entirely, or pass a non-blank value.",
 59            context={"blank_args": blank_args},
 60        )
 61
 62    provided = sorted(key for key, value in all_args.items() if value)
 63    if len(provided) != 1:
 64        raise PyAirbyteInputError(
 65            message="Exactly one workspace lookup argument is required.",
 66            guidance="Pass a workspace ID or name positionally, or as `workspace_id` or `name`.",
 67            context={"provided": provided},
 68        )
 69
 70    if id_or_name:
 71        return _WorkspaceLookup(workspace_id=id_or_name, name=id_or_name)
 72
 73    return _WorkspaceLookup(workspace_id=workspace_id, name=name)
 74
 75
 76class AgentOrganization:
 77    """An organization on the Airbyte Agents platform.
 78
 79    ```python
 80    from airbyte import agents
 81
 82    organization = agents.AgentOrganization.from_env()
 83    workspace = organization.get_workspace("my-workspace")  # by ID or name (case insensitive)
 84    ```
 85    """
 86
 87    def __init__(
 88        self,
 89        *,
 90        organization_id: str | None = None,
 91        client_id: str | SecretString | None = None,
 92        client_secret: str | SecretString | None = None,
 93        bearer_token: str | SecretString | None = None,
 94    ) -> None:
 95        """Initialize an `AgentOrganization`.
 96
 97        Credentials fall back to the `AIRBYTE_CLOUD_*` environment variables when they are
 98        not passed explicitly. The organization ID is optional: the Agents API infers it
 99        when the credentials belong to exactly one organization.
100        """
101        self._credentials = _AirbyteCredentials.from_auth(
102            organization_id=organization_id,
103            client_id=client_id,
104            client_secret=client_secret,
105            bearer_token=bearer_token,
106            # Mirrors `CloudWorkspace.__init__`: any explicit credential disables env
107            # fallback, since an env bearer token plus explicit client creds is rejected
108            # as mutually exclusive auth.
109            env_vars=not (client_id or client_secret or bearer_token),
110        )
111
112        self.organization_id: str | None = self._credentials.organization_id
113        """The organization ID, when known."""
114
115    @classmethod
116    def from_env(cls, organization_id: str | None = None) -> AgentOrganization:
117        """Create an `AgentOrganization` from the `AIRBYTE_CLOUD_*` environment variables."""
118        return cls(organization_id=organization_id)
119
120    def list_workspaces(self) -> list[AgentWorkspace]:
121        """List the workspaces visible to these credentials in the Agents API."""
122        return [
123            self._workspace_from_info(info)
124            for info in (
125                AgentWorkspaceInfo.model_validate(record)
126                for record in _api_util.list_agent_workspaces(
127                    credentials=self._credentials,
128                    organization_id=self.organization_id,
129                )
130            )
131        ]
132
133    def get_workspace(
134        self,
135        id_or_name: str | None = None,
136        /,
137        *,
138        workspace_id: str | None = None,
139        name: str | None = None,
140    ) -> AgentWorkspace:
141        """Get a workspace in this organization, by ID or by name.
142
143        Pass a single positional value to look the workspace up by either its ID or its
144        name, or name the argument to be explicit.
145
146        Lookup by an explicit `workspace_id` does not call the Agents API. Every other form
147        lists the organization's workspaces and matches on ID first, then on an exact name,
148        ignoring case.
149        """
150        lookup = _resolve_workspace_lookup(
151            id_or_name,
152            workspace_id=workspace_id,
153            name=name,
154        )
155
156        if lookup.workspace_id and not lookup.name:
157            return self._workspace_from_info(AgentWorkspaceInfo(id=lookup.workspace_id))
158
159        workspaces = self.list_workspaces()
160        if lookup.workspace_id:
161            id_matches = [
162                workspace
163                for workspace in workspaces
164                if workspace.workspace_id == lookup.workspace_id
165            ]
166            if id_matches:
167                return id_matches[0]
168
169        name_lower = (lookup.name or "").lower()
170        matches = [
171            workspace
172            for workspace in workspaces
173            if workspace.name and workspace.name.lower() == name_lower
174        ]
175        if not matches:
176            raise AirbyteError(
177                message="No workspace found with the given ID or name.",
178                guidance="Use `list_workspaces()` to see the available workspaces.",
179                context={"name": lookup.name},
180            )
181        if len(matches) > 1:
182            raise AirbyteError(
183                message="Multiple workspaces matched the given name.",
184                guidance="Pass `workspace_id` instead of `name`.",
185                context={"name": lookup.name, "match_count": len(matches)},
186            )
187        return matches[0]
188
189    def as_cloud_organization(self) -> CloudOrganization:
190        """Return this organization as an `airbyte.cloud.CloudOrganization`.
191
192        Every Agents organization is also a Cloud organization, so this conversion needs no
193        API call. It requires a known organization ID, and raises `PyAirbyteInputError`
194        when the ID is unknown.
195        """
196        if not self.organization_id:
197            raise PyAirbyteInputError(
198                message="Organization ID is required to convert to a `CloudOrganization`.",
199                guidance=(
200                    "Provide `organization_id`, or set the `AIRBYTE_CLOUD_ORGANIZATION_ID` "
201                    "environment variable."
202                ),
203            )
204        return CloudOrganization(
205            organization_id=self.organization_id,
206            client_id=self._credentials.client_id,
207            client_secret=self._credentials.client_secret,
208            bearer_token=self._credentials.bearer_token,
209            public_api_root=self._credentials.public_api_root,
210            config_api_root=self._credentials.config_api_root,
211        )
212
213    @classmethod
214    def from_cloud_organization(
215        cls,
216        cloud_organization: CloudOrganization,
217    ) -> AgentOrganization:
218        """Return a Cloud organization as an `AgentOrganization`.
219
220        Whether the organization can actually execute connector actions depends on its
221        Airbyte Agents subscription, which is only knowable per workspace. Use
222        `AgentWorkspace.from_cloud_workspace()` for an authoritative eligibility check.
223
224        Raises `PyAirbyteInputError` when the Cloud organization uses non-public Cloud API
225        roots, since an `AgentOrganization` cannot carry them.
226        """
227        credentials = cloud_organization._credentials  # noqa: SLF001  # Same-domain conversion.
228        _api_util.check_public_cloud_api_roots(credentials)
229        return cls(
230            organization_id=cloud_organization.organization_id,
231            client_id=credentials.client_id,
232            client_secret=credentials.client_secret,
233            bearer_token=credentials.bearer_token,
234        )
235
236    def _workspace_from_info(self, info: AgentWorkspaceInfo) -> AgentWorkspace:
237        """Build an `AgentWorkspace` from workspace info, reusing these credentials."""
238        return AgentWorkspace(
239            workspace_id=info.id,
240            organization_id=info.organization_id or self.organization_id,
241            name=info.name,
242            client_id=self._credentials.client_id,
243            client_secret=self._credentials.client_secret,
244            bearer_token=self._credentials.bearer_token,
245        )