plenum/backend/services/eval_conversation_logger.py

"""Opt-in conversation logging to Postgres for LLM-based evaluation.

Activated ONLY when the first user message of a conversation starts with
"TEST " (case-sensitive, followed by the real question). The prefix is
stripped before the LLM sees it so the logged conversation behaves like a
real one. Normal conversations are never stored.

Never raises — logging must not break the chat (same contract as
event_logger).
"""

import re
import threading
import time
import traceback as _tb
import uuid
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple

_MIGRATION = Path(__file__).resolve().parents[2] / "_postgres/migrations/add_eval_conversations.sql"

_TEST_PREFIX = re.compile(r"^TEST\s+(?=\S)")
_STR_MAX = 20_000
_MAX_EVENTS = 1_000
_MAX_DOC_BYTES = 4_000_000

_table_ready = False
_ready_lock = threading.Lock()


def detect_and_strip_test_prefix(messages: list) -> Tuple[list, bool]:
    """Return (messages, is_eval), inspecting only the FIRST user message.

    On match, returns a new list where the first user message dict is a copy
    with the "TEST " prefix removed. The input is never mutated. Bare "TEST"
    (no remainder) and "TESTAR..." do not activate; "TEST" in later user
    messages is ignored.
    """
    try:
        for idx, msg in enumerate(messages):
            if msg.get("role") != "user":
                continue
            content = msg.get("content") or ""
            if not isinstance(content, str) or not _TEST_PREFIX.match(content):
                return messages, False
            stripped = dict(msg)
            stripped["content"] = _TEST_PREFIX.sub("", content, count=1)
            return list(messages[:idx]) + [stripped] + list(messages[idx + 1 :]), True
        return messages, False
    except Exception:
        return messages, False


def sanitize_provider(provider_override) -> Optional[Dict[str, Any]]:
    """Whitelist provider metadata. api_key is ephemeral and NEVER stored."""
    if provider_override is None:
        return None
    out = {}
    for field in ("provider_id", "smart_model", "fast_model", "editor_model"):
        value = getattr(provider_override, field, None)
        if value:
            out[field] = value
    return out or None


def _truncate(value):
    """Recursively cap strings at _STR_MAX chars (event_logger style)."""
    if isinstance(value, str) and len(value) > _STR_MAX:
        return value[:_STR_MAX] + "…[truncated]"
    if isinstance(value, dict):
        return {k: _truncate(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)):
        return [_truncate(v) for v in value]
    return value


def _insert(doc: dict) -> None:
    global _table_ready
    import json

    from postgres_client import pg

    with _ready_lock:
        if not _table_ready:
            pg.execute_void(_MIGRATION.read_text())
            _table_ready = True
    pg.execute_void(
        """
        INSERT INTO eval_conversations
            (session_id, turn_index, stream, started_at, finished_at,
             duration_s, iterations, has_error, doc)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
        """,
        (
            doc["session_id"],
            doc["turn_index"],
            doc["stream"],
            doc["started_at"],
            doc["finished_at"],
            doc["duration_s"],
            doc["iterations"],
            doc["error"] is not None,
            json.dumps(doc, default=str),
        ),
    )


class ConversationRecorder:
    """Collects one chat turn's events and persists them on finish().

    record() may be called concurrently from the tool-loop thread and the
    shadow-communicator daemon threads; events arriving after finish() are
    dropped.
    """

    def __init__(
        self,
        session_id: Optional[str],
        messages: list,
        request_meta: Dict[str, Any],
        stream: bool = False,
    ):
        self.session_id = session_id or f"anon-{uuid.uuid4().hex[:12]}"
        self.messages = messages
        self.request_meta = request_meta
        self.stream = stream
        self.turn_index = sum(1 for m in messages if m.get("role") == "user")
        self.started_at = datetime.now(timezone.utc)
        self._t0 = time.monotonic()
        self._events: List[Dict[str, Any]] = []
        self._lock = threading.Lock()
        self._iter = 0
        self._finished = False

    def wrap(
        self, downstream: Optional[Callable[[Dict[str, Any]], None]]
    ) -> Callable[[Dict[str, Any]], None]:
        """Record every event, forwarding non-eval-only events downstream."""

        def callback(event: Dict[str, Any]) -> None:
            self.record(event)
            if downstream is not None and not event.get("_eval_only"):
                downstream(event)

        return callback

    def record(self, event: Dict[str, Any]) -> None:
        try:
            with self._lock:
                if self._finished or len(self._events) >= _MAX_EVENTS:
                    return
                compact = {
                    "i": len(self._events),
                    "t_ms": int((time.monotonic() - self._t0) * 1000),
                    **{k: v for k, v in event.items() if k != "_eval_only"},
                }
                if event.get("type") == "tool_call":
                    self._iter += 1
                    compact["iteration"] = self._iter
                self._events.append(_truncate(compact))
        except Exception:
            pass

    def finish(
        self,
        result: Optional[dict] = None,
        error: Optional[BaseException] = None,
    ) -> None:
        try:
            with self._lock:
                if self._finished:
                    return
                self._finished = True
                events = self._events
            error_block = None
            if error is not None:
                error_block = {
                    "type": type(error).__name__,
                    "message": str(error),
                    "traceback": "".join(
                        _tb.format_exception(type(error), error, error.__traceback__)
                    )[:_STR_MAX],
                }
            doc = {
                "kind": "chat_turn",
                "session_id": self.session_id,
                "turn_index": self.turn_index,
                "stream": self.stream,
                "started_at": self.started_at.isoformat(),
                "finished_at": datetime.now(timezone.utc).isoformat(),
                "duration_s": round(time.monotonic() - self._t0, 2),
                "request": {"messages": _truncate(self.messages), **self.request_meta},
                "events": events,
                "iterations": self._iter,
                "result": _truncate(result) if result is not None else None,
                "error": error_block,
            }
            if self._doc_size(doc) > _MAX_DOC_BYTES:
                doc = self._degrade(doc)
            _insert(doc)
        except Exception:
            pass

    @staticmethod
    def _doc_size(doc: dict) -> int:
        import json

        return len(json.dumps(doc, default=str))

    @staticmethod
    def _degrade(doc: dict) -> dict:
        """Shrink an oversized doc: keep search_card hit ids only, trim sources."""
        for event in doc.get("events", []):
            if event.get("type") == "search_card" and isinstance(
                event.get("results"), list
            ):
                event["results"] = [
                    r.get("_id") for r in event["results"] if isinstance(r, dict)
                ]
                event["_degraded"] = True
        result = doc.get("result")
        if isinstance(result, dict) and isinstance(result.get("sources"), list):
            for src in result["sources"]:
                if isinstance(src, dict) and isinstance(src.get("snippet"), str):
                    src["snippet"] = src["snippet"][:500]
        return doc