Revert "fix: import session handling"

This reverts commit deb257dc10.
This commit is contained in:
Ishaan Jaff 2025-05-29 17:54:24 -07:00
parent deb257dc10
commit 0d4c85c419
3 changed files with 47 additions and 65 deletions

View file

@ -1,19 +1,17 @@
import json
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._logging import verbose_proxy_logger
from typing import Optional, List, Union
import json
from litellm.types.utils import ModelResponse, Message
from litellm.types.llms.openai import (
AllMessageValues,
ChatCompletionResponseMessage,
GenericChatCompletionMessage,
ResponseInputParam,
)
from litellm.types.utils import ChatCompletionMessageToolCall, Message, ModelResponse
from litellm.types.utils import ChatCompletionMessageToolCall
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.responses.litellm_completion_transformation.transformation import ChatCompletionSession
class _ENTERPRISE_ResponsesSessionHandler:
@ -24,16 +22,9 @@ 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")
@ -48,16 +39,14 @@ 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
############################################################
@ -66,17 +55,15 @@ class _ENTERPRISE_ResponsesSessionHandler:
if isinstance(_response_input_param, str):
response_input_param = _response_input_param
elif isinstance(_response_input_param, dict):
response_input_param = cast(
ResponseInputParam, _response_input_param
)
response_input_param = 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
############################################################
@ -86,22 +73,17 @@ class _ENTERPRISE_ResponsesSessionHandler:
model_response = ModelResponse(**_response_output)
for choice in model_response.choices:
if hasattr(choice, "message"):
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),
)
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))
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
@ -112,15 +94,8 @@ 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 []
@ -136,12 +111,21 @@ 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

View file

@ -1,8 +1,4 @@
model_list:
- model_name: anthropic/*
- model_name: openai/*
litellm_params:
model: anthropic/*
general_settings:
store_prompts_in_spend_logs: true
model: openai/*

View file

@ -7,17 +7,16 @@ from typing import Any, Dict, List, Optional, Union
from openai.types.responses.tool_param import FunctionToolParam
from typing_extensions import TypedDict
from litellm._logging import verbose_logger
HAS_ENTERPRISE_DIRECTORY = False
try:
from litellm_enterprise.enterprise_callbacks.session_handler import (
from enterprise.enterprise_hooks.session_handler import (
_ENTERPRISE_ResponsesSessionHandler,
)
except Exception as e:
verbose_logger.debug(
f"[Non-Blocking] Unable to import _ENTERPRISE_ResponsesSessionHandler - LiteLLM Enterprise Feature - {str(e)}"
)
_ENTERPRISE_ResponsesSessionHandler = None
HAS_ENTERPRISE_DIRECTORY = True
except ImportError:
_ENTERPRISE_ResponsesSessionHandler = None # type: ignore
HAS_ENTERPRISE_DIRECTORY = False
from litellm.caching import InMemoryCache
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
@ -198,7 +197,10 @@ class LiteLLMCompletionResponsesConfig:
"""
Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages
"""
if _ENTERPRISE_ResponsesSessionHandler is not None:
if (
HAS_ENTERPRISE_DIRECTORY is True
and _ENTERPRISE_ResponsesSessionHandler is not None
):
chat_completion_session = ChatCompletionSession(
messages=[], litellm_session_id=None
)