token_tracking_logic

agent.py

@@ -21,6 +21,9 @@ from smolagents import (
     FinalAnswerTool,
 )
 
+# LiteLLM
+from litellm import token_counter
+
 # Import your custom tools (to be used in app, not in local notebook)
 from tools.gemini_native_tools import analyze_video, analyze_image, analyze_audio
 from tools.download_file       import download_file_from_url
@@ -183,12 +186,20 @@ class GeminiAgent:
             model          = self.model,
             tools          = self.tools,
             add_base_tools = True,        # probably redundant, but it does not hurt
-            max_steps      = 6, 
+            max_steps      = 4, 
             additional_authorized_imports = AUTHORIZED_IMPORTS,
             verbosity_level = 1,
-            max_print_outputs_length=1_000_000
+            max_print_outputs_length=2_000
             )
         
+        def check_token_safety(self, question):
+            messages = self.gemini_agent.memory.get_messages()
+            count = token_counter(model="gemma3", messages=messages)
+
+            print(f"📊 Current Context Size: {count} tokens")
+            return count
+
+        
         print(f"✅ Gemini agent initialized with model: {model_id}")
                 
     def __call__(self, question: str, file_path: Optional[str] = None) -> str:

app.py

@@ -110,6 +110,11 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
     # 3. Run your Agent
     results_log = []
     answers_payload = []
+
+    # track tokens used in the LAST 60 seconds
+    tokens_used_this_minute = 0
+    TPM_LIMIT = 14000 # Keep a 1k safety buffer
+
     print(f"Running agent on {len(questions_data)} questions...")
     for item in questions_data:
         # --- Check the "Kill Switch" at the start of every loop iteration ---
@@ -133,6 +138,14 @@ def run_and_submit_all( profile: gr.OAuthProfile | None):
             continue
 
         try:
+            current_tokens = agent.check_token_safety(question_text)
+            if (tokens_used_this_minute + current_tokens) > TPM_LIMIT:
+                print(f"⏳TPM limit imminent ({tokens_used_this_minute}/{TPM_LIMIT}) Cooling down...")
+                time.sleep(60)
+                tokens_used_this_minute = 0 # Reset
+            last_usage = agent.model.last_input_token_count + agent.model.last_output_token_count
+            tokens_used_this_minute += last_usage
+
             if agent_type == "GeminiAgent" and request_count >= MAX_TOTAL_REQUESTS:
                 raise RateLimitError("Global request cap reached")