From d465389e9e1c4d1c51f589d9e485a66a4fb2d012 Mon Sep 17 00:00:00 2001 From: HAL9000 Date: Sat, 18 Apr 2026 19:03:12 +0000 Subject: [PATCH 1/3] feat(budget): implement CostTracker service for LLM spending tracking --- features/cost_tracker_service.feature | 93 +++++++ features/steps/cost_tracker_steps.py | 162 +++++++++++++ features/steps/lsp_actor_service_steps.py | 3 + .../infrastructure/database/cost_tracker.py | 228 ++++++++++++++++++ 4 files changed, 486 insertions(+) create mode 100644 features/cost_tracker_service.feature create mode 100644 features/steps/cost_tracker_steps.py create mode 100644 features/steps/lsp_actor_service_steps.py create mode 100644 src/cleveragents/infrastructure/database/cost_tracker.py diff --git a/features/cost_tracker_service.feature b/features/cost_tracker_service.feature new file mode 100644 index 000000000..b15d7eca7 --- /dev/null +++ b/features/cost_tracker_service.feature @@ -0,0 +1,93 @@ +Feature: CostTracker service for per-session and per-plan spending tracking + As a system administrator + I want to track LLM API spending per session and per plan + So that I can monitor costs and enforce budget limits + + Background: + Given a temporary data directory for cost tracking + And a CostTracker instance with default model pricing + + Scenario: Record usage for a single LLM call + When I record usage for session "sess-001" and plan "plan-001" + | tokens_in | tokens_out | model | + | 100 | 50 | gpt-3.5-turbo | + Then the recorded cost should be approximately 0.000125 USD + And the cost record should have session_id "sess-001" + And the cost record should have plan_id "plan-001" + + Scenario: Calculate cost for different models + When I record usage for session "sess-002" and plan "plan-002" + | tokens_in | tokens_out | model | + | 1000 | 1000 | gpt-4 | + Then the recorded cost should be approximately 0.09 USD + When I record usage for session "sess-002" and plan "plan-002" + | tokens_in | tokens_out | model | + | 1000 | 1000 | gpt-3.5-turbo | + Then the recorded cost should be approximately 0.002 USD + + Scenario: Get total session cost + When I record multiple usages for session "sess-003" + | plan_id | tokens_in | tokens_out | model | + | plan-001 | 100 | 50 | gpt-3.5-turbo | + | plan-002 | 200 | 100 | gpt-3.5-turbo | + | plan-003 | 300 | 150 | gpt-3.5-turbo | + Then the total session cost for "sess-003" should be approximately 0.00035 USD + + Scenario: Get total plan cost + When I record multiple usages for plan "plan-004" + | session_id | tokens_in | tokens_out | model | + | sess-001 | 100 | 50 | gpt-3.5-turbo | + | sess-002 | 200 | 100 | gpt-3.5-turbo | + | sess-003 | 300 | 150 | gpt-3.5-turbo | + Then the total plan cost for "plan-004" should be approximately 0.00035 USD + + Scenario: Cost persistence across tracker instances + When I record usage for session "sess-004" and plan "plan-005" + | tokens_in | tokens_out | model | + | 500 | 250 | gpt-3.5-turbo | + And I create a new CostTracker instance with the same data directory + Then the total session cost for "sess-004" should be approximately 0.000625 USD + And the total plan cost for "plan-005" should be approximately 0.000625 USD + + Scenario: Get session entries + When I record multiple usages for session "sess-005" + | plan_id | tokens_in | tokens_out | model | + | plan-001 | 100 | 50 | gpt-3.5-turbo | + | plan-002 | 200 | 100 | gpt-3.5-turbo | + Then I should get 2 entries for session "sess-005" + And each entry should have the correct session_id and plan_id + + Scenario: Get plan entries + When I record multiple usages for plan "plan-006" + | session_id | tokens_in | tokens_out | model | + | sess-001 | 100 | 50 | gpt-3.5-turbo | + | sess-002 | 200 | 100 | gpt-3.5-turbo | + Then I should get 2 entries for plan "plan-006" + And each entry should have the correct session_id and plan_id + + Scenario: Custom model pricing + Given a CostTracker instance with custom model pricing + | model | input_price | output_price | + | custom-llm | 0.001 | 0.002 | + When I record usage for session "sess-006" and plan "plan-007" + | tokens_in | tokens_out | model | + | 1000 | 1000 | custom-llm | + Then the recorded cost should be approximately 0.003 USD + + Scenario: Zero cost for unknown model + When I record usage for session "sess-007" and plan "plan-008" + | tokens_in | tokens_out | model | + | 1000 | 1000 | unknown-llm | + Then the recorded cost should be approximately 0.0 USD + + Scenario: Multiple sessions and plans isolation + When I record usage for session "sess-008" and plan "plan-009" + | tokens_in | tokens_out | model | + | 100 | 50 | gpt-3.5-turbo | + And I record usage for session "sess-009" and plan "plan-010" + | tokens_in | tokens_out | model | + | 200 | 100 | gpt-3.5-turbo | + Then the total session cost for "sess-008" should be approximately 0.000125 USD + And the total session cost for "sess-009" should be approximately 0.00025 USD + And the total plan cost for "plan-009" should be approximately 0.000125 USD + And the total plan cost for "plan-010" should be approximately 0.00025 USD diff --git a/features/steps/cost_tracker_steps.py b/features/steps/cost_tracker_steps.py new file mode 100644 index 000000000..f63c153ea --- /dev/null +++ b/features/steps/cost_tracker_steps.py @@ -0,0 +1,162 @@ +"""Step definitions for CostTracker service tests.""" + +from __future__ import annotations + +import tempfile +from pathlib import Path + +from behave import given, then, when + +from cleveragents.infrastructure.database.cost_tracker import CostRecord, CostTracker + + +@given("a temporary data directory for cost tracking") +def step_create_temp_dir(context): + """Create a temporary directory for cost tracking.""" + context.temp_dir = tempfile.mkdtemp() + context.data_dir = Path(context.temp_dir) + + +@given("a CostTracker instance with default model pricing") +def step_create_cost_tracker(context): + """Create a CostTracker instance with default pricing.""" + context.cost_tracker = CostTracker(context.data_dir) + context.recorded_records = [] + + +@when('I record usage for session "{session_id}" and plan "{plan_id}"') +def step_record_usage(context, session_id, plan_id): + """Record usage for a session and plan.""" + for row in context.table: + tokens_in = int(row["tokens_in"]) + tokens_out = int(row["tokens_out"]) + model = row["model"] + + record = context.cost_tracker.record_usage( + session_id=session_id, + plan_id=plan_id, + tokens_in=tokens_in, + tokens_out=tokens_out, + model=model, + ) + context.recorded_records.append(record) + + +@then("the recorded cost should be approximately {cost} USD") +def step_check_recorded_cost(context, cost): + """Check the recorded cost.""" + expected_cost = float(cost) + actual_cost = context.recorded_records[-1].cost_usd + # Allow for small floating point differences + assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + + +@then('the cost record should have session_id "{session_id}"') +def step_check_session_id(context, session_id): + """Check the session_id of the recorded cost.""" + assert context.recorded_records[-1].session_id == session_id + + +@then('the cost record should have plan_id "{plan_id}"') +def step_check_plan_id(context, plan_id): + """Check the plan_id of the recorded cost.""" + assert context.recorded_records[-1].plan_id == plan_id + + +@when('I record multiple usages for session "{session_id}"') +def step_record_multiple_usages_session(context, session_id): + """Record multiple usages for a session.""" + for row in context.table: + plan_id = row["plan_id"] + tokens_in = int(row["tokens_in"]) + tokens_out = int(row["tokens_out"]) + model = row["model"] + + record = context.cost_tracker.record_usage( + session_id=session_id, + plan_id=plan_id, + tokens_in=tokens_in, + tokens_out=tokens_out, + model=model, + ) + context.recorded_records.append(record) + + +@then('the total session cost for "{session_id}" should be approximately {cost} USD') +def step_check_session_cost(context, session_id, cost): + """Check the total cost for a session.""" + expected_cost = float(cost) + actual_cost = context.cost_tracker.get_session_cost(session_id) + assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + + +@when('I record multiple usages for plan "{plan_id}"') +def step_record_multiple_usages_plan(context, plan_id): + """Record multiple usages for a plan.""" + for row in context.table: + session_id = row["session_id"] + tokens_in = int(row["tokens_in"]) + tokens_out = int(row["tokens_out"]) + model = row["model"] + + record = context.cost_tracker.record_usage( + session_id=session_id, + plan_id=plan_id, + tokens_in=tokens_in, + tokens_out=tokens_out, + model=model, + ) + context.recorded_records.append(record) + + +@then('the total plan cost for "{plan_id}" should be approximately {cost} USD') +def step_check_plan_cost(context, plan_id, cost): + """Check the total cost for a plan.""" + expected_cost = float(cost) + actual_cost = context.cost_tracker.get_plan_cost(plan_id) + assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + + +@when("I create a new CostTracker instance with the same data directory") +def step_create_new_tracker(context): + """Create a new CostTracker instance with the same data directory.""" + context.cost_tracker = CostTracker(context.data_dir) + + +@given("a CostTracker instance with custom model pricing") +def step_create_tracker_custom_pricing(context): + """Create a CostTracker instance with custom pricing.""" + pricing = {} + for row in context.table: + model = row["model"] + input_price = float(row["input_price"]) + output_price = float(row["output_price"]) + pricing[model] = {"input": input_price, "output": output_price} + + context.cost_tracker = CostTracker(context.data_dir, model_pricing=pricing) + context.recorded_records = [] + + +@then('I should get {count} entries for session "{session_id}"') +def step_check_session_entries_count(context, count, session_id): + """Check the number of entries for a session.""" + entries = context.cost_tracker.get_session_entries(session_id) + assert len(entries) == int(count), f"Expected {count} entries, got {len(entries)}" + context.session_entries = entries + + +@then("each entry should have the correct session_id and plan_id") +def step_check_entries_ids(context): + """Check that each entry has the correct IDs.""" + for entry in context.session_entries: + assert isinstance(entry, CostRecord) + assert entry.session_id is not None + assert entry.plan_id is not None + + +@then('I should get {count} entries for plan "{plan_id}"') +def step_check_plan_entries_count(context, count, plan_id): + """Check the number of entries for a plan.""" + entries = context.cost_tracker.get_plan_entries(plan_id) + assert len(entries) == int(count), f"Expected {count} entries, got {len(entries)}" + context.plan_entries = entries diff --git a/features/steps/lsp_actor_service_steps.py b/features/steps/lsp_actor_service_steps.py new file mode 100644 index 000000000..6aeb9e11e --- /dev/null +++ b/features/steps/lsp_actor_service_steps.py @@ -0,0 +1,3 @@ +"""Stub for LSP actor service steps.""" + +# This is a placeholder file to prevent import errors diff --git a/src/cleveragents/infrastructure/database/cost_tracker.py b/src/cleveragents/infrastructure/database/cost_tracker.py new file mode 100644 index 000000000..12f38643f --- /dev/null +++ b/src/cleveragents/infrastructure/database/cost_tracker.py @@ -0,0 +1,228 @@ +"""Cost tracking service for per-session and per-plan LLM spending.""" + +from __future__ import annotations + +from dataclasses import dataclass +from datetime import datetime +from pathlib import Path + +from sqlalchemy import Column, Float, Integer, String, create_engine +from sqlalchemy.orm import declarative_base, sessionmaker + +Base = declarative_base() + + +@dataclass +class CostRecord: + """Record of LLM API usage and cost.""" + + session_id: str + plan_id: str + tokens_in: int + tokens_out: int + model: str + cost_usd: float + timestamp: datetime + + +class CostEntry(Base): + """SQLAlchemy model for cost tracking.""" + + __tablename__ = "cost_entries" + + id = Column(Integer, primary_key=True) + session_id = Column(String, nullable=False, index=True) + plan_id = Column(String, nullable=False, index=True) + tokens_in = Column(Integer, nullable=False) + tokens_out = Column(Integer, nullable=False) + model = Column(String, nullable=False) + cost_usd = Column(Float, nullable=False) + timestamp = Column(String, nullable=False) + + +class CostTracker: + """Service for tracking LLM API spending per session and per plan.""" + + def __init__(self, data_dir: Path, model_pricing: dict[str, dict[str, float]] | None = None) -> None: + """Initialize the CostTracker. + + Args: + data_dir: Directory where costs.db will be stored + model_pricing: Dict mapping model names to pricing dicts with 'input' and 'output' keys + (cost per 1K tokens) + """ + self.data_dir = Path(data_dir) + self.data_dir.mkdir(parents=True, exist_ok=True) + self.db_path = self.data_dir / "costs.db" + + # Default pricing for common models (cost per 1K tokens) + self.model_pricing = model_pricing or { + "gpt-4": {"input": 0.03, "output": 0.06}, + "gpt-4-turbo": {"input": 0.01, "output": 0.03}, + "gpt-3.5-turbo": {"input": 0.0005, "output": 0.0015}, + "claude-3-opus": {"input": 0.015, "output": 0.075}, + "claude-3-sonnet": {"input": 0.003, "output": 0.015}, + "claude-3-haiku": {"input": 0.00025, "output": 0.00125}, + } + + # Initialize database + self._init_db() + + def _init_db(self) -> None: + """Initialize the SQLite database.""" + engine = create_engine(f"sqlite:///{self.db_path}") + Base.metadata.create_all(engine) + self.engine = engine + self.SessionLocal = sessionmaker(bind=engine) + + def _calculate_cost(self, tokens_in: int, tokens_out: int, model: str) -> float: + """Calculate cost for given tokens and model. + + Args: + tokens_in: Number of input tokens + tokens_out: Number of output tokens + model: Model name + + Returns: + Cost in USD + """ + pricing = self.model_pricing.get(model, {"input": 0.0, "output": 0.0}) + input_cost = (tokens_in / 1000) * pricing.get("input", 0.0) + output_cost = (tokens_out / 1000) * pricing.get("output", 0.0) + return round(input_cost + output_cost, 6) + + def record_usage( + self, + session_id: str, + plan_id: str, + tokens_in: int, + tokens_out: int, + model: str, + ) -> CostRecord: + """Record LLM API usage. + + Args: + session_id: Session identifier + plan_id: Plan identifier + tokens_in: Number of input tokens + tokens_out: Number of output tokens + model: Model name + + Returns: + CostRecord with calculated cost + """ + cost_usd = self._calculate_cost(tokens_in, tokens_out, model) + timestamp = datetime.utcnow() + + # Store in database + session = self.SessionLocal() + try: + entry = CostEntry( + session_id=session_id, + plan_id=plan_id, + tokens_in=tokens_in, + tokens_out=tokens_out, + model=model, + cost_usd=cost_usd, + timestamp=timestamp.isoformat(), + ) + session.add(entry) + session.commit() + finally: + session.close() + + return CostRecord( + session_id=session_id, + plan_id=plan_id, + tokens_in=tokens_in, + tokens_out=tokens_out, + model=model, + cost_usd=cost_usd, + timestamp=timestamp, + ) + + def get_session_cost(self, session_id: str) -> float: + """Get total cost for a session. + + Args: + session_id: Session identifier + + Returns: + Total cost in USD for the session + """ + session = self.SessionLocal() + try: + total = session.query(CostEntry).filter(CostEntry.session_id == session_id).all() + return round(sum(entry.cost_usd for entry in total), 6) + finally: + session.close() + + def get_plan_cost(self, plan_id: str) -> float: + """Get total cost for a plan. + + Args: + plan_id: Plan identifier + + Returns: + Total cost in USD for the plan + """ + session = self.SessionLocal() + try: + total = session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() + return round(sum(entry.cost_usd for entry in total), 6) + finally: + session.close() + + def get_session_entries(self, session_id: str) -> list[CostRecord]: + """Get all cost entries for a session. + + Args: + session_id: Session identifier + + Returns: + List of CostRecord objects for the session + """ + session = self.SessionLocal() + try: + entries = session.query(CostEntry).filter(CostEntry.session_id == session_id).all() + return [ + CostRecord( + session_id=entry.session_id, + plan_id=entry.plan_id, + tokens_in=entry.tokens_in, + tokens_out=entry.tokens_out, + model=entry.model, + cost_usd=entry.cost_usd, + timestamp=datetime.fromisoformat(entry.timestamp), + ) + for entry in entries + ] + finally: + session.close() + + def get_plan_entries(self, plan_id: str) -> list[CostRecord]: + """Get all cost entries for a plan. + + Args: + plan_id: Plan identifier + + Returns: + List of CostRecord objects for the plan + """ + session = self.SessionLocal() + try: + entries = session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() + return [ + CostRecord( + session_id=entry.session_id, + plan_id=entry.plan_id, + tokens_in=entry.tokens_in, + tokens_out=entry.tokens_out, + model=entry.model, + cost_usd=entry.cost_usd, + timestamp=datetime.fromisoformat(entry.timestamp), + ) + for entry in entries + ] + finally: + session.close() -- 2.52.0 From 6b9e422814c16425068da0c2b0e9f06fb66f551f Mon Sep 17 00:00:00 2001 From: HAL9000 Date: Thu, 23 Apr 2026 11:56:18 +0000 Subject: [PATCH 2/3] fix(budget): resolve lint and typecheck failures in CostTracker service - Add __allow_unmapped__ = True and type: ignore[misc] to CostEntry model - Import typing.cast and use cast() for all SQLAlchemy column attribute accesses - Wrap long __init__ signature to stay within 88-char line limit - Wrap long docstring lines to stay within 88-char line limit - Wrap long query chains to stay within 88-char line limit --- .../infrastructure/database/cost_tracker.py | 74 +++++++++++++------ 1 file changed, 50 insertions(+), 24 deletions(-) diff --git a/src/cleveragents/infrastructure/database/cost_tracker.py b/src/cleveragents/infrastructure/database/cost_tracker.py index 12f38643f..4e2b2c873 100644 --- a/src/cleveragents/infrastructure/database/cost_tracker.py +++ b/src/cleveragents/infrastructure/database/cost_tracker.py @@ -5,6 +5,7 @@ from __future__ import annotations from dataclasses import dataclass from datetime import datetime from pathlib import Path +from typing import cast from sqlalchemy import Column, Float, Integer, String, create_engine from sqlalchemy.orm import declarative_base, sessionmaker @@ -25,9 +26,10 @@ class CostRecord: timestamp: datetime -class CostEntry(Base): +class CostEntry(Base): # type: ignore[misc] """SQLAlchemy model for cost tracking.""" + __allow_unmapped__ = True __tablename__ = "cost_entries" id = Column(Integer, primary_key=True) @@ -43,13 +45,17 @@ class CostEntry(Base): class CostTracker: """Service for tracking LLM API spending per session and per plan.""" - def __init__(self, data_dir: Path, model_pricing: dict[str, dict[str, float]] | None = None) -> None: + def __init__( + self, + data_dir: Path, + model_pricing: dict[str, dict[str, float]] | None = None, + ) -> None: """Initialize the CostTracker. Args: data_dir: Directory where costs.db will be stored - model_pricing: Dict mapping model names to pricing dicts with 'input' and 'output' keys - (cost per 1K tokens) + model_pricing: Dict mapping model names to pricing dicts with + 'input' and 'output' keys (cost per 1K tokens) """ self.data_dir = Path(data_dir) self.data_dir.mkdir(parents=True, exist_ok=True) @@ -152,8 +158,14 @@ class CostTracker: """ session = self.SessionLocal() try: - total = session.query(CostEntry).filter(CostEntry.session_id == session_id).all() - return round(sum(entry.cost_usd for entry in total), 6) + total = ( + session.query(CostEntry) + .filter(CostEntry.session_id == session_id) + .all() + ) + return round( + sum(cast(float, entry.cost_usd) for entry in total), 6 + ) finally: session.close() @@ -168,8 +180,14 @@ class CostTracker: """ session = self.SessionLocal() try: - total = session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() - return round(sum(entry.cost_usd for entry in total), 6) + total = ( + session.query(CostEntry) + .filter(CostEntry.plan_id == plan_id) + .all() + ) + return round( + sum(cast(float, entry.cost_usd) for entry in total), 6 + ) finally: session.close() @@ -184,16 +202,20 @@ class CostTracker: """ session = self.SessionLocal() try: - entries = session.query(CostEntry).filter(CostEntry.session_id == session_id).all() + entries = ( + session.query(CostEntry) + .filter(CostEntry.session_id == session_id) + .all() + ) return [ CostRecord( - session_id=entry.session_id, - plan_id=entry.plan_id, - tokens_in=entry.tokens_in, - tokens_out=entry.tokens_out, - model=entry.model, - cost_usd=entry.cost_usd, - timestamp=datetime.fromisoformat(entry.timestamp), + session_id=cast(str, entry.session_id), + plan_id=cast(str, entry.plan_id), + tokens_in=cast(int, entry.tokens_in), + tokens_out=cast(int, entry.tokens_out), + model=cast(str, entry.model), + cost_usd=cast(float, entry.cost_usd), + timestamp=datetime.fromisoformat(cast(str, entry.timestamp)), ) for entry in entries ] @@ -211,16 +233,20 @@ class CostTracker: """ session = self.SessionLocal() try: - entries = session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() + entries = ( + session.query(CostEntry) + .filter(CostEntry.plan_id == plan_id) + .all() + ) return [ CostRecord( - session_id=entry.session_id, - plan_id=entry.plan_id, - tokens_in=entry.tokens_in, - tokens_out=entry.tokens_out, - model=entry.model, - cost_usd=entry.cost_usd, - timestamp=datetime.fromisoformat(entry.timestamp), + session_id=cast(str, entry.session_id), + plan_id=cast(str, entry.plan_id), + tokens_in=cast(int, entry.tokens_in), + tokens_out=cast(int, entry.tokens_out), + model=cast(str, entry.model), + cost_usd=cast(float, entry.cost_usd), + timestamp=datetime.fromisoformat(cast(str, entry.timestamp)), ) for entry in entries ] -- 2.52.0 From aee4e2585316e646e27aed12311a30e97f3f11d6 Mon Sep 17 00:00:00 2001 From: CleverThis Date: Thu, 4 Jun 2026 08:37:16 -0400 Subject: [PATCH 3/3] fix(budget): fix type safety, formatting, and test correctness in CostTracker MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Convert CostRecord from @dataclass to Pydantic BaseModel (architecture policy) - Replace declarative_base() + # type: ignore[misc] with DeclarativeBase subclass - Use Mapped/mapped_column for proper SQLAlchemy 2.0 typed columns - Remove all cast() calls — types now flow from Mapped[T] annotations - Change timestamp column from String to DateTime; use datetime.now(UTC) - Fix wrong expected totals in cost_tracker_service.feature (0.00035 → 0.00075) - Fix step_check_entries_ids AttributeError: read plan_entries when session_entries absent - Apply ruff format to both changed files ISSUES CLOSED: #5248 --- features/cost_tracker_service.feature | 4 +- features/steps/cost_tracker_steps.py | 18 +++- .../infrastructure/database/cost_tracker.py | 94 ++++++++----------- 3 files changed, 56 insertions(+), 60 deletions(-) diff --git a/features/cost_tracker_service.feature b/features/cost_tracker_service.feature index b15d7eca7..0d5604c5e 100644 --- a/features/cost_tracker_service.feature +++ b/features/cost_tracker_service.feature @@ -31,7 +31,7 @@ Feature: CostTracker service for per-session and per-plan spending tracking | plan-001 | 100 | 50 | gpt-3.5-turbo | | plan-002 | 200 | 100 | gpt-3.5-turbo | | plan-003 | 300 | 150 | gpt-3.5-turbo | - Then the total session cost for "sess-003" should be approximately 0.00035 USD + Then the total session cost for "sess-003" should be approximately 0.00075 USD Scenario: Get total plan cost When I record multiple usages for plan "plan-004" @@ -39,7 +39,7 @@ Feature: CostTracker service for per-session and per-plan spending tracking | sess-001 | 100 | 50 | gpt-3.5-turbo | | sess-002 | 200 | 100 | gpt-3.5-turbo | | sess-003 | 300 | 150 | gpt-3.5-turbo | - Then the total plan cost for "plan-004" should be approximately 0.00035 USD + Then the total plan cost for "plan-004" should be approximately 0.00075 USD Scenario: Cost persistence across tracker instances When I record usage for session "sess-004" and plan "plan-005" diff --git a/features/steps/cost_tracker_steps.py b/features/steps/cost_tracker_steps.py index f63c153ea..d02a5339e 100644 --- a/features/steps/cost_tracker_steps.py +++ b/features/steps/cost_tracker_steps.py @@ -47,8 +47,9 @@ def step_check_recorded_cost(context, cost): """Check the recorded cost.""" expected_cost = float(cost) actual_cost = context.recorded_records[-1].cost_usd - # Allow for small floating point differences - assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + assert abs(actual_cost - expected_cost) < 0.000001, ( + f"Expected {expected_cost}, got {actual_cost}" + ) @then('the cost record should have session_id "{session_id}"') @@ -87,7 +88,9 @@ def step_check_session_cost(context, session_id, cost): """Check the total cost for a session.""" expected_cost = float(cost) actual_cost = context.cost_tracker.get_session_cost(session_id) - assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + assert abs(actual_cost - expected_cost) < 0.000001, ( + f"Expected {expected_cost}, got {actual_cost}" + ) @when('I record multiple usages for plan "{plan_id}"') @@ -114,7 +117,9 @@ def step_check_plan_cost(context, plan_id, cost): """Check the total cost for a plan.""" expected_cost = float(cost) actual_cost = context.cost_tracker.get_plan_cost(plan_id) - assert abs(actual_cost - expected_cost) < 0.000001, f"Expected {expected_cost}, got {actual_cost}" + assert abs(actual_cost - expected_cost) < 0.000001, ( + f"Expected {expected_cost}, got {actual_cost}" + ) @when("I create a new CostTracker instance with the same data directory") @@ -148,7 +153,10 @@ def step_check_session_entries_count(context, count, session_id): @then("each entry should have the correct session_id and plan_id") def step_check_entries_ids(context): """Check that each entry has the correct IDs.""" - for entry in context.session_entries: + entries = getattr(context, "session_entries", None) or getattr( + context, "plan_entries", [] + ) + for entry in entries: assert isinstance(entry, CostRecord) assert entry.session_id is not None assert entry.plan_id is not None diff --git a/src/cleveragents/infrastructure/database/cost_tracker.py b/src/cleveragents/infrastructure/database/cost_tracker.py index 4e2b2c873..c358b8107 100644 --- a/src/cleveragents/infrastructure/database/cost_tracker.py +++ b/src/cleveragents/infrastructure/database/cost_tracker.py @@ -2,19 +2,19 @@ from __future__ import annotations -from dataclasses import dataclass -from datetime import datetime +from datetime import UTC, datetime from pathlib import Path -from typing import cast -from sqlalchemy import Column, Float, Integer, String, create_engine -from sqlalchemy.orm import declarative_base, sessionmaker - -Base = declarative_base() +from pydantic import BaseModel +from sqlalchemy import DateTime, Float, Integer, String, create_engine +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker -@dataclass -class CostRecord: +class _Base(DeclarativeBase): + pass + + +class CostRecord(BaseModel): """Record of LLM API usage and cost.""" session_id: str @@ -26,20 +26,19 @@ class CostRecord: timestamp: datetime -class CostEntry(Base): # type: ignore[misc] +class CostEntry(_Base): """SQLAlchemy model for cost tracking.""" - __allow_unmapped__ = True __tablename__ = "cost_entries" - id = Column(Integer, primary_key=True) - session_id = Column(String, nullable=False, index=True) - plan_id = Column(String, nullable=False, index=True) - tokens_in = Column(Integer, nullable=False) - tokens_out = Column(Integer, nullable=False) - model = Column(String, nullable=False) - cost_usd = Column(Float, nullable=False) - timestamp = Column(String, nullable=False) + id: Mapped[int] = mapped_column(Integer, primary_key=True) + session_id: Mapped[str] = mapped_column(String, nullable=False, index=True) + plan_id: Mapped[str] = mapped_column(String, nullable=False, index=True) + tokens_in: Mapped[int] = mapped_column(Integer, nullable=False) + tokens_out: Mapped[int] = mapped_column(Integer, nullable=False) + model: Mapped[str] = mapped_column(String, nullable=False) + cost_usd: Mapped[float] = mapped_column(Float, nullable=False) + timestamp: Mapped[datetime] = mapped_column(DateTime, nullable=False) class CostTracker: @@ -61,7 +60,6 @@ class CostTracker: self.data_dir.mkdir(parents=True, exist_ok=True) self.db_path = self.data_dir / "costs.db" - # Default pricing for common models (cost per 1K tokens) self.model_pricing = model_pricing or { "gpt-4": {"input": 0.03, "output": 0.06}, "gpt-4-turbo": {"input": 0.01, "output": 0.03}, @@ -71,13 +69,12 @@ class CostTracker: "claude-3-haiku": {"input": 0.00025, "output": 0.00125}, } - # Initialize database self._init_db() def _init_db(self) -> None: """Initialize the SQLite database.""" engine = create_engine(f"sqlite:///{self.db_path}") - Base.metadata.create_all(engine) + _Base.metadata.create_all(engine) self.engine = engine self.SessionLocal = sessionmaker(bind=engine) @@ -118,9 +115,8 @@ class CostTracker: CostRecord with calculated cost """ cost_usd = self._calculate_cost(tokens_in, tokens_out, model) - timestamp = datetime.utcnow() + timestamp = datetime.now(UTC).replace(tzinfo=None) - # Store in database session = self.SessionLocal() try: entry = CostEntry( @@ -130,7 +126,7 @@ class CostTracker: tokens_out=tokens_out, model=model, cost_usd=cost_usd, - timestamp=timestamp.isoformat(), + timestamp=timestamp, ) session.add(entry) session.commit() @@ -158,14 +154,12 @@ class CostTracker: """ session = self.SessionLocal() try: - total = ( + entries = ( session.query(CostEntry) .filter(CostEntry.session_id == session_id) .all() ) - return round( - sum(cast(float, entry.cost_usd) for entry in total), 6 - ) + return round(sum(entry.cost_usd for entry in entries), 6) finally: session.close() @@ -180,14 +174,10 @@ class CostTracker: """ session = self.SessionLocal() try: - total = ( - session.query(CostEntry) - .filter(CostEntry.plan_id == plan_id) - .all() - ) - return round( - sum(cast(float, entry.cost_usd) for entry in total), 6 + entries = ( + session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() ) + return round(sum(entry.cost_usd for entry in entries), 6) finally: session.close() @@ -209,13 +199,13 @@ class CostTracker: ) return [ CostRecord( - session_id=cast(str, entry.session_id), - plan_id=cast(str, entry.plan_id), - tokens_in=cast(int, entry.tokens_in), - tokens_out=cast(int, entry.tokens_out), - model=cast(str, entry.model), - cost_usd=cast(float, entry.cost_usd), - timestamp=datetime.fromisoformat(cast(str, entry.timestamp)), + session_id=entry.session_id, + plan_id=entry.plan_id, + tokens_in=entry.tokens_in, + tokens_out=entry.tokens_out, + model=entry.model, + cost_usd=entry.cost_usd, + timestamp=entry.timestamp, ) for entry in entries ] @@ -234,19 +224,17 @@ class CostTracker: session = self.SessionLocal() try: entries = ( - session.query(CostEntry) - .filter(CostEntry.plan_id == plan_id) - .all() + session.query(CostEntry).filter(CostEntry.plan_id == plan_id).all() ) return [ CostRecord( - session_id=cast(str, entry.session_id), - plan_id=cast(str, entry.plan_id), - tokens_in=cast(int, entry.tokens_in), - tokens_out=cast(int, entry.tokens_out), - model=cast(str, entry.model), - cost_usd=cast(float, entry.cost_usd), - timestamp=datetime.fromisoformat(cast(str, entry.timestamp)), + session_id=entry.session_id, + plan_id=entry.plan_id, + tokens_in=entry.tokens_in, + tokens_out=entry.tokens_out, + model=entry.model, + cost_usd=entry.cost_usd, + timestamp=entry.timestamp, ) for entry in entries ] -- 2.52.0