agent-task-executor/tests/test_llm.py
zhukang 0cfcd3c576 test: improve async test setup
- Add pytest-asyncio for handling async tests
- Update test_llm.py with proper async test decorators
- Add more assertions to LLM test
2025-01-14 20:56:32 +08:00

61 lines
2.0 KiB
Python

"""Test LLM API call."""
import asyncio
import pytest
from config.config_loader import load_config
import httpx
@pytest.mark.asyncio
async def test_llm_call():
"""Test LLM API call with a simple prompt."""
config = load_config()
llm_config = config["llm"]
headers = {
"Authorization": f"Bearer {llm_config['api_key']}",
"Content-Type": "application/json"
}
# Deepseek chat completion payload
payload = {
"model": llm_config["model"],
"messages": [
{"role": "user", "content": "Say hello and introduce yourself in one sentence."}
],
"temperature": llm_config.get("temperature", 0.7),
"max_tokens": llm_config.get("max_tokens", 2000)
}
print("Making API call to Deepseek...")
print(f"Base URL: {llm_config['api_base']}")
print(f"Model: {llm_config['model']}")
async with httpx.AsyncClient() as client:
try:
response = await client.post(
f"{llm_config['api_base']}/v1/chat/completions",
headers=headers,
json=payload,
timeout=llm_config["timeout"]
)
response.raise_for_status()
result = response.json()
print("\nResponse:")
print(result["choices"][0]["message"]["content"])
assert response.status_code == 200
assert "choices" in result
assert len(result["choices"]) > 0
assert "message" in result["choices"][0]
assert "content" in result["choices"][0]["message"]
content = result["choices"][0]["message"]["content"]
assert isinstance(content, str)
assert len(content) > 0
except httpx.HTTPStatusError as e:
print(f"HTTP error occurred: {e.response.text}")
except Exception as e:
print(f"An error occurred: {str(e)}")
if __name__ == "__main__":
asyncio.run(test_llm_call())