Updates for pydantic 2 compatibility

Makefile

@@ -7,7 +7,7 @@ EVAL_DIR = src/$(PACKAGE)/evaluation
 TEMPLATES = $(notdir $(basename $(wildcard $(TEMPLATE_DIR)/*.yaml)))
 ENTRY_CLASSES = recipe.Recipe gocam.GoCamAnnotations reaction.ReactionDocument ctd.ChemicalToDiseaseDocument
 
-all: all_pydantic all_projects
+all: all_pydantic
 
 all_pydantic: $(patsubst %, $(TEMPLATE_DIR)/%.py, $(TEMPLATES))
 all_projects: $(patsubst %, projects/%, $(TEMPLATES))
@@ -28,10 +28,10 @@ get_version:
 	$(RUN) python -c "import ontogpt;print('.'.join((ontogpt.__version__).split('.', 3)[:3]))"
 
 $(TEMPLATE_DIR)/%.py: src/$(PACKAGE)/templates/%.yaml
-	$(RUN) gen-pydantic $< > [email protected] && mv [email protected] $@
+	$(RUN) gen-pydantic --pydantic_version 2 $< > [email protected] && mv [email protected] $@
 
 %.py: %.yaml
-	$(RUN) gen-pydantic $< > $@
+	$(RUN) gen-pydantic --pydantic_version 2 $< > $@
 
 #all_images: $(patsubst %, docs/images/%.png, $(ENTRY_CLASSES))
 #docs/images/%.png:
@@ -46,7 +46,6 @@ docs/index.md: README.md
 docs/%/index.md: src/$(PACKAGE)/templates/%.yaml
 	$(RUN) gen-doc --include-top-level-diagram --diagram-type er_diagram $< -d docs/$*
 
-
 serve:
 	$(RUN) mkdocs serve
 

src/ontogpt/engines/enrichment.py

@@ -9,7 +9,7 @@
 
 from jinja2 import Template
 from oaklib import BasicOntologyInterface, get_adapter
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt import MODELS
 from ontogpt.engines.knowledge_engine import KnowledgeEngine

src/ontogpt/engines/generic_engine.py

@@ -6,7 +6,7 @@
 from typing import Iterator, List
 
 from jinja2 import Template
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt.engines.knowledge_engine import KnowledgeEngine
 from ontogpt.prompts.qa import GENERIC_QA_PROMPT

src/ontogpt/engines/gpt4all_engine.py

@@ -14,7 +14,7 @@
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
-import pydantic.v1
+import pydantic
 from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
 from oaklib import BasicOntologyInterface
 
@@ -202,7 +202,7 @@ def _remove_parenthetical_context(s: str):
                 f.write(dump_minimal_yaml(db))
 
     def generalize(
-        self, object: Union[pydantic.v1.BaseModel, dict], examples: List[EXAMPLE]
+        self, object: Union[pydantic.BaseModel, dict], examples: List[EXAMPLE]
     ) -> ExtractionResult:
         """
         Generalize the given examples.
@@ -217,7 +217,7 @@ def generalize(
         for example in examples:
             prompt += f"{self.serialize_object(example)}\n\n"
         prompt += "\n\n===\n\n"
-        if isinstance(object, pydantic.v1.BaseModel):
+        if isinstance(object, pydantic.BaseModel):
             object = object.dict()
         for k, v in object.items():
             if v:
@@ -309,7 +309,7 @@ def serialize_object(self, example: EXAMPLE, cls: ClassDefinition = None) -> str
             cls = self.template_class
         if isinstance(example, str):
             return example
-        if isinstance(example, pydantic.v1.BaseModel):
+        if isinstance(example, pydantic.BaseModel):
             example = example.dict()
         lines = []
         sv = self.schemaview
@@ -388,7 +388,7 @@ def get_completion_prompt(
         if object:
             if cls is None:
                 cls = self.template_class
-            if isinstance(object, pydantic.v1.BaseModel):
+            if isinstance(object, pydantic.BaseModel):
                 object = object.dict()
             for k, v in object.items():
                 if v:
@@ -509,7 +509,7 @@ def _parse_line_to_dict(
 
     def parse_completion_payload(
         self, results: str, cls: ClassDefinition = None, object: dict = None
-    ) -> pydantic.v1.BaseModel:
+    ) -> pydantic.BaseModel:
         """
         Parse the completion payload into a pydantic class.
 
@@ -526,7 +526,7 @@ def parse_completion_payload(
 
     def ground_annotation_object(
         self, ann: RESPONSE_DICT, cls: ClassDefinition = None
-    ) -> Optional[pydantic.v1.BaseModel]:
+    ) -> Optional[pydantic.BaseModel]:
         """Ground the direct parse of the OpenAI payload.
 
         The raw openAI payload is a YAML-like string, which is parsed to

src/ontogpt/engines/halo_engine.py

@@ -10,7 +10,7 @@
 from typing import Any, Dict, List, Optional, Set
 
 import openai
-import pydantic.v1
+import pydantic
 import tiktoken
 import yaml
 from linkml.utils.schema_fixer import uncamel
@@ -39,7 +39,7 @@
 """
 
 
-class StructuredPrompt(pydantic.v1.BaseModel):
+class StructuredPrompt(pydantic.BaseModel):
     header: str = None
     body: str = None
     main_prompt: str = None
@@ -317,7 +317,7 @@ def integrate_object(self, obj: Dict[str, Any], strict=True) -> Optional[Ontolog
         obj = self.repair_dict(obj)
         try:
             elt = OntologyElement(**obj)
-        except pydantic.v1.ValidationError as e:
+        except pydantic.ValidationError as e:
             logger.warning(f"## COULD NOT PARSE: {obj} /// {e}")
             if strict:
                 raise e
@@ -369,7 +369,7 @@ def old_integrate_payload(self, prompt: StructuredPrompt, payload: Dict[str, Any
             obj = {k: v for k, v in obj.items() if k in allowed_slots}
             try:
                 elt = OntologyElement(**obj)
-            except pydantic.v1.ValidationError as e:
+            except pydantic.ValidationError as e:
                 logger.info(f"## COULD NOT PARSE: {obj} /// {e}")
                 return added
             logger.info(f"Elt: {elt.name} // {slots_populated} // {obj}")

src/ontogpt/engines/hfhub_engine.py

@@ -16,7 +16,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-import pydantic.v1
+import pydantic
 from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
 
 from ontogpt.engines.knowledge_engine import (
@@ -116,7 +116,7 @@ def _extract_from_text_to_dict(self, text: str, cls: ClassDefinition = None) ->
         return self._parse_response_to_dict(raw_text, cls)
 
     def generalize(
-        self, object: Union[pydantic.v1.BaseModel, dict], examples: List[EXAMPLE]
+        self, object: Union[pydantic.BaseModel, dict], examples: List[EXAMPLE]
     ) -> ExtractionResult:
         """
         Generalize the given examples.
@@ -131,7 +131,7 @@ def generalize(
         for example in examples:
             prompt += f"{self.serialize_object(example)}\n\n"
         prompt += "\n\n===\n\n"
-        if isinstance(object, pydantic.v1.BaseModel):
+        if isinstance(object, pydantic.BaseModel):
             object = object.dict()
         for k, v in object.items():
             if v:
@@ -219,7 +219,7 @@ def serialize_object(self, example: EXAMPLE, cls: ClassDefinition = None) -> str
             cls = self.template_class
         if isinstance(example, str):
             return example
-        if isinstance(example, pydantic.v1.BaseModel):
+        if isinstance(example, pydantic.BaseModel):
             example = example.dict()
         lines = []
         sv = self.schemaview
@@ -298,7 +298,7 @@ def get_completion_prompt(
         if object:
             if cls is None:
                 cls = self.template_class
-            if isinstance(object, pydantic.v1.BaseModel):
+            if isinstance(object, pydantic.BaseModel):
                 object = object.dict()
             for k, v in object.items():
                 if v:
@@ -419,7 +419,7 @@ def _parse_line_to_dict(
 
     def parse_completion_payload(
         self, results: str, cls: ClassDefinition = None, object: dict = None
-    ) -> pydantic.v1.BaseModel:
+    ) -> pydantic.BaseModel:
         """
         Parse the completion payload into a pydantic class.
 
@@ -436,7 +436,7 @@ def parse_completion_payload(
 
     def ground_annotation_object(
         self, ann: RESPONSE_DICT, cls: ClassDefinition = None
-    ) -> Optional[pydantic.v1.BaseModel]:
+    ) -> Optional[pydantic.BaseModel]:
         """Ground the direct parse of the OpenAI payload.
 
         The raw openAI payload is a YAML-like string, which is parsed to

src/ontogpt/engines/knowledge_engine.py

@@ -11,7 +11,7 @@
 
 import inflection
 import openai
-import pydantic.v1
+import pydantic
 import tiktoken
 import yaml
 from linkml_runtime import SchemaView
@@ -31,7 +31,7 @@
 logger = logging.getLogger(__name__)
 
 
-OBJECT = Union[str, pydantic.v1.BaseModel, dict]
+OBJECT = Union[str, pydantic.BaseModel, dict]
 EXAMPLE = OBJECT
 FIELD = str
 TEMPLATE_NAME = str
@@ -177,7 +177,7 @@ def extract_from_text(
     ) -> ExtractionResult:
         raise NotImplementedError
 
-    def extract_from_file(self, file: Union[str, Path, TextIO]) -> pydantic.v1.BaseModel:
+    def extract_from_file(self, file: Union[str, Path, TextIO]) -> pydantic.BaseModel:
         """
         Extract annotations from the given text.
 
@@ -216,7 +216,7 @@ def synthesize(self, cls: ClassDefinition = None, object: OBJECT = None) -> Extr
         raise NotImplementedError
 
     def generalize(
-        self, object: Union[pydantic.v1.BaseModel, dict], examples: List[EXAMPLE]
+        self, object: Union[pydantic.BaseModel, dict], examples: List[EXAMPLE]
     ) -> ExtractionResult:
         raise NotImplementedError
 

src/ontogpt/engines/mapping_engine.py

@@ -11,7 +11,7 @@
 from oaklib.datamodels.vocabulary import IS_A, SKOS_RELATED_MATCH
 from oaklib.interfaces import MappingProviderInterface
 from oaklib.types import CURIE
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 from sssom.parsers import parse_sssom_table, to_mapping_set_document
 from sssom_schema import Mapping
 

src/ontogpt/engines/pheno_engine.py

@@ -9,7 +9,7 @@
 from oaklib import get_adapter
 from oaklib.datamodels.text_annotator import TextAnnotationConfiguration
 from oaklib.interfaces import MappingProviderInterface, TextAnnotatorInterface
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt.engines.knowledge_engine import KnowledgeEngine
 from ontogpt.io.yaml_wrapper import dump_minimal_yaml

src/ontogpt/engines/reasoner_engine.py

@@ -6,7 +6,7 @@
 from typing import List, Optional, Union
 
 from jinja2 import Template
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt import MODELS
 from ontogpt.engines.knowledge_engine import KnowledgeEngine

src/ontogpt/engines/spires_engine.py

@@ -15,7 +15,7 @@
 from pathlib import Path
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
-import pydantic.v1
+import pydantic
 import yaml
 from linkml_runtime.linkml_model import ClassDefinition, SlotDefinition
 from oaklib import BasicOntologyInterface
@@ -212,7 +212,7 @@ def _remove_parenthetical_context(s: str):
 
     def generalize(
         self,
-        object: Union[pydantic.v1.BaseModel, dict],
+        object: Union[pydantic.BaseModel, dict],
         examples: List[EXAMPLE],
         show_prompt: bool = False,
     ) -> ExtractionResult:
@@ -229,7 +229,7 @@ def generalize(
         for example in examples:
             prompt += f"{self.serialize_object(example)}\n\n"
         prompt += "\n\n===\n\n"
-        if isinstance(object, pydantic.v1.BaseModel):
+        if isinstance(object, pydantic.BaseModel):
             object = object.dict()
         for k, v in object.items():
             if v:
@@ -323,7 +323,7 @@ def serialize_object(self, example: EXAMPLE, cls: ClassDefinition = None) -> str
             cls = self.template_class
         if isinstance(example, str):
             return example
-        if isinstance(example, pydantic.v1.BaseModel):
+        if isinstance(example, pydantic.BaseModel):
             example = example.dict()
         lines = []
         sv = self.schemaview
@@ -404,7 +404,7 @@ def get_completion_prompt(
         if object:
             if cls is None:
                 cls = self.template_class
-            if isinstance(object, pydantic.v1.BaseModel):
+            if isinstance(object, pydantic.BaseModel):
                 object = object.dict()
             for k, v in object.items():
                 if v:
@@ -525,7 +525,7 @@ def _parse_line_to_dict(
 
     def parse_completion_payload(
         self, results: str, cls: ClassDefinition = None, object: dict = None
-    ) -> pydantic.v1.BaseModel:
+    ) -> pydantic.BaseModel:
         """
         Parse the completion payload into a pydantic class.
 
@@ -558,7 +558,7 @@ def _auto_add_ids(self, ann: RESPONSE_DICT, cls: ClassDefinition = None) -> None
 
     def ground_annotation_object(
         self, ann: RESPONSE_DICT, cls: ClassDefinition = None
-    ) -> Optional[pydantic.v1.BaseModel]:
+    ) -> Optional[pydantic.BaseModel]:
         """Ground the direct parse of the OpenAI payload.
 
         The raw openAI payload is a YAML-like string, which is parsed to

src/ontogpt/evaluation/ctd/eval_ctd.py

@@ -27,7 +27,7 @@
 import yaml
 from bioc import biocxml
 from oaklib import BasicOntologyInterface, get_implementation_from_shorthand
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt.engines.knowledge_engine import chunk_text
 from ontogpt.engines.spires_engine import SPIRESEngine

src/ontogpt/evaluation/drugmechdb/datamodel/drugmechdb.py

@@ -3,8 +3,8 @@
 
 from typing import List, Optional
 
-from pydantic.v1 import BaseModel as BaseModel
-from pydantic.v1 import Field
+from pydantic import BaseModel as BaseModel
+from pydantic import Field
 
 metamodel_version = "None"
 version = "None"

src/ontogpt/evaluation/drugmechdb/eval_drugmechdb.py

@@ -13,7 +13,7 @@
 
 import yaml
 from oaklib import get_implementation_from_shorthand
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 import ontogpt.evaluation.drugmechdb.datamodel.drugmechdb as source_datamodel
 import ontogpt.templates.drug as target_datamodel

src/ontogpt/evaluation/enrichment/eval_enrichment.py

@@ -20,7 +20,7 @@
 )
 from oaklib.interfaces.obograph_interface import OboGraphInterface
 from oaklib.parsers.association_parser_factory import get_association_parser
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 from tiktoken import Encoding
 
 from ontogpt.engines import create_engine

src/ontogpt/evaluation/evaluation_engine.py

@@ -14,7 +14,7 @@
 from typing import List, Optional, Set
 
 from oaklib import BasicOntologyInterface
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt.engines.spires_engine import SPIRESEngine
 

src/ontogpt/evaluation/go/eval_go.py

@@ -9,7 +9,7 @@
 from oaklib.datamodels.obograph import LogicalDefinitionAxiom
 from oaklib.datamodels.vocabulary import IS_A
 from oaklib.interfaces.obograph_interface import OboGraphInterface
-from pydantic.v1 import BaseModel
+from pydantic import BaseModel
 
 from ontogpt.engines.spires_engine import SPIRESEngine
 from ontogpt.evaluation.evaluation_engine import SimilarityScore, SPIRESEvaluationEngine