mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-10-11 03:40:03 +00:00
67 lines
2.4 KiB
Python
67 lines
2.4 KiB
Python
"""Read-only evaluation helpers for application job execution statistics."""
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from ..components.job import BackgroundJob, BaseJob, CronJob, StreamJob
|
|
from .counter import global_counter_get
|
|
|
|
if TYPE_CHECKING:
|
|
from ..components.application_context import ApplicationContext
|
|
|
|
|
|
_JOB_ENTRY_CLASSES = {CronJob, StreamJob, BackgroundJob, BaseJob}
|
|
|
|
|
|
def check_job_count(job_name: str, app_context: "ApplicationContext") -> int:
|
|
"""Return the application-lifetime execution count for a registered job.
|
|
|
|
``app_context`` scopes the lookup because ReMe does not maintain a global
|
|
current Application instance. Unknown job names use the same ``KeyError``
|
|
contract as :meth:`Application.run_job`.
|
|
"""
|
|
job = app_context.jobs.get(job_name)
|
|
if job is None:
|
|
raise KeyError(f"Job '{job_name}' not found")
|
|
|
|
entry_class = next((cls for cls in type(job).mro() if cls in _JOB_ENTRY_CLASSES), None)
|
|
if entry_class is None:
|
|
raise TypeError(f"Job '{job_name}' does not inherit from a supported job implementation")
|
|
# pylint: disable-next=protected-access
|
|
return global_counter_get(app_context.metadata, job._counter_key(entry_class))
|
|
|
|
|
|
class JobCountTracker:
|
|
"""Measure registered job calls made while this context is active."""
|
|
|
|
def __init__(self, job_names: list[str], app_context: "ApplicationContext") -> None:
|
|
self.job_names = list(dict.fromkeys(job_names))
|
|
self.app_context = app_context
|
|
self._start_counts: dict[str, int] = {}
|
|
self.counts: dict[str, int] = {}
|
|
|
|
def __enter__(self) -> dict[str, int]:
|
|
self._start_counts = {name: check_job_count(name, self.app_context) for name in self.job_names}
|
|
return self.counts
|
|
|
|
def __exit__(self, exc_type, exc_value, traceback) -> bool:
|
|
self.counts.update(
|
|
{
|
|
name: check_job_count(name, self.app_context) - start_count
|
|
for name, start_count in self._start_counts.items()
|
|
},
|
|
)
|
|
return False
|
|
|
|
|
|
def track_job_counts(job_names: list[str], app_context: "ApplicationContext") -> JobCountTracker:
|
|
"""Return a context manager that reports call deltas for ``job_names``.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
with track_job_counts(["search"], app.context) as counts:
|
|
await app.run_job("agentic_answer", query="...")
|
|
assert counts == {"search": 2}
|
|
"""
|
|
return JobCountTracker(job_names, app_context)
|