From ded1db642c0102af7fffd24285425ddfbb8122ef Mon Sep 17 00:00:00 2001
From: Debanjum <debanjum@gmail.com>
Date: Fri, 30 May 2025 16:44:01 -0700
Subject: [PATCH] Get max context for user, operator model pair for context
 compression

---
 src/khoj/processor/operator/__init__.py       | 20 +++++++++++++++++--
 .../processor/operator/operator_agent_base.py |  5 ++---
 .../operator/operator_agent_binary.py         |  2 ++
 3 files changed, 22 insertions(+), 5 deletions(-)

diff --git a/src/khoj/processor/operator/__init__.py b/src/khoj/processor/operator/__init__.py
index 34f4ad7d..59bd04e9 100644
--- a/src/khoj/processor/operator/__init__.py
+++ b/src/khoj/processor/operator/__init__.py
@@ -63,15 +63,30 @@ async def operate_environment(
     chat_history = construct_chat_history_for_operator(conversation_log)
 
     # Initialize Agent
+    max_context = await ConversationAdapters.aget_max_context_size(reasoning_model, user) or 20000
     max_iterations = int(os.getenv("KHOJ_OPERATOR_ITERATIONS", 100))
     operator_agent: OperatorAgent
     if is_operator_model(reasoning_model.name) == ChatModel.ModelType.OPENAI:
         operator_agent = OpenAIOperatorAgent(
-            query, reasoning_model, environment_type, max_iterations, chat_history, previous_trajectory, tracer
+            query,
+            reasoning_model,
+            environment_type,
+            max_iterations,
+            max_context,
+            chat_history,
+            previous_trajectory,
+            tracer,
         )
     elif is_operator_model(reasoning_model.name) == ChatModel.ModelType.ANTHROPIC:
         operator_agent = AnthropicOperatorAgent(
-            query, reasoning_model, environment_type, max_iterations, chat_history, previous_trajectory, tracer
+            query,
+            reasoning_model,
+            environment_type,
+            max_iterations,
+            max_context,
+            chat_history,
+            previous_trajectory,
+            tracer,
         )
     else:
         grounding_model_name = "ui-tars-1.5"
@@ -88,6 +103,7 @@ async def operate_environment(
             grounding_model,
             environment_type,
             max_iterations,
+            max_context,
             chat_history,
             previous_trajectory,
             tracer,
diff --git a/src/khoj/processor/operator/operator_agent_base.py b/src/khoj/processor/operator/operator_agent_base.py
index a04b44d4..3528348a 100644
--- a/src/khoj/processor/operator/operator_agent_base.py
+++ b/src/khoj/processor/operator/operator_agent_base.py
@@ -34,6 +34,7 @@ class OperatorAgent(ABC):
         vision_model: ChatModel,
         environment_type: EnvironmentType,
         max_iterations: int,
+        max_context: int,
         chat_history: List[AgentMessage] = [],
         previous_trajectory: Optional[OperatorRun] = None,
         tracer: dict = {},
@@ -56,9 +57,7 @@ class OperatorAgent(ABC):
         # Context compression parameters
         self.context_compress_trigger = 2e3  # heuristic to determine compression trigger
         # turns after which compression triggered. scales with model max context size. Minimum 5 turns.
-        self.message_limit = 2 * max(
-            5, int(self.vision_model.subscribed_max_prompt_size / self.context_compress_trigger)
-        )
+        self.message_limit = 2 * max(5, int(max_context / self.context_compress_trigger))
         # compression ratio determines how many messages to compress down to one
         # e.g. if 5 messages, a compress ratio of 4/5 means compress 5 messages into 1 + keep 1 uncompressed
         self.message_compress_ratio = 4 / 5
diff --git a/src/khoj/processor/operator/operator_agent_binary.py b/src/khoj/processor/operator/operator_agent_binary.py
index ade98176..7b1ec086 100644
--- a/src/khoj/processor/operator/operator_agent_binary.py
+++ b/src/khoj/processor/operator/operator_agent_binary.py
@@ -40,6 +40,7 @@ class BinaryOperatorAgent(OperatorAgent):
         grounding_model: ChatModel,
         environment_type: EnvironmentType,
         max_iterations: int,
+        max_context: int,
         chat_history: List[AgentMessage] = [],
         previous_trajectory: Optional[OperatorRun] = None,
         tracer: dict = {},
@@ -49,6 +50,7 @@ class BinaryOperatorAgent(OperatorAgent):
             reasoning_model,
             environment_type,
             max_iterations,
+            max_context,
             chat_history,
             previous_trajectory,
             tracer,