Merge pull request #2223 from BerriAI/litellm_clickhouse_create_table_views

[FEAT] Clickhouse - Create Analytics Tables

enterprise/utils.py

@@ -110,3 +110,138 @@ async def view_spend_logs_from_clickhouse(
             "log_count": num_rows,
         }
         return response_data
+
+
+def _create_clickhouse_material_views(client=None, table_names=[]):
+    # Create Materialized Views if they don't exist
+    # Materialized Views send new inserted rows to the aggregate tables
+
+    verbose_logger.debug("Clickhouse: Creating Materialized Views")
+    if "daily_aggregated_spend_per_model_mv" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_per_model_mv")
+        client.command(
+            """
+            CREATE MATERIALIZED VIEW daily_aggregated_spend_per_model_mv
+            TO daily_aggregated_spend_per_model
+            AS
+            SELECT
+                toDate(startTime) as day,
+                sumState(spend) AS DailySpend,
+                model as model
+            FROM spend_logs
+            GROUP BY
+                day, model
+            """
+        )
+    if "daily_aggregated_spend_per_api_key_mv" not in table_names:
+        verbose_logger.debug(
+            "Clickhouse: Creating daily_aggregated_spend_per_api_key_mv"
+        )
+        client.command(
+            """
+            CREATE MATERIALIZED VIEW daily_aggregated_spend_per_api_key_mv
+            TO daily_aggregated_spend_per_api_key
+            AS
+            SELECT
+                toDate(startTime) as day,
+                sumState(spend) AS DailySpend,
+                api_key as api_key
+            FROM spend_logs
+            GROUP BY
+                day, api_key
+            """
+        )
+    if "daily_aggregated_spend_per_user_mv" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_per_user_mv")
+        client.command(
+            """
+            CREATE MATERIALIZED VIEW daily_aggregated_spend_per_user_mv
+            TO daily_aggregated_spend_per_user
+            AS
+            SELECT
+                toDate(startTime) as day,
+                sumState(spend) AS DailySpend,
+                user as user
+            FROM spend_logs
+            GROUP BY
+                day, user
+            """
+        )
+    if "daily_aggregated_spend_mv" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_mv")
+        client.command(
+            """
+            CREATE MATERIALIZED VIEW daily_aggregated_spend_mv
+            TO daily_aggregated_spend
+            AS
+            SELECT
+                toDate(startTime) as day,
+                sumState(spend) AS DailySpend
+            FROM spend_logs
+            GROUP BY
+                day
+            """
+        )
+
+
+def _create_clickhouse_aggregate_tables(client=None, table_names=[]):
+    # Basic Logging works without this - this is only used for low latency reporting apis
+    verbose_logger.debug("Clickhouse: Creating Aggregate Tables")
+
+    # Create Aggregeate Tables if they don't exist
+    if "daily_aggregated_spend_per_model" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_per_model")
+        client.command(
+            """
+            CREATE TABLE daily_aggregated_spend_per_model
+            (
+                `day` Date,
+                `DailySpend` AggregateFunction(sum, Float64),
+                `model` String
+            )
+            ENGINE = SummingMergeTree()
+            ORDER BY (day, model);
+            """
+        )
+    if "daily_aggregated_spend_per_api_key" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_per_api_key")
+        client.command(
+            """
+            CREATE TABLE daily_aggregated_spend_per_api_key
+            (
+                `day` Date,
+                `DailySpend` AggregateFunction(sum, Float64),
+                `api_key` String
+            )
+            ENGINE = SummingMergeTree()
+            ORDER BY (day, api_key);
+            """
+        )
+    if "daily_aggregated_spend_per_user" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend_per_user")
+        client.command(
+            """
+            CREATE TABLE daily_aggregated_spend_per_user
+            (
+                `day` Date,
+                `DailySpend` AggregateFunction(sum, Float64),
+                `user` String
+            )
+            ENGINE = SummingMergeTree()
+            ORDER BY (day, user);
+            """
+        )
+    if "daily_aggregated_spend" not in table_names:
+        verbose_logger.debug("Clickhouse: Creating daily_aggregated_spend")
+        client.command(
+            """
+            CREATE TABLE daily_aggregated_spend
+            (
+                `day` Date,
+                `DailySpend` AggregateFunction(sum, Float64),
+            )
+            ENGINE = SummingMergeTree()
+            ORDER BY (day);
+            """
+        )
+    return

litellm/integrations/clickhouse.py

@@ -86,6 +86,14 @@ def _start_clickhouse():
             response = client.query("DESCRIBE default.spend_logs")
             verbose_logger.debug(f"spend logs schema ={response.result_rows}")
 
+        # RUN Enterprise Clickhouse Setup
+        # TLDR: For Enterprise - we create views / aggregate tables for low latency reporting APIs
+        from litellm.proxy.enterprise.utils import _create_clickhouse_aggregate_tables
+        from litellm.proxy.enterprise.utils import _create_clickhouse_material_views
+
+        _create_clickhouse_aggregate_tables(client=client, table_names=table_names)
+        _create_clickhouse_material_views(client=client, table_names=table_names)
+
 
 class ClickhouseLogger:
     # Class variables or attributes

litellm/proxy/proxy_config.yaml

@@ -43,7 +43,7 @@ model_list:
       api_key: os.environ/OPENAI_API_KEY
 litellm_settings:
   fallbacks: [{"openai-gpt-3.5": ["azure-gpt-3.5"]}]
-  success_callback: ['langfuse']
+  success_callback: ['clickhouse', 'langfuse']
   # setting callback class
   # callbacks: custom_callbacks.proxy_handler_instance # sets litellm.callbacks = [proxy_handler_instance]