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

This commit is contained in:
Amruth Pillai
2026-02-02 00:43:31 +01:00
committed by GitHub
parent f5a9ffb776
commit da9a3c0b12
63 changed files with 349 additions and 348 deletions
+12 -12
View File
@@ -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: {} }] })],
});
+25 -121
View File
@@ -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;
+148
View File
@@ -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,
};
@@ -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",
+5 -6
View File
@@ -1,7 +1,6 @@
import { SmartCoercionPlugin } from "@orpc/json-schema";
import { OpenAPIGenerator } from "@orpc/openapi";
import { OpenAPIHandler } from "@orpc/openapi/fetch";
import { onError } from "@orpc/server";
import { RequestHeadersPlugin } from "@orpc/server/plugins";
import { ZodToJsonSchemaConverter } from "@orpc/zod/zod4";
import { createFileRoute } from "@tanstack/react-router";
@@ -10,11 +9,11 @@ import { env } from "@/utils/env";
import { getLocale } from "@/utils/locale";
const openAPIHandler = new OpenAPIHandler(router, {
interceptors: [
onError((error) => {
console.error(error);
}),
],
// interceptors: [
// onError((error) => {
// console.error(error);
// }),
// ],
plugins: [
new RequestHeadersPlugin(),
new SmartCoercionPlugin({
+5 -6
View File
@@ -1,4 +1,3 @@
import { onError } from "@orpc/server";
import { RPCHandler } from "@orpc/server/fetch";
import { BatchHandlerPlugin, RequestHeadersPlugin } from "@orpc/server/plugins";
import { createFileRoute } from "@tanstack/react-router";
@@ -6,11 +5,11 @@ import router from "@/integrations/orpc/router";
import { getLocale } from "@/utils/locale";
const rpcHandler = new RPCHandler(router, {
interceptors: [
onError((error) => {
console.error(error);
}),
],
// interceptors: [
// onError((error) => {
// console.error(error);
// }),
// ],
plugins: [new BatchHandlerPlugin(), new RequestHeadersPlugin()],
});
@@ -82,9 +82,8 @@ export function BuilderDock() {
try {
const { url } = await printResumeAsPDF({ id: resume.id });
downloadFromUrl(url, filename);
} catch (error) {
} catch {
toast.error(t`There was a problem while generating the PDF, please try again in some time.`);
console.error("[Error from printResumeAsPDF]:", error);
} finally {
toast.dismiss(toastId);
}
@@ -34,9 +34,8 @@ export function ExportSectionBuilder() {
try {
const { url } = await printResumeAsPDF({ id: resume.id });
downloadFromUrl(url, filename);
} catch (error) {
} catch {
toast.error(t`There was a problem while generating the PDF, please try again in some time.`);
console.error("[Error from printResumeAsPDF]:", error);
} finally {
toast.dismiss(toastId);
}
+1 -1
View File
@@ -179,7 +179,7 @@ export function DashboardSidebar() {
animate={{ y: 0, height: "auto", opacity: 1 }}
exit={{ y: 50, height: 0, opacity: 0 }}
>
<Copyright className="shrink-0 break-words whitespace-normal p-2" />
<Copyright className="wrap-break-word shrink-0 whitespace-normal p-2" />
</motion.div>
)}
</AnimatePresence>
+20 -15
View File
@@ -5,6 +5,7 @@ import { useMutation } from "@tanstack/react-query";
import { createFileRoute } from "@tanstack/react-router";
import { motion } from "motion/react";
import { useMemo } from "react";
import { toast } from "sonner";
import { useIsClient } from "usehooks-ts";
import { Button } from "@/components/ui/button";
import { Combobox, type ComboboxOption } from "@/components/ui/combobox";
@@ -14,7 +15,7 @@ import { Separator } from "@/components/ui/separator";
import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import { type AIProvider, useAIStore } from "@/integrations/ai/store";
import { client } from "@/integrations/orpc/client";
import { orpc } from "@/integrations/orpc/client";
import { cn } from "@/utils/style";
import { DashboardHeader } from "../-components/header";
@@ -62,19 +63,7 @@ function AIForm() {
return providerOptions.find((option) => option.value === provider);
}, [provider]);
const { mutate: testConnection, isPending: isTesting } = useMutation({
mutationFn: async () => {
if (testStatus === "success") return;
const stream = await client.ai.testConnection({ provider, model, apiKey, baseURL });
let result = "";
for await (const chunk of stream) {
result += chunk;
}
set((draft) => {
draft.testStatus = result === "1" ? "success" : "failure";
});
},
});
const { mutate: testConnection, isPending: isTesting } = useMutation(orpc.ai.testConnection.mutationOptions());
const handleProviderChange = (value: AIProvider | null) => {
if (!value) return;
@@ -101,6 +90,22 @@ function AIForm() {
});
};
const handleTestConnection = () => {
testConnection(
{ provider, model, apiKey, baseURL },
{
onSuccess: (data) => {
set((draft) => {
draft.testStatus = data ? "success" : "failure";
});
},
onError: (error) => {
toast.error(error.message);
},
},
);
};
return (
<div className="grid gap-6 sm:grid-cols-2">
<div className="flex flex-col gap-y-2">
@@ -158,7 +163,7 @@ function AIForm() {
</div>
<div>
<Button variant="outline" disabled={isTesting || enabled} onClick={() => testConnection()}>
<Button variant="outline" disabled={isTesting || enabled} onClick={handleTestConnection}>
{isTesting ? (
<Spinner />
) : testStatus === "success" ? (