mindroot 9.24.0__py3-none-any.whl → 10.2.0__py3-none-any.whl

mindroot/lib/chatlog.py

@@ -281,7 +281,7 @@ async def find_child_logs_by_parent_id(parent_log_id: str) -> List[str]:
 
 def extract_delegate_task_log_ids(messages: List[Dict]) -> List[str]:
     """
-    Extract log IDs from delegate_task commands in messages.
+    Extract log IDs from delegate_task and delegate_subtask commands in messages.
     
     Args:
         messages: List of chat messages
@@ -312,10 +312,15 @@ def extract_delegate_task_log_ids(messages: List[Dict]) -> List[str]:
                     for key, value in cmd.items():
                         if key == 'delegate_task' and 'log_id' in value:
                             log_ids.append(value['log_id'])
+                        elif key == 'delegate_subtask' 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)
+                # Also check for delegate_subtask commands
+                matches = re.findall(r'"delegate_subtask"\s*:\s*{\s*"log_id"\s*:\s*"([^"]+)"', text)
+                log_ids.extend(matches)
     
     return log_ids
 
@@ -389,7 +394,119 @@ async def save_token_counts_to_cache(log_id: str, token_counts: Dict[str, int])
     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]:
+async def build_token_hierarchy(log_id: str, user: str = None, visited: set = None) -> Dict:
+    """
+    Build a hierarchical token count structure for a log and its children.
+    
+    Args:
+        log_id: The log ID to build hierarchy for
+        user: User for the log
+        visited: Set of already visited log IDs to prevent infinite recursion
+        
+    Returns:
+        Dictionary with hierarchical structure containing:
+        - log_id: The log ID
+        - individual_counts: Token counts for this log only
+        - cumulative_counts: Token counts including all children
+        - children: List of child hierarchies
+    """
+    if visited is None:
+        visited = set()
+    
+    if log_id in visited:
+        return None  # Prevent infinite recursion
+    
+    visited.add(log_id)
+    
+    # Find the chatlog file
+    chatlog_path = await find_chatlog_file(log_id)
+    if not chatlog_path:
+        return None
+    
+    # Load the chat log
+    async with aiofiles.open(chatlog_path, 'r') as f:
+        content = await f.read()
+        log_data = json.loads(content)
+    
+    # Check if we have cached individual counts for this specific session
+    cached_individual = await get_cached_token_counts(log_id, chatlog_path)
+    if cached_individual and 'input_tokens_sequence' in cached_individual:
+        print(f"Using cached individual token counts for session {log_id}")
+        individual_counts = {
+            'input_tokens_sequence': cached_individual['input_tokens_sequence'],
+            'output_tokens_sequence': cached_individual['output_tokens_sequence'],
+            'input_tokens_total': cached_individual['input_tokens_total']
+        }
+    else:
+        # Calculate individual counts for this session
+        if user is None:
+            try:
+                path_parts = chatlog_path.split(os.sep)
+                if len(path_parts) >= 4 and path_parts[-4] == 'chat':
+                    user = path_parts[-3]
+                else:
+                    user = "system"
+            except Exception:
+                user = "system"
+        
+        temp_log = ChatLog(log_id=log_id, user=user, agent=log_data.get('agent', 'unknown'))
+        temp_log.messages = log_data.get('messages', [])
+        
+        # Count tokens for this log only
+        individual_counts = temp_log.count_tokens()
+        
+        # Cache the individual session counts
+        individual_cache_data = {
+            'input_tokens_sequence': individual_counts['input_tokens_sequence'],
+            'output_tokens_sequence': individual_counts['output_tokens_sequence'],
+            'input_tokens_total': individual_counts['input_tokens_total']
+        }
+        await save_token_counts_to_cache(log_id, individual_cache_data)
+        print(f"Cached individual token counts for session {log_id}")
+    
+    # Find all child log IDs
+    if user is None:
+        try:
+            path_parts = chatlog_path.split(os.sep)
+            if len(path_parts) >= 4 and path_parts[-4] == 'chat':
+                user = path_parts[-3]
+            else:
+                user = "system"
+        except Exception:
+            user = "system"
+    
+    temp_log = ChatLog(log_id=log_id, user=user, agent=log_data.get('agent', 'unknown'))
+    temp_log.messages = log_data.get('messages', [])
+    delegated_log_ids = extract_delegate_task_log_ids(temp_log.messages)
+    child_logs_by_parent = await find_child_logs_by_parent_id(log_id)
+    all_child_log_ids = list(set(delegated_log_ids) | set(child_logs_by_parent))
+    
+    # Build child hierarchies
+    children = []
+    cumulative_counts = {
+        'input_tokens_sequence': individual_counts['input_tokens_sequence'],
+        'output_tokens_sequence': individual_counts['output_tokens_sequence'],
+        'input_tokens_total': individual_counts['input_tokens_total']
+    }
+    
+    for child_id in all_child_log_ids:
+        child_hierarchy = await build_token_hierarchy(child_id, user, visited.copy())
+        if child_hierarchy:
+            children.append(child_hierarchy)
+            # Add child's cumulative counts to our cumulative counts
+            cumulative_counts['input_tokens_sequence'] += child_hierarchy['cumulative_counts']['input_tokens_sequence']
+            cumulative_counts['output_tokens_sequence'] += child_hierarchy['cumulative_counts']['output_tokens_sequence']
+            cumulative_counts['input_tokens_total'] += child_hierarchy['cumulative_counts']['input_tokens_total']
+    
+    return {
+        'log_id': log_id,
+        'agent': log_data.get('agent', 'unknown'),
+        'individual_counts': individual_counts,
+        'cumulative_counts': cumulative_counts,
+        'children': children
+    }
+
+async def count_tokens_for_log_id(log_id: str, user: str = None, hierarchical: bool = False) -> Dict[str, int]:
     """
     Count tokens for a chat log identified by log_id, including any delegated tasks.
     
@@ -397,13 +514,32 @@ async def count_tokens_for_log_id(log_id: str) -> Dict[str, int]:
         log_id: The log ID to count tokens for
         
     Returns:
-        Dictionary with token counts or None if log not found
+        Dictionary with token counts or None if log not found.
+        If hierarchical=True, includes 'hierarchy' key with tree structure.
+        If hierarchical=False (default), returns flat structure for backwards compatibility.
     """
     # Find the chatlog file
     chatlog_path = await find_chatlog_file(log_id)
     if not chatlog_path:
         return None
     
+    # If hierarchical structure is requested, build and return it
+    if hierarchical:
+        # Check cache first for hierarchical data
+        cached_counts = await get_cached_token_counts(log_id, chatlog_path)
+        if cached_counts and 'hierarchy' in cached_counts:
+            print(f"Using cached hierarchical token counts for {log_id}")
+            return cached_counts
+        
+        print(f"Calculating hierarchical token counts for {log_id}")
+        hierarchy = await build_token_hierarchy(log_id, user)
+        if hierarchy:
+            result = {'hierarchy': hierarchy}
+            # Save hierarchical data to cache
+            await save_token_counts_to_cache(log_id, result)
+            return result
+        return None
+    
     # Check cache first
     cached_counts = await get_cached_token_counts(log_id, chatlog_path)
     if cached_counts:
@@ -421,7 +557,21 @@ async def count_tokens_for_log_id(log_id: str) -> Dict[str, int]:
     parent_log_id = log_data.get('parent_log_id')
         
     # Create a temporary ChatLog instance to count tokens
-    temp_log = ChatLog(log_id=log_id, user="system", agent=log_data.get('agent', 'unknown'))
+    # Use provided user or try to determine from chatlog path or fallback to "system"
+    if user is None:
+        # Try to extract user from chatlog path: data/chat/{user}/{agent}/chatlog_{log_id}.json
+        try:
+            path_parts = chatlog_path.split(os.sep)
+            if len(path_parts) >= 4 and path_parts[-4] == 'chat':
+                extracted_user = path_parts[-3]  # User is third from the end
+                user = extracted_user
+                print(f"Extracted user '{user}' from chatlog path: {chatlog_path}")
+            else:
+                user = "system"  # Default fallback
+        except Exception as e:
+            print(f"Error extracting user from path {chatlog_path}: {e}")
+            user = "system"  # Default fallback
+    temp_log = ChatLog(log_id=log_id, user=user, agent=log_data.get('agent', 'unknown'))
     temp_log.messages = log_data.get('messages', [])
     
     # Count tokens for this log
@@ -449,7 +599,7 @@ async def count_tokens_for_log_id(log_id: str) -> Dict[str, int]:
     
     # Recursively count tokens for all child tasks
     for child_id in all_child_log_ids:
-        delegated_counts = await count_tokens_for_log_id(child_id)
+        delegated_counts = await count_tokens_for_log_id(child_id, user=user)
         if delegated_counts:
             combined_counts['input_tokens_sequence'] += delegated_counts['input_tokens_sequence']
             combined_counts['output_tokens_sequence'] += delegated_counts['output_tokens_sequence']

{mindroot-9.24.0.dist-info → mindroot-10.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mindroot
-Version: 9.24.0
+Version: 10.2.0
 Summary: MindRoot AI Agent Framework
 Requires-Python: >=3.9
 License-File: LICENSE