mirror of
https://github.com/browser-use/browser-use
synced 2026-05-06 17:52:15 +02:00
91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
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()
|