fix(token-count): include Antigravity preprompt tokens in count

Token counting endpoints (/v1/token-count and /v1/messages/count_tokens)
were returning inaccurate counts because they didn't include the
Antigravity preprompts that get injected during actual API calls.

- Add get_antigravity_preprompt_text() helper to expose preprompt text
- Update RotatingClient.token_count() to add preprompt tokens for
  Antigravity provider models

Signed-off-by: Moeeze Hassan <[email protected]>

Author: FammasMaz

src/rotator_library/client.py

@@ -2464,21 +2464,46 @@ def aembedding(
         )
 
     def token_count(self, **kwargs) -> int:
-        """Calculates the number of tokens for a given text or list of messages."""
+        """Calculates the number of tokens for a given text or list of messages.
+
+        For Antigravity provider models, this also includes the preprompt tokens
+        that get injected during actual API calls (agent instruction + identity override).
+        This ensures token counts match actual usage.
+        """
         kwargs = self._convert_model_params(**kwargs)
         model = kwargs.get("model")
         text = kwargs.get("text")
         messages = kwargs.get("messages")
 
         if not model:
             raise ValueError("'model' is a required parameter.")
+
+        # Calculate base token count
         if messages:
-            return token_counter(model=model, messages=messages)
+            base_count = token_counter(model=model, messages=messages)
         elif text:
-            return token_counter(model=model, text=text)
+            base_count = token_counter(model=model, text=text)
         else:
             raise ValueError("Either 'text' or 'messages' must be provided.")
 
+        # Add preprompt tokens for Antigravity provider
+        # The Antigravity provider injects system instructions during actual API calls,
+        # so we need to account for those tokens in the count
+        provider = model.split("/")[0] if "/" in model else ""
+        if provider == "antigravity":
+            try:
+                from .providers.antigravity_provider import get_antigravity_preprompt_text
+
+                preprompt_text = get_antigravity_preprompt_text()
+                if preprompt_text:
+                    preprompt_tokens = token_counter(model=model, text=preprompt_text)
+                    base_count += preprompt_tokens
+            except ImportError:
+                # Provider not available, skip preprompt token counting
+                pass
+
+        return base_count
+
     async def get_available_models(self, provider: str) -> List[str]:
         """Returns a list of available models for a specific provider, with caching."""
         lib_logger.info(f"Getting available models for provider: {provider}")

src/rotator_library/providers/antigravity_provider.py

@@ -449,6 +449,42 @@ def _is_valid_thinking_signature(signature):
 # =============================================================================
 
 
+def get_antigravity_preprompt_text() -> str:
+    """
+    Get the combined Antigravity preprompt text that gets injected into requests.
+
+    This function returns the exact text that gets prepended to system instructions
+    during actual API calls. It respects the current configuration settings:
+    - PREPEND_INSTRUCTION: Whether to include any preprompt at all
+    - USE_SHORT_ANTIGRAVITY_PROMPTS: Whether to use short or full versions
+    - INJECT_IDENTITY_OVERRIDE: Whether to include the identity override
+
+    This is useful for accurate token counting - the token count endpoints should
+    include these preprompts to match what actually gets sent to the API.
+
+    Returns:
+        The combined preprompt text, or empty string if prepending is disabled.
+    """
+    if not PREPEND_INSTRUCTION:
+        return ""
+
+    # Choose prompt versions based on USE_SHORT_ANTIGRAVITY_PROMPTS setting
+    if USE_SHORT_ANTIGRAVITY_PROMPTS:
+        agent_instruction = ANTIGRAVITY_AGENT_SYSTEM_INSTRUCTION_SHORT
+        override_instruction = ANTIGRAVITY_IDENTITY_OVERRIDE_INSTRUCTION_SHORT
+    else:
+        agent_instruction = ANTIGRAVITY_AGENT_SYSTEM_INSTRUCTION
+        override_instruction = ANTIGRAVITY_IDENTITY_OVERRIDE_INSTRUCTION
+
+    # Build the combined preprompt
+    parts = [agent_instruction]
+
+    if INJECT_IDENTITY_OVERRIDE:
+        parts.append(override_instruction)
+
+    return "\n".join(parts)
+
+
 def _sanitize_headers(headers: Dict[str, str]) -> Dict[str, str]:
     """
     Strip identifiable client headers for privacy/security.