Fix object encode

src/knowledge/data/chunk.py

@@ -1,9 +1,9 @@
-from enum import Enum
+from enum import IntEnum
 from ..object import ObjectID
 
 ChunkID = ObjectID
 
-class PositionType(Enum):
+class PositionType(IntEnum):
     Unknown = 1
     Device = 2
     File = 3

src/knowledge/object/object.py

@@ -1,4 +1,3 @@
-
 # define a object type enum
 from abc import ABC, abstractmethod
 from enum import Enum
@@ -6,6 +5,16 @@ from .object_id import ObjectID, ObjectType
 import hashlib
 import json
 import pickle
+from typing import Any
+
+
+class ObjectEnhancedJSONEncoder(json.JSONEncoder):
+    def default(self, o: Any) -> Any:
+        if isinstance(o, ObjectID):
+            return o.to_base58()
+
+        return super().default(o)
+
 
 class KnowledgeObject(ABC):
     def __init__(self, object_type: ObjectType, desc: dict = {}, body: dict = {}):
@@ -18,7 +27,7 @@ class KnowledgeObject(ABC):
 
     def object_id(self) -> ObjectID:
         return self.calculate_id()
-    
+
     def set_desc_with_key_value(self, key, value):
         self.desc[key] = value
 
@@ -27,7 +36,7 @@ class KnowledgeObject(ABC):
 
     def get_desc(self) -> dict:
         return self.desc
-    
+
     def set_body_with_key_value(self, key, value):
         self.body[key] = value
 
@@ -36,10 +45,13 @@ class KnowledgeObject(ABC):
 
     def get_body(self) -> dict:
         return self.body
-    
+
     def calculate_id(self):
         # Convert the object_type and desc to string and compute the SHA256 hash
-        data = json.dumps({"object_type": self.object_type, "desc": self.desc})
+        data = json.dumps(
+            {"object_type": self.object_type, "desc": self.desc},
+            cls=ObjectEnhancedJSONEncoder,
+        )
         sha256 = hashlib.sha256()
         sha256.update(data.encode())
         return ObjectID(sha256.digest())
@@ -50,5 +62,3 @@ class KnowledgeObject(ABC):
     @staticmethod
     def decode(data: bytes):
         return pickle.loads(data)
-
-    

src/knowledge/object/object_id.py

@@ -1,12 +1,12 @@
 # define a object type enum
 from abc import ABC, abstractmethod
-from enum import Enum
+from enum import IntEnum
 from .hash import HashValue
 import base58
 import base36
 
 
-class ObjectType(Enum):
+class ObjectType(IntEnum):
     Chunk = 7
     Image = 101
     Video = 102

src/knowledge/object/relation.py

@@ -4,10 +4,10 @@ import sqlite3
 from typing import List, Tuple, Optional
 import logging
 import os
-from enum import Enum
+from enum import IntEnum
 
 
-class ObjectRelationType(Enum):
+class ObjectRelationType(IntEnum):
     Parent = 1
 
 

test/test_object.py

@@ -23,6 +23,7 @@ from knowledge import (
     EmailObjectBuilder,
     ObjectRelationStore,
     KnowledgeStore,
+    EmailObject,
 )
 import asyncio
 import unittest
@@ -42,6 +43,20 @@ class TestVectorSTorage(unittest.TestCase):
 
         email_folder = "F:\\system\\Downloads\\8081ffdb80925f5bff9c6ab9c4756c7d"
         email_object = EmailObjectBuilder({}, email_folder).build()
+        
+        id = email_object.calculate_id()
+        print(f"got email object: {id.to_base58()}")
+        
+        # test encode & decode
+        ret = email_object.encode()
+        obj = EmailObject.decode(ret)
+        id2 = obj.calculate_id()
+        print(f"got email object: {id2.to_base58()}")
+        self.assertEqual(id.to_base58(), id2.to_base58())
+        
+        ret2 = obj.encode()
+        self.assertEqual(ret, ret2)
+        
 
     def test_relation(self):
         obj1 = ObjectID.hash_data("12345".encode("utf-8"))