You need to pass in a pydantic class that represents a graph, but it has some pretty specific requirements.
Under the hood we are creating a model on the fly based on the options you pass in
If you were to do it manually, it might look like
class Entity(BaseModel):
"""A representation of a entity in text."""
type: Literal["PERSON", "PLACE", "THING"]
name: str
class Relation(BaseModel):
"""A type of relationship between two entities."""
type: Literal["BORN_IN", "RELATED_TO"]
class Triplet(BaseModel):
subject: Entity
relation: Relation
object: Entity
class KGSchema(BaseModel):
"""A graph represented by a list of triplets."""
triplets: List[Triplet]
kg_schema_cls = KGSchema
