imbi-plugin-openai 2.44.0__tar.gz

imbi_plugin_openai-2.44.0/.gitignore

@@ -0,0 +1,30 @@
+__pycache__
+*.egg-info
+*.pyc
+*.env
+.env
+.env.test
+build/
+dist/
+docs/site/
+helm/imbi/charts/
+runtime/state/
+.venv/
+.coverage
+.mypy_cache/
+.pytest_cache/
+.ruff_cache/
+coverage.xml
+site/
+
+# moon
+.moon/cache
+.moon/docker
+
+# design-sync
+.design-sync/sb-reference/
+.design-sync/learnings/
+.design-sync/.cache/
+.design-sync/node_modules
+.ds-sync/
+ds-bundle/

imbi_plugin_openai-2.44.0/PKG-INFO

@@ -0,0 +1,54 @@
+Metadata-Version: 2.5
+Name: imbi-plugin-openai
+Version: 2.44.0
+Summary: OpenAI workload identity federation for Imbi
+Author-email: "Gavin M. Roy" <gavinr@aweber.com>
+License-Expression: BSD-3-Clause
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Natural Language :: English
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.14
+Requires-Dist: imbi-common==2.44.0
+Requires-Dist: pydantic>=2
+Description-Content-Type: text/markdown
+
+# imbi-plugin-openai
+
+OpenAI workload identity federation (WIF) for Imbi. The plugin adds:
+
+- The OpenAI token exchanger. An OpenAI AI Provider can then use workload
+  identity instead of a stored key (Admin, AI Models).
+- The `openai_wif` secret provider. A dynamic binding then mints a
+  short-lived OpenAI token for the agent of the task. The harness does
+  not use it yet: the Claude Agent SDK calls only the Anthropic API.
+
+## How it works
+
+1. The organization selects an identity source (for example `aws_sts`) in
+   Admin, Secrets, Identity.
+2. In the OpenAI dashboard, create a workload identity provider for the
+   issuer of that identity source. Set its audience to
+   `https://api.openai.com/v1`. Imbi always requests this audience.
+   See the [OpenAI guide](https://developers.openai.com/api/reference/workload-identity-federation).
+3. Create a service account and a mapping. The mapping matches the role
+   ARN (`sub`) and the request tag `imbi_ai_provider` (an AI Provider) or
+   `imbi_agent` (an agent). AWS request tags are nested in the JWT, so
+   expose them with a CEL transformation, for example
+   `assertion["https://sts.amazonaws.com/"]["request_tags"]["imbi_ai_provider"]`.
+4. Give Imbi the ids: on the AI Provider (Admin, AI Models, the key icon)
+   or in a dynamic `openai_wif` binding.
+5. At use time, Imbi gets a new identity JWT that lasts 55 minutes, and
+   exchanges it at `https://auth.openai.com/oauth/token` (RFC 8693). The
+   token never outlives the JWT.
+
+## Target
+
+| Field                  | Required | Description                               |
+| ---------------------- | -------- | ----------------------------------------- |
+| `identity_provider_id` | yes      | The OpenAI workload identity provider id. |
+| `service_account_id`   | yes      | The OpenAI service account id.            |
+
+The provider takes no configuration.

imbi_plugin_openai-2.44.0/README.md

@@ -0,0 +1,37 @@
+# imbi-plugin-openai
+
+OpenAI workload identity federation (WIF) for Imbi. The plugin adds:
+
+- The OpenAI token exchanger. An OpenAI AI Provider can then use workload
+  identity instead of a stored key (Admin, AI Models).
+- The `openai_wif` secret provider. A dynamic binding then mints a
+  short-lived OpenAI token for the agent of the task. The harness does
+  not use it yet: the Claude Agent SDK calls only the Anthropic API.
+
+## How it works
+
+1. The organization selects an identity source (for example `aws_sts`) in
+   Admin, Secrets, Identity.
+2. In the OpenAI dashboard, create a workload identity provider for the
+   issuer of that identity source. Set its audience to
+   `https://api.openai.com/v1`. Imbi always requests this audience.
+   See the [OpenAI guide](https://developers.openai.com/api/reference/workload-identity-federation).
+3. Create a service account and a mapping. The mapping matches the role
+   ARN (`sub`) and the request tag `imbi_ai_provider` (an AI Provider) or
+   `imbi_agent` (an agent). AWS request tags are nested in the JWT, so
+   expose them with a CEL transformation, for example
+   `assertion["https://sts.amazonaws.com/"]["request_tags"]["imbi_ai_provider"]`.
+4. Give Imbi the ids: on the AI Provider (Admin, AI Models, the key icon)
+   or in a dynamic `openai_wif` binding.
+5. At use time, Imbi gets a new identity JWT that lasts 55 minutes, and
+   exchanges it at `https://auth.openai.com/oauth/token` (RFC 8693). The
+   token never outlives the JWT.
+
+## Target
+
+| Field                  | Required | Description                               |
+| ---------------------- | -------- | ----------------------------------------- |
+| `identity_provider_id` | yes      | The OpenAI workload identity provider id. |
+| `service_account_id`   | yes      | The OpenAI service account id.            |
+
+The provider takes no configuration.

imbi_plugin_openai-2.44.0/moon.yml

@@ -0,0 +1,22 @@
+# yaml-language-server: $schema=https://moonrepo.dev/schemas/project.json
+layer: 'library'
+stack: 'backend'
+language: 'python'
+tags:
+  - 'python'
+
+project:
+  name: 'imbi-plugin-openai'
+  description: 'OpenAI workload identity plugin for Imbi'
+
+dependsOn:
+  - 'common'
+
+# Inherits lint / format / typecheck / test from .moon/tasks/python.yml. The
+# shared test task points --cov-config at this member's pyproject, where
+# [tool.coverage] scopes source + threshold to the member.
+tasks:
+  test:
+    args:
+      - '--cov'
+      - '--cov-report=term-missing'

imbi_plugin_openai-2.44.0/pyproject.toml

@@ -0,0 +1,44 @@
+[project]
+name = "imbi-plugin-openai"
+version = "2.44.0"
+description = "OpenAI workload identity federation for Imbi"
+readme = "README.md"
+requires-python = ">=3.14"
+license = "BSD-3-Clause"
+authors = [
+  { name = "Gavin M. Roy", email = "gavinr@aweber.com" },
+]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Natural Language :: English",
+  "Operating System :: OS Independent",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.14",
+]
+dependencies = [
+  "imbi-common==2.44.0",
+  "pydantic>=2",
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+# Coverage for the per-member `moon run openai:test` run. Scoped to this
+# member so the report/threshold reflect only imbi.plugins.openai; the aggregate
+# `root:coverage` run uses the root pyproject config instead.
+[tool.coverage.run]
+branch = true
+source_pkgs = ["imbi.plugins.openai"]
+
+[tool.coverage.report]
+exclude_also = ["typing.TYPE_CHECKING"]
+fail_under = 85
+show_missing = true
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/imbi"]
+include = [
+  "src/imbi/plugins/openai/py.typed",
+]

imbi_plugin_openai-2.44.0/src/imbi/plugins/openai/__init__.py

@@ -0,0 +1,8 @@
+"""Imbi OpenAI plugin: workload identity federation."""
+
+from imbi.plugins.openai.plugin import OpenAIPlugin
+from imbi.plugins.openai.wif import OpenAIExchanger, OpenAIWif
+
+PLUGIN = OpenAIPlugin
+
+__all__ = ['PLUGIN', 'OpenAIExchanger', 'OpenAIPlugin', 'OpenAIWif']

imbi_plugin_openai-2.44.0/src/imbi/plugins/openai/plugin.py

@@ -0,0 +1,21 @@
+"""The OpenAI plugin manifest.
+
+The plugin has no Integration capability. It adds the ``openai_wif``
+secret provider and the OpenAI token exchanger (for AI Providers).
+"""
+
+import imbi.common.plugins as plugins
+from imbi.plugins.openai import wif
+
+
+class OpenAIPlugin(plugins.Plugin):
+    manifest = plugins.PluginManifest(
+        slug='openai',
+        name='OpenAI',
+        description='OpenAI workload identity federation for model calls',
+        icon='si-openai',
+        auth_type='none',
+        capabilities=[],
+        secret_providers=[wif.OpenAIWif],
+        token_exchangers=[wif.OpenAIExchanger],
+    )

imbi_plugin_openai-2.44.0/src/imbi/plugins/openai/wif.py

@@ -0,0 +1,83 @@
+"""OpenAI workload identity federation (WIF plan 3.3, 4.3, 4.5).
+
+- :class:`OpenAIExchanger` exchanges an identity JWT for an OpenAI
+  access token at a fixed endpoint (D8), with an RFC 8693 token
+  exchange.
+- :class:`OpenAIWif` is the ``openai_wif`` secret provider. A dynamic
+  binding holds the exchanger target. Each resolution mints a new token
+  for the agent of the task through core (``vendor_token``).
+
+The subject token and the access token are never logged.
+"""
+
+import typing
+
+import pydantic
+
+from imbi.common.plugins import identity
+
+#: Fixed in code (D8). No tenant setting can change it.
+TOKEN_URL = 'https://auth.openai.com/oauth/token'
+_GRANT_TYPE = 'urn:ietf:params:oauth:grant-type:token-exchange'
+_SUBJECT_TOKEN_TYPE = 'urn:ietf:params:oauth:token-type:jwt'
+
+_Id = typing.Annotated[
+    str,
+    pydantic.StringConstraints(
+        min_length=1, max_length=128, pattern=r'^[A-Za-z0-9_-]+$'
+    ),
+]
+
+
+class _Target(pydantic.BaseModel):
+    """The OpenAI ids of one workload identity provider and account."""
+
+    model_config = pydantic.ConfigDict(extra='forbid')
+
+    identity_provider_id: _Id
+    service_account_id: _Id
+
+
+class OpenAIExchanger(identity.TokenExchanger):
+    vendor = 'openai'
+    label = 'OpenAI'
+    target_help = (
+        f'{identity.target_keys(_Target)}, the ids from the OpenAI dashboard.'
+    )
+    #: The audience of the OpenAI workload identity provider. The
+    #: provider in the OpenAI dashboard must use this value.
+    audience = 'https://api.openai.com/v1'
+    #: An OpenAI token never outlives its subject token. 3300 seconds
+    #: keeps the ``aws_sts`` role session at 3600 seconds, the AWS
+    #: limit for role chaining (imbi-api on IRSA or EKS Pod Identity).
+    jwt_seconds = 3300
+
+    def validate_target(
+        self, target: dict[str, typing.Any]
+    ) -> dict[str, typing.Any]:
+        return identity.validate_model(_Target, target)
+
+    async def exchange(
+        self, target: dict[str, typing.Any], identity_token: str
+    ) -> identity.AccessToken:
+        return await identity.post_exchange(
+            'OpenAI',
+            TOKEN_URL,
+            {
+                'grant_type': _GRANT_TYPE,
+                'subject_token': identity_token,
+                'subject_token_type': _SUBJECT_TOKEN_TYPE,
+                **target,
+            },
+        )
+
+
+EXCHANGER = OpenAIExchanger()
+
+
+class OpenAIWif(identity.WifSecretProvider):
+    """``openai_wif``: an OpenAI token for the agent of the task."""
+
+    type = 'openai_wif'
+    label = 'OpenAI workload identity'
+    exchanger = EXCHANGER

imbi_plugin_openai-2.44.0/tests/test_wif.py

@@ -0,0 +1,160 @@
+"""Tests for the OpenAI exchanger and the ``openai_wif`` provider.
+
+The token endpoint is an httpx mock transport: no call leaves the test.
+The error mapping is tested with the Anthropic exchanger, which uses
+the same :func:`identity.post_exchange`.
+"""
+
+import datetime
+import functools
+import json
+import unittest
+import uuid
+from unittest import mock
+
+import httpx
+
+from imbi.common.plugins import identity, secrets
+from imbi.plugins.openai import plugin, wif
+
+TARGET = {'identity_provider_id': 'idp_1', 'service_account_id': 'sa_1'}
+
+
+class ExchangerTestCase(unittest.IsolatedAsyncioTestCase):
+    def setUp(self) -> None:
+        self.requests: list[httpx.Request] = []
+        self.reply = httpx.Response(
+            200,
+            json={
+                'access_token': 'openai-token',
+                'issued_token_type': (
+                    'urn:ietf:params:oauth:token-type:access_token'
+                ),
+                'token_type': 'Bearer',
+                'expires_in': 3600,
+                'expires_at': 1_800_000_000,
+            },
+        )
+        patcher = mock.patch.object(
+            identity.httpx,
+            'AsyncClient',
+            functools.partial(
+                httpx.AsyncClient, transport=httpx.MockTransport(self._handle)
+            ),
+        )
+        patcher.start()
+        self.addCleanup(patcher.stop)
+
+    def _handle(self, request: httpx.Request) -> httpx.Response:
+        self.requests.append(request)
+        return self.reply
+
+    async def test_exchange(self) -> None:
+        before = datetime.datetime.now(datetime.UTC)
+        token = await wif.EXCHANGER.exchange(TARGET, 'eyJ.jwt')
+        self.assertEqual(token.token, 'openai-token')
+        self.assertIsNone(token.scope)
+        self.assertGreaterEqual(
+            token.expires_at, before + datetime.timedelta(seconds=3600)
+        )
+        request = self.requests[0]
+        self.assertEqual(str(request.url), wif.TOKEN_URL)
+        self.assertEqual(request.headers['content-type'], 'application/json')
+        self.assertEqual(
+            json.loads(request.content),
+            {
+                'grant_type': (
+                    'urn:ietf:params:oauth:grant-type:token-exchange'
+                ),
+                'subject_token': 'eyJ.jwt',
+                'subject_token_type': 'urn:ietf:params:oauth:token-type:jwt',
+                **TARGET,
+            },
+        )
+
+    async def test_refused(self) -> None:
+        self.reply = httpx.Response(400, json={'error': 'invalid_grant'})
+        with self.assertRaises(identity.IdentityError) as ctx:
+            await wif.EXCHANGER.exchange(TARGET, 'eyJ.jwt')
+        self.assertEqual(ctx.exception.reason, 'exchange_refused')
+
+    def test_target(self) -> None:
+        self.assertEqual(wif.EXCHANGER.validate_target(TARGET), TARGET)
+        for bad in (
+            {**TARGET, 'extra': 'x'},
+            {'identity_provider_id': 'idp_1'},
+            {**TARGET, 'service_account_id': '../x'},
+        ):
+            with self.assertRaises(identity.InvalidConfiguration):
+                wif.EXCHANGER.validate_target(bad)
+        # The help lists the keys of the target model.
+        self.assertTrue(
+            wif.EXCHANGER.target_help.startswith(
+                json.dumps(dict.fromkeys(TARGET, '...'))
+            )
+        )
+
+    def test_claims(self) -> None:
+        self.assertEqual(wif.EXCHANGER.audience, 'https://api.openai.com/v1')
+        self.assertEqual(wif.EXCHANGER.jwt_seconds, 3300)
+
+
+class ProviderTestCase(unittest.IsolatedAsyncioTestCase):
+    def setUp(self) -> None:
+        self.host = mock.Mock()
+        self.expires = datetime.datetime.now(datetime.UTC)
+        self.host.vendor_token = mock.AsyncMock(
+            return_value=identity.AccessToken('openai-token', self.expires)
+        )
+        self.provider = wif.OpenAIWif(self.host)
+
+    def _context(self, agent_id: str | None) -> secrets.ResolveContext:
+        return secrets.ResolveContext(
+            organization_id='org',
+            task_id=uuid.uuid4(),
+            project_id=None,
+            purpose='model_inference',
+            ttl_seconds=900,
+            credential_class='dynamic',
+            agent_id=agent_id,
+        )
+
+    async def test_resolve_mints_for_the_agent(self) -> None:
+        material = await self.provider.resolve(
+            {}, TARGET, None, self._context('agent-1')
+        )
+        self.assertEqual(material.value, b'openai-token')
+        self.assertEqual(material.credential_expires_at, self.expires)
+        self.assertTrue(material.bearer)
+        self.host.vendor_token.assert_awaited_once_with(
+            'org', wif.EXCHANGER, TARGET, agent_id='agent-1'
+        )
+
+    async def test_resolve_refusals(self) -> None:
+        for field, agent_id, code in (
+            ('api-key', 'agent-1', 'secret_field_not_allowed'),
+            (None, None, 'secret_access_denied'),
+        ):
+            with self.subTest(code=code):
+                with self.assertRaises(secrets.BrokerError) as ctx:
+                    await self.provider.resolve(
+                        {}, TARGET, field, self._context(agent_id)
+                    )
+                self.assertEqual(ctx.exception.code, code)
+        self.host.vendor_token.assert_not_awaited()
+
+    def test_config_and_locator(self) -> None:
+        self.assertEqual(self.provider.validate_config({}), {})
+        with self.assertRaises(secrets.InvalidConfiguration):
+            self.provider.validate_config({'api_key': 'x'})
+        self.assertEqual(
+            self.provider.validate_locator(TARGET, 'dynamic'), TARGET
+        )
+        with self.assertRaises(secrets.InvalidConfiguration):
+            self.provider.validate_locator(TARGET, 'static')
+
+    def test_manifest(self) -> None:
+        manifest = plugin.OpenAIPlugin.manifest
+        self.assertEqual(manifest.capabilities, [])
+        self.assertEqual(manifest.secret_providers, [wif.OpenAIWif])
+        self.assertEqual(manifest.token_exchangers, [wif.OpenAIExchanger])