| | |
| |
|
| | 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() |
| |
|