diff --git a/enterprise/enterprise_hooks/session_handler.py b/enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py similarity index 73% rename from enterprise/enterprise_hooks/session_handler.py rename to enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py index b9d7eab877..62e1b5b0a7 100644 --- a/enterprise/enterprise_hooks/session_handler.py +++ b/enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py @@ -1,17 +1,19 @@ -from litellm.proxy._types import SpendLogsPayload -from litellm._logging import verbose_proxy_logger -from typing import Optional, List, Union import json -from litellm.types.utils import ModelResponse, Message +from typing import List, Optional, Union, cast + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import SpendLogsPayload +from litellm.responses.litellm_completion_transformation.transformation import ( + ChatCompletionSession, +) +from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionResponseMessage, GenericChatCompletionMessage, ResponseInputParam, ) -from litellm.types.utils import ChatCompletionMessageToolCall -from litellm.responses.utils import ResponsesAPIRequestUtils -from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession +from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse class _ENTERPRISE_ResponsesSessionHandler: @@ -22,9 +24,16 @@ class _ENTERPRISE_ResponsesSessionHandler: """ Return the chat completion message history for a previous response id """ - from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig - all_spend_logs: List[SpendLogsPayload] = await _ENTERPRISE_ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id(previous_response_id) - + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + all_spend_logs: List[ + SpendLogsPayload + ] = await _ENTERPRISE_ResponsesSessionHandler.get_all_spend_logs_for_previous_response_id( + previous_response_id + ) + litellm_session_id: Optional[str] = None if len(all_spend_logs) > 0: litellm_session_id = all_spend_logs[0].get("session_id") @@ -39,14 +48,16 @@ class _ENTERPRISE_ResponsesSessionHandler: ] ] = [] for spend_log in all_spend_logs: - proxy_server_request: Union[str, dict] = spend_log.get("proxy_server_request") or "{}" + proxy_server_request: Union[str, dict] = ( + spend_log.get("proxy_server_request") or "{}" + ) proxy_server_request_dict: Optional[dict] = None response_input_param: Optional[Union[str, ResponseInputParam]] = None if isinstance(proxy_server_request, dict): proxy_server_request_dict = proxy_server_request else: proxy_server_request_dict = json.loads(proxy_server_request) - + ############################################################ # Add Input messages for this Spend Log ############################################################ @@ -55,15 +66,17 @@ class _ENTERPRISE_ResponsesSessionHandler: if isinstance(_response_input_param, str): response_input_param = _response_input_param elif isinstance(_response_input_param, dict): - response_input_param = ResponseInputParam(**_response_input_param) - + response_input_param = cast( + ResponseInputParam, _response_input_param + ) + if response_input_param: chat_completion_messages = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( input=response_input_param, - responses_api_request=proxy_server_request_dict or {} + responses_api_request=proxy_server_request_dict or {}, ) chat_completion_message_history.extend(chat_completion_messages) - + ############################################################ # Add Output messages for this Spend Log ############################################################ @@ -73,17 +86,22 @@ class _ENTERPRISE_ResponsesSessionHandler: model_response = ModelResponse(**_response_output) for choice in model_response.choices: if hasattr(choice, "message"): - chat_completion_message_history.append(choice.message) - - verbose_proxy_logger.debug("chat_completion_message_history %s", json.dumps(chat_completion_message_history, indent=4, default=str)) + chat_completion_message_history.append( + getattr(choice, "message") + ) + + verbose_proxy_logger.debug( + "chat_completion_message_history %s", + json.dumps(chat_completion_message_history, indent=4, default=str), + ) return ChatCompletionSession( messages=chat_completion_message_history, - litellm_session_id=litellm_session_id + litellm_session_id=litellm_session_id, ) @staticmethod async def get_all_spend_logs_for_previous_response_id( - previous_response_id: str + previous_response_id: str, ) -> List[SpendLogsPayload]: """ Get all spend logs for a previous response id @@ -94,8 +112,15 @@ class _ENTERPRISE_ResponsesSessionHandler: SELECT session_id FROM spend_logs WHERE response_id = previous_response_id, SELECT * FROM spend_logs WHERE session_id = session_id """ from litellm.proxy.proxy_server import prisma_client - decoded_response_id = ResponsesAPIRequestUtils._decode_responses_api_response_id(previous_response_id) - previous_response_id = decoded_response_id.get("response_id", previous_response_id) + + decoded_response_id = ( + ResponsesAPIRequestUtils._decode_responses_api_response_id( + previous_response_id + ) + ) + previous_response_id = decoded_response_id.get( + "response_id", previous_response_id + ) if prisma_client is None: return [] @@ -111,21 +136,12 @@ class _ENTERPRISE_ResponsesSessionHandler: ORDER BY "endTime" ASC; """ - spend_logs = await prisma_client.db.query_raw( - query, - previous_response_id - ) + spend_logs = await prisma_client.db.query_raw(query, previous_response_id) verbose_proxy_logger.debug( "Found the following spend logs for previous response id %s: %s", previous_response_id, - json.dumps(spend_logs, indent=4, default=str) + json.dumps(spend_logs, indent=4, default=str), ) - return spend_logs - - - - - diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index 5af761a96d..c8a2eb3899 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -1,4 +1,8 @@ model_list: - - model_name: openai/* + - model_name: anthropic/* litellm_params: - model: openai/* \ No newline at end of file + model: anthropic/* + + +general_settings: + store_prompts_in_spend_logs: true \ No newline at end of file diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index baffca3ac6..c7b8ed88ef 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -7,16 +7,17 @@ from typing import Any, Dict, List, Optional, Union from openai.types.responses.tool_param import FunctionToolParam from typing_extensions import TypedDict -HAS_ENTERPRISE_DIRECTORY = False +from litellm._logging import verbose_logger + try: - from enterprise.enterprise_hooks.session_handler import ( + from litellm_enterprise.enterprise_callbacks.session_handler import ( _ENTERPRISE_ResponsesSessionHandler, ) - - HAS_ENTERPRISE_DIRECTORY = True -except ImportError: - _ENTERPRISE_ResponsesSessionHandler = None # type: ignore - HAS_ENTERPRISE_DIRECTORY = False +except Exception as e: + verbose_logger.debug( + f"[Non-Blocking] Unable to import _ENTERPRISE_ResponsesSessionHandler - LiteLLM Enterprise Feature - {str(e)}" + ) + _ENTERPRISE_ResponsesSessionHandler = None from litellm.caching import InMemoryCache from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -197,10 +198,7 @@ class LiteLLMCompletionResponsesConfig: """ Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages """ - if ( - HAS_ENTERPRISE_DIRECTORY is True - and _ENTERPRISE_ResponsesSessionHandler is not None - ): + if _ENTERPRISE_ResponsesSessionHandler is not None: chat_completion_session = ChatCompletionSession( messages=[], litellm_session_id=None )