import { ollama } from "ollama-ai-provider";
import { generateObject, generateText } from "ai";
import { z } from "zod";
import { globby } from "globby";
import { parseArgs } from "node:util";
import { readFile } from "node:fs/promises";
import path from "node:path";
const { positionals } = parseArgs({
allowPositionals: true,
});
const model = ollama("gemma3:latest");
const objectPrompt = positionals[0];
const fileTypes = ["md", "txt"] as const;
// generateObject to pull out the file type
const { object: fileTypeObject } = await generateObject({
model,
schema: z.object({
fileType: z
.string()
.describe(
`The file type to process. My file types are: ${fileTypes.join(", ")}}`
)
.optional(),
}),
prompt: `
Based on the following user request, determine the most relevant file type to search.
If no file type if mentioned or implied, please return an empty string.
User Request: "${objectPrompt}"
`,
});
if (!fileTypeObject.fileType) {
console.error(
`No file type found. Only one of the following file types is allowed: ${fileTypes.join(
", "
)}`
);
process.exit(1);
}
const files = await globby(`*.${fileTypeObject.fileType}`, { gitignore: true });
console.log(files);
const isNonEmpty = (files: string[]): files is [string, ...string[]] => {
return files.length > 0;
};
if (!isNonEmpty(files)) {
console.error("No files found");
process.exit(1);
}
const commandFilePaths = await globby("commands/*.md", {
gitignore: true,
});
const commandMap = new Map<string, string>();
for (const commandFilePath of commandFilePaths) {
const commandName = path.parse(commandFilePath).name;
const commandContent = await readFile(commandFilePath, "utf-8");
commandMap.set(commandName, commandContent);
}
const validCommands = Array.from(commandMap.keys()) as [string, ...string[]];
const { object } = await generateObject({
model,
schema: z.object({
command: z.enum(validCommands).describe(`The command to execute}`),
filePath: z.enum(files).describe("The path to the files to prcess"),
}),
prompt: objectPrompt,
}).catch((err) => {
console.error(`Model generated an invalid reponse for ${objectPrompt}`);
process.exit(1);
});
type Command = (typeof validCommands)[number];
const isValidCommand = (cmd: unknown): cmd is Command => {
return (
typeof cmd === "string" &&
(validCommands as readonly string[]).includes(cmd)
);
};
if (!isValidCommand(object.command)) {
console.error("Invalid command");
process.exit(1);
}
const content = await readFile(object.filePath, "utf-8");
const textPrompt = `
<command>
${object.command}
</command>
<instructions>
${commandMap.get(object.command)}
</instructions>
<content>
${content}
</content>
`;
const { text } = await generateText({
model,
prompt: textPrompt,
});
console.log(text);