File size: 3,286 Bytes
f02fdcc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 | import aiosqlite
from datetime import datetime
from pathlib import Path
from typing import List, Optional
from .models import MemoryItem
DB_PATH = Path("data/memory.db")
async def init_db():
"""Initialize database and create tables securely."""
Path("data").mkdir(exist_ok=True)
async with aiosqlite.connect(DB_PATH) as db:
await db.execute("""
CREATE TABLE IF NOT EXISTS memory_items (
item_id TEXT PRIMARY KEY,
tier INTEGER NOT NULL DEFAULT 5,
title TEXT NOT NULL,
content TEXT NOT NULL,
tags TEXT NOT NULL,
confidence REAL NOT NULL,
usage_count INTEGER NOT NULL,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
metadata TEXT NOT NULL
)
""")
await db.execute("CREATE INDEX IF NOT EXISTS idx_tags ON memory_items(tags)")
await db.execute("CREATE INDEX IF NOT EXISTS idx_tier ON memory_items(tier)")
await db.commit()
class DatabaseMemoryStore:
def __init__(self):
self.db_path = DB_PATH
async def add(self, item: MemoryItem) -> MemoryItem:
async with aiosqlite.connect(self.db_path) as db:
await db.execute("""
INSERT OR REPLACE INTO memory_items
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""", (
item.item_id,
item.tier.value if hasattr(item.tier, 'value') else (getattr(item, 'tier', 5) if hasattr(item, 'tier') else 5),
item.title,
item.content,
",".join(item.tags),
item.confidence,
item.usage_count,
item.created_at.isoformat(),
item.updated_at.isoformat(),
str(item.metadata)
))
await db.commit()
return item
async def get(self, item_id: str) -> Optional[MemoryItem]:
async with aiosqlite.connect(self.db_path) as db:
async with db.execute("SELECT * FROM memory_items WHERE item_id = ?", (item_id,)) as cursor:
row = await cursor.fetchone()
if row:
return self._row_to_item(row)
return None
async def search(self, query: str, limit: int = 10) -> List[MemoryItem]:
async with aiosqlite.connect(self.db_path) as db:
async with db.execute("""
SELECT * FROM memory_items
WHERE tags LIKE ? OR content LIKE ?
ORDER BY tier DESC, usage_count DESC
LIMIT ?
""", (f"%{query}%", f"%{query}%", limit)) as cursor:
rows = await cursor.fetchall()
return [self._row_to_item(row) for row in rows]
def _row_to_item(self, row) -> MemoryItem:
return MemoryItem(
item_id=row[0],
tier=row[1],
title=row[2],
content=row[3],
tags=row[4].split(",") if row[4] else [],
confidence=row[5],
usage_count=row[6],
created_at=datetime.fromisoformat(row[7]),
updated_at=datetime.fromisoformat(row[8]),
metadata=eval(row[9]) if row[9] else {}
) |