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

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

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

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()