mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-08-10 16:32:28 +00:00
63d4202432
LiteLLM and flow event paths can pass raw Anthropic usage where input_tokens excludes cache counters; fold them into prompt and total tokens without double-counting native provider payloads. Co-authored-by: Cursor <cursoragent@cursor.com>
420 lines
14 KiB
Python
420 lines
14 KiB
Python
from typing import Any
|
|
from unittest.mock import patch
|
|
|
|
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from crewai.events.event_bus import CrewAIEventsBus
|
|
from crewai.events.types.llm_events import LLMCallCompletedEvent, LLMCallType
|
|
from crewai.llm import LLM
|
|
from crewai.llms.base_llm import BaseLLM
|
|
|
|
|
|
class TestLLMCallCompletedEventUsageField:
|
|
def test_accepts_usage_dict(self):
|
|
event = LLMCallCompletedEvent(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
call_id="test-id",
|
|
usage={"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30},
|
|
)
|
|
assert event.usage == {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
def test_usage_defaults_to_none(self):
|
|
event = LLMCallCompletedEvent(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
call_id="test-id",
|
|
)
|
|
assert event.usage is None
|
|
|
|
def test_accepts_none_usage(self):
|
|
event = LLMCallCompletedEvent(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
call_id="test-id",
|
|
usage=None,
|
|
)
|
|
assert event.usage is None
|
|
|
|
def test_accepts_nested_usage_dict(self):
|
|
usage = {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"total_tokens": 300,
|
|
"prompt_tokens_details": {"cached_tokens": 50},
|
|
}
|
|
event = LLMCallCompletedEvent(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
call_id="test-id",
|
|
usage=usage,
|
|
)
|
|
assert event.usage["prompt_tokens_details"]["cached_tokens"] == 50
|
|
|
|
|
|
class TestUsageToDict:
|
|
def test_none_returns_none(self):
|
|
assert LLM._usage_to_dict(None) is None
|
|
|
|
def test_dict_without_nested_shapes_is_returned_unchanged(self):
|
|
usage = {"prompt_tokens": 10, "total_tokens": 30}
|
|
result = LLM._usage_to_dict(usage)
|
|
assert result == usage
|
|
# The input dict is copied, not mutated, so derived keys are not added.
|
|
assert "cached_prompt_tokens" not in result
|
|
|
|
@pytest.mark.parametrize(
|
|
("usage", "expected"),
|
|
[
|
|
pytest.param(
|
|
{"prompt_tokens": 100, "prompt_tokens_details": {"cached_tokens": 40}},
|
|
{"cached_prompt_tokens": 40},
|
|
id="openai-nested-cached-tokens",
|
|
),
|
|
pytest.param(
|
|
{"prompt_tokens": 100, "cached_tokens": 30},
|
|
{"cached_prompt_tokens": 30},
|
|
id="flat-cached-tokens",
|
|
),
|
|
pytest.param(
|
|
{"input_tokens": 100, "cache_read_input_tokens": 25},
|
|
{"cached_prompt_tokens": 25},
|
|
id="anthropic-cache-read-input-tokens",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"completion_tokens": 200,
|
|
"completion_tokens_details": {"reasoning_tokens": 60},
|
|
},
|
|
{"reasoning_tokens": 60},
|
|
id="openai-nested-reasoning-tokens",
|
|
),
|
|
pytest.param(
|
|
{"input_tokens": 100, "cache_creation_input_tokens": 70},
|
|
{"cache_creation_tokens": 70},
|
|
id="anthropic-cache-creation-input-tokens",
|
|
),
|
|
pytest.param(
|
|
{
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"prompt_tokens_details": {"cached_tokens": 40},
|
|
"completion_tokens_details": {"reasoning_tokens": 60},
|
|
"cache_creation_input_tokens": 10,
|
|
},
|
|
{
|
|
"cached_prompt_tokens": 40,
|
|
"reasoning_tokens": 60,
|
|
"cache_creation_tokens": 10,
|
|
},
|
|
id="all-buckets-from-nested-shapes",
|
|
),
|
|
],
|
|
)
|
|
def test_normalizes_nested_litellm_buckets(self, usage, expected):
|
|
result = LLM._usage_to_dict(usage)
|
|
for key, value in expected.items():
|
|
assert result[key] == value
|
|
|
|
def test_does_not_alter_core_token_counts(self):
|
|
usage = {
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 200,
|
|
"total_tokens": 300,
|
|
"prompt_tokens_details": {"cached_tokens": 40},
|
|
}
|
|
result = LLM._usage_to_dict(usage)
|
|
assert result["prompt_tokens"] == 100
|
|
assert result["completion_tokens"] == 200
|
|
assert result["total_tokens"] == 300
|
|
|
|
def test_absent_buckets_are_not_added(self):
|
|
usage = {"prompt_tokens": 100, "completion_tokens": 200, "total_tokens": 300}
|
|
result = LLM._usage_to_dict(usage)
|
|
assert "cached_prompt_tokens" not in result
|
|
assert "reasoning_tokens" not in result
|
|
assert "cache_creation_tokens" not in result
|
|
|
|
def test_pydantic_model_uses_model_dump(self):
|
|
class Usage(BaseModel):
|
|
prompt_tokens: int = 10
|
|
completion_tokens: int = 20
|
|
total_tokens: int = 30
|
|
|
|
result = LLM._usage_to_dict(Usage())
|
|
assert result == {
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 30,
|
|
}
|
|
|
|
def test_object_with_dict_attr(self):
|
|
class UsageObj:
|
|
def __init__(self):
|
|
self.prompt_tokens = 5
|
|
self.completion_tokens = 15
|
|
self.total_tokens = 20
|
|
|
|
result = LLM._usage_to_dict(UsageObj())
|
|
assert result == {
|
|
"prompt_tokens": 5,
|
|
"completion_tokens": 15,
|
|
"total_tokens": 20,
|
|
}
|
|
|
|
def test_object_with_dict_excludes_private_attrs(self):
|
|
class UsageObj:
|
|
def __init__(self):
|
|
self.total_tokens = 42
|
|
self._internal = "hidden"
|
|
|
|
result = LLM._usage_to_dict(UsageObj())
|
|
assert result == {"total_tokens": 42}
|
|
assert "_internal" not in result
|
|
|
|
def test_unsupported_type_returns_none(self):
|
|
assert LLM._usage_to_dict(42) is None
|
|
assert LLM._usage_to_dict("string") is None
|
|
|
|
|
|
class _StubLLM(BaseLLM):
|
|
"""Minimal concrete BaseLLM for testing event emission."""
|
|
|
|
model: str = "test-model"
|
|
|
|
def call(self, *args: Any, **kwargs: Any) -> str:
|
|
return ""
|
|
|
|
async def acall(self, *args: Any, **kwargs: Any) -> str:
|
|
return ""
|
|
|
|
def supports_function_calling(self) -> bool:
|
|
return False
|
|
|
|
def supports_stop_words(self) -> bool:
|
|
return True
|
|
|
|
|
|
class TestEmitCallCompletedEventPassesUsage:
|
|
@pytest.fixture
|
|
def mock_emit(self):
|
|
with patch.object(CrewAIEventsBus, "emit") as mock:
|
|
yield mock
|
|
|
|
@pytest.fixture
|
|
def llm(self):
|
|
return _StubLLM(model="test-model")
|
|
|
|
def test_usage_is_passed_to_event(self, mock_emit, llm):
|
|
usage_data = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
|
|
|
|
llm._emit_call_completed_event(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
messages="test prompt",
|
|
usage=usage_data,
|
|
)
|
|
|
|
mock_emit.assert_called_once()
|
|
event = mock_emit.call_args[1]["event"]
|
|
assert isinstance(event, LLMCallCompletedEvent)
|
|
assert event.usage == usage_data
|
|
|
|
def test_none_usage_is_passed_to_event(self, mock_emit, llm):
|
|
llm._emit_call_completed_event(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
messages="test prompt",
|
|
usage=None,
|
|
)
|
|
|
|
mock_emit.assert_called_once()
|
|
event = mock_emit.call_args[1]["event"]
|
|
assert isinstance(event, LLMCallCompletedEvent)
|
|
assert event.usage is None
|
|
|
|
def test_usage_omitted_defaults_to_none(self, mock_emit, llm):
|
|
llm._emit_call_completed_event(
|
|
response="hello",
|
|
call_type=LLMCallType.LLM_CALL,
|
|
messages="test prompt",
|
|
)
|
|
|
|
mock_emit.assert_called_once()
|
|
event = mock_emit.call_args[1]["event"]
|
|
assert isinstance(event, LLMCallCompletedEvent)
|
|
assert event.usage is None
|
|
|
|
class TestUsageMetricsNewFields:
|
|
def test_add_usage_metrics_aggregates_reasoning_and_cache_creation(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics1 = UsageMetrics(
|
|
total_tokens=100,
|
|
prompt_tokens=60,
|
|
completion_tokens=40,
|
|
cached_prompt_tokens=10,
|
|
reasoning_tokens=15,
|
|
cache_creation_tokens=5,
|
|
successful_requests=1,
|
|
)
|
|
metrics2 = UsageMetrics(
|
|
total_tokens=200,
|
|
prompt_tokens=120,
|
|
completion_tokens=80,
|
|
cached_prompt_tokens=20,
|
|
reasoning_tokens=25,
|
|
cache_creation_tokens=10,
|
|
successful_requests=1,
|
|
)
|
|
|
|
metrics1.add_usage_metrics(metrics2)
|
|
|
|
assert metrics1.total_tokens == 300
|
|
assert metrics1.prompt_tokens == 180
|
|
assert metrics1.completion_tokens == 120
|
|
assert metrics1.cached_prompt_tokens == 30
|
|
assert metrics1.reasoning_tokens == 40
|
|
assert metrics1.cache_creation_tokens == 15
|
|
assert metrics1.successful_requests == 2
|
|
|
|
def test_new_fields_default_to_zero(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics()
|
|
assert metrics.reasoning_tokens == 0
|
|
assert metrics.cache_creation_tokens == 0
|
|
|
|
def test_model_dump_includes_new_fields(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics(reasoning_tokens=10, cache_creation_tokens=5)
|
|
dumped = metrics.model_dump()
|
|
assert dumped["reasoning_tokens"] == 10
|
|
assert dumped["cache_creation_tokens"] == 5
|
|
|
|
|
|
class TestFromProviderDictAnthropicCacheTokens:
|
|
def test_cache_read_tokens_included_in_prompt_and_total(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 3,
|
|
"output_tokens": 44,
|
|
"cache_read_input_tokens": 2061,
|
|
}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 2064
|
|
assert metrics.completion_tokens == 44
|
|
assert metrics.total_tokens == 2108
|
|
assert metrics.cached_prompt_tokens == 2061
|
|
|
|
def test_cache_creation_tokens_included_in_prompt_and_total(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_creation_input_tokens": 20,
|
|
}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 120
|
|
assert metrics.total_tokens == 170
|
|
|
|
def test_cache_read_and_creation_tokens_both_included(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 30,
|
|
"cache_creation_input_tokens": 20,
|
|
}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 150
|
|
assert metrics.total_tokens == 200
|
|
assert metrics.cached_prompt_tokens == 30
|
|
|
|
def test_missing_cache_fields_preserve_non_cached_totals(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{"input_tokens": 100, "output_tokens": 50}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 100
|
|
assert metrics.total_tokens == 150
|
|
assert metrics.cached_prompt_tokens == 0
|
|
|
|
def test_reconciled_native_dict_is_not_double_counted(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 150,
|
|
"output_tokens": 50,
|
|
"cached_prompt_tokens": 30,
|
|
}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 150
|
|
assert metrics.total_tokens == 200
|
|
|
|
def test_openai_cached_prompt_tokens_are_not_added_twice(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
metrics = UsageMetrics.from_provider_dict(
|
|
{
|
|
"prompt_tokens": 100,
|
|
"completion_tokens": 50,
|
|
"prompt_tokens_details": {"cached_tokens": 30},
|
|
}
|
|
)
|
|
|
|
assert metrics is not None
|
|
assert metrics.prompt_tokens == 100
|
|
assert metrics.total_tokens == 150
|
|
assert metrics.cached_prompt_tokens == 30
|
|
|
|
def test_cumulative_usage_via_add_usage_metrics(self):
|
|
from crewai.types.usage_metrics import UsageMetrics
|
|
|
|
first = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 100,
|
|
"output_tokens": 50,
|
|
"cache_read_input_tokens": 30,
|
|
}
|
|
)
|
|
second = UsageMetrics.from_provider_dict(
|
|
{
|
|
"input_tokens": 40,
|
|
"output_tokens": 20,
|
|
}
|
|
)
|
|
|
|
assert first is not None and second is not None
|
|
first.add_usage_metrics(second)
|
|
|
|
assert first.prompt_tokens == 170
|
|
assert first.completion_tokens == 70
|
|
assert first.total_tokens == 240
|
|
assert first.cached_prompt_tokens == 30
|
|
assert first.successful_requests == 2
|