Gan, Jimmy
241 lines
8.0 KiB
Python
241 lines
8.0 KiB
Python
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import Dict, List, Any, Optional
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def parse_jsonl_file(file_path: str, start_line: int = 0, end_line: int = None) -> List[Dict[str, Any]]:
|
|
"""Parse JSONL file from a specific line number to end_line (or EOF if None)"""
|
|
messages = []
|
|
try:
|
|
with open(file_path, 'r', encoding='utf-8') as f:
|
|
for i, line in enumerate(f):
|
|
if i < start_line:
|
|
continue
|
|
if end_line is not None and i >= end_line:
|
|
break
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
obj = json.loads(line)
|
|
messages.append(obj)
|
|
except json.JSONDecodeError as e:
|
|
logger.warning(f"Failed to parse line {i} in {file_path}: {e}")
|
|
continue
|
|
except Exception as e:
|
|
logger.error(f"Failed to read {file_path}: {e}")
|
|
return messages
|
|
|
|
|
|
def extract_text_content(content: Any) -> str:
|
|
"""Extract text from message content"""
|
|
if isinstance(content, str):
|
|
return content
|
|
|
|
if isinstance(content, list):
|
|
text_parts = []
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if block.get('type') == 'text':
|
|
text_parts.append(block.get('text', ''))
|
|
elif block.get('type') == 'tool_result':
|
|
# Include tool results in searchable text
|
|
result = block.get('content', '')
|
|
if isinstance(result, str):
|
|
text_parts.append(result[:500]) # Limit length
|
|
return '\n'.join(text_parts)
|
|
|
|
return ''
|
|
|
|
|
|
def extract_tool_calls(content: Any, message_uuid: str, timestamp: str) -> List[Dict[str, Any]]:
|
|
"""Extract tool use blocks from message content"""
|
|
tool_calls = []
|
|
|
|
if not isinstance(content, list):
|
|
return tool_calls
|
|
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get('type') == 'tool_use':
|
|
tool_calls.append({
|
|
'message_uuid': message_uuid,
|
|
'tool_use_id': block.get('id'),
|
|
'tool_name': block.get('name'),
|
|
'tool_input': json.dumps(block.get('input', {})),
|
|
'tool_result': None,
|
|
'is_error': False,
|
|
'timestamp': timestamp
|
|
})
|
|
|
|
return tool_calls
|
|
|
|
|
|
def extract_tool_results(content: Any, tool_calls_map: Dict[str, Dict]) -> None:
|
|
"""Extract tool results and match them to tool calls"""
|
|
if not isinstance(content, list):
|
|
return
|
|
|
|
for block in content:
|
|
if isinstance(block, dict) and block.get('type') == 'tool_result':
|
|
tool_use_id = block.get('tool_use_id')
|
|
if tool_use_id and tool_use_id in tool_calls_map:
|
|
result_content = block.get('content', '')
|
|
if isinstance(result_content, str):
|
|
tool_calls_map[tool_use_id]['tool_result'] = result_content[:5000] # Limit length
|
|
tool_calls_map[tool_use_id]['is_error'] = block.get('is_error', False)
|
|
|
|
|
|
def extract_message_data(obj: Dict[str, Any], session_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Extract relevant fields from a message object"""
|
|
msg_type = obj.get('type')
|
|
|
|
# Only process user and assistant messages
|
|
if msg_type not in ['user', 'assistant']:
|
|
return None
|
|
|
|
message = obj.get('message', {})
|
|
role = message.get('role')
|
|
|
|
if not role or role not in ['user', 'assistant']:
|
|
return None
|
|
|
|
content = message.get('content', '')
|
|
has_tool_use = False
|
|
has_thinking = False
|
|
|
|
# Check for tool use and thinking blocks
|
|
if isinstance(content, list):
|
|
for block in content:
|
|
if isinstance(block, dict):
|
|
if block.get('type') == 'tool_use':
|
|
has_tool_use = True
|
|
elif block.get('type') == 'thinking':
|
|
has_thinking = True
|
|
|
|
# Extract usage metrics
|
|
usage = message.get('usage', {})
|
|
input_tokens = usage.get('input_tokens', 0)
|
|
output_tokens = usage.get('output_tokens', 0)
|
|
|
|
return {
|
|
'session_id': session_id,
|
|
'message_uuid': obj.get('uuid'),
|
|
'parent_uuid': obj.get('parentUuid'),
|
|
'role': role,
|
|
'content_text': extract_text_content(content),
|
|
'timestamp': obj.get('timestamp'),
|
|
'prompt_id': obj.get('promptId'),
|
|
'model': message.get('model'),
|
|
'input_tokens': input_tokens,
|
|
'output_tokens': output_tokens,
|
|
'has_tool_use': has_tool_use,
|
|
'has_thinking': has_thinking,
|
|
'stop_reason': message.get('stop_reason')
|
|
}
|
|
|
|
|
|
def calculate_conversation_stats(messages: List[Dict[str, Any]]) -> Dict[str, Any]:
|
|
"""Calculate aggregate statistics for a conversation"""
|
|
stats = {
|
|
'message_count': 0,
|
|
'user_message_count': 0,
|
|
'assistant_message_count': 0,
|
|
'total_input_tokens': 0,
|
|
'total_output_tokens': 0,
|
|
'started_at': None,
|
|
'last_updated_at': None,
|
|
'model': None
|
|
}
|
|
|
|
for msg in messages:
|
|
if msg.get('role') == 'user':
|
|
stats['user_message_count'] += 1
|
|
elif msg.get('role') == 'assistant':
|
|
stats['assistant_message_count'] += 1
|
|
if not stats['model'] and msg.get('model'):
|
|
stats['model'] = msg.get('model')
|
|
|
|
stats['message_count'] += 1
|
|
stats['total_input_tokens'] += msg.get('input_tokens', 0)
|
|
stats['total_output_tokens'] += msg.get('output_tokens', 0)
|
|
|
|
timestamp = msg.get('timestamp')
|
|
if timestamp:
|
|
if not stats['started_at'] or timestamp < stats['started_at']:
|
|
stats['started_at'] = timestamp
|
|
if not stats['last_updated_at'] or timestamp > stats['last_updated_at']:
|
|
stats['last_updated_at'] = timestamp
|
|
|
|
return stats
|
|
|
|
|
|
def parse_conversation_file(file_path: str, start_line: int = 0, end_line: int = None) -> Dict[str, Any]:
|
|
"""Parse a conversation file and extract all data"""
|
|
raw_messages = parse_jsonl_file(file_path, start_line, end_line)
|
|
|
|
if not raw_messages:
|
|
return None
|
|
|
|
# Extract session metadata from first message
|
|
session_id = None
|
|
project_path = None
|
|
git_branch = None
|
|
slug = None
|
|
|
|
for obj in raw_messages:
|
|
if obj.get('sessionId'):
|
|
session_id = obj['sessionId']
|
|
if obj.get('cwd'):
|
|
project_path = obj['cwd']
|
|
if obj.get('gitBranch'):
|
|
git_branch = obj['gitBranch']
|
|
if obj.get('slug'):
|
|
slug = obj['slug']
|
|
|
|
if session_id and project_path:
|
|
break
|
|
|
|
if not session_id:
|
|
logger.warning(f"No session ID found in {file_path}")
|
|
return None
|
|
|
|
# Extract messages
|
|
messages = []
|
|
tool_calls_map = {} # Map tool_use_id to tool call data
|
|
|
|
for obj in raw_messages:
|
|
msg_data = extract_message_data(obj, session_id)
|
|
if msg_data:
|
|
messages.append(msg_data)
|
|
|
|
# Extract tool calls from assistant messages
|
|
if msg_data['role'] == 'assistant' and msg_data['has_tool_use']:
|
|
message_content = obj.get('message', {}).get('content', [])
|
|
tool_calls = extract_tool_calls(message_content, msg_data['message_uuid'], msg_data['timestamp'])
|
|
for tc in tool_calls:
|
|
tool_calls_map[tc['tool_use_id']] = tc
|
|
|
|
# Extract tool results from user messages
|
|
if obj.get('type') == 'user':
|
|
message_content = obj.get('message', {}).get('content', [])
|
|
extract_tool_results(message_content, tool_calls_map)
|
|
|
|
if not messages:
|
|
return None
|
|
|
|
# Calculate stats
|
|
stats = calculate_conversation_stats(messages)
|
|
|
|
return {
|
|
'session_id': session_id,
|
|
'project_path': project_path or 'unknown',
|
|
'git_branch': git_branch,
|
|
'slug': slug,
|
|
'file_path': file_path,
|
|
'messages': messages,
|
|
'tool_calls': list(tool_calls_map.values()),
|
|
**stats
|
|
}
|