Anthropic Tool Support

package.json

@@ -111,7 +111,7 @@
 		"zod": "^3.22.3"
 	},
 	"optionalDependencies": {
-		"@anthropic-ai/sdk": "^0.25.0",
+		"@anthropic-ai/sdk": "^0.32.1",
 		"@anthropic-ai/vertex-sdk": "^0.4.1",
 		"@aws-sdk/client-bedrock-runtime": "^3.631.0",
 		"@google-cloud/vertexai": "^1.1.0",

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

@@ -3,7 +3,19 @@ import type { Endpoint } from "../endpoints";
 import { env } from "$env/dynamic/private";
 import type { TextGenerationStreamOutput } from "@huggingface/inference";
 import { createImageProcessorOptionsValidator } from "../images";
-import { endpointMessagesToAnthropicMessages } from "./utils";
+import { endpointMessagesToAnthropicMessages, addToolResults } from "./utils";
+import { createDocumentProcessorOptionsValidator } from "../document";
+import type {
+	Tool,
+	ToolCall,
+	ToolInput,
+	ToolInputFile,
+	ToolInputFixed,
+	ToolInputOptional,
+} from "$lib/types/Tool";
+import type Anthropic from "@anthropic-ai/sdk";
+import type { MessageParam } from "@anthropic-ai/sdk/resources/messages.mjs";
+import directlyAnswer from "$lib/server/tools/directlyAnswer";
 
 export const endpointAnthropicParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -23,6 +35,10 @@ export const endpointAnthropicParametersSchema = z.object({
 				maxWidth: 4096,
 				maxHeight: 4096,
 			}),
+			document: createDocumentProcessorOptionsValidator({
+				supportedMimeTypes: ["application/pdf"],
+				maxSizeInMB: 32,
+			}),
 		})
 		.default({}),
 });
@@ -46,7 +62,14 @@ export async function endpointAnthropic(
 		defaultQuery,
 	});
 
-	return async ({ messages, preprompt, generateSettings }) => {
+	return async ({
+		messages,
+		preprompt,
+		generateSettings,
+		conversationId,
+		tools = [],
+		toolResults = [],
+	}) => {
 		let system = preprompt;
 		if (messages?.[0]?.from === "system") {
 			system = messages[0].content;
@@ -59,7 +82,13 @@ export async function endpointAnthropic(
 		return (async function* () {
 			const stream = anthropic.messages.stream({
 				model: model.id ?? model.name,
-				messages: await endpointMessagesToAnthropicMessages(messages, multimodal),
+				tools: createAnthropicTools(tools),
+				tool_choice:
+					tools.length > 0 ? { type: "auto", disable_parallel_tool_use: false } : undefined,
+				messages: addToolResults(
+					await endpointMessagesToAnthropicMessages(messages, multimodal, conversationId),
+					toolResults
+				) as MessageParam[],
 				max_tokens: parameters?.max_new_tokens,
 				temperature: parameters?.temperature,
 				top_p: parameters?.top_p,
@@ -70,21 +99,40 @@ export async function endpointAnthropic(
 			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;
+					if ("tool_use" === stream.receivedMessages[0].stop_reason) {
+						// this should really create a new "Assistant" message with the tool id in it.
+						const toolCalls: ToolCall[] = stream.receivedMessages[0].content
+							.filter(
+								(block): block is Anthropic.Messages.ContentBlock & { type: "tool_use" } =>
+									block.type === "tool_use"
+							)
+							.map((block) => ({
+								name: block.name,
+								parameters: block.input as Record<string, string | number | boolean>,
+								id: block.id,
+							}));
+
+						yield {
+							token: { id: tokenId, text: "", logprob: 0, special: false, toolCalls },
+							generated_text: null,
+							details: null,
+						};
+					} else {
+						yield {
+							token: {
+								id: tokenId++,
+								text: "",
+								logprob: 0,
+								special: true,
+							},
+							generated_text: await stream.finalText(),
+							details: null,
+						} satisfies TextGenerationStreamOutput;
+					}
+
 					return;
 				}
-
 				// Text delta
 				yield {
 					token: {
@@ -100,3 +148,66 @@ export async function endpointAnthropic(
 		})();
 	};
 }
+
+function createAnthropicTools(tools: Tool[]): Anthropic.Messages.Tool[] {
+	return tools
+		.filter((tool) => tool.name !== directlyAnswer.name)
+		.map((tool) => {
+			const properties = tool.inputs.reduce((acc, input) => {
+				acc[input.name] = convertToolInputToJSONSchema(input);
+				return acc;
+			}, {} as Record<string, unknown>);
+
+			const required = tool.inputs
+				.filter((input) => input.paramType === "required")
+				.map((input) => input.name);
+
+			return {
+				name: tool.name,
+				description: tool.description,
+				input_schema: {
+					type: "object",
+					properties,
+					required: required.length > 0 ? required : undefined,
+				},
+			};
+		});
+}
+
+function convertToolInputToJSONSchema(input: ToolInput): Record<string, unknown> {
+	const baseSchema: Record<string, unknown> = {};
+	if ("description" in input) {
+		baseSchema["description"] = input.description || "";
+	}
+	switch (input.paramType) {
+		case "optional":
+			baseSchema["default"] = (input as ToolInputOptional).default;
+			break;
+		case "fixed":
+			baseSchema["const"] = (input as ToolInputFixed).value;
+			break;
+	}
+
+	if (input.type === "file") {
+		baseSchema["type"] = "string";
+		baseSchema["format"] = "binary";
+		baseSchema["mimeTypes"] = (input as ToolInputFile).mimeTypes;
+	} else {
+		switch (input.type) {
+			case "str":
+				baseSchema["type"] = "string";
+				break;
+			case "int":
+				baseSchema["type"] = "integer";
+				break;
+			case "float":
+				baseSchema["type"] = "number";
+				break;
+			case "bool":
+				baseSchema["type"] = "boolean";
+				break;
+		}
+	}
+
+	return baseSchema;
+}

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

@@ -3,6 +3,7 @@ import type { Endpoint } from "../endpoints";
 import type { TextGenerationStreamOutput } from "@huggingface/inference";
 import { createImageProcessorOptionsValidator } from "../images";
 import { endpointMessagesToAnthropicMessages } from "./utils";
+import type { MessageParam } from "@anthropic-ai/sdk/resources/messages.mjs";
 
 export const endpointAnthropicVertexParametersSchema = z.object({
 	weight: z.number().int().positive().default(1),
@@ -56,7 +57,10 @@ export async function endpointAnthropicVertex(
 		return (async function* () {
 			const stream = anthropic.messages.stream({
 				model: model.id ?? model.name,
-				messages: await endpointMessagesToAnthropicMessages(messages, multimodal),
+				messages: (await endpointMessagesToAnthropicMessages(
+					messages,
+					multimodal
+				)) as MessageParam[],
 				max_tokens: model.parameters?.max_new_tokens,
 				temperature: model.parameters?.temperature,
 				top_p: model.parameters?.top_p,

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

@@ -1,13 +1,22 @@
 import { makeImageProcessor, type ImageProcessorOptions } from "../images";
+import { makeDocumentProcessor, type FileProcessorOptions } from "../document";
 import type { EndpointMessage } from "../endpoints";
 import type { MessageFile } from "$lib/types/Message";
-import type { ImageBlockParam, MessageParam } from "@anthropic-ai/sdk/resources/messages.mjs";
+import type {
+	BetaImageBlockParam,
+	BetaMessageParam,
+	BetaBase64PDFBlock,
+} from "@anthropic-ai/sdk/resources/beta/messages/messages.mjs";
+import type { ToolResult } from "$lib/types/Tool";
+import { downloadFile } from "$lib/server/files/downloadFile";
+import type { ObjectId } from "mongodb";
 
 export async function fileToImageBlock(
 	file: MessageFile,
 	opts: ImageProcessorOptions<"image/png" | "image/jpeg" | "image/webp">
-): Promise<ImageBlockParam> {
+): Promise<BetaImageBlockParam> {
 	const processor = makeImageProcessor(opts);
+
 	const { image, mime } = await processor(file);
 
 	return {
@@ -20,25 +29,92 @@ export async function fileToImageBlock(
 	};
 }
 
-type NonSystemMessage = EndpointMessage & { from: "user" | "assistant" };
+export async function fileToDocumentBlock(
+	file: MessageFile,
+	opts: FileProcessorOptions<"application/pdf">
+): Promise<BetaBase64PDFBlock> {
+	const processor = makeDocumentProcessor(opts);
+	const { file: document, mime } = await processor(file);
 
+	return {
+		type: "document",
+		source: {
+			type: "base64",
+			media_type: mime,
+			data: document.toString("base64"),
+		},
+	};
+}
+
+type NonSystemMessage = EndpointMessage & { from: "user" | "assistant" };
 export async function endpointMessagesToAnthropicMessages(
 	messages: EndpointMessage[],
-	multimodal: { image: ImageProcessorOptions<"image/png" | "image/jpeg" | "image/webp"> }
-): Promise<MessageParam[]> {
+	multimodal: {
+		image: ImageProcessorOptions<"image/png" | "image/jpeg" | "image/webp">;
+		document?: FileProcessorOptions<"application/pdf">;
+	},
+	conversationId?: ObjectId | undefined
+): Promise<BetaMessageParam[]> {
 	return await Promise.all(
 		messages
 			.filter((message): message is NonSystemMessage => message.from !== "system")
-			.map<Promise<MessageParam>>(async (message) => {
+			.map<Promise<BetaMessageParam>>(async (message) => {
 				return {
 					role: message.from,
 					content: [
-						...(await Promise.all(
-							(message.files ?? []).map((file) => fileToImageBlock(file, multimodal.image))
-						)),
+						...(message.from === "user"
+							? await Promise.all(
+									(message.files ?? []).map(async (file) => {
+										if (file.type === "hash" && conversationId) {
+											file = await downloadFile(file.value, conversationId);
+										}
+
+										if (file.mime.startsWith("image/")) {
+											return fileToImageBlock(file, multimodal.image);
+										} else if (file.mime === "application/pdf" && multimodal.document) {
+											return fileToDocumentBlock(file, multimodal.document);
+										} else {
+											throw new Error(`Unsupported file type: ${file.mime}`);
+										}
+									})
+							  )
+							: []),
 						{ type: "text", text: message.content },
 					],
 				};
 			})
 	);
 }
+
+export function addToolResults(
+	messages: BetaMessageParam[],
+	toolResults: ToolResult[]
+): BetaMessageParam[] {
+	const id = crypto.randomUUID();
+	if (toolResults.length === 0) {
+		return messages;
+	}
+	return [
+		...messages,
+		{
+			role: "assistant",
+			content: toolResults.map((result, index) => ({
+				type: "tool_use",
+				id: `tool_${index}_${id}`,
+				name: result.call.name,
+				input: result.call.parameters,
+			})),
+		},
+		{
+			role: "user",
+			content: toolResults.map((result, index) => ({
+				type: "tool_result",
+				tool_use_id: `tool_${index}_${id}`,
+				is_error: result.status === "error",
+				content: JSON.stringify(
+					result.status === "error" ? result.message : "outputs" in result ? result.outputs : ""
+				),
+			})),
+		},
+	];
+}

src/lib/server/endpoints/document.ts

@@ -34,23 +34,28 @@ export type DocumentProcessor<TMimeType extends string = string> = (file: Messag
 	mime: TMimeType;
 };
 
+export type AsyncDocumentProcessor<TMimeType extends string = string> = (
+	file: MessageFile
+) => Promise<{
+	file: Buffer;
+	mime: TMimeType;
+}>;
+
 export function makeDocumentProcessor<TMimeType extends string = string>(
 	options: FileProcessorOptions<TMimeType>
-): DocumentProcessor<TMimeType> {
-	return (file) => {
+): AsyncDocumentProcessor<TMimeType> {
+	return async (file) => {
 		const { supportedMimeTypes, maxSizeInMB } = options;
 		const { mime, value } = file;
 
 		const buffer = Buffer.from(value, "base64");
-
 		const tooLargeInBytes = buffer.byteLength > maxSizeInMB * 1000 * 1000;
 
 		if (tooLargeInBytes) {
 			throw Error("Document is too large");
 		}
 
 		const outputMime = validateMimeType(supportedMimeTypes, mime);
-
 		return { file: buffer, mime: outputMime };
 	};
 }