diff --git a/anansyTypes/schemas.py b/anansyTypes/schemas.py index 9140fe2..0b10c9b 100644 --- a/anansyTypes/schemas.py +++ b/anansyTypes/schemas.py @@ -1,33 +1,13 @@ from pydantic import BaseModel as PydanticBaseModel -from typing import List, Any, Dict,Union +from typing import List, Any, Dict, Optional,Union from .structures import Embedding,Word, Symbol,Phrase from pydantic import field_validator, ConfigDict from pydantic_tensor import Tensor import torch from typing import Any, Literal class ModifiedBaseModel(PydanticBaseModel): - model_config = ConfigDict(arbitrary_types_allowed=True) - - @field_validator('*', mode='before') - @classmethod - def convert_custom_types(cls, v): - if isinstance(v, Word): - return { - 'id': v.id, - 'content': v.content, - 'embedding': {'data': v.embedding.data.tolist()} - } - if isinstance(v, Phrase): - return { - 'id': v.id, - 'content': v.content, - 'embedding': {'data': v.embedding.data.tolist()} - } - if isinstance(v, Embedding): - return {'data': v.data.tolist()} - if isinstance(v, list): - return [cls.convert_custom_types(item) for item in v] - return v + pass + #model_config = ConfigDict(arbitrary_types_allowed=True) class EmbeddingRequest(ModifiedBaseModel): @@ -36,7 +16,7 @@ class EmbeddingRequest(ModifiedBaseModel): word_ids: List[int] class EmbeddingResponse(ModifiedBaseModel): - embeddings: List[dict] + embeddings: List[Embedding] class ResponseRequest(ModifiedBaseModel): text: str @@ -54,20 +34,20 @@ class GetSymbolRequest(ModifiedBaseModel): lexeme_id: int class GetSymbolResponse(ModifiedBaseModel): - symbol: Union[Symbol, None] + symbol: Optional[Symbol] class GetWordRequest(ModifiedBaseModel): word_text: str stemmed: bool = False class GetWordResponse(ModifiedBaseModel): - word: Union[dict, None] + word: Optional[Word] class GetSimilarWordRequest(ModifiedBaseModel): - embedding: dict + embedding: Embedding class GetAllPhrasesResponse(ModifiedBaseModel): - phrases: List[dict] + phrases: List[Phrase] class GetTranscriptionRequest(ModifiedBaseModel): audio_data: bytes diff --git a/anansyTypes/structures/embedding.py b/anansyTypes/structures/embedding.py index 5ea6050..3cbffd0 100644 --- a/anansyTypes/structures/embedding.py +++ b/anansyTypes/structures/embedding.py @@ -26,3 +26,16 @@ class Embedding: @classmethod def from_dict(cls, data: dict) -> "Embedding": return cls(data=np.array(data["data"], dtype=np.float32)) + + @classmethod + def __get_pydantic_core_schema__(cls, source_type: Any, handler): + from pydantic_core import core_schema + + return core_schema.with_info_plain_validator_function( + lambda value, _: cls.from_dict(value) if isinstance(value, dict) else value, + serialization=core_schema.plain_serializer_function_ser_schema( + lambda x: {"data": x.data.tolist()}, + when_used="json", + ), + ) +