Gan, Jimmy
131 lines
4.6 KiB
Python
131 lines
4.6 KiB
Python
import os
|
|
import logging
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
import asyncio
|
|
|
|
import db
|
|
import parser
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
CONVERSATIONS_PATH = os.environ.get("CONVERSATIONS_PATH", "/app/conversations")
|
|
COLLECT_INTERVAL = int(os.environ.get("COLLECT_INTERVAL", "60"))
|
|
|
|
|
|
async def discover_jsonl_files():
|
|
"""Discover all .jsonl files in conversations directory"""
|
|
conversations_dir = Path(CONVERSATIONS_PATH)
|
|
if not conversations_dir.exists():
|
|
logger.warning(f"Conversations directory does not exist: {CONVERSATIONS_PATH}")
|
|
return []
|
|
|
|
jsonl_files = []
|
|
for file_path in conversations_dir.rglob("*.jsonl"):
|
|
if file_path.is_file():
|
|
jsonl_files.append(str(file_path))
|
|
|
|
return jsonl_files
|
|
|
|
|
|
async def process_file(file_path: str):
|
|
"""Process a single conversation file incrementally"""
|
|
try:
|
|
# Get checkpoint
|
|
checkpoint = await db.get_checkpoint(file_path)
|
|
start_line = checkpoint['last_line_count'] if checkpoint else 0
|
|
|
|
# Get file modification time
|
|
file_stat = os.stat(file_path)
|
|
file_mtime = datetime.fromtimestamp(file_stat.st_mtime).isoformat()
|
|
|
|
# Check if file has been modified
|
|
if checkpoint and checkpoint['last_modified'] == file_mtime:
|
|
# File hasn't changed, skip
|
|
return 0
|
|
|
|
# Count total lines
|
|
with open(file_path, 'r') as f:
|
|
total_lines = sum(1 for _ in f)
|
|
|
|
if total_lines <= start_line:
|
|
# No new lines
|
|
return 0
|
|
|
|
# Process in batches for large files (max 200 lines at a time)
|
|
batch_size = 200
|
|
end_line = min(start_line + batch_size, total_lines)
|
|
|
|
# Parse new content
|
|
logger.info(f"Processing {file_path} from line {start_line} to {end_line} (total: {total_lines})")
|
|
conversation_data = parser.parse_conversation_file(file_path, start_line, end_line)
|
|
|
|
if not conversation_data:
|
|
logger.warning(f"No data extracted from {file_path}")
|
|
return 0
|
|
|
|
# Store in database
|
|
await db.upsert_conversation(
|
|
conversation_data['session_id'],
|
|
{
|
|
'project_path': conversation_data['project_path'],
|
|
'git_branch': conversation_data.get('git_branch'),
|
|
'started_at': conversation_data['started_at'],
|
|
'last_updated_at': conversation_data['last_updated_at'],
|
|
'message_count': conversation_data['message_count'],
|
|
'user_message_count': conversation_data['user_message_count'],
|
|
'assistant_message_count': conversation_data['assistant_message_count'],
|
|
'total_input_tokens': conversation_data['total_input_tokens'],
|
|
'total_output_tokens': conversation_data['total_output_tokens'],
|
|
'model': conversation_data.get('model'),
|
|
'slug': conversation_data.get('slug'),
|
|
'file_path': file_path
|
|
}
|
|
)
|
|
|
|
# Store messages
|
|
for msg in conversation_data['messages']:
|
|
await db.insert_message(msg)
|
|
|
|
# Store tool calls
|
|
for tool_call in conversation_data['tool_calls']:
|
|
await db.insert_tool_call(tool_call)
|
|
|
|
# Update checkpoint (use end_line instead of total_lines for batch processing)
|
|
await db.update_checkpoint(file_path, end_line, file_mtime)
|
|
|
|
logger.info(f"Processed {len(conversation_data['messages'])} messages from {file_path} (batch {start_line}-{end_line}/{total_lines})")
|
|
return len(conversation_data['messages'])
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error processing {file_path}: {e}", exc_info=True)
|
|
return 0
|
|
|
|
|
|
async def collect_conversations():
|
|
"""Main collection loop"""
|
|
while True:
|
|
try:
|
|
logger.info("Starting conversation collection cycle")
|
|
|
|
# Discover files
|
|
files = await discover_jsonl_files()
|
|
logger.info(f"Found {len(files)} conversation files")
|
|
|
|
# Process each file
|
|
total_messages = 0
|
|
for file_path in files:
|
|
messages_processed = await process_file(file_path)
|
|
total_messages += messages_processed
|
|
|
|
if total_messages > 0:
|
|
logger.info(f"Collection cycle complete: processed {total_messages} new messages")
|
|
else:
|
|
logger.debug("Collection cycle complete: no new messages")
|
|
|
|
except Exception as e:
|
|
logger.error(f"Error in collection cycle: {e}", exc_info=True)
|
|
|
|
# Wait before next cycle
|
|
await asyncio.sleep(COLLECT_INTERVAL)
|