diff --git a/features/cost_tracker_service.feature b/features/cost_tracker_service.feature new file mode 100644 index 000000000..0d5604c5e --- /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.00075 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.00075 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..d02a5339e --- /dev/null +++ b/features/steps/cost_tracker_steps.py @@ -0,0 +1,170 @@ +"""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 + 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.""" + 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 + + +@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..c358b8107 --- /dev/null +++ b/src/cleveragents/infrastructure/database/cost_tracker.py @@ -0,0 +1,242 @@ +"""Cost tracking service for per-session and per-plan LLM spending.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from pathlib import Path + +from pydantic import BaseModel +from sqlalchemy import DateTime, Float, Integer, String, create_engine +from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, sessionmaker + + +class _Base(DeclarativeBase): + pass + + +class CostRecord(BaseModel): + """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: 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: + """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" + + 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}, + } + + 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.now(UTC).replace(tzinfo=None) + + 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, + ) + 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: + entries = ( + session.query(CostEntry) + .filter(CostEntry.session_id == session_id) + .all() + ) + return round(sum(entry.cost_usd for entry in entries), 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: + 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() + + 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=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=entry.timestamp, + ) + for entry in entries + ] + finally: + session.close()