topoteretes--cognee/cognee/modules/storage/utils/__init__.py

import json
import copy
from uuid import UUID
from decimal import Decimal
from datetime import datetime
from pydantic_core import PydanticUndefined
from pydantic import create_model, ConfigDict, BaseModel, Field

from cognee.infrastructure.engine import DataPoint


class JSONEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime):
            return obj.isoformat()  # Convert datetime to ISO 8601 string
        elif isinstance(obj, UUID):
            # if the obj is uuid, we simply return the value of uuid
            return str(obj)
        elif isinstance(obj, Decimal):
            return float(obj)
        return json.JSONEncoder.default(self, obj)


def copy_model(
    model: DataPoint, include_fields: dict | None = None, exclude_fields: list | None = None
):
    if include_fields is None:
        include_fields = {}
    if exclude_fields is None:
        exclude_fields = []
    fields = {}
    for name, field in model.model_fields.items():
        if name in exclude_fields:
            continue

        if hasattr(field, "is_required") and field.is_required():
            fields[name] = (field.annotation, PydanticUndefined)
        elif field.default_factory is not None:
            fields[name] = (field.annotation, Field(default_factory=field.default_factory))
        else:
            # Preserve explicit None defaults (and other scalar/container defaults)
            # so optional fields remain optional in copied models.
            fields[name] = (field.annotation, copy.deepcopy(field.default))

    final_fields = {**fields, **include_fields}

    # Create a base class with the same configuration as DataPoint
    class ConfiguredBase(BaseModel):
        model_config = ConfigDict(arbitrary_types_allowed=True)

    # Create the model inheriting from the configured base
    new_model = create_model(model.__name__, __base__=ConfiguredBase, **final_fields)

    new_model.model_rebuild()
    return new_model


def get_own_properties(data_point: DataPoint):
    properties = {}

    for field_name, field_value in data_point:
        if (
            field_name == "metadata"
            or isinstance(field_value, dict)
            or isinstance(field_value, DataPoint)
            or (
                isinstance(field_value, list)
                and len(field_value) > 0
                and isinstance(field_value[0], DataPoint)
            )
        ):
            continue

        properties[field_name] = field_value

    return properties