From dd46bcabc27370c776ecb6634d5b6de09d89a151 Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Wed, 23 Apr 2025 17:31:29 +0530
Subject: [PATCH] Track gpt-4.1 model costs. Set prompt size of new gemini,
 openai models

---
 src/khoj/processor/conversation/utils.py | 7 +++++++
 src/khoj/utils/constants.py              | 3 +++
 2 files changed, 10 insertions(+)

diff --git a/src/khoj/processor/conversation/utils.py b/src/khoj/processor/conversation/utils.py
index b8eea907..9601a5cd 100644
--- a/src/khoj/processor/conversation/utils.py
+++ b/src/khoj/processor/conversation/utils.py
@@ -53,10 +53,17 @@ model_to_prompt_size = {
     # OpenAI Models
     "gpt-4o": 60000,
     "gpt-4o-mini": 120000,
+    "gpt-4.1": 60000,
+    "gpt-4.1-mini": 120000,
+    "gpt-4.1-nano": 120000,
     "o1": 20000,
+    "o3": 30000,
     "o1-mini": 60000,
     "o3-mini": 60000,
+    "o4-mini": 60000,
     # Google Models
+    "gemini-2.5-flash-preview-04-17": 120000,
+    "gemini-2.5-pro-preview-03-25": 60000,
     "gemini-2.0-flash": 120000,
     "gemini-2.0-flash-lite": 120000,
     "gemini-1.5-flash": 120000,
diff --git a/src/khoj/utils/constants.py b/src/khoj/utils/constants.py
index daf8469f..a7f460df 100644
--- a/src/khoj/utils/constants.py
+++ b/src/khoj/utils/constants.py
@@ -38,6 +38,9 @@ model_to_cost: Dict[str, Dict[str, float]] = {
     # OpenAI Pricing: https://openai.com/api/pricing/
     "gpt-4o": {"input": 2.50, "output": 10.00},
     "gpt-4o-mini": {"input": 0.15, "output": 0.60},
+    "gpt-4.1": {"input": 2.00, "output": 8.00},
+    "gpt-4.1-mini": {"input": 0.40, "output": 1.60},
+    "gpt-4.1-nano": {"input": 0.10, "output": 0.40},
     "o1": {"input": 15.0, "output": 60.00},
     "o3": {"input": 10.0, "output": 40.00},
     "o1-mini": {"input": 3.0, "output": 12.0},