Abdelrahman Abdallah
42 lines
1.5 KiB
Python
42 lines
1.5 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
from src.agent_context_usage import collect_context_usage, format_context_usage
|
|
from src.agent_session import AgentSessionState
|
|
|
|
|
|
class AgentContextUsageTests(unittest.TestCase):
|
|
def test_collect_context_usage_formats_breakdown(self) -> None:
|
|
session = AgentSessionState.create(
|
|
['# Intro\nhello', '# System\nworld'],
|
|
'inspect repo',
|
|
user_context={'currentDate': "Today's date is 2026-04-01."},
|
|
system_context={'gitStatus': 'Current branch: main'},
|
|
)
|
|
session.append_assistant(
|
|
'Reading files',
|
|
(
|
|
{
|
|
'id': 'call_1',
|
|
'type': 'function',
|
|
'function': {'name': 'read_file', 'arguments': '{"path":"a.py"}'},
|
|
},
|
|
),
|
|
)
|
|
session.append_tool('read_file', 'call_1', '{"ok": true, "content": "print(1)"}')
|
|
|
|
report = collect_context_usage(
|
|
session=session,
|
|
model='Qwen/Qwen3-Coder-30B-A3B-Instruct',
|
|
strategy='test session',
|
|
)
|
|
rendered = format_context_usage(report)
|
|
|
|
self.assertGreater(report.total_tokens, 0)
|
|
self.assertIn('## Context Usage', rendered)
|
|
self.assertIn('**Token counter:**', rendered)
|
|
self.assertIn('### System Prompt Sections', rendered)
|
|
self.assertIn('### Message Breakdown', rendered)
|
|
self.assertIn('#### Top Tools', rendered)
|