From 1e0ee2244543ff2fea0414467c71aca7813bde07 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Sat, 26 Apr 2025 19:04:59 -0700
Subject: [PATCH] fix build from pip tests

---
 .../enterprise_hooks/session_handler.py       |  6 +--
 .../transformation.py                         | 40 +++++++++++--------
 litellm/types/responses/main.py               | 31 +++++++++++++-
 3 files changed, 54 insertions(+), 23 deletions(-)

diff --git a/enterprise/enterprise_hooks/session_handler.py b/enterprise/enterprise_hooks/session_handler.py
index e3756c386bf..1ac419a7a1c 100644
--- a/enterprise/enterprise_hooks/session_handler.py
+++ b/enterprise/enterprise_hooks/session_handler.py
@@ -11,13 +11,9 @@ from litellm.types.llms.openai import (
     ResponseInputParam,
 )
 from litellm.types.utils import ChatCompletionMessageToolCall
-
+from litellm.types.responses.main import ChatCompletionSession
 from litellm.responses.utils import ResponsesAPIRequestUtils
-from typing import TypedDict
 
-class ChatCompletionSession(TypedDict, total=False):
-    messages: List[Union[AllMessageValues, GenericChatCompletionMessage, ChatCompletionMessageToolCall, ChatCompletionResponseMessage, Message]]
-    litellm_session_id: Optional[str]
 
 class _ENTERPRISE_ResponsesSessionHandler:
     @staticmethod
diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py
index 719fc204a67..e9e6aff26b4 100644
--- a/litellm/responses/litellm_completion_transformation/transformation.py
+++ b/litellm/responses/litellm_completion_transformation/transformation.py
@@ -6,10 +6,17 @@ from typing import Any, Dict, List, Optional, Union
 
 from openai.types.responses.tool_param import FunctionToolParam
 
-from enterprise.enterprise_hooks.session_handler import (
-    ChatCompletionSession,
-    _ENTERPRISE_ResponsesSessionHandler,
-)
+from litellm.types.responses.main import ChatCompletionSession
+
+try:
+    import enterprise
+    from enterprise.enterprise_hooks.session_handler import (
+        _ENTERPRISE_ResponsesSessionHandler,
+    )
+except ImportError:
+    enterprise = None
+    _ENTERPRISE_ResponsesSessionHandler = None
+
 from litellm.caching import InMemoryCache
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.types.llms.openai import (
@@ -174,19 +181,20 @@ class LiteLLMCompletionResponsesConfig:
         """
         Async hook to get the chain of previous input and output pairs and return a list of Chat Completion messages
         """
-        chat_completion_session: ChatCompletionSession = ChatCompletionSession(
-            messages=[], litellm_session_id=None
-        )
-        if previous_response_id:
-            chat_completion_session = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
-                previous_response_id=previous_response_id
+        if _ENTERPRISE_ResponsesSessionHandler and ChatCompletionSession:
+            chat_completion_session = ChatCompletionSession(
+                messages=[], litellm_session_id=None
+            )
+            if previous_response_id:
+                chat_completion_session = await _ENTERPRISE_ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
+                    previous_response_id=previous_response_id
+                )
+            _messages = litellm_completion_request.get("messages") or []
+            session_messages = chat_completion_session.get("messages") or []
+            litellm_completion_request["messages"] = session_messages + _messages
+            litellm_completion_request["litellm_trace_id"] = (
+                chat_completion_session.get("litellm_session_id")
             )
-        _messages = litellm_completion_request.get("messages") or []
-        session_messages = chat_completion_session.get("messages") or []
-        litellm_completion_request["messages"] = session_messages + _messages
-        litellm_completion_request["litellm_trace_id"] = chat_completion_session.get(
-            "litellm_session_id"
-        )
         return litellm_completion_request
 
     @staticmethod
diff --git a/litellm/types/responses/main.py b/litellm/types/responses/main.py
index b85df206bc6..a73a7780785 100644
--- a/litellm/types/responses/main.py
+++ b/litellm/types/responses/main.py
@@ -1,10 +1,24 @@
-from typing import Literal
+from typing import List, Literal, Optional, Union
 
 from pydantic import PrivateAttr
-from typing_extensions import Any, List, Optional, TypedDict
+from typing_extensions import TYPE_CHECKING, Any, List, Optional, TypedDict
 
 from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject
 
+if TYPE_CHECKING:
+    from litellm.types.llms.openai import (
+        AllMessageValues,
+        ChatCompletionResponseMessage,
+        GenericChatCompletionMessage,
+    )
+    from litellm.types.utils import ChatCompletionMessageToolCall, Message
+else:
+    Message = Any
+    ChatCompletionMessageToolCall = Any
+    AllMessageValues = Any
+    ChatCompletionResponseMessage = Any
+    GenericChatCompletionMessage = Any
+
 
 class GenericResponseOutputItemContentAnnotation(BaseLiteLLMOpenAIResponseObject):
     """Annotation for content in a message"""
@@ -74,3 +88,16 @@ class DecodedResponseId(TypedDict, total=False):
     custom_llm_provider: Optional[str]
     model_id: Optional[str]
     response_id: str
+
+
+class ChatCompletionSession(TypedDict, total=False):
+    messages: List[
+        Union[
+            AllMessageValues,
+            GenericChatCompletionMessage,
+            ChatCompletionMessageToolCall,
+            ChatCompletionResponseMessage,
+            Message,
+        ]
+    ]
+    litellm_session_id: Optional[str]