Changes in schemas.py

Trying to use proper types now
Changes in embedding.py
    Implemented serialisation

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

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",
+            ),
+        )
+