mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-11 03:38:38 +00:00
test: add mocked langfuse discovery test in correct directory
This commit is contained in:
Thiago Riemma Carbonera
parent
fc4715eeab
commit
2db4c8b267
17 changed files with 491 additions and 310 deletions
|
|
@ -14,53 +14,74 @@ from litellm.types.utils import StandardCallbackDynamicParams
|
|||
class EnterpriseCallbackControls:
|
||||
@staticmethod
|
||||
def is_callback_disabled_dynamically(
|
||||
callback: litellm.CALLBACK_TYPES,
|
||||
litellm_params: dict,
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a callback is disabled via the x-litellm-disable-callbacks header or via `litellm_disabled_callbacks` in standard_callback_dynamic_params.
|
||||
|
||||
Args:
|
||||
callback: The callback to check (can be string, CustomLogger instance, or callable)
|
||||
litellm_params: Parameters containing proxy server request info
|
||||
|
||||
Returns:
|
||||
bool: True if the callback should be disabled, False otherwise
|
||||
"""
|
||||
from litellm.litellm_core_utils.custom_logger_registry import (
|
||||
CustomLoggerRegistry,
|
||||
)
|
||||
callback: litellm.CALLBACK_TYPES,
|
||||
litellm_params: dict,
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
"""
|
||||
Check if a callback is disabled via the x-litellm-disable-callbacks header or via `litellm_disabled_callbacks` in standard_callback_dynamic_params.
|
||||
|
||||
Args:
|
||||
callback: The callback to check (can be string, CustomLogger instance, or callable)
|
||||
litellm_params: Parameters containing proxy server request info
|
||||
|
||||
Returns:
|
||||
bool: True if the callback should be disabled, False otherwise
|
||||
"""
|
||||
from litellm.litellm_core_utils.custom_logger_registry import (
|
||||
CustomLoggerRegistry,
|
||||
)
|
||||
|
||||
try:
|
||||
disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(
|
||||
litellm_params, standard_callback_dynamic_params
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}"
|
||||
)
|
||||
verbose_logger.debug(
|
||||
f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}"
|
||||
)
|
||||
if disabled_callbacks is not None:
|
||||
#########################################################
|
||||
# premium user check
|
||||
#########################################################
|
||||
if (
|
||||
not EnterpriseCallbackControls._should_allow_dynamic_callback_disabling()
|
||||
):
|
||||
return False
|
||||
#########################################################
|
||||
if isinstance(callback, str):
|
||||
if callback.lower() in disabled_callbacks:
|
||||
verbose_logger.debug(
|
||||
f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}"
|
||||
)
|
||||
return True
|
||||
elif isinstance(callback, CustomLogger):
|
||||
# get the string name of the callback
|
||||
callback_str = (
|
||||
CustomLoggerRegistry.get_callback_str_from_class_type(
|
||||
callback.__class__
|
||||
)
|
||||
)
|
||||
if (
|
||||
callback_str is not None
|
||||
and callback_str.lower() in disabled_callbacks
|
||||
):
|
||||
verbose_logger.debug(
|
||||
f"Not logging to {callback_str} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}"
|
||||
)
|
||||
return True
|
||||
return False
|
||||
except Exception as e:
|
||||
verbose_logger.debug(f"Error checking disabled callbacks header: {str(e)}")
|
||||
return False
|
||||
|
||||
try:
|
||||
disabled_callbacks = EnterpriseCallbackControls.get_disabled_callbacks(litellm_params, standard_callback_dynamic_params)
|
||||
verbose_logger.debug(f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}")
|
||||
verbose_logger.debug(f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}")
|
||||
if disabled_callbacks is not None:
|
||||
#########################################################
|
||||
# premium user check
|
||||
#########################################################
|
||||
if not EnterpriseCallbackControls._should_allow_dynamic_callback_disabling():
|
||||
return False
|
||||
#########################################################
|
||||
if isinstance(callback, str):
|
||||
if callback.lower() in disabled_callbacks:
|
||||
verbose_logger.debug(f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}")
|
||||
return True
|
||||
elif isinstance(callback, CustomLogger):
|
||||
# get the string name of the callback
|
||||
callback_str = CustomLoggerRegistry.get_callback_str_from_class_type(callback.__class__)
|
||||
if callback_str is not None and callback_str.lower() in disabled_callbacks:
|
||||
verbose_logger.debug(f"Not logging to {callback_str} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}")
|
||||
return True
|
||||
return False
|
||||
except Exception as e:
|
||||
verbose_logger.debug(
|
||||
f"Error checking disabled callbacks header: {str(e)}"
|
||||
)
|
||||
return False
|
||||
@staticmethod
|
||||
def get_disabled_callbacks(litellm_params: dict, standard_callback_dynamic_params: StandardCallbackDynamicParams) -> Optional[List[str]]:
|
||||
def get_disabled_callbacks(
|
||||
litellm_params: dict,
|
||||
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||
) -> Optional[List[str]]:
|
||||
"""
|
||||
Get the disabled callbacks from the standard callback dynamic params.
|
||||
"""
|
||||
|
|
@ -71,18 +92,24 @@ class EnterpriseCallbackControls:
|
|||
request_headers = get_proxy_server_request_headers(litellm_params)
|
||||
disabled_callbacks = request_headers.get(X_LITELLM_DISABLE_CALLBACKS, None)
|
||||
if disabled_callbacks is not None:
|
||||
disabled_callbacks = set([cb.strip().lower() for cb in disabled_callbacks.split(",")])
|
||||
disabled_callbacks = set(
|
||||
[cb.strip().lower() for cb in disabled_callbacks.split(",")]
|
||||
)
|
||||
return list(disabled_callbacks)
|
||||
|
||||
|
||||
#########################################################
|
||||
# check if disabled via request body
|
||||
#########################################################
|
||||
if standard_callback_dynamic_params.get("litellm_disabled_callbacks", None) is not None:
|
||||
return standard_callback_dynamic_params.get("litellm_disabled_callbacks", None)
|
||||
|
||||
if (
|
||||
standard_callback_dynamic_params.get("litellm_disabled_callbacks", None)
|
||||
is not None
|
||||
):
|
||||
return standard_callback_dynamic_params.get(
|
||||
"litellm_disabled_callbacks", None
|
||||
)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
@staticmethod
|
||||
def _should_allow_dynamic_callback_disabling():
|
||||
import litellm
|
||||
|
|
@ -90,10 +117,14 @@ class EnterpriseCallbackControls:
|
|||
|
||||
# Check if admin has disabled this feature
|
||||
if litellm.allow_dynamic_callback_disabling is not True:
|
||||
verbose_logger.debug("Dynamic callback disabling is disabled by admin via litellm.allow_dynamic_callback_disabling")
|
||||
verbose_logger.debug(
|
||||
"Dynamic callback disabling is disabled by admin via litellm.allow_dynamic_callback_disabling"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
if premium_user:
|
||||
return True
|
||||
verbose_logger.warning(f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}")
|
||||
return False
|
||||
verbose_logger.warning(
|
||||
f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}"
|
||||
)
|
||||
return False
|
||||
|
|
|
|||
|
|
@ -612,9 +612,9 @@ class BaseEmailLogger(CustomLogger):
|
|||
else "LiteLLM Notification"
|
||||
)
|
||||
|
||||
recipient_email: Optional[str] = (
|
||||
user_email or await self._lookup_user_email_from_db(user_id=user_id)
|
||||
)
|
||||
recipient_email: Optional[
|
||||
str
|
||||
] = user_email or await self._lookup_user_email_from_db(user_id=user_id)
|
||||
if recipient_email is None:
|
||||
raise ValueError(
|
||||
f"User email not found for user_id: {user_id}. User email is required to send email."
|
||||
|
|
|
|||
|
|
@ -79,4 +79,4 @@ class SendGridEmailLogger(BaseEmailLogger):
|
|||
verbose_logger.debug(
|
||||
f"SendGrid response status={response.status_code}, body={response.text}"
|
||||
)
|
||||
return
|
||||
return
|
||||
|
|
|
|||
|
|
@ -7,4 +7,4 @@ including custom SSO handlers and advanced authentication features.
|
|||
|
||||
from .custom_sso_handler import EnterpriseCustomSSOHandler
|
||||
|
||||
__all__ = ["EnterpriseCustomSSOHandler"]
|
||||
__all__ = ["EnterpriseCustomSSOHandler"]
|
||||
|
|
|
|||
|
|
@ -26,12 +26,12 @@ from litellm.proxy.management_endpoints.types import CustomOpenID
|
|||
class EnterpriseCustomSSOHandler:
|
||||
"""
|
||||
Enterprise Custom SSO Handler for LiteLLM Proxy
|
||||
|
||||
|
||||
This class provides methods for handling custom SSO authentication flows
|
||||
where users can implement their own authentication logic by processing
|
||||
request headers and returning user information in OpenID format.
|
||||
"""
|
||||
|
||||
|
||||
@staticmethod
|
||||
async def handle_custom_ui_sso_sign_in(
|
||||
request: Request,
|
||||
|
|
@ -40,16 +40,16 @@ class EnterpriseCustomSSOHandler:
|
|||
Allow a user to execute their custom code to parse incoming request headers and return a OpenID object
|
||||
|
||||
Use this when you have an OAuth proxy in front of LiteLLM (where the OAuth proxy has already authenticated the user)
|
||||
|
||||
|
||||
Args:
|
||||
request: The FastAPI request object containing headers and other request data
|
||||
|
||||
|
||||
Returns:
|
||||
RedirectResponse: Redirect response that sends the user to the LiteLLM UI with authentication token
|
||||
|
||||
|
||||
Raises:
|
||||
ValueError: If custom_ui_sso_sign_in_handler is not configured
|
||||
|
||||
|
||||
Example:
|
||||
This method is typically called when a user has already been authenticated by an
|
||||
external OAuth proxy and the proxy has added custom headers containing user information.
|
||||
|
|
@ -63,24 +63,31 @@ class EnterpriseCustomSSOHandler:
|
|||
premium_user,
|
||||
user_custom_ui_sso_sign_in_handler,
|
||||
)
|
||||
|
||||
if premium_user is not True:
|
||||
raise ValueError(CommonProxyErrors.not_premium_user.value)
|
||||
|
||||
|
||||
if user_custom_ui_sso_sign_in_handler is None:
|
||||
raise ValueError("custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings.")
|
||||
|
||||
custom_sso_login_handler = cast(CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler)
|
||||
openid_response: OpenID = await custom_sso_login_handler.handle_custom_ui_sso_sign_in(
|
||||
request=request,
|
||||
raise ValueError(
|
||||
"custom_ui_sso_sign_in_handler is not configured. Please set it in general_settings."
|
||||
)
|
||||
|
||||
custom_sso_login_handler = cast(
|
||||
CustomSSOLoginHandler, user_custom_ui_sso_sign_in_handler
|
||||
)
|
||||
|
||||
openid_response: OpenID = (
|
||||
await custom_sso_login_handler.handle_custom_ui_sso_sign_in(
|
||||
request=request,
|
||||
)
|
||||
)
|
||||
|
||||
# Import here to avoid circular imports
|
||||
from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler
|
||||
|
||||
|
||||
return await SSOAuthenticationHandler.get_redirect_response_from_openid(
|
||||
result=openid_response,
|
||||
request=request,
|
||||
received_response=None,
|
||||
generic_client_id=None,
|
||||
ui_access_mode=None,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -53,7 +53,9 @@ class CheckBatchCost:
|
|||
"user_api_key_alias": getattr(user_row, "user_alias", None),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}")
|
||||
verbose_proxy_logger.error(
|
||||
f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}"
|
||||
)
|
||||
return {}
|
||||
|
||||
async def _cleanup_stale_managed_objects(self) -> None:
|
||||
|
|
@ -62,11 +64,22 @@ class CheckBatchCost:
|
|||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(
|
||||
days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS
|
||||
)
|
||||
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
|
||||
"status": {
|
||||
"not_in": [
|
||||
"completed",
|
||||
"complete",
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"stale_expired",
|
||||
]
|
||||
},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
|
|
@ -152,7 +165,11 @@ class CheckBatchCost:
|
|||
order={"created_at": "asc"},
|
||||
)
|
||||
except Exception as query_err:
|
||||
if "batch_processed" not in str(query_err).lower() and "unknown column" not in str(query_err).lower() and "does not exist" not in str(query_err).lower():
|
||||
if (
|
||||
"batch_processed" not in str(query_err).lower()
|
||||
and "unknown column" not in str(query_err).lower()
|
||||
and "does not exist" not in str(query_err).lower()
|
||||
):
|
||||
raise
|
||||
# Permanent schema gap — cache the result so future cycles skip straight to fallback
|
||||
self._has_batch_processed_column = False
|
||||
|
|
@ -205,10 +222,7 @@ class CheckBatchCost:
|
|||
continue
|
||||
|
||||
## RETRIEVE THE BATCH JOB OUTPUT FILE
|
||||
if (
|
||||
response.status == "completed"
|
||||
and response.output_file_id is not None
|
||||
):
|
||||
if response.status == "completed" and response.output_file_id is not None:
|
||||
verbose_proxy_logger.info(
|
||||
f"Batch ID: {batch_id} is complete, tracking cost and usage"
|
||||
)
|
||||
|
|
@ -235,20 +249,25 @@ class CheckBatchCost:
|
|||
decoded = _is_base64_encoded_unified_file_id(raw_output_file_id)
|
||||
if decoded:
|
||||
try:
|
||||
raw_output_file_id = decoded.split("llm_output_file_id,")[1].split(";")[0]
|
||||
raw_output_file_id = decoded.split("llm_output_file_id,")[
|
||||
1
|
||||
].split(";")[0]
|
||||
except (IndexError, AttributeError):
|
||||
pass
|
||||
|
||||
credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
credentials = (
|
||||
self.llm_router.get_deployment_credentials_with_provider(model_id)
|
||||
or {}
|
||||
)
|
||||
_file_content = await afile_content(
|
||||
file_id=raw_output_file_id,
|
||||
**credentials,
|
||||
)
|
||||
|
||||
# Access content - handle both direct attribute and method call
|
||||
if hasattr(_file_content, 'content'):
|
||||
if hasattr(_file_content, "content"):
|
||||
content_bytes = _file_content.content # type: ignore[union-attr]
|
||||
elif hasattr(_file_content, 'read'):
|
||||
elif hasattr(_file_content, "read"):
|
||||
content_bytes = await _file_content.read() # type: ignore[misc]
|
||||
else:
|
||||
content_bytes = _file_content # type: ignore[assignment]
|
||||
|
|
@ -273,14 +292,20 @@ class CheckBatchCost:
|
|||
|
||||
# Pass deployment model_info so custom batch pricing
|
||||
# (input_cost_per_token_batches etc.) is used for cost calc
|
||||
deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}
|
||||
batch_cost, batch_usage, batch_models = (
|
||||
await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=file_content_as_dict,
|
||||
custom_llm_provider=llm_provider, # type: ignore
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info, # type: ignore[arg-type]
|
||||
)
|
||||
deployment_model_info = (
|
||||
deployment_info.model_info.model_dump()
|
||||
if deployment_info.model_info
|
||||
else {}
|
||||
)
|
||||
(
|
||||
batch_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
) = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=file_content_as_dict,
|
||||
custom_llm_provider=llm_provider, # type: ignore
|
||||
model_name=model_name,
|
||||
model_info=deployment_model_info, # type: ignore[arg-type]
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=batch_models[0],
|
||||
|
|
|
|||
|
|
@ -38,11 +38,22 @@ class CheckResponsesCost:
|
|||
in non-terminal states as 'stale_expired'. These will never complete and
|
||||
should not be polled.
|
||||
"""
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS)
|
||||
cutoff = datetime.now(timezone.utc) - timedelta(
|
||||
days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS
|
||||
)
|
||||
result = await self.prisma_client.db.litellm_managedobjecttable.update_many(
|
||||
where={
|
||||
"file_purpose": "response",
|
||||
"status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]},
|
||||
"status": {
|
||||
"not_in": [
|
||||
"completed",
|
||||
"complete",
|
||||
"failed",
|
||||
"expired",
|
||||
"cancelled",
|
||||
"stale_expired",
|
||||
]
|
||||
},
|
||||
"created_at": {"lt": cutoff},
|
||||
},
|
||||
data={"status": "stale_expired"},
|
||||
|
|
@ -76,7 +87,7 @@ class CheckResponsesCost:
|
|||
take=MAX_OBJECTS_PER_POLL_CYCLE,
|
||||
order={"created_at": "asc"},
|
||||
)
|
||||
|
||||
|
||||
verbose_proxy_logger.debug(f"Found {len(jobs)} response jobs to check")
|
||||
completed_jobs = []
|
||||
|
||||
|
|
@ -91,29 +102,35 @@ class CheckResponsesCost:
|
|||
# Get the stored response object to extract model information
|
||||
stored_response = job.file_object
|
||||
model_name = stored_response.get("model", None)
|
||||
|
||||
|
||||
# Decrypt the response ID
|
||||
responses_id_security, _, _ = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
|
||||
|
||||
(
|
||||
responses_id_security,
|
||||
_,
|
||||
_,
|
||||
) = ResponsesIDSecurity()._decrypt_response_id(unified_object_id)
|
||||
|
||||
# Prepare metadata with model information for cost tracking
|
||||
litellm_metadata = {
|
||||
"user_api_key_user_id": job.created_by or "default-user-id",
|
||||
}
|
||||
|
||||
|
||||
# Add model information if available
|
||||
if model_name:
|
||||
litellm_metadata["model"] = model_name
|
||||
litellm_metadata["model_group"] = model_name # Use same value for model_group
|
||||
|
||||
litellm_metadata[
|
||||
"model_group"
|
||||
] = model_name # Use same value for model_group
|
||||
|
||||
response = await litellm.aget_responses(
|
||||
response_id=responses_id_security,
|
||||
litellm_metadata=litellm_metadata,
|
||||
)
|
||||
|
||||
|
||||
verbose_proxy_logger.debug(
|
||||
f"Response {unified_object_id} status: {response.status}, model: {model_name}"
|
||||
)
|
||||
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.info(
|
||||
f"Skipping job {unified_object_id} due to error: {e}"
|
||||
|
|
@ -126,7 +143,7 @@ class CheckResponsesCost:
|
|||
f"Response {unified_object_id} is complete. Cost automatically tracked by aget_responses."
|
||||
)
|
||||
completed_jobs.append(job)
|
||||
|
||||
|
||||
elif response.status in ["failed", "cancelled"]:
|
||||
verbose_proxy_logger.info(
|
||||
f"Response {unified_object_id} has status {response.status}, marking as complete"
|
||||
|
|
@ -142,4 +159,3 @@ class CheckResponsesCost:
|
|||
verbose_proxy_logger.info(
|
||||
f"Marked {len(completed_jobs)} response jobs as completed"
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -118,7 +118,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
db_data["storage_backend"] = hidden_params["storage_backend"]
|
||||
if "storage_url" in hidden_params:
|
||||
db_data["storage_url"] = hidden_params["storage_url"]
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Storage metadata: storage_backend={db_data.get('storage_backend')}, "
|
||||
f"storage_url={db_data.get('storage_url')}"
|
||||
|
|
@ -278,28 +278,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
raise Exception(
|
||||
"Filtering by 'target_model_names' is not supported when using managed batches."
|
||||
)
|
||||
|
||||
|
||||
where_clause: Dict[str, Any] = {"file_purpose": "batch"}
|
||||
|
||||
|
||||
# Filter by user who created the batch
|
||||
if user_api_key_dict.user_id:
|
||||
where_clause["created_by"] = user_api_key_dict.user_id
|
||||
|
||||
|
||||
if after:
|
||||
where_clause["id"] = {"gt": after}
|
||||
|
||||
|
||||
# Fetch more than needed to allow for post-fetch filtering
|
||||
fetch_limit = limit or 20
|
||||
if target_model_names:
|
||||
# Fetch extra to account for filtering
|
||||
fetch_limit = max(fetch_limit * 3, 100)
|
||||
|
||||
|
||||
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where=where_clause,
|
||||
take=fetch_limit,
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
|
||||
|
||||
batch_objects: List[LiteLLMBatch] = []
|
||||
for batch in batches:
|
||||
try:
|
||||
|
|
@ -307,7 +307,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if len(batch_objects) >= (limit or 20):
|
||||
break
|
||||
|
||||
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
|
||||
batch_data = (
|
||||
json.loads(batch.file_object)
|
||||
if isinstance(batch.file_object, str)
|
||||
else batch.file_object
|
||||
)
|
||||
batch_obj = LiteLLMBatch(**batch_data)
|
||||
batch_obj.id = batch.unified_object_id
|
||||
batch_objects.append(batch_obj)
|
||||
|
|
@ -317,7 +321,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
f"Failed to parse batch object {batch.unified_object_id}: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
|
||||
return {
|
||||
"object": "list",
|
||||
"data": batch_objects,
|
||||
|
|
@ -370,11 +374,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
"""
|
||||
Check if the user has access to a list of file IDs.
|
||||
Only checks managed (unified) file IDs.
|
||||
|
||||
|
||||
Args:
|
||||
file_ids: List of file IDs to check access for
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Raises:
|
||||
HTTPException: If user doesn't have access to any of the files
|
||||
"""
|
||||
|
|
@ -412,10 +416,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
### HANDLE TRANSFORMATIONS ###
|
||||
# Check both completion and acompletion call types
|
||||
is_completion_call = (
|
||||
call_type == CallTypes.completion.value
|
||||
call_type == CallTypes.completion.value
|
||||
or call_type == CallTypes.acompletion.value
|
||||
)
|
||||
|
||||
|
||||
if is_completion_call:
|
||||
messages = data.get("messages")
|
||||
model = data.get("model", "")
|
||||
|
|
@ -424,22 +428,27 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if file_ids:
|
||||
# Check user has access to all managed files
|
||||
await self.check_file_ids_access(file_ids, user_api_key_dict)
|
||||
|
||||
|
||||
# Check if any files are stored in storage backends and need base64 conversion
|
||||
# This is needed for Vertex AI/Gemini which requires base64 content
|
||||
is_vertex_ai = model and ("vertex_ai" in model or "gemini" in model.lower())
|
||||
is_vertex_ai = model and (
|
||||
"vertex_ai" in model or "gemini" in model.lower()
|
||||
)
|
||||
if is_vertex_ai:
|
||||
await self._convert_storage_files_to_base64(
|
||||
messages=messages,
|
||||
file_ids=file_ids,
|
||||
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
|
||||
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
file_ids, user_api_key_dict.parent_otel_span
|
||||
)
|
||||
data["model_file_id_mapping"] = model_file_id_mapping
|
||||
elif call_type == CallTypes.aresponses.value or call_type == CallTypes.responses.value:
|
||||
elif (
|
||||
call_type == CallTypes.aresponses.value
|
||||
or call_type == CallTypes.responses.value
|
||||
):
|
||||
# Handle managed files in responses API input and tools
|
||||
file_ids = []
|
||||
|
||||
|
|
@ -604,7 +613,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if model_id is None:
|
||||
model_id = cast(
|
||||
Optional[str],
|
||||
kwargs.get("litellm_metadata", {}).get("model_info", {}).get("id", None),
|
||||
kwargs.get("litellm_metadata", {})
|
||||
.get("model_info", {})
|
||||
.get("id", None),
|
||||
)
|
||||
mapped_file_id: Optional[str] = None
|
||||
if input_file_id and model_file_id_mapping and model_id:
|
||||
|
|
@ -641,7 +652,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> List[str]:
|
||||
"""
|
||||
Gets file ids from responses API input.
|
||||
|
||||
|
||||
The input can be:
|
||||
- A string (no files)
|
||||
- A list of input items, where each item can have:
|
||||
|
|
@ -649,32 +660,35 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
- content: a list that can contain items with type: "input_file" and file_id
|
||||
"""
|
||||
file_ids: List[str] = []
|
||||
|
||||
|
||||
if isinstance(input, str):
|
||||
return file_ids
|
||||
|
||||
|
||||
if not isinstance(input, list):
|
||||
return file_ids
|
||||
|
||||
|
||||
for item in input:
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
|
||||
|
||||
# Check for direct input_file type
|
||||
if item.get("type") == "input_file":
|
||||
file_id = item.get("file_id")
|
||||
if file_id:
|
||||
file_ids.append(file_id)
|
||||
|
||||
|
||||
# Check for input_file in content array
|
||||
content = item.get("content")
|
||||
if isinstance(content, list):
|
||||
for content_item in content:
|
||||
if isinstance(content_item, dict) and content_item.get("type") == "input_file":
|
||||
if (
|
||||
isinstance(content_item, dict)
|
||||
and content_item.get("type") == "input_file"
|
||||
):
|
||||
file_id = content_item.get("file_id")
|
||||
if file_id:
|
||||
file_ids.append(file_id)
|
||||
|
||||
|
||||
return file_ids
|
||||
|
||||
def get_file_ids_from_responses_tools(
|
||||
|
|
@ -682,7 +696,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> List[str]:
|
||||
"""
|
||||
Gets file ids from responses API tools parameter.
|
||||
|
||||
|
||||
The tools can contain code_interpreter with container.file_ids:
|
||||
[
|
||||
{
|
||||
|
|
@ -692,14 +706,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
]
|
||||
"""
|
||||
file_ids: List[str] = []
|
||||
|
||||
|
||||
if not isinstance(tools, list):
|
||||
return file_ids
|
||||
|
||||
|
||||
for tool in tools:
|
||||
if not isinstance(tool, dict):
|
||||
continue
|
||||
|
||||
|
||||
# Check for code_interpreter with container file_ids
|
||||
if tool.get("type") == "code_interpreter":
|
||||
container = tool.get("container")
|
||||
|
|
@ -709,7 +723,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
for file_id in container_file_ids:
|
||||
if isinstance(file_id, str):
|
||||
file_ids.append(file_id)
|
||||
|
||||
|
||||
return file_ids
|
||||
|
||||
def get_vector_store_ids_from_file_search_tools(
|
||||
|
|
@ -1041,16 +1055,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
model_name=resolved_model_name,
|
||||
)
|
||||
setattr(response, file_attr, unified_file_id)
|
||||
|
||||
|
||||
# Use llm_router credentials when available. Without credentials,
|
||||
# Azure and other auth-required providers return 500/401.
|
||||
file_object = None
|
||||
try:
|
||||
# Import module and use getattr for better testability with mocks
|
||||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
_llm_router = getattr(proxy_server_module, 'llm_router', None)
|
||||
|
||||
_llm_router = getattr(
|
||||
proxy_server_module, "llm_router", None
|
||||
)
|
||||
if _llm_router is not None and model_id:
|
||||
_creds = _llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
_creds = (
|
||||
_llm_router.get_deployment_credentials_with_provider(
|
||||
model_id
|
||||
)
|
||||
or {}
|
||||
)
|
||||
file_object = await litellm.afile_retrieve(
|
||||
file_id=original_file_id,
|
||||
**_creds,
|
||||
|
|
@ -1067,7 +1089,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
verbose_logger.warning(
|
||||
f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand."
|
||||
)
|
||||
|
||||
|
||||
await self.store_unified_file_id(
|
||||
file_id=unified_file_id,
|
||||
file_object=file_object,
|
||||
|
|
@ -1142,7 +1164,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
# Case 1 : This is not a managed file
|
||||
if not stored_file_object:
|
||||
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
|
||||
|
||||
|
||||
# Case 2: Managed file and the file object exists in the database
|
||||
# The stored file_object has the raw provider ID. Replace with the unified ID
|
||||
# so callers see a consistent ID (matching Case 3 which does response.id = file_id).
|
||||
|
|
@ -1160,13 +1182,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
|
||||
try:
|
||||
model_id, model_file_id = next(iter(stored_file_object.model_mappings.items()))
|
||||
credentials = llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
response = await litellm.afile_retrieve(file_id=model_file_id, **credentials)
|
||||
model_id, model_file_id = next(
|
||||
iter(stored_file_object.model_mappings.items())
|
||||
)
|
||||
credentials = (
|
||||
llm_router.get_deployment_credentials_with_provider(model_id) or {}
|
||||
)
|
||||
response = await litellm.afile_retrieve(
|
||||
file_id=model_file_id, **credentials
|
||||
)
|
||||
response.id = file_id # Replace with unified ID
|
||||
return response
|
||||
except Exception as e:
|
||||
raise Exception(f"Failed to retrieve file {file_id} from provider: {str(e)}") from e
|
||||
raise Exception(
|
||||
f"Failed to retrieve file {file_id} from provider: {str(e)}"
|
||||
) from e
|
||||
|
||||
async def afile_list(
|
||||
self,
|
||||
|
|
@ -1188,19 +1218,19 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
import litellm.proxy.proxy_server as proxy_server_module
|
||||
|
||||
# Check if the scheduler has the batch cost checking job registered
|
||||
scheduler = getattr(proxy_server_module, 'scheduler', None)
|
||||
scheduler = getattr(proxy_server_module, "scheduler", None)
|
||||
if scheduler is None:
|
||||
return False
|
||||
|
||||
|
||||
# Check if the check_batch_cost_job exists in the scheduler
|
||||
try:
|
||||
job = scheduler.get_job('check_batch_cost_job')
|
||||
job = scheduler.get_job("check_batch_cost_job")
|
||||
if job is not None:
|
||||
return True
|
||||
except Exception:
|
||||
# Job not found or scheduler doesn't support get_job
|
||||
pass
|
||||
|
||||
|
||||
return False
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
|
|
@ -1208,28 +1238,26 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
)
|
||||
return False
|
||||
|
||||
async def _get_batches_referencing_file(
|
||||
self, file_id: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
async def _get_batches_referencing_file(self, file_id: str) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Find batches that reference this file and still need cost tracking.
|
||||
Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost.
|
||||
Args:
|
||||
file_id: The unified file ID to check
|
||||
|
||||
|
||||
Returns:
|
||||
List of batch objects referencing this file in non-terminal state
|
||||
(max 10 for error message display)
|
||||
"""
|
||||
# Prepare list of file IDs to check (both unified and provider IDs)
|
||||
file_ids_to_check = [file_id]
|
||||
|
||||
|
||||
# Get model-specific file IDs for this unified file ID if it's a managed file
|
||||
try:
|
||||
model_file_id_mapping = await self.get_model_file_id_mapping(
|
||||
[file_id], litellm_parent_otel_span=None
|
||||
)
|
||||
|
||||
|
||||
if model_file_id_mapping and file_id in model_file_id_mapping:
|
||||
# Add all provider file IDs for this unified file
|
||||
provider_file_ids = list(model_file_id_mapping[file_id].values())
|
||||
|
|
@ -1239,59 +1267,67 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
f"Could not get model file ID mapping for {file_id}: {e}. "
|
||||
f"Will only check unified file ID."
|
||||
)
|
||||
MAX_MATCHES_TO_RETURN = 10
|
||||
|
||||
MAX_MATCHES_TO_RETURN = 10
|
||||
|
||||
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
|
||||
where={
|
||||
"file_purpose": "batch",
|
||||
"batch_processed": False,
|
||||
"status": {"not_in": ["failed", "expired", "cancelled"]}
|
||||
"status": {"not_in": ["failed", "expired", "cancelled"]},
|
||||
},
|
||||
take=MAX_MATCHES_TO_RETURN,
|
||||
order={"created_at": "desc"},
|
||||
)
|
||||
|
||||
|
||||
referencing_batches = []
|
||||
for batch in batches:
|
||||
try:
|
||||
# Parse the batch file_object to check for file references
|
||||
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object
|
||||
|
||||
batch_data = (
|
||||
json.loads(batch.file_object)
|
||||
if isinstance(batch.file_object, str)
|
||||
else batch.file_object
|
||||
)
|
||||
|
||||
# Extract file IDs from batch
|
||||
# Batches typically reference the unified file ID in input_file_id
|
||||
# Output and error files are generated by the provider
|
||||
input_file_id = batch_data.get("input_file_id")
|
||||
output_file_id = batch_data.get("output_file_id")
|
||||
error_file_id = batch_data.get("error_file_id")
|
||||
|
||||
referenced_file_ids = [fid for fid in [input_file_id, output_file_id, error_file_id] if fid]
|
||||
|
||||
|
||||
referenced_file_ids = [
|
||||
fid for fid in [input_file_id, output_file_id, error_file_id] if fid
|
||||
]
|
||||
|
||||
# Check if any referenced file ID matches the file we're trying to delete
|
||||
if any(ref_id in file_ids_to_check for ref_id in referenced_file_ids):
|
||||
referencing_batches.append({
|
||||
"batch_id": batch.unified_object_id,
|
||||
"status": batch.status,
|
||||
"created_at": batch.created_at,
|
||||
})
|
||||
referencing_batches.append(
|
||||
{
|
||||
"batch_id": batch.unified_object_id,
|
||||
"status": batch.status,
|
||||
"created_at": batch.created_at,
|
||||
}
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Error parsing batch object {batch.unified_object_id}: {e}"
|
||||
)
|
||||
continue
|
||||
|
||||
|
||||
return referencing_batches
|
||||
|
||||
async def _check_file_deletion_allowed(self, file_id: str) -> None:
|
||||
"""
|
||||
Check if file deletion should be blocked due to batch references.
|
||||
|
||||
|
||||
Blocks deletion if:
|
||||
1. File is referenced by any batch in non-terminal state, AND
|
||||
2. Batch polling is configured (user wants cost tracking)
|
||||
|
||||
|
||||
Args:
|
||||
file_id: The unified file ID to check
|
||||
|
||||
|
||||
Raises:
|
||||
HTTPException: If file deletion should be blocked
|
||||
"""
|
||||
|
|
@ -1299,39 +1335,45 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if not self._is_batch_polling_enabled():
|
||||
# Batch polling not configured, allow deletion
|
||||
return
|
||||
|
||||
|
||||
# Check if file is referenced by any non-terminal batches
|
||||
referencing_batches = await self._get_batches_referencing_file(file_id)
|
||||
|
||||
|
||||
if referencing_batches:
|
||||
# File is referenced by non-terminal batches and polling is enabled
|
||||
MAX_BATCHES_IN_ERROR = 5 # Limit batches shown in error message for readability
|
||||
|
||||
MAX_BATCHES_IN_ERROR = (
|
||||
5 # Limit batches shown in error message for readability
|
||||
)
|
||||
|
||||
# Show up to MAX_BATCHES_IN_ERROR in the error message
|
||||
batches_to_show = referencing_batches[:MAX_BATCHES_IN_ERROR]
|
||||
batch_statuses = [f"{b['batch_id']}: {b['status']}" for b in batches_to_show]
|
||||
|
||||
batch_statuses = [
|
||||
f"{b['batch_id']}: {b['status']}" for b in batches_to_show
|
||||
]
|
||||
|
||||
# Determine the count message
|
||||
count_message = f"{len(referencing_batches)}"
|
||||
if len(referencing_batches) >= 10: # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
|
||||
if (
|
||||
len(referencing_batches) >= 10
|
||||
): # MAX_MATCHES_TO_RETURN from _get_batches_referencing_file
|
||||
count_message = "10+"
|
||||
|
||||
|
||||
error_message = (
|
||||
f"Cannot delete file {file_id}. "
|
||||
f"The file is referenced by {count_message} batch(es) in non-terminal state"
|
||||
)
|
||||
|
||||
|
||||
# Add specific batch details if not too many
|
||||
if len(referencing_batches) <= MAX_BATCHES_IN_ERROR:
|
||||
error_message += f": {', '.join(batch_statuses)}. "
|
||||
else:
|
||||
error_message += f" (showing {MAX_BATCHES_IN_ERROR} most recent): {', '.join(batch_statuses)}. "
|
||||
|
||||
|
||||
error_message += (
|
||||
f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. "
|
||||
f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)."
|
||||
)
|
||||
|
||||
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=error_message,
|
||||
|
|
@ -1344,7 +1386,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
llm_router: Router,
|
||||
**data: Dict,
|
||||
) -> OpenAIFileObject:
|
||||
|
||||
# Check if file deletion should be blocked due to batch references
|
||||
await self._check_file_deletion_allowed(file_id)
|
||||
|
||||
|
|
@ -1357,7 +1398,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
|
||||
if specific_model_file_id_mapping:
|
||||
# Remove conflicting keys from data to avoid duplicate keyword arguments
|
||||
filtered_data = {k: v for k, v in data.items() if k not in ("model", "file_id")}
|
||||
filtered_data = {
|
||||
k: v for k, v in data.items() if k not in ("model", "file_id")
|
||||
}
|
||||
for model_id, model_file_id in specific_model_file_id_mapping.items():
|
||||
delete_response = await llm_router.afile_delete(model=model_id, file_id=model_file_id, **filtered_data) # type: ignore
|
||||
|
||||
|
|
@ -1412,7 +1455,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> None:
|
||||
"""
|
||||
Convert files stored in storage backends to base64 format for Vertex AI/Gemini.
|
||||
|
||||
|
||||
This method checks if any managed files are stored in storage backends,
|
||||
downloads them, and converts them to base64 format in the messages.
|
||||
"""
|
||||
|
|
@ -1420,29 +1463,29 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
for file_id in file_ids:
|
||||
# Check if this is a base64 encoded unified file ID
|
||||
decoded_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
|
||||
|
||||
|
||||
if not decoded_unified_file_id:
|
||||
continue
|
||||
|
||||
|
||||
# Check database for storage backend info
|
||||
# IMPORTANT: The database stores the base64 encoded unified_file_id (not the decoded version)
|
||||
# So we query with the original file_id (which is base64 encoded)
|
||||
db_file = await self.prisma_client.db.litellm_managedfiletable.find_first(
|
||||
where={"unified_file_id": file_id}
|
||||
)
|
||||
|
||||
|
||||
if not db_file or not db_file.storage_backend or not db_file.storage_url:
|
||||
continue
|
||||
|
||||
|
||||
# File is stored in a storage backend, download and convert to base64
|
||||
try:
|
||||
from litellm.llms.base_llm.files.storage_backend_factory import (
|
||||
get_storage_backend,
|
||||
)
|
||||
|
||||
|
||||
storage_backend_name = db_file.storage_backend
|
||||
storage_url = db_file.storage_url
|
||||
|
||||
|
||||
# Get storage backend (uses same env vars as callback)
|
||||
try:
|
||||
storage_backend = get_storage_backend(storage_backend_name)
|
||||
|
|
@ -1451,18 +1494,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
f"Storage backend '{storage_backend_name}' error for file {file_id}: {str(e)}"
|
||||
)
|
||||
continue
|
||||
|
||||
|
||||
file_content = await storage_backend.download_file(storage_url)
|
||||
|
||||
|
||||
# Determine content type from file object
|
||||
content_type = self._get_content_type_from_file_object(db_file.file_object)
|
||||
|
||||
content_type = self._get_content_type_from_file_object(
|
||||
db_file.file_object
|
||||
)
|
||||
|
||||
# Convert to base64
|
||||
base64_data = base64.b64encode(file_content).decode("utf-8")
|
||||
base64_data_uri = f"data:{content_type};base64,{base64_data}"
|
||||
|
||||
|
||||
# Update messages to use base64 instead of file_id
|
||||
self._update_messages_with_base64_data(messages, file_id, base64_data_uri, content_type)
|
||||
self._update_messages_with_base64_data(
|
||||
messages, file_id, base64_data_uri, content_type
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_logger.exception(
|
||||
f"Error converting file {file_id} from storage backend to base64: {str(e)}"
|
||||
|
|
@ -1473,21 +1520,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
def _get_content_type_from_file_object(self, file_object: Optional[Any]) -> str:
|
||||
"""
|
||||
Determine content type from file object.
|
||||
|
||||
|
||||
Uses the MIME type utility for consistent detection and normalization.
|
||||
|
||||
|
||||
Args:
|
||||
file_object: The file object from the database (can be dict, JSON string, or None)
|
||||
|
||||
|
||||
Returns:
|
||||
str: MIME type (defaults to "application/octet-stream" if cannot be determined)
|
||||
"""
|
||||
# Use utility function for detection
|
||||
content_type = get_content_type_from_file_object(file_object)
|
||||
|
||||
|
||||
# Normalize for Gemini/Vertex AI (requires image/jpeg, not image/jpg)
|
||||
content_type = normalize_mime_type_for_provider(content_type, provider="gemini")
|
||||
|
||||
|
||||
return content_type
|
||||
|
||||
def _update_messages_with_base64_data(
|
||||
|
|
@ -1499,7 +1546,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
) -> None:
|
||||
"""
|
||||
Update messages to replace file_id with base64 data URI.
|
||||
|
||||
|
||||
Args:
|
||||
messages: List of messages to update
|
||||
file_id: The file ID to replace
|
||||
|
|
@ -1514,7 +1561,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
if element.get("type") == "file":
|
||||
file_element = cast(ChatCompletionFileObject, element)
|
||||
file_element_file = file_element.get("file", {})
|
||||
|
||||
|
||||
if file_element_file.get("file_id") == file_id:
|
||||
# Replace file_id with base64 data
|
||||
file_element_file["file_data"] = base64_data_uri
|
||||
|
|
@ -1522,7 +1569,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
|
|||
file_element_file["format"] = content_type
|
||||
# Remove file_id to ensure only file_data is used
|
||||
file_element_file.pop("file_id", None)
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Converted file {file_id} from storage backend to base64 with format {content_type}"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -41,7 +41,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
):
|
||||
"""
|
||||
Managed vector stores with target_model_names support.
|
||||
|
||||
|
||||
This class provides functionality to:
|
||||
- Create vector stores across multiple models
|
||||
- Retrieve vector stores by unified ID
|
||||
|
|
@ -77,14 +77,14 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> str:
|
||||
"""
|
||||
Generate the format string for the unified vector store ID.
|
||||
|
||||
|
||||
Format:
|
||||
litellm_proxy:vector_store;unified_id,<uuid>;target_model_names,<models>;resource_id,<vs_id>;model_id,<model_id>
|
||||
"""
|
||||
# VectorStoreCreateResponse is a TypedDict, so resource_object is a dictionary
|
||||
# Extract provider resource ID from the response
|
||||
provider_resource_id = resource_object.get("id", "")
|
||||
|
||||
|
||||
# Model ID is stored in hidden params if the response object supports it
|
||||
# For TypedDict responses, we need to check if _hidden_params was added
|
||||
hidden_params: Dict[str, Any] = {}
|
||||
|
|
@ -109,20 +109,18 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> VectorStoreCreateResponse:
|
||||
"""
|
||||
Create a vector store for a specific model.
|
||||
|
||||
|
||||
Args:
|
||||
llm_router: LiteLLM router instance
|
||||
model: Model name to create vector store for
|
||||
request_data: Request data for vector store creation
|
||||
litellm_parent_otel_span: OpenTelemetry span for tracing
|
||||
|
||||
|
||||
Returns:
|
||||
VectorStoreCreateResponse from the provider
|
||||
"""
|
||||
# Use the router to create the vector store
|
||||
response = await llm_router.avector_store_create(
|
||||
model=model, **request_data
|
||||
)
|
||||
response = await llm_router.avector_store_create(model=model, **request_data)
|
||||
return response
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -139,14 +137,14 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> VectorStoreCreateResponse:
|
||||
"""
|
||||
Create a vector store across multiple models.
|
||||
|
||||
|
||||
Args:
|
||||
create_request: Vector store creation request parameters
|
||||
llm_router: LiteLLM router instance
|
||||
target_model_names_list: List of target model names
|
||||
litellm_parent_otel_span: OpenTelemetry span for tracing
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
VectorStoreCreateResponse with unified ID
|
||||
"""
|
||||
|
|
@ -196,7 +194,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# VectorStoreCreateResponse is a TypedDict, so we need to create a new dict with the unified ID
|
||||
response = responses[0].copy()
|
||||
response["id"] = unified_id
|
||||
|
||||
|
||||
verbose_logger.info(
|
||||
f"Successfully created managed vector store with unified ID: {unified_id}"
|
||||
)
|
||||
|
|
@ -212,13 +210,13 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> Dict[str, Any]:
|
||||
"""
|
||||
List vector stores created by a user.
|
||||
|
||||
|
||||
Args:
|
||||
user_api_key_dict: User API key authentication details
|
||||
limit: Maximum number of vector stores to return
|
||||
after: Cursor for pagination
|
||||
order: Sort order ('asc' or 'desc')
|
||||
|
||||
|
||||
Returns:
|
||||
Dictionary with list of vector stores and pagination info
|
||||
"""
|
||||
|
|
@ -238,23 +236,23 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> bool:
|
||||
"""
|
||||
Check if user has access to a vector store.
|
||||
|
||||
|
||||
Args:
|
||||
vector_store_id: The unified vector store ID
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
True if user has access, False otherwise
|
||||
"""
|
||||
is_unified_id = is_base64_encoded_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if is_unified_id:
|
||||
# Check access for managed vector store
|
||||
return await self.can_user_access_unified_resource_id(
|
||||
vector_store_id,
|
||||
user_api_key_dict,
|
||||
)
|
||||
|
||||
|
||||
# Not a managed vector store, allow access
|
||||
return True
|
||||
|
||||
|
|
@ -263,24 +261,22 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> bool:
|
||||
"""
|
||||
Check if user has access to a managed vector store in request data.
|
||||
|
||||
|
||||
Args:
|
||||
data: Request data containing vector_store_id
|
||||
user_api_key_dict: User API key authentication details
|
||||
|
||||
|
||||
Returns:
|
||||
True if this is a managed vector store and user has access
|
||||
|
||||
|
||||
Raises:
|
||||
HTTPException: If user doesn't have access
|
||||
"""
|
||||
vector_store_id = cast(Optional[str], data.get("vector_store_id"))
|
||||
is_unified_id = (
|
||||
is_base64_encoded_unified_id(vector_store_id)
|
||||
if vector_store_id
|
||||
else False
|
||||
is_base64_encoded_unified_id(vector_store_id) if vector_store_id else False
|
||||
)
|
||||
|
||||
|
||||
if is_unified_id and vector_store_id:
|
||||
if await self.can_user_access_unified_resource_id(
|
||||
vector_store_id, user_api_key_dict
|
||||
|
|
@ -291,7 +287,7 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
status_code=403,
|
||||
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
|
||||
)
|
||||
|
||||
|
||||
return False
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -307,18 +303,18 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> Union[Exception, str, Dict, None]:
|
||||
"""
|
||||
Pre-call hook to handle vector store operations.
|
||||
|
||||
|
||||
This hook intercepts vector store requests and:
|
||||
- Validates access for managed vector stores
|
||||
- Transforms unified IDs to provider-specific IDs
|
||||
- Adds model routing information
|
||||
|
||||
|
||||
Args:
|
||||
user_api_key_dict: User API key authentication details
|
||||
cache: Cache instance
|
||||
data: Request data
|
||||
call_type: Type of call being made
|
||||
|
||||
|
||||
Returns:
|
||||
Modified request data or None
|
||||
"""
|
||||
|
|
@ -330,40 +326,40 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
# Handle vector store search operations
|
||||
if call_type == "avector_store_search":
|
||||
vector_store_id = data.get("vector_store_id")
|
||||
|
||||
|
||||
if vector_store_id:
|
||||
# Check if it's a managed vector store ID
|
||||
decoded_id = is_base64_encoded_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if decoded_id:
|
||||
verbose_logger.debug(
|
||||
f"Processing managed vector store search: {vector_store_id}"
|
||||
)
|
||||
|
||||
|
||||
# Check access
|
||||
has_access = await self.can_user_access_unified_resource_id(
|
||||
vector_store_id, user_api_key_dict
|
||||
)
|
||||
|
||||
|
||||
if not has_access:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail=f"User {user_api_key_dict.user_id} does not have access to vector store {vector_store_id}",
|
||||
)
|
||||
|
||||
|
||||
# Parse the unified ID to extract components
|
||||
parsed_id = parse_unified_id(vector_store_id)
|
||||
|
||||
|
||||
if parsed_id:
|
||||
# Extract the model ID and provider resource ID
|
||||
model_id = parsed_id.get("model_id")
|
||||
provider_resource_id = parsed_id.get("provider_resource_id")
|
||||
target_model_names = parsed_id.get("target_model_names", [])
|
||||
|
||||
|
||||
verbose_logger.debug(
|
||||
f"Decoded vector store - model_id: {model_id}, provider_resource_id: {provider_resource_id}, target_model_names: {target_model_names}"
|
||||
)
|
||||
|
||||
|
||||
# Determine which model to use for routing
|
||||
# Priority: model_id (deployment ID) > first target_model_name
|
||||
routing_model = None
|
||||
|
|
@ -371,28 +367,28 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
routing_model = model_id
|
||||
elif target_model_names and len(target_model_names) > 0:
|
||||
routing_model = target_model_names[0]
|
||||
|
||||
|
||||
# Set the model for routing
|
||||
if routing_model:
|
||||
data["model"] = routing_model
|
||||
verbose_logger.info(
|
||||
f"Routing vector store search to model: {routing_model}"
|
||||
)
|
||||
|
||||
|
||||
# Replace the unified ID with the provider-specific ID
|
||||
if provider_resource_id:
|
||||
data["vector_store_id"] = provider_resource_id
|
||||
verbose_logger.debug(
|
||||
f"Replaced unified ID with provider resource ID: {provider_resource_id}"
|
||||
)
|
||||
|
||||
|
||||
# Handle vector store retrieve/delete operations
|
||||
elif call_type in ("avector_store_retrieve", "avector_store_delete"):
|
||||
await self.check_managed_vector_store_access(data, user_api_key_dict)
|
||||
|
||||
|
||||
# If it's a managed vector store, we'll handle it in the endpoint
|
||||
# No need to transform here as the endpoint will route to the hook
|
||||
|
||||
|
||||
return data
|
||||
|
||||
# ============================================================================
|
||||
|
|
@ -407,15 +403,15 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> Any:
|
||||
"""
|
||||
Post-call hook to transform responses.
|
||||
|
||||
|
||||
This hook can be used to transform responses if needed.
|
||||
For now, it just passes through the response.
|
||||
|
||||
|
||||
Args:
|
||||
data: Request data
|
||||
user_api_key_dict: User API key authentication details
|
||||
response: Response from the provider
|
||||
|
||||
|
||||
Returns:
|
||||
Potentially modified response
|
||||
"""
|
||||
|
|
@ -436,21 +432,21 @@ class _PROXY_LiteLLMManagedVectorStores(
|
|||
) -> List[Dict]:
|
||||
"""
|
||||
Filter deployments based on vector store availability.
|
||||
|
||||
|
||||
This is used by the router to select only deployments that have
|
||||
the vector store available.
|
||||
|
||||
|
||||
Note: This method signature is a compromise between CustomLogger and BaseManagedResource
|
||||
parent classes which have incompatible signatures. The type: ignore[override] is necessary
|
||||
due to this multiple inheritance conflict.
|
||||
|
||||
|
||||
Args:
|
||||
model: Model name
|
||||
healthy_deployments: List of healthy deployments
|
||||
messages: Messages (unused for vector stores, required by CustomLogger interface)
|
||||
request_kwargs: Request kwargs containing vector_store_id and mappings
|
||||
parent_otel_span: OpenTelemetry span for tracing
|
||||
|
||||
|
||||
Returns:
|
||||
Filtered list of deployments
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -36,7 +36,6 @@ def apply_enterprise_key_management_params(
|
|||
data: GenerateKeyRequest,
|
||||
team_table: Optional[LiteLLM_TeamTable],
|
||||
) -> GenerateKeyRequest:
|
||||
|
||||
data = add_team_member_key_duration(team_table, data)
|
||||
data = add_team_organization_id(team_table, data)
|
||||
return data
|
||||
|
|
|
|||
|
|
@ -147,12 +147,12 @@ async def list_vector_stores(
|
|||
vector_stores_from_db = await VectorStoreRegistry._get_vector_stores_from_db(
|
||||
prisma_client=prisma_client
|
||||
)
|
||||
|
||||
|
||||
# Also clean up in-memory registry to remove any deleted vector stores
|
||||
if litellm.vector_store_registry is not None:
|
||||
db_vector_store_ids = {
|
||||
vs.get("vector_store_id")
|
||||
for vs in vector_stores_from_db
|
||||
vs.get("vector_store_id")
|
||||
for vs in vector_stores_from_db
|
||||
if vs.get("vector_store_id")
|
||||
}
|
||||
# Remove any in-memory vector stores that no longer exist in database
|
||||
|
|
|
|||
|
|
@ -39,15 +39,23 @@ class EmailEvent(str, enum.Enum):
|
|||
soft_budget_crossed = "Soft Budget Crossed"
|
||||
max_budget_alert = "Max Budget Alert"
|
||||
|
||||
|
||||
class EmailEventSettings(BaseModel):
|
||||
event: EmailEvent
|
||||
enabled: bool
|
||||
|
||||
|
||||
class EmailEventSettingsUpdateRequest(BaseModel):
|
||||
settings: List[EmailEventSettings]
|
||||
|
||||
|
||||
class EmailEventSettingsResponse(BaseModel):
|
||||
settings: List[EmailEventSettings]
|
||||
|
||||
|
||||
class DefaultEmailSettings(BaseModel):
|
||||
"""Default settings for email events"""
|
||||
|
||||
settings: Dict[EmailEvent, bool] = Field(
|
||||
default_factory=lambda: {
|
||||
EmailEvent.virtual_key_created: True, # On by default
|
||||
|
|
@ -57,10 +65,12 @@ class DefaultEmailSettings(BaseModel):
|
|||
EmailEvent.max_budget_alert: True, # On by default
|
||||
}
|
||||
)
|
||||
|
||||
def to_dict(self) -> Dict[str, bool]:
|
||||
"""Convert to dictionary with string keys for storage"""
|
||||
return {event.value: enabled for event, enabled in self.settings.items()}
|
||||
|
||||
@classmethod
|
||||
def get_defaults(cls) -> Dict[str, bool]:
|
||||
"""Get the default settings as a dictionary with string keys"""
|
||||
return cls().to_dict()
|
||||
return cls().to_dict()
|
||||
|
|
|
|||
|
|
@ -9,8 +9,8 @@ from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Uni
|
|||
from packaging.version import Version
|
||||
from typing_extensions import TypeAlias
|
||||
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.prompt_management_base import PromptManagementClient
|
||||
from ..custom_prompt_management import CustomPromptManagement
|
||||
from litellm.litellm_core_utils.asyncify import run_async_function
|
||||
from litellm.types.llms.openai import AllMessageValues, ChatCompletionSystemMessage
|
||||
from litellm.types.prompts.init_prompts import PromptSpec
|
||||
|
|
@ -19,7 +19,6 @@ from litellm.types.utils import StandardCallbackDynamicParams, StandardLoggingPa
|
|||
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
|
||||
|
||||
|
|
@ -108,7 +107,7 @@ def langfuse_client_init(
|
|||
return client
|
||||
|
||||
|
||||
class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogger):
|
||||
class LangfusePromptManagement(CustomPromptManagement, LangFuseLogger):
|
||||
def __init__(
|
||||
self,
|
||||
langfuse_public_key=None,
|
||||
|
|
@ -175,6 +174,47 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
optional_params[k] = v
|
||||
return optional_params
|
||||
|
||||
def get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
messages: List[AllMessageValues],
|
||||
non_default_params: dict,
|
||||
prompt_id: Optional[str],
|
||||
prompt_variables: Optional[dict],
|
||||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
prompt_spec: Optional[PromptSpec] = None,
|
||||
prompt_label: Optional[str] = None,
|
||||
prompt_version: Optional[int] = None,
|
||||
ignore_prompt_manager_model: Optional[bool] = False,
|
||||
ignore_prompt_manager_optional_params: Optional[bool] = False,
|
||||
tools: Optional[List[Dict]] = None,
|
||||
) -> Tuple[str, List[AllMessageValues], dict]:
|
||||
"""
|
||||
Overrides the base method to ensure Langfuse-specific prompt management is triggered.
|
||||
"""
|
||||
|
||||
if prompt_id is None and not model.startswith("langfuse/"):
|
||||
return model, messages, non_default_params
|
||||
|
||||
prompt_template = self.compile_prompt(
|
||||
prompt_id=prompt_id or model.replace("langfuse/", ""),
|
||||
prompt_variables=prompt_variables,
|
||||
client_messages=messages,
|
||||
dynamic_callback_params=dynamic_callback_params,
|
||||
prompt_label=prompt_label,
|
||||
prompt_version=prompt_version,
|
||||
prompt_spec=prompt_spec,
|
||||
)
|
||||
|
||||
return self.post_compile_prompt_processing(
|
||||
prompt_template=prompt_template,
|
||||
messages=messages,
|
||||
non_default_params=non_default_params,
|
||||
model=model,
|
||||
ignore_prompt_manager_model=ignore_prompt_manager_model,
|
||||
ignore_prompt_manager_optional_params=ignore_prompt_manager_optional_params,
|
||||
)
|
||||
|
||||
async def async_get_chat_completion_prompt(
|
||||
self,
|
||||
model: str,
|
||||
|
|
@ -212,18 +252,15 @@ class LangfusePromptManagement(LangFuseLogger, PromptManagementBase, CustomLogge
|
|||
dynamic_callback_params: StandardCallbackDynamicParams,
|
||||
) -> bool:
|
||||
if prompt_id is None:
|
||||
return False
|
||||
langfuse_client = langfuse_client_init(
|
||||
langfuse_public_key=dynamic_callback_params.get("langfuse_public_key"),
|
||||
langfuse_secret=dynamic_callback_params.get("langfuse_secret"),
|
||||
langfuse_secret_key=dynamic_callback_params.get("langfuse_secret_key"),
|
||||
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
|
||||
return True
|
||||
|
||||
if (
|
||||
prompt_spec is not None
|
||||
and getattr(prompt_spec, "prompt_id", None) is not None
|
||||
):
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def _compile_prompt_helper(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -86,9 +86,9 @@ class OpenAIChatCompletionsHandler(BaseTranslation):
|
|||
if tool_calls_to_check:
|
||||
inputs["tool_calls"] = tool_calls_to_check # type: ignore
|
||||
if messages:
|
||||
inputs["structured_messages"] = (
|
||||
messages # pass the openai /chat/completions messages to the guardrail, as-is
|
||||
)
|
||||
inputs[
|
||||
"structured_messages"
|
||||
] = messages # pass the openai /chat/completions messages to the guardrail, as-is
|
||||
# Pass tools (function definitions) to the guardrail
|
||||
tools = data.get("tools")
|
||||
if tools:
|
||||
|
|
|
|||
|
|
@ -51,7 +51,11 @@ def _build_audit_log_payload(
|
|||
if request_data.updated_at is not None:
|
||||
updated_at = request_data.updated_at.isoformat()
|
||||
|
||||
table_name_str: str = request_data.table_name.value if isinstance(request_data.table_name, LitellmTableNames) else str(request_data.table_name)
|
||||
table_name_str: str = (
|
||||
request_data.table_name.value
|
||||
if isinstance(request_data.table_name, LitellmTableNames)
|
||||
else str(request_data.table_name)
|
||||
)
|
||||
|
||||
return StandardAuditLogPayload(
|
||||
id=request_data.id,
|
||||
|
|
@ -89,7 +93,9 @@ async def _dispatch_audit_log_to_callbacks(
|
|||
|
||||
for callback in litellm.audit_log_callbacks:
|
||||
try:
|
||||
resolved: Optional[CustomLogger] = callback if isinstance(callback, CustomLogger) else None
|
||||
resolved: Optional[CustomLogger] = (
|
||||
callback if isinstance(callback, CustomLogger) else None
|
||||
)
|
||||
if isinstance(callback, str):
|
||||
resolved = _resolve_audit_log_callback(callback)
|
||||
if resolved is None:
|
||||
|
|
@ -138,9 +144,7 @@ async def create_object_audit_log(
|
|||
return
|
||||
|
||||
_changed_by = (
|
||||
litellm_changed_by
|
||||
or user_api_key_dict.user_id
|
||||
or litellm_proxy_admin_name
|
||||
litellm_changed_by or user_api_key_dict.user_id or litellm_proxy_admin_name
|
||||
)
|
||||
|
||||
await create_audit_log_for_update(
|
||||
|
|
|
|||
|
|
@ -1,24 +0,0 @@
|
|||
import pytest
|
||||
from litellm.proxy.prompts.prompt_registry import get_prompt_initializer_from_integrations
|
||||
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
|
||||
|
||||
def test_langfuse_discovery_and_init():
|
||||
"""
|
||||
Tests whether Langfuse is dynamically discovered and can be initialized.
|
||||
This validates the fix in __init__.py at litellm/integrations/langfuse/
|
||||
"""
|
||||
|
||||
registry = get_prompt_initializer_from_integrations()
|
||||
assert "langfuse" in registry, "Error: Langfuse wasn't discovered by the prompt_registry!"
|
||||
|
||||
init_func = registry["langfuse"]
|
||||
params = PromptLiteLLMParams(
|
||||
prompt_integration="langfuse",
|
||||
langfuse_public_key="test-key",
|
||||
langfuse_secret="test-secret"
|
||||
)
|
||||
spec = PromptSpec(prompt_id="test-id", litellm_params=params)
|
||||
|
||||
obj = init_func(params, spec)
|
||||
assert obj is not None, "Initiation function returned None!"
|
||||
assert obj.integration_name == "langfuse"
|
||||
|
|
@ -0,0 +1,33 @@
|
|||
"""
|
||||
Tests whether Langfuse is dynamically discovered and can be initialized without real network calls.
|
||||
"""
|
||||
from unittest.mock import patch
|
||||
from litellm.proxy.prompts.prompt_registry import get_prompt_initializer_from_integrations
|
||||
from litellm.types.prompts.init_prompts import PromptLiteLLMParams, PromptSpec
|
||||
|
||||
def test_langfuse_discovery_and_init():
|
||||
"""
|
||||
Validates the fix in __init__.py by ensuring Langfuse is registered and initialized (mocked).
|
||||
"""
|
||||
# 1. Verify if the integration was registered in the global dictionary
|
||||
registry = get_prompt_initializer_from_integrations()
|
||||
assert "langfuse" in registry, "Error: Langfuse wasn't discovered by the prompt_registry!"
|
||||
|
||||
# 2. Test initialization using MOCK to avoid real network calls
|
||||
# Patch the function that creates the real Langfuse client
|
||||
with patch("litellm.integrations.langfuse.langfuse_prompt_management.langfuse_client_init") as mocked_client:
|
||||
mocked_client.return_value = None # No real client needed for this test
|
||||
|
||||
init_func = registry["langfuse"]
|
||||
params = PromptLiteLLMParams(
|
||||
prompt_integration="langfuse",
|
||||
langfuse_public_key="test-key",
|
||||
langfuse_secret="test-secret"
|
||||
)
|
||||
spec = PromptSpec(prompt_id="test-id", litellm_params=params)
|
||||
|
||||
# Initialize the class (this would call the network, but is now mocked)
|
||||
obj = init_func(params, spec)
|
||||
|
||||
assert obj is not None, "Initiation function returned None!"
|
||||
assert obj.integration_name == "langfuse"
|
||||
Loading…
Add table
Reference in a new issue