ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/browser-use-rpa/browser-use/browser_use/tokens/tests/test_cost.py

@@ -0,0 +1,102 @@
+"""
+Simple test for token cost tracking with real LLM calls.
+
+Tests ChatOpenAI and ChatGoogle by iteratively generating countries.
+"""
+
+import asyncio
+import logging
+
+from browser_use.llm import ChatOpenAI
+from browser_use.llm.messages import AssistantMessage, SystemMessage, UserMessage
+from browser_use.tokens.service import TokenCost
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+async def test_iterative_country_generation():
+	"""Test token cost tracking with iterative country generation"""
+
+	# Initialize token cost service
+	tc = TokenCost(include_cost=True)
+
+	# System prompt that explains the iterative task
+	system_prompt = """You are a country name generator. When asked, you will provide exactly ONE country name and nothing else.
+Each time you're asked to continue, provide the next country name that hasn't been mentioned yet.
+Keep track of which countries you've already said and don't repeat them.
+Only output the country name, no numbers, no punctuation, just the name."""
+
+	# Test with different models
+	models = [
+		ChatOpenAI(model='gpt-4.1'),
+		# ChatGoogle(model='gemini-2.0-flash-exp'),
+	]
+
+	print('\n🌍 Iterative Country Generation Test')
+	print('=' * 80)
+
+	for llm in models:
+		print(f'\n📍 Testing {llm.model}')
+		print('-' * 60)
+
+		# Register the LLM for automatic tracking
+		tc.register_llm(llm)
+
+		# Initialize conversation
+		messages = [SystemMessage(content=system_prompt), UserMessage(content='Give me a country name')]
+
+		countries = []
+
+		# Generate 10 countries iteratively
+		for i in range(10):
+			# Call the LLM
+			result = await llm.ainvoke(messages)
+			country = result.completion.strip()
+			countries.append(country)
+
+			# Add the response to messages
+			messages.append(AssistantMessage(content=country))
+
+			# Add the next request (except for the last iteration)
+			if i < 9:
+				messages.append(UserMessage(content='Next country please'))
+
+			print(f'  Country {i + 1}: {country}')
+
+		print(f'\n  Generated countries: {", ".join(countries)}')
+
+	# Display cost summary
+	print('\n💰 Cost Summary')
+	print('=' * 80)
+
+	summary = await tc.get_usage_summary()
+	print(f'Total calls: {summary.entry_count}')
+	print(f'Total tokens: {summary.total_tokens:,}')
+	print(f'Total cost: ${summary.total_cost:.6f}')
+
+	expected_cost = 0
+	expected_invocations = 0
+
+	print('\n📊 Cost breakdown by model:')
+	for model, stats in summary.by_model.items():
+		expected_cost += stats.cost
+		expected_invocations += stats.invocations
+
+		print(f'\n{model}:')
+		print(f'  Calls: {stats.invocations}')
+		print(f'  Prompt tokens: {stats.prompt_tokens:,}')
+		print(f'  Completion tokens: {stats.completion_tokens:,}')
+		print(f'  Total tokens: {stats.total_tokens:,}')
+		print(f'  Cost: ${stats.cost:.6f}')
+		print(f'  Average tokens per call: {stats.average_tokens_per_invocation:.1f}')
+
+	assert summary.entry_count == expected_invocations, f'Expected {expected_invocations} invocations, got {summary.entry_count}'
+	assert abs(summary.total_cost - expected_cost) < 1e-6, (
+		f'Expected total cost ${expected_cost:.6f}, got ${summary.total_cost:.6f}'
+	)
+
+
+if __name__ == '__main__':
+	# Run the test
+	asyncio.run(test_iterative_country_generation())