Has anyone tried defining a custom `kg_

from pydantic.v1 import BaseModel, Field, validator, root_validator
from typing import Literal

class Entity(BaseModel):
  """An entity in a graph."""
  type: Literal["PERSON", "PLACE", "THING"] = Field(description="Entity in a knowledge graph. Only extract entities with types that are listed as valid: PERSON, PLACE, OR THING.")

class Relation(BaseModel):
  """A relation connecting to entities in a graph."""
  type: Literal["HAS", "PART_OF"] = Field(description="Relation in a knowledge graph. Only extract relations with types that are listed as valid: HAS, PART_OF.")

class Triplet(BaseModel):
  """A triplet of two entities and a relation."""
  subject: Entity
  relation: Relation
  object: Entity

class Triplets(BaseModel):
  """Knowledge Graph Schema."""
  triplets: list[Triplet]

  @validator("triplets", pre=True)
  def validate_triplets(v, values):
    passing_triplets = []
    for i, triplet in enumerate(v):
      # cleanup
      try:
          for key in triplet:
              triplet[key]["type"] = triplet[key]["type"].replace(" ", "_")
              triplet[key]["type"] = triplet[key]["type"].upper()

              # validate, skip if invalid
              _ = triplet_cls(**triplet)
              passing_triplets.append(v[i])
      except (KeyError, ValueError):
          continue

      return passing_triplets


kg_schema_cls = Triplets