gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add versioning to the data point model

Browse source
This commit is contained in:
vasilije 2024-12-17 20:06:29 +01:00
parent 6fb3b4abec
commit 87bc5d8266
1 changed files with 10 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
cognee/infrastructure/engine/models/DataPoint.py
View file

@ -24,22 +24,20 @@ class DataPoint(BaseModel):
type: Optional[str] = "text" # "text", "file", "image", "video" type: Optional[str] = "text" # "text", "file", "image", "video"
topological_rank: Optional[int] = 0 topological_rank: Optional[int] = 0
extra: Optional[Dict[str, Any]] = None # For additional properties extra: Optional[Dict[str, Any]] = None # For additional properties
_metadata: Optional[MetaData] = Field( _metadata: Optional[MetaData] = {
default={"index_fields": [], "type": "DataPoint"} "index_fields": [],
) "type": "DataPoint"
}
# Override the Pydantic configuration # Override the Pydantic configuration
class Config: class Config:
underscore_attrs_are_private = True underscore_attrs_are_private = True
@classmethod @classmethod
def get_embeddable_data(cls, data_point): @classmethod
"""Retrieve embeddable data based on metadata's index_fields.""" def get_embeddable_data(self, data_point):
if ( if data_point._metadata and len(data_point._metadata["index_fields"]) > 0 \
data_point._metadata and hasattr(data_point, data_point._metadata["index_fields"][0]):
and len(data_point._metadata["index_fields"]) > 0
and hasattr(data_point, data_point._metadata["index_fields"][0])
):
attribute = getattr(data_point, data_point._metadata["index_fields"][0]) attribute = getattr(data_point, data_point._metadata["index_fields"][0])
if isinstance(attribute, str): if isinstance(attribute, str):
@ -47,14 +45,14 @@ class DataPoint(BaseModel):
return attribute return attribute
@classmethod @classmethod
def get_embeddable_properties(cls, data_point): def get_embeddable_properties(self, data_point):
"""Retrieve all embeddable properties.""" """Retrieve all embeddable properties."""
if data_point._metadata and len(data_point._metadata["index_fields"]) > 0: if data_point._metadata and len(data_point._metadata["index_fields"]) > 0:
return [getattr(data_point, field, None) for field in data_point._metadata["index_fields"]] return [getattr(data_point, field, None) for field in data_point._metadata["index_fields"]]
return [] return []
@classmethod @classmethod
def get_embeddable_property_names(cls, data_point): def get_embeddable_property_names(self, data_point):
"""Retrieve names of embeddable properties.""" """Retrieve names of embeddable properties."""
return data_point._metadata["index_fields"] or [] return data_point._metadata["index_fields"] or []