Spaces:
Sleeping
Sleeping
| import sys | |
| import tempfile | |
| import unittest | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| SRC = ROOT / "src" | |
| if str(SRC) not in sys.path: | |
| sys.path.insert(0, str(SRC)) | |
| try: | |
| from langchain_core.language_models.fake_chat_models import GenericFakeChatModel | |
| from langchain_core.messages import AIMessage | |
| except ImportError: # the basic-education extra is optional | |
| GenericFakeChatModel = None | |
| from elbench.utils.spend import BudgetExceeded, SpendLimits, spend_limits | |
| class RuntimeSpendCallbackTest(unittest.TestCase): | |
| def _model(self): | |
| from elbench.basic_education_runtime.spend_callback import SpendCallback | |
| replies = iter([AIMessage(content="hi", usage_metadata={"input_tokens": 1000, "output_tokens": 100, "total_tokens": 1100})] * 3) | |
| model = GenericFakeChatModel(messages=replies) | |
| model.callbacks = [SpendCallback("student")] | |
| return model | |
| def test_runtime_calls_are_recorded_and_stopped_at_the_role_limit(self): | |
| with tempfile.TemporaryDirectory() as tmpdir: | |
| limits = SpendLimits(per_role_usd={"student": 0.002}, prices={"student": (1e-6, 2e-6)}) | |
| with spend_limits(limits, Path(tmpdir) / "spend.jsonl") as ledger: | |
| model = self._model() | |
| model.invoke("hello") | |
| self.assertAlmostEqual(ledger.spent()["student"], 0.0012) | |
| model.invoke("hello") | |
| with self.assertRaises(BudgetExceeded): | |
| model.invoke("hello") | |
| if __name__ == "__main__": | |
| unittest.main() | |