|
|
|
|
|
|
|
|
import sqlite3 |
|
|
from datetime import datetime |
|
|
|
|
|
DEFAULT_DB_PATH = "agent_storage.db" |
|
|
|
|
|
class Storage: |
|
|
def __init__(self, config=None): |
|
|
self.config = config or {} |
|
|
db_path = self.config.get("db_path", DEFAULT_DB_PATH) |
|
|
self.conn = sqlite3.connect(db_path) |
|
|
self._init_db() |
|
|
|
|
|
def _init_db(self): |
|
|
c = self.conn.cursor() |
|
|
|
|
|
c.execute(''' |
|
|
CREATE TABLE IF NOT EXISTS diary_entries ( |
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT, |
|
|
text TEXT NOT NULL, |
|
|
tags TEXT, |
|
|
timestamp TEXT NOT NULL |
|
|
) |
|
|
''') |
|
|
|
|
|
c.execute(''' |
|
|
CREATE TABLE IF NOT EXISTS concepts ( |
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT, |
|
|
name TEXT NOT NULL, |
|
|
description TEXT |
|
|
) |
|
|
''') |
|
|
|
|
|
c.execute(''' |
|
|
CREATE TABLE IF NOT EXISTS links ( |
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT, |
|
|
source_id INTEGER, |
|
|
target_id INTEGER, |
|
|
relation TEXT, |
|
|
FOREIGN KEY(source_id) REFERENCES concepts(id), |
|
|
FOREIGN KEY(target_id) REFERENCES concepts(id) |
|
|
) |
|
|
''') |
|
|
self.conn.commit() |
|
|
|
|
|
|
|
|
def write_entry(self, text, tags=None): |
|
|
timestamp = datetime.utcnow().isoformat() |
|
|
tag_str = ",".join(tags) if tags else "" |
|
|
self.conn.execute( |
|
|
'INSERT INTO diary_entries (text, tags, timestamp) VALUES (?, ?, ?)', |
|
|
(text, tag_str, timestamp) |
|
|
) |
|
|
self.conn.commit() |
|
|
|
|
|
def read_entries(self, limit=10, tag_filter=None): |
|
|
cursor = self.conn.cursor() |
|
|
if tag_filter: |
|
|
|
|
|
if isinstance(tag_filter, list): |
|
|
tag_filter = ",".join(tag_filter) |
|
|
like_expr = f"%{tag_filter}%" |
|
|
cursor.execute( |
|
|
'SELECT * FROM diary_entries WHERE tags LIKE ? ORDER BY id DESC LIMIT ?', |
|
|
(like_expr, limit) |
|
|
) |
|
|
else: |
|
|
cursor.execute( |
|
|
'SELECT * FROM diary_entries ORDER BY id DESC LIMIT ?', |
|
|
(limit,) |
|
|
) |
|
|
return cursor.fetchall() |
|
|
|
|
|
def search_entries_by_time(self, from_ts, to_ts): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute( |
|
|
'SELECT * FROM diary_entries WHERE timestamp BETWEEN ? AND ? ORDER BY timestamp DESC', |
|
|
(from_ts, to_ts) |
|
|
) |
|
|
return cursor.fetchall() |
|
|
|
|
|
|
|
|
def add_concept(self, name, description=None): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute( |
|
|
'INSERT INTO concepts (name, description) VALUES (?, ?)', |
|
|
(name, description) |
|
|
) |
|
|
self.conn.commit() |
|
|
return cursor.lastrowid |
|
|
|
|
|
def query_concept(self, name_substr): |
|
|
cursor = self.conn.execute( |
|
|
'SELECT id, name, description FROM concepts WHERE name LIKE ?', |
|
|
(f"%{name_substr}%",) |
|
|
) |
|
|
return cursor.fetchall() |
|
|
|
|
|
def add_link(self, source_id, target_id, relation): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute( |
|
|
'INSERT INTO links (source_id, target_id, relation) VALUES (?, ?, ?)', |
|
|
(source_id, target_id, relation) |
|
|
) |
|
|
self.conn.commit() |
|
|
return cursor.lastrowid |
|
|
|
|
|
def list_concepts(self): |
|
|
return self.conn.execute('SELECT * FROM concepts').fetchall() |
|
|
|
|
|
def list_links(self): |
|
|
return self.conn.execute('SELECT * FROM links').fetchall() |
|
|
|
|
|
def expand_graph(self, start_id, depth): |
|
|
visited = set() |
|
|
results = [] |
|
|
|
|
|
def dfs(node_id, level): |
|
|
if level > depth or node_id in visited: |
|
|
return |
|
|
visited.add(node_id) |
|
|
cursor = self.conn.execute( |
|
|
'SELECT source_id, target_id, relation FROM links WHERE source_id=?', |
|
|
(node_id,) |
|
|
) |
|
|
for row in cursor.fetchall(): |
|
|
results.append(row) |
|
|
dfs(row[1], level + 1) |
|
|
|
|
|
dfs(start_id, 0) |
|
|
return results |
|
|
|
|
|
def delete_concept(self, concept_id): |
|
|
self.conn.execute('DELETE FROM concepts WHERE id = ?', (concept_id,)) |
|
|
self.conn.execute('DELETE FROM links WHERE source_id = ? OR target_id = ?', (concept_id, concept_id)) |
|
|
self.conn.commit() |
|
|
|
|
|
def delete_link(self, link_id): |
|
|
self.conn.execute('DELETE FROM links WHERE id = ?', (link_id,)) |
|
|
self.conn.commit() |
|
|
|
|
|
def delete_entry(self, entry_id): |
|
|
self.conn.execute('DELETE FROM diary_entries WHERE id = ?', (entry_id,)) |
|
|
self.conn.commit() |
|
|
|
|
|
def export_diary(self): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute('SELECT id, text, tags, timestamp FROM diary_entries ORDER BY id ASC') |
|
|
return cursor.fetchall() |
|
|
|
|
|
def export_graph(self): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute('SELECT id, name, description FROM concepts ORDER BY id ASC') |
|
|
concepts = cursor.fetchall() |
|
|
|
|
|
cursor.execute('SELECT id, source_id, target_id, relation FROM links ORDER BY id ASC') |
|
|
links = cursor.fetchall() |
|
|
|
|
|
return {"concepts": concepts, "links": links} |
|
|
|
|
|
def update_concept(self, concept_id, name=None, description=None): |
|
|
cursor = self.conn.cursor() |
|
|
if name is not None: |
|
|
cursor.execute('UPDATE concepts SET name = ? WHERE id = ?', (name, concept_id)) |
|
|
if description is not None: |
|
|
cursor.execute('UPDATE concepts SET description = ? WHERE id = ?', (description, concept_id)) |
|
|
self.conn.commit() |
|
|
|
|
|
def get_tag_stats(self): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute('SELECT tags FROM diary_entries') |
|
|
tag_counts = {} |
|
|
for row in cursor.fetchall(): |
|
|
tags = row[0].split(",") if row[0] else [] |
|
|
for tag in tags: |
|
|
tag = tag.strip() |
|
|
if tag: |
|
|
tag_counts[tag] = tag_counts.get(tag, 0) + 1 |
|
|
return tag_counts |
|
|
|
|
|
def search_links_by_relation(self, relation): |
|
|
cursor = self.conn.cursor() |
|
|
cursor.execute( |
|
|
'SELECT id, source_id, target_id, relation FROM links WHERE relation LIKE ?', |
|
|
(f"%{relation}%",) |
|
|
) |
|
|
return cursor.fetchall() |
|
|
|
|
|
def search_concepts(self, query): |
|
|
cursor = self.conn.execute( |
|
|
'''SELECT id, name, description FROM concepts |
|
|
WHERE name LIKE ? OR description LIKE ?''', |
|
|
(f"%{query}%", f"%{query}%") |
|
|
) |
|
|
return cursor.fetchall() |
|
|
|
|
|
def merge_concepts(self, source_id, target_id): |
|
|
cursor = self.conn.cursor() |
|
|
|
|
|
cursor.execute('UPDATE links SET source_id = ? WHERE source_id = ?', (target_id, source_id)) |
|
|
cursor.execute('UPDATE links SET target_id = ? WHERE target_id = ?', (target_id, source_id)) |
|
|
|
|
|
self.delete_concept(source_id) |
|
|
self.conn.commit() |
|
|
|
|
|
def find_concept_id_by_name(self, name): |
|
|
cursor = self.conn.execute('SELECT id FROM concepts WHERE name = ?', (name,)) |
|
|
row = cursor.fetchone() |
|
|
return row[0] if row else None |
|
|
|
|
|
def close(self): |
|
|
self.conn.close() |
|
|
|
