diff --git a/docs/my-website/docs/proxy/prompt_management.md b/docs/my-website/docs/proxy/prompt_management.md index 693477ed740..980043f4555 100644 --- a/docs/my-website/docs/proxy/prompt_management.md +++ b/docs/my-website/docs/proxy/prompt_management.md @@ -2,7 +2,13 @@ import Image from '@theme/IdealImage'; import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# Prompt Management +# [BETA] Prompt Management + +:::info + +This feature is currently in beta, and might change unexpectedly. We expect this to be more stable by next month (February 2025). + +::: Run experiments or change the specific model (e.g. from gpt-4o to gpt4o-mini finetune) from your prompt management tool (e.g. Langfuse) instead of making changes in the application. diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index 483fd3333e2..c4434969428 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -179,6 +179,7 @@ class LangFuseLogger: optional_params = copy.deepcopy(kwargs.get("optional_params", {})) prompt = {"messages": kwargs.get("messages")} + functions = optional_params.pop("functions", None) tools = optional_params.pop("tools", None) if functions is not None: @@ -462,15 +463,27 @@ class LangFuseLogger: if standard_logging_object is None: end_user_id = None + prompt_management_metadata: Optional[dict] = None else: end_user_id = standard_logging_object["metadata"].get( "user_api_key_end_user_id", None ) + prompt_management_metadata = cast( + Optional[dict], + standard_logging_object["metadata"].get( + "prompt_management_metadata", None + ), + ) + # Clean Metadata before logging - never log raw metadata # the raw metadata can contain circular references which leads to infinite recursion # we clean out all extra litellm metadata params before logging - clean_metadata = {} + clean_metadata: Dict[str, Any] = {} + if prompt_management_metadata is not None: + clean_metadata["prompt_management_metadata"] = ( + prompt_management_metadata + ) if isinstance(metadata, dict): for key, value in metadata.items(): # generate langfuse tags - Default Tags sent to Langfuse from LiteLLM Proxy @@ -498,10 +511,10 @@ class LangFuseLogger: ) session_id = clean_metadata.pop("session_id", None) - trace_name = clean_metadata.pop("trace_name", None) + trace_name = cast(Optional[str], clean_metadata.pop("trace_name", None)) trace_id = clean_metadata.pop("trace_id", litellm_call_id) existing_trace_id = clean_metadata.pop("existing_trace_id", None) - update_trace_keys = clean_metadata.pop("update_trace_keys", []) + update_trace_keys = cast(list, clean_metadata.pop("update_trace_keys", [])) debug = clean_metadata.pop("debug_langfuse", None) mask_input = clean_metadata.pop("mask_input", False) mask_output = clean_metadata.pop("mask_output", False) @@ -514,7 +527,7 @@ class LangFuseLogger: trace_name = f"litellm-{kwargs.get('call_type', 'completion')}" if existing_trace_id is not None: - trace_params = {"id": existing_trace_id} + trace_params: Dict[str, Any] = {"id": existing_trace_id} # Update the following keys for this trace for metadata_param_key in update_trace_keys: @@ -656,8 +669,12 @@ class LangFuseLogger: # if `generation_name` is None, use sensible default values # If using litellm proxy user `key_alias` if not None # If `key_alias` is None, just log `litellm-{call_type}` as the generation name - _user_api_key_alias = clean_metadata.get("user_api_key_alias", None) - generation_name = f"litellm-{kwargs.get('call_type', 'completion')}" + _user_api_key_alias = cast( + Optional[str], clean_metadata.get("user_api_key_alias", None) + ) + generation_name = ( + f"litellm-{cast(str, kwargs.get('call_type', 'completion'))}" + ) if _user_api_key_alias is not None: generation_name = f"litellm:{_user_api_key_alias}" diff --git a/litellm/integrations/langfuse/langfuse_prompt_management.py b/litellm/integrations/langfuse/langfuse_prompt_management.py index 6d13045597f..2bb8d55e6a3 100644 --- a/litellm/integrations/langfuse/langfuse_prompt_management.py +++ b/litellm/integrations/langfuse/langfuse_prompt_management.py @@ -10,12 +10,14 @@ from packaging.version import Version from typing_extensions import TypeAlias from litellm.integrations.custom_logger import CustomLogger -from litellm.types.llms.openai import AllMessageValues +from litellm.integrations.prompt_management_base import PromptManagementClient +from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPayload from ...litellm_core_utils.specialty_caches.dynamic_logging_cache import ( DynamicLoggingCache, ) +from ..prompt_management_base import PromptManagementBase from .langfuse import LangFuseLogger from .langfuse_handler import LangFuseHandler @@ -97,7 +99,7 @@ def langfuse_client_init( return client -class LangfusePromptManagement(LangFuseLogger, CustomLogger): +class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger): def __init__( self, langfuse_public_key=None, @@ -112,6 +114,10 @@ class LangfusePromptManagement(LangFuseLogger, CustomLogger): flush_interval=flush_interval, ) + @property + def integration_name(self): + return "langfuse" + def _get_prompt_from_id( self, langfuse_prompt_id: str, langfuse_client: LangfuseClass ) -> PROMPT_CLIENT: @@ -122,7 +128,7 @@ class LangfusePromptManagement(LangFuseLogger, CustomLogger): langfuse_prompt_client: PROMPT_CLIENT, langfuse_prompt_variables: Optional[dict], call_type: Union[Literal["completion"], Literal["text_completion"]], - ) -> Optional[Union[str, list]]: + ) -> List[AllMessageValues]: compiled_prompt: Optional[Union[str, list]] = None if langfuse_prompt_variables is None: @@ -130,16 +136,14 @@ class LangfusePromptManagement(LangFuseLogger, CustomLogger): compiled_prompt = langfuse_prompt_client.compile(**langfuse_prompt_variables) - return compiled_prompt - - def _get_model_from_prompt( - self, langfuse_prompt_client: PROMPT_CLIENT, model: str - ) -> str: - config = langfuse_prompt_client.config - if "model" in config: - return config["model"] + if isinstance(compiled_prompt, str): + compiled_prompt = [ + ChatCompletionSystemMessage(role="system", content=compiled_prompt) + ] else: - return model.replace("langfuse/", "") + compiled_prompt = cast(List[AllMessageValues], compiled_prompt) + + return compiled_prompt def _get_optional_params_from_langfuse( self, langfuse_prompt_client: PROMPT_CLIENT @@ -173,23 +177,27 @@ class LangfusePromptManagement(LangFuseLogger, CustomLogger): dynamic_callback_params, ) - def get_chat_completion_prompt( + def should_run_prompt_management( + self, + prompt_id: str, + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> bool: + langfuse_client = langfuse_client_init( + langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"), + langfuse_secret=dynamic_callback_params.get("langfuse_secret"), + langfuse_host=dynamic_callback_params.get("langfuse_host"), + ) + langfuse_prompt_client = self._get_prompt_from_id( + langfuse_prompt_id=prompt_id, langfuse_client=langfuse_client + ) + return langfuse_prompt_client is not None + + def _compile_prompt_helper( self, - model: str, - messages: List[AllMessageValues], - non_default_params: dict, prompt_id: str, prompt_variables: Optional[dict], dynamic_callback_params: StandardCallbackDynamicParams, - ) -> Tuple[ - str, - List[AllMessageValues], - dict, - ]: - if prompt_id is None: - raise ValueError( - "Langfuse prompt id is required. Pass in as parameter 'langfuse_prompt_id'" - ) + ) -> PromptManagementClient: langfuse_client = langfuse_client_init( langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"), langfuse_secret=dynamic_callback_params.get("langfuse_secret"), @@ -206,25 +214,19 @@ class LangfusePromptManagement(LangFuseLogger, CustomLogger): call_type="completion", ) - if compiled_prompt is None: - raise ValueError(f"Langfuse prompt not found. Prompt id={prompt_id}") - if isinstance(compiled_prompt, list): - messages = compiled_prompt - elif isinstance(compiled_prompt, str): - messages = [{"role": "user", "content": compiled_prompt}] - else: - raise ValueError( - f"Langfuse prompt is not a list or string. Prompt id={prompt_id}, compiled_prompt type={type(compiled_prompt)}" - ) + template_model = langfuse_prompt_client.config.get("model") - ## SET MODEL - model = self._get_model_from_prompt(langfuse_prompt_client, model) - - optional_params = self._get_optional_params_from_langfuse( + template_optional_params = self._get_optional_params_from_langfuse( langfuse_prompt_client ) - return model, messages, optional_params + return PromptManagementClient( + prompt_id=prompt_id, + prompt_template=compiled_prompt, + prompt_template_model=template_model, + prompt_template_optional_params=template_optional_params, + completed_messages=None, + ) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): standard_callback_dynamic_params = kwargs.get( diff --git a/litellm/integrations/prompt_management_base.py b/litellm/integrations/prompt_management_base.py new file mode 100644 index 00000000000..3fe3b31ed8b --- /dev/null +++ b/litellm/integrations/prompt_management_base.py @@ -0,0 +1,118 @@ +from abc import ABC, abstractmethod +from typing import Any, Dict, List, Optional, Tuple, TypedDict + +from litellm.types.llms.openai import AllMessageValues +from litellm.types.utils import StandardCallbackDynamicParams + + +class PromptManagementClient(TypedDict): + prompt_id: str + prompt_template: List[AllMessageValues] + prompt_template_model: Optional[str] + prompt_template_optional_params: Optional[Dict[str, Any]] + completed_messages: Optional[List[AllMessageValues]] + + +class PromptManagementBase(ABC): + + @property + @abstractmethod + def integration_name(self) -> str: + pass + + @abstractmethod + def should_run_prompt_management( + self, + prompt_id: str, + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> bool: + pass + + @abstractmethod + def _compile_prompt_helper( + self, + prompt_id: str, + prompt_variables: Optional[dict], + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> PromptManagementClient: + pass + + def merge_messages( + self, + prompt_template: List[AllMessageValues], + client_messages: List[AllMessageValues], + ) -> List[AllMessageValues]: + return prompt_template + client_messages + + def compile_prompt( + self, + prompt_id: str, + prompt_variables: Optional[dict], + client_messages: List[AllMessageValues], + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> PromptManagementClient: + compiled_prompt_client = self._compile_prompt_helper( + prompt_id=prompt_id, + prompt_variables=prompt_variables, + dynamic_callback_params=dynamic_callback_params, + ) + + try: + messages = compiled_prompt_client["prompt_template"] + client_messages + except Exception as e: + raise ValueError( + f"Error compiling prompt: {e}. Prompt id={prompt_id}, prompt_variables={prompt_variables}, client_messages={client_messages}, dynamic_callback_params={dynamic_callback_params}" + ) + + compiled_prompt_client["completed_messages"] = messages + return compiled_prompt_client + + def _get_model_from_prompt( + self, prompt_management_client: PromptManagementClient, model: str + ) -> str: + if prompt_management_client["prompt_template_model"] is not None: + return prompt_management_client["prompt_template_model"] + else: + return model.replace("{}/".format(self.integration_name), "") + + def get_chat_completion_prompt( + self, + model: str, + messages: List[AllMessageValues], + non_default_params: dict, + prompt_id: str, + prompt_variables: Optional[dict], + dynamic_callback_params: StandardCallbackDynamicParams, + ) -> Tuple[ + str, + List[AllMessageValues], + dict, + ]: + if not self.should_run_prompt_management( + prompt_id=prompt_id, dynamic_callback_params=dynamic_callback_params + ): + return model, messages, non_default_params + + prompt_template = self.compile_prompt( + prompt_id=prompt_id, + prompt_variables=prompt_variables, + client_messages=messages, + dynamic_callback_params=dynamic_callback_params, + ) + + completed_messages = prompt_template["completed_messages"] or messages + + prompt_template_optional_params = ( + prompt_template["prompt_template_optional_params"] or {} + ) + + updated_non_default_params = { + **non_default_params, + **prompt_template_optional_params, + } + + model = self._get_model_from_prompt( + prompt_management_client=prompt_template, model=model + ) + + return model, completed_messages, updated_non_default_params diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 3faad6fbbec..f066e95b1f3 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -59,6 +59,7 @@ from litellm.types.utils import ( StandardLoggingPayload, StandardLoggingPayloadErrorInformation, StandardLoggingPayloadStatus, + StandardLoggingPromptManagementMetadata, TextCompletionResponse, TranscriptionResponse, Usage, @@ -424,6 +425,7 @@ class Logging(LiteLLMLoggingBaseClass): dynamic_callback_params=self.standard_callback_dynamic_params, ) ) + self.messages = messages return model, messages, non_default_params @@ -431,6 +433,7 @@ class Logging(LiteLLMLoggingBaseClass): """ Common helper function across the sync + async pre-call function """ + self.model_call_details["input"] = input self.model_call_details["api_key"] = api_key self.model_call_details["additional_args"] = additional_args @@ -2628,7 +2631,7 @@ class StandardLoggingPayloadSetup: @staticmethod def get_standard_logging_metadata( - metadata: Optional[Dict[str, Any]] + metadata: Optional[Dict[str, Any]], litellm_params: Optional[dict] = None ) -> StandardLoggingMetadata: """ Clean and filter the metadata dictionary to include only the specified keys in StandardLoggingMetadata. @@ -2643,6 +2646,20 @@ class StandardLoggingPayloadSetup: - If the input metadata is None or not a dictionary, an empty StandardLoggingMetadata object is returned. - If 'user_api_key' is present in metadata and is a valid SHA256 hash, it's stored as 'user_api_key_hash'. """ + prompt_management_metadata: Optional[ + StandardLoggingPromptManagementMetadata + ] = None + if litellm_params is not None: + prompt_id = cast(Optional[str], litellm_params.get("prompt_id", None)) + prompt_variables = cast( + Optional[dict], litellm_params.get("prompt_variables", None) + ) + + prompt_management_metadata = StandardLoggingPromptManagementMetadata( + prompt_id=prompt_id, + prompt_variables=prompt_variables, + ) + # Initialize with default values clean_metadata = StandardLoggingMetadata( user_api_key_hash=None, @@ -2655,6 +2672,7 @@ class StandardLoggingPayloadSetup: requester_ip_address=None, requester_metadata=None, user_api_key_end_user_id=None, + prompt_management_metadata=prompt_management_metadata, ) if isinstance(metadata, dict): # Filter the metadata dictionary to include only the specified keys @@ -2949,7 +2967,7 @@ def get_standard_logging_object_payload( ) # clean up litellm metadata clean_metadata = StandardLoggingPayloadSetup.get_standard_logging_metadata( - metadata=metadata + metadata=metadata, litellm_params=litellm_params ) saved_cache_cost: float = 0.0 @@ -2966,6 +2984,7 @@ def get_standard_logging_object_payload( ## Get model cost information ## base_model = _get_base_model_from_metadata(model_call_details=kwargs) custom_pricing = use_custom_pricing_for_model(litellm_params=litellm_params) + model_cost_information = StandardLoggingPayloadSetup.get_model_cost_information( base_model=base_model, custom_pricing=custom_pricing, @@ -3072,6 +3091,7 @@ def get_standard_logging_metadata( requester_ip_address=None, requester_metadata=None, user_api_key_end_user_id=None, + prompt_management_metadata=None, ) if isinstance(metadata, dict): # Filter the metadata dictionary to include only the specified keys diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index eb13dbb8b05..8109c8bf612 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -82,7 +82,7 @@ def _process_gemini_image(image_url: str) -> PartType: ): file_data = FileDataType(file_uri=image_url, mime_type=image_type) return PartType(file_data=file_data) - elif "https://" in image_url or "base64" in image_url: + elif "http://" in image_url or "https://" in image_url or "base64" in image_url: # https links for unsupported mime types and base64 images image = convert_to_anthropic_image_obj(image_url) _blob = BlobType(data=image["data"], mime_type=image["media_type"]) diff --git a/litellm/main.py b/litellm/main.py index c97188b4522..082244091cb 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -1077,6 +1077,8 @@ def completion( # type: ignore # noqa: PLR0915 litellm_metadata=kwargs.get("litellm_metadata"), disable_add_transform_inline_image_block=disable_add_transform_inline_image_block, drop_params=kwargs.get("drop_params"), + prompt_id=prompt_id, + prompt_variables=prompt_variables, ) logging.update_environment_variables( model=model, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index c9762c36e7e..6bf00fb16af 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -10,9 +10,8 @@ model_list: api_key: os.environ/OPENAI_API_KEY - model_name: chatbot_actions litellm_params: - model: langfuse/azure/gpt-4o - api_base: "os.environ/AZURE_API_BASE" - api_key: "os.environ/AZURE_API_KEY" + model: langfuse/gpt-3.5-turbo + api_key: os.environ/OPENAI_API_KEY tpm: 1000000 prompt_id: "jokes" - model_name: openai-deepseek diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 8ae59e38c54..7dc3ae1be2f 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2565,7 +2565,6 @@ class ProxyConfig: for response in responses: if response is not None: param_name = getattr(response, "param_name", None) - verbose_proxy_logger.info(f"loading {param_name} settings from db") if param_name == "litellm_settings": verbose_proxy_logger.info( f"litellm_settings: {response.param_value}" diff --git a/litellm/router.py b/litellm/router.py index bfa7708854e..9b0d34f92c8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -47,9 +47,6 @@ from litellm.caching.caching import DualCache, InMemoryCache, RedisCache from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging -from litellm.litellm_core_utils.litellm_logging import ( - _init_custom_logger_compatible_class, -) from litellm.router_strategy.budget_limiter import RouterBudgetLimiting from litellm.router_strategy.least_busy import LeastBusyLoggingHandler from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler @@ -120,6 +117,8 @@ from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, ModelResponse, + Rules, + function_setup, get_llm_provider, get_non_default_completion_params, get_secret, @@ -1457,6 +1456,17 @@ class Router: messages: List[AllMessageValues], kwargs: Dict[str, Any], ): + litellm_logging_object = kwargs.get("litellm_logging_obj", None) + if litellm_logging_object is None: + litellm_logging_object, kwargs = function_setup( + **{ + "original_function": "acompletion", + "rules_obj": Rules(), + "start_time": get_utc_datetime(), + **kwargs, + } + ) + litellm_logging_object = cast(LiteLLMLogging, litellm_logging_object) prompt_management_deployment = self.get_available_deployment( model=model, messages=[{"role": "user", "content": "prompt"}], @@ -1475,38 +1485,31 @@ class Router: "prompt_variables", None ) - if litellm_model is None or "/" not in litellm_model: + if prompt_id is None or not isinstance(prompt_id, str): raise ValueError( - f"Model is not a custom logger compatible callback. Got={litellm_model}" + f"Prompt ID is not set or not a string. Got={prompt_id}, type={type(prompt_id)}" + ) + if prompt_variables is not None and not isinstance(prompt_variables, dict): + raise ValueError( + f"Prompt variables is set but not a dictionary. Got={prompt_variables}, type={type(prompt_variables)}" ) - custom_logger_compatible_callback = litellm_model.split("/", 1)[0] - split_litellm_model = litellm_model.split("/", 1)[1] - - custom_logger = _init_custom_logger_compatible_class( - logging_integration=custom_logger_compatible_callback, - internal_usage_cache=None, - llm_router=None, - ) - - if custom_logger is None: - raise ValueError( - f"Custom logger is not initialized. Got={custom_logger_compatible_callback}" - ) model, messages, optional_params = ( - await custom_logger.async_get_chat_completion_prompt( - model=split_litellm_model, + litellm_logging_object.get_chat_completion_prompt( + model=litellm_model, messages=messages, non_default_params=get_non_default_completion_params(kwargs=kwargs), prompt_id=prompt_id, prompt_variables=prompt_variables, - dynamic_callback_params={}, ) ) kwargs = {**kwargs, **optional_params} kwargs["model"] = model kwargs["messages"] = messages + kwargs["litellm_logging_obj"] = litellm_logging_object + kwargs["prompt_id"] = prompt_id + kwargs["prompt_variables"] = prompt_variables _model_list = self.get_model_list(model_name=model) if _model_list is None or len(_model_list) == 0: # if direct call to model diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 58719687cd6..39a329f8893 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1457,9 +1457,14 @@ class StandardLoggingUserAPIKeyMetadata(TypedDict): user_api_key_end_user_id: Optional[str] +class StandardLoggingPromptManagementMetadata(TypedDict): + prompt_id: Optional[str] + prompt_variables: Optional[dict] + + class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): """ - Specific metadata k,v pairs logged to integration for easier cost tracking + Specific metadata k,v pairs logged to integration for easier cost tracking and prompt management """ spend_logs_metadata: Optional[ @@ -1467,6 +1472,7 @@ class StandardLoggingMetadata(StandardLoggingUserAPIKeyMetadata): ] # special param to log k,v pairs to spendlogs for a call requester_ip_address: Optional[str] requester_metadata: Optional[dict] + prompt_management_metadata: Optional[StandardLoggingPromptManagementMetadata] class StandardLoggingAdditionalHeaders(TypedDict, total=False): diff --git a/litellm/utils.py b/litellm/utils.py index 6051a873272..5aae94e4f03 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2034,6 +2034,8 @@ def get_litellm_params( litellm_metadata: Optional[dict] = None, disable_add_transform_inline_image_block: Optional[bool] = None, drop_params: Optional[bool] = None, + prompt_id: Optional[str] = None, + prompt_variables: Optional[dict] = None, ): litellm_params = { "acompletion": acompletion, @@ -2068,6 +2070,8 @@ def get_litellm_params( "litellm_metadata": litellm_metadata, "disable_add_transform_inline_image_block": disable_add_transform_inline_image_block, "drop_params": drop_params, + "prompt_id": prompt_id, + "prompt_variables": prompt_variables, } return litellm_params diff --git a/tests/llm_translation/test_vertex.py b/tests/llm_translation/test_vertex.py index ee820e99e14..15d2df71519 100644 --- a/tests/llm_translation/test_vertex.py +++ b/tests/llm_translation/test_vertex.py @@ -5,6 +5,10 @@ import traceback from dotenv import load_dotenv +import litellm.litellm_core_utils +import litellm.litellm_core_utils.prompt_templates +import litellm.litellm_core_utils.prompt_templates.factory + load_dotenv() import io from unittest.mock import AsyncMock, MagicMock, patch @@ -16,6 +20,8 @@ import pytest import litellm from litellm import get_optional_params from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_image +from litellm.types.llms.vertex_ai import PartType, BlobType import httpx @@ -1240,3 +1246,53 @@ def test_vertex_embedding_url(model, expected_url): assert url == expected_url assert endpoint == "predict" + + +import pytest +from unittest.mock import Mock, patch +from typing import Dict, Any + +# Import your actual module here +# from your_module import _process_gemini_image, PartType, FileDataType, BlobType + + +@pytest.fixture +def mock_convert_url_to_base64(): + with patch( + "litellm.litellm_core_utils.prompt_templates.factory.convert_url_to_base64", + ) as mock: + # Setup the mock to return a valid image object + mock.return_value = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." + yield mock + + +@pytest.fixture +def mock_blob(): + return Mock(spec=BlobType) + + +@pytest.mark.parametrize( + "http_url", + [ + "http://img1.etsystatic.com/260/0/7813604/il_fullxfull.4226713999_q86e.jpg", + "http://example.com/image.jpg", + "http://subdomain.domain.com/path/to/image.png", + ], +) +def test_process_gemini_image_http_url( + http_url: str, mock_convert_url_to_base64: Mock, mock_blob: Mock +) -> None: + """ + Test that _process_gemini_image correctly handles HTTP URLs. + + Args: + http_url: Test HTTP URL + mock_convert_to_anthropic: Mocked convert_to_anthropic_image_obj function + mock_blob: Mocked BlobType instance + """ + # Arrange + expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." + mock_convert_url_to_base64.return_value = expected_image_data + + # Act + result = _process_gemini_image(http_url) diff --git a/tests/local_testing/test_utils.py b/tests/local_testing/test_utils.py index 768c20834e4..2d19d5eccc5 100644 --- a/tests/local_testing/test_utils.py +++ b/tests/local_testing/test_utils.py @@ -500,6 +500,37 @@ def test_get_supported_openai_params() -> None: assert get_supported_openai_params("nonexistent") is None +def test_get_chat_completion_prompt(): + """ + Unit test to ensure get_chat_completion_prompt updates messages in logging object. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + + litellm_logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + litellm_call_id="1234", + start_time=datetime.now(), + function_id="1234", + ) + + updated_message = "hello world" + + litellm_logging_obj.get_chat_completion_prompt( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": updated_message}], + non_default_params={}, + prompt_id="1234", + prompt_variables=None, + ) + + assert litellm_logging_obj.messages == [ + {"role": "user", "content": updated_message} + ] + + def test_redact_msgs_from_logs(): """ Tests that turn_off_message_logging does not modify the response_obj diff --git a/tests/logging_callback_tests/test_otel_logging.py b/tests/logging_callback_tests/test_otel_logging.py index a9a11cfba21..9c19c9d261a 100644 --- a/tests/logging_callback_tests/test_otel_logging.py +++ b/tests/logging_callback_tests/test_otel_logging.py @@ -259,6 +259,7 @@ def validate_redacted_message_span_attributes(span): "gen_ai.response.id", "gen_ai.response.model", "llm.usage.total_tokens", + "metadata.prompt_management_metadata", "gen_ai.usage.completion_tokens", "gen_ai.usage.prompt_tokens", "metadata.user_api_key_hash",