pdf-flashcards-autogen / llm_utils.py
LE Quoc Dat
wip
4e763c9
raw
history blame
1.18 kB
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