Support passing prompt_label to langfuse (#11018)

* fix: add prompt label support to prompt management hook

* feat: support 'prompt_label' parameter for langfuse prompt management

Closes https://github.com/BerriAI/litellm/discussions/9003#discussioncomment-13221555

* fix(litellm_logging.py): deep copy optional params to avoid mutation while logging

* fix(log-consistent-optional-param-values-across-providers): ensures params can be used for finetuning from providers

* fix: fix linting error

* test: update test

* test: update langfuse tests

* fix(litellm_logging.py): avoid deepcopying optional params

might contain thread object
This commit is contained in:
Krish Dholakia 2025-05-21 22:27:36 -07:00 committed by GitHub
parent cd496fee2e
commit 2b50b43ae2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
31 changed files with 116 additions and 87 deletions

View File

@ -28,6 +28,7 @@ class AnthropicCacheControlHook(CustomPromptManagement):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Apply cache control directives based on specified injection points.
@ -79,10 +80,10 @@ class AnthropicCacheControlHook(CustomPromptManagement):
# Case 1: Target by specific index
if targetted_index is not None:
if 0 <= targetted_index < len(messages):
messages[targetted_index] = (
AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
messages[
targetted_index
] = AnthropicCacheControlHook._safe_insert_cache_control_in_message(
messages[targetted_index], control
)
# Case 2: Target by role
elif targetted_role is not None:

View File

@ -87,6 +87,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Returns:
@ -104,6 +105,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Returns:

View File

@ -18,6 +18,7 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Returns:
@ -43,6 +44,7 @@ class CustomPromptManagement(CustomLogger, PromptManagementBase):
prompt_id: str,
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> PromptManagementClient:
raise NotImplementedError(
"Custom prompt management does not support compile prompt helper"

View File

@ -155,11 +155,8 @@ class HumanloopLogger(CustomLogger):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
) -> Tuple[
str,
List[AllMessageValues],
dict,
]:
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict,]:
humanloop_api_key = dynamic_callback_params.get(
"humanloop_api_key"
) or get_secret_str("HUMANLOOP_API_KEY")

View File

@ -130,9 +130,12 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
return "langfuse"
def _get_prompt_from_id(
self, langfuse_prompt_id: str, langfuse_client: LangfuseClass
self,
langfuse_prompt_id: str,
langfuse_client: LangfuseClass,
prompt_label: Optional[str] = None,
) -> PROMPT_CLIENT:
return langfuse_client.get_prompt(langfuse_prompt_id)
return langfuse_client.get_prompt(langfuse_prompt_id, label=prompt_label)
def _compile_prompt(
self,
@ -176,11 +179,8 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
) -> Tuple[
str,
List[AllMessageValues],
dict,
]:
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict,]:
return self.get_chat_completion_prompt(
model,
messages,
@ -188,6 +188,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
prompt_id,
prompt_variables,
dynamic_callback_params,
prompt_label=prompt_label,
)
def should_run_prompt_management(
@ -211,6 +212,7 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
prompt_id: str,
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> PromptManagementClient:
langfuse_client = langfuse_client_init(
langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"),
@ -219,7 +221,9 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
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
langfuse_prompt_id=prompt_id,
langfuse_client=langfuse_client,
prompt_label=prompt_label,
)
## SET PROMPT

View File

@ -33,6 +33,7 @@ class PromptManagementBase(ABC):
prompt_id: str,
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> PromptManagementClient:
pass
@ -49,11 +50,13 @@ class PromptManagementBase(ABC):
prompt_variables: Optional[dict],
client_messages: List[AllMessageValues],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> PromptManagementClient:
compiled_prompt_client = self._compile_prompt_helper(
prompt_id=prompt_id,
prompt_variables=prompt_variables,
dynamic_callback_params=dynamic_callback_params,
prompt_label=prompt_label,
)
try:
@ -82,6 +85,7 @@ class PromptManagementBase(ABC):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
if prompt_id is None:
raise ValueError("prompt_id is required for Prompt Management Base class")
@ -95,6 +99,7 @@ class PromptManagementBase(ABC):
prompt_variables=prompt_variables,
client_messages=messages,
dynamic_callback_params=dynamic_callback_params,
prompt_label=prompt_label,
)
completed_messages = prompt_template["completed_messages"] or messages

View File

@ -75,6 +75,7 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
dynamic_callback_params: StandardCallbackDynamicParams,
litellm_logging_obj: LiteLLMLoggingObj,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Retrieves the context from the Bedrock Knowledge Base and appends it to the messages.
@ -99,10 +100,11 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
f"Bedrock Knowledge Base Response: {bedrock_kb_response}"
)
context_message, context_string = (
self.get_chat_completion_message_from_bedrock_kb_response(
bedrock_kb_response
)
(
context_message,
context_string,
) = self.get_chat_completion_message_from_bedrock_kb_response(
bedrock_kb_response
)
if context_message is not None:
messages.append(context_message)
@ -126,9 +128,9 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
)
)
litellm_logging_obj.model_call_details["vector_store_request_metadata"] = (
vector_store_request_metadata
)
litellm_logging_obj.model_call_details[
"vector_store_request_metadata"
] = vector_store_request_metadata
return model, messages, non_default_params
@ -140,9 +142,9 @@ class BedrockVectorStore(BaseVectorStore, BaseAWSLLM):
"""
Transform a BedrockKBResponse to a VectorStoreSearchResponse
"""
retrieval_results: Optional[List[BedrockKBRetrievalResult]] = (
bedrock_kb_response.get("retrievalResults", None)
)
retrieval_results: Optional[
List[BedrockKBRetrievalResult]
] = bedrock_kb_response.get("retrievalResults", None)
vector_store_search_response: VectorStoreSearchResponse = (
VectorStoreSearchResponse(search_query=query, data=[])
)

View File

@ -539,6 +539,7 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
prompt_management_logger: Optional[CustomLogger] = None,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
custom_logger = (
prompt_management_logger
@ -559,6 +560,7 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_id=prompt_id,
prompt_variables=prompt_variables,
dynamic_callback_params=self.standard_callback_dynamic_params,
prompt_label=prompt_label,
)
self.messages = messages
return model, messages, non_default_params
@ -572,6 +574,7 @@ class Logging(LiteLLMLoggingBaseClass):
prompt_variables: Optional[dict],
prompt_management_logger: Optional[CustomLogger] = None,
tools: Optional[List[Dict]] = None,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
custom_logger = (
prompt_management_logger
@ -594,6 +597,7 @@ class Logging(LiteLLMLoggingBaseClass):
dynamic_callback_params=self.standard_callback_dynamic_params,
litellm_logging_obj=self,
tools=tools,
prompt_label=prompt_label,
)
self.messages = messages
return model, messages, non_default_params

View File

@ -97,6 +97,7 @@ from litellm.utils import (
get_optional_params_image_gen,
get_optional_params_transcription,
get_secret,
get_standard_openai_params,
mock_completion_streaming_obj,
read_config_args,
supports_httpx_timeout,
@ -428,6 +429,7 @@ async def acompletion(
prompt_id=kwargs.get("prompt_id", None),
prompt_variables=kwargs.get("prompt_variables", None),
tools=tools,
prompt_label=kwargs.get("prompt_label", None),
)
#########################################################
@ -983,6 +985,7 @@ def completion( # type: ignore # noqa: PLR0915
assistant_continue_message=assistant_continue_message,
)
######## end of unpacking kwargs ###########
standard_openai_params = get_standard_openai_params(params=args)
non_default_params = get_non_default_completion_params(kwargs=kwargs)
litellm_params = {} # used to prevent unbound var errors
## PROMPT MANAGEMENT HOOKS ##
@ -1001,6 +1004,7 @@ def completion( # type: ignore # noqa: PLR0915
non_default_params=non_default_params,
prompt_id=prompt_id,
prompt_variables=prompt_variables,
prompt_label=kwargs.get("prompt_label", None),
)
try:
@ -1234,10 +1238,13 @@ def completion( # type: ignore # noqa: PLR0915
max_retries=max_retries,
timeout=timeout,
)
logging.update_environment_variables(
cast(LiteLLMLoggingObj, logging).update_environment_variables(
model=model,
user=user,
optional_params=optional_params,
optional_params={
**standard_openai_params,
**non_default_params,
}, # [IMPORTANT] - using standard_openai_params ensures consistent params logged to langfuse for finetuning / eval datasets.
litellm_params=litellm_params,
custom_llm_provider=custom_llm_provider,
)

View File

@ -1,8 +1,8 @@
model_list:
- model_name: "gemini-2.0-flash"
- model_name: "gemini-2.0-flash-gemini"
litellm_params:
model: gemini/gemini-2.0-flash-live-001
- model_name: "gpt-4.1-openai"
model: gemini/gemini-2.0-flash
- model_name: "gpt-4o-mini-openai"
litellm_params:
model: gpt-4.1-mini-2025-04-14
api_key: os.environ/OPENAI_API_KEY
@ -71,6 +71,16 @@ model_list:
model: mistral/*
api_key: os.environ/MISTRAL_API_KEY
access_groups: ["beta-models"]
- model_name: my-langfuse-model
litellm_params:
model: langfuse/gpt-3.5-turbo
prompt_id: "jokes"
prompt_label: "latest"
api_key: os.environ/OPENAI_API_KEY
litellm_settings:
cache: true
callbacks: ["langfuse"]
general_settings:
store_model_in_db: true
store_prompts_in_spend_logs: true

View File

@ -670,15 +670,20 @@ class UserObjectCache:
- update user object in cache
"""
if isinstance(user_object, LiteLLM_UserTable):
user_object = user_object.model_dump()
for k, v in user_object.items():
if isinstance(v, datetime):
user_object[k] = v.isoformat()
await self.user_api_key_cache.async_set_cache(key=user_id, value=user_object)
user_object_dict = user_object.model_dump()
else:
user_object_dict = user_object
for k, v in user_object_dict.items():
if isinstance(v, datetime):
user_object_dict[k] = v.isoformat()
await self.user_api_key_cache.async_set_cache(
key=user_id, value=user_object_dict
)
if self.internal_usage_cache is not None:
await self.internal_usage_cache.async_set_cache(
key=user_id,
value=user_object,
value=user_object_dict,
litellm_parent_otel_span=litellm_parent_otel_span,
)

View File

@ -15,6 +15,7 @@ class X42PromptManagement(CustomPromptManagement):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str] = None,
) -> Tuple[str, List[AllMessageValues], dict]:
"""
Returns:

View File

@ -1700,9 +1700,13 @@ class Router:
specific_deployment=kwargs.pop("specific_deployment", None),
)
litellm_model = prompt_management_deployment["litellm_params"].get(
"model", None
self._update_kwargs_with_deployment(
deployment=prompt_management_deployment, kwargs=kwargs
)
data = prompt_management_deployment["litellm_params"].copy()
litellm_model = data.get("model", None)
prompt_id = kwargs.get("prompt_id") or prompt_management_deployment[
"litellm_params"
].get("prompt_id", None)
@ -1711,6 +1715,9 @@ class Router:
) or prompt_management_deployment["litellm_params"].get(
"prompt_variables", None
)
prompt_label = kwargs.get("prompt_label", None) or prompt_management_deployment[
"litellm_params"
].get("prompt_label", None)
if prompt_id is None or not isinstance(prompt_id, str):
raise ValueError(
@ -1731,14 +1738,16 @@ class Router:
non_default_params=get_non_default_completion_params(kwargs=kwargs),
prompt_id=prompt_id,
prompt_variables=prompt_variables,
prompt_label=prompt_label,
)
kwargs = {**kwargs, **optional_params}
kwargs = {**data, **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
kwargs["prompt_label"] = prompt_label
_model_list = self.get_model_list(model_name=model)
if _model_list is None or len(_model_list) == 0: # if direct call to model

View File

@ -2084,6 +2084,7 @@ all_litellm_params = [
"allowed_openai_params",
"litellm_session_id",
"use_litellm_proxy",
"prompt_label",
] + list(StandardCallbackDynamicParams.__annotations__.keys())

View File

@ -6835,6 +6835,14 @@ def _add_path_to_api_base(api_base: str, ending_path: str) -> str:
return str(modified_url.copy_with(params=original_url.params))
def get_standard_openai_params(params: dict) -> dict:
return {
k: v
for k, v in params.items()
if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS and v is not None
}
def get_non_default_completion_params(kwargs: dict) -> dict:
openai_params = litellm.OPENAI_CHAT_COMPLETION_PARAMS
default_params = openai_params + all_litellm_params

View File

@ -33,6 +33,7 @@ class TestCustomPromptManagement(CustomPromptManagement):
prompt_id: Optional[str],
prompt_variables: Optional[dict],
dynamic_callback_params: StandardCallbackDynamicParams,
prompt_label: Optional[str],
) -> Tuple[str, List[AllMessageValues], dict]:
print(
"TestCustomPromptManagement: running get_chat_completion_prompt for prompt_id: ",

View File

@ -62,9 +62,7 @@
"endTime": "2025-01-16T11:28:55.124353-08:00",
"completionStartTime": "2025-01-16T11:28:55.124353-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -103,9 +103,7 @@
"endTime": "2025-01-22T09:27:51.702048-08:00",
"completionStartTime": "2025-01-22T09:27:51.702048-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -81,9 +81,7 @@
"endTime": "2025-01-22T09:19:11.234200-08:00",
"completionStartTime": "2025-01-22T09:19:11.234200-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -52,9 +52,7 @@
"endTime": "2025-02-06T16:23:27.644253-08:00",
"completionStartTime": "2025-02-06T16:23:27.644253-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 10,

View File

@ -71,9 +71,7 @@
"endTime": "2025-01-22T07:31:28.962389-08:00",
"completionStartTime": "2025-01-22T07:31:28.962389-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -71,9 +71,7 @@
"endTime": "2025-01-22T08:38:26.015666-08:00",
"completionStartTime": "2025-01-22T08:38:26.015666-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -78,9 +78,7 @@
"endTime": "2025-01-22T09:59:39.365756-08:00",
"completionStartTime": "2025-01-22T09:59:39.365756-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -70,9 +70,7 @@
"endTime": "2025-01-22T10:06:50.958374-08:00",
"completionStartTime": "2025-01-22T10:06:50.958374-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -64,9 +64,7 @@
"endTime": "2025-01-22T09:59:32.880691-08:00",
"completionStartTime": "2025-01-22T09:59:32.880691-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -64,9 +64,7 @@
"endTime": "2025-01-22T09:59:36.161959-08:00",
"completionStartTime": "2025-01-22T09:59:36.161959-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -64,9 +64,7 @@
"endTime": "2025-01-22T09:59:32.880691-08:00",
"completionStartTime": "2025-01-22T09:59:32.880691-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -70,9 +70,7 @@
"endTime": "2025-01-22T09:55:28.853979-08:00",
"completionStartTime": "2025-01-22T09:55:28.853979-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -70,9 +70,7 @@
"endTime": "2025-01-22T09:53:53.753431-08:00",
"completionStartTime": "2025-01-22T09:53:53.753431-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -74,9 +74,7 @@
"endTime": "2025-01-22T09:56:35.476236-08:00",
"completionStartTime": "2025-01-22T09:56:35.476236-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,

View File

@ -78,9 +78,7 @@
"endTime": "2025-01-22T09:56:38.785762-08:00",
"completionStartTime": "2025-01-22T09:56:38.785762-08:00",
"model": "gpt-3.5-turbo",
"modelParameters": {
"extra_body": "{}"
},
"modelParameters": {},
"usage": {
"input": 10,
"output": 20,