diff --git a/src/aios_kernel/knowledge_base.py b/src/aios_kernel/knowledge_base.py index ebbfca8..18d78a3 100644 --- a/src/aios_kernel/knowledge_base.py +++ b/src/aios_kernel/knowledge_base.py @@ -1,6 +1,5 @@ # define a knowledge base class import json -import pickle from . import AgentPrompt, ComputeKernel from ..knowledge import * diff --git a/src/knowledge/object/relation.py b/src/knowledge/object/relation.py index 1b56e48..f6eca2e 100644 --- a/src/knowledge/object/relation.py +++ b/src/knowledge/object/relation.py @@ -1,16 +1,64 @@ # define a relation store class +import sqlite3 +import os +import logging from .object_id import ObjectID class ObjectRelationStore: def __init__(self, root_dir: str): - pass + if not os.path.exists(root_dir): + os.makedirs(root_dir) + file = os.path.join(root_dir, "relationships.db") + logging.info(f"will init chunk tracker, db={file}") + + self.conn = sqlite3.connect(file) + self.cursor = self.conn.cursor() + self.cursor.execute( + """ + CREATE TABLE IF NOT EXISTS relationships ( + id TEXT NOT NULL, + parent TEXT NOT NULL, + PRIMARY KEY(id, parent) + ) + """ + ) + self.conn.commit() def add_relation(self, object: ObjectID, parent: ObjectID): - pass + self.cursor.execute( + """ + INSERT INTO relationships (id, parent) + VALUES (?, ?, ?, ?, ?) + """, + ( + str(object), + str(parent), + ), + ) + self.conn.commit() def get_related_objects(self, object: ObjectID) -> [ObjectID]: - pass + parents = [] + cur = object + while True: + self.cursor.execute( + """ + SELECT id, parent FROM relationships WHERE id = ? + """, + (str(cur),), + ) + parent = self.cursor.fetchone() + if parent is None: + break + parents.append(parent) + cur = parent def delete_relation(self, object_id: ObjectID) -> [ObjectID]: - pass \ No newline at end of file + self.cursor.execute( + """ + DELETE FROM relationships WHERE id = ? + """, + (str(object_id),), + ) + self.conn.commit() \ No newline at end of file