Implementation and tests for DocumentChunk

Author: dhirupandey

src/coherence/ai.py

@@ -6,9 +6,12 @@
 
 import base64
 from abc import ABC
-from typing import Optional
+from collections import OrderedDict
+from typing import Any, Dict, Optional, cast
 
-from coherence.serialization import proxy
+import jsonpickle
+
+from coherence.serialization import JavaProxyUnpickler, proxy
 
 
 class Vector(ABC):
@@ -54,3 +57,73 @@ class FloatVector(Vector):
     def __init__(self, float_array: list[float]):
         super().__init__()
         self.array = float_array
+
+
+class AbstractEvolvable(ABC):
+    def __init__(self, data_version: int = 0, bin_future: Optional[Any] = None):
+        self.dataVersion = data_version
+        self.binFuture = bin_future
+
+
+@proxy("ai.DocumentChunk")
+class DocumentChunk(AbstractEvolvable):
+    def __init__(
+        self,
+        text: str,
+        metadata: Optional[dict[str, Any] | OrderedDict[str, Any]] = None,
+        vector: Optional[Vector] = None,
+    ):
+        super().__init__()
+        self.text = text
+        if metadata is None:
+            self.metadata: Dict[str, Any] = OrderedDict()
+        else:
+            self.metadata = metadata
+        self.vector = vector
+
+
+@jsonpickle.handlers.register(DocumentChunk)
+class DocumentChunkHandler(jsonpickle.handlers.BaseHandler):
+    def flatten(self, obj: object, data: dict[str, Any]) -> dict[str, Any]:
+        dc: DocumentChunk = cast(DocumentChunk, obj)
+        result_dict: Dict[Any, Any] = dict()
+        result_dict["@class"] = "ai.DocumentChunk"
+        result_dict["dataVersion"] = dc.dataVersion
+        if hasattr(dc, "binFuture"):
+            if dc.binFuture is not None:
+                result_dict["binFuture"] = dc.binFuture
+        if hasattr(dc, "metadata"):
+            if dc.metadata is not None:
+                result_dict["metadata"] = dict()
+                if isinstance(dc.metadata, OrderedDict):
+                    result_dict["metadata"]["@ordered"] = True
+                entries = list()
+                for k, v in dc.metadata.items():
+                    entries.append({"key": k, "value": v})
+                result_dict["metadata"]["entries"] = entries
+        if hasattr(dc, "vector"):
+            v = dc.vector
+            if v is not None:
+                if isinstance(v, BitVector):
+                    result_dict["vector"] = dict()
+                    result_dict["vector"]["@class"] = "ai.BitVector"
+                    # noinspection PyUnresolvedReferences
+                    result_dict["vector"]["bits"] = v.bits
+                elif isinstance(v, ByteVector):
+                    result_dict["vector"] = dict()
+                    result_dict["vector"]["@class"] = "ai.Int8Vector"
+                    # noinspection PyUnresolvedReferences
+                    result_dict["vector"]["array"] = v.array
+                elif isinstance(v, FloatVector):
+                    result_dict["vector"] = dict()
+                    result_dict["vector"]["@class"] = "ai.Float32Vector"
+                    # noinspection PyUnresolvedReferences
+                    result_dict["vector"]["array"] = v.array
+        result_dict["text"] = dc.text
+        return result_dict
+
+    def restore(self, obj: dict[str, Any]) -> DocumentChunk:
+        jpu = JavaProxyUnpickler()
+        d = DocumentChunk("")
+        o = jpu._restore_from_dict(obj, d)
+        return o

src/coherence/serialization.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import collections
 from abc import ABC, abstractmethod
 from decimal import Decimal
 from typing import Any, Callable, Dict, Final, Optional, Type, TypeVar, cast
@@ -18,6 +19,7 @@
 _META_CLASS: Final[str] = "@class"
 _META_VERSION: Final[str] = "@version"
 _META_ENUM: Final[str] = "enum"
+_META_ORDERED: Final[str] = "@ordered"
 
 _JSON_KEY = "key"
 _JSON_VALUE = "value"
@@ -205,10 +207,10 @@ class JavaProxyUnpickler(jsonpickle.Unpickler):
     # noinspection PyUnresolvedReferences
     def _restore(self, obj: Any) -> Any:
         if isinstance(obj, dict):
-            metadata: str = obj.get(_META_CLASS, None)
+            metadata: Any = obj.get(_META_CLASS, None)
             if metadata is not None:
                 type_: Optional[Type[Any]] = _type_for(metadata)
-                actual: dict[str, Any] = dict()
+                actual: dict[Any, Any] = dict()
                 if type_ is None:
                     if "map" in metadata.lower():
                         for entry in obj[_JSON_ENTRIES]:
@@ -228,6 +230,25 @@ def _restore(self, obj: Any) -> Any:
 
                 return super().restore(actual, reset=False)
 
+            # When "@Ordered" set to true which converts to OrderedDict()
+            metadata = obj.get(_META_ORDERED, False)
+            if metadata is True:
+                o = collections.OrderedDict()
+                entries = obj.get(_JSON_ENTRIES, None)
+                if entries is not None:
+                    for entry in obj[_JSON_ENTRIES]:
+                        o[entry[_JSON_KEY]] = entry[_JSON_VALUE]
+                return o
+
+            #  When there is no "@Ordered" set. Only "entries" list exists
+            if len(obj) == 1:
+                entries = obj.get(_JSON_ENTRIES, None)
+                if entries is not None:
+                    actual = dict()
+                    for entry in obj[_JSON_ENTRIES]:
+                        actual[entry[_JSON_KEY]] = entry[_JSON_VALUE]
+                    return super().restore(actual, reset=False)
+
         return super()._restore(obj)
 
 

tests/test_ai.py

@@ -2,7 +2,7 @@
 # Licensed under the Universal Permissive License v 1.0 as shown at
 # https://oss.oracle.com/licenses/upl.
 
-from coherence.ai import BitVector, ByteVector, FloatVector
+from coherence.ai import BitVector, ByteVector, FloatVector, DocumentChunk
 from coherence.serialization import JSONSerializer, SerializerRegistry
 
 s = SerializerRegistry.serializer(JSONSerializer.SER_FORMAT)
@@ -12,27 +12,63 @@ def test_BitVector_serialization() -> None:
     coh_bv = BitVector(hex_string="AABBCC")
     ser = s.serialize(coh_bv)
     assert ser == b'\x15{"@class": "ai.BitVector", "bits": "0xAABBCC"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, BitVector)
 
     coh_bv = BitVector(hex_string="0xAABBCC")
     ser = s.serialize(coh_bv)
     assert ser == b'\x15{"@class": "ai.BitVector", "bits": "0xAABBCC"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, BitVector)
 
     coh_bv = BitVector(hex_string=None, byte_array=bytes([1, 2, 10]))
     ser = s.serialize(coh_bv)
     assert ser == b'\x15{"@class": "ai.BitVector", "bits": "0x01020a"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, BitVector)
 
     coh_bv = BitVector(hex_string=None, int_array=[1234, 1235])
     ser = s.serialize(coh_bv)
     assert ser == b'\x15{"@class": "ai.BitVector", "bits": "0x4d24d3"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, BitVector)
 
 
 def test_ByteVector_serialization() -> None:
     coh_int8v = ByteVector(bytes([1, 2, 3, 4]))
     ser = s.serialize(coh_int8v)
     assert ser == b'\x15{"@class": "ai.Int8Vector", "array": "AQIDBA=="}'
+    o = s.deserialize(ser)
+    assert isinstance(o, ByteVector)
 
 
 def test_FloatVector_serialization() -> None:
     coh_fv = FloatVector([1.0, 2.0, 3.0])
     ser = s.serialize(coh_fv)
     assert ser == b'\x15{"@class": "ai.Float32Vector", "array": [1.0, 2.0, 3.0]}'
+    o = s.deserialize(ser)
+    assert isinstance(o, FloatVector)
+
+
+def test_DocumentChunk_serialization() -> None:
+    dc = DocumentChunk("test")
+    ser = s.serialize(dc)
+    assert ser == b'\x15{"@class": "ai.DocumentChunk", "dataVersion": 0, "metadata": {"@ordered": true, "entries": []}, "text": "test"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, DocumentChunk)
+
+    d = {"one":"one-value", "two": "two-value"}
+    dc = DocumentChunk("test",d)
+    ser = s.serialize(dc)
+    assert ser == b'\x15{"@class": "ai.DocumentChunk", "dataVersion": 0, "metadata": {"entries": [{"key": "one", "value": "one-value"}, {"key": "two", "value": "two-value"}]}, "text": "test"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, DocumentChunk)
+
+    coh_fv = FloatVector([1.0, 2.0, 3.0])
+    d = {"one":"one-value", "two": "two-value"}
+    dc = DocumentChunk("test",d, coh_fv)
+    ser = s.serialize(dc)
+    assert ser == b'\x15{"@class": "ai.DocumentChunk", "dataVersion": 0, "metadata": {"entries": [{"key": "one", "value": "one-value"}, {"key": "two", "value": "two-value"}]}, "vector": {"@class": "ai.Float32Vector", "array": [1.0, 2.0, 3.0]}, "text": "test"}'
+    o = s.deserialize(ser)
+    assert isinstance(o, DocumentChunk)
+