heaven_base.utils.token_counter

  1"""
  2Token counting utilities for precise context management.
  3
  4Provides accurate token counting for OpenAI models using tiktoken,
  5with fallback approximations for other models.
  6"""
  7
  8import tiktoken
  9from typing import Optional, List, Union
 10from langchain_core.messages import BaseMessage
 11
 12
 13def count_tokens(
 14    text: str, 
 15    model: str = "gpt-4o-mini"
 16) -> int:
 17    """
 18    Count tokens in text using tiktoken for OpenAI models.
 19    
 20    Args:
 21        text: Text to count tokens for
 22        model: Model name for encoding selection
 23        
 24    Returns:
 25        int: Number of tokens in the text
 26    """
 27    try:
 28        # Map model names to tiktoken encodings
 29        if "gpt-4" in model.lower():
 30            encoding = tiktoken.encoding_for_model("gpt-4")
 31        elif "gpt-3.5" in model.lower():
 32            encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
 33        elif "o1" in model.lower() or "o3" in model.lower() or "o4" in model.lower():
 34            # Use gpt-4 encoding for o-series models
 35            encoding = tiktoken.encoding_for_model("gpt-4")
 36        else:
 37            # Default to gpt-4 encoding for other models (Claude, Gemini, etc.)
 38            encoding = tiktoken.encoding_for_model("gpt-4")
 39        
 40        return len(encoding.encode(text))
 41        
 42    except Exception as e:
 43        # Fallback to rough approximation if tiktoken fails
 44        # OpenAI's rule of thumb: ~4 characters per token
 45        return max(1, len(text) // 4)
 46
 47
 48def count_tokens_in_messages(
 49    messages: List[Union[BaseMessage, dict]], 
 50    model: str = "gpt-4o-mini"
 51) -> int:
 52    """
 53    Count tokens in a list of messages.
 54    
 55    Args:
 56        messages: List of LangChain messages or dict messages
 57        model: Model name for encoding selection
 58        
 59    Returns:
 60        int: Total number of tokens in all messages
 61    """
 62    total_tokens = 0
 63    
 64    for message in messages:
 65        if isinstance(message, BaseMessage):
 66            # LangChain message
 67            text = str(message.content)
 68        elif isinstance(message, dict):
 69            # Dict message (uni-api format)
 70            text = message.get("content", "")
 71        else:
 72            # Fallback: convert to string
 73            text = str(message)
 74        
 75        total_tokens += count_tokens(text, model)
 76    
 77    return total_tokens
 78
 79
 80def estimate_tokens_for_model(text: str, model: str) -> int:
 81    """
 82    Estimate tokens for different model families with family-specific approximations.
 83    
 84    Args:
 85        text: Text to estimate tokens for
 86        model: Model name
 87        
 88    Returns:
 89        int: Estimated number of tokens
 90    """
 91    model_lower = model.lower()
 92    
 93    if "gpt" in model_lower or "openai" in model_lower:
 94        # Use tiktoken for OpenAI models
 95        return count_tokens(text, model)
 96    elif "claude" in model_lower:
 97        # Claude tokenization is similar to GPT-4 but slightly different
 98        # Rough approximation: 3.8 chars per token
 99        return max(1, len(text) // 4)
100    elif "gemini" in model_lower:
101        # Gemini has different tokenization
102        # Rough approximation: 4.2 chars per token  
103        return max(1, len(text) // 4)
104    else:
105        # Generic fallback
106        return max(1, len(text) // 4)
107
108
109class TokenBudget:
110    """
111    Utility class for managing token budgets across conversation sections.
112    """
113    
114    def __init__(self, total_budget: int, model: str = "gpt-4o-mini"):
115        self.total_budget = total_budget
116        self.model = model
117        self.allocations = {}
118        self.used_tokens = {}
119    
120    def allocate(self, section: str, tokens: int) -> None:
121        """Allocate tokens to a section."""
122        self.allocations[section] = tokens
123        if section not in self.used_tokens:
124            self.used_tokens[section] = 0
125    
126    def use_tokens(self, section: str, text: str) -> bool:
127        """
128        Use tokens for a section. Returns True if within budget.
129        
130        Args:
131            section: Section name
132            text: Text to count tokens for
133            
134        Returns:
135            bool: True if within section budget, False otherwise
136        """
137        tokens_needed = count_tokens(text, self.model)
138        
139        if section not in self.allocations:
140            return False
141        
142        if self.used_tokens.get(section, 0) + tokens_needed <= self.allocations[section]:
143            self.used_tokens[section] = self.used_tokens.get(section, 0) + tokens_needed
144            return True
145        
146        return False
147    
148    def get_remaining(self, section: str) -> int:
149        """Get remaining tokens for a section."""
150        allocated = self.allocations.get(section, 0)
151        used = self.used_tokens.get(section, 0)
152        return max(0, allocated - used)
153    
154    def get_total_used(self) -> int:
155        """Get total tokens used across all sections."""
156        return sum(self.used_tokens.values())
157    
158    def get_total_remaining(self) -> int:
159        """Get total remaining budget."""
160        return max(0, self.total_budget - self.get_total_used())
161    
162    def reset_section(self, section: str) -> None:
163        """Reset usage for a section."""
164        self.used_tokens[section] = 0
165    
166    def get_budget_summary(self) -> dict:
167        """Get a summary of budget usage."""
168        return {
169            "total_budget": self.total_budget,
170            "total_used": self.get_total_used(),
171            "total_remaining": self.get_total_remaining(),
172            "sections": {
173                section: {
174                    "allocated": self.allocations.get(section, 0),
175                    "used": self.used_tokens.get(section, 0),
176                    "remaining": self.get_remaining(section)
177                }
178                for section in set(list(self.allocations.keys()) + list(self.used_tokens.keys()))
179            }
180        }
181
182
183# Convenience functions for common use cases
184def count_conversation_tokens(messages: List[Union[BaseMessage, dict]], model: str = "gpt-4o-mini") -> int:
185    """Count tokens in a full conversation."""
186    return count_tokens_in_messages(messages, model)
187
188
189def check_token_limit(text: str, limit: int, model: str = "gpt-4o-mini") -> bool:
190    """Check if text is within token limit."""
191    return count_tokens(text, model) <= limit
192
193
194def truncate_to_token_limit(text: str, limit: int, model: str = "gpt-4o-mini") -> str:
195    """
196    Truncate text to fit within token limit.
197    
198    This is a rough approximation - for precise truncation,
199    you'd need to decode tokens back to text.
200    """
201    current_tokens = count_tokens(text, model)
202    
203    if current_tokens <= limit:
204        return text
205    
206    # Rough truncation based on character ratio
207    char_ratio = len(text) / current_tokens
208    target_chars = int(limit * char_ratio * 0.9)  # 10% safety margin
209    
210    return text[:target_chars] + "..."
211
212
213if __name__ == "__main__":
214    # Example usage
215    test_text = "Hello world! This is a test of the token counting system."
216    
217    print(f"Text: {test_text}")
218    print(f"Tokens (GPT-4): {count_tokens(test_text, 'gpt-4o-mini')}")
219    print(f"Tokens (Claude): {estimate_tokens_for_model(test_text, 'claude-3-5-sonnet')}")
220    
221    # Test budget system
222    budget = TokenBudget(1000, "gpt-4o-mini")
223    budget.allocate("summary", 300)
224    budget.allocate("history", 600)
225    budget.allocate("prompt", 100)
226    
227    print(f"\nBudget summary: {budget.get_budget_summary()}")