diff --git a/litellm/integrations/greenscale.py b/litellm/integrations/greenscale.py new file mode 100644 index 00000000000..3ff808ddbba --- /dev/null +++ b/litellm/integrations/greenscale.py @@ -0,0 +1,51 @@ +import requests +import json +import traceback +from datetime import datetime, timezone + +class GreenscaleLogger: + def __init__(self): + import os + self.greenscale_api_key = os.getenv("GREENSCALE_API_KEY") + self.headers = { + "api-key": self.greenscale_api_key, + "Content-Type": "application/json" + } + self.greenscale_logging_url = os.getenv("GREENSCALE_ENDPOINT") + + def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): + try: + response_json = response_obj.model_dump() if response_obj else {} + data = { + "modelId": kwargs.get("model"), + "inputTokenCount": response_json.get("usage", {}).get("prompt_tokens"), + "outputTokenCount": response_json.get("usage", {}).get("completion_tokens"), + } + data["timestamp"] = datetime.now(timezone.utc).strftime('%Y-%m-%dT%H:%M:%SZ') + + if type(end_time) == datetime and type(start_time) == datetime: + data["invocationLatency"] = int((end_time - start_time).total_seconds() * 1000) + + + # Add additional metadata keys to tags + tags = [] + metadata = kwargs.get("litellm_params", {}).get("metadata", {}) + for key, value in metadata.items(): + if key.startswith("greenscale"): + if key == "greenscale_project": + data["project"] = value + elif key == "greenscale_application": + data["application"] = value + else: + tags.append({"key": key.replace("greenscale_", ""), "value": str(value)}) + + data["tags"] = tags + + response = requests.post(self.greenscale_logging_url, headers=self.headers, data=json.dumps(data, default=str)) + if response.status_code != 200: + print_verbose(f"Greenscale Logger Error - {response.text}, {response.status_code}") + else: + print_verbose(f"Greenscale Logger Succeeded - {response.text}") + except Exception as e: + print_verbose(f"Greenscale Logger Error - {e}, Stack trace: {traceback.format_exc()}") + pass \ No newline at end of file diff --git a/litellm/utils.py b/litellm/utils.py index 17a31751d96..12ad0fc2b78 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -70,6 +70,7 @@ from .integrations.prometheus import PrometheusLogger from .integrations.dynamodb import DyanmoDBLogger from .integrations.s3 import S3Logger from .integrations.clickhouse import ClickhouseLogger +from .integrations.greenscale import GreenscaleLogger from .integrations.litedebugger import LiteDebugger from .proxy._types import KeyManagementSystem from openai import OpenAIError as OriginalError @@ -129,6 +130,7 @@ dynamoLogger = None s3Logger = None genericAPILogger = None clickHouseLogger = None +greenscaleLogger = None lunaryLogger = None aispendLogger = None berrispendLogger = None @@ -1717,6 +1719,33 @@ class Logging: user_id=kwargs.get("user", None), print_verbose=print_verbose, ) + if callback == "greenscale": + kwargs = {} + for k, v in self.model_call_details.items(): + if ( + k != "original_response" + ): # copy.deepcopy raises errors as this could be a coroutine + kwargs[k] = v + # this only logs streaming once, complete_streaming_response exists i.e when stream ends + if self.stream: + verbose_logger.debug( + f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}" + ) + if complete_streaming_response is None: + break + else: + print_verbose( + "reaches greenscale for streaming logging!" + ) + result = kwargs["complete_streaming_response"] + + greenscaleLogger.log_event( + kwargs=kwargs, + response_obj=result, + start_time=start_time, + end_time=end_time, + print_verbose=print_verbose, + ) if callback == "cache" and litellm.cache is not None: # this only logs streaming once, complete_streaming_response exists i.e when stream ends print_verbose("success_callback: reaches cache for logging!") @@ -6304,7 +6333,7 @@ def validate_environment(model: Optional[str] = None) -> dict: def set_callbacks(callback_list, function_id=None): - global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, langsmithLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger + global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, lunaryLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, langsmithLogger, dynamoLogger, s3Logger, dataDogLogger, prometheusLogger, greenscaleLogger try: for callback in callback_list: @@ -6391,6 +6420,9 @@ def set_callbacks(callback_list, function_id=None): elif callback == "supabase": print_verbose(f"instantiating supabase") supabaseClient = Supabase() + elif callback == "greenscale": + greenscaleLogger = GreenscaleLogger() + print_verbose("Initialized Greenscale Logger") elif callback == "lite_debugger": print_verbose(f"instantiating lite_debugger") if function_id: