Skip to content

Added tool for the brand config files #1831

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions packages/ai/src/tools/helpers.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import fg from 'fast-glob';
import { access } from 'fs/promises';

export const IGNORE_PATHS = [
'node_modules/**',
Expand Down Expand Up @@ -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,
Expand All @@ -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);
}
21 changes: 20 additions & 1 deletion packages/ai/src/tools/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down Expand Up @@ -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,
};
72 changes: 72 additions & 0 deletions packages/ai/test/tools/brand-config-file.test.ts
Original file line number Diff line number Diff line change
@@ -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}`);
});
});