diff --git a/packages/ai/src/tools/helpers.ts b/packages/ai/src/tools/helpers.ts index 877e1ccb95..10d634b643 100644 --- a/packages/ai/src/tools/helpers.ts +++ b/packages/ai/src/tools/helpers.ts @@ -1,4 +1,5 @@ import fg from 'fast-glob'; +import { access } from 'fs/promises'; export const IGNORE_PATHS = [ 'node_modules/**', @@ -29,6 +30,15 @@ export async function getAllFiles( }, ): Promise<{ success: boolean; files?: string[]; error?: string }> { try { + const exists = await access(dirPath) + .then(() => true) + .catch(() => false); + if (!exists) { + return { + success: false, + error: `Directory does not exist: ${dirPath}`, + }; + } const files = await fg(options.patterns, { cwd: dirPath, ignore: options.ignore, @@ -40,3 +50,13 @@ export async function getAllFiles( return { success: false, error: error instanceof Error ? error.message : 'Unknown error' }; } } +export async function getBrandConfigFiles( + dirPath: string, + options: FileFilterOptions = { + patterns: ['**/globals.css', '**/tailwind.config.{js,ts,mjs}'], + ignore: IGNORE_PATHS, + maxDepth: 5, + }, +): Promise<{ success: boolean; files?: string[]; error?: string }> { + return getAllFiles(dirPath, options); +} diff --git a/packages/ai/src/tools/index.ts b/packages/ai/src/tools/index.ts index d16e1ead8e..f17d97a755 100644 --- a/packages/ai/src/tools/index.ts +++ b/packages/ai/src/tools/index.ts @@ -3,7 +3,7 @@ import { tool, type ToolSet } from 'ai'; import { readFile } from 'fs/promises'; import { z } from 'zod'; import { ONLOOK_PROMPT } from '../prompt/onlook'; -import { getAllFiles } from './helpers'; +import { getAllFiles, getBrandConfigFiles } from './helpers'; export const listFilesTool = tool({ description: 'List all files in the current directory, including subdirectories', @@ -130,8 +130,27 @@ export const getStrReplaceEditorTool = (handlers: FileOperationHandlers) => { return strReplaceEditorTool; }; +export const getBrandConfigTool = tool({ + description: 'Get the brand config of the current project', + parameters: z.object({ + path: z + .string() + .describe( + 'The absolute path to the directory to get files from. This should be the root directory of the project.', + ), + }), + execute: async ({ path }) => { + const res = await getBrandConfigFiles(path); + if (!res.success) { + return { error: res.error }; + } + return res.files; + }, +}); + export const chatToolSet: ToolSet = { list_files: listFilesTool, read_files: readFilesTool, onlook_instructions: onlookInstructionsTool, + get_brand_config: getBrandConfigTool, }; diff --git a/packages/ai/test/tools/brand-config-file.test.ts b/packages/ai/test/tools/brand-config-file.test.ts new file mode 100644 index 0000000000..31525bc463 --- /dev/null +++ b/packages/ai/test/tools/brand-config-file.test.ts @@ -0,0 +1,72 @@ +import { getBrandConfigFiles } from '@onlook/ai/src/tools/helpers'; +import { afterEach, beforeEach, describe, expect, test } from 'bun:test'; +import { mkdirSync, rmSync, writeFileSync } from 'fs'; +import { join } from 'path'; + +describe('getBrandConfigFiles', () => { + const testDir = join(__dirname, 'test-files'); + + beforeEach(() => { + // Create test directory structure + mkdirSync(testDir); + mkdirSync(join(testDir, 'src')); + mkdirSync(join(testDir, 'styles')); + mkdirSync(join(testDir, 'node_modules')); + + // Create test files + writeFileSync(join(testDir, 'tailwind.config.js'), 'content'); + writeFileSync(join(testDir, 'styles/globals.css'), 'content'); + writeFileSync(join(testDir, 'src/tailwind.config.ts'), 'content'); + writeFileSync(join(testDir, 'node_modules/globals.css'), 'content'); + }); + + afterEach(() => { + // Cleanup test directory + rmSync(testDir, { recursive: true, force: true }); + }); + + test('should find all brand config files', async () => { + const { files } = await getBrandConfigFiles(testDir, { + patterns: ['**/globals.css', '**/tailwind.config.{js,ts,mjs}'], + ignore: ['node_modules/**'], + }); + expect(files?.length).toBe(3); + expect(files?.some((f) => f.endsWith('tailwind.config.js'))).toBe(true); + expect(files?.some((f) => f.endsWith('tailwind.config.ts'))).toBe(true); + expect(files?.some((f) => f.endsWith('globals.css'))).toBe(true); + }); + + test('should exclude node_modules', async () => { + const { files } = await getBrandConfigFiles(testDir, { + patterns: ['**/globals.css', '**/tailwind.config.{js,ts,mjs}'], + ignore: ['node_modules/**'], + }); + expect(files?.length).toBe(3); + expect(files?.every((f) => !f.includes('node_modules'))).toBe(true); + }); + + test('should find only tailwind config files', async () => { + const { files } = await getBrandConfigFiles(testDir, { + patterns: ['**/tailwind.config.{js,ts,mjs}'], + ignore: [], + }); + expect(files?.length).toBe(2); + expect(files?.every((f) => f.includes('tailwind.config'))).toBe(true); + }); + + test('should find only globals.css files', async () => { + const { files } = await getBrandConfigFiles(testDir, { + patterns: ['**/globals.css'], + ignore: ['node_modules/**'], + }); + expect(files?.length).toBe(1); + expect(files?.every((f) => f.endsWith('globals.css'))).toBe(true); + }); + + test('should handle non-existent directory', async () => { + const nonExistentDir = join(testDir, 'does-not-exist'); + const result = await getBrandConfigFiles(nonExistentDir); + expect(result.success).toBe(false); + expect(result.error).toBe(`Directory does not exist: ${nonExistentDir}`); + }); +});