from qwen_agent.agents import Assistant
from qwen_agent.utils.output_beautify import typewriter_print
import time
import tiktoken  # 需提前 pip install tiktoken

# Define LLM
llm_cfg = {
    'model': 'qwen3-8b',

    # Use the endpoint provided by Alibaba Model Studio:
    # 'model_type': 'qwen_dashscope',
    # 'api_key': os.getenv('DASHSCOPE_API_KEY'),

    # Use a custom endpoint compatible with OpenAI API:
    'model_server': 'http://192.168.16.116:18088/v1',  # api_base
    'api_key': 'EMPTY',

    # Other parameters:
    # 'generate_cfg': {
    #         # Add: When the response content is `<think>this is the thought</think>this is the answer;
    #         # Do not add: When the response has been separated by reasoning_content and content.
    #         'thought_in_content': True,
    #     },
}

# Define Tools
tools = [
  'code_interpreter',  # Built-in tools
]

# Define Agent
bot = Assistant(llm=llm_cfg, function_list=tools)

query = """
# Instruction
- You are a helpful problem solver who is good at using python interpreter to solve complex problems efficiently.
- Notice: If the code interpreter returns exception, don't repeat submitting the same wrong code. Fix it first, then retry.
- Notice: don't count by your fingures or in mind, use the code interpreter and get the accurate anwser.
- IMPORTANT!! IF THERE ARE CONFLICTS BETWEEN YOUR COUNTING AND THE CODE OUPPUT, TRUST THE CODE, DON’T TRUST YOUR COUNTING! 

# Goal
Write code to solve the following problem. Tell me the customer who completed the fifth most orders, if there are more than one customer are the same amount, provide all of them.

# Context
Interval        Customer        Orders
2022    Jane Doe        5
2022    Jane Smith      5
2022    John Smith      4
2022    Ava Brown       4
2022    Jennifer White  4
2022    Sarah Miller    3
2022    Matt Baker      3
2022    Lucy Garcia     3
2022    Alex Martin     3
2022    Jason Miller    3
2022    Adam Garcia     3
2022    Bob Jones       2
2022    Mary Martin     2
2022    Samantha Jones  2
2022    Lily Potter     2
2022    Grace Nguyen    2
2022    Olivia Lee      2
2022    Alexander Thomas        2
2022    Julia Williams  1
2022    Bob Johnson     1
2022    Jane Smith      1
2022    Daniel Jackson  1
2022    Lisa Kim        1
2022    John Doe        1
2022    Sophie Taylor   1
2022    Alex Johnson    1
2022    Emma Davis      1
2022    Lisa Green      1
2022    Michael Nguyen  1
2022    Katie Wong      1
2022    Samantha Nguyen 1
2023    Sarah Miller    5
2023    Grace Nguyen    5
2023    Michael Nguyen  5
2023    Alex Johnson    4
2023    Matt Baker      3
2023    John Smith      2
2023    Jane Doe        2
2023    Bob Jones       2
2023    Samantha Jones  2
2023    Lily Potter     2
2023    Ava Brown       2
2023    Adam Garcia     2
2023    Alex Martin     2
2023    Julia Williams  1
2023    Bob Johnson     1
2023    Mary Martin     1
2023    John Lee        1
2023    Jane Smith      1
2023    Lucy Garcia     1
2023    Jennifer White  1
2023    Lisa Green      1
2023    Jason Miller    1
2023    Katie Wong      1
2023    Alexander Thomas        1
"""

response_plain_text = ''
print('bot response:')

# 记录开始时间
start_time = time.time()

messages = [{'role': 'user', 'content': f'{query}'}]
try:
    for responses in bot.run(messages=messages):
        response_plain_text = typewriter_print(responses, response_plain_text)
except Exception as e:
    print(f'get exception {e}')

print(responses)

# 记录结束时间
end_time = time.time()
elapsed_time = end_time - start_time
print(f"耗时: {elapsed_time:.2f}秒")

# 统计 token 数
encoding = tiktoken.get_encoding("cl100k_base")
token_count = len(encoding.encode(response_plain_text))

print(f"消耗token数: {token_count}")