(Bug fix) - Langfuse / Callback settings stored in DB (#8251)

* fix _decrypt_and_set_db_env_variables

* fix proxy config

* test callbacks in DB

* test langfuse callbacks in db

* test_e2e_langfuse_callbacks_in_db

* proxy_store_model_in_db_tests

* fix proxy_store_model_in_db_tests

* proxy_store_model_in_db_tests

* fix store_model_db_config.yaml

* fix check_langfuse_request

* fix test langfuse base url

* ci/cd run again

.circleci/config.yml

@@ -1464,6 +1464,87 @@ jobs:
       # Store test results
       - store_test_results:
           path: test-results
+
+  proxy_store_model_in_db_tests:
+    machine:
+      image: ubuntu-2204:2023.10.1
+    resource_class: xlarge
+    working_directory: ~/project
+    steps:
+      - checkout
+      - run:
+          name: Install Docker CLI (In case it's not already installed)
+          command: |
+            sudo apt-get update
+            sudo apt-get install -y docker-ce docker-ce-cli containerd.io
+      - run:
+          name: Install Python 3.9
+          command: |
+            curl https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh --output miniconda.sh
+            bash miniconda.sh -b -p $HOME/miniconda
+            export PATH="$HOME/miniconda/bin:$PATH"
+            conda init bash
+            source ~/.bashrc
+            conda create -n myenv python=3.9 -y
+            conda activate myenv
+            python --version
+      - run:
+          name: Install Dependencies
+          command: |
+            pip install "pytest==7.3.1"
+            pip install "pytest-asyncio==0.21.1"
+            pip install aiohttp
+            python -m pip install --upgrade pip
+            python -m pip install -r requirements.txt
+            pip install "pytest==7.3.1"
+            pip install "pytest-retry==1.6.3"
+            pip install "pytest-mock==3.12.0"
+            pip install "pytest-asyncio==0.21.1"
+      - run:
+          name: Build Docker image
+          command: docker build -t my-app:latest -f ./docker/Dockerfile.database .
+      - run:
+          name: Run Docker container
+          # intentionally give bad redis credentials here
+          # the OTEL test - should get this as a trace
+          command: |
+            docker run -d \
+              -p 4000:4000 \
+              -e DATABASE_URL=$PROXY_DATABASE_URL \
+              -e STORE_MODEL_IN_DB="True" \
+              -e LITELLM_MASTER_KEY="sk-1234" \
+              -e LITELLM_LICENSE=$LITELLM_LICENSE \
+              --name my-app \
+              -v $(pwd)/litellm/proxy/example_config_yaml/store_model_db_config.yaml:/app/config.yaml \
+              my-app:latest \
+              --config /app/config.yaml \
+              --port 4000 \
+              --detailed_debug \
+      - run:
+          name: Install curl and dockerize
+          command: |
+            sudo apt-get update
+            sudo apt-get install -y curl
+            sudo wget https://github.com/jwilder/dockerize/releases/download/v0.6.1/dockerize-linux-amd64-v0.6.1.tar.gz
+            sudo tar -C /usr/local/bin -xzvf dockerize-linux-amd64-v0.6.1.tar.gz
+            sudo rm dockerize-linux-amd64-v0.6.1.tar.gz
+      - run:
+          name: Start outputting logs
+          command: docker logs -f my-app
+          background: true
+      - run:
+          name: Wait for app to be ready
+          command: dockerize -wait http://localhost:4000 -timeout 5m
+      - run:
+          name: Run tests
+          command: |
+            pwd
+            ls
+            python -m pytest -vv tests/store_model_in_db_tests -x --junitxml=test-results/junit.xml --durations=5
+          no_output_timeout:
+            120m
+            # Clean up first container
+      
   proxy_build_from_pip_tests:
     # Change from docker to machine executor
     machine:
@@ -2012,6 +2093,12 @@ workflows:
               only:
                 - main
                 - /litellm_.*/
+      - proxy_store_model_in_db_tests:
+          filters:
+            branches:
+              only:
+                - main
+                - /litellm_.*/
       - proxy_build_from_pip_tests:
           filters:
             branches:
@@ -2135,6 +2222,7 @@ workflows:
             - installing_litellm_on_python
             - installing_litellm_on_python_3_13
             - proxy_logging_guardrails_model_info_tests
+            - proxy_store_model_in_db_tests
             - proxy_build_from_pip_tests
             - proxy_pass_through_endpoint_tests
             - check_code_and_doc_quality

litellm/proxy/example_config_yaml/store_model_db_config.yaml

@@ -0,0 +1,10 @@
+model_list:
+  - model_name: fake-openai-endpoint
+    litellm_params:
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
+
+general_settings:
+  store_model_in_db: true
+

litellm/proxy/proxy_config.yaml

@@ -1,11 +1,9 @@
 model_list:
-  - model_name: bedrock/*
+  - model_name: fake-openai-endpoint
     litellm_params:
-      model: bedrock/*
-
-litellm_settings:
-  callbacks: ["datadog"]
-
+      model: openai/my-fake-model
+      api_key: my-fake-key
+      api_base: https://exampleopenaiendpoint-production.up.railway.app/
 general_settings:
-  store_prompts_in_spend_logs: true
+  store_model_in_db: true
 

litellm/proxy/proxy_server.py

@@ -1689,7 +1689,7 @@ class ProxyConfig:
             # default to file
             config = await self._get_config_from_file(config_file_path=config_file_path)
         ## UPDATE CONFIG WITH DB
-        if prisma_client is not None:
+        if prisma_client is not None and store_model_in_db is True:
             config = await self._update_config_from_db(
                 config=config,
                 prisma_client=prisma_client,
@@ -2525,6 +2525,16 @@ class ProxyConfig:
         Adds environment variables from DB config to litellm
         """
         environment_variables = config_data.get("environment_variables", {})
+        self._decrypt_and_set_db_env_variables(environment_variables)
+
+    def _decrypt_and_set_db_env_variables(self, environment_variables: dict) -> None:
+        """
+        Decrypts a dictionary of environment variables and then sets them in the environment
+
+        Args:
+            environment_variables: dict - dictionary of environment variables to decrypt and set
+            eg. `{"LANGFUSE_PUBLIC_KEY": "kFiKa1VZukMmD8RB6WXB9F......."}`
+        """
         for k, v in environment_variables.items():
             try:
                 decrypted_value = decrypt_value_helper(value=v)
@@ -2653,19 +2663,43 @@ class ProxyConfig:
     def _update_config_fields(
         self,
         current_config: dict,
-        param_name: str,
+        param_name: Literal[
+            "general_settings",
+            "router_settings",
+            "litellm_settings",
+            "environment_variables",
+        ],
         db_param_value: Any,
     ) -> dict:
+        """
+        Updates the config fields with the new values from the DB
+
+        Args:
+            current_config (dict): Current configuration dictionary to update
+            param_name (Literal): Name of the parameter to update
+            db_param_value (Any): New value from the database
+
+        Returns:
+            dict: Updated configuration dictionary
+        """
+        if param_name == "environment_variables":
+            self._decrypt_and_set_db_env_variables(db_param_value)
+            return current_config
+
+        # If param doesn't exist in config, add it
+        if param_name not in current_config:
+            current_config[param_name] = db_param_value
+            return current_config
+
+        # For dictionary values, update only non-empty values
         if isinstance(current_config[param_name], dict):
-            # if dict exists (e.g. litellm_settings),
-            # go through each key and value,
-            # and update if new value is not None/empty dict
-            for key, value in db_param_value.items():
-                if value:
-                    current_config[param_name][key] = value
+            # Only keep non None values from db_param_value
+            non_empty_values = {k: v for k, v in db_param_value.items() if v}
+
+            # Update the config with non-empty values
+            current_config[param_name].update(non_empty_values)
         else:
             current_config[param_name] = db_param_value
-
         return current_config
 
     async def _update_config_from_db(
@@ -2674,7 +2708,6 @@ class ProxyConfig:
         config: dict,
         store_model_in_db: Optional[bool],
     ):
-
         if store_model_in_db is not True:
             verbose_proxy_logger.info(
                 "'store_model_in_db' is not True, skipping db updates"
@@ -2696,24 +2729,21 @@ class ProxyConfig:
 
         responses = await asyncio.gather(*_tasks)
         for response in responses:
-            if response is not None:
-                param_name = getattr(response, "param_name", None)
-                if param_name == "litellm_settings":
-                    verbose_proxy_logger.info(
-                        f"litellm_settings: {response.param_value}"
-                    )
-                param_value = getattr(response, "param_value", None)
-                if param_name is not None and param_value is not None:
-                    # check if param_name is already in the config
-                    if param_name in config:
-                        config = self._update_config_fields(
-                            current_config=config,
-                            param_name=param_name,
-                            db_param_value=param_value,
-                        )
-                    else:
-                        # if it's not in the config - then add it
-                        config[param_name] = param_value
+            if response is None:
+                continue
+
+            param_name = getattr(response, "param_name", None)
+            param_value = getattr(response, "param_value", None)
+            verbose_proxy_logger.debug(
+                f"param_name={param_name}, param_value={param_value}"
+            )
+
+            if param_name is not None and param_value is not None:
+                config = self._update_config_fields(
+                    current_config=config,
+                    param_name=param_name,
+                    db_param_value=param_value,
+                )
 
         return config
 
@@ -5786,7 +5816,7 @@ async def token_counter(request: TokenCountRequest):
         model=model_to_use,
         text=prompt,
         messages=messages,
-        custom_tokenizer=_tokenizer_used,
+        custom_tokenizer=_tokenizer_used,  # type: ignore
     )
     return TokenCountResponse(
         total_tokens=total_tokens,

tests/store_model_in_db_tests/test_callbacks_in_db.py

@@ -0,0 +1,93 @@
+"""
+PROD TEST - DO NOT Delete this Test
+
+e2e test for langfuse callback in DB 
+- Add langfuse callback to DB - with /config/update
+- wait 20 seconds for the callback to be loaded into the instance 
+- Make a /chat/completions request to the proxy
+- Check if the request is logged in Langfuse
+"""
+
+import pytest
+import asyncio
+import aiohttp
+import os
+import dotenv
+from dotenv import load_dotenv
+import pytest
+from openai import AsyncOpenAI
+from openai.types.chat import ChatCompletion
+
+load_dotenv()
+
+# used for testing
+LANGFUSE_BASE_URL = "https://exampleopenaiendpoint-production-c715.up.railway.app"
+
+
+async def config_update(session, routing_strategy=None):
+    url = "http://0.0.0.0:4000/config/update"
+    headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
+    print("routing_strategy: ", routing_strategy)
+    data = {
+        "litellm_settings": {"success_callback": ["langfuse"]},
+        "environment_variables": {
+            "LANGFUSE_PUBLIC_KEY": "any-public-key",
+            "LANGFUSE_SECRET_KEY": "any-secret-key",
+            "LANGFUSE_HOST": LANGFUSE_BASE_URL,
+        },
+    }
+
+    async with session.post(url, headers=headers, json=data) as response:
+        status = response.status
+        response_text = await response.text()
+
+        print(response_text)
+        print("status: ", status)
+
+        if status != 200:
+            raise Exception(f"Request did not return a 200 status code: {status}")
+        return await response.json()
+
+
+async def check_langfuse_request(response_id: str):
+    async with aiohttp.ClientSession() as session:
+        url = f"{LANGFUSE_BASE_URL}/langfuse/trace/{response_id}"
+        async with session.get(url) as response:
+            response_json = await response.json()
+            assert response.status == 200, f"Expected status 200, got {response.status}"
+            assert (
+                response_json["exists"] == True
+            ), f"Request {response_id} not found in Langfuse traces"
+            assert response_json["request_id"] == response_id, f"Request ID mismatch"
+
+
+async def make_chat_completions_request() -> ChatCompletion:
+    client = AsyncOpenAI(api_key="sk-1234", base_url="http://0.0.0.0:4000")
+    response = await client.chat.completions.create(
+        model="fake-openai-endpoint",
+        messages=[{"role": "user", "content": "Hello, world!"}],
+    )
+    print(response)
+    return response
+
+
+@pytest.mark.asyncio
+async def test_e2e_langfuse_callbacks_in_db():
+
+    session = aiohttp.ClientSession()
+
+    # add langfuse callback to DB
+    await config_update(session)
+
+    # wait 20 seconds for the callback to be loaded into the instance
+    await asyncio.sleep(20)
+
+    # make a /chat/completions request to the proxy
+    response = await make_chat_completions_request()
+    print(response)
+    response_id = response.id
+    print("response_id: ", response_id)
+
+    await asyncio.sleep(11)
+    # check if the request is logged in Langfuse
+    await check_langfuse_request(response_id)