""" title: Token Clip Filter author: houxin author_url: https://github.com/hx173149 funding_url: https://github.com/hx173149 version: 0.1 """ from pydantic import BaseModel, Field from typing import Optional import tiktoken class Filter: class Valves(BaseModel): priority: int = Field( default=0, description="Priority level for the filter operations." ) n_token_limit: int = Field( default=7000, description="Number of token limit to retain." ) pass class UserValves(BaseModel): pass def __init__(self): self.valves = self.Valves() self.encoding = tiktoken.get_encoding("cl100k_base") pass def inlet(self, body: dict, __user__: Optional[dict] = None) -> dict: messages = body["messages"] sys_msgs = [message for message in messages if message.get("role") == "system"] if len(sys_msgs) > 0: sys_msg = sys_msgs[0] else: sys_msg = {"role": "system", "content": ""} token_len = len(self.encoding.encode(sys_msg["content"])) filter_messages = [] remain_messages = [ message for message in messages if message.get("role") != "system" ] for msg in remain_messages[::-1]: if ( len(self.encoding.encode(msg["content"])) + token_len > self.valves.n_token_limit ): break filter_messages.append(msg) token_len += len(self.encoding.encode(msg["content"])) body["messages"] = [sys_msg] + filter_messages[::-1] return body