pdf-flashcards-autogen / test_llm_utils.py
LE Quoc Dat
wip
4e763c9
import os
import yaml
from llm_utils import generate_completion
def test_generate_completion():
# Load config to get the environment variable name
with open('models.yaml', 'r') as file:
config = yaml.safe_load(file)
# Get the first environment variable name
env_var_name = list(config['models'][0].keys())[0]
# Get API key from environment variable
api_key = os.getenv(env_var_name)
if not api_key:
raise ValueError(f"Please set {env_var_name} environment variable")
# Test prompt
test_prompt = "What is 2+2? Answer in one word."
try:
# Test with explicit API key
response = generate_completion(test_prompt, api_key)
print(f"Test prompt: {test_prompt}")
print(f"Response with explicit API key: {response}")
assert isinstance(response, str)
assert len(response) > 0
# Test with environment variable
response = generate_completion(test_prompt)
print(f"Response with environment variable: {response}")
assert isinstance(response, str)
assert len(response) > 0
print("Test passed successfully!")
except Exception as e:
print(f"Test failed with error: {str(e)}")
if __name__ == "__main__":
test_generate_completion()