|
# retoor <retoor@molodetz.nl>
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any, Callable, Optional
|
|
|
|
from ..errors import ToolInputError
|
|
from .tracker import get_tracker
|
|
|
|
|
|
class CostController:
|
|
def __init__(self, quota_provider: Optional[Callable[[], dict[str, Any]]] = None) -> None:
|
|
self._quota_provider = quota_provider
|
|
|
|
async def dispatch(self, name: str, arguments: dict[str, Any]) -> str:
|
|
if name == "usage_quota":
|
|
if self._quota_provider is None:
|
|
return json.dumps(
|
|
{"status": "unavailable", "message": "No quota information is available for this session."}
|
|
)
|
|
return json.dumps({"status": "success", **self._quota_provider()}, ensure_ascii=False)
|
|
if name == "cost_stats":
|
|
tracker = get_tracker()
|
|
if tracker is None:
|
|
return json.dumps(
|
|
{"status": "unavailable", "message": "No cost tracker is active for this session."}
|
|
)
|
|
return json.dumps({"status": "success", **tracker.stats()}, ensure_ascii=False)
|
|
raise ToolInputError(f"Unknown cost tool: {name}")
|