Spaces:
Runtime error
Runtime error
File size: 1,320 Bytes
4e763c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
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() |