1 files changed, 28 insertions, 0 deletions
diff --git a/packages/core/src/core/tokenLimits.ts b/packages/core/src/core/tokenLimits.ts
new file mode 100644
index 00000000..56882c21
--- /dev/null
+++ b/packages/core/src/core/tokenLimits.ts
@@ -0,0 +1,28 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+type Model = string;
+type TokenCount = number;
+
+export const DEFAULT_TOKEN_LIMIT = 1_048_576;
+
+export function tokenLimit(model: Model): TokenCount {
+  // Add other models as they become relevant or if specified by config
+  // Pulled from https://ai.google.dev/gemini-api/docs/models
+  switch (model) {
+    case 'gemini-1.5-pro':
+      return 2_097_152;
+    case 'gemini-1.5-flash':
+    case 'gemini-2.5-pro-preview-05-06':
+    case 'gemini-2.5-flash-preview-05-20':
+    case 'gemini-2.0-flash':
+      return 1_048_576;
+    case 'gemini-2.0-flash-preview-image-generation':
+      return 32_000;
+    default:
+      return DEFAULT_TOKEN_LIMIT;
+  }
+}