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