Merge branch 'main' of https://github.com/continuedev/continue into snyk-upgrade-3095d9fd5e6b0170e79d83641d320c82

Author: RomneyDa

core/config/yaml/loadYaml.ts

@@ -1,9 +1,9 @@
 import {
   AssistantUnrolled,
+  AssistantUnrolledNonNullable,
   BLOCK_TYPES,
   ConfigResult,
   ConfigValidationError,
-  isAssistantUnrolledNonNullable,
   mergeConfigYamlRequestOptions,
   mergeUnrolledAssistants,
   ModelRole,
@@ -145,8 +145,8 @@ async function loadConfigYaml(options: {
     }
   }
 
-  if (config && isAssistantUnrolledNonNullable(config)) {
-    errors.push(...validateConfigYaml(config));
+  if (config) {
+    errors.push(...validateConfigYaml(nonNullifyConfigYaml(config)));
   }
 
   if (errors?.some((error) => error.fatal)) {
@@ -165,15 +165,30 @@ async function loadConfigYaml(options: {
   };
 }
 
+function nonNullifyConfigYaml(
+  unrolledAssistant: AssistantUnrolled,
+): AssistantUnrolledNonNullable {
+  return {
+    ...unrolledAssistant,
+    data: unrolledAssistant.data?.filter((k) => !!k),
+    context: unrolledAssistant.context?.filter((k) => !!k),
+    docs: unrolledAssistant.docs?.filter((k) => !!k),
+    mcpServers: unrolledAssistant.mcpServers?.filter((k) => !!k),
+    models: unrolledAssistant.models?.filter((k) => !!k),
+    prompts: unrolledAssistant.prompts?.filter((k) => !!k),
+    rules: unrolledAssistant.rules?.filter((k) => !!k).map((k) => k!),
+  };
+}
+
 export async function configYamlToContinueConfig(options: {
-  config: AssistantUnrolled;
+  unrolledAssistant: AssistantUnrolled;
   ide: IDE;
   ideInfo: IdeInfo;
   uniqueId: string;
   llmLogger: ILLMLogger;
   workOsAccessToken: string | undefined;
 }): Promise<{ config: ContinueConfig; errors: ConfigValidationError[] }> {
-  let { config, ide, ideInfo, uniqueId, llmLogger } = options;
+  let { unrolledAssistant, ide, ideInfo, uniqueId, llmLogger } = options;
 
   const localErrors: ConfigValidationError[] = [];
 
@@ -203,22 +218,10 @@ export async function configYamlToContinueConfig(options: {
       subagent: null,
     },
     rules: [],
-    requestOptions: { ...config.requestOptions },
+    requestOptions: { ...unrolledAssistant.requestOptions },
   };
 
-  // Right now, if there are any missing packages in the config, then we will just throw an error
-  if (!isAssistantUnrolledNonNullable(config)) {
-    return {
-      config: continueConfig,
-      errors: [
-        {
-          message:
-            "Failed to load config due to missing blocks, see which blocks are missing below",
-          fatal: true,
-        },
-      ],
-    };
-  }
+  const config = nonNullifyConfigYaml(unrolledAssistant);
 
   for (const rule of config.rules ?? []) {
     const convertedRule = convertYamlRuleToContinueRule(rule);
@@ -447,7 +450,7 @@ export async function loadContinueConfigFromYaml(options: {
 
   const { config: continueConfig, errors: localErrors } =
     await configYamlToContinueConfig({
-      config: configYamlResult.config,
+      unrolledAssistant: configYamlResult.config,
       ide,
       ideInfo,
       uniqueId,

core/llm/countTokens.ts

@@ -18,6 +18,7 @@ import {
 import { renderChatMessage } from "../util/messageContent.js";
 import { AsyncEncoder, LlamaAsyncEncoder } from "./asyncEncoder.js";
 import { DEFAULT_PRUNING_LENGTH } from "./constants.js";
+import { getAdjustedTokenCountFromModel } from "./getAdjustedTokenCount.js";
 import llamaTokenizer from "./llamaTokenizer.js";
 interface Encoding {
   encode: Tiktoken["encode"];
@@ -114,8 +115,9 @@ function countTokens(
   modelName = "llama2",
 ): number {
   const encoding = encodingForModel(modelName);
+  let baseTokens = 0;
   if (Array.isArray(content)) {
-    return content.reduce((acc, part) => {
+    baseTokens = content.reduce((acc, part) => {
       return (
         acc +
         (part.type === "text"
@@ -124,8 +126,9 @@ function countTokens(
       );
     }, 0);
   } else {
-    return encoding.encode(content ?? "", "all", []).length;
+    baseTokens = encoding.encode(content ?? "", "all", []).length;
   }
+  return getAdjustedTokenCountFromModel(baseTokens, modelName);
 }
 
 // https://community.openai.com/t/how-to-calculate-the-tokens-when-using-function-call/266573/10

core/llm/getAdjustedTokenCount.test.ts

@@ -0,0 +1,48 @@
+import { getAdjustedTokenCountFromModel } from "./getAdjustedTokenCount";
+
+describe("getAdjustedTokenCountFromModel", () => {
+  it("should return base tokens for non-special models", () => {
+    expect(getAdjustedTokenCountFromModel(100, "gpt-4")).toBe(100);
+    expect(getAdjustedTokenCountFromModel(100, "llama2")).toBe(100);
+    expect(getAdjustedTokenCountFromModel(100, "random-model")).toBe(100);
+  });
+
+  it("should apply multiplier for Claude models", () => {
+    expect(getAdjustedTokenCountFromModel(100, "claude-3-opus")).toBe(123);
+    expect(getAdjustedTokenCountFromModel(100, "claude-3.5-sonnet")).toBe(123);
+    expect(getAdjustedTokenCountFromModel(100, "CLAUDE-2")).toBe(123);
+    expect(getAdjustedTokenCountFromModel(50, "claude")).toBe(62); // 50 * 1.23 = 61.5, ceiled to 62
+  });
+
+  it("should apply multiplier for Gemini models", () => {
+    expect(getAdjustedTokenCountFromModel(100, "gemini-pro")).toBe(118);
+    expect(getAdjustedTokenCountFromModel(100, "gemini-1.5-pro")).toBe(118);
+    expect(getAdjustedTokenCountFromModel(100, "GEMINI-flash")).toBe(118);
+    expect(getAdjustedTokenCountFromModel(50, "gemini")).toBe(59); // 50 * 1.18 = 59
+  });
+
+  it("should apply multiplier for Mistral family models", () => {
+    expect(getAdjustedTokenCountFromModel(100, "mistral-large")).toBe(126);
+    expect(getAdjustedTokenCountFromModel(100, "mixtral-8x7b")).toBe(126);
+    expect(getAdjustedTokenCountFromModel(100, "devstral")).toBe(126);
+    expect(getAdjustedTokenCountFromModel(100, "CODESTRAL")).toBe(126);
+    expect(getAdjustedTokenCountFromModel(50, "mistral")).toBe(63); // 50 * 1.26 = 63
+  });
+
+  it("should handle edge cases", () => {
+    expect(getAdjustedTokenCountFromModel(0, "claude")).toBe(0);
+    expect(getAdjustedTokenCountFromModel(1, "gemini")).toBe(2); // 1 * 1.18 = 1.18, ceiled to 2
+    expect(getAdjustedTokenCountFromModel(1000, "mixtral")).toBe(1260);
+  });
+
+  it("should handle empty or undefined model names", () => {
+    expect(getAdjustedTokenCountFromModel(100, "")).toBe(100);
+    expect(getAdjustedTokenCountFromModel(100, undefined as any)).toBe(100);
+  });
+
+  it("should be case-insensitive", () => {
+    expect(getAdjustedTokenCountFromModel(100, "ClAuDe-3-OpUs")).toBe(123);
+    expect(getAdjustedTokenCountFromModel(100, "GeMiNi-PrO")).toBe(118);
+    expect(getAdjustedTokenCountFromModel(100, "MiXtRaL")).toBe(126);
+  });
+});

core/llm/getAdjustedTokenCount.ts

@@ -0,0 +1,38 @@
+// Importing a bunch of tokenizers can be very resource intensive (MB-scale per tokenizer)
+// Using token counting APIs (e.g. for anthropic) can be complicated and unreliable in many environments
+// So for now we will just use super fast gpt-tokenizer and apply safety buffers
+// I'm using rough estimates from this article to apply safety buffers to common tokenizers
+// which will have HIGHER token counts than gpt. Roughly using token ratio from article + 10%
+// https://medium.com/@disparate-ai/not-all-tokens-are-created-equal-7347d549af4d
+const ANTHROPIC_TOKEN_MULTIPLIER = 1.23;
+const GEMINI_TOKEN_MULTIPLIER = 1.18;
+const MISTRAL_TOKEN_MULTIPLIER = 1.26;
+
+/**
+ * Adjusts token count based on model-specific tokenizer differences.
+ * Since we use llama tokenizer (~= gpt tokenizer) for all models, we apply
+ * multipliers for models known to have higher token counts.
+ *
+ * @param baseTokens - Token count from llama/gpt tokenizer
+ * @param modelName - Name of the model
+ * @returns Adjusted token count with safety buffer
+ */
+export function getAdjustedTokenCountFromModel(
+  baseTokens: number,
+  modelName: string,
+): number {
+  let multiplier = 1;
+  const lowerModelName = modelName?.toLowerCase() ?? "";
+  if (lowerModelName.includes("claude")) {
+    multiplier = ANTHROPIC_TOKEN_MULTIPLIER;
+  } else if (lowerModelName.includes("gemini")) {
+    multiplier = GEMINI_TOKEN_MULTIPLIER;
+  } else if (
+    lowerModelName.includes("stral") ||
+    lowerModelName.includes("mixtral")
+  ) {
+    // Mistral family models: mistral, mixtral, codestral, devstral, etc
+    multiplier = MISTRAL_TOKEN_MULTIPLIER;
+  }
+  return Math.ceil(baseTokens * multiplier);
+}