ed/manual_slop
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

never ends

Browse Source
This commit is contained in:
Edward R. Gonzalez
2026-03-05 20:39:56 -05:00
parent 2c5476dc5d
commit d2481b2de7
7 changed files with 129 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files

70
tests/test_token_usage.py
View File

@@ -10,43 +10,33 @@ sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from src import ai_client
def test_token_usage_tracking() -> None:
ai_client.reset_session()
with patch("src.ai_client._ensure_gemini_client"), \
patch("src.ai_client._gemini_client") as mock_client:
mock_chat = MagicMock()
mock_client.chats.create.return_value = mock_chat
mock_usage = SimpleNamespace(
prompt_token_count=100,
candidates_token_count=50,
total_token_count=150,
cached_content_token_count=20
)
mock_part = SimpleNamespace(text="Mock Response", function_call=None)
mock_content = SimpleNamespace(parts=[mock_part])
mock_candidate = SimpleNamespace()
mock_candidate.content = mock_content
mock_candidate.finish_reason = SimpleNamespace(name="STOP")
mock_response = SimpleNamespace()
mock_response.candidates = [mock_candidate]
mock_response.usage_metadata = mock_usage
mock_response.text = "Mock Response"
mock_chat.send_message.return_value = mock_response
ai_client.set_provider("gemini", "gemini-2.5-flash-lite")
ai_client.send("Context", "Hello")
comms = ai_client.get_comms_log()
response_entries = [e for e in comms if e.get("direction") == "IN" and e["kind"] == "response"]
assert len(response_entries) > 0
usage = response_entries[0]["payload"]["usage"]
assert usage["input_tokens"] == 100
assert usage["output_tokens"] == 50
assert usage["cache_read_input_tokens"] == 20
ai_client.reset_session()
with patch("src.ai_client._ensure_gemini_client"), \
patch("src.ai_client._gemini_client") as mock_client:
mock_chat = MagicMock()
mock_client.chats.create.return_value = mock_chat
mock_usage = SimpleNamespace(
prompt_token_count=100,
candidates_token_count=50,
total_token_count=150,
cached_content_token_count=20
)
mock_part = SimpleNamespace(text="Mock Response", function_call=None)
mock_content = SimpleNamespace(parts=[mock_part])
mock_candidate = SimpleNamespace()
mock_candidate.content = mock_content
mock_candidate.finish_reason = SimpleNamespace(name="STOP")
mock_response = SimpleNamespace()
mock_response.candidates = [mock_candidate]
mock_response.usage_metadata = mock_usage
mock_response.text = "Mock Response"
mock_chat.send_message.return_value = mock_response
ai_client.set_provider("gemini", "gemini-2.5-flash-lite")
ai_client.send("Context", "Hello")
comms = ai_client.get_comms_log()
response_entries = [e for e in comms if e.get("direction") == "IN" and e["kind"] == "response"]
assert len(response_entries) > 0
usage = response_entries[0]["payload"]["usage"]
assert usage["input_tokens"] == 100
assert usage["output_tokens"] == 50
assert usage["cache_read_input_tokens"] == 20