tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-14 14:26:43 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: import session handling

Browse Source
This commit is contained in:
Ishaan Jaff
2025-05-29 17:48:34 -07:00
parent 5e6f6ddc52
commit deb257dc10
3 changed files with 65 additions and 47 deletions
Download Patch File Download Diff File Expand all files Collapse all files
enterprise/enterprise_hooks/session_handler.py → enterprise/litellm_enterprise/enterprise_callbacks/session_handler.py
+50 -34
View File
@@ -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
litellm/proxy/proxy_config.yaml
+6 -2
View File
@@ -1,4 +1,8 @@
model_list:
- model_name: openai/*
- model_name: anthropic/*
litellm_params:
model: openai/*
model: anthropic/*
general_settings:
store_prompts_in_spend_logs: true
litellm/responses/litellm_completion_transformation/transformation.py
+9 -11
View File
@@ -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
)