Skip to content
Back to top

# Source code for mindroot.lib.token_counter

    import os
    import json
    import re
    import time
    import asyncio
    import aiofiles
    import aiofiles.os
    from typing import Dict, List
    from mindroot.lib.chatlog import ChatLog

    async def find_chatlog_file(log_id: str) -> str:
        """
        Find a chatlog file by its log_id.

        Args:
            log_id: The log ID to search for

        Returns:
            The full path to the chatlog file if found, None otherwise
        """
        chat_dir = os.environ.get('CHATLOG_DIR', 'data/chat')

        # Use os.walk to search through all subdirectories
        for root, dirs, files in await asyncio.to_thread(os.walk, chat_dir):
            for file in files:
                if file == f"chatlog_{log_id}.json":
                    return os.path.join(root, file)

        return None

    def extract_delegate_task_log_ids(messages: List[Dict]) -> List[str]:
        """
        Extract log IDs from delegate_task commands in messages.

        Args:
            messages: List of chat messages

        Returns:
            List of log IDs found in delegate_task commands
        """
        log_ids = []

        for message in messages:
            if message['role'] == 'assistant':
                content = message['content']
                # Handle both string and list content formats
                if isinstance(content, str):
                    text = content
                elif isinstance(content, list) and len(content) > 0 and 'text' in content[0]:
                    text = content[0]['text']
                else:
                    continue

                # Try to parse as JSON
                try:
                    commands = json.loads(text)
                    if not isinstance(commands, list):
                        commands = [commands]

                    for cmd in commands:
                        for key, value in cmd.items():
                            if key == 'delegate_task' and 'log_id' in value:
                                log_ids.append(value['log_id'])
                except (json.JSONDecodeError, TypeError, KeyError):
                    # If not JSON, try regex to find log_ids in delegate_task commands
                    matches = re.findall(r'"delegate_task"\s*:\s*{\s*"log_id"\s*:\s*"([^"]+)"', text)
                    log_ids.extend(matches)

        return log_ids

    async def get_cache_dir() -> str:
        """
        Get the directory for token count cache files.
        Creates the directory if it doesn't exist.
        """
        cache_dir = os.environ.get('TOKEN_CACHE_DIR', 'data/token_cache')
        if not await aiofiles.os.path.exists(cache_dir):
            await aiofiles.os.makedirs(cache_dir)
        return cache_dir

    async def get_cache_path(log_id: str) -> str:
        """
        Get the path to the cache file for a specific log_id.
        """
        cache_dir = await get_cache_dir()
        return os.path.join(cache_dir, f"tokens_{log_id}.json")

    async def get_cached_token_counts(log_id: str, log_path: str) -> Dict[str, int]:
        """
        Get cached token counts if available and valid.

        Args:
            log_id: The log ID
            log_path: Path to the actual log file

        Returns:
            Cached token counts if valid, None otherwise
        """
        cache_path = await get_cache_path(log_id)

        # If cache doesn't exist, return None
        if not await aiofiles.os.path.exists(cache_path):
            return None

        try:
            # Get modification times
            log_mtime = await aiofiles.os.path.getmtime(log_path)
            cache_mtime = await aiofiles.os.path.getmtime(cache_path)
            current_time = time.time()

            # If log was modified after cache was created, cache is invalid
            if log_mtime > cache_mtime:
                return None

            # Don't recalculate sooner than 3 minutes after last calculation
            if current_time - cache_mtime < 180:  # 3 minutes in seconds
                async with aiofiles.open(cache_path, 'r') as f:
                    content = await f.read()
                    return json.loads(content)

            # For logs that haven't been modified in over an hour, consider them "finished"
            # and use the cache regardless of when it was last calculated
            if current_time - log_mtime > 3600:  # 1 hour in seconds
                async with aiofiles.open(cache_path, 'r') as f:
                    content = await f.read()
                    return json.loads(content)

        except (json.JSONDecodeError, IOError) as e:
            print(f"Error reading token cache: {e}")

        return None

    async def save_token_counts_to_cache(log_id: str, token_counts: Dict[str, int]) -> None:
        """
        Save token counts to cache.
        """
        cache_path = await get_cache_path(log_id)
        async with aiofiles.open(cache_path, 'w') as f:
            await f.write(json.dumps(token_counts))

    async def count_tokens_for_log_id(log_id: str) -> Dict[str, int]:
        """
        Count tokens for a chat log identified by log_id, including any delegated tasks.

        Args:
            log_id: The log ID to count tokens for

        Returns:
            Dictionary with token counts or None if log not found
        """
        # Find the chatlog file
        chatlog_path = await find_chatlog_file(log_id)
        if not chatlog_path:
            return None

        # Check cache first
        cached_counts = await get_cached_token_counts(log_id, chatlog_path)
        if cached_counts:
            print(f"Using cached token counts for {log_id}")
            return cached_counts

        print(f"Calculating token counts for {log_id}")

        # Load the chat log
        async with aiofiles.open(chatlog_path, 'r') as f:
            content = await f.read()
            log_data = json.loads(content)

        # Create a temporary ChatLog instance to count tokens
        temp_log = ChatLog(log_id=log_id, user="system", agent=log_data.get('agent', 'unknown'))
        temp_log.messages = log_data.get('messages', [])

        # Count tokens for this log
        parent_counts = temp_log.count_tokens()

        # Create combined counts (starting with parent counts)
        combined_counts = {}
        combined_counts['input_tokens_sequence'] = parent_counts['input_tokens_sequence']
        combined_counts['output_tokens_sequence'] = parent_counts['output_tokens_sequence']
        combined_counts['input_tokens_total'] = parent_counts['input_tokens_total']

        # Find delegated task log IDs
        delegated_log_ids = extract_delegate_task_log_ids(temp_log.messages)

        # Recursively count tokens for delegated tasks
        for delegated_id in delegated_log_ids:
            delegated_counts = await count_tokens_for_log_id(delegated_id)
            if delegated_counts:
                combined_counts['input_tokens_sequence'] += delegated_counts['input_tokens_sequence']
                combined_counts['output_tokens_sequence'] += delegated_counts['output_tokens_sequence']
                combined_counts['input_tokens_total'] += delegated_counts['input_tokens_total']

        # Create final result with both parent and combined counts
        token_counts = {}
        # Parent session only counts
        token_counts['input_tokens_sequence'] = parent_counts['input_tokens_sequence']
        token_counts['output_tokens_sequence'] = parent_counts['output_tokens_sequence']
        token_counts['input_tokens_total'] = parent_counts['input_tokens_total']
        # Combined counts (parent + all subtasks)
        token_counts['combined_input_tokens_sequence'] = combined_counts['input_tokens_sequence']
        token_counts['combined_output_tokens_sequence'] = combined_counts['output_tokens_sequence']
        token_counts['combined_input_tokens_total'] = combined_counts['input_tokens_total']

        # Save to cache
        await save_token_counts_to_cache(log_id, token_counts)

        return token_counts