feat: add support for anthropic on vertex (huggingface#958)

Co-authored-by: Nathan Sarrazin <sarrazin.nathan@gmail.com>

Author: cheahjs

README.md

@@ -480,8 +480,8 @@ MODELS=`[
           // optionals
           "apiKey": "sk-ant-...",
           "baseURL": "https://api.anthropic.com",
-          defaultHeaders: {},
-          defaultQuery: {}
+          "defaultHeaders": {},
+          "defaultQuery": {}
         }
       ]
   },
@@ -498,8 +498,51 @@ MODELS=`[
           // optionals
           "apiKey": "sk-ant-...",
           "baseURL": "https://api.anthropic.com",
-          defaultHeaders: {},
-          defaultQuery: {}
+          "defaultHeaders": {},
+          "defaultQuery": {}
+        }
+      ]
+  }
+]`
+```
+
+We also support using Anthropic models running on Vertex AI. Authentication is done using Google Application Default Credentials. Project ID can be provided through the `endpoints.projectId` as per the following example:
+
+```
+MODELS=`[
+  {
+      "name": "claude-3-sonnet@20240229",
+      "displayName": "Claude 3 Sonnet",
+      "description": "Ideal balance of intelligence and speed",
+      "parameters": {
+        "max_new_tokens": 4096,
+      },
+      "endpoints": [
+        {
+          "type": "anthropic-vertex",
+          "region": "us-central1",
+          "projectId": "gcp-project-id",
+          // optionals
+          "defaultHeaders": {},
+          "defaultQuery": {}
+        }
+      ]
+  },
+  {
+      "name": "claude-3-haiku@20240307",
+      "displayName": "Claude 3 Haiku",
+      "description": "Fastest, most compact model for near-instant responsiveness",
+      "parameters": {
+         "max_new_tokens": 4096
+      },
+      "endpoints": [
+        {
+          "type": "anthropic-vertex",
+          "region": "us-central1",
+          "projectId": "gcp-project-id",
+          // optionals
+          "defaultHeaders": {},
+          "defaultQuery": {}
         }
       ]
   }

package-lock.json

@@ -89,6 +89,7 @@
 	},
 	"optionalDependencies": {
 		"@anthropic-ai/sdk": "^0.17.1",
+		"@anthropic-ai/vertex-sdk": "^0.3.0",
 		"@google-cloud/vertexai": "^1.1.0",
 		"aws4fetch": "^1.0.17",
 		"cohere-ai": "^7.9.0",

package.json

@@ -0,0 +1,95 @@
+import { z } from "zod";
+import type { Endpoint } from "../endpoints";
+import type { TextGenerationStreamOutput } from "@huggingface/inference";
+
+export const endpointAnthropicVertexParametersSchema = z.object({
+	weight: z.number().int().positive().default(1),
+	model: z.any(),
+	type: z.literal("anthropic-vertex"),
+	region: z.string().default("us-central1"),
+	projectId: z.string(),
+	defaultHeaders: z.record(z.string()).optional(),
+	defaultQuery: z.record(z.string()).optional(),
+});
+
+export async function endpointAnthropicVertex(
+	input: z.input<typeof endpointAnthropicVertexParametersSchema>
+): Promise<Endpoint> {
+	const { region, projectId, model, defaultHeaders, defaultQuery } =
+		endpointAnthropicVertexParametersSchema.parse(input);
+	let AnthropicVertex;
+	try {
+		AnthropicVertex = (await import("@anthropic-ai/vertex-sdk")).AnthropicVertex;
+	} catch (e) {
+		throw new Error("Failed to import @anthropic-ai/vertex-sdk", { cause: e });
+	}
+
+	const anthropic = new AnthropicVertex({
+		baseURL: `https://${region}-aiplatform.googleapis.com/v1`,
+		region,
+		projectId,
+		defaultHeaders,
+		defaultQuery,
+	});
+
+	return async ({ messages, preprompt }) => {
+		let system = preprompt;
+		if (messages?.[0]?.from === "system") {
+			system = messages[0].content;
+		}
+
+		const messagesFormatted = messages
+			.filter((message) => message.from !== "system")
+			.map((message) => ({
+				role: message.from,
+				content: message.content,
+			})) as unknown as {
+			role: "user" | "assistant";
+			content: string;
+		}[];
+
+		let tokenId = 0;
+		return (async function* () {
+			const stream = anthropic.messages.stream({
+				model: model.id ?? model.name,
+				messages: messagesFormatted,
+				max_tokens: model.parameters?.max_new_tokens,
+				temperature: model.parameters?.temperature,
+				top_p: model.parameters?.top_p,
+				top_k: model.parameters?.top_k,
+				stop_sequences: model.parameters?.stop,
+				system,
+			});
+			while (true) {
+				const result = await Promise.race([stream.emitted("text"), stream.emitted("end")]);
+
+				// Stream end
+				if (result === undefined) {
+					yield {
+						token: {
+							id: tokenId++,
+							text: "",
+							logprob: 0,
+							special: true,
+						},
+						generated_text: await stream.finalText(),
+						details: null,
+					} satisfies TextGenerationStreamOutput;
+					return;
+				}
+
+				// Text delta
+				yield {
+					token: {
+						id: tokenId++,
+						text: result as unknown as string,
+						special: false,
+						logprob: 0,
+					},
+					generated_text: null,
+					details: null,
+				} satisfies TextGenerationStreamOutput;
+			}
+		})();
+	};
+}

src/lib/server/endpoints/anthropic/endpointAnthropicVertex.ts

@@ -12,6 +12,10 @@ import {
 	endpointAnthropic,
 	endpointAnthropicParametersSchema,
 } from "./anthropic/endpointAnthropic";
+import {
+	endpointAnthropicVertex,
+	endpointAnthropicVertexParametersSchema,
+} from "./anthropic/endpointAnthropicVertex";
 import type { Model } from "$lib/types/Model";
 import endpointCloudflare, {
 	endpointCloudflareParametersSchema,
@@ -44,6 +48,7 @@ export type EndpointGenerator<T extends CommonEndpoint> = (parameters: T) => End
 export const endpoints = {
 	tgi: endpointTgi,
 	anthropic: endpointAnthropic,
+	anthropicvertex: endpointAnthropicVertex,
 	aws: endpointAws,
 	openai: endpointOai,
 	llamacpp: endpointLlamacpp,
@@ -56,6 +61,7 @@ export const endpoints = {
 
 export const endpointSchema = z.discriminatedUnion("type", [
 	endpointAnthropicParametersSchema,
+	endpointAnthropicVertexParametersSchema,
 	endpointAwsParametersSchema,
 	endpointOAIParametersSchema,
 	endpointTgiParametersSchema,

src/lib/server/endpoints/endpoints.ts

@@ -159,6 +159,8 @@ const addEndpoint = (m: Awaited<ReturnType<typeof processModel>>) => ({
 						return endpoints.tgi(args);
 					case "anthropic":
 						return endpoints.anthropic(args);
+					case "anthropic-vertex":
+						return endpoints.anthropicvertex(args);
 					case "aws":
 						return await endpoints.aws(args);
 					case "openai":