TalkingBot/ai_agent/engine_llm.py

"""LiveKit LLM adapter that streams from the internal chat engine."""

import asyncio
import json
import logging
import os
import re
import threading
import time
import uuid

from livekit.agents import llm
from livekit.agents.llm import ChatChunk, ChoiceDelta

from chat_engine_client import ChatEngineClient

logger = logging.getLogger("voice-agent")
FILLER = "Let me check that in the book."
APOLOGY = "I lost my place for a moment. Please ask that again."


def parse_sse_line(line):
    text = line.decode("utf-8").strip() if isinstance(line, bytes) else str(line).strip()
    if not text.startswith("data:"):
        return None
    return json.loads(text[5:].strip())


tts_sanitizer_hits = 0
_COMPLETE_MARKER = re.compile(r"\[\[[^\]]{0,24}\]\]")
_OPEN_MARKER = re.compile(r"\[\[[^\]]{0,24}$")


def sanitize_for_tts(text):
    """Strip leaked source markers, then leftover markdown, before speech."""
    global tts_sanitizer_hits
    cleaned = re.sub(r"```.*?```", " ", text or "", flags=re.S)
    cleaned, complete_hits = _COMPLETE_MARKER.subn(" ", cleaned)
    cleaned, open_hits = _OPEN_MARKER.subn(" ", cleaned)
    tts_sanitizer_hits += complete_hits + open_hits
    cleaned = re.sub(r"[*_#>`\[\]]", "", cleaned)
    cleaned = re.sub(r"(?m)^\s*-\s+", "", cleaned)
    cleaned = re.sub(r"\s+", " ", cleaned).strip()
    return cleaned


def filler_due(started_at, now, text_sent, filler_sent, tool_seen, filler_after_ms):
    if not tool_seen or text_sent or filler_sent:
        return False
    return (now - started_at) * 1000 >= filler_after_ms


def validate_metadata(data):
    if not isinstance(data, dict) or not data.get("sessionId"):
        raise ValueError("room metadata is missing sessionId")
    return data


def parse_metadata(raw):
    if not raw:
        return None
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None
    if not data.get("sessionId"):
        return None
    return data


def truncate_spoken(text, playback_position, chars_per_second=14):
    if playback_position is None:
        return text
    count = max(0, int(playback_position * chars_per_second))
    return (text or "")[:count]


class EngineLLM(llm.LLM):
    def __init__(self, client, session_id, filler_after_ms=800):
        super().__init__()
        self.client = client
        self.session_id = session_id
        self.filler_after_ms = filler_after_ms
        self.last_text = ""

    @property
    def model(self):
        return "chat-engine"

    @property
    def provider(self):
        return "talkingbooklet"

    def chat(self, *, chat_ctx, tools=None, conn_options=None, parallel_tool_calls=None, tool_choice=None, extra_kwargs=None):
        from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS

        if conn_options is None:
            conn_options = DEFAULT_API_CONNECT_OPTIONS
        if tools is None:
            tools = []
        return EngineLLMStream(
            self,
            chat_ctx=chat_ctx,
            tools=tools,
            conn_options=conn_options,
        )


class EngineLLMStream(llm.LLMStream):
    async def _run(self) -> None:
        user_text = latest_user_text(self._chat_ctx)
        if not user_text:
            return
        started = time.perf_counter()
        text_sent = False
        filler_sent = False
        tool_seen = False
        produced = []
        attempt = 0
        while attempt < 2:
            attempt += 1
            try:
                async for event in aiter_events(self._llm.client, self._llm.session_id, user_text):
                    event_type = event.get("type")
                    if event_type == "tool_start":
                        tool_seen = True
                        if filler_due(started, time.perf_counter(), text_sent, filler_sent, True, self._llm.filler_after_ms):
                            filler_sent = True
                            self._send(FILLER)
                        continue
                    if event_type == "text_delta":
                        piece = sanitize_for_tts(event.get("text") or "")
                        if not piece:
                            continue
                        text_sent = True
                        produced.append(piece)
                        self._send(piece)
                        if attempt == 1:
                            self._llm.last_text = "".join(produced)
                    elif event_type == "error":
                        if event.get("retryable") and attempt == 1 and not text_sent:
                            break
                        self._send(APOLOGY)
                        return
                    elif event_type == "final":
                        elapsed = int((time.perf_counter() - started) * 1000)
                        logger.info(
                            "sessionId=%s engine_ttft_ms=%s interrupted=%s",
                            self._llm.session_id,
                            elapsed,
                            bool(event.get("interrupted")),
                        )
                        self._llm.last_text = "".join(produced)
                        return
                else:
                    self._llm.last_text = "".join(produced)
                    return
            except asyncio.CancelledError:
                self._llm.client.close()
                raise
        self._send(APOLOGY)

    def _send(self, text):
        self._event_ch.send_nowait(
            ChatChunk(
                id=str(uuid.uuid4()),
                delta=ChoiceDelta(role="assistant", content=text),
            )
        )


def latest_user_text(chat_ctx):
    items = getattr(chat_ctx, "items", None) or getattr(chat_ctx, "messages", None) or []
    for item in reversed(list(items)):
        role = getattr(item, "role", None)
        if role == "user":
            text = getattr(item, "text_content", None)
            if text:
                return text.strip()
            content = getattr(item, "content", "")
            if isinstance(content, str):
                return content.strip()
    return ""


async def aiter_events(client, session_id, text):
    loop = asyncio.get_running_loop()
    queue = asyncio.Queue()

    def worker():
        try:
            for event in client.iter_events(session_id, text):
                loop.call_soon_threadsafe(queue.put_nowait, ("event", event))
        except Exception as error:
            loop.call_soon_threadsafe(queue.put_nowait, ("error", error))
        finally:
            loop.call_soon_threadsafe(queue.put_nowait, ("end", None))

    threading.Thread(target=worker, daemon=True).start()
    while True:
        kind, payload = await queue.get()
        if kind == "end":
            return
        if kind == "error":
            raise payload
        yield payload


def on_playback_interrupted(client, session_id, full_text, playback_position):
    spoken = truncate_spoken(full_text, playback_position)
    client.close()
    if spoken:
        client.mark_interrupted(session_id, spoken)
    return spoken


def engine_client_from_env():
    return ChatEngineClient(
        base_url=os.environ["CHAT_ENGINE_URL"],
        token=os.environ["CHAT_ENGINE_TOKEN"],
    )