implement object relation store with sqlite

This commit is contained in:
tsukasa
2023-09-14 16:03:56 +08:00
parent aaf0f15228
commit 9cbf29e0b0
2 changed files with 52 additions and 5 deletions
-1
View File
@@ -1,6 +1,5 @@
# define a knowledge base class
import json
import pickle
from . import AgentPrompt, ComputeKernel
from ..knowledge import *
+52 -4
View File
@@ -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
self.cursor.execute(
"""
DELETE FROM relationships WHERE id = ?
""",
(str(object_id),),
)
self.conn.commit()