EvoScientist/EvoScientist/usage/callback.py

"""LangChain terminal callback that emits one UsageEvent per real model call."""

from __future__ import annotations

import copy
import logging
import os
import threading
from collections.abc import Mapping
from dataclasses import dataclass, field
from datetime import UTC, datetime
from typing import Any
from uuid import UUID

from langchain_core.callbacks import BaseCallbackHandler

from .schema import MAX_SAFE_TOKEN_INTEGER, UsageEventV1, UsageScope
from .spool import get_usage_spool, mark_tracking_degraded

logger = logging.getLogger(__name__)


@dataclass(frozen=True)
class UsageModelIdentity:
    provider_profile_id: str
    provider_revision: str | None
    provider_adapter: str
    model_alias: str
    upstream_model_id: str


@dataclass
class _CallState:
    started_at: datetime
    parent_run_id: str | None
    metadata: dict[str, Any] = field(default_factory=dict)
    last_usage: dict[str, Any] | None = None
    provider_request_id: str | None = None


def usage_tracking_requested() -> bool:
    """Return whether this process was launched in usage-tracking mode."""
    enabled = os.getenv("EVOSCIENTIST_USAGE_TRACKING", "").strip().lower()
    return enabled in {"1", "true", "yes", "on"}


def usage_tracking_enabled() -> bool:
    """Return whether the integrated launcher supplied a complete usage sink."""
    return usage_tracking_requested() and all(
        os.getenv(name, "").strip()
        for name in (
            "EVOSCIENTIST_USAGE_SINK_URL",
            "EVOSCIENTIST_USAGE_SINK_TOKEN",
            "EVOSCIENTIST_USAGE_DEPLOYMENT_ID",
            "EVOSCIENTIST_WORKSPACE_ID",
            "EVOSCIENTIST_USAGE_SPOOL_DIR",
        )
    )


def _text(value: Any, *, limit: int = 256) -> str | None:
    if value is None:
        return None
    rendered = "".join(
        character
        for character in str(value)
        if ord(character) > 31 and ord(character) != 127
    )
    return rendered[:limit] if rendered else None


def _metadata_value(metadata: Mapping[str, Any], *keys: str) -> Any:
    for key in keys:
        value = metadata.get(key)
        if value is not None:
            return value
    return None


def _scope(metadata: Mapping[str, Any]) -> UsageScope:
    explicit = metadata.get("usage_scope")
    allowed = {
        "main", "sync_subagent", "async_subagent", "tool_selector", "summarizer",
        "memory", "scheduler", "autoskills", "diagnostic", "skill_eval", "unattributed",
    }
    if explicit in allowed:
        return explicit  # type: ignore[return-value]
    if metadata.get("lc_source") == "summarization":
        return "summarizer"
    run_kind = metadata.get("run_kind")
    if run_kind == "scheduled_task":
        return "scheduler"
    if run_kind == "evomemory_autoskills":
        return "autoskills"
    if isinstance(run_kind, str) and run_kind.startswith("evomemory_"):
        return "memory"
    # deepagents attaches this stable metadata field to every compiled agent
    # graph. Remote async runs carry an explicit usage_scope and have already
    # returned above; remaining non-main names are synchronous subagents.
    agent_name = metadata.get("lc_agent_name")
    if agent_name == "EvoScientist":
        return "main"
    if isinstance(agent_name, str) and agent_name:
        return "sync_subagent"
    if metadata.get("async_task_id") or metadata.get("source_session_id"):
        return "async_subagent"
    if metadata.get("thread_id") or metadata.get("langgraph_thread_id"):
        return "main"
    return "unattributed"


def _clean_details(value: Any) -> dict[str, int]:
    if not isinstance(value, Mapping):
        return {}
    result: dict[str, int] = {}
    for key, count in value.items():
        if (
            isinstance(key, str)
            and 0 < len(key) <= 128
            and isinstance(count, int)
            and not isinstance(count, bool)
            and 0 <= count <= MAX_SAFE_TOKEN_INTEGER
        ):
            result[key] = count
    return result


def _normalize_usage(value: Any) -> dict[str, Any] | None:
    if not isinstance(value, Mapping):
        return None
    input_tokens = value.get("input_tokens", value.get("prompt_tokens"))
    output_tokens = value.get("output_tokens", value.get("completion_tokens"))
    if not all(
        isinstance(token, int)
        and not isinstance(token, bool)
        and 0 <= token <= MAX_SAFE_TOKEN_INTEGER
        for token in (input_tokens, output_tokens)
    ):
        return None
    if input_tokens + output_tokens > MAX_SAFE_TOKEN_INTEGER:
        return None
    provider_total = value.get("total_tokens")
    if not (
        isinstance(provider_total, int)
        and not isinstance(provider_total, bool)
        and 0 <= provider_total <= MAX_SAFE_TOKEN_INTEGER
    ):
        provider_total = None
    return {
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
        "provider_total_tokens": provider_total,
        "input_token_details": _clean_details(value.get("input_token_details")),
        "output_token_details": _clean_details(value.get("output_token_details")),
    }


def _message_from_chunk(chunk: Any) -> Any:
    return getattr(chunk, "message", chunk)


def _message_usage(message: Any) -> dict[str, Any] | None:
    return _normalize_usage(getattr(message, "usage_metadata", None))


def _request_id(message: Any) -> str | None:
    metadata = getattr(message, "response_metadata", None)
    if not isinstance(metadata, Mapping):
        return None
    return _text(_metadata_value(metadata, "request_id", "id", "x_request_id"))


class UsageCaptureCallback(BaseCallbackHandler):
    """Capture final provider usage without influencing the model invocation."""

    def __init__(self, identity: UsageModelIdentity) -> None:
        self.identity = identity
        self._lock = threading.Lock()
        self._calls: dict[str, _CallState] = {}

    def on_chat_model_start(
        self,
        serialized: dict[str, Any],
        messages: list[list[Any]],
        *,
        run_id: UUID,
        parent_run_id: UUID | None = None,
        metadata: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> None:
        try:
            with self._lock:
                self._calls[str(run_id)] = _CallState(
                    started_at=datetime.now(UTC),
                    parent_run_id=str(parent_run_id) if parent_run_id else None,
                    metadata=dict(metadata or {}),
                )
        except Exception:
            logger.exception("Usage callback could not initialize run state")

    def on_llm_new_token(
        self, token: Any, *, run_id: UUID, chunk: Any = None, **kwargs: Any
    ) -> None:
        try:
            message = _message_from_chunk(chunk)
            usage = _message_usage(message)
            if usage is None:
                return
            with self._lock:
                state = self._calls.get(str(run_id))
                if state:
                    state.last_usage = usage
                    state.provider_request_id = _request_id(message) or state.provider_request_id
        except Exception:
            logger.exception("Usage callback could not inspect a stream chunk")

    def on_llm_end(self, response: Any, *, run_id: UUID, **kwargs: Any) -> None:
        usage: dict[str, Any] | None = None
        request_id: str | None = None
        try:
            for group in getattr(response, "generations", []) or []:
                for generation in group or []:
                    message = getattr(generation, "message", generation)
                    usage = _message_usage(message) or usage
                    request_id = _request_id(message) or request_id
            if usage is None:
                output = getattr(response, "llm_output", None)
                if isinstance(output, Mapping):
                    usage = _normalize_usage(output.get("token_usage") or output.get("usage"))
        except Exception:
            logger.exception("Usage callback could not inspect terminal usage")
        try:
            self._finish(str(run_id), usage=usage, request_id=request_id)
        except Exception:
            logger.exception("Usage callback could not emit terminal usage")

    def on_llm_error(self, error: BaseException, *, run_id: UUID, **kwargs: Any) -> None:
        try:
            self._finish(str(run_id), usage=None, request_id=None)
        except Exception:
            logger.exception("Usage callback could not emit unknown terminal usage")

    def _finish(
        self, run_id: str, *, usage: dict[str, Any] | None, request_id: str | None
    ) -> None:
        with self._lock:
            state = self._calls.pop(run_id, None)
        if state is None:
            return
        usage = usage or state.last_usage
        completed = datetime.now(UTC)
        metadata = state.metadata
        confirmed = usage is not None
        event = UsageEventV1(
            schema_version=1,
            event_id=(
                f"{os.environ['EVOSCIENTIST_USAGE_DEPLOYMENT_ID']}:{run_id}:callback_final:1"
            ),
            event_type="usage_observed",
            source="callback_final",
            authority_class="observed_final",
            revision=1,
            deployment_id=os.environ["EVOSCIENTIST_USAGE_DEPLOYMENT_ID"],
            workspace_id=os.environ["EVOSCIENTIST_WORKSPACE_ID"],
            model_call_id=run_id,
            parent_run_id=state.parent_run_id,
            provider_request_id=request_id or state.provider_request_id,
            thread_id=_text(
                _metadata_value(metadata, "thread_id", "langgraph_thread_id")
            ),
            source_session_id=_text(
                _metadata_value(
                    metadata, "source_session_id", "evomemory_source_session_id"
                )
            ),
            turn_id=_text(
                _metadata_value(metadata, "turn_id", "evomemory_source_turn_id")
            ),
            workspace_dir=_text(
                _metadata_value(metadata, "workspace_dir"), limit=4096
            )
            or _text(os.getenv("EVOSCIENTIST_WORKSPACE_DIR"), limit=4096),
            scope=_scope(metadata),
            source_agent=_text(
                _metadata_value(metadata, "source_agent", "evomemory_source_agent")
            ),
            provider_profile_id=_text(
                self.identity.provider_profile_id, limit=512
            )
            or "unknown",
            provider_revision=_text(self.identity.provider_revision),
            provider_adapter=_text(self.identity.provider_adapter, limit=512)
            or "unknown",
            model_alias=_text(self.identity.model_alias, limit=512) or "unknown",
            upstream_model_id=_text(
                self.identity.upstream_model_id, limit=512
            )
            or "unknown",
            usage_status="confirmed" if confirmed else "unknown",
            input_tokens=usage["input_tokens"] if usage else None,
            output_tokens=usage["output_tokens"] if usage else None,
            provider_total_tokens=usage["provider_total_tokens"] if usage else None,
            input_token_details=usage["input_token_details"] if usage else {},
            output_token_details=usage["output_token_details"] if usage else {},
            started_at=state.started_at,
            observed_at=completed,
            completed_at=completed,
        )
        get_usage_spool().enqueue(event)


def _merged_callbacks(existing: Any, callback: UsageCaptureCallback) -> Any:
    if existing is None:
        return [callback]
    if isinstance(existing, (list, tuple)):
        return [*existing, callback]
    if hasattr(existing, "add_handler"):
        manager = copy.copy(existing)
        manager.add_handler(callback, inherit=True)
        return manager
    raise TypeError("unsupported callbacks value")


def attach_usage_callback(model: Any, identity: UsageModelIdentity) -> Any:
    """Return a callback-enabled model copy, or the original model on degradation."""

    if not usage_tracking_enabled():
        return model
    try:
        from langchain_core.language_models import BaseChatModel

        if not isinstance(model, BaseChatModel):
            raise TypeError("model is not a BaseChatModel")
        callback = UsageCaptureCallback(identity)
        updated = model.model_copy(
            update={"callbacks": _merged_callbacks(model.callbacks, callback)}
        )
        if not isinstance(updated, BaseChatModel):
            raise TypeError("callback model_copy did not return a BaseChatModel")
        # Tool Selector depends on a second metadata-only model copy. Verify the
        # contract before enabling tracking for this model so a selector call
        # can never be silently counted as a main-agent call.
        probe = updated.model_copy(
            update={"metadata": {**(updated.metadata or {}), "usage_scope": "tool_selector"}}
        )
        if not isinstance(probe, BaseChatModel):
            raise TypeError("selector metadata model_copy did not return a BaseChatModel")
        # Start capability negotiation and heartbeat at backend startup rather
        # than waiting for the first completed model call. This lets the UI
        # distinguish a healthy empty database from unsupported tracking.
        get_usage_spool()
        return updated
    except Exception:
        mark_tracking_degraded("selector_model_copy_unsupported")
        logger.exception("Usage callback injection failed; model behavior is unchanged")
        return model