""" title: Token Budget Context Filter author: aaron-helper version: 0.2 """ from pydantic import BaseModel, Field from typing import Optional, List, Dict, Any import json try: import tiktoken _enc_cache: Dict[str, Any] = {} def get_encoder(model: str): # pick a close encoder; for GPT-5 use cl100k_base fallback if model not in _enc_cache: try: _enc_cache[model] = tiktoken.encoding_for_model(model) except Exception: _enc_cache[model] = tiktoken.get_encoding("cl100k_base") return _enc_cache[model] def count_tokens(model: str, text: str) -> int: return len(get_encoder(model).encode(text)) except Exception: # Fallback: rough char/4 estimate if tiktoken isn't available def count_tokens(model: str, text: str) -> int: return max(1, len(text) // 4) class Filter: class Valves(BaseModel): priority: int = Field(0, description="Execution priority.") token_budget: int = Field(110000, description="Max input tokens allowed.") reserve_for_completion: int = Field(8000, description="Headroom for the reply.") min_last_messages: int = Field( 2, description="Always keep at least this many last messages (assistant+user).", ) drop_tool_messages: bool = Field( True, description="Prefer to drop tool/function messages first." ) drop_images: bool = Field( True, description="Prefer to drop image parts first (if present)." ) class UserValves(BaseModel): pass def __init__(self): self.valves = self.Valves() def _msg_text(self, m: Dict[str, Any]) -> str: role = m.get("role", "") content = m.get("content", "") # Handle non-string content (tools, images, dicts) if isinstance(content, list): # Strip images if requested; keep text parts = [] for part in content: if isinstance(part, dict) and part.get("type") in ( "image_url", "image", "input_image", ): if self.valves.drop_images: continue parts.append( json.dumps(part, ensure_ascii=False) if not isinstance(part, str) else part ) content = "\n".join(parts) elif not isinstance(content, str): content = json.dumps(content, ensure_ascii=False) return f"{role}: {content}" def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict: model = body.get("model", "gpt-5-chat-latest") messages: List[Dict[str, Any]] = body.get("messages", []) # Always keep one system prompt (first encountered) system_idx = next( (i for i, m in enumerate(messages) if m.get("role") == "system"), None ) system_msg = messages[system_idx] if system_idx is not None else None # Build a working set without the system message working = [m for i, m in enumerate(messages) if i != system_idx] # Prefer to drop tool/function messages early (they get big) if self.valves.drop_tool_messages: tools = [m for m in working if m.get("role") in ("tool", "function")] non_tools = [ m for m in working if m.get("role") not in ("tool", "function") ] else: tools, non_tools = [], working # Start from the most recent tail and grow backward under the token budget tail: List[Dict[str, Any]] = [] # Always seed with a minimum number of last messages (from non-tools) min_keep = max(0, self.valves.min_last_messages) seed = non_tools[-min_keep:] if min_keep > 0 else [] remaining_non_tools = non_tools[: len(non_tools) - len(seed)] tail = seed[:] # Compute current token usage with system + seed def total_tokens(msgs: List[Dict[str, Any]]) -> int: txts = [self._msg_text(m) for m in msgs] joined = "\n\n".join(txts) return count_tokens(model, joined) budget = max( 1000, self.valves.token_budget - max(0, self.valves.reserve_for_completion) ) chosen: List[Dict[str, Any]] = [] # Add recent non-tool messages backward until we hit the budget for m in reversed(remaining_non_tools): candidate = ([system_msg] if system_msg else []) + [m] + chosen + tail if total_tokens(candidate) <= budget: chosen.insert(0, m) else: break # If room remains, try to add recent tool messages if tools: for m in reversed(tools): candidate = ([system_msg] if system_msg else []) + chosen + [m] + tail if total_tokens(candidate) <= budget: tail.insert(0, m) # put tools just before the final tail else: break final_msgs = ([] if system_msg is None else [system_msg]) + chosen + tail body["messages"] = final_msgs body.setdefault("metadata", {})["token_pruned"] = True body["metadata"]["input_token_budget"] = budget return body