diff --git a/src/xagent/core/model/asr/usage.py b/src/xagent/core/model/asr/usage.py new file mode 100644 index 0000000000..00dde17144 --- /dev/null +++ b/src/xagent/core/model/asr/usage.py @@ -0,0 +1,115 @@ +"""Media-usage recording for ASR. + +ASR is reached from several entry points that do not go through ``audio_tool`` +— the Telegram channel and the ``/speech/transcribe`` API both call +``transcribe`` directly — so this module is the single place that knows how to +turn a transcription result into a usage record. Every ASR caller routes here, +including ``audio_tool``: keeping one implementation is what stops the copies +from drifting (an earlier pair differed in whether they logged at all). +""" + +from __future__ import annotations + +import logging +from typing import Any, Optional, Sequence, Union + +from ...tools.core.media_usage import coerce_duration, record_media_seconds +from ..chat.token_context import MediaCallType +from .base import ASRResult + +logger = logging.getLogger(__name__) + +# Keys providers use for the total length of the submitted audio. +_DURATION_KEYS = ("duration", "audio_duration", "duration_seconds") + + +def duration_from_raw_response(raw_response: Any) -> Optional[float]: + """Provider-reported total audio duration, if it exposed one. + + Preferred over segment timings: the end of the last spoken segment + undercounts a recording with trailing silence, which the provider still + processed and billed for. + """ + if not isinstance(raw_response, dict): + return None + for key in _DURATION_KEYS: + seconds = coerce_duration(raw_response.get(key)) + if seconds is not None: + return seconds + return None + + +def duration_from_segments(segments: Optional[Sequence[Any]]) -> Optional[float]: + """Transcribed duration inferred from the end of the last timed segment. + + Accepts both ``ASRSegment`` objects and the plain dicts ``audio_tool`` + builds from them. + """ + if not segments: + return None + last_end = 0.0 + for segment in segments: + end = ( + segment.get("end") + if isinstance(segment, dict) + else getattr(segment, "end", None) + ) + seconds = coerce_duration(end) + if seconds is not None: + last_end = max(last_end, seconds) + return last_end if last_end > 0 else None + + +def resolve_asr_seconds( + raw_response: Any = None, + segments: Optional[Sequence[Any]] = None, +) -> Optional[float]: + """Best available transcribed-audio duration, provider field first.""" + return duration_from_raw_response(raw_response) or duration_from_segments(segments) + + +def record_asr_seconds( + seconds: Optional[float], + *, + model_name: str = "", + model_id: str = "", +) -> None: + """Record one transcription from an already-resolved duration. + + Thin wrapper over the shared :func:`record_media_seconds`: ASR is one of + several duration-billed modalities, and keeping a second implementation is + how the two drifted before. The shared helper owns the invariants — unit + always seconds (a call whose duration is unknown records 0 rather than + switching units, which would make a (model, unit) price table unusable), + the unmeasured-call warning, and swallowing errors so accounting never + breaks a transcription. + """ + record_media_seconds( + seconds, + model=model_name, + model_id=model_id, + call_type=MediaCallType.ASR, + ) + + +def record_asr_usage( + result: Union[str, ASRResult], + *, + model_name: str = "", + model_id: str = "", +) -> None: + """Record one transcription from a provider result. + + A non-verbose call returns a bare string with no timing information, so it + is recorded as unmeasured rather than silently skipped. + """ + raw_response = None + segments = None + if isinstance(result, ASRResult): + raw_response = result.raw_response + segments = result.segments + record_asr_seconds( + resolve_asr_seconds(raw_response, segments), + model_name=model_name, + model_id=model_id, + ) diff --git a/src/xagent/core/tools/core/audio_tool.py b/src/xagent/core/tools/core/audio_tool.py index 168b92e38e..66e9604695 100644 --- a/src/xagent/core/tools/core/audio_tool.py +++ b/src/xagent/core/tools/core/audio_tool.py @@ -17,6 +17,8 @@ from ...file_ref import build_workspace_file_ref from ...model.asr.base import ASRResult, BaseASR +from ...model.asr.usage import record_asr_seconds, resolve_asr_seconds +from ...model.chat.token_context import MediaCallType from ...model.tts.base import BaseTTS, TTSResult from ...workspace import TaskWorkspace from .audio_tool_descriptions import ( @@ -27,6 +29,7 @@ SYNTHESIZE_SPEECH_JSON_DESCRIPTION, TRANSCRIBE_AUDIO_DESCRIPTION, ) +from .media_usage import record_media_usage, resolve_billing_model logger = logging.getLogger(__name__) @@ -572,6 +575,33 @@ def _get_tts_model_id(self, tts_model: BaseTTS) -> str: return model_id return "default" + @staticmethod + def _resolve_billing_model( + models: Dict[str, Any], model: Any, model_id: Optional[str] + ) -> str: + """Identify the model that actually served a call, for usage records. + + Resolves by configured key, then by object identity (the batch path + hands back a different instance than the registry dict holds), then + delegates to the shared :func:`resolve_billing_model` for the + provider-attribute fallback and the placeholder filter — so the + Xinference default model, which exposes no ``model_name``, is never + billed under the literal string ``"default"``. + """ + if model_id and model_id in models: + return model_id + if model is not None: + for configured_id, configured_model in models.items(): + if configured_model is model: + return configured_id + # The provider class name is the last identity that still attributes + # cost to something real. The shared helper's own fallback is the + # literal "default", which the metering invariants forbid as a billing + # identity — the Xinference default model reaches exactly that case, + # exposing no model_name and never matching by identity above. + fallback = type(model).__name__ if model is not None else "default" + return resolve_billing_model(None, model, fallback=fallback) + def _get_provider_tts_model( self, *, @@ -713,14 +743,15 @@ async def transcribe_audio( ) # Determine the actual model used - actual_model_id = ( - model_id if model_id and model_id in self._asr_models else "default" + actual_model_id = self._resolve_billing_model( + self._asr_models, asr_model, model_id ) # Handle different result types text = None raw_segments = None language_detected = None + raw_response: dict[str, Any] = {} if isinstance(result, str): text = result @@ -741,6 +772,35 @@ async def transcribe_audio( else None ) language_detected = result.language + if isinstance(result.raw_response, dict): + raw_response = result.raw_response + + # Shared with /speech/transcribe and the Telegram channel so the + # duration rule (provider's own total first, segment timings only + # as a fallback) has exactly one implementation. + audio_seconds = resolve_asr_seconds(raw_response, raw_segments) + + # Recorded here, before _aggregate_segments and the rest of the + # post-processing below. Those steps can raise (a segment missing + # start/end raises ValueError), and the broad handler at the bottom + # turns any raise into success:False -- so recording afterwards + # meant a provider call that succeeded and was billed went + # unmetered. That is the bug class this metering exists to remove, + # and TTS in this same file already records before unpacking. + # Routed through the shared ASR recorder so every entry point + # meters identically. + record_asr_seconds( + audio_seconds, + # model_id is deliberately left unset. This tool is handed + # name-keyed registries (model_service keys _asr_models by + # model_name), so the only identity in scope here is a name — + # writing a name into model_id would persist it under an id + # field. The aggregator groups on `model_id or model`, so + # leaving it empty lets these rows key on the name + # consistently instead of inventing a third identity that can + # never be reconciled. + model_name=str(actual_model_id), + ) segment_view = "raw" if verbose else "processed" segments = raw_segments @@ -916,8 +976,19 @@ async def synthesize_speech( ) # Determine the actual model used - actual_model_id = ( - model_id if model_id and model_id in self._tts_models else "default" + actual_model_id = self._resolve_billing_model( + self._tts_models, tts_model, model_id + ) + + # Meter TTS by input characters (how providers like ElevenLabs + # bill). Recorded before the result is unpacked below: the provider + # call already happened and is billable regardless of what fails + # afterwards. model_id stays unset — see synthesize/ASR above, the + # registries here are name-keyed so only a name is ever in scope. + record_media_usage( + MediaCallType.TTS, + len(text or ""), + model=str(actual_model_id), ) audio_data: Optional[bytes] = None @@ -1873,6 +1944,20 @@ async def _synthesize_single_segment( **kwargs, ) + # Meter TTS by input characters, matching synthesize_speech so batch + # synthesis is metered the same way as single-shot synthesis. + # _resolve_billing_model rather than _get_tts_model_id: the latter + # bottoms out at the literal "default" with no model_name fallback, + # so an implicitly-resolved model would bill a phantom name here. + batch_model_id = self._resolve_billing_model( + self._tts_models, tts_model, None + ) + record_media_usage( + MediaCallType.TTS, + len(text or ""), + model=batch_model_id, + ) + # Handle result if isinstance(audio_data, bytes): audio_binary = audio_data diff --git a/src/xagent/core/tools/core/music_tool.py b/src/xagent/core/tools/core/music_tool.py index 1ea0baf387..0cf9cc23d1 100644 --- a/src/xagent/core/tools/core/music_tool.py +++ b/src/xagent/core/tools/core/music_tool.py @@ -8,8 +8,10 @@ from typing import Any, Optional from ...file_ref import build_workspace_file_ref +from ...model.chat.token_context import MediaCallType from ...model.music import BaseMusicModel, MusicResult from ...workspace import TaskWorkspace +from .media_usage import coerce_duration, record_media_seconds, resolve_billing_model logger = logging.getLogger(__name__) @@ -137,6 +139,45 @@ async def generate_music( force_instrumental=force_instrumental, output_format=output_format, ) + # Billing policy for the media tools: a provider call that has + # already happened is billable regardless of what fails afterwards, + # so usage is recorded before the response is validated. Metering + # after the checks below would silently drop every call that + # succeeded at the HTTP level but came back empty or malformed — + # which the provider still charges for. + # + # Music is duration-billed: always meter in seconds, even when the + # length was auto-selected and no duration came back. `result` is + # not yet known to be a MusicResult at this point, so read through + # getattr rather than attribute access. + raw_response = getattr(result, "raw_response", None) + reported_length = ( + raw_response.get("music_length_seconds") + if isinstance(raw_response, dict) + else None + ) + seconds = coerce_duration(reported_length) or coerce_duration( + music_length_seconds + ) + record_media_seconds( + seconds, + # `model` carries the provider's name, `model_id` the + # configured id. Passing configured_model_id as the first + # argument here would return it unchanged, writing the same id + # into both fields and losing the provider name entirely. + # None is passed instead so the resolver falls through to the + # provider's model_name. The fallback keeps the configured id + # ahead of the class name -- a provider exposing no model_name + # (Xinference's default) would otherwise be billed under a + # Python class name while its real id sat unused in scope -- + # and never the placeholder "default". + model=resolve_billing_model( + None, model, fallback=configured_model_id or type(model).__name__ + ), + model_id=configured_model_id or "", + call_type=MediaCallType.MUSIC, + ) + if not isinstance(result, MusicResult): raise RuntimeError(f"Unexpected music response: {type(result)}") if not result.audio: diff --git a/src/xagent/core/tools/core/sound_effect_tool.py b/src/xagent/core/tools/core/sound_effect_tool.py index d404551007..40cabfbc9d 100644 --- a/src/xagent/core/tools/core/sound_effect_tool.py +++ b/src/xagent/core/tools/core/sound_effect_tool.py @@ -8,8 +8,10 @@ from typing import Any, Optional from ...file_ref import build_workspace_file_ref +from ...model.chat.token_context import MediaCallType from ...model.sound_effect import BaseSoundEffectModel, SoundEffectResult from ...workspace import TaskWorkspace +from .media_usage import coerce_duration, record_media_seconds, resolve_billing_model logger = logging.getLogger(__name__) @@ -151,6 +153,39 @@ async def generate_sound_effect( loop=loop, output_format=output_format, ) + # Billing policy for the media tools: a provider call that has + # already happened is billable regardless of what fails afterwards, + # so usage is recorded before the response is validated — a call + # that returned empty or malformed audio was still charged for. + # + # ElevenLabs prices sound effects by duration, so seconds is the + # only meaningful unit here — an auto-length call with no reported + # duration records 0 seconds rather than switching to characters, + # which would be wrong in kind, not just in magnitude. `result` is + # not yet known to be a SoundEffectResult, so read via getattr. + raw_response = getattr(result, "raw_response", None) + reported_duration = ( + raw_response.get("duration_seconds") + if isinstance(raw_response, dict) + else None + ) + seconds = coerce_duration(reported_duration) or coerce_duration( + duration_seconds + ) + record_media_seconds( + seconds, + # `model` is the provider name, `model_id` the configured id — + # see music_tool: passing configured_model_id as the first + # argument would put the same id in both fields and drop the + # provider name, while the fallback keeps it ahead of the + # class name for providers exposing no model_name. + model=resolve_billing_model( + None, model, fallback=configured_model_id or type(model).__name__ + ), + model_id=configured_model_id or "", + call_type=MediaCallType.SOUND_EFFECT, + ) + if not isinstance(result, SoundEffectResult): raise RuntimeError(f"Unexpected sound effect response: {type(result)}") if not result.audio: diff --git a/src/xagent/core/tools/core/video_tool.py b/src/xagent/core/tools/core/video_tool.py index b995d74771..443c4edf50 100644 --- a/src/xagent/core/tools/core/video_tool.py +++ b/src/xagent/core/tools/core/video_tool.py @@ -22,10 +22,16 @@ from pydantic import Field from ...file_ref import build_workspace_file_ref, guess_mime_type, parse_file_id_ref +from ...model.chat.token_context import MediaCallType from ...model.video.ark import ArkVideoModel from ...model.video.base import BaseVideoModel from ...model.video.xinference import XinferenceVideoModel from ...workspace import TaskWorkspace +from .media_usage import ( + coerce_duration, + record_media_seconds, + resolve_billing_model, +) logger = logging.getLogger(__name__) @@ -721,6 +727,41 @@ async def generate_video( result = await video_model.generate_video(**generate_params) + # Video is duration-billed, so always meter in seconds — never + # switch units when the provider omits a duration (async tasks + # started with wait_for_result=False have none yet). + # + # The provider reports one duration but generates n videos and + # bills for all of them, so the billable total is duration * n. + # Ark rejects n>1; Xinference does not, and previously only the + # first video's duration was recorded. + per_video_seconds = coerce_duration(result.get("duration")) + billable_count = max(1, int(n or 1)) + billing_model_id = str(actual_model_id) + record_media_seconds( + per_video_seconds * billable_count + if per_video_seconds is not None + else None, + # `model` is the provider's name and `model_id` the configured + # id. Writing the configured id into `model` and leaving + # `model_id` empty loses the canonical name for display and + # external consumers; the aggregator groups on + # `model_id or model`, so populating both keeps the same row + # identity while carrying the name. + model=resolve_billing_model( + None, video_model, fallback=str(actual_model_id) + ), + # _model_id_for_model bottoms out at the literal "default" when + # a model exposes neither an id nor a name, and that string is + # a forbidden billing identity. Drop it rather than persist it, + # matching music/sound_effect which pass `configured_id or ""`. + # The aggregator groups on `model_id or model`, so an empty id + # falls through to the resolved name instead of inventing a + # phantom model shared by every unidentifiable provider. + model_id=("" if billing_model_id == "default" else billing_model_id), + call_type=MediaCallType.VIDEO, + ) + video_url = result.get("video_url") video_path = None video_file_id: Optional[str] = None diff --git a/src/xagent/web/api/model.py b/src/xagent/web/api/model.py index 77a162503a..da96b9e4bf 100644 --- a/src/xagent/web/api/model.py +++ b/src/xagent/web/api/model.py @@ -980,6 +980,10 @@ async def transcribe_speech_input( """Transcribe a short UI voice input clip with an accessible ASR model.""" from xagent.core.model.asr.adapter import get_asr_model_instance + from xagent.core.model.asr.usage import record_asr_usage + from xagent.core.tools.core.media_usage import resolve_billing_model + + from ..tracking.standalone_usage import usage_scope try: audio_bytes = await _read_transcribe_upload_with_size_limit(file) @@ -991,14 +995,46 @@ async def transcribe_speech_input( db_model = _resolve_asr_model_for_transcription(db, user, model_id) asr_model = get_asr_model_instance(db_model) try: - result = await asyncio.wait_for( - asr_model.transcribe( - audio_bytes, - language=language, - format=_format_from_upload(file), - ), - timeout=180.0, - ) + # usage_scope binds a TokenUsage and reports it to the quota hook on + # exit. Recording alone is not enough here: this endpoint has no + # TaskTracker, so without a bound context the usage would land in a + # throwaway object and the transcription would bill nothing. + with usage_scope(int(user.id) if user and user.id is not None else None): + result = await asyncio.wait_for( + asr_model.transcribe( + audio_bytes, + language=language, + format=_format_from_upload(file), + # Without this the provider returns a bare string carrying + # no timings, so the usage record below would be a + # 0-second — unbillable — entry on every call. The extra + # fields are discarded; only `text` is returned to the UI. + verbose=True, + ), + timeout=180.0, + ) + # model_id is deliberately left unset even though the real DB id is + # in scope here. The aggregator groups on `model_id or model`, and + # the other two ASR entry points (audio_tool, the Telegram channel) + # only ever see name-keyed model registries — model_service keys + # them by model_name — so no DB id is available there. Setting it + # on this path alone would split one physical model into two + # billing rows that can never be reconciled once written. Name is + # the one identity all three share. + record_asr_usage( + result, + # Routed through the shared resolver rather than str() on the + # DB column: model_name is nullable=False but not constrained + # to be non-placeholder, so an empty or literally "default" + # name would otherwise be billed as a model identity, which + # the module invariants forbid. The configured id is the + # fallback here -- it is the one real identity in scope -- and + # it only ever reaches `model`, never `model_id`, so the + # single-identity rule above still holds. + model_name=resolve_billing_model( + None, db_model, fallback=str(db_model.model_id) + ), + ) except asyncio.TimeoutError as exc: raise HTTPException(status_code=504, detail="Transcription timed out") from exc except Exception as exc: diff --git a/src/xagent/web/channels/telegram/bot.py b/src/xagent/web/channels/telegram/bot.py index 55b8f335dc..80e47dd51c 100644 --- a/src/xagent/web/channels/telegram/bot.py +++ b/src/xagent/web/channels/telegram/bot.py @@ -1368,6 +1368,9 @@ async def _transcribe_uploaded_voice_files( uploaded_info: list[dict[str, Any]], asr_model: Any, ) -> dict[str, str]: + from ....core.model.asr.usage import record_asr_usage + from ....core.tools.core.media_usage import resolve_billing_model + uploaded_by_source_id = { str(info.get("telegram_file_id")): info for info in uploaded_info @@ -1387,9 +1390,30 @@ async def _transcribe_uploaded_voice_files( asr_model.transcribe( audio=str(file_info["path"]), format=self._audio_format_from_file_info(file_info), + # Without this the provider returns a bare string + # carrying no timings, making every voice message a + # 0-second — unbillable — usage record. Only `text` is + # used below; the extra fields are discarded. + verbose=True, ), timeout=self.voice_transcription_timeout_seconds, ) + # Telegram calls the ASR provider directly rather than going + # through audio_tool, so meter here too. Identity goes through + # the shared resolver so all three ASR entry points agree on + # the name the aggregator groups by. The explicit fallback + # matters: resolve_billing_model's own default is the literal + # "default", which the module invariants forbid as a billing + # identity, so a provider exposing no model_name would + # otherwise be billed under exactly the placeholder this is + # meant to avoid. The provider class name is the last identity + # that still attributes cost to something real. + record_asr_usage( + result, + model_name=resolve_billing_model( + None, asr_model, fallback=type(asr_model).__name__ + ), + ) except asyncio.TimeoutError as exc: raise TelegramVoiceTranscriptionError( "Telegram voice transcription timed out" diff --git a/src/xagent/web/services/quota_hooks.py b/src/xagent/web/services/quota_hooks.py index a215b5c7e5..e6dbf5a92e 100644 --- a/src/xagent/web/services/quota_hooks.py +++ b/src/xagent/web/services/quota_hooks.py @@ -23,9 +23,36 @@ # callbacks may rely on that affinity and must remain non-blocking. _run_gate_hook: Callable[[Any, Any], str | Mapping[str, Any] | None] | None = None # (db, user_id, delta_details, delta_actions) -> None; best-effort post-run -# metering. delta_details is this turn's per-model token breakdown (list of -# {"type","tokens","model"}) for cost-based credits; delta_actions counts tool -# calls (one billable action per tool invocation). +# metering. delta_details is this turn's per-model usage breakdown for +# cost-based credits; delta_actions counts tool calls (one billable action per +# tool invocation). Entry shapes in delta_details: +# - LLM tokens: {"type":"input"|"output", "tokens", "model", "model_id", +# "call_type", ...cache fields} +# - Non-LLM media (image/video/tts/asr/embedding/rerank/...): +# {"type":"media", "unit":"images"|"seconds"|"characters"|"texts"| +# "requests", "quantity", "provider_tokens", "provider_input_tokens", +# "provider_output_tokens", "tokens_estimated", "model", "model_id", +# "call_type", "resolution"} +# The app layer should price media entries by their "unit"/"quantity". The unit +# is stable for a given (model, call_type) — a duration-billed modality always +# reports "seconds", recording quantity=0 when the provider gave no duration, +# rather than switching units. "requests" always carries quantity=1; a batch of +# N embedded texts reports unit="texts" with quantity=N and calls=1. +# For image models whose price varies by resolution, "resolution" ("1K"/"2K"/ +# "4K" or "1024x1024") keys a per-(model, resolution) price table. Providers +# that report real image tokens (Gemini, OpenAI gpt-image) also fill +# "provider_tokens", recorded raw for a future consumer. That is NOT a pricing +# rule: no precedence between token-based and unit/resolution-based pricing is +# defined here, the row schema carries no discriminator to express one, and the +# aggregate groups purely by (model, unit, call_type, resolution). See the +# authoritative contract on ``add_media_usage`` in +# ``core/model/chat/token_context.py``; defining that precedence is tracked in +# #1461. When a hook does price "provider_tokens", note "tokens_estimated" +# marks counts that are local heuristics (embedding/rerank) rather than +# provider-reported, so they must not be priced as measured tokens. +# Media token counts are deliberately NOT under the "tokens" key, so a consumer +# summing "tokens" across entries cannot double-count them. +# Unknown entry types must be ignored, not summed as tokens. # # TRANSACTION CONTRACT: the hook is invoked from TaskTracker.complete_tracking # only after the run/runner-fenced token-usage update commits. The hook owns @@ -41,7 +68,8 @@ # in-flight run's live-so-far usage would push the team over a run-gated quota, # else None. Polled per step (each LLM reply / tool call) during a run so a # single long/expensive run is stopped mid-flight instead of only being metered -# at completion. +# at completion. delta_details carries the same entry shapes documented on the +# metering hook above (LLM token entries plus type:"media" entries). # # CONTRACT: invoked SYNCHRONOUSLY on the event loop once per step. It MUST NOT # block (no synchronous network/DB round-trips per call) — blocking work stalls @@ -90,6 +118,18 @@ def check_run_gate(db: Any, user_id: Any) -> str | Mapping[str, Any] | None: return _run_gate_hook(db, user_id) +def has_usage_record_hook() -> bool: + """Whether a usage-record hook is installed. + + Lets a caller skip the work of preparing a report — notably checking out a + DB session — when :func:`record_usage` would be a no-op anyway. In the + stock open-source configuration no hook is registered, so every ingest and + transcription would otherwise pay a pool checkout, transaction and close + for nothing. + """ + return _usage_record_hook is not None + + def record_usage( db: Any, user_id: Any, delta_details: list, delta_actions: int ) -> None: diff --git a/src/xagent/web/tracking/standalone_usage.py b/src/xagent/web/tracking/standalone_usage.py new file mode 100644 index 0000000000..eff592b2fa --- /dev/null +++ b/src/xagent/web/tracking/standalone_usage.py @@ -0,0 +1,103 @@ +"""Usage metering for work that is not a tracked agent task. + +``TaskTracker`` binds a ``TokenUsage`` for chat/agent runs and reports the +delta to the quota hook when the run completes. Everything else — KB ingestion +over HTTP or Celery, ``/speech/transcribe``, Telegram voice — records usage +with no context bound, so ``get_token_usage()`` lazily creates a throwaway +object that nothing ever reads. The calls succeed, the provider bills, and the +usage silently evaporates. + +This module is the equivalent sink for those paths: bind a context around the +work, then report whatever was recorded. + +Why a shared helper rather than a `TokenContextManager` at each call site: the +quota hook has a transaction contract that is easy to violate by accident. It +must not be handed a caller's request Session (it manages its own durability +and must not commit or leave writes pending on someone else's session), so the +report step opens and disposes a short-lived compatibility Session of its own. +Reproducing that at four call sites would mean four chances to get it wrong. +""" + +from __future__ import annotations + +import logging +from contextlib import contextmanager +from typing import Iterator, Optional + +from ...core.model.chat.token_context import TokenUsage, set_token_usage + +logger = logging.getLogger(__name__) + + +def _report(user_id: Optional[int], usage: TokenUsage) -> None: + """Hand this unit of work's usage to the quota hook, best-effort.""" + details = [dict(item) for item in usage.details if isinstance(item, dict)] + if not details: + return + try: + from ..services.quota_hooks import has_usage_record_hook + + # Imported lazily, like the hook above: this module is deliberately + # light so importing it costs nothing at startup, and task_tracker + # pulls in the whole tracking chain. + from .task_tracker import _record_usage_on_event_loop + + # Check the hook before checking out a session: with no hook installed + # (the stock configuration) record_usage is a guaranteed no-op, and a + # pool checkout + transaction + close per ingest/transcription is pure + # overhead. The sibling path in task_tracker predates this check and + # still pays it; extending the short-circuit there is a separate change + # to a hot path this PR does not otherwise touch. + if not has_usage_record_hook(): + return + + # Reuses task_tracker's helper rather than repeating its session + # lifecycle: it already owns the "hand the hook a short-lived + # compatibility Session, never leave it holding a transaction" + # contract, and two copies of that is how they drift. + # delta_actions=0: these paths make provider calls, not agent tool + # calls, and tool invocations are what that counter bills for. + _record_usage_on_event_loop(user_id, details, 0) + except Exception as e: # noqa: BLE001 + # Metering must never break the work it is measuring. + logger.warning("Standalone usage recording failed: %s", e) + + +@contextmanager +def usage_scope(user_id: Optional[int]) -> Iterator[TokenUsage]: + """Bind a usage context for one unit of non-task work and report it after. + + Usage is reported even when the body raises: a provider call that already + happened is billable regardless of what fails afterwards. + + Note the body must not cross a thread boundary that drops contextvars. + ``asyncio.to_thread`` copies the context and is safe; a bare + ``ThreadPoolExecutor`` or ``run_in_executor`` is not, and would need the + caller's usage bound explicitly inside the worker. + """ + from ...core.model.chat.token_context import token_context + + usage = TokenUsage() + # Restore whatever was bound before (usually None) so a nested scope cannot + # leak its usage object into the caller's. + previous = token_context.get(None) + set_token_usage(usage) + try: + yield usage + finally: + try: + # Only restore if we are still the bound context. A TaskTracker + # started *inside* this scope also calls set_token_usage, and + # clobbering that would silently detach it from its own run. + # Unreachable at today's call sites, but cheap to get right before + # a sixth one appears. + if token_context.get(None) is usage: + set_token_usage(previous) # type: ignore[arg-type] + else: + logger.debug( + "usage_scope exiting with a different context bound; " + "leaving it in place rather than detaching its owner" + ) + except Exception: # noqa: BLE001 + pass + _report(user_id, usage) diff --git a/tests/core/tools/core/test_media_tool_billing_policy.py b/tests/core/tools/core/test_media_tool_billing_policy.py new file mode 100644 index 0000000000..db3730d2c8 --- /dev/null +++ b/tests/core/tools/core/test_media_tool_billing_policy.py @@ -0,0 +1,484 @@ +"""The media tools bill a provider call that happened, even if it then fails. + +Music, sound effect, TTS and ASR are otherwise-symmetric tools that had three +different implicit "is this call billable" rules. The policy is now uniform: +usage is recorded as soon as the provider call returns, before the response is +validated, because a call that succeeded at the HTTP level but came back empty +or malformed was still charged for. +""" + +from pathlib import Path +from typing import Any, Optional + +import pytest + +from xagent.core.model.chat.token_context import TokenContextManager +from xagent.core.model.music.base import MusicResult +from xagent.core.tools.core.music_tool import MusicToolCore +from xagent.core.tools.core.sound_effect_tool import SoundEffectToolCore + +# Sentinel distinguishing "this fake has no model_name attribute at all" +# (Xinference's default model behaves this way) from "its name is empty". +_NO_NAME = object() + +# The two ways a provider call can succeed at the HTTP level and still yield +# nothing usable. Both are billed, which is the policy this module guards. +_EMPTY_MUSIC = MusicResult(audio=b"", format="mp3", raw_response={}) +_GARBAGE = {"not": "a result object"} +_PLAYABLE_MUSIC = MusicResult(audio=b"x", format="mp3", raw_response={}) + + +def _music_model(result: Any = _GARBAGE, *, name: Any = "music-a") -> Any: + """A music provider returning ``result``; ``name=_NO_NAME`` exposes none.""" + + class _FakeMusicModel: + async def generate_music(self, **kwargs: Any) -> Any: + _ = kwargs + return result + + if name is not _NO_NAME: + _FakeMusicModel.model_name = name # type: ignore[attr-defined] + return _FakeMusicModel() + + +def _sound_effect_model(*, name: Any = "sfx-a") -> Any: + """A sound-effect provider returning the wrong type; still billed.""" + + class _FakeSoundEffectModel: + async def generate_sound_effect(self, **kwargs: Any) -> Any: + _ = kwargs + return _GARBAGE + + if name is not _NO_NAME: + _FakeSoundEffectModel.model_name = name # type: ignore[attr-defined] + return _FakeSoundEffectModel() + + +def _media_entries(manager: TokenContextManager) -> list[dict]: + return [d for d in manager.get_usage().details if d.get("type") == "media"] + + +@pytest.mark.parametrize( + ("bad_result", "seconds"), + [ + pytest.param(_EMPTY_MUSIC, 30.0, id="well-formed-but-empty"), + pytest.param(_GARBAGE, 12.0, id="wrong-type-entirely"), + ], +) +@pytest.mark.asyncio +async def test_music_bills_a_call_that_returned_nothing_usable( + bad_result: Any, seconds: float +) -> None: + """Both ways a call can succeed at the HTTP level and yield nothing. + + The unit stays seconds in both: a (model, unit) price table breaks if the + unit varies with how complete the response happened to be. + """ + tool = MusicToolCore(models={"music-a": _music_model(bad_result)}) + + with TokenContextManager() as manager: + result = await tool.generate_music(prompt="p", music_length_seconds=seconds) + entries = _media_entries(manager) + + # The tool still reports failure to the caller... + assert result["success"] is False + # ...but the provider call happened and was billed, so it is metered. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == seconds + assert entries[0]["call_type"] == "music" + + +@pytest.mark.asyncio +async def test_sound_effect_bills_malformed_response() -> None: + tool = SoundEffectToolCore(models={"sfx-a": _sound_effect_model()}) + + with TokenContextManager() as manager: + result = await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert result["success"] is False + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 4.0 + assert entries[0]["call_type"] == "sound_effect" + + +class _RecordingASR: + """Captures the verbose flag the caller passed.""" + + model_name = "asr-a" + + def __init__(self) -> None: + self.verbose_seen: Optional[bool] = None + + async def transcribe(self, *args: Any, verbose: bool = False, **kwargs: Any) -> Any: + _ = args, kwargs + self.verbose_seen = verbose + if not verbose: + return "bare string, no timings" + from xagent.core.model.asr.base import ASRResult, ASRSegment + + return ASRResult( + text="hi", + segments=[ASRSegment(text="hi", start=0.0, end=2.5, confidence=1.0)], + language="eng", + ) + + +@pytest.mark.asyncio +async def test_asr_usage_is_unmeasured_without_verbose() -> None: + # Guards the reason /speech/transcribe and the Telegram channel now pass + # verbose=True: without it the provider returns a bare string with no + # timings, and every call is metered as an unbillable 0 seconds. + from xagent.core.model.asr.usage import record_asr_usage + + with TokenContextManager() as manager: + record_asr_usage("bare string, no timings", model_name="asr-a") + entries = _media_entries(manager) + + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 0.0 + + +@pytest.mark.asyncio +async def test_asr_usage_meters_duration_with_verbose() -> None: + from xagent.core.model.asr.base import ASRResult, ASRSegment + from xagent.core.model.asr.usage import record_asr_usage + + result = ASRResult( + text="hi", + segments=[ASRSegment(text="hi", start=0.0, end=2.5, confidence=1.0)], + language="eng", + ) + + with TokenContextManager() as manager: + record_asr_usage(result, model_name="asr-a") + entries = _media_entries(manager) + + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 2.5 + # Name-keyed identity across all three ASR entry points; see PR body. + assert entries[0]["model"] == "asr-a" + assert entries[0]["model_id"] == "" + + +# --- ASR duration precedence (resolve_asr_seconds) --------------------------- +# +# The producer tests above cover a bare string and one ordered segment. The +# resolver has several more branches — each provider total field, the fallback +# to segment ends, unsorted segments and unusable values — and a regression in +# any of them would silently change the billed quantity while leaving the +# suite green. + + +@pytest.mark.parametrize("field", ["duration", "audio_duration", "duration_seconds"]) +def test_provider_total_duration_wins_over_segments(field: str) -> None: + """A provider total covers trailing silence the last segment end misses.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds( + {field: 30.0}, + [{"start": 0.0, "end": 2.5}], + ) + assert seconds == 30.0 + + +def test_segment_ends_are_used_when_no_provider_total() -> None: + from xagent.core.model.asr.usage import resolve_asr_seconds + + assert resolve_asr_seconds({}, [{"start": 0.0, "end": 7.5}]) == 7.5 + + +def test_unsorted_segments_take_the_maximum_end() -> None: + """Segment order is not guaranteed, so the last element is not the end.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds( + None, + [{"start": 5.0, "end": 9.0}, {"start": 0.0, "end": 3.0}], + ) + assert seconds == 9.0 + + +@pytest.mark.parametrize("bad", [None, "abc", float("inf"), float("nan"), -1.0, True]) +def test_unusable_provider_total_falls_back_to_segments(bad: Any) -> None: + """A non-finite/negative/boolean total must not be billed as a duration.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + seconds = resolve_asr_seconds({"duration": bad}, [{"start": 0.0, "end": 4.0}]) + assert seconds == 4.0 + + +def test_no_usable_timing_reports_none() -> None: + """None (not 0.0) so the caller can tell "unmeasured" from "zero-length"; + record_media_seconds turns it into a 0-second row plus a warning.""" + from xagent.core.model.asr.usage import resolve_asr_seconds + + assert resolve_asr_seconds({}, []) is None + assert resolve_asr_seconds(None, None) is None + assert resolve_asr_seconds({"duration": None}, [{"start": 0.0, "end": 0.0}]) is None + + +# --- Media identity field shape --------------------------------------------- +# +# Convention: `model` carries the human-readable provider name, `model_id` the +# configured id. Writing the configured id into both (or leaving model_id +# empty) loses the canonical name for display and external consumers. Each +# case below uses a provider name and a configured id that differ, so a +# resolver that returns the id for both fields cannot pass. + + +@pytest.mark.asyncio +async def test_music_records_provider_name_and_configured_id_separately() -> None: + """`model` carries the provider's name, `model_id` the configured id. + + Writing the id into both fields loses the canonical name for display and + for external consumers. + """ + model = _music_model(_PLAYABLE_MUSIC, name="music-provider-name") + tool = MusicToolCore(models={"configured-music-id": model}) + + with TokenContextManager() as manager: + await tool.generate_music(prompt="p", music_length_seconds=10) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "music-provider-name" + assert entries[0]["model_id"] == "configured-music-id" + + +@pytest.mark.asyncio +async def test_sound_effect_records_provider_name_and_configured_id_separately() -> ( + None +): + model = _sound_effect_model(name="sfx-provider-name") + tool = SoundEffectToolCore(models={"configured-sfx-id": model}) + + with TokenContextManager() as manager: + await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "sfx-provider-name" + assert entries[0]["model_id"] == "configured-sfx-id" + + +@pytest.mark.asyncio +async def test_music_falls_back_to_configured_id_not_the_class_name() -> None: + """A provider with no model_name must still bill under its configured id. + + The resolver's fallback is what decides this, and a Python class name is + not a billing identity: it is not configurable, not unique across + providers, and means nothing to a price table. + """ + model = _music_model(_PLAYABLE_MUSIC, name=_NO_NAME) + tool = MusicToolCore(models={"configured-music-id": model}) + + with TokenContextManager() as manager: + await tool.generate_music(prompt="p", music_length_seconds=10) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "configured-music-id" + assert entries[0]["model"] != type(model).__name__ + assert entries[0]["model_id"] == "configured-music-id" + + +@pytest.mark.asyncio +async def test_sound_effect_falls_back_to_configured_id_not_the_class_name() -> None: + model = _sound_effect_model(name=_NO_NAME) + tool = SoundEffectToolCore(models={"configured-sfx-id": model}) + + with TokenContextManager() as manager: + await tool.generate_sound_effect(text="p", duration_seconds=4) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "configured-sfx-id" + assert entries[0]["model"] != type(model).__name__ + + +# --- Video billing -------------------------------------------------------- +# +# video_tool is duration-billed like music/sound effect, but bills +# `duration * n` because the provider reports one duration while generating +# and charging for n videos. These cover the multiplication, the identity +# fields, and the async path that reports no duration. + + +def _video_model(duration: Any, *, name: str = "video-provider") -> Any: + from unittest.mock import AsyncMock, Mock + + from xagent.core.model.video.base import BaseVideoModel + + model = Mock(spec=BaseVideoModel) + model.has_ability = Mock(return_value=True) + model.abilities = ["generate"] + model.model_name = name + result = {"task_id": "t-1", "status": "succeeded", "video_url": "", "ratio": "16:9"} + if duration is not None: + result["duration"] = duration + model.generate_video = AsyncMock(return_value=result) + return model + + +@pytest.mark.asyncio +async def test_video_bills_duration_times_count() -> None: + """The provider reports one duration but generates and bills for n videos.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore(video_models={"cfg-video-id": _video_model(5)}) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p", n=3) + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 15.0 + assert entries[0]["call_type"] == "video" + + +@pytest.mark.asyncio +async def test_video_records_provider_name_and_configured_id_separately() -> None: + """`model` is the provider name, `model_id` the configured id.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore( + video_models={"cfg-video-id": _video_model(5, name="video-provider")} + ) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p") + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["model"] == "video-provider" + assert entries[0]["model_id"] == "cfg-video-id" + + +@pytest.mark.asyncio +async def test_video_without_duration_is_recorded_as_unmeasured_seconds() -> None: + """An async task reports no duration yet. The unit must stay seconds with + quantity 0 rather than degrading to another unit -- reconciling that row is + tracked in #1583.""" + from xagent.core.tools.core.video_tool import VideoGenerationToolCore + + tool = VideoGenerationToolCore(video_models={"cfg-video-id": _video_model(None)}) + + with TokenContextManager() as manager: + await tool.generate_video(prompt="p") + entries = _media_entries(manager) + + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 0.0 + + +# --- ASR records before post-processing ----------------------------------- + + +class _MalformedSegmentASR: + """Provider call succeeds and is billed; a segment lacks an end time. + + _aggregate_segments raises ValueError on this input, which the tool's + broad handler turns into success:False. The provider still ran and + charged, so the row must already be recorded by then. + """ + + model_name = "asr-provider" + + @property + def abilities(self) -> list: + return ["asr", "timestamps"] + + async def transcribe(self, audio: Any, **kwargs: Any) -> Any: + _ = (audio, kwargs) + from xagent.core.model.asr.base import ASRResult, ASRSegment + + return ASRResult( + text="hello world", + raw_response={"duration": 42.0}, + segments=[ + ASRSegment(text="hello", start=0.0, end=1.0, confidence=1.0), + ASRSegment(text="world", start=1.0, end=None, confidence=1.0), + ], + ) + + +@pytest.mark.asyncio +async def test_asr_bills_when_post_processing_raises() -> None: + """Metering must not depend on post-processing succeeding. + + Recording after _aggregate_segments meant a provider call that succeeded + and was billed went entirely unmetered whenever segment data was + malformed -- the exact bug class this module's policy exists to prevent. + """ + from xagent.core.tools.core.audio_tool import AudioToolCore + + tool = AudioToolCore(asr_models={"asr-provider": _MalformedSegmentASR()}) + + with TokenContextManager() as manager: + result = await tool.transcribe_audio( + audio_file_path="/tmp/does-not-exist.wav", verbose=False + ) + entries = _media_entries(manager) + + # The tool still reports failure to the caller... + assert result["success"] is False + # ...but the provider call happened and was billed, so it is metered. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 42.0 + assert entries[0]["call_type"] == "asr" + + +# --- verbose=True reaches the provider ------------------------------------ +# +# _RecordingASR exists to capture the flag. Without verbose=True the provider +# returns a bare string with no timings and every call meters as an +# unbillable 0 seconds, so the flag reaching the provider is what makes ASR +# billable at all. + + +@pytest.mark.asyncio +async def test_telegram_passes_verbose_and_bills_a_real_duration( + tmp_path: Path, +) -> None: + """Drive the real Telegram method, not a re-implementation of it. + + Asserting on a locally-issued transcribe() call would pass with + verbose=True deleted from bot.py, which is the only thing worth guarding + here: without it the provider returns a bare string with no timings and + every voice message meters as an unbillable 0 seconds. + """ + from xagent.web.channels.telegram.bot import TelegramBotInstance + + provider = _RecordingASR() + # __new__ rather than the real constructor: the method only touches a + # class-level timeout and a staticmethod, so a fully wired bot instance + # would add setup without adding coverage. + bot = TelegramBotInstance.__new__(TelegramBotInstance) + + audio = tmp_path / "voice.ogg" + audio.write_bytes(b"fake-ogg") + + with TokenContextManager() as manager: + transcripts = await bot._transcribe_uploaded_voice_files( + ["file-1"], + [{"telegram_file_id": "file-1", "path": str(audio), "name": "voice.ogg"}], + provider, + ) + entries = _media_entries(manager) + + # bot.py passed verbose through to the provider... + assert provider.verbose_seen is True + assert transcripts == {"file-1": "hi"} + # ...so the row carries a real duration instead of an unbillable zero. + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["quantity"] == 2.5 + # and the identity is the provider name, never the forbidden placeholder. + assert entries[0]["model"] == "asr-a" + assert entries[0]["model"] != "default" diff --git a/tests/core/tools/test_audio_tool.py b/tests/core/tools/test_audio_tool.py index 80200ec0d5..4716980338 100644 --- a/tests/core/tools/test_audio_tool.py +++ b/tests/core/tools/test_audio_tool.py @@ -949,6 +949,71 @@ async def test_synthesize_speech_json_merges_default_and_segment_options() -> No ] +async def test_synthesize_speech_json_records_media_usage_per_segment() -> None: + from xagent.core.model.chat.token_context import TokenContextManager + + tts = FakeTTS() + tool = AudioToolCore(tts_models={"fake": tts}) + + with TokenContextManager() as manager: + result = await tool.synthesize_speech_json( + json_data={ + "segments": [ + {"text": "First line"}, + {"text": "Second longer line"}, + ], + "default_voice": "voice-1", + }, + model_id="fake", + ) + usage = manager.get_usage() + + assert result["success"] is True + media_entries = [d for d in usage.details if d.get("type") == "media"] + # One TTS media entry per synthesized segment, metered by input characters. + assert usage.media_calls == 2 + assert [(d["unit"], d["quantity"], d["call_type"]) for d in media_entries] == [ + ("characters", len("First line"), "tts"), + ("characters", len("Second longer line"), "tts"), + ] + + +async def test_tts_usage_records_real_model_not_default() -> None: + """Omitting model_id is the documented common case, so the usage entry must + still name the model — billing the literal string "default" would destroy + per-model cost attribution.""" + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(tts_models={"tts-a": FakeTTS()}) + + with TokenContextManager() as manager: + await tool.synthesize_speech(text="hello") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert len(entries) == 1 + assert entries[0]["model"] == "tts-a" + assert entries[0]["model"] != "default" + + +async def test_asr_usage_meters_seconds_from_segments() -> None: + """ASR is duration-billed, so the unit is seconds and the quantity comes + from the transcribed audio's timing.""" + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(asr_models={"asr-a": FakeASR()}) + + with TokenContextManager() as manager: + result = await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert result["success"] is True + assert len(entries) == 1 + assert entries[0]["unit"] == "seconds" + assert entries[0]["call_type"] == "asr" + assert entries[0]["quantity"] > 0 + assert entries[0]["model"] == "asr-a" + + async def test_synthesize_speech_json_rejects_non_object_json() -> None: tool = AudioToolCore(tts_models={"fake": FakeTTS()}) @@ -963,3 +1028,59 @@ async def test_synthesize_speech_json_rejects_non_object_json() -> None: "failed": 0, "errors": ["JSON data must be an object"], } + + +class _NamelessASR(FakeASR): + """The Xinference default-model shape: no ``model_name`` attribute. + + Its default-getter builds a separate instance from the one in the registry + dict, so identity-by-``is`` never matches either. + """ + + +async def test_asr_default_model_is_not_billed_as_placeholder() -> None: + # Previously fell through to the literal string "default" — the exact + # placeholder the metering invariants forbid as a billing identity. + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore( + asr_models={"asr-a": _NamelessASR()}, + default_asr_model=_NamelessASR(), + ) + + with TokenContextManager() as manager: + await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert len(entries) == 1 + assert entries[0]["model"] not in {"default", "none", "null", ""} + + +async def test_asr_usage_leaves_model_id_unset() -> None: + # The tool only ever sees name-keyed registries, so writing the name into + # model_id would persist a name under an id field and still not match the + # real DB model_id /speech/transcribe records. The aggregator groups on + # `model_id or model`, so all three ASR entry points must key on the name. + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(asr_models={"asr-a": FakeASR()}) + + with TokenContextManager() as manager: + await tool.transcribe_audio(audio_file_path="x.wav") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert entries[0]["model"] == "asr-a" + assert entries[0]["model_id"] == "" + + +async def test_tts_usage_leaves_model_id_unset() -> None: + from xagent.core.model.chat.token_context import TokenContextManager + + tool = AudioToolCore(tts_models={"tts-a": FakeTTS()}) + + with TokenContextManager() as manager: + await tool.synthesize_speech(text="hello") + entries = [d for d in manager.get_usage().details if d.get("type") == "media"] + + assert entries[0]["model"] == "tts-a" + assert entries[0]["model_id"] == "" diff --git a/tests/web/test_connector_runtime_entrypoints_e2e.py b/tests/web/test_connector_runtime_entrypoints_e2e.py index 84f4cd3145..c2b85feb2d 100644 --- a/tests/web/test_connector_runtime_entrypoints_e2e.py +++ b/tests/web/test_connector_runtime_entrypoints_e2e.py @@ -1101,7 +1101,14 @@ class _FakeASR: def __init__(self) -> None: self.closed = False - async def transcribe(self, *, audio: str, format: str | None = None) -> str: + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: assert audio == "/workspace/input/voice.oga" assert format == "ogg" return "今晚有世界杯比赛吗？" diff --git a/tests/web/test_model_api.py b/tests/web/test_model_api.py index 129d9b1097..d88a505a96 100644 --- a/tests/web/test_model_api.py +++ b/tests/web/test_model_api.py @@ -1010,9 +1010,19 @@ def __init__(self, model_name: str): self.calls = [] self.closed = False - async def transcribe(self, audio, language=None, format=None): + async def transcribe( + self, audio, language=None, format=None, verbose=False, **kwargs + ): + # Mirrors the real BaseASR signature, which takes verbose and + # **kwargs. A double narrower than the interface it stands in + # for is how the missing verbose=True went unnoticed. self.calls.append( - {"audio": audio, "language": language, "format": format} + { + "audio": audio, + "language": language, + "format": format, + "verbose": verbose, + } ) return f"text from {self.model_name}" @@ -1042,8 +1052,16 @@ def create_fake_asr(db_model): assert data["text"] == "text from asr-second" assert data["model_id"] == "test-asr-second" assert len(created_models) == 1 + # verbose=True is required: without it the provider returns a bare + # string with no timings and the ASR usage record is an unbillable + # 0 seconds. assert created_models[0].calls == [ - {"audio": b"audio-bytes", "language": "en", "format": "webm"} + { + "audio": b"audio-bytes", + "language": "en", + "format": "webm", + "verbose": True, + } ] assert created_models[0].closed is True diff --git a/tests/web/test_telegram_message_queue.py b/tests/web/test_telegram_message_queue.py index 562e65c404..543a8a79e1 100644 --- a/tests/web/test_telegram_message_queue.py +++ b/tests/web/test_telegram_message_queue.py @@ -166,8 +166,15 @@ class FakeASR: def __init__(self) -> None: self.calls: list[dict[str, str | None]] = [] - async def transcribe(self, *, audio: str, format: str | None = None) -> str: - self.calls.append({"audio": audio, "format": format}) + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: + self.calls.append({"audio": audio, "format": format, "verbose": verbose}) return "今晚有世界杯比赛吗？" asr = FakeASR() @@ -189,8 +196,14 @@ async def transcribe(self, *, audio: str, format: str | None = None) -> str: ) assert transcripts == {"voice-file-id": "今晚有世界杯比赛吗？"} + # verbose=True is required: without it the provider returns a bare string + # with no timings and the ASR usage record is an unbillable 0 seconds. assert asr.calls == [ - {"audio": "/workspace/input/voice-file-id.oga", "format": "ogg"} + { + "audio": "/workspace/input/voice-file-id.oga", + "format": "ogg", + "verbose": True, + } ] assert uploaded_info[0]["file_id"] == "workspace-file-id" @@ -201,7 +214,12 @@ async def test_transcribe_uploaded_voice_files_extracts_result_text() -> None: class ResultASR: async def transcribe( - self, *, audio: str, format: str | None = None + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, ) -> SimpleNamespace: return SimpleNamespace(text="今晚有世界杯比赛吗？") @@ -242,7 +260,14 @@ async def test_transcribe_uploaded_voice_files_rejects_empty_result() -> None: bot = make_bot() class EmptyASR: - async def transcribe(self, *, audio: str, format: str | None = None) -> str: + async def transcribe( + self, + *, + audio: str, + format: str | None = None, + verbose: bool = False, + **kwargs: object, + ) -> str: return " " with pytest.raises( diff --git a/tests/web/tracking/test_standalone_usage.py b/tests/web/tracking/test_standalone_usage.py new file mode 100644 index 0000000000..dd7777dce1 --- /dev/null +++ b/tests/web/tracking/test_standalone_usage.py @@ -0,0 +1,394 @@ +"""The usage sink for work that has no TaskTracker. + +Recording usage is only half of metering: something has to bind a TokenUsage +and hand it to the quota hook. Entry points such as ``/speech/transcribe`` have +no TaskTracker, so without this their recorded usage lands in a throwaway +object. +""" + +import asyncio +from typing import Any, Optional + +import pytest + +from xagent.core.model.chat.token_context import ( + MediaCallType, + add_media_usage, + get_token_usage, +) +from xagent.web.tracking.standalone_usage import usage_scope + + +def _async_return(value: Any): + """An async callable ignoring its args and returning ``value``.""" + + async def _fn(*_a: Any, **_k: Any) -> Any: + return value + + return _fn + + +class _FakeSession: + """Stand-in for the short-lived compatibility Session ``_report`` opens. + + ``in_transaction`` is configurable and the calls are recorded in order: + with it hard-coded to False the rollback branch is unreachable, so + deleting the rollback from _report would leave every test green. + """ + + def __init__(self, in_transaction: bool = False) -> None: + self._in_transaction = in_transaction + self.rolled_back = False + self.closed = False + self.calls: list[str] = [] + + def in_transaction(self) -> bool: + return self._in_transaction + + def rollback(self) -> None: + self.rolled_back = True + self._in_transaction = False + self.calls.append("rollback") + + def close(self) -> None: + self.closed = True + self.calls.append("close") + + +@pytest.fixture +def captured(monkeypatch: pytest.MonkeyPatch): + """Capture what reaches quota_hooks.record_usage.""" + calls: list[dict[str, Any]] = [] + session = _FakeSession() + + def _record(db: Any, user_id: Any, details: list, actions: int) -> None: + calls.append( + {"db": db, "user_id": user_id, "details": details, "actions": actions} + ) + + from xagent.web.models import database + from xagent.web.services import quota_hooks + + monkeypatch.setattr(quota_hooks, "record_usage", _record) + # _report short-circuits before touching a session when no hook is + # installed, so the stock configuration would otherwise report nothing. + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + return calls, session + + +def test_usage_recorded_in_scope_reaches_the_quota_hook(captured) -> None: + calls, _ = captured + with usage_scope(42): + add_media_usage(MediaCallType.ASR, 3, model="asr-1") + + assert len(calls) == 1 + assert calls[0]["user_id"] == 42 + assert calls[0]["details"][0]["call_type"] == "asr" + # The unit is derived from the call type, never passed in. + assert calls[0]["details"][0]["unit"] == "seconds" + assert calls[0]["details"][0]["quantity"] == 3.0 + # These paths make provider calls, not agent tool calls. + assert calls[0]["actions"] == 0 + + +def test_scope_reports_even_when_the_body_raises(captured) -> None: + """A provider call that already happened is billable regardless of what + fails afterwards.""" + calls, _ = captured + with pytest.raises(RuntimeError): + with usage_scope(7): + add_media_usage(MediaCallType.TTS, 12, model="t") + raise RuntimeError("downstream failure") + + assert len(calls) == 1 + assert calls[0]["details"][0]["call_type"] == "tts" + assert calls[0]["details"][0]["unit"] == "characters" + + +def test_no_usage_means_no_hook_call(captured) -> None: + calls, _ = captured + with usage_scope(1): + pass + assert calls == [] + + +def test_compatibility_session_is_disposed(captured) -> None: + """The hook manages its own durability; the session handed to it must be + closed by us and never left holding a transaction.""" + calls, session = captured + with usage_scope(5): + add_media_usage(MediaCallType.ASR, 2, model="e") + + assert session.closed is True + assert calls[0]["db"] is session + + +def test_active_transaction_is_rolled_back_before_close( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The hook owns its own durability and must not be left holding one. + + With in_transaction() hard-coded False the rollback branch is unreachable, + so removing the rollback from _report would leave the suite green. + """ + from xagent.web.models import database + from xagent.web.services import quota_hooks + + session = _FakeSession(in_transaction=True) + monkeypatch.setattr(quota_hooks, "record_usage", lambda *a, **k: None) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is True + # Ordering matters: rolling back after close would raise on a real Session. + assert session.calls == ["rollback", "close"] + + +def test_active_transaction_is_rolled_back_when_the_hook_raises( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The finally-branch must dispose the session on the failure path too.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + session = _FakeSession(in_transaction=True) + + def _boom(*_a: Any, **_k: Any) -> None: + raise RuntimeError("quota backend down") + + monkeypatch.setattr(quota_hooks, "record_usage", _boom) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: lambda: session) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is True + assert session.calls == ["rollback", "close"] + + +def test_inactive_transaction_is_closed_without_rollback(captured) -> None: + """The common case: nothing pending, so no rollback is issued.""" + _calls, session = captured + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert session.rolled_back is False + assert session.calls == ["close"] + + +def test_scope_restores_the_previous_context(captured) -> None: + """A nested scope must not leak its usage object into the outer one.""" + _calls, _ = captured + with usage_scope(1) as outer: + add_media_usage(MediaCallType.TTS, 1, model="a") + with usage_scope(2): + add_media_usage(MediaCallType.TTS, 1, model="b") + # Back on the outer usage, and the inner call did not land here. + assert get_token_usage() is outer + assert len(outer.details) == 1 + + +def test_none_user_id_is_tolerated(captured) -> None: + """record_usage no-ops on a missing user; the scope must not raise.""" + calls, _ = captured + user_id: Optional[int] = None + with usage_scope(user_id): + add_media_usage(MediaCallType.ASR, 1, model="m") + assert calls[0]["user_id"] is None + + +def test_hook_failure_does_not_break_the_work(monkeypatch: pytest.MonkeyPatch) -> None: + """Metering must never break the operation it measures.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + def _boom(*_a: Any, **_k: Any) -> None: + raise RuntimeError("quota backend down") + + monkeypatch.setattr(quota_hooks, "record_usage", _boom) + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: True) + monkeypatch.setattr(database, "get_session_local", lambda: _FakeSession) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + # No exception escaped. + + +def test_no_hook_installed_skips_the_session_checkout( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The stock configuration has no hook, so _report must not pay for a pool + checkout per transcription.""" + from xagent.web.models import database + from xagent.web.services import quota_hooks + + checkouts: list[int] = [] + + monkeypatch.setattr(quota_hooks, "has_usage_record_hook", lambda: False) + monkeypatch.setattr( + database, "get_session_local", lambda: checkouts.append(1) or _FakeSession + ) + + with usage_scope(1): + add_media_usage(MediaCallType.ASR, 1, model="m") + + assert checkouts == [] + + +def test_transcribe_endpoint_reports_asr_usage_for_the_authenticated_owner( + captured, monkeypatch: pytest.MonkeyPatch +) -> None: + """Drive the real endpoint and observe the quota hook. + + Asserting that the literal "usage_scope(" appears in the source would pass + with the wrong user, the wrong scope boundary, or no recording at all. This + observes what actually reaches the hook: owner, quantity, unit and that it + is reported exactly once. + """ + calls, _ = captured + + from xagent.core.model.asr import adapter + from xagent.core.model.asr.base import ASRResult + from xagent.web.api import model as model_api + + class _FakeASR: + model_name = "asr-provider" + + async def transcribe(self, audio: Any, **kwargs: Any) -> ASRResult: + _ = (audio, kwargs) + # A real ASRResult, not a dict: record_asr_usage reads + # raw_response/segments off ASRResult, so a dict meters as 0s. + return ASRResult(text="hello", raw_response={"duration": 12.5}) + + class _DBModel: + # Mirrors the real DBModel: `id` is the integer PK and `model_id` the + # configured string id. The endpoint returns both model_id and + # model_name, so a fake carrying only one of them fails on the + # response build rather than on anything this test is about. + id = 1 + model_id = "configured-asr-id" + model_name = "asr-provider" + category = "speech" + + class _User: + id = 4242 + + class _Upload: + filename = "clip.wav" + content_type = "audio/wav" + + async def read(self, *_a: Any, **_k: Any) -> bytes: + return b"RIFFfake" + + async def close(self) -> None: + return None + + monkeypatch.setattr( + model_api, "_resolve_asr_model_for_transcription", lambda *a, **k: _DBModel() + ) + monkeypatch.setattr(adapter, "get_asr_model_instance", lambda *a, **k: _FakeASR()) + monkeypatch.setattr( + model_api, + "_read_transcribe_upload_with_size_limit", + _async_return(b"RIFFfake"), + ) + + response = asyncio.run( + model_api.transcribe_speech_input( + file=_Upload(), + language=None, + model_id=None, + db=object(), + user=_User(), + ) + ) + + # Assert the response too: it is built from db_model after the scope + # closes, so a fake that diverges from the real schema shows up here + # instead of as an AttributeError deep in the endpoint. + assert response["text"] == "hello" + assert response["model_id"] == "configured-asr-id" + assert response["model_name"] == "asr-provider" + + # Exactly one report, for the authenticated owner, with the ASR row. + assert len(calls) == 1, calls + assert calls[0]["user_id"] == 4242 + media = [d for d in calls[0]["details"] if d.get("type") == "media"] + assert len(media) == 1, media + assert media[0]["call_type"] == "asr" + assert media[0]["unit"] == "seconds" + assert media[0]["quantity"] == 12.5 + assert media[0]["model"] == "asr-provider" + + +def test_transcribe_falls_back_to_configured_id_for_a_placeholder_db_name( + captured, monkeypatch: pytest.MonkeyPatch +) -> None: + """model_name is nullable=False but not constrained to be meaningful. + + A row whose name is empty or literally "default" must not be billed under + that string -- the module invariants forbid placeholder identities -- so + the resolver falls back to the configured id. + """ + calls, _ = captured + + from xagent.core.model.asr import adapter + from xagent.core.model.asr.base import ASRResult + from xagent.web.api import model as model_api + + class _FakeASR: + model_name = "asr-provider" + + async def transcribe(self, audio: Any, **kwargs: Any) -> ASRResult: + _ = (audio, kwargs) + return ASRResult(text="hello", raw_response={"duration": 3.0}) + + class _PlaceholderNameDBModel: + id = 1 + model_id = "configured-asr-id" + model_name = "default" + category = "speech" + + class _User: + id = 7 + + class _Upload: + filename = "clip.wav" + content_type = "audio/wav" + + async def read(self, *_a: Any, **_k: Any) -> bytes: + return b"RIFFfake" + + async def close(self) -> None: + return None + + monkeypatch.setattr( + model_api, + "_resolve_asr_model_for_transcription", + lambda *a, **k: _PlaceholderNameDBModel(), + ) + monkeypatch.setattr(adapter, "get_asr_model_instance", lambda *a, **k: _FakeASR()) + monkeypatch.setattr( + model_api, + "_read_transcribe_upload_with_size_limit", + _async_return(b"RIFFfake"), + ) + + asyncio.run( + model_api.transcribe_speech_input( + file=_Upload(), + language=None, + model_id=None, + db=object(), + user=_User(), + ) + ) + + media = [d for d in calls[0]["details"] if d.get("type") == "media"] + assert media[0]["model"] == "configured-asr-id" + assert media[0]["model"] != "default"