From 3a8ad9050935d36e06c940df2257ae6e1cb37f48 Mon Sep 17 00:00:00 2001
From: Eugene Yurtsev <eyurtsev@gmail.com>
Date: Mon, 8 Jan 2024 16:53:10 -0500
Subject: [PATCH] langchain(patch): Fix output type for pydantic output parser
 (#15714)

This PR fixes the output type for the pydantic output parser.

Fix for: https://github.com/langchain-ai/langserve/issues/301
---
 .../langchain/output_parsers/pydantic.py      | 11 +++++++-
 .../output_parsers/test_pydantic_parser.py    | 25 +++++++++++++++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/libs/langchain/langchain/output_parsers/pydantic.py b/libs/langchain/langchain/output_parsers/pydantic.py
index b4fe96b5644f3..00424018531a1 100644
--- a/libs/langchain/langchain/output_parsers/pydantic.py
+++ b/libs/langchain/langchain/output_parsers/pydantic.py
@@ -15,7 +15,11 @@ class PydanticOutputParser(BaseOutputParser[T]):
     """Parse an output using a pydantic model."""
 
     pydantic_object: Type[T]
-    """The pydantic model to parse."""
+    """The pydantic model to parse.
+    
+    Attention: To avoid potential compatibility issues, it's recommended to use
+        pydantic <2 or leverage the v1 namespace in pydantic >= 2.
+    """
 
     def parse(self, text: str) -> T:
         try:
@@ -51,3 +55,8 @@ def get_format_instructions(self) -> str:
     @property
     def _type(self) -> str:
         return "pydantic"
+
+    @property
+    def OutputType(self) -> Type[T]:
+        """Return the pydantic model."""
+        return self.pydantic_object
diff --git a/libs/langchain/tests/unit_tests/output_parsers/test_pydantic_parser.py b/libs/langchain/tests/unit_tests/output_parsers/test_pydantic_parser.py
index 50565292857e0..9230792f7641f 100644
--- a/libs/langchain/tests/unit_tests/output_parsers/test_pydantic_parser.py
+++ b/libs/langchain/tests/unit_tests/output_parsers/test_pydantic_parser.py
@@ -76,3 +76,28 @@ def test_pydantic_output_parser_fail() -> None:
         assert "Failed to parse TestModel from completion" in str(e)
     else:
         assert False, "Expected OutputParserException"
+
+
+def test_pydantic_output_parser_type_inference() -> None:
+    """Test pydantic output parser type inference."""
+
+    class SampleModel(BaseModel):
+        foo: int
+        bar: str
+
+    # Ignoring mypy error that appears in python 3.8, but not 3.11.
+    # This seems to be functionally correct, so we'll ignore the error.
+    pydantic_parser = PydanticOutputParser(
+        pydantic_object=SampleModel  # type: ignore[var-annotated]
+    )
+    schema = pydantic_parser.get_output_schema().schema()
+
+    assert schema == {
+        "properties": {
+            "bar": {"title": "Bar", "type": "string"},
+            "foo": {"title": "Foo", "type": "integer"},
+        },
+        "required": ["foo", "bar"],
+        "title": "SampleModel",
+        "type": "object",
+    }