parser.py

"""
Log parsers for different request formats.

Supported formats:
  openai  — {"messages": [{"role": "...", "content": "..."}]}  (OpenAI chat API)
  vllm    — {"prompt": "...", "request_id": "..."}            (vLLM request log)
  plain   — one prompt per line (plaintext)
  jsonl   — generic JSONL with a configurable prompt field

All parsers return Request objects or None if the line is invalid/empty.
"""

import json
from pathlib import Path
from typing import Callable

from .models import LogFormat, Request
from .security import (
    validate_input_path,
    check_line_size,
    check_prompt_length,
    safe_json_loads,
    SecurityError,
)


# ── Individual line parsers ──────────────────────────────────────────────────

def _parse_openai_jsonl(line: str) -> Request | None:
    """
    Parse OpenAI chat completion request format.

    Concatenates all message roles+content into a single string so the trie
    can find shared prefixes across the system+few-shot portion of requests.

    Handles both text-only and multimodal (list-of-parts) content.

    Session ID is extracted from the first matching field:
      session_id > conversation_id > thread_id
    """
    try:
        data = safe_json_loads(line)
        messages: list[dict] = data.get("messages", [])
        if not messages:
            return None

        parts: list[str] = []
        for msg in messages:
            role = msg.get("role", "")
            content = msg.get("content", "")
            if isinstance(content, str):
                parts.append(f"<|{role}|>\n{content}")
            elif isinstance(content, list):
                # Multimodal: extract text parts only
                text = " ".join(
                    c["text"] for c in content if c.get("type") == "text"
                )
                parts.append(f"<|{role}|>\n{text}")

        request_id = data.get("request_id") or data.get("id")
        session_id = (
            data.get("session_id")
            or data.get("conversation_id")
            or data.get("thread_id")
        )
        prompt = check_prompt_length("\n".join(parts), request_id)
        return Request(
            prompt=prompt,
            request_id=request_id,
            timestamp=data.get("timestamp"),
            session_id=str(session_id) if session_id is not None else None,
        )
    except (json.JSONDecodeError, KeyError, TypeError, SecurityError):
        return None


def _parse_vllm(line: str) -> Request | None:
    """
    Parse vLLM request log format.

    vLLM logs vary by version; we probe common field names.
    Session ID is extracted from `session_id` if present.
    """
    try:
        data = safe_json_loads(line)
        prompt = (
            data.get("prompt")
            or data.get("inputs")
            or data.get("text")
            or data.get("input")
        )
        if not prompt:
            return None
        request_id = data.get("request_id")
        session_id = data.get("session_id")
        return Request(
            prompt=check_prompt_length(str(prompt), request_id),
            request_id=request_id,
            timestamp=data.get("timestamp") or data.get("arrival_time"),
            session_id=str(session_id) if session_id is not None else None,
        )
    except (json.JSONDecodeError, KeyError, TypeError, SecurityError):
        return None


def _parse_plain(line: str) -> Request | None:
    """One prompt per line — strip and skip blanks."""
    stripped = line.strip()
    return Request(prompt=stripped) if stripped else None


def _make_jsonl_parser(
    prompt_field: str,
    session_field: str | None = None,
) -> Callable[[str], Request | None]:
    """Return a parser for generic JSONL with user-specified prompt and session fields."""
    def _parse(line: str) -> Request | None:
        try:
            data = safe_json_loads(line)
            prompt = data.get(prompt_field)
            if not prompt:
                return None
            request_id = data.get("id") or data.get("request_id")
            session_id = data.get(session_field) if session_field else None
            return Request(
                prompt=check_prompt_length(str(prompt), request_id),
                request_id=request_id,
                timestamp=data.get("timestamp"),
                session_id=str(session_id) if session_id is not None else None,
            )
        except (json.JSONDecodeError, KeyError, TypeError, SecurityError):
            return None
    return _parse


# ── Public loader ────────────────────────────────────────────────────────────

def load_requests(
    path: Path | str,
    format: LogFormat = LogFormat.OPENAI_JSONL,
    prompt_field: str = "prompt",
    limit: int | None = None,
    session_field: str | None = None,
) -> list[Request]:
    """
    Load and parse requests from a log file.

    Args:
        path:          Path to the log file.
        format:        One of the LogFormat enum values.
        prompt_field:  For LogFormat.JSONL — which JSON key holds the prompt.
        limit:         If set, stop after this many successfully parsed lines.
        session_field: For LogFormat.JSONL — which JSON key holds the session ID.
                       For openai/vllm formats, session_id is auto-detected from
                       common field names (session_id, conversation_id, thread_id).

    Returns:
        List of Request objects. Invalid lines are silently skipped.
    """
    validate_input_path(Path(path))

    parser: Callable[[str], Request | None] = {
        LogFormat.OPENAI_JSONL: _parse_openai_jsonl,
        LogFormat.VLLM:         _parse_vllm,
        LogFormat.PLAIN:        _parse_plain,
        LogFormat.JSONL:        _make_jsonl_parser(prompt_field, session_field),
    }[format]

    requests: list[Request] = []

    with open(path, "r", encoding="utf-8") as f:
        for line_number, line in enumerate(f, start=1):
            if limit and len(requests) >= limit:
                break
            raw = line.rstrip("\n")
            check_line_size(raw, line_number)
            req = parser(raw)
            if req is not None:
                requests.append(req)

    return requests