Spaces:
Runtime error
Runtime error
from litellm import completion | |
import yaml | |
import os | |
def load_model_config(): | |
with open('models.yaml', 'r') as file: | |
return yaml.safe_load(file) | |
def generate_completion(prompt: str, api_key: str = None) -> str: | |
""" | |
Generate completion using LiteLLM with the configured model | |
Args: | |
prompt (str): The input prompt | |
api_key (str, optional): Override API key. If not provided, will use environment variable | |
Returns: | |
str: The generated completion text | |
""" | |
config = load_model_config() | |
# Get the first environment variable and its models | |
first_env_var = list(config['models'][0].keys())[0] | |
model_name = config['models'][0][first_env_var][0] | |
# If no API key provided, get from environment | |
if api_key is None: | |
api_key = os.getenv(first_env_var) | |
if not api_key: | |
raise ValueError(f"Please set {first_env_var} environment variable") | |
messages = [{"role": "user", "content": prompt}] | |
response = completion( | |
model=model_name, | |
messages=messages, | |
api_key=api_key | |
) | |
return response.choices[0].message.content |