import json
import os
import tiktoken
from browser_use.agent.views import AgentOutput
from browser_use.controller.service import Controller
from browser_use.llm.schema import SchemaOptimizer
def test_optimized_schema():
"""Test the optimized schema generation and save to file."""
# Create controller and get all registered actions
controller = Controller()
ActionModel = controller.registry.create_action_model()
# Create the agent output model with custom actions
agent_output_model = AgentOutput.type_with_custom_actions(ActionModel)
# Get original schema for comparison
original_schema = agent_output_model.model_json_schema()
# Create the optimized schema
optimized_schema = SchemaOptimizer.create_optimized_json_schema(agent_output_model)
# Create tmp directory if it doesn't exist
os.makedirs('./tmp', exist_ok=True)
# Save optimized schema
with open('./tmp/optimized_schema.json', 'w') as f:
json.dump(optimized_schema, f, separators=(',', ':'), indent=2)
print('✅ Optimized schema generated and saved to ./tmp/optimized_schema.json')
# Compare token counts of both
try:
enc = tiktoken.encoding_for_model('gpt-4o')
except KeyError:
enc = tiktoken.get_encoding('cl100k_base')
original_tokens = len(enc.encode(json.dumps(original_schema)))
optimized_tokens = len(enc.encode(json.dumps(optimized_schema, separators=(',', ':'))))
savings = original_tokens - optimized_tokens
savings_percentage = (savings / original_tokens * 100) if original_tokens > 0 else 0
print('\n📊 Token Count Comparison:')
print(f' Original schema: {original_tokens:,} tokens')
print(f' Optimized schema: {optimized_tokens:,} tokens')
print(f' Token savings: {savings:,} tokens ({savings_percentage:.1f}% reduction)')
# Count tokens per action in optimized schema
print('\n🔍 Tokens per Action in Optimized Schema:')
if 'properties' in optimized_schema and 'action' in optimized_schema['properties']:
action_prop = optimized_schema['properties']['action']
if 'items' in action_prop and 'anyOf' in action_prop['items']:
actions = action_prop['items']['anyOf']
total_action_tokens = 0
for i, action in enumerate(actions):
action_json = json.dumps(action, separators=(',', ':'))
action_tokens = len(enc.encode(action_json))
total_action_tokens += action_tokens
# Try to get action name from the schema
action_name = 'Unknown'
if 'properties' in action:
# Get the first property that's not common ones like 'index', 'reasoning'
for prop_name in action['properties'].keys():
if prop_name not in ['index', 'reasoning']:
action_name = prop_name
break
print(f' Action {i + 1} ({action_name}): {action_tokens:,} tokens')
print('\n📈 Summary:')
print(f' Total actions: {len(actions)}')
print(f' Total action tokens: {total_action_tokens:,} tokens')
print(f' Average tokens per action: {total_action_tokens // len(actions):,} tokens')
print(f' Action tokens as % of total: {(total_action_tokens / optimized_tokens * 100):.1f}%')
else:
print(' No actions found in expected schema structure')
else:
print(' No action property found in optimized schema')
if __name__ == '__main__':
test_optimized_schema()
