Oh, definitely. I recommend you go for contracts. I've used something similar for a contract that iteratively "stitched together" a broken ontology graph. Here's some of the data models for inspiration -- you could have something similar for your ops, and write the contract to solve for one op, then apply the op, etc.
---
class Merge(LLMDataModel):
indexes: list[int] = Field(description="The indices of the clusters that are being merged.")
relations: list[SubClassRelation] = Field(
description="A list of superclass-subclass relations chosen from the existing two clusters in such a way that they merge."
)
@field_validator("indexes")
@classmethod
def is_binary(cls, v):
if len(v) != 2:
raise ValueError(
f"Binary op error: Invalid number of clusters: {len(v)}. The merge operation requires exactly two clusters."
)
return v
class Bridge(LLMDataModel):
indexes: list[int] = Field(description="The indices of the clusters that are being bridged.")
relations: list[SubClassRelation] = Field(
description="A list of new superclass-subclass relations used to bridge the two clusters from the ontology."
)
@field_validator("indexes")
@classmethod
def is_binary(cls, v):
if len(v) != 2:
raise ValueError(
f"Binary op error: Invalid number of clusters: {len(v)}. The merge operation requires exactly two clusters."
)
return v
class Prune(LLMDataModel):
indexes: list[int] = Field(description="The indices of the clusters that are being pruned.")
classes: list[str] = Field(description="A list of classes that are being pruned from the ontology.")
@field_validator("indexes")
@classmethod
def is_unary(cls, v):
if len(v) > 1:
raise ValueError(
f"Unary op error: Invalid number of clusters: {len(v)}. The prune operation requires exactly one cluster."
)
return v
class Operation(LLMDataModel):
type: Merge | Bridge | Prune = Field(description="The type of operation to perform.")
---