From 35203d69d0628e977876a29b1f145865e7d0af5a Mon Sep 17 00:00:00 2001
From: sulijia <984115358@qq.com>
Date: Wed, 24 Jun 2026 15:34:55 +0800
Subject: [PATCH] Add Public AI (TEE) inference provider

Register `publicai-tee` as a new inference provider backed by an
OpenAI-compatible, TEE-hosted endpoint (https://tee.publicai.io).

- providers/publicai-tee.ts: PublicAITeeConversationalTask (BaseConversationalTask)
- getProviderHelper.ts / types.ts / consts.ts: register the provider
- README.md: list the provider and its supported-models endpoint
- test: add PublicAI TEE conversational integration tests

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 packages/inference/README.md                  |  2 +
 .../inference/src/lib/getProviderHelper.ts    |  4 ++
 packages/inference/src/providers/consts.ts    |  1 +
 .../inference/src/providers/publicai-tee.ts   |  7 +++
 packages/inference/src/types.ts               |  2 +
 .../inference/test/InferenceClient.spec.ts    | 53 +++++++++++++++++++
 6 files changed, 69 insertions(+)
 create mode 100644 packages/inference/src/providers/publicai-tee.ts

diff --git a/packages/inference/README.md b/packages/inference/README.md
index 6f139e87d3..3e30d0eade 100644
--- a/packages/inference/README.md
+++ b/packages/inference/README.md
@@ -58,6 +58,7 @@ Currently, we support the following providers:
 - [NVIDIA](https://build.nvidia.com/)
 - [OVHcloud](https://endpoints.ai.cloud.ovh.net/)
 - [Public AI](https://publicai.co)
+- [Public AI (TEE)](https://publicai.io)
 - [Replicate](https://replicate.com)
 - [Sambanova](https://sambanova.ai)
 - [Scaleway](https://www.scaleway.com/en/generative-apis/)
@@ -99,6 +100,7 @@ Only a subset of models are supported when requesting third-party providers. You
 - [Nscale supported models](https://huggingface.co/api/partners/nscale/models)
 - [NVIDIA supported models](https://huggingface.co/api/partners/nvidia/models)
 - [OVHcloud supported models](https://huggingface.co/api/partners/ovhcloud/models)
+- [Public AI (TEE) supported models](https://huggingface.co/api/partners/publicai-tee/models)
 - [Replicate supported models](https://huggingface.co/api/partners/replicate/models)
 - [Sambanova supported models](https://huggingface.co/api/partners/sambanova/models)
 - [Scaleway supported models](https://huggingface.co/api/partners/scaleway/models)
diff --git a/packages/inference/src/lib/getProviderHelper.ts b/packages/inference/src/lib/getProviderHelper.ts
index dc10b6ba4c..92c7a3af0b 100644
--- a/packages/inference/src/lib/getProviderHelper.ts
+++ b/packages/inference/src/lib/getProviderHelper.ts
@@ -17,6 +17,7 @@ import * as Nvidia from "../providers/nvidia.js";
 import * as OpenAI from "../providers/openai.js";
 import * as OvhCloud from "../providers/ovhcloud.js";
 import * as PublicAI from "../providers/publicai.js";
+import * as PublicAITee from "../providers/publicai-tee.js";
 import type {
 	AudioClassificationTaskHelper,
 	AudioToAudioTaskHelper,
@@ -166,6 +167,9 @@ export const PROVIDERS: Record<InferenceProvider, Partial<Record<InferenceTask,
 	publicai: {
 		conversational: new PublicAI.PublicAIConversationalTask(),
 	},
+	"publicai-tee": {
+		conversational: new PublicAITee.PublicAITeeConversationalTask(),
+	},
 	replicate: {
 		"text-to-image": new Replicate.ReplicateTextToImageTask(),
 		"text-to-speech": new Replicate.ReplicateTextToSpeechTask(),
diff --git a/packages/inference/src/providers/consts.ts b/packages/inference/src/providers/consts.ts
index 8e7a4ab498..4c4c333dff 100644
--- a/packages/inference/src/providers/consts.ts
+++ b/packages/inference/src/providers/consts.ts
@@ -36,6 +36,7 @@ export const HARDCODED_MODEL_INFERENCE_MAPPING: Record<
 	nvidia: {},
 	openai: {},
 	publicai: {},
+	"publicai-tee": {},
 	ovhcloud: {},
 	replicate: {},
 	sambanova: {},
diff --git a/packages/inference/src/providers/publicai-tee.ts b/packages/inference/src/providers/publicai-tee.ts
new file mode 100644
index 0000000000..975344b6da
--- /dev/null
+++ b/packages/inference/src/providers/publicai-tee.ts
@@ -0,0 +1,7 @@
+import { BaseConversationalTask } from "./providerHelper.js";
+
+export class PublicAITeeConversationalTask extends BaseConversationalTask {
+	constructor() {
+		super("publicai-tee", "https://tee.publicai.io");
+	}
+}
diff --git a/packages/inference/src/types.ts b/packages/inference/src/types.ts
index ebd100e8a7..30ee8c5e7b 100644
--- a/packages/inference/src/types.ts
+++ b/packages/inference/src/types.ts
@@ -64,6 +64,7 @@ export const INFERENCE_PROVIDERS = [
 	"openai",
 	"ovhcloud",
 	"publicai",
+	"publicai-tee",
 	"replicate",
 	"sambanova",
 	"scaleway",
@@ -103,6 +104,7 @@ export const PROVIDERS_HUB_ORGS: Record<InferenceProvider, string> = {
 	openai: "openai",
 	ovhcloud: "ovhcloud",
 	publicai: "publicai",
+	"publicai-tee": "publicai-tee",
 	replicate: "replicate",
 	sambanova: "sambanovasystems",
 	scaleway: "scaleway",
diff --git a/packages/inference/test/InferenceClient.spec.ts b/packages/inference/test/InferenceClient.spec.ts
index 5aea175a42..4c7c15d19b 100644
--- a/packages/inference/test/InferenceClient.spec.ts
+++ b/packages/inference/test/InferenceClient.spec.ts
@@ -2611,6 +2611,59 @@ describe.skip("InferenceClient", () => {
 		TIMEOUT,
 	);
 
+	describe.concurrent(
+		"PublicAI TEE",
+		() => {
+			const client = new InferenceClient(env.HF_PUBLICAI_TEE_KEY ?? "dummy");
+
+			HARDCODED_MODEL_INFERENCE_MAPPING["publicai-tee"] = {
+				"MiniMaxAI/MiniMax-M2": {
+					provider: "publicai-tee",
+					hfModelId: "MiniMaxAI/MiniMax-M2",
+					providerId: "MiniMax-M2",
+					status: "live",
+					task: "conversational",
+				},
+			};
+
+			it("chatCompletion", async () => {
+				const res = await client.chatCompletion({
+					model: "MiniMaxAI/MiniMax-M2",
+					provider: "publicai-tee",
+					messages: [{ role: "user", content: "Complete this sentence with words, one plus one is equal " }],
+				});
+				if (res.choices && res.choices.length > 0) {
+					const completion = res.choices[0].message?.content;
+					expect(completion).toContain("two");
+				}
+			});
+
+			it("chatCompletion stream", async () => {
+				const stream = client.chatCompletionStream({
+					model: "MiniMaxAI/MiniMax-M2",
+					provider: "publicai-tee",
+					messages: [{ role: "user", content: "Say 'this is a test'" }],
+					stream: true,
+				}) as AsyncGenerator<ChatCompletionStreamOutput>;
+
+				let fullResponse = "";
+				for await (const chunk of stream) {
+					if (chunk.choices && chunk.choices.length > 0) {
+						const content = chunk.choices[0].delta?.content;
+						if (content) {
+							fullResponse += content;
+						}
+					}
+				}
+
+				// Verify we got a meaningful response
+				expect(fullResponse).toBeTruthy();
+				expect(fullResponse.length).toBeGreaterThan(0);
+			});
+		},
+		TIMEOUT,
+	);
+
 	describe.concurrent(
 		"Baseten",
 		() => {