Spaces:
Sleeping
Sleeping
File size: 1,653 Bytes
5db30ac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | 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
@unittest.skipIf(GenericFakeChatModel is None, "langchain-core is not installed")
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()
|