chat_parser.py

"""Chat log parsing utilities for SillyLoresmith."""

import json
import re
from helpers import STATUS_LOG


def parse_sillytavern_jsonl(content):
    """
    Parse SillyTavern JSONL chat format.
    Line 1 = metadata (chat info)
    Lines 2+ = individual messages
    
    Returns (messages, metadata)
    """
    lines = content.strip().split('\n')
    if not lines:
        return [], {}
    
    metadata = {}
    messages = []
    
    for i, line in enumerate(lines):
        if not line.strip():
            continue
        
        try:
            data = json.loads(line)
            
            # First line is usually metadata
            if i == 0 and "user_name" in data:
                metadata = data
                continue
            
            # Parse message
            msg = {
                "name": data.get("name", data.get("sender", "Unknown")),
                "text": data.get("mes", data.get("message", data.get("content", ""))),
                "is_user": data.get("is_user", data.get("role") == "user"),
                "timestamp": data.get("send_date", data.get("timestamp", "")),
            }
            
            if msg["text"]:
                messages.append(msg)
                
        except json.JSONDecodeError:
            # Not valid JSON, skip
            continue
    
    STATUS_LOG.info(f"Parsed JSONL: {len(messages)} messages")
    return messages, metadata


def parse_sillytavern_json(content):
    """
    Parse SillyTavern JSON chat format (array or object with messages).
    
    Returns (messages, metadata)
    """
    try:
        data = json.loads(content)
    except json.JSONDecodeError as e:
        STATUS_LOG.error(f"Invalid JSON: {e}")
        return [], {}
    
    metadata = {}
    messages = []
    
    # Handle different formats
    if isinstance(data, list):
        # Array of messages
        msg_list = data
    elif isinstance(data, dict):
        # Object with messages array
        metadata = {k: v for k, v in data.items() if k not in ("messages", "chat", "data")}
        msg_list = data.get("messages", data.get("chat", data.get("data", [])))
        if isinstance(msg_list, dict):
            msg_list = list(msg_list.values())
    else:
        return [], {}
    
    for item in msg_list:
        if isinstance(item, dict):
            msg = {
                "name": item.get("name", item.get("sender", item.get("role", "Unknown"))),
                "text": item.get("mes", item.get("message", item.get("content", ""))),
                "is_user": item.get("is_user", item.get("role") == "user"),
                "timestamp": item.get("send_date", item.get("timestamp", "")),
            }
            if msg["text"]:
                messages.append(msg)
        elif isinstance(item, str):
            # Plain string message
            messages.append({"name": "Unknown", "text": item, "is_user": False})
    
    STATUS_LOG.info(f"Parsed JSON: {len(messages)} messages")
    return messages, metadata


def parse_plain_text_chat(content):
    """
    Parse plain text chat format (Name: Message).
    
    Returns (messages, metadata)
    """
    messages = []
    
    # Common patterns for chat messages
    patterns = [
        # "Name: Message" format
        r'^([A-Za-z0-9_\-\s]{1,30}):\s*(.+?)(?=\n[A-Za-z0-9_\-\s]{1,30}:|\Z)',
        # "[Name] Message" format
        r'^\[([A-Za-z0-9_\-\s]{1,30})\]\s*(.+?)(?=\n\[[A-Za-z0-9_\-\s]{1,30}\]|\Z)',
        # "<Name> Message" format
        r'^<([A-Za-z0-9_\-\s]{1,30})>\s*(.+?)(?=\n<[A-Za-z0-9_\-\s]{1,30}>|\Z)',
    ]
    
    for pattern in patterns:
        matches = re.findall(pattern, content, re.MULTILINE | re.DOTALL)
        if matches and len(matches) > 2:  # Found at least a few messages
            for name, text in matches:
                text = text.strip()
                if text:
                    messages.append({
                        "name": name.strip(),
                        "text": text,
                        "is_user": False,  # Can't determine from plain text
                        "timestamp": "",
                    })
            STATUS_LOG.info(f"Parsed plain text: {len(messages)} messages")
            return messages, {}
    
    # Fallback: Split by double newlines and treat as alternating messages
    paragraphs = [p.strip() for p in content.split('\n\n') if p.strip()]
    for i, para in enumerate(paragraphs):
        messages.append({
            "name": f"Speaker {i % 2 + 1}",
            "text": para,
            "is_user": i % 2 == 0,
            "timestamp": "",
        })
    
    STATUS_LOG.info(f"Parsed as paragraphs: {len(messages)} messages")
    return messages, {}


def parse_chat_log(content, format_hint=None):
    """
    Auto-detect and parse chat log format.
    
    format_hint: "jsonl", "json", "text", or None for auto-detect
    
    Returns (messages, metadata)
    """
    content = content.strip()
    
    if not content:
        return [], {}
    
    # Try JSONL first (most specific)
    if format_hint == "jsonl" or (format_hint is None and content.startswith('{')):
        try:
            # Check if it's JSONL (multiple JSON objects per line)
            lines = content.split('\n')
            if len(lines) > 1:
                json.loads(lines[0])
                json.loads(lines[1]) if len(lines) > 1 else None
                messages, metadata = parse_sillytavern_jsonl(content)
                if messages:
                    return messages, metadata
        except Exception:
            pass
    
    # Try JSON
    if format_hint == "json" or (format_hint is None and content.startswith(('[', '{'))):
        try:
            messages, metadata = parse_sillytavern_json(content)
            if messages:
                return messages, metadata
        except Exception:
            pass
    
    # Fall back to plain text
    return parse_plain_text_chat(content)


def extract_character_names(messages):
    """Extract unique character names from messages."""
    names = set()
    for msg in messages:
        name = msg.get("name", "").strip()
        if name and name.lower() not in ("unknown", "system", "narrator"):
            names.add(name)
    return sorted(names)


def filter_messages_by_character(messages, character_name):
    """Filter messages to only those from a specific character."""
    return [m for m in messages 
            if m.get("name", "").lower() == character_name.lower()]


def get_message_context(messages, index, window=3):
    """Get context around a specific message."""
    start = max(0, index - window)
    end = min(len(messages), index + window + 1)
    return messages[start:end]