fixes #2663: validate AI output for test connection endpoint (#2674)

src/integrations/orpc/client.ts

@@ -1,4 +1,4 @@
-import { createORPCClient, onError } from "@orpc/client";
+import { createORPCClient } from "@orpc/client";
 import { RPCLink } from "@orpc/client/fetch";
 import { BatchLinkPlugin } from "@orpc/client/plugins";
 import { createRouterClient, type InferRouterInputs, type InferRouterOutputs, type RouterClient } from "@orpc/server";
@@ -11,11 +11,11 @@ import { getLocale } from "@/utils/locale";
 export const getORPCClient = createIsomorphicFn()
 	.server((): RouterClient<typeof router> => {
 		return createRouterClient(router, {
-			interceptors: [
-				onError((error) => {
-					console.error(error);
-				}),
-			],
+			// interceptors: [
+			// 	onError((error) => {
+			// 		console.error(error);
+			// 	}),
+			// ],
 			context: async () => {
 				const locale = await getLocale();
 				const reqHeaders = getRequestHeaders();
@@ -33,12 +33,12 @@ export const getORPCClient = createIsomorphicFn()
 			fetch: (request, init) => {
 				return fetch(request, { ...init, credentials: "include" });
 			},
-			interceptors: [
-				onError((error) => {
-					if (error instanceof DOMException) return;
-					console.error(error);
-				}),
-			],
+			// interceptors: [
+			// 	onError((error) => {
+			// 		if (error instanceof DOMException) return;
+			// 		console.error(error);
+			// 	}),
+			// ],
 			plugins: [new BatchLinkPlugin({ groups: [{ condition: () => true, context: {} }] })],
 		});
 

src/integrations/orpc/router/ai.ts

@@ -1,52 +1,8 @@
-import { createAnthropic } from "@ai-sdk/anthropic";
-import { createGoogleGenerativeAI } from "@ai-sdk/google";
-import { createOpenAI } from "@ai-sdk/openai";
-import { createGateway, generateText, Output, streamText } from "ai";
-import { createOllama } from "ai-sdk-ollama";
-import { match } from "ts-pattern";
-import z, { flattenError, ZodError } from "zod";
-import docxParserSystemPrompt from "@/integrations/ai/prompts/docx-parser-system.md?raw";
-import docxParserUserPrompt from "@/integrations/ai/prompts/docx-parser-user.md?raw";
-import pdfParserSystemPrompt from "@/integrations/ai/prompts/pdf-parser-system.md?raw";
-import pdfParserUserPrompt from "@/integrations/ai/prompts/pdf-parser-user.md?raw";
-import { defaultResumeData, resumeDataSchema } from "@/schema/resume/data";
+import { ORPCError } from "@orpc/client";
+import { AISDKError } from "ai";
+import z, { ZodError } from "zod";
 import { protectedProcedure } from "../context";
-
-const aiProviderSchema = z.enum(["ollama", "openai", "gemini", "anthropic", "vercel-ai-gateway"]);
-
-type AIProvider = z.infer<typeof aiProviderSchema>;
-
-type GetModelInput = {
-	provider: AIProvider;
-	model: string;
-	apiKey: string;
-	baseURL: string;
-};
-
-function getModel(input: GetModelInput) {
-	const { provider, model, apiKey } = input;
-	const baseURL = input.baseURL || undefined;
-
-	return match(provider)
-		.with("openai", () => createOpenAI({ apiKey, baseURL }).languageModel(model))
-		.with("ollama", () => createOllama({ apiKey, baseURL }).languageModel(model))
-		.with("anthropic", () => createAnthropic({ apiKey, baseURL }).languageModel(model))
-		.with("vercel-ai-gateway", () => createGateway({ apiKey, baseURL }).languageModel(model))
-		.with("gemini", () => createGoogleGenerativeAI({ apiKey, baseURL }).languageModel(model))
-		.exhaustive();
-}
-
-const aiCredentialsSchema = z.object({
-	provider: aiProviderSchema,
-	model: z.string(),
-	apiKey: z.string(),
-	baseURL: z.string(),
-});
-
-const fileInputSchema = z.object({
-	name: z.string(),
-	data: z.string(), // base64 encoded
-});
+import { aiCredentialsSchema, aiProviderSchema, aiService, fileInputSchema, formatZodError } from "../services/ai";
 
 export const aiRouter = {
 	testConnection: protectedProcedure
@@ -58,14 +14,16 @@ export const aiRouter = {
 				baseURL: z.string(),
 			}),
 		)
-		.handler(async function* ({ input }) {
-			const stream = streamText({
-				temperature: 0,
-				model: getModel(input),
-				messages: [{ role: "user", content: 'Respond with "1"' }],
-			});
+		.handler(async ({ input }) => {
+			try {
+				return await aiService.testConnection(input);
+			} catch (error) {
+				if (error instanceof AISDKError) {
+					throw new ORPCError("BAD_GATEWAY", { message: error.message });
+				}
 
-			yield* stream.textStream;
+				throw error;
+			}
 		}),
 
 	parsePdf: protectedProcedure
@@ -77,44 +35,15 @@ export const aiRouter = {
 		)
 		.handler(async ({ input }) => {
 			try {
-				const model = getModel(input);
-
-				const result = await generateText({
-					model,
-					maxRetries: 0,
-					output: Output.object({ schema: resumeDataSchema }),
-					messages: [
-						{
-							role: "system",
-							content: pdfParserSystemPrompt,
-						},
-						{
-							role: "user",
-							content: [
-								{ type: "text", text: pdfParserUserPrompt },
-								{
-									type: "file",
-									filename: input.file.name,
-									mediaType: "application/pdf",
-									data: input.file.data,
-								},
-							],
-						},
-					],
-				});
-
-				return resumeDataSchema.parse({
-					...result.output,
-					customSections: [],
-					picture: defaultResumeData.picture,
-					metadata: defaultResumeData.metadata,
-				});
+				return await aiService.parsePdf(input);
 			} catch (error) {
-				if (error instanceof ZodError) {
-					const errors = flattenError(error);
-					throw new Error(JSON.stringify(errors));
+				if (error instanceof AISDKError) {
+					throw new ORPCError("BAD_GATEWAY", { message: error.message });
 				}
 
+				if (error instanceof ZodError) {
+					throw new Error(formatZodError(error));
+				}
 				throw error;
 			}
 		}),
@@ -132,39 +61,14 @@ export const aiRouter = {
 		)
 		.handler(async ({ input }) => {
 			try {
-				const model = getModel(input);
-
-				const result = await generateText({
-					model,
-					maxRetries: 0,
-					output: Output.object({ schema: resumeDataSchema }),
-					messages: [
-						{ role: "system", content: docxParserSystemPrompt },
-						{
-							role: "user",
-							content: [
-								{ type: "text", text: docxParserUserPrompt },
-								{
-									type: "file",
-									filename: input.file.name,
-									mediaType: input.mediaType,
-									data: input.file.data,
-								},
-							],
-						},
-					],
-				});
-
-				return resumeDataSchema.parse({
-					...result.output,
-					customSections: [],
-					picture: defaultResumeData.picture,
-					metadata: defaultResumeData.metadata,
-				});
+				return await aiService.parseDocx(input);
 			} catch (error) {
+				if (error instanceof AISDKError) {
+					throw new ORPCError("BAD_GATEWAY", { message: error.message });
+				}
+
 				if (error instanceof ZodError) {
-					const errors = flattenError(error);
-					throw new Error(JSON.stringify(errors));
+					throw new Error(formatZodError(error));
 				}
 
 				throw error;

src/integrations/orpc/services/ai.ts

@@ -0,0 +1,148 @@
+import { createAnthropic } from "@ai-sdk/anthropic";
+import { createGoogleGenerativeAI } from "@ai-sdk/google";
+import { createOpenAI } from "@ai-sdk/openai";
+import { createGateway, generateText, Output } from "ai";
+import { createOllama } from "ai-sdk-ollama";
+import { match } from "ts-pattern";
+import type { ZodError } from "zod";
+import z, { flattenError } from "zod";
+import docxParserSystemPrompt from "@/integrations/ai/prompts/docx-parser-system.md?raw";
+import docxParserUserPrompt from "@/integrations/ai/prompts/docx-parser-user.md?raw";
+import pdfParserSystemPrompt from "@/integrations/ai/prompts/pdf-parser-system.md?raw";
+import pdfParserUserPrompt from "@/integrations/ai/prompts/pdf-parser-user.md?raw";
+import type { ResumeData } from "@/schema/resume/data";
+import { defaultResumeData, resumeDataSchema } from "@/schema/resume/data";
+
+export const aiProviderSchema = z.enum(["ollama", "openai", "gemini", "anthropic", "vercel-ai-gateway"]);
+
+export type AIProvider = z.infer<typeof aiProviderSchema>;
+
+export type GetModelInput = {
+	provider: AIProvider;
+	model: string;
+	apiKey: string;
+	baseURL: string;
+};
+
+function getModel(input: GetModelInput) {
+	const { provider, model, apiKey } = input;
+	const baseURL = input.baseURL || undefined;
+
+	return match(provider)
+		.with("openai", () => createOpenAI({ apiKey, baseURL }).languageModel(model))
+		.with("ollama", () => createOllama({ apiKey, baseURL }).languageModel(model))
+		.with("anthropic", () => createAnthropic({ apiKey, baseURL }).languageModel(model))
+		.with("vercel-ai-gateway", () => createGateway({ apiKey, baseURL }).languageModel(model))
+		.with("gemini", () => createGoogleGenerativeAI({ apiKey, baseURL }).languageModel(model))
+		.exhaustive();
+}
+
+export const aiCredentialsSchema = z.object({
+	provider: aiProviderSchema,
+	model: z.string(),
+	apiKey: z.string(),
+	baseURL: z.string(),
+});
+
+export const fileInputSchema = z.object({
+	name: z.string(),
+	data: z.string(), // base64 encoded
+});
+
+export type TestConnectionInput = z.infer<typeof aiCredentialsSchema>;
+
+export async function testConnection(input: TestConnectionInput): Promise<boolean> {
+	const RESPONSE_OK = "1";
+
+	const result = await generateText({
+		model: getModel(input),
+		output: Output.choice({ options: [RESPONSE_OK] }),
+		messages: [{ role: "user", content: `Respond with "${RESPONSE_OK}"` }],
+	});
+
+	return result.output === RESPONSE_OK;
+}
+
+export type ParsePdfInput = z.infer<typeof aiCredentialsSchema> & {
+	file: z.infer<typeof fileInputSchema>;
+};
+
+export async function parsePdf(input: ParsePdfInput): Promise<ResumeData> {
+	const model = getModel(input);
+
+	const result = await generateText({
+		model,
+		output: Output.object({ schema: resumeDataSchema }),
+		messages: [
+			{
+				role: "system",
+				content: pdfParserSystemPrompt,
+			},
+			{
+				role: "user",
+				content: [
+					{ type: "text", text: pdfParserUserPrompt },
+					{
+						type: "file",
+						filename: input.file.name,
+						mediaType: "application/pdf",
+						data: input.file.data,
+					},
+				],
+			},
+		],
+	});
+
+	return resumeDataSchema.parse({
+		...result.output,
+		customSections: [],
+		picture: defaultResumeData.picture,
+		metadata: defaultResumeData.metadata,
+	});
+}
+
+export type ParseDocxInput = z.infer<typeof aiCredentialsSchema> & {
+	file: z.infer<typeof fileInputSchema>;
+	mediaType: "application/msword" | "application/vnd.openxmlformats-officedocument.wordprocessingml.document";
+};
+
+export async function parseDocx(input: ParseDocxInput): Promise<ResumeData> {
+	const model = getModel(input);
+
+	const result = await generateText({
+		model,
+		output: Output.object({ schema: resumeDataSchema }),
+		messages: [
+			{ role: "system", content: docxParserSystemPrompt },
+			{
+				role: "user",
+				content: [
+					{ type: "text", text: docxParserUserPrompt },
+					{
+						type: "file",
+						filename: input.file.name,
+						mediaType: input.mediaType,
+						data: input.file.data,
+					},
+				],
+			},
+		],
+	});
+
+	return resumeDataSchema.parse({
+		...result.output,
+		customSections: [],
+		picture: defaultResumeData.picture,
+		metadata: defaultResumeData.metadata,
+	});
+}
+
+export function formatZodError(error: ZodError): string {
+	return JSON.stringify(flattenError(error));
+}
+
+export const aiService = {
+	testConnection,
+	parsePdf,
+	parseDocx,
+};

src/integrations/orpc/services/storage.ts

@@ -286,8 +286,6 @@ class S3StorageService implements StorageService {
 				message: "S3 storage is accessible and credentials are valid.",
 			};
 		} catch (error: unknown) {
-			console.error(error);
-
 			return {
 				type: "s3",
 				status: "unhealthy",