diff --git a/anansyTypes/structures/embedding.py b/anansyTypes/structures/embedding.py index 9553cdb..ef50d69 100644 --- a/anansyTypes/structures/embedding.py +++ b/anansyTypes/structures/embedding.py @@ -30,18 +30,26 @@ class Embedding: @classmethod def __get_pydantic_core_schema__(cls, source_type: Any, handler): from pydantic_core import core_schema - + def validate_embedding(value): if isinstance(value, cls): return value if isinstance(value, dict): return cls.from_dict(value) raise ValueError(f"Cannot convert {type(value)} to Embedding") - - return core_schema.with_info_plain_validator_function( + + python_schema = core_schema.with_info_plain_validator_function( lambda value, _: validate_embedding(value), serialization=core_schema.plain_serializer_function_ser_schema( lambda x: {"data": x.data.tolist()}, when_used="json", ), ) + + return core_schema.json_or_python_schema( + json_schema=core_schema.typed_dict_schema({ + 'data': core_schema.typed_dict_field(core_schema.list_schema()) + }), + python_schema=python_schema, + ) + diff --git a/anansyTypes/structures/lexeme.py b/anansyTypes/structures/lexeme.py index 2d27e8d..b26edf6 100644 --- a/anansyTypes/structures/lexeme.py +++ b/anansyTypes/structures/lexeme.py @@ -19,6 +19,45 @@ class Lexeme(ABC): def embedding(self) -> Embedding: """Embedding(s) associated with the lexeme.""" pass - - - + + @classmethod + @abstractmethod + def from_dict(cls, data: dict) -> "Lexeme": + """Create instance from dictionary.""" + pass + + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler): + from pydantic_core import core_schema + + def validate_lexeme(value): + if isinstance(value, cls): + return value + if isinstance(value, dict): + return cls.from_dict(value) + raise ValueError(f"Cannot convert {type(value)} to {cls.__name__}") + + python_schema = core_schema.with_info_plain_validator_function( + lambda value, _: validate_lexeme(value), + serialization=core_schema.plain_serializer_function_ser_schema( + lambda x: { + "id": x.id, + "content": x.content, + "embedding": {"data": x.embedding.data.tolist()}, + }, + when_used="json", + ), + ) + + return core_schema.json_or_python_schema( + json_schema=core_schema.typed_dict_schema( + { + "id": core_schema.typed_dict_field(core_schema.int_schema()), + "content": core_schema.typed_dict_field(core_schema.str_schema()), + "embedding": core_schema.typed_dict_field( + handler.generate_schema(Embedding) + ), + } + ), + python_schema=python_schema, + ) diff --git a/anansyTypes/structures/phrase.py b/anansyTypes/structures/phrase.py index a5dba06..92d56ed 100644 --- a/anansyTypes/structures/phrase.py +++ b/anansyTypes/structures/phrase.py @@ -25,27 +25,4 @@ class Phrase(Lexeme): id=data["id"], content=data["content"], embedding=Embedding.from_dict(data["embedding"]), - ) - - @classmethod - def __get_pydantic_core_schema__(cls, source_type: Any, handler): - from pydantic_core import core_schema - - def validate_phrase(value): - if isinstance(value, cls): - return value - if isinstance(value, dict): - return cls.from_dict(value) - raise ValueError(f"Cannot convert {type(value)} to Phrase") - - return core_schema.with_info_plain_validator_function( - lambda value, _: validate_phrase(value), - serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: { - "id": x.id, - "content": x.content, - "embedding": {"data": x.embedding.data.tolist()}, - }, - when_used="json", - ), - ) + ) \ No newline at end of file diff --git a/anansyTypes/structures/word.py b/anansyTypes/structures/word.py index 9f1128c..f95b2df 100644 --- a/anansyTypes/structures/word.py +++ b/anansyTypes/structures/word.py @@ -31,26 +31,3 @@ class Word(Lexeme): content=data["content"], embedding=Embedding.from_dict(data["embedding"]), ) - - @classmethod - def __get_pydantic_core_schema__(cls, source_type: Any, handler): - from pydantic_core import core_schema - - def validate_word(value): - if isinstance(value, cls): - return value - if isinstance(value, dict): - return cls.from_dict(value) - raise ValueError(f"Cannot convert {type(value)} to Word") - - return core_schema.with_info_plain_validator_function( - lambda value, _: validate_word(value), - serialization=core_schema.plain_serializer_function_ser_schema( - lambda x: { - "id": x.id, - "content": x.content, - "embedding": {"data": x.embedding.data.tolist()}, - }, - when_used="json", - ), - )