From 91890adfef9d60702f31f8ad76385bc2c1fd01b1 Mon Sep 17 00:00:00 2001 From: Advait Jayant Date: Mon, 26 Oct 2026 20:02:42 +0000 Subject: [PATCH] Add abliteration.ai provider and Abliterated Large v2 --- .env.example | 1 + .github/workflows/test.yml | 3 +- CLAUDE.md | 2 ++ README.md | 5 ++- scripts/run-enclave.sh | 5 ++- tee_gateway/__main__.py | 6 ++++ tee_gateway/config.py | 3 ++ tee_gateway/llm_backend.py | 36 +++++++++++++++++++ tee_gateway/model_registry.py | 20 ++++++++++- .../test/test_provider_usage_integration.py | 4 +++ tests/test_pricing.py | 26 ++++++++++++++ 11 files changed, 107 insertions(+), 4 deletions(-) diff --git a/.env.example b/.env.example index 3d0aa0b..2ffed49 100644 --- a/.env.example +++ b/.env.example @@ -10,6 +10,7 @@ ANTHROPIC_API_KEY= XAI_API_KEY= ARK_API_KEY= OPENROUTER_API_KEY= +ABLITERATION_API_KEY= WAVESPEED_API_KEY= # Blockchain heartbeat configuration (optional). diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 41ecb61..246bffc 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -41,6 +41,7 @@ jobs: # accept either secret name so the job runs whichever the repo still holds. OPENROUTER_API_KEY: ${{ secrets.OPENROUTER_API_KEY || secrets.NOUS_API_KEY }} ZAI_API_KEY: ${{ secrets.ZAI_API_KEY }} + ABLITERATION_API_KEY: ${{ secrets.ABLITERATION_API_KEY }} WAVESPEED_API_KEY: ${{ secrets.WAVESPEED_API_KEY }} steps: - uses: actions/checkout@v4 @@ -50,7 +51,7 @@ jobs: run: | for name in \ OPENAI_API_KEY ANTHROPIC_API_KEY GOOGLE_API_KEY XAI_API_KEY \ - ARK_API_KEY OPENROUTER_API_KEY ZAI_API_KEY WAVESPEED_API_KEY + ARK_API_KEY OPENROUTER_API_KEY ZAI_API_KEY ABLITERATION_API_KEY WAVESPEED_API_KEY do if [[ -z "${!name}" ]]; then echo "Missing required repository secret: $name" >&2 diff --git a/CLAUDE.md b/CLAUDE.md index 63fd030..07ad917 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -75,6 +75,7 @@ API keys (injected at runtime via POST /v1/keys — do NOT bake into the image): - `ARK_API_KEY` (BytePlus / ByteDance ModelArk; injected as `bytedance_api_key`) - `OPENROUTER_API_KEY` (OpenRouter; injected as `openrouter_api_key`) - `ZAI_API_KEY` (Z.ai Model API; injected as `zai_api_key`) +- `ABLITERATION_API_KEY` (abliteration.ai; injected as `abliteration_api_key`) - `WAVESPEED_API_KEY` (WaveSpeed; injected as `wavespeed_api_key`) - `EXA_API_KEY` (Exa search; injected as `exa_api_key`) — backs the in-enclave `/v1/web_search` endpoint, not an LLM provider. Without it the endpoint @@ -133,6 +134,7 @@ prefixes determine routing: - **xAI**: grok-4.7, grok-4.6, grok-4.5, grok-4.3, grok-4.20-reasoning, grok-4.20-non-reasoning; image generation: grok-2-image, grok-imagine-image-2.0 - **ByteDance** (BytePlus ModelArk, OpenAI-compatible, ap-southeast): seed-1.6, seed-1.8, seed-2.0-lite, deepseek-v4-flash, deepseek-v4-pro, glm-5.2 (Z.ai's model served via a ModelArk deployment endpoint); image generation: seedream-4.0, seedream-5.0-lite, seedance-4.5, seedance-5.0 - **OpenRouter** (OpenAI-compatible): hy4-preview, hermes-4-405b, hy3 +- **abliteration.ai** (OpenAI-compatible): abliterated-model-large-v2 - **Z.ai** (Model API, OpenAI-compatible): image generation: glm-image (glm-5.2 chat is routed through BytePlus ModelArk, see ByteDance above) - **WaveSpeed** (async prediction API): image generation: krea-3 (Krea's model, always `berry` at `reasoning_level: medium`), qwen-image-3.0-pro (Alibaba's model); each is two WaveSpeed endpoints, text-to-image and edit diff --git a/README.md b/README.md index d2090a0..8940230 100644 --- a/README.md +++ b/README.md @@ -18,7 +18,7 @@ The gateway solves this by running inside a hardware-isolated Nitro Enclave wher ## Features -- **Multi-provider routing** - OpenAI, Anthropic, Google Gemini, xAI Grok, ByteDance (BytePlus ModelArk), OpenRouter +- **Multi-provider routing** - OpenAI, Anthropic, Google Gemini, xAI Grok, ByteDance (BytePlus ModelArk), OpenRouter, abliteration.ai - **Remote attestation** - AWS Nitro attestation documents with PCR measurements - **Response signing** - RSA-PSS signatures on all inference results - **Request integrity** - SHA256 hash of original request included in signed response @@ -39,6 +39,7 @@ The gateway solves this by running inside a hardware-isolated Nitro Enclave wher | xAI | grok-4.7, grok-4.6, grok-4.5, grok-4.3, grok-4.20-reasoning, grok-4.20-non-reasoning | | ByteDance | seed-1.6, seed-1.8, seed-2.0-lite, dola-seed-2.0-mini, deepseek-v4-flash, deepseek-v4-pro, glm-5.2 | | OpenRouter | hy4-preview, hermes-4-405b, hy3 | +| abliteration.ai | abliterated-model-large-v2 | Image generation: krea-3 and qwen-image-3.0-pro (via WaveSpeed), gpt-image-2, gpt-image-2.5-flare, gpt-image-2.5-sunburst, gemini-2.5-flash-image, @@ -65,6 +66,7 @@ export XAI_API_KEY=... export ARK_API_KEY=... # BytePlus / ByteDance ModelArk export OPENROUTER_API_KEY=... # OpenRouter export ZAI_API_KEY=... # Z.ai Model API +export ABLITERATION_API_KEY=... # abliteration.ai export WAVESPEED_API_KEY=... # WaveSpeed (krea-3, qwen-image-3.0-pro) # Run server (starts the Flask/connexion app on port 8000) @@ -419,6 +421,7 @@ Clients use an x402-compatible client (e.g., the [x402 SDK](https://github.com/o | `ARK_API_KEY` | - | BytePlus / ByteDance ModelArk API key (injected as `bytedance_api_key`) | | `OPENROUTER_API_KEY` | - | OpenRouter API key (injected as `openrouter_api_key`) | | `ZAI_API_KEY` | - | Z.ai Model API key (injected as `zai_api_key`) | +| `ABLITERATION_API_KEY` | - | abliteration.ai API key (injected as `abliteration_api_key`) | | `WAVESPEED_API_KEY` | - | WaveSpeed API key (injected as `wavespeed_api_key`) | | `EVM_PAYMENT_ADDRESS` | - | Wallet address to receive x402 payments | | `FACILITATOR_URL` | see `tee_gateway/__main__.py` | x402 payment facilitator endpoint | diff --git a/scripts/run-enclave.sh b/scripts/run-enclave.sh index 7e74dda..afe8502 100755 --- a/scripts/run-enclave.sh +++ b/scripts/run-enclave.sh @@ -92,6 +92,7 @@ if [ -f "$ENV_FILE" ]; then ARK_API_KEY="$(grep -E '^ARK_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" OPENROUTER_API_KEY="$(grep -E '^OPENROUTER_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" ZAI_API_KEY="$(grep -E '^ZAI_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" + ABLITERATION_API_KEY="$(grep -E '^ABLITERATION_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" WAVESPEED_API_KEY="$(grep -E '^WAVESPEED_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" # Backs the in-enclave `web_search` tool (Exa), not an LLM provider. EXA_API_KEY="$(grep -E '^EXA_API_KEY=' "$ENV_FILE" | cut -d'=' -f2-)" @@ -113,6 +114,7 @@ if [ -f "$ENV_FILE" ]; then --arg bytedance "$ARK_API_KEY" \ --arg openrouter "$OPENROUTER_API_KEY" \ --arg zai "$ZAI_API_KEY" \ + --arg abliteration "$ABLITERATION_API_KEY" \ --arg wavespeed "$WAVESPEED_API_KEY" \ --arg exa "$EXA_API_KEY" \ --arg hb_contract "$HEARTBEAT_CONTRACT_ADDRESS" \ @@ -126,6 +128,7 @@ if [ -f "$ENV_FILE" ]; then bytedance_api_key: $bytedance, openrouter_api_key: $openrouter, zai_api_key: $zai, + abliteration_api_key: $abliteration, wavespeed_api_key: $wavespeed, exa_api_key: $exa } @@ -158,7 +161,7 @@ if [ -f "$ENV_FILE" ]; then # Clear key variables from this shell immediately after use unset OPENAI_API_KEY GOOGLE_API_KEY ANTHROPIC_API_KEY XAI_API_KEY ARK_API_KEY - unset OPENROUTER_API_KEY ZAI_API_KEY WAVESPEED_API_KEY EXA_API_KEY + unset OPENROUTER_API_KEY ZAI_API_KEY ABLITERATION_API_KEY WAVESPEED_API_KEY EXA_API_KEY unset HEARTBEAT_CONTRACT_ADDRESS FACILITATOR_URL TEE_HEARTBEAT_INTERVAL fi else diff --git a/tee_gateway/__main__.py b/tee_gateway/__main__.py index 23ac5bf..b567544 100644 --- a/tee_gateway/__main__.py +++ b/tee_gateway/__main__.py @@ -429,6 +429,7 @@ def set_provider_keys(): bytedance_api_key=body.get("bytedance_api_key") or None, openrouter_api_key=body.get("openrouter_api_key") or None, zai_api_key=body.get("zai_api_key") or None, + abliteration_api_key=body.get("abliteration_api_key") or None, wavespeed_api_key=body.get("wavespeed_api_key") or None, exa_api_key=body.get("exa_api_key") or None, ) @@ -497,6 +498,10 @@ def _set(val: str | None) -> str: logger.info( " zai_api_key : %s", _set(provider_config.zai_api_key) ) + logger.info( + " abliteration_api_key : %s", + _set(provider_config.abliteration_api_key), + ) logger.info( " wavespeed_api_key : %s", _set(provider_config.wavespeed_api_key), @@ -538,6 +543,7 @@ def _set(val: str | None) -> str: "bytedance": provider_config.bytedance_api_key, "openrouter": provider_config.openrouter_api_key, "zai": provider_config.zai_api_key, + "abliteration": provider_config.abliteration_api_key, "wavespeed": provider_config.wavespeed_api_key, }.items() if k diff --git a/tee_gateway/config.py b/tee_gateway/config.py index 772f88b..f6b6433 100644 --- a/tee_gateway/config.py +++ b/tee_gateway/config.py @@ -30,6 +30,7 @@ class ProviderConfig: bytedance_api_key: Optional[str] = None openrouter_api_key: Optional[str] = None zai_api_key: Optional[str] = None + abliteration_api_key: Optional[str] = None wavespeed_api_key: Optional[str] = None # Exa, which backs the `web_search` tool the gateway executes inside the # enclave for every model (see web_search.py). Not an LLM provider, so it is @@ -54,6 +55,8 @@ def initialized_providers(self) -> list[str]: providers.append("openrouter") if self.zai_api_key: providers.append("zai") + if self.abliteration_api_key: + providers.append("abliteration") if self.wavespeed_api_key: providers.append("wavespeed") return providers diff --git a/tee_gateway/llm_backend.py b/tee_gateway/llm_backend.py index 6ea591f..9853971 100644 --- a/tee_gateway/llm_backend.py +++ b/tee_gateway/llm_backend.py @@ -62,6 +62,11 @@ # the subscription Coding Plan endpoint at /api/coding/paas/v4. ZAI_BASE_URL = "https://api.z.ai/api/paas/v4" +# abliteration.ai OpenAI-compatible endpoint; ChatOpenAI appends +# /chat/completions. Streams report usage (incl. reasoning tokens) in a final +# chunk when stream_options.include_usage is set, which stream_usage does. +ABLITERATION_BASE_URL = "https://api.abliteration.ai/v1" + # WaveSpeed REST API. Image models are submitted to /{model_id} as async # predictions and polled at /predictions/{id}/result (see image_generation). WAVESPEED_BASE_URL = "https://api.wavespeed.ai/api/v3" @@ -73,6 +78,7 @@ bytedance_http_client: Optional[httpx.Client] = None openrouter_http_client: Optional[httpx.Client] = None zai_http_client: Optional[httpx.Client] = None +abliteration_http_client: Optional[httpx.Client] = None wavespeed_http_client: Optional[httpx.Client] = None @@ -83,12 +89,14 @@ def set_provider_config(config: ProviderConfig) -> None: """Store the provider config and rebuild HTTP clients. Called once after key injection.""" global _provider_config, openai_http_client, xai_http_client, bytedance_http_client global openrouter_http_client, zai_http_client, wavespeed_http_client + global abliteration_http_client old_openai = openai_http_client old_xai = xai_http_client old_bytedance = bytedance_http_client old_openrouter = openrouter_http_client old_zai = zai_http_client + old_abliteration = abliteration_http_client old_wavespeed = wavespeed_http_client openai_http_client = httpx.Client( @@ -138,6 +146,14 @@ def set_provider_config(config: ProviderConfig) -> None: http2=True, follow_redirects=False, ) + abliteration_http_client = httpx.Client( + base_url=ABLITERATION_BASE_URL, + headers={"Authorization": f"Bearer {config.abliteration_api_key or ''}"}, + timeout=_TIMEOUT, + limits=_LIMITS, + http2=True, + follow_redirects=False, + ) wavespeed_http_client = httpx.Client( base_url=WAVESPEED_BASE_URL, headers={"Authorization": f"Bearer {config.wavespeed_api_key or ''}"}, @@ -168,6 +184,8 @@ def set_provider_config(config: ProviderConfig) -> None: old_openrouter.close() if old_zai is not None: old_zai.close() + if old_abliteration is not None: + old_abliteration.close() if old_wavespeed is not None: old_wavespeed.close() @@ -390,6 +408,24 @@ def get_chat_model_cached( stream_usage=True, ) # type: ignore [call-arg] + elif provider == "abliteration": + if not config.abliteration_api_key: + raise ValueError("abliteration_api_key not set in ProviderConfig") + + if abliteration_http_client is None: + raise RuntimeError("abliteration.ai HTTP client has not been initialized") + + return ChatOpenAI( + model=api_name, + temperature=effective_temp, + max_tokens=max_tokens, + http_client=abliteration_http_client, + api_key=SecretStr(config.abliteration_api_key), + base_url=ABLITERATION_BASE_URL, + streaming=True, + stream_usage=True, + ) # type: ignore [call-arg] + else: raise ValueError(f"Unsupported provider: {provider}") diff --git a/tee_gateway/model_registry.py b/tee_gateway/model_registry.py index 53a4553..3b21f97 100644 --- a/tee_gateway/model_registry.py +++ b/tee_gateway/model_registry.py @@ -43,7 +43,7 @@ class ImageResolutionTier: @dataclass(frozen=True) class ModelConfig: # "openai" | "anthropic" | "google" | "x-ai" | "bytedance" | - # "openrouter" | "zai" | "wavespeed" + # "openrouter" | "zai" | "abliteration" | "wavespeed" provider: str api_name: str # model name sent to provider API input_price_usd: Decimal # USD per token @@ -1176,6 +1176,20 @@ class SupportedModel(Enum): image_aspect_ratio_param="size", ) + # ── abliteration.ai (OpenAI-compatible) ───────────────────────────── + # Abliterated Large v2: 1M-token context, text input, a reasoning model + # (thinking arrives in `reasoning`, counted in reasoning_tokens). Prices + # from abliteration.ai/pricing and its /v1/models: $3.00 input, $5.00 + # output per 1M tokens. /v1/models lists temperature among the supported + # sampling parameters (0-2), and temperature 0, tools and streamed usage + # were confirmed against the live API, so no restriction flags. + ABLITERATED_LARGE_V2 = ModelConfig( + provider="abliteration", + api_name="abliterated-model-large-v2", + input_price_usd=Decimal("0.000003"), + output_price_usd=Decimal("0.000005"), + ) + # ── WaveSpeed (async prediction API) ──────────────────────────────── # Krea 3 — Krea's image editing model, served through WaveSpeed as two # model endpoints (text-to-image, and edit with 1-8 references). Every @@ -1408,6 +1422,10 @@ class SupportedModel(Enum): "glm-5.2": SupportedModel.GLM_5_2, "ep-20260803211658-fwpzs": SupportedModel.GLM_5_2, "glm-image": SupportedModel.GLM_IMAGE, + # abliteration.ai + "abliterated-model-large-v2": SupportedModel.ABLITERATED_LARGE_V2, + "abliterated-large-v2": SupportedModel.ABLITERATED_LARGE_V2, + "abliteration/abliterated-model-large-v2": SupportedModel.ABLITERATED_LARGE_V2, # WaveSpeed "krea-3": SupportedModel.KREA_3, "krea-v3": SupportedModel.KREA_3, diff --git a/tee_gateway/test/test_provider_usage_integration.py b/tee_gateway/test/test_provider_usage_integration.py index 144e5d6..2b902be 100644 --- a/tee_gateway/test/test_provider_usage_integration.py +++ b/tee_gateway/test/test_provider_usage_integration.py @@ -56,6 +56,9 @@ class _ProviderCase: _ProviderCase("ByteDance", "deepseek-v4-flash", "ARK_API_KEY"), _ProviderCase("OpenRouter", "hy3", "OPENROUTER_API_KEY"), _ProviderCase("Z.ai", "glm-5.2", "ZAI_API_KEY"), + _ProviderCase( + "abliteration.ai", "abliterated-model-large-v2", "ABLITERATION_API_KEY" + ), ) # The newest model per provider — one each, so the run stays cheap. Replace a @@ -259,6 +262,7 @@ def setUpClass(cls) -> None: bytedance_api_key=os.environ["ARK_API_KEY"], openrouter_api_key=os.environ["OPENROUTER_API_KEY"], zai_api_key=os.environ["ZAI_API_KEY"], + abliteration_api_key=os.environ["ABLITERATION_API_KEY"], wavespeed_api_key=os.environ["WAVESPEED_API_KEY"], ) ) diff --git a/tests/test_pricing.py b/tests/test_pricing.py index 27b154e..d70debd 100644 --- a/tests/test_pricing.py +++ b/tests/test_pricing.py @@ -718,6 +718,25 @@ def test_glm_image_resolves(self): self.assertTrue(cfg.image_generation) self.assertEqual(cfg.per_image_price_usd, Decimal("0.015")) + # ── abliteration.ai ────────────────────────────────────────────────────── + + def test_abliterated_large_v2_resolves(self): + cfg = get_model_config("abliterated-model-large-v2") + self.assertEqual(cfg.provider, "abliteration") + self.assertEqual(cfg.api_name, "abliterated-model-large-v2") + self.assertEqual(cfg.input_price_usd, Decimal("0.000003")) + self.assertEqual(cfg.output_price_usd, Decimal("0.000005")) + + def test_abliterated_large_v2_aliases_resolve(self): + for alias in ( + "abliterated-large-v2", + "abliteration/abliterated-model-large-v2", + ): + self.assertEqual( + get_model_config(alias), + get_model_config("abliterated-model-large-v2"), + ) + # ── WaveSpeed ──────────────────────────────────────────────────────────── def test_qwen_image_3_0_pro_resolves(self): @@ -1200,6 +1219,13 @@ def test_hy3_cost(self): 396_000_000_000_000, ) + def test_abliterated_large_v2_cost(self): + # 1000*0.000003 + 500*0.000005 = 0.003 + 0.0025 = 0.0055 USD + self.assertEqual( + self._calc("abliterated-model-large-v2", 1000, 500), + 5_500_000_000_000_000, + ) + def test_hy4_preview_cost(self): self.assertEqual( self._calc("hy4-preview", 1000, 500),