Updates loading of config variables and litellm variables

config/configuration_variables.py

@@ -2,13 +2,14 @@
 from pathlib import Path
 
 from labs.logger import setup_logger
-
+from .models import Config
 
 setup_logger()
 
-
 PROJ_ROOT = Path(__file__).resolve().parents[1]
 
+configs = Config.objects.all()
+
 DATA_DIR = PROJ_ROOT / "data"
 RAW_DATA_DIR = DATA_DIR / "raw"
 INTERIM_DATA_DIR = DATA_DIR / "interim"
@@ -27,16 +28,16 @@
 
 TEST_ENVIRONMENT = eval(os.environ.get("TEST_ENVIRONMENT", "False"))
 
-LITELLM_MASTER_KEY = os.environ.get("LITELLM_MASTER_KEY")
-LITELLM_API_KEY = os.environ.get("LITELLM_API_KEY")
+LITELLM_MASTER_KEY = configs.get(key="LITELLM_MASTER_KEY") or os.environ.get("LITELLM_MASTER_KEY")
+LITELLM_API_KEY = configs.get(key="LITELLM_API_KEY") or os.environ.get("LITELLM_API_KEY")
 
 CLONE_DESTINATION_DIR = os.getenv("CLONE_DESTINATION_DIR", "/tmp/")
 LLM_MODEL_NAME = os.getenv("LLM_MODEL_NAME", "openai/gpt-4o")
 
-OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY")
-GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
-COHERE_API_KEY = os.environ.get("COHERE_API_KEY")
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+OPENAI_API_KEY = configs.get(key="OPENAI_API_KEY") or os.environ.get("OPENAI_API_KEY")
+GEMINI_API_KEY = configs.get(key="GEMINI_API_KEY") or os.environ.get("GEMINI_API_KEY")
+COHERE_API_KEY = configs.get(key="COHERE_API_KEY") or os.environ.get("COHERE_API_KEY")
+GROQ_API_KEY = configs.get(key="GROQ_API_KEY") or os.environ.get("GROQ_API_KEY")
 
 DATABASE_USER = os.environ.get("DATABASE_USER", "postgres")
 DATABASE_PASS = os.environ.get("DATABASE_PASS", "postgres")

config/lite_llm_database_config.py

@@ -0,0 +1,45 @@
+from django.db.models.signals import post_save
+from django.dispatch import receiver
+from .models import Config
+import yaml
+
+
+configs= {}
+
+def load_config_from_db():
+  config_objs = Config.objects.all()
+
+  for config in config_objs:
+    configs[config.key] = config.value
+
+  return configs
+
+def rewrite_config_yaml():
+  yaml_config = {}
+
+  with open("litellm.yaml", "r") as file:
+    yaml_config = yaml.safe_load(file)
+
+  for model in yaml_config['model_list']:
+    for key, value in model['litellm_params']:
+      model['litellm_params'][key] = get_variable(value)
+
+  with open("tmp_litellm.yaml", "w") as file:
+    yaml.dump(yaml_config, file)
+
+def get_variable(value):
+  if 'os.environ/' in value:
+    env_name = value.split('/')[1]
+    return configs[env_name]
+  else: 
+    return value
+
+@receiver(post_save, sender=Config)
+def update_api_keys(sender, instance, **kwargs):
+  global configs
+  configs = load_config_from_db()
+  rewrite_config_yaml()
+  print(f"Configuration reloaded.")
+
+
+

config/management/commands/reload_litellm_keys.py

@@ -0,0 +1,6 @@
+from django.core.management.base import BaseCommand
+from config.models import Config
+from config.lite_llm_database_config import load_config_from_db, rewrite_config_yaml
+
+class Command(BaseCommand):
+  help = 'Reload API Keys for lite llm from the database'