simbarag/blueprints/conversation/runtime.py

"""Shared agent invocation for authenticated chat and messaging channels."""

import asyncio
import json
import logging
import time
import uuid
import weakref
from contextlib import asynccontextmanager
from dataclasses import dataclass

from langchain_core.messages import convert_to_messages
from quart import current_app

from metrics import AGENT_REQUESTS_TOTAL, AGENT_RESPONSE_SECONDS
from .checkpoints import prune_thread_checkpoints
from .logic import get_agent_for
from .memory import build_system_prompt_with_memories


# Concurrent turns on one checkpoint thread race on read-modify-write of its
# history, silently dropping a turn. Hold this around preparation + invocation.
_thread_locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
    weakref.WeakValueDictionary()
)


@asynccontextmanager
async def thread_lock(thread_id: str):
    lock = _thread_locks.get(thread_id)
    if lock is None:
        lock = _thread_locks[thread_id] = asyncio.Lock()
    async with lock:
        yield


@dataclass
class PreparedTurn:
    agent: object
    payload: dict
    config: dict


async def prepare_summary(agent, payload: dict, config: dict) -> None:
    store = getattr(current_app, "conversation_summary_store", None)
    if store is None:
        return
    configurable = config["configurable"]
    try:
        checkpoint = await agent.aget_state(config)
        history = convert_to_messages(checkpoint.values.get("messages", []))
        state = await store.update(
            configurable["user_id"],
            configurable["thread_id"],
            history + convert_to_messages(payload["messages"]),
        )
        configurable["conversation_summary"] = state.as_message()
    except Exception:
        logging.exception("Failed to prepare conversation summary")


async def prepare_agent_turn(
    user,
    query: str,
    thread_id: str,
    *,
    channel: str = "web",
    messages: list | None = None,
    instructions: str = "",
) -> PreparedTurn:
    """Callers resolve and authorize the user/thread before preparing a turn.

    Only authenticated server-side user objects determine permissions. Caller
    input cannot supply a config that overrides identity or administration.
    """
    if user is None or not getattr(user, "id", None):
        raise ValueError("An authenticated user is required for agent invocation")
    if not isinstance(query, str) or not isinstance(thread_id, str) or not thread_id:
        raise ValueError("query and thread_id must be valid strings")
    prompt = await build_system_prompt_with_memories(str(user.id), query=query)
    if instructions:
        prompt += "\n\n" + instructions
    config = {
        "configurable": {
            "user_id": str(user.id),
            "is_admin": bool(user.is_admin()),
            "thread_id": thread_id,
            "system_prompt": prompt,
            "channel": channel,
            "request_id": str(uuid.uuid4()),
        }
    }
    payload = {
        "messages": list(messages)
        if messages is not None
        else [{"role": "user", "content": query}]
    }
    agent = get_agent_for(user)
    await prepare_summary(agent, payload, config)
    return PreparedTurn(agent, payload, config)


def _record(turn: PreparedTurn, started: float, result: str) -> None:
    cfg = turn.config["configurable"]
    seconds = time.monotonic() - started
    AGENT_REQUESTS_TOTAL.labels(channel=cfg["channel"], result=result).inc()
    AGENT_RESPONSE_SECONDS.labels(channel=cfg["channel"]).observe(seconds)
    logging.info(
        json.dumps(
            {
                "event": "agent_invocation",
                "request_id": cfg["request_id"],
                "channel": cfg["channel"],
                "thread_id": cfg["thread_id"],
                "user_id": cfg["user_id"],
                "result": result,
                "total_ms": round(seconds * 1000),
            }
        )
    )


async def _prune_after_turn(turn: PreparedTurn) -> None:
    checkpointer = getattr(turn.agent, "checkpointer", None)
    if checkpointer is None:
        return
    try:
        await prune_thread_checkpoints(
            checkpointer, turn.config["configurable"]["thread_id"]
        )
    except Exception:
        logging.exception("Failed to prune checkpoints")


async def invoke_prepared(turn: PreparedTurn) -> dict:
    started = time.monotonic()
    result = "success"
    try:
        return await turn.agent.ainvoke(turn.payload, config=turn.config)
    except asyncio.CancelledError:
        result = "cancelled"
        raise
    except Exception:
        result = "error"
        raise
    finally:
        _record(turn, started, result)
        await _prune_after_turn(turn)


async def invoke_agent_for_user(
    user,
    query: str,
    thread_id: str,
    *,
    channel: str = "web",
    messages: list | None = None,
    instructions: str = "",
) -> dict:
    async with thread_lock(thread_id):
        turn = await prepare_agent_turn(
            user,
            query,
            thread_id,
            channel=channel,
            messages=messages,
            instructions=instructions,
        )
        return await invoke_prepared(turn)


async def stream_prepared(turn: PreparedTurn):
    started = time.monotonic()
    result = "cancelled"
    try:
        async for event in turn.agent.astream_events(
            turn.payload, version="v2", config=turn.config
        ):
            yield event
        result = "success"
    except Exception:
        result = "error"
        raise
    finally:
        _record(turn, started, result)
        await _prune_after_turn(turn)