[GenAI] Dynamically Inferring File Types from User Prompts in Local AI Scripts

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);