from types import SimpleNamespace from uuid import uuid4 import pytest from prometheus_client import generate_latest from investibook.core.agent_runtime import AgentEventBridge from investibook.core.token_usage import TokenCostRates, record_backend_llm_usage from investibook.db import queries as dbq def test_web_acp_client_accumulates_reasoning_tokens_from_meta_and_typed_usage(): client = AgentEventBridge(workspace_id="ws-reasoning") client._accumulate_usage( { "usage": { "inputTokens": 10, "outputTokens": 4, "thoughtTokens": 3, "cacheReadInputTokens": 2, "cacheWriteInputTokens": 1, } } ) client._accumulate_usage_from_resp( SimpleNamespace( input_tokens=5, output_tokens=6, reasoning_tokens=7, cached_read_tokens=8, cached_write_tokens=9, ) ) usage = client.get_and_reset_usage() assert usage["input_tokens"] == 15 assert usage["output_tokens"] == 10 assert usage["reasoning_tokens"] == 10 assert usage["cache_read_tokens"] == 10 assert usage["cache_write_tokens"] == 10 assert client.get_and_reset_usage()["reasoning_tokens"] == 0 def test_db_insert_and_read_token_usage_reasoning_tokens(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() dbq._db_insert_token_usage( workspace_id="ws-reasoning", session_id="session-1", model_tier="standard", channel="chat", user_email="user@example.com", input_tokens=100, output_tokens=20, reasoning_tokens=11, cache_read_tokens=3, cache_write_tokens=4, cost_amount=0.123, cost_currency="USD", ) rows = dbq._db_get_workspace_token_usage("ws-reasoning") assert len(rows) == 1 assert rows[0]["input_tokens"] == 100 assert rows[0]["output_tokens"] == 20 assert rows[0]["reasoning_tokens"] == 11 def test_backend_owned_usage_prices_and_persists_reasoning_tokens(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() result = record_backend_llm_usage( workspace_id="ws-backend", session_id="backend-session", model_tier="evidence-postpass", channel="evidence-postpass", raw_usage={ "inputTokens": 1_000_000, "outputTokens": 1_000_000, "thoughtTokens": 500_000, }, rates=TokenCostRates(input_per_million=1.0, output_per_million=2.0), user_email="backend@example.com", ) assert result["reasoning_tokens"] == 500_000 assert result["cost_amount"] == pytest.approx(4.0) rows = dbq._db_get_workspace_token_usage("ws-backend") assert rows[0]["reasoning_tokens"] == 500_000 assert rows[0]["cost_amount"] == pytest.approx(4.0) def test_backend_owned_usage_exports_runtime_aware_cost_metric(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() workspace_id = f"ws-metric-{uuid4().hex}" record_backend_llm_usage( workspace_id=workspace_id, session_id=f"image-parse:{uuid4().hex}", model_tier="image-parse", channel="image-send", raw_usage={"inputTokens": 1_000_000}, rates=TokenCostRates(input_per_million=1.0, output_per_million=2.0), user_email="metric@example.com", usage_scenario="image_parse", ) metrics_text = generate_latest().decode("utf-8") assert "investibook_acp_interaction_cost_total" in metrics_text assert f'workspace_id="{workspace_id}"' in metrics_text assert 'runtime="image_parse"' in metrics_text assert 'cost_semantics="interaction_increment"' in metrics_text def test_workspace_token_usage_totals_keep_opencode_cumulative_cost(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() contributions = [] for input_tokens, cost_amount in [(100, 1.25), (300, 2.75)]: contributions.append(dbq._db_insert_token_usage( workspace_id="ws-opencode-cost", session_id="ses_opencode_1", model_tier="standard", channel="web", user_email="user@example.com", input_tokens=input_tokens, output_tokens=10, cost_amount=cost_amount, cost_currency="USD", )) usage = dbq._db_workspace_token_usage_totals("ws-opencode-cost") assert contributions == pytest.approx([1.25, 1.5]) assert usage["session_count"] == 1 assert usage["input_tokens"] == 300 assert usage["cost_amount"] == pytest.approx(2.75) assert dbq._db_workspace_total_recorded_cost("ws-opencode-cost") == pytest.approx(2.75) def test_workspace_token_usage_totals_sum_pi_incremental_cost(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() pi_session_id = str(uuid4()) contributions = [] for input_tokens, cost_amount in [(100, 1.25), (300, 2.75)]: contributions.append(dbq._db_insert_token_usage( workspace_id="ws-pi-cost", session_id=pi_session_id, model_tier="standard", channel="web", user_email="user@example.com", input_tokens=input_tokens, output_tokens=10, cost_amount=cost_amount, cost_currency="USD", )) usage = dbq._db_workspace_token_usage_totals("ws-pi-cost") assert contributions == pytest.approx([1.25, 2.75]) assert usage["session_count"] == 1 assert usage["input_tokens"] == 400 assert usage["cost_amount"] == pytest.approx(4.0) assert dbq._db_workspace_total_recorded_cost("ws-pi-cost") == pytest.approx(4.0) def test_workspace_token_usage_totals_sum_backend_owned_incremental_cost(tmp_path, monkeypatch): monkeypatch.setattr(dbq, "DB_PATH", tmp_path / "investibook.db") dbq._db_init() contributions = [] for attachment_id, cost_amount in [("img_1", 0.1), ("img_2", 0.2)]: contributions.append(dbq._db_insert_token_usage( workspace_id="ws-image-cost", session_id=f"image-parse:{attachment_id}", model_tier="image-parse", channel="image-send", user_email="user@example.com", input_tokens=10, output_tokens=2, cost_amount=cost_amount, cost_currency="USD", )) usage = dbq._db_workspace_token_usage_totals("ws-image-cost") assert contributions == pytest.approx([0.1, 0.2]) assert usage["session_count"] == 2 assert usage["input_tokens"] == 20 assert usage["cost_amount"] == pytest.approx(0.3) assert dbq._db_workspace_total_recorded_cost("ws-image-cost") == pytest.approx(0.3)