4 Commits

Author SHA1 Message Date
geoffsee
ee137774d3 wip 2025-06-13 14:03:53 -04:00
geoffsee
39ca507911 Add .toak-ignore to .gitignore and implement updateGitignore method
Enhanced MarkdownGenerator to update .gitignore programmatically when required, introduced tests for updateGitignore functionality, and ensured `.toak-ignore` is added.
2025-06-13 13:02:44 -04:00
geoffsee
6f9d37a292 Add CLI argument parsing tests and improve argument handling
- Introduced `test/cli.test.ts` with unit tests for CLI argument parsing.
- Enhanced CLI argument handling in `cli.ts` for better error reporting and help message support.
- Updated `MarkdownGenerator` and related files for improved robustness.
2025-06-13 16:41:03 +00:00
geoffsee
c8b0215435 wip 2025-06-13 16:41:03 +00:00
10 changed files with 377 additions and 10 deletions

2
.gitignore vendored
View File

@@ -3,3 +3,5 @@
/dist/
prompt.md
todo
.toak-ignore
/google_gemma-3-1b-it-Q6_K.llamafile

BIN
bun.lockb

Binary file not shown.

View File

@@ -45,7 +45,9 @@
"lint:fix": "eslint src/ --fix",
"format": "prettier --write \"**/*.{js,jsx,ts,tsx,json,md,yml,yaml}\"",
"fix": "bun format && bun lint:fix",
"release": "bunx release-it"
"release": "bunx release-it",
"run:inference": "./google_gemma-3-1b-it-Q6_K.llamafile --server --nobrowser",
"download:inference": "./download_optimization_model.sh"
},
"dependencies": {
"glob": "^11.0.1",

View File

@@ -0,0 +1,4 @@
MODEL_URL=https://huggingface.co/Mozilla/gemma-3-1b-it-llamafile/resolve/main/google_gemma-3-1b-it-Q6_K.llamafile?download=true
wget "${MODEL_URL}"

30
scripts/optimize-ignore.ts Executable file
View File

@@ -0,0 +1,30 @@
export async function optimizeToakIgnore(content: string) {
const inferenceProcess = Bun.spawn(['bun', 'run:inference']);
await new Promise(resolve => setTimeout(resolve, 5000));
const prompt = `You are a helpful assistant.
## Context
~~~
${content}
~~~
Respond with a list of files that should be added to the .toak-ignore file to reduce noise in the context. No extra text or explanations.`;
async function run() {
try {
const response = await fetch('http://127.0.0.1:8080/completion', {
method: 'POST',
body: JSON.stringify({
prompt,
n_predict: 512,
}),
});
const data = await response.json();
console.log(data.content);
} catch (error) {
console.error('Error:', error);
}
}
await run();
inferenceProcess.kill();
}

View File

@@ -8,9 +8,10 @@ import fileTypeExclusions from './fileTypeExclusions.js';
import fileExclusions from './fileExclusions.js';
import { readFileSync } from 'node:fs';
import { glob } from 'glob';
import { isPreset, type PresetPrompt, prompts } from './prompts.ts';
interface MarkdownGeneratorOptions {
export interface MarkdownGeneratorOptions {
dir?: string;
outputFilePath?: string;
fileTypeExclusions?: Set<string>;
@@ -18,6 +19,7 @@ interface MarkdownGeneratorOptions {
customPatterns?: Record<string, any>;
customSecretPatterns?: Record<string, any>;
verbose?: boolean;
todoPrompt?: string
}
/**
@@ -33,6 +35,7 @@ export class MarkdownGenerator {
private tokenCleaner: TokenCleaner;
private verbose: boolean;
private initialized: boolean;
private todoPrompt: string;
/**
* Creates an instance of MarkdownGenerator.
@@ -44,10 +47,12 @@ export class MarkdownGenerator {
this.fileTypeExclusions = new Set(
options.fileTypeExclusions || fileTypeExclusions,
);
this.fileExclusions = options.fileExclusions || fileExclusions;
this.fileExclusions = options.fileExclusions || [...fileExclusions];
// @ts-ignore - options.customPatterns signature is valid
this.tokenCleaner = new TokenCleaner(options.customPatterns, options.customSecretPatterns);
this.verbose = options.verbose !== undefined ? options.verbose : true;
this.initialized = false;
this.todoPrompt = prompts.getPrompt(options.todoPrompt)
}
/**
@@ -59,6 +64,7 @@ export class MarkdownGenerator {
private async initialize(): Promise<void> {
if (!this.initialized) {
await this.loadNestedIgnoreFiles();
await this.updateGitignore();
this.initialized = true;
}
}
@@ -255,6 +261,58 @@ export class MarkdownGenerator {
}
}
async updateGitignore(): Promise<void> {
const gitignorePath = path.join(this.dir, '.gitignore');
try {
let content = '';
try {
content = await readFile(gitignorePath, 'utf-8');
} catch (error: any) {
if (error.code === 'ENOENT') {
// .gitignore doesn't exist, create it
if (this.verbose) {
console.log('File not found, creating a \'.gitignore\' file.');
}
content = '';
} else {
throw error;
}
}
// Check if entries already exist
const lines = content.split('\n');
const needsPromptMd = !lines.some(line => line.trim() === 'prompt.md');
const needsToakIgnore = !lines.some(line => line.trim() === '.toak-ignore');
// Add entries if needed
if (needsPromptMd || needsToakIgnore) {
if (this.verbose) {
console.log('Updating .gitignore with prompt.md and .toak-ignore');
}
let newContent = content;
if (newContent && !newContent.endsWith('\n')) {
newContent += '\n';
}
if (needsPromptMd) {
newContent += 'prompt.md\n';
}
if (needsToakIgnore) {
newContent += '.toak-ignore\n';
}
await writeFile(gitignorePath, newContent);
}
} catch (error) {
if (this.verbose) {
console.error('Error updating .gitignore:', error);
}
throw error;
}
}
/**
* Creates a complete markdown document combining code documentation and todos.
* @async
@@ -277,7 +335,7 @@ export class MarkdownGenerator {
console.log({ total_tokens: totalTokens });
}
return { success: true, tokenCount: llama3Tokenizer.encode(markdown).length };
} catch (error) {
} catch (error: any) {
if (this.verbose) {
console.error('Error writing markdown document:', error);
}

View File

@@ -1,14 +1,64 @@
#!/usr/bin/env node
console.log('RUNNING TOKENIZER');
import { MarkdownGenerator } from './MarkdownGenerator.js';
import type { PresetPrompt } from './prompts';
console.log('RUNNING TOKENIZER');
import { MarkdownGenerator, type MarkdownGeneratorOptions } from './MarkdownGenerator';
import { optimizeToakIgnore } from '../scripts/optimize-ignore.ts';
let output = '';
process.stdout.write = (write => {
return (str: string | Uint8Array, ...args: any) => {
output += str.toString();
return write.apply(process.stdout, [str, ...args]);
};
})(process.stdout.write);
let optimizeIgnore = false;
const args = process.argv.slice(2);
const options: { prompt?: PresetPrompt; } & MarkdownGeneratorOptions = {
};
type ValidArg = keyof MarkdownGeneratorOptions;
for (let i = 0; i < args.length; i++) {
if (args[i] === '--help') {
console.log(`${Object.keys(options).map(item => "--" + item).join(', ')}`);
}
if (args[i] === '--prompt') {
options["todoPrompt"] = args[i + 1]
i++;
}
if (args[i] === '--optimize-ignore') {
optimizeIgnore = true;
i++;
}
const arg = args[i].replace(/^--/, '');
if (arg as any satisfies ValidArg) {
// @ts-ignore - arg can't be used to index options
options[arg] = args[i + 1]
i++;
} else {
console.log(`Invalid argument specified: ${arg}`);
console.log(`Possible arguments: ${Object.keys(options).map(item => "--" + item).join(', ')}`);
}
}
const generator = new MarkdownGenerator(options);
const generator = new MarkdownGenerator();
generator
.createMarkdownDocument()
.then((result: { success: boolean }) => {
.then(async (result: { success: boolean }) => {
if (!result.success) {
process.exit(1);
}
if (optimizeIgnore) {
await optimizeToakIgnore(output);
}
})
.catch((error: any) => {
console.error('Error:', error);

39
src/prompts.ts Normal file
View File

@@ -0,0 +1,39 @@
const taskConditionsStandard_FixErrors = `
## Task
1. Fix these errors.
## Conditions
2. Output labeled and fully fixed files only, no diffs.
## Standard
3. Respond with the files, no examples or excessive explanations.
~~~console
clean up your errors and put them here
~~~
`;
export const customPrompts: Record<string, string> = {
"tcs:fix:errors": taskConditionsStandard_FixErrors,
};
export type PresetPrompt = keyof typeof customPrompts;
export function isPreset(key: string): boolean {
return key in customPrompts;
}
export const prompts = {
...customPrompts,
default: customPrompts["tcs:fix:errors"],
getPrompt(key?: string) {
if (!key) return prompts.default;
if (!isPreset(key)) return prompts.default;
return customPrompts[key];
}
};

100
test/cli.test.ts Normal file
View File

@@ -0,0 +1,100 @@
// test/cli.test.ts
import { describe, it, expect, beforeEach, afterEach, spyOn, mock } from 'bun:test';
import { MarkdownGenerator, type MarkdownGeneratorOptions } from '../src/MarkdownGenerator';
import type { PresetPrompt } from '../src/prompts';
// Function to process CLI arguments similar to cli.ts
function processArgs(args: string[]): { prompt?: PresetPrompt; } & MarkdownGeneratorOptions {
const options: { prompt?: PresetPrompt; } & MarkdownGeneratorOptions = {};
for (let i = 0; i < args.length; i++) {
if (args[i] === '--help') {
console.log(`${Object.keys(options).map(item => "--" + item).join(', ')}`);
continue;
}
if (args[i] === '--prompt') {
options["todoPrompt"] = args[i + 1] as PresetPrompt;
i++;
continue;
}
const arg = args[i].replace(/^--/, '');
if (['dir', 'outputFilePath', 'verbose', 'todoPrompt', 'fileTypeExclusions', 'fileExclusions', 'customPatterns', 'customSecretPatterns'].includes(arg)) {
// @ts-ignore - dynamic property access
options[arg] = args[i + 1];
i++;
} else {
console.log(`Invalid argument specified: ${arg}`);
console.log(`Possible arguments: ${Object.keys(options).map(item => "--" + item).join(', ')}`);
}
}
return options;
}
describe('CLI', () => {
describe('argument handling', () => {
it('should process --prompt argument correctly', () => {
// Set up test arguments
const args = ['--prompt', 'tcs:fix:errors'];
// Process arguments
const options = processArgs(args);
// Verify options
expect(options).toEqual(expect.objectContaining({
todoPrompt: 'tcs:fix:errors'
}));
});
it('should process other valid arguments correctly', () => {
// Set up test arguments
const args = ['--dir', './src', '--outputFilePath', './custom.md', '--verbose', 'false'];
// Process arguments
const options = processArgs(args);
// Verify options
expect(options).toEqual(expect.objectContaining({
dir: './src',
outputFilePath: './custom.md',
verbose: 'false'
}));
});
it('should handle invalid arguments', () => {
// Mock console.log to capture output
const consoleLogSpy = spyOn(console, 'log');
// Set up test arguments
const args = ['--invalidArg', 'value'];
// Process arguments
processArgs(args);
// Verify error message was logged
expect(consoleLogSpy).toHaveBeenCalledWith(expect.stringContaining('Invalid argument specified: invalidArg'));
// Restore original function
consoleLogSpy.mockRestore();
});
it('should display help when --help argument is provided', () => {
// Mock console.log to capture output
const consoleLogSpy = spyOn(console, 'log');
// Set up test arguments
const args = ['--help'];
// Process arguments
processArgs(args);
// Verify help message was logged
// At this point, the options object is empty, so we just check that console.log was called
expect(consoleLogSpy).toHaveBeenCalled();
// Restore original function
consoleLogSpy.mockRestore();
});
});
});

View File

@@ -50,8 +50,8 @@ const a = 1;`;
it('should trim whitespace and empty lines', () => {
const code = `const a = 1;
const b = 2; `;
const expected = `const a = 1;
const b = 2;`;
@@ -381,6 +381,88 @@ const a = 1;
});
});
describe('updateGitignore', () => {
it('should update .gitignore with prompt.md and .toak-ignore on first run', async () => {
const gitignorePath = path.join('.', '.gitignore');
// Mock readFile to simulate .gitignore exists but doesn't have the entries
const readFileSpy = spyOn(fs, 'readFile').mockResolvedValue('node_modules\ndist\n');
// Spy on fs.writeFile
const writeFileSpy = spyOn(fs, 'writeFile').mockResolvedValue(undefined);
// Call the method
await markdownGenerator.updateGitignore();
// Verify readFile was called
expect(readFileSpy).toHaveBeenCalledWith(gitignorePath, 'utf-8');
// Verify writeFile was called with correct content
expect(writeFileSpy).toHaveBeenCalledWith(
gitignorePath,
'node_modules\ndist\nprompt.md\n.toak-ignore\n'
);
// Restore the original implementations
readFileSpy.mockRestore();
writeFileSpy.mockRestore();
});
it('should not update .gitignore if entries already exist', async () => {
const gitignorePath = path.join('.', '.gitignore');
// Mock readFile to simulate .gitignore already has the entries
const readFileSpy = spyOn(fs, 'readFile')
.mockResolvedValue('node_modules\ndist\nprompt.md\n.toak-ignore\n');
// Spy on fs.writeFile
const writeFileSpy = spyOn(fs, 'writeFile').mockResolvedValue(undefined);
// Call the method
await markdownGenerator.updateGitignore();
// Verify readFile was called
expect(readFileSpy).toHaveBeenCalledWith(gitignorePath, 'utf-8');
// Verify writeFile was NOT called
expect(writeFileSpy).not.toHaveBeenCalled();
// Restore the original implementations
readFileSpy.mockRestore();
writeFileSpy.mockRestore();
});
it('should create .gitignore if it does not exist', async () => {
const gitignorePath = path.join('.', '.gitignore');
// Mock readFile to throw ENOENT error
const readFileSpy = spyOn(fs, 'readFile').mockImplementation(() => {
const error: any = new Error('File not found');
error.code = 'ENOENT';
return Promise.reject(error);
});
// Spy on fs.writeFile
const writeFileSpy = spyOn(fs, 'writeFile').mockResolvedValue(undefined);
// Call the method
await markdownGenerator.updateGitignore();
// Verify readFile was called
expect(readFileSpy).toHaveBeenCalledWith(gitignorePath, 'utf-8');
// Verify writeFile was called with correct content
expect(writeFileSpy).toHaveBeenCalledWith(
gitignorePath,
'prompt.md\n.toak-ignore\n'
);
// Restore the original implementations
readFileSpy.mockRestore();
writeFileSpy.mockRestore();
});
});
describe('createMarkdownDocument', () => {
it('should create markdown document successfully', async () => {
const mockContent = '# Project Files\n\n## test.txt\n~~~\ntest\n~~~\n\n';