Merge branch 'main' of https://github.com/stuartgo/anansyTypes

anansyTypes/enums.py

@@ -1,21 +1,33 @@
 from enum import Enum
+from typing import Union
+
+
 class ModelType(Enum):
     EMBEDDING = "embedding"
     TRANSCRIPTION = "transcription"
     RESPONSE = "response"
-    
+
+
 class EmbeddingModelName(Enum):
     NB_BERT = "nb_bert"
 
+
 class TranscriptionModelName(Enum):
     NB_WHISPER_SMALL = "nb_whisper_small"
 
+
 class ResponseModelName(Enum):
-    GENERIC = "generic"  
-    
+    GENERIC = "generic"
+
+
+ModelName= Union[EmbeddingModelName, TranscriptionModelName, ResponseModelName]
+
+
 class PreprocessorType(Enum):
     NORWEGIAN = "norwegian"
     DEFAULT = "norwegian"
+
+
 class SymbolProvider(Enum):
     TOBII_DYNAVOX = "tobii_dynavox"
 
@@ -24,4 +36,3 @@ class Language(Enum):
     NORWEGIAN = "norwegian"
     ENGLISH = "english"
     DEFAULT = "norwegian"
-

anansyTypes/schemas.py

@@ -21,12 +21,6 @@ class ModifiedBaseModel(PydanticBaseModel):
                 'content': v.content,
                 'embedding': {'data': v.embedding.data.tolist()}
             }
-        if isinstance(v, Symbol):
-            return {
-                'id': v.id,
-                'label': v.label,
-                'image': v.image.tolist()
-            }
         if isinstance(v, Embedding):
             return {'data': v.data.tolist()}
         if isinstance(v, list):
@@ -56,7 +50,7 @@ class GetSymbolRequest(ModifiedBaseModel):
     lexeme_id: int
 
 class GetSymbolResponse(ModifiedBaseModel):
-    symbol: dict
+    symbol: Union[Symbol, None]
 
 class GetWordRequest(ModifiedBaseModel):
     word_text: str

anansyTypes/structures/symbol.py

@@ -4,6 +4,7 @@ import numpy as np
 from pydantic_tensor import Tensor
 from typing import Any, Literal
 from pydantic_core import core_schema
+import base64
 
 
 @dataclass
@@ -13,40 +14,42 @@ class Symbol:
     image: Tensor[np.ndarray[Any, Any], tuple, Literal["int32", "int64"]]
 
     @classmethod
-    def from_dict(cls, data: dict) -> 'Symbol':
+    def from_dict(cls, data: dict) -> "Symbol":
+        if isinstance(data["image"], str):
+            # Decode from base64
+            image_bytes = base64.b64decode(data["image"])
+            image_array = np.frombuffer(image_bytes, dtype=np.int32)
+        else:
+            # From list/array
+            image_array = np.array(data["image"], dtype=np.int32)
+
         return cls(
-            id=data['id'],
-            label=data['label'],
-            image=np.array(data['image'], dtype=np.int32)
+            id=data["id"],
+            label=data["label"],
+            image=image_array,
         )
 
     @classmethod
     def __get_pydantic_core_schema__(cls, source_type: Any, handler):
         from pydantic_core import core_schema
-        
-        return core_schema.no_info_after_validator_function(
-            cls,
-            core_schema.dataclass_schema(
-                cls=cls,
-                schema=core_schema.dataclass_args_schema(
-                    cls.__name__,
-                    [
-                        core_schema.dataclass_field(name='id', schema=handler.generate_schema(int)),
-                        core_schema.dataclass_field(name='label', schema=handler.generate_schema(str)),
-                        core_schema.dataclass_field(
-                            name='image',
-                            schema=handler.generate_schema(Tensor[np.ndarray[Any, Any], tuple, Literal["int32", "int64"]])
-                        ),
-                    ],
-                ),
-                fields=['id', 'label', 'image'],
-            ),
+
+        def validate_symbol(value):
+            if isinstance(value, cls):
+                return value
+            if isinstance(value, dict):
+                return cls.from_dict(value)
+            raise ValueError(f"Cannot convert {type(value)} to Symbol")
+
+        return core_schema.with_info_plain_validator_function(
+            lambda value, _: validate_symbol(value),
             serialization=core_schema.plain_serializer_function_ser_schema(
                 lambda x: {
-                    'id': x.id,
-                    'label': x.label,
-                    'image': x.image
+                    "id": x.id,
+                    "label": x.label,
+                    "image": base64.b64encode(np.array(x.image).tobytes()).decode(
+                        "utf-8"
+                    ),
                 },
-                when_used='json',
+                when_used="json",
             ),
         )