diff --git a/messages/en.json b/messages/en.json index 1ac73a758..baaccf633 100644 --- a/messages/en.json +++ b/messages/en.json @@ -247,7 +247,8 @@ "visualization": "Data Visualization", "webSearch": "Search the Web", "http": "HTTP Request", - "code": "Code Execution" + "code": "Code Execution", + "fileGenerator": "File Generator" } }, "VoiceChat": { diff --git a/src/app/api/chat/route.ts b/src/app/api/chat/route.ts index 9f6e4e21c..6b62ce516 100644 --- a/src/app/api/chat/route.ts +++ b/src/app/api/chat/route.ts @@ -286,6 +286,7 @@ export async function POST(request: Request) { const vercelAITooles = safe({ ...MCP_TOOLS, ...WORKFLOW_TOOLS, + ...IMAGE_TOOL, }) .map((t) => { const bindingTools = diff --git a/src/app/api/storage/config/route.ts b/src/app/api/storage/config/route.ts new file mode 100644 index 000000000..5d4038fab --- /dev/null +++ b/src/app/api/storage/config/route.ts @@ -0,0 +1,7 @@ +import { isFileStorageConfigured } from "@/lib/file-storage/is-storage-configured"; +import { NextResponse } from "next/server"; + +export async function GET() { + const configured = isFileStorageConfigured(); + return NextResponse.json({ configured }); +} diff --git a/src/app/layout.tsx b/src/app/layout.tsx index ea07350af..6b551678e 100644 --- a/src/app/layout.tsx +++ b/src/app/layout.tsx @@ -9,6 +9,7 @@ import { import { Toaster } from "ui/sonner"; import { NextIntlClientProvider } from "next-intl"; import { getLocale } from "next-intl/server"; +import { FileStorageInitializer } from "@/components/file-storage-initializer"; const geistSans = Geist({ variable: "--font-geist-sans", subsets: ["latin"], @@ -45,6 +46,7 @@ export default async function RootLayout({ > +
{children} diff --git a/src/app/store/index.ts b/src/app/store/index.ts index 44a7ba541..fb8a1c59c 100644 --- a/src/app/store/index.ts +++ b/src/app/store/index.ts @@ -88,6 +88,8 @@ const initialState: AppState = { allowedAppDefaultToolkit: [ AppDefaultToolkit.Code, AppDefaultToolkit.Visualization, + AppDefaultToolkit.WebSearch, + // FileGenerator will be added dynamically if storage is configured ], toolPresets: [], chatModel: undefined, diff --git a/src/components/file-storage-initializer.tsx b/src/components/file-storage-initializer.tsx new file mode 100644 index 000000000..0b76f1960 --- /dev/null +++ b/src/components/file-storage-initializer.tsx @@ -0,0 +1,50 @@ +"use client"; + +import { appStore } from "@/app/store"; +import { AppDefaultToolkit } from "@/lib/ai/tools"; +import { useEffect } from "react"; +import { useShallow } from "zustand/shallow"; + +/** + * Initializes file-generator tool if file storage is configured. + * This runs once on app mount and adds FileGenerator to the default toolkit. + */ +export function FileStorageInitializer() { + const [allowedAppDefaultToolkit, appStoreMutate] = appStore( + useShallow((state) => [state.allowedAppDefaultToolkit, state.mutate]), + ); + + useEffect(() => { + // Only run once on mount + const checkStorageAndInit = async () => { + try { + const response = await fetch("/api/storage/config"); + const data = await response.json(); + + if (data.configured) { + // Check if FileGenerator is already in the list + const hasFileGenerator = allowedAppDefaultToolkit?.includes( + AppDefaultToolkit.FileGenerator, + ); + + if (!hasFileGenerator) { + // Add FileGenerator to the default toolkit + appStoreMutate((prev) => ({ + allowedAppDefaultToolkit: [ + ...(prev.allowedAppDefaultToolkit || []), + AppDefaultToolkit.FileGenerator, + ], + })); + } + } + } catch (error) { + console.error("Failed to check file storage config:", error); + } + }; + + checkStorageAndInit(); + // eslint-disable-next-line react-hooks/exhaustive-deps + }, []); // Only run once on mount + + return null; // This component doesn't render anything +} diff --git a/src/components/message-parts.tsx b/src/components/message-parts.tsx index 366d989c8..37826eaf8 100644 --- a/src/components/message-parts.tsx +++ b/src/components/message-parts.tsx @@ -51,7 +51,11 @@ import { VercelAIWorkflowToolStreamingResultTag, } from "app-types/workflow"; import { Avatar, AvatarFallback, AvatarImage } from "ui/avatar"; -import { DefaultToolName, ImageToolName } from "lib/ai/tools"; +import { + DefaultToolName, + ImageToolName, + FileGeneratorToolName, +} from "lib/ai/tools"; import { Shortcut, getShortcutKeyList, @@ -726,6 +730,17 @@ const ImageGeneratorToolInvocation = dynamic( }, ); +const FileGeneratorToolInvocation = dynamic( + () => + import("./tool-invocation/file-generator").then( + (mod) => mod.FileGeneratorToolInvocation, + ), + { + ssr: false, + loading, + }, +); + // Local shortcuts for tool invocation approval/rejection const approveToolInvocationShortcut: Shortcut = { description: "approveToolInvocation", @@ -880,6 +895,10 @@ export const ToolMessagePart = memo( return ; } + if (toolName === FileGeneratorToolName) { + return ; + } + if (toolName === DefaultToolName.JavascriptExecution) { return ( { + return !part.state.startsWith("output"); + }, [part.state]); + + const result = useMemo(() => { + if (!part.state.startsWith("output")) return null; + return part.output as FileGeneratorResult; + }, [part.state, part.output]); + + const files = useMemo(() => { + return result?.files || []; + }, [result]); + + const hasError = useMemo(() => { + return ( + part.state === "output-error" || + (part.state === "output-available" && result?.files.length === 0) + ); + }, [part.state, result]); + + // Loading state + if (isGenerating) { + return ( +
+ Generating file... +
+
+ + Creating your file... +
+
+
+ ); + } + + return ( +
+
+ {!hasError && } + + {hasError + ? "File generation failed" + : files.length === 1 + ? "File generated" + : `${files.length} files generated`} + +
+ + {hasError ? ( +
+ {part.errorText ?? "Failed to generate file. Please try again."} +
+ ) : ( +
+ {result?.description && ( +

+ {result.description} +

+ )} +
+ {files.map((file, index) => { + const fileExtension = + file.filename.split(".").pop()?.toUpperCase() || "FILE"; + + return ( +
+
+
+ +
+
+

+ {file.filename} +

+
+ + {fileExtension} + + {formatFileSize(file.size)} + {file.mimeType && ( + + {file.mimeType} + + )} +
+
+ +
+
+ ); + })} +
+
+ )} +
+ ); +} + +export const FileGeneratorToolInvocation = memo( + PureFileGeneratorToolInvocation, + (prev, next) => { + return equal(prev.part, next.part); + }, +); diff --git a/src/components/tool-select-dropdown.tsx b/src/components/tool-select-dropdown.tsx index 44f0ee2b9..f83aa9153 100644 --- a/src/components/tool-select-dropdown.tsx +++ b/src/components/tool-select-dropdown.tsx @@ -9,6 +9,7 @@ import { ChartColumn, ChevronRight, CodeIcon, + FileIcon, GlobeIcon, HardDriveUploadIcon, ImagesIcon, @@ -887,6 +888,9 @@ function AppDefaultToolKitSelector() { case AppDefaultToolkit.Code: icon = CodeIcon; break; + case AppDefaultToolkit.FileGenerator: + icon = FileIcon; + break; } return { label, diff --git a/src/lib/ai/tools/file-generator/index.test.ts b/src/lib/ai/tools/file-generator/index.test.ts new file mode 100644 index 000000000..fd41352fe --- /dev/null +++ b/src/lib/ai/tools/file-generator/index.test.ts @@ -0,0 +1,405 @@ +import { describe, expect, it, vi, beforeEach } from "vitest"; +import type { FileStorage } from "lib/file-storage/file-storage.interface"; + +// Mock the server file storage +const mockUpload = vi.fn(); +const mockGetDownloadUrl = vi.fn(); + +vi.mock("lib/file-storage", () => ({ + serverFileStorage: { + upload: mockUpload, + getDownloadUrl: mockGetDownloadUrl, + } as Partial, +})); + +// Mock logger +vi.mock("logger", () => ({ + default: { + error: vi.fn(), + info: vi.fn(), + }, +})); + +// Import after mocking +const { fileGeneratorTool } = await import("./index"); + +type FileGeneratorToolResult = { + files: { + url: string; + filename: string; + mimeType: string; + size: number; + }[]; + description?: string; + guide?: string; +}; + +describe("fileGeneratorTool.execute", () => { + beforeEach(() => { + vi.clearAllMocks(); + }); + + describe("MIME type inference", () => { + const testCases = [ + { ext: "csv", expected: "text/csv" }, + { ext: "json", expected: "application/json" }, + { ext: "xml", expected: "application/xml" }, + { ext: "yaml", expected: "text/yaml" }, + { ext: "yml", expected: "text/yaml" }, + { ext: "txt", expected: "text/plain" }, + { ext: "md", expected: "text/markdown" }, + { ext: "html", expected: "text/html" }, + { ext: "css", expected: "text/css" }, + { ext: "js", expected: "text/javascript" }, + { ext: "ts", expected: "text/typescript" }, + { ext: "py", expected: "text/x-python" }, + { ext: "sh", expected: "application/x-sh" }, + { ext: "sql", expected: "application/sql" }, + { ext: "env", expected: "text/plain" }, + { ext: "log", expected: "text/plain" }, + { ext: "unknown", expected: "text/plain" }, // fallback + ]; + + testCases.forEach(({ ext, expected }) => { + it(`should infer ${expected} for .${ext} files`, async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file", + metadata: { + key: "test-key", + filename: `test.${ext}`, + contentType: expected, + size: 100, + uploadedAt: new Date(), + }, + }); + + if (fileGeneratorTool.execute) { + await fileGeneratorTool.execute( + { + files: [ + { + filename: `test.${ext}`, + content: "test content", + }, + ], + }, + {} as any, // ToolCallOptions + ); + } + + expect(mockUpload).toHaveBeenCalledWith( + expect.any(Buffer), + expect.objectContaining({ + filename: `test.${ext}`, + contentType: expected, + }), + ); + }); + }); + + it("should use provided mimeType over inferred one", async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file", + metadata: { + key: "test-key", + filename: "test.csv", + contentType: "application/custom", + size: 100, + uploadedAt: new Date(), + }, + }); + + if (fileGeneratorTool.execute) { + await fileGeneratorTool.execute( + { + files: [ + { + filename: "test.csv", + content: "test content", + mimeType: "application/custom", + }, + ], + }, + {} as any, + ); + } + + expect(mockUpload).toHaveBeenCalledWith( + expect.any(Buffer), + expect.objectContaining({ + contentType: "application/custom", + }), + ); + }); + }); + + describe("file upload", () => { + it("should upload file content as Buffer", async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file", + metadata: { + key: "test-key", + filename: "test.txt", + contentType: "text/plain", + size: 12, + uploadedAt: new Date(), + }, + }); + + const content = "test content"; + if (fileGeneratorTool.execute) { + await fileGeneratorTool.execute( + { + files: [ + { + filename: "test.txt", + content, + }, + ], + }, + {} as any, + ); + } + + expect(mockUpload).toHaveBeenCalledTimes(1); + const uploadedBuffer = mockUpload.mock.calls[0][0] as Buffer; + expect(uploadedBuffer.toString("utf-8")).toBe(content); + }); + + it("should use presigned URL when available", async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://s3.example.com/file", + metadata: { + key: "test-key", + filename: "test.txt", + contentType: "text/plain", + size: 12, + uploadedAt: new Date(), + }, + }); + + mockGetDownloadUrl.mockResolvedValue( + "https://s3.example.com/file?presigned=true", + ); + + let result: + | FileGeneratorToolResult + | AsyncIterable + | undefined; + if (fileGeneratorTool.execute) { + result = await fileGeneratorTool.execute( + { + files: [ + { + filename: "test.txt", + content: "test", + }, + ], + }, + {} as any, + ); + } + + expect(mockGetDownloadUrl).toHaveBeenCalledWith("test-key"); + if (result && !(Symbol.asyncIterator in result)) { + expect(result.files[0].url).toBe( + "https://s3.example.com/file?presigned=true", + ); + } + }); + + it("should fall back to sourceUrl when getDownloadUrl returns null", async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file", + metadata: { + key: "test-key", + filename: "test.txt", + contentType: "text/plain", + size: 12, + uploadedAt: new Date(), + }, + }); + + mockGetDownloadUrl.mockResolvedValue(null); + + let result: + | FileGeneratorToolResult + | AsyncIterable + | undefined; + if (fileGeneratorTool.execute) { + result = await fileGeneratorTool.execute( + { + files: [ + { + filename: "test.txt", + content: "test", + }, + ], + }, + {} as any, + ); + } + + if (result && !(Symbol.asyncIterator in result)) { + expect(result.files[0].url).toBe("https://example.com/file"); + } + }); + + it("should upload multiple files in parallel", async () => { + mockUpload.mockImplementation(async (buffer, options) => ({ + key: `key-${options?.filename}`, + sourceUrl: `https://example.com/${options?.filename}`, + metadata: { + key: `key-${options?.filename}`, + filename: options?.filename || "file", + contentType: options?.contentType || "text/plain", + size: (buffer as Buffer).length, + uploadedAt: new Date(), + }, + })); + + let result: + | FileGeneratorToolResult + | AsyncIterable + | undefined; + if (fileGeneratorTool.execute) { + result = await fileGeneratorTool.execute( + { + files: [ + { filename: "file1.txt", content: "content1" }, + { filename: "file2.txt", content: "content2" }, + { filename: "file3.txt", content: "content3" }, + ], + }, + {} as any, + ); + } + + expect(mockUpload).toHaveBeenCalledTimes(3); + if (result && !(Symbol.asyncIterator in result)) { + expect(result.files).toHaveLength(3); + expect(result.files[0].filename).toBe("file1.txt"); + expect(result.files[1].filename).toBe("file2.txt"); + expect(result.files[2].filename).toBe("file3.txt"); + } + }); + }); + + describe("result format", () => { + it("should return correct result structure", async () => { + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file.csv", + metadata: { + key: "test-key", + filename: "data.csv", + contentType: "text/csv", + size: 100, + uploadedAt: new Date(), + }, + }); + + let result: + | FileGeneratorToolResult + | AsyncIterable + | undefined; + if (fileGeneratorTool.execute) { + result = await fileGeneratorTool.execute( + { + files: [ + { + filename: "data.csv", + content: "col1,col2\nval1,val2", + mimeType: "text/csv", + }, + ], + description: "Test CSV file", + }, + {} as any, + ); + } + + if (result && !(Symbol.asyncIterator in result)) { + expect(result).toMatchObject({ + files: [ + { + url: "https://example.com/file.csv", + filename: "data.csv", + mimeType: "text/csv", + size: 100, + }, + ], + description: "Test CSV file", + }); + expect(result.guide).toBeDefined(); + } + }); + + it("should include file size in result", async () => { + const content = "test content with some length"; + + mockUpload.mockResolvedValue({ + key: "test-key", + sourceUrl: "https://example.com/file", + metadata: { + key: "test-key", + filename: "test.txt", + contentType: "text/plain", + size: Buffer.from(content).length, + uploadedAt: new Date(), + }, + }); + + let result: + | FileGeneratorToolResult + | AsyncIterable + | undefined; + if (fileGeneratorTool.execute) { + result = await fileGeneratorTool.execute( + { + files: [ + { + filename: "test.txt", + content, + }, + ], + }, + {} as any, + ); + } + + if (result && !(Symbol.asyncIterator in result)) { + expect(result.files[0].size).toBe(Buffer.from(content).length); + } + }); + }); + + describe("error handling", () => { + it("should throw error when upload fails", async () => { + mockUpload.mockRejectedValue(new Error("Upload failed")); + + if (fileGeneratorTool.execute) { + await expect( + fileGeneratorTool.execute( + { + files: [ + { + filename: "test.txt", + content: "test", + }, + ], + }, + {} as any, + ), + ).rejects.toThrow( + "File generation was successful, but file upload failed", + ); + } + }); + }); +}); diff --git a/src/lib/ai/tools/file-generator/index.ts b/src/lib/ai/tools/file-generator/index.ts new file mode 100644 index 000000000..1373e78fb --- /dev/null +++ b/src/lib/ai/tools/file-generator/index.ts @@ -0,0 +1,148 @@ +import { tool as createTool } from "ai"; +import { serverFileStorage } from "lib/file-storage"; +import z from "zod"; +import { FileGeneratorToolName } from ".."; +import logger from "logger"; + +export type FileGeneratorToolResult = { + files: { + url: string; + filename: string; + mimeType: string; + size: number; + }[]; + description?: string; + guide?: string; +}; + +export const fileGeneratorTool = createTool({ + name: FileGeneratorToolName, + description: `Create and save files with specified content. Use this tool when the user requests downloadable files such as: +- Data files: CSV, JSON, XML, YAML +- Documents: Markdown, text files, configuration files +- Code files: Python, JavaScript, HTML, CSS, etc. +- Structured data exports + +The tool will generate the file, upload it to storage, and provide a download link. Do not use this for images (use image-manager instead) or for simple text responses that don't need to be downloaded. + +IMPORTANT: After invoking this tool, do NOT include file URLs in your text response. The download links are automatically displayed in the UI above your message. Simply acknowledge that the file(s) have been generated.`, + inputSchema: z.object({ + files: z + .array( + z.object({ + filename: z + .string() + .describe( + "The name of the file including extension (e.g., data.csv, script.py)", + ), + content: z.string().describe("The complete content of the file"), + mimeType: z + .string() + .optional() + .describe( + "MIME type (e.g., 'text/csv', 'application/json', 'text/plain'). If not provided, will be inferred from filename.", + ), + }), + ) + .min(1) + .max(5) + .describe("Array of files to generate (1-5 files)"), + description: z + .string() + .optional() + .describe( + "Optional description of what the files contain or how to use them", + ), + }), + execute: async ({ files, description }) => { + try { + const uploadedFiles = await Promise.all( + files.map(async (file) => { + // Convert content to Buffer + const buffer = Buffer.from(file.content, "utf-8"); + + // Infer MIME type from filename if not provided + let mimeType = file.mimeType; + if (!mimeType) { + const extension = file.filename.split(".").pop()?.toLowerCase(); + const mimeTypeMap: Record = { + csv: "text/csv", + json: "application/json", + xml: "application/xml", + yaml: "text/yaml", + yml: "text/yaml", + txt: "text/plain", + md: "text/markdown", + html: "text/html", + css: "text/css", + js: "text/javascript", + ts: "text/typescript", + py: "text/x-python", + sh: "application/x-sh", + sql: "application/sql", + env: "text/plain", + log: "text/plain", + }; + mimeType = mimeTypeMap[extension || ""] || "text/plain"; + } + + // Upload to storage (same pattern as image tool) + const uploaded = await serverFileStorage.upload(buffer, { + filename: file.filename, + contentType: mimeType, + }); + + logger.info("File uploaded:", { + key: uploaded.key, + sourceUrl: uploaded.sourceUrl, + hasGetDownloadUrl: !!serverFileStorage.getDownloadUrl, + }); + + // Get presigned URL for private buckets if available + let downloadUrl: string; + if (serverFileStorage.getDownloadUrl) { + const presignedUrl = await serverFileStorage.getDownloadUrl( + uploaded.key, + ); + downloadUrl = presignedUrl || uploaded.sourceUrl; + logger.info("Presigned URL generated:", { + downloadUrl, + isPresigned: downloadUrl.includes("X-Amz"), + }); + } else { + downloadUrl = uploaded.sourceUrl; + logger.info("Using source URL (no getDownloadUrl):", { + downloadUrl, + }); + } + + return { + url: downloadUrl, + filename: uploaded.metadata.filename || file.filename, + mimeType: uploaded.metadata.contentType || mimeType, + size: uploaded.metadata.size || buffer.length, + }; + }), + ); + + const fileCount = uploadedFiles.length; + const guide = + fileCount > 0 + ? fileCount === 1 + ? "Your file has been generated successfully and is ready for download above. You can click the download button to save it to your device." + : `Your ${fileCount} files have been generated successfully and are ready for download above. You can click the download buttons to save them to your device.` + : undefined; + + return { + files: uploadedFiles, + description, + guide, + }; + } catch (e) { + logger.error(e); + throw new Error( + "File generation was successful, but file upload failed. Please check your file upload configuration and try again.", + ); + } + }, +}); diff --git a/src/lib/ai/tools/index.ts b/src/lib/ai/tools/index.ts index 233683d7a..4fe84cf93 100644 --- a/src/lib/ai/tools/index.ts +++ b/src/lib/ai/tools/index.ts @@ -3,6 +3,7 @@ export enum AppDefaultToolkit { WebSearch = "webSearch", Http = "http", Code = "code", + FileGenerator = "fileGenerator", } export enum DefaultToolName { @@ -20,3 +21,5 @@ export enum DefaultToolName { export const SequentialThinkingToolName = "sequential-thinking"; export const ImageToolName = "image-manager"; + +export const FileGeneratorToolName = "file-generator"; diff --git a/src/lib/ai/tools/tool-kit.ts b/src/lib/ai/tools/tool-kit.ts index 22623a8e6..d2a409937 100644 --- a/src/lib/ai/tools/tool-kit.ts +++ b/src/lib/ai/tools/tool-kit.ts @@ -3,11 +3,12 @@ import { createBarChartTool } from "./visualization/create-bar-chart"; import { createLineChartTool } from "./visualization/create-line-chart"; import { createTableTool } from "./visualization/create-table"; import { exaSearchTool, exaContentsTool } from "./web/web-search"; -import { AppDefaultToolkit, DefaultToolName } from "."; +import { AppDefaultToolkit, DefaultToolName, FileGeneratorToolName } from "."; import { Tool } from "ai"; import { httpFetchTool } from "./http/fetch"; import { jsExecutionTool } from "./code/js-run-tool"; import { pythonExecutionTool } from "./code/python-run-tool"; +import { fileGeneratorTool } from "./file-generator"; export const APP_DEFAULT_TOOL_KIT: Record< AppDefaultToolkit, @@ -30,4 +31,7 @@ export const APP_DEFAULT_TOOL_KIT: Record< [DefaultToolName.JavascriptExecution]: jsExecutionTool, [DefaultToolName.PythonExecution]: pythonExecutionTool, }, + [AppDefaultToolkit.FileGenerator]: { + [FileGeneratorToolName]: fileGeneratorTool, + }, }; diff --git a/src/lib/file-storage/is-storage-configured.ts b/src/lib/file-storage/is-storage-configured.ts new file mode 100644 index 000000000..fb75f0016 --- /dev/null +++ b/src/lib/file-storage/is-storage-configured.ts @@ -0,0 +1,25 @@ +import "server-only"; + +/** + * Checks if file storage is properly configured. + * Returns true if either Vercel Blob or S3 credentials are available. + */ +export function isFileStorageConfigured(): boolean { + // Check for Vercel Blob token + if (process.env.BLOB_READ_WRITE_TOKEN) { + return true; + } + + // Check for S3 configuration + const hasS3Bucket = Boolean(process.env.FILE_STORAGE_S3_BUCKET); + const hasS3Type = process.env.FILE_STORAGE_TYPE === "s3"; + + // S3 is considered configured if: + // 1. FILE_STORAGE_TYPE is set to "s3" AND bucket is specified + // 2. Will work with either explicit credentials OR IAM role (in AWS environment) + if (hasS3Type && hasS3Bucket) { + return true; // Assume IAM role or credentials are available + } + + return false; +}