x

eyurtsev · Sep 16, 2024 · 6581734 · 6581734
1 parent 55d02ef
commit 6581734
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 34 deletions.
diff --git a/kor/extraction/parser.py b/kor/extraction/parser.py
@@ -3,7 +3,7 @@
 from typing import List, Optional
 
 from langchain_core.output_parsers import BaseOutputParser
-from pydantic import Extra
+from pydantic import ConfigDict
 
 from kor.encoders import Encoder
 from kor.exceptions import ParseError
@@ -22,6 +22,7 @@ class KorParser(BaseOutputParser[Extraction]):
     encoder: Encoder
     schema_: Object
     validator: Optional[Validator] = None
+    model_config = ConfigDict(arbitrary_types_allowed=True)
 
     @property
     def _type(self) -> str:
@@ -66,8 +67,7 @@ def parse(self, text: str) -> Extraction:
             "validated_data": validated_data,
         }
 
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(
+        extra="forbid",
+        arbitrary_types_allowed=True,
+    )
diff --git a/kor/prompts.py b/kor/prompts.py
@@ -6,6 +6,7 @@
 from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
 from langchain_core.prompt_values import PromptValue
 from langchain_core.prompts import BasePromptTemplate, PromptTemplate
+from pydantic import ConfigDict
 
 from kor.encoders import Encoder
 from kor.encoders.encode import InputFormatter, encode_examples, format_text
@@ -14,12 +15,6 @@
 from kor.nodes import Object
 from kor.type_descriptors import TypeDescriptor
 
-try:
-    # Use pydantic v1 namespace since working with langchain
-    from pydantic.v1 import Extra  # type: ignore[assignment]
-except ImportError:
-    from pydantic import Extra  # type: ignore[assignment]
-
 from .validators import Validator
 
 DEFAULT_INSTRUCTION_TEMPLATE = PromptTemplate(
@@ -41,11 +36,10 @@ class ExtractionPromptValue(PromptValue):
     string: str
     messages: List[BaseMessage]
 
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(
+        extra="forbid",
+        arbitrary_types_allowed=True,
+    )
 
     def to_string(self) -> str:
         """Format the prompt to a string."""
@@ -65,11 +59,10 @@ class ExtractionPromptTemplate(BasePromptTemplate):
     input_formatter: InputFormatter
     instruction_template: PromptTemplate
 
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(
+        extra="forbid",
+        arbitrary_types_allowed=True,
+    )
 
     def format_prompt(  # type: ignore[override]
         self,

diff --git a/tests/utils.py b/tests/utils.py
@@ -7,23 +7,16 @@
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage
 from langchain_core.outputs import ChatGeneration, ChatResult
-
-from kor._pydantic import PYDANTIC_MAJOR_VERSION
-
-if PYDANTIC_MAJOR_VERSION == 1:
-    from pydantic import Extra  # type: ignore[assignment]
-else:
-    from pydantic.v1 import Extra  # type: ignore[assignment,no-redef]
+from pydantic import ConfigDict
 
 
 class ToyChatModel(BaseChatModel):
     response: str
 
-    class Config:
-        """Configuration for this pydantic object."""
-
-        extra = Extra.forbid
-        arbitrary_types_allowed = True
+    model_config = ConfigDict(
+        extra="forbid",
+        arbitrary_types_allowed=True,
+    )
 
     def _generate(
         self,