WealthySmart/backend/app/agent/chat_audit.py

"""Durable diagnostics for individual AG-UI agent runs."""

from __future__ import annotations

import asyncio
import contextvars
from datetime import timedelta
from typing import Any

from sqlmodel import Session, delete, select

from app.config import settings
from app.db import engine
from app.models.models import ChatRunLog
from app.timeutil import utcnow


CHAT_LOG_RETENTION_DAYS = 90
_chat_run_id: contextvars.ContextVar[int | None] = contextvars.ContextVar(
    "chat_run_id", default=None
)


def set_current_chat_run(run_id: int) -> contextvars.Token:
    return _chat_run_id.set(run_id)


def reset_current_chat_run(token: contextvars.Token) -> None:
    _chat_run_id.reset(token)


def current_chat_run_id() -> int | None:
    return _chat_run_id.get()


def _last_user_prompt(messages: list[dict[str, Any]]) -> str | None:
    for message in reversed(messages):
        if message.get("role") == "user" and isinstance(message.get("content"), str):
            return message["content"]
    return None


def create_chat_run(
    *,
    request_id: str,
    agui_run_id: str | None,
    thread_id: str | None,
    messages: list[dict[str, Any]],
    client_ip: str | None = None,
    user_agent: str | None = None,
    browser: str | None = None,
) -> int:
    now = utcnow()
    with Session(engine) as session:
        row = ChatRunLog(
            request_id=request_id,
            agui_run_id=agui_run_id,
            requested_thread_id=thread_id,
            model=settings.AGENT_MODEL,
            client_ip=client_ip,
            user_agent=user_agent,
            browser=browser,
            user_prompt=_last_user_prompt(messages),
            expires_at=now + timedelta(days=CHAT_LOG_RETENTION_DAYS),
        )
        session.add(row)
        session.commit()
        session.refresh(row)
        assert row.id is not None
        return row.id


def _tool_activity(messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
    calls: dict[str, dict[str, Any]] = {}
    order: list[str] = []
    for message in messages:
        if message.get("role") == "assistant":
            for call in message.get("toolCalls") or message.get("tool_calls") or []:
                call_id = str(call.get("id") or "")
                if not call_id or call_id in calls:
                    continue
                function = call.get("function") or {}
                calls[call_id] = {
                    "id": call_id,
                    "name": function.get("name"),
                    "arguments": function.get("arguments"),
                }
                order.append(call_id)
        elif message.get("role") == "tool":
            call_id = str(message.get("toolCallId") or message.get("tool_call_id") or "")
            if call_id in calls:
                calls[call_id]["result"] = message.get("content")
                if message.get("error") is not None:
                    calls[call_id]["error"] = message["error"]
    return [calls[call_id] for call_id in order]


def _last_assistant_response(messages: list[dict[str, Any]]) -> str | None:
    for message in reversed(messages):
        if message.get("role") == "assistant" and isinstance(message.get("content"), str):
            return message["content"]
    return None


def finalize_chat_run_in_session(
    session: Session, *, run_id: int, response_thread_id: str, messages: list[dict[str, Any]]
) -> None:
    row = session.get(ChatRunLog, run_id)
    if row is None:
        return
    row.response_thread_id = response_thread_id
    row.tool_calls = _tool_activity(messages)
    row.assistant_response = _last_assistant_response(messages)
    row.status = "completed"
    row.completed_at = utcnow()
    session.add(row)


def mark_chat_run_failed(run_id: int, error: str) -> None:
    with Session(engine) as session:
        row = session.get(ChatRunLog, run_id)
        if row is None or row.status == "completed":
            return
        row.status = "failed"
        row.error = error[:4000]
        row.completed_at = utcnow()
        session.add(row)
        session.commit()


def finish_pending_chat_run(run_id: int) -> None:
    with Session(engine) as session:
        row = session.get(ChatRunLog, run_id)
        if row is None or row.status != "started":
            return
        row.status = "completed"
        row.completed_at = utcnow()
        session.add(row)
        session.commit()


def purge_expired_chat_runs() -> int:
    with Session(engine) as session:
        result = session.exec(delete(ChatRunLog).where(ChatRunLog.expires_at < utcnow()))
        session.commit()
        return result.rowcount or 0


async def purge_expired_chat_runs_periodically() -> None:
    """Purge at startup and then once per day for a bounded audit footprint."""
    while True:
        await asyncio.to_thread(purge_expired_chat_runs)
        await asyncio.sleep(24 * 60 * 60)