fix(o1_transformation.py): return tool calling/response_format in sup…

…ported params if model map says so Fixes #7292
BerriAI · Dec 18, 2024 · 0715ccc · 0715ccc
1 parent 9a0d6db
commit 0715ccc
Show file tree

Hide file tree

Showing 7 changed files with 188 additions and 172 deletions.
diff --git a/litellm/llms/ollama/completion/transformation.py b/litellm/llms/ollama/completion/transformation.py
@@ -23,6 +23,7 @@
 from litellm.types.utils import (
     GenericStreamingChunk,
     ModelInfo,
+    ModelInfoBase,
     ModelResponse,
     ProviderField,
     StreamingChoices,
@@ -198,7 +199,7 @@ def _get_max_tokens(self, ollama_model_info: dict) -> Optional[int]:
                 return v
         return None
 
-    def get_model_info(self, model: str) -> ModelInfo:
+    def get_model_info(self, model: str) -> ModelInfoBase:
         """
         curl http://localhost:11434/api/show -d '{
           "name": "mistral"
@@ -222,11 +223,10 @@ def get_model_info(self, model: str) -> ModelInfo:
 
         _max_tokens: Optional[int] = self._get_max_tokens(model_info)
 
-        return ModelInfo(
+        return ModelInfoBase(
             key=model,
             litellm_provider="ollama",
             mode="chat",
-            supported_openai_params=self.get_supported_openai_params(model=model),
             supports_function_calling=self._supports_function_calling(model_info),
             input_cost_per_token=0.0,
             output_cost_per_token=0.0,

diff --git a/litellm/llms/openai/chat/o1_transformation.py b/litellm/llms/openai/chat/o1_transformation.py
@@ -16,7 +16,11 @@
 
 import litellm
 from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage
-from litellm.utils import supports_system_messages
+from litellm.utils import (
+    supports_function_calling,
+    supports_response_schema,
+    supports_system_messages,
+)
 
 from .gpt_transformation import OpenAIGPTConfig
 
@@ -39,26 +43,34 @@ def get_supported_openai_params(self, model: str) -> list:
         all_openai_params = super().get_supported_openai_params(model=model)
         non_supported_params = [
             "logprobs",
-            "tools",
-            "tool_choice",
-            "parallel_tool_calls",
-            "function_call",
-            "functions",
             "top_p",
             "presence_penalty",
             "frequency_penalty",
             "top_logprobs",
         ]
 
         supported_streaming_models = ["o1-preview", "o1-mini"]
+        _supports_function_calling = supports_function_calling(model, "openai")
+        _supports_response_schema = supports_response_schema(model, "openai")
 
         if model not in supported_streaming_models:
             non_supported_params.append("stream")
             non_supported_params.append("stream_options")
 
-        return [
+        if not _supports_function_calling:
+            non_supported_params.append("tools")
+            non_supported_params.append("tool_choice")
+            non_supported_params.append("parallel_tool_calls")
+            non_supported_params.append("function_call")
+            non_supported_params.append("functions")
+
+        if not _supports_response_schema:
+            non_supported_params.append("response_format")
+
+        returned_params = [
             param for param in all_openai_params if param not in non_supported_params
         ]
+        return returned_params
 
     def map_openai_params(
         self,

diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
@@ -137,7 +137,8 @@
         "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
     "o1-mini": {
         "max_tokens": 65536,
@@ -148,8 +149,6 @@
         "cache_read_input_token_cost": 0.0000015,
         "litellm_provider": "openai",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true
     },
@@ -162,8 +161,6 @@
         "cache_read_input_token_cost": 0.0000015,
         "litellm_provider": "openai",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true
     },
@@ -176,8 +173,6 @@
         "cache_read_input_token_cost": 0.0000075,
         "litellm_provider": "openai",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true
     },
@@ -190,8 +185,6 @@
         "cache_read_input_token_cost": 0.0000075,
         "litellm_provider": "openai",
         "mode": "chat",
-        "supports_function_calling": true,
-        "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true
     },
@@ -208,7 +201,8 @@
         "supports_parallel_function_calling": true,
         "supports_vision": false,
         "supports_prompt_caching": true,
-        "supports_system_messages": true
+        "supports_system_messages": true,
+        "supports_response_schema": true
     },
     "chatgpt-4o-latest": {
         "max_tokens": 4096,

diff --git a/litellm/types/utils.py b/litellm/types/utils.py
@@ -74,11 +74,7 @@ class ProviderField(TypedDict):
     field_value: str
 
 
-class ModelInfo(TypedDict, total=False):
-    """
-    Model info for a given model, this is information found in litellm.model_prices_and_context_window.json
-    """
-
+class ModelInfoBase(TypedDict, total=False):
     key: Required[str]  # the key in litellm.model_cost which is returned
 
     max_tokens: Required[Optional[int]]
@@ -119,7 +115,6 @@ class ModelInfo(TypedDict, total=False):
             "completion", "embedding", "image_generation", "chat", "audio_transcription"
         ]
     ]
-    supported_openai_params: Required[Optional[List[str]]]
     supports_system_messages: Optional[bool]
     supports_response_schema: Optional[bool]
     supports_vision: Optional[bool]
@@ -133,6 +128,14 @@ class ModelInfo(TypedDict, total=False):
     rpm: Optional[int]
 
 
+class ModelInfo(ModelInfoBase, total=False):
+    """
+    Model info for a given model, this is information found in litellm.model_prices_and_context_window.json
+    """
+
+    supported_openai_params: Required[Optional[List[str]]]
+
+
 class GenericStreamingChunk(TypedDict, total=False):
     text: Required[str]
     tool_use: Optional[ChatCompletionToolCallChunk]