from fastapi.testclient import TestClient
from app import app
import json

client = TestClient(app)

print('Non-streaming test')
payload = {
    'model': 'rwkv-latest',
    'prompt': 'Who is the president of France today?',
    'stream': False,
    'max_tokens': 64,
    'temperature': 0.2,
    'include_usage': True,
}
res = client.post('/api/v1/chat/completions', json=payload)
print('Status', res.status_code)
try:
    print(json.dumps(res.json(), indent=2))
except Exception as e:
    print('Response not JSON or parse failed', e)

print('\nTools calc test')
payload = {
    'model': 'rwkv-latest',
    'prompt': 'Calculate 2+3*4 and explain the result.',
    'stream': False,
    'tools': [{'name': 'calc', 'args': {'expression': '2+3*4'}}],
}
res = client.post('/api/v1/chat/completions', json=payload)
print('Status', res.status_code)
try:
    print(json.dumps(res.json(), indent=2))
except Exception as e:
    print('Response not JSON or parse failed', e)

print('\nTools web_search test')
payload = {
    'model': 'rwkv-latest',
    'prompt': 'Who is the current president of France?',
    'stream': False,
    'web_search': True,
    'search_top_k': 2,
}
res = client.post('/api/v1/chat/completions', json=payload)
print('Status', res.status_code)
try:
    print(json.dumps(res.json(), indent=2))
except Exception as e:
    print('Response not JSON or parse failed', e)