fmt

langchain-ai · Aug 7, 2024 · 7416720 · 7416720
1 parent 839b9b0
commit 7416720
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 32 deletions.
diff --git a/libs/partners/openai/langchain_openai/chat_models/base.py b/libs/partners/openai/langchain_openai/chat_models/base.py
@@ -28,7 +28,6 @@
     TypeVar,
     Union,
     cast,
-    overload,
 )
 from urllib.parse import urlparse
 
@@ -1100,33 +1099,6 @@ def bind_tools(
             kwargs["tool_choice"] = tool_choice
         return super().bind(tools=formatted_tools, **kwargs)
 
-    # TODO: Fix typing.
-    @overload  # type: ignore[override]
-    def with_structured_output(
-        self,
-        schema: Optional[_DictOrPydanticClass] = None,
-        *,
-        method: Literal[
-            "function_calling", "json_mode", "json_schema"
-        ] = "function_calling",
-        include_raw: Literal[True] = True,
-        strict: Optional[bool] = None,
-        **kwargs: Any,
-    ) -> Runnable[LanguageModelInput, _AllReturnType]: ...
-
-    @overload
-    def with_structured_output(
-        self,
-        schema: Optional[_DictOrPydanticClass] = None,
-        *,
-        method: Literal[
-            "function_calling", "json_mode", "json_schema"
-        ] = "function_calling",
-        include_raw: Literal[False] = False,
-        strict: Optional[bool] = None,
-        **kwargs: Any,
-    ) -> Runnable[LanguageModelInput, _DictOrPydantic]: ...
-
     def with_structured_output(
         self,
         schema: Optional[_DictOrPydanticClass] = None,

diff --git a/libs/partners/openai/tests/unit_tests/chat_models/test_base.py b/libs/partners/openai/tests/unit_tests/chat_models/test_base.py
@@ -1,7 +1,7 @@
 """Test OpenAI Chat API wrapper."""
 
 import json
-from typing import Any, List, Literal, Optional, Type, Union
+from typing import Any, Dict, List, Literal, Optional, Type, Union
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
@@ -357,8 +357,8 @@ def test_bind_tools_tool_choice(tool_choice: Any, strict: Optional[bool]) -> Non
 @pytest.mark.parametrize("include_raw", [True, False])
 @pytest.mark.parametrize("strict", [True, False, None])
 def test_with_structured_output(
-    schema: Union[Type[BaseModel], dict],
-    method: Literal["json_schema", "function_calling", "json_mode"],
+    schema: Union[Type, Dict[str, Any], None],
+    method: Literal["function_calling", "json_mode", "json_schema"],
     include_raw: bool,
     strict: Optional[bool],
 ) -> None:
@@ -367,7 +367,7 @@ def test_with_structured_output(
         strict = None
     llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
     llm.with_structured_output(
-        schema, method=method, include_raw=include_raw, strict=strict
+        schema, method=method, strict=strict, include_raw=include_raw
     )