Merge pull request #41 from 0xturboblitz/master

feat: model selection logic

Author: samheutmaker

src/cli/commands/estimate/index.ts

@@ -14,6 +14,9 @@ export const estimate = async ({
   root,
   output,
   llms,
+  priority,
+  maxConcurrentCalls,
+  addQuestions,
   ignore,
   filePrompt,
   folderPrompt,
@@ -37,6 +40,9 @@ export const estimate = async ({
       root,
       output: json,
       llms,
+      priority,
+      maxConcurrentCalls,
+      addQuestions,
       ignore,
       filePrompt,
       folderPrompt,

src/cli/commands/index/index.ts

@@ -11,6 +11,9 @@ export const index = async ({
   root,
   output,
   llms,
+  priority,
+  maxConcurrentCalls,
+  addQuestions,
   ignore,
   filePrompt,
   folderPrompt,
@@ -35,6 +38,9 @@ export const index = async ({
     root,
     output: json,
     llms,
+    priority,
+    maxConcurrentCalls,
+    addQuestions,
     ignore,
     filePrompt,
     folderPrompt,
@@ -56,6 +62,9 @@ export const index = async ({
     root: json,
     output: markdown,
     llms,
+    priority,
+    maxConcurrentCalls,
+    addQuestions,
     ignore,
     filePrompt,
     folderPrompt,
@@ -73,6 +82,9 @@ export const index = async ({
     root: markdown,
     output: data,
     llms,
+    priority,
+    maxConcurrentCalls,
+    addQuestions,
     ignore,
     filePrompt,
     folderPrompt,

src/cli/commands/index/processRepository.ts

@@ -30,6 +30,7 @@ import {
   githubFolderUrl,
 } from '../../utils/FileUtil.js';
 import { models } from '../../utils/LLMUtil.js';
+import { selectModel } from './selectModel.js';
 
 export const processRepository = async (
   {
@@ -38,6 +39,9 @@ export const processRepository = async (
     root: inputRoot,
     output: outputRoot,
     llms,
+    priority,
+    maxConcurrentCalls,
+    addQuestions,
     ignore,
     filePrompt,
     folderPrompt,
@@ -47,8 +51,7 @@ export const processRepository = async (
   }: AutodocRepoConfig,
   dryRun?: boolean,
 ) => {
-  const encoding = encoding_for_model('gpt-3.5-turbo');
-  const rateLimit = new APIRateLimit(25);
+  const rateLimit = new APIRateLimit(maxConcurrentCalls);
 
   const callLLM = async (
     prompt: string,
@@ -91,6 +94,7 @@ export const processRepository = async (
 
     const markdownFilePath = path.join(outputRoot, filePath);
     const url = githubFileUrl(repositoryUrl, inputRoot, filePath, linkHosted);
+
     const summaryPrompt = createCodeFileSummary(
       projectName,
       projectName,
@@ -105,51 +109,28 @@ export const processRepository = async (
       contentType,
       targetAudience,
     );
-    const summaryLength = encoding.encode(summaryPrompt).length;
-    const questionLength = encoding.encode(questionsPrompt).length;
-    const max = Math.max(questionLength, summaryLength);
 
-    /**
-     * TODO: Encapsulate logic for selecting the best model
-     * TODO: Allow for different selection strategies based
-     * TODO: preference for cost/performace
-     * TODO: When this is re-written, it should use the correct
-     * TODO: TikToken encoding for each model
-     */
+    const prompts = addQuestions
+      ? [summaryPrompt, questionsPrompt]
+      : [summaryPrompt];
 
-    const model: LLMModelDetails | null = (() => {
-      if (
-        models[LLMModels.GPT3].maxLength > max &&
-        llms.includes(LLMModels.GPT3)
-      ) {
-        return models[LLMModels.GPT3];
-      } else if (
-        models[LLMModels.GPT4].maxLength > max &&
-        llms.includes(LLMModels.GPT4)
-      ) {
-        return models[LLMModels.GPT4];
-      } else if (
-        models[LLMModels.GPT432k].maxLength > max &&
-        llms.includes(LLMModels.GPT432k)
-      ) {
-        return models[LLMModels.GPT432k];
-      } else {
-        return null;
-      }
-    })();
+    const model = selectModel(prompts, llms, models, priority);
 
     if (!isModel(model)) {
       // console.log(`Skipped ${filePath} | Length ${max}`);
       return;
     }
 
+    const encoding = encoding_for_model(model.name);
+    const summaryLength = encoding.encode(summaryPrompt).length;
+    const questionLength = encoding.encode(questionsPrompt).length;
+
     try {
       if (!dryRun) {
         /** Call LLM */
-        const [summary, questions] = await Promise.all([
-          callLLM(summaryPrompt, model.llm),
-          callLLM(questionsPrompt, model.llm),
-        ]);
+        const response = await Promise.all(
+          prompts.map(async (prompt) => callLLM(prompt, model.llm)),
+        );
 
         /**
          * Create file and save to disk
@@ -158,8 +139,8 @@ export const processRepository = async (
           fileName,
           filePath,
           url,
-          summary,
-          questions,
+          summary: response[0],
+          questions: addQuestions ? response[1] : '',
           checksum: newChecksum,
         };
 
@@ -186,7 +167,8 @@ export const processRepository = async (
       /**
        * Track usage for end of run summary
        */
-      model.inputTokens += summaryLength + questionLength;
+      model.inputTokens += summaryLength;
+      if (addQuestions) model.inputTokens += questionLength;
       model.total++;
       model.outputTokens += 1000;
       model.succeeded++;
@@ -236,7 +218,12 @@ export const processRepository = async (
     }
 
     // eslint-disable-next-line prettier/prettier
-    const url = githubFolderUrl(repositoryUrl, inputRoot, folderPath, linkHosted);
+    const url = githubFolderUrl(
+      repositoryUrl,
+      inputRoot,
+      folderPath,
+      linkHosted,
+    );
     const allFiles: (FileSummary | null)[] = await Promise.all(
       contents.map(async (fileName) => {
         const entryPath = path.join(folderPath, fileName);
@@ -279,18 +266,24 @@ export const processRepository = async (
         (folder): folder is FolderSummary => folder !== null,
       );
 
-      const summary = await callLLM(
-        folderSummaryPrompt(
-          folderPath,
-          projectName,
-          files,
-          folders,
-          contentType,
-          folderPrompt,
-        ),
-        models[LLMModels.GPT4].llm,
+      const summaryPrompt = folderSummaryPrompt(
+        folderPath,
+        projectName,
+        files,
+        folders,
+        contentType,
+        folderPrompt,
       );
 
+      const model = selectModel([summaryPrompt], llms, models, priority);
+
+      if (!isModel(model)) {
+        // console.log(`Skipped ${filePath} | Length ${max}`);
+        return;
+      }
+
+      const summary = await callLLM(summaryPrompt, model.llm);
+
       const folderSummary: FolderSummary = {
         folderName,
         folderPath,

src/cli/commands/index/selectModel.ts

@@ -0,0 +1,57 @@
+import { encoding_for_model } from '@dqbd/tiktoken';
+import { LLMModelDetails, LLMModels, Priority } from '../../../types.js';
+
+export const selectModel = (
+  prompts: string[],
+  llms: LLMModels[],
+  models: Record<LLMModels, LLMModelDetails>,
+  priority: Priority,
+): LLMModelDetails | null => {
+  if (priority === Priority.COST) {
+    if (
+      llms.includes(LLMModels.GPT3) &&
+      models[LLMModels.GPT3].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT3)
+    ) {
+      return models[LLMModels.GPT3];
+    } else if (
+      llms.includes(LLMModels.GPT4) &&
+      models[LLMModels.GPT4].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT4)
+    ) {
+      return models[LLMModels.GPT4];
+    } else if (
+      llms.includes(LLMModels.GPT432k) &&
+      models[LLMModels.GPT432k].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT432k)
+    ) {
+      return models[LLMModels.GPT432k];
+    } else {
+      return null;
+    }
+  } else {
+    if (llms.includes(LLMModels.GPT4)) {
+      if (
+        models[LLMModels.GPT4].maxLength >
+        getMaxPromptLength(prompts, LLMModels.GPT4)
+      ) {
+        return models[LLMModels.GPT4];
+      } else if (
+        llms.includes(LLMModels.GPT432k) &&
+        models[LLMModels.GPT432k].maxLength >
+          getMaxPromptLength(prompts, LLMModels.GPT432k)
+      ) {
+        return models[LLMModels.GPT432k];
+      } else {
+        return null;
+      }
+    } else {
+      return models[LLMModels.GPT3];
+    }
+  }
+
+  function getMaxPromptLength(prompts: string[], model: LLMModels) {
+    const encoding = encoding_for_model(model);
+    return Math.max(...prompts.map((p) => encoding.encode(p).length));
+  }
+};

src/cli/commands/init/index.ts

@@ -2,7 +2,7 @@ import chalk from 'chalk';
 import inquirer from 'inquirer';
 import fs from 'node:fs';
 import path from 'node:path';
-import { AutodocRepoConfig, LLMModels } from '../../../types.js';
+import { AutodocRepoConfig, LLMModels, Priority } from '../../../types.js';
 
 export const makeConfigTemplate = (
   config?: AutodocRepoConfig,
@@ -16,6 +16,9 @@ export const makeConfigTemplate = (
       config?.llms?.length ?? 0 > 0
         ? (config as AutodocRepoConfig).llms
         : [LLMModels.GPT3],
+    priority: Priority.COST,
+    maxConcurrentCalls: 25,
+    addQuestions: true,
     ignore: [
       '.*',
       '*package-lock.json',

src/cli/utils/LLMUtil.ts

@@ -4,7 +4,7 @@ import { LLMModelDetails, LLMModels } from '../../types.js';
 export const models: Record<LLMModels, LLMModelDetails> = {
   [LLMModels.GPT3]: {
     name: LLMModels.GPT3,
-    inputCostPer1KTokens: 0.002,
+    inputCostPer1KTokens: 0.0015,
     outputCostPer1KTokens: 0.002,
     maxLength: 3050,
     llm: new OpenAIChat({
@@ -61,7 +61,7 @@ export const printModelDetails = (models: LLMModelDetails[]): void => {
       Failed: model.failed,
       Tokens: model.inputTokens + model.outputTokens,
       Cost:
-        (model.total / 1000) * model.inputCostPer1KTokens +
+        (model.inputTokens / 1000) * model.inputCostPer1KTokens +
         (model.outputTokens / 1000) * model.outputCostPer1KTokens,
     };
   });
@@ -95,7 +95,7 @@ export const totalIndexCostEstimate = (models: LLMModelDetails[]): number => {
   const totalCost = models.reduce((cur, model) => {
     return (
       cur +
-      (model.total / 1000) * model.inputCostPer1KTokens +
+      (model.inputTokens / 1000) * model.inputCostPer1KTokens +
       (model.outputTokens / 1000) * model.outputCostPer1KTokens
     );
   }, 0);

src/types.ts

@@ -10,6 +10,9 @@ export type AutodocRepoConfig = {
   root: string;
   output: string;
   llms: LLMModels[];
+  priority: Priority;
+  maxConcurrentCalls: number;
+  addQuestions: boolean;
   ignore: string[];
   filePrompt: string;
   folderPrompt: string;
@@ -24,7 +27,7 @@ export type FileSummary = {
   filePath: string;
   url: string;
   summary: string;
-  questions: string;
+  questions?: string;
   checksum: string;
 };
 
@@ -96,3 +99,8 @@ export type LLMModelDetails = {
   failed: number;
   total: number;
 };
+
+export enum Priority {
+  COST = 'cost',
+  PERFORMANCE = 'performance',
+}