wehub-resource-sync
4a19d70af1
CodeQL Advanced / Analyze (actions) (push) Waiting to run
CodeQL Advanced / Analyze (python) (push) Waiting to run
Tests / database-integration-tests (push) Waiting to run
MCP Server Tests / live-mcp-tests (push) Waiting to run
Server Tests / live-server-tests (push) Waiting to run
Pyright Type Check / pyright (push) Waiting to run
Tests / unit-tests (push) Waiting to run
Lint with Ruff / ruff (push) Has been cancelled
215 lines
7.4 KiB
Python
215 lines
7.4 KiB
Python
"""
|
|
Copyright 2024, Zep Software, Inc.
|
|
|
|
Licensed under the Apache License, Version 2.0 (the "License");
|
|
you may not use this file except in compliance with the License.
|
|
You may obtain a copy of the License at
|
|
|
|
http://www.apache.org/licenses/LICENSE-2.0
|
|
|
|
Unless required by applicable law or agreed to in writing, software
|
|
distributed under the License is distributed on an "AS IS" BASIS,
|
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
See the License for the specific language governing permissions and
|
|
limitations under the License.
|
|
"""
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
|
|
from graphiti_core.llm_client.token_tracker import (
|
|
PromptTokenUsage,
|
|
TokenUsage,
|
|
TokenUsageTracker,
|
|
)
|
|
|
|
|
|
class TestTokenUsage:
|
|
def test_total_tokens(self):
|
|
"""Test that total_tokens correctly sums input and output tokens."""
|
|
usage = TokenUsage(input_tokens=100, output_tokens=50)
|
|
assert usage.total_tokens == 150
|
|
|
|
def test_default_values(self):
|
|
"""Test that default values are zero."""
|
|
usage = TokenUsage()
|
|
assert usage.input_tokens == 0
|
|
assert usage.output_tokens == 0
|
|
assert usage.total_tokens == 0
|
|
|
|
|
|
class TestPromptTokenUsage:
|
|
def test_total_tokens(self):
|
|
"""Test that total_tokens correctly sums input and output tokens."""
|
|
usage = PromptTokenUsage(
|
|
prompt_name='test',
|
|
call_count=5,
|
|
total_input_tokens=1000,
|
|
total_output_tokens=500,
|
|
)
|
|
assert usage.total_tokens == 1500
|
|
|
|
def test_avg_input_tokens(self):
|
|
"""Test average input tokens calculation."""
|
|
usage = PromptTokenUsage(
|
|
prompt_name='test',
|
|
call_count=4,
|
|
total_input_tokens=1000,
|
|
total_output_tokens=500,
|
|
)
|
|
assert usage.avg_input_tokens == 250.0
|
|
|
|
def test_avg_output_tokens(self):
|
|
"""Test average output tokens calculation."""
|
|
usage = PromptTokenUsage(
|
|
prompt_name='test',
|
|
call_count=4,
|
|
total_input_tokens=1000,
|
|
total_output_tokens=500,
|
|
)
|
|
assert usage.avg_output_tokens == 125.0
|
|
|
|
def test_avg_tokens_zero_calls(self):
|
|
"""Test that average returns 0 when call_count is zero."""
|
|
usage = PromptTokenUsage(
|
|
prompt_name='test',
|
|
call_count=0,
|
|
total_input_tokens=0,
|
|
total_output_tokens=0,
|
|
)
|
|
assert usage.avg_input_tokens == 0
|
|
assert usage.avg_output_tokens == 0
|
|
|
|
|
|
class TestTokenUsageTracker:
|
|
def test_record_new_prompt(self):
|
|
"""Test recording usage for a new prompt."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('extract_nodes', 100, 50)
|
|
|
|
usage = tracker.get_usage()
|
|
assert 'extract_nodes' in usage
|
|
assert usage['extract_nodes'].call_count == 1
|
|
assert usage['extract_nodes'].total_input_tokens == 100
|
|
assert usage['extract_nodes'].total_output_tokens == 50
|
|
|
|
def test_record_existing_prompt(self):
|
|
"""Test that multiple calls accumulate correctly."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('extract_nodes', 100, 50)
|
|
tracker.record('extract_nodes', 200, 100)
|
|
|
|
usage = tracker.get_usage()
|
|
assert usage['extract_nodes'].call_count == 2
|
|
assert usage['extract_nodes'].total_input_tokens == 300
|
|
assert usage['extract_nodes'].total_output_tokens == 150
|
|
|
|
def test_record_none_prompt_name(self):
|
|
"""Test that None prompt_name is recorded as 'unknown'."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record(None, 100, 50)
|
|
|
|
usage = tracker.get_usage()
|
|
assert 'unknown' in usage
|
|
assert usage['unknown'].call_count == 1
|
|
|
|
def test_record_multiple_prompts(self):
|
|
"""Test recording usage for multiple different prompts."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('extract_nodes', 100, 50)
|
|
tracker.record('dedupe_nodes', 200, 100)
|
|
tracker.record('extract_edges', 150, 75)
|
|
|
|
usage = tracker.get_usage()
|
|
assert len(usage) == 3
|
|
assert 'extract_nodes' in usage
|
|
assert 'dedupe_nodes' in usage
|
|
assert 'extract_edges' in usage
|
|
|
|
def test_get_usage_returns_copy(self):
|
|
"""Test that get_usage returns a copy, not the internal dict."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('test', 100, 50)
|
|
|
|
usage1 = tracker.get_usage()
|
|
usage1['test'].total_input_tokens = 9999
|
|
|
|
usage2 = tracker.get_usage()
|
|
assert usage2['test'].total_input_tokens == 100 # Original unchanged
|
|
|
|
def test_get_total_usage(self):
|
|
"""Test getting total usage across all prompts."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('extract_nodes', 100, 50)
|
|
tracker.record('dedupe_nodes', 200, 100)
|
|
tracker.record('extract_edges', 150, 75)
|
|
|
|
total = tracker.get_total_usage()
|
|
assert total.input_tokens == 450
|
|
assert total.output_tokens == 225
|
|
assert total.total_tokens == 675
|
|
|
|
def test_get_total_usage_empty(self):
|
|
"""Test getting total usage when no records exist."""
|
|
tracker = TokenUsageTracker()
|
|
total = tracker.get_total_usage()
|
|
assert total.input_tokens == 0
|
|
assert total.output_tokens == 0
|
|
|
|
def test_reset(self):
|
|
"""Test that reset clears all tracked usage."""
|
|
tracker = TokenUsageTracker()
|
|
tracker.record('extract_nodes', 100, 50)
|
|
tracker.record('dedupe_nodes', 200, 100)
|
|
|
|
tracker.reset()
|
|
|
|
usage = tracker.get_usage()
|
|
assert len(usage) == 0
|
|
|
|
total = tracker.get_total_usage()
|
|
assert total.total_tokens == 0
|
|
|
|
def test_thread_safety(self):
|
|
"""Test that concurrent access from multiple threads is safe."""
|
|
tracker = TokenUsageTracker()
|
|
num_threads = 10
|
|
calls_per_thread = 100
|
|
|
|
def record_tokens(thread_id):
|
|
for _ in range(calls_per_thread):
|
|
tracker.record(f'prompt_{thread_id}', 10, 5)
|
|
|
|
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
|
futures = [executor.submit(record_tokens, i) for i in range(num_threads)]
|
|
for f in futures:
|
|
f.result()
|
|
|
|
usage = tracker.get_usage()
|
|
assert len(usage) == num_threads
|
|
|
|
total = tracker.get_total_usage()
|
|
expected_input = num_threads * calls_per_thread * 10
|
|
expected_output = num_threads * calls_per_thread * 5
|
|
assert total.input_tokens == expected_input
|
|
assert total.output_tokens == expected_output
|
|
|
|
def test_concurrent_same_prompt(self):
|
|
"""Test concurrent access to the same prompt name."""
|
|
tracker = TokenUsageTracker()
|
|
num_threads = 10
|
|
calls_per_thread = 100
|
|
|
|
def record_tokens():
|
|
for _ in range(calls_per_thread):
|
|
tracker.record('shared_prompt', 10, 5)
|
|
|
|
with ThreadPoolExecutor(max_workers=num_threads) as executor:
|
|
futures = [executor.submit(record_tokens) for _ in range(num_threads)]
|
|
for f in futures:
|
|
f.result()
|
|
|
|
usage = tracker.get_usage()
|
|
assert usage['shared_prompt'].call_count == num_threads * calls_per_thread
|
|
assert usage['shared_prompt'].total_input_tokens == num_threads * calls_per_thread * 10
|
|
assert usage['shared_prompt'].total_output_tokens == num_threads * calls_per_thread * 5
|