Function create_input_token_counter

augmentoolkit/utils/observers.py:77–149

Creates an input observer that counts tokens and calculates cost. Args: counter: A dictionary to track token counts and costs cost_per_million: Cost per million tokens count_tokens_fn: Function to count tokens in text persistence_path: Optional path to persi

(
    counter: Dict[str, float],
    cost_per_million: float,
    count_tokens_fn: Callable,
    persistence_path: str = None,
)

def create_input_token_counter(
    counter: Dict[str, float],
    cost_per_million: float,
    count_tokens_fn: Callable,
    persistence_path: str = None,
) -> Callable:
    """
    Creates an input observer that counts tokens and calculates cost.

    Args:
        counter: A dictionary to track token counts and costs
        cost_per_million: Cost per million tokens
        count_tokens_fn: Function to count tokens in text
        persistence_path: Optional path to persist counter data to a JSON file

    Returns:
        A callback function that updates the counter
    """
    # Initialize lock for file access
    json_path = None

    if persistence_path:
        os.makedirs(os.path.dirname(persistence_path), exist_ok=True)
        json_path = persistence_path

        # Load existing counter data if available
        if os.path.exists(json_path):
            try:
                with open(json_path, "r") as f:
                    saved_data = json.load(f)
                    # Update the counter with saved values
                    counter.update(saved_data)
            except (json.JSONDecodeError, FileNotFoundError):
                # If file is corrupted or doesn't exist, start fresh
                pass

    def input_token_counter(
        input_data: Union[str, List[Dict[str, str]]],
        completion_mode: bool,
        *args,
        **kwargs,
    ) -> None:
        """
        Counts tokens in the input and updates the counter.

        Args:
            input_data: Either a prompt string (completion mode) or messages list (chat mode)
            completion_mode: Whether this was a completion (True) or chat (False) request
        """
        if completion_mode:
            # For completion mode, input is a string
            token_count = count_tokens_fn(input_data)
        else:
            # For chat mode, input is a list of message dicts
            token_count = 0
            for message in input_data:
                token_count += count_tokens_fn(message["content"])