import vertexai
from vertexai.generative_models import GenerativeModel

# export GOOGLE_APPLICATION_CREDENTIALS="/home/weifengsun/tangou1/step2/gemini.json"

# You can find your project ID by running `gcloud config get-value project` in the terminal
vertexai.init(project="tangou")
model = GenerativeModel("gemini-2.5-flash")
response = model.generate_content("hello")
print(response.text)

# 打印Token使用情况
print("\n" + "="*50)
print("Token Usage Information:")
print("="*50)

# 获取usage metadata
usage_metadata = response.usage_metadata
print(f"输入Token数 (Prompt Tokens): {usage_metadata.prompt_token_count}")
print(f"输出Token数 (Completion Tokens): {usage_metadata.candidates_token_count}")
print(f"总Token数 (Total Tokens): {usage_metadata.total_token_count}")

# 计算API价格 (Gemini 2.5 Flash定价)
# 输入: $0.075 per 1M tokens
# 输出: $0.30 per 1M tokens
input_price_per_million = 0.075
output_price_per_million = 0.30

input_cost = (usage_metadata.prompt_token_count / 1_000_000) * input_price_per_million
output_cost = (usage_metadata.candidates_token_count / 1_000_000) * output_price_per_million
total_cost = input_cost + output_cost

print("\n" + "="*50)
print("API Price Calculation:")
print("="*50)
print(f"输入Token价格: ${input_cost:.8f}")
print(f"输出Token价格: ${output_cost:.8f}")
print(f"总价格: ${total_cost:.8f}")
print("="*50)