Add audit and memory UI
This commit is contained in:
@@ -0,0 +1,179 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import aiosqlite
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from duck_core.tasks.store import utc_now
|
||||
|
||||
|
||||
class MemoryRecord(BaseModel):
|
||||
id: int | None = None
|
||||
memory_id: str
|
||||
text: str
|
||||
workspace: str
|
||||
conversation_id: str | None = None
|
||||
memory_type: str = "note"
|
||||
importance: float = 0.5
|
||||
metadata: dict[str, Any] = Field(default_factory=dict)
|
||||
created_at: str
|
||||
updated_at: str
|
||||
|
||||
|
||||
class MemoryStore:
|
||||
def __init__(self, db_path: str):
|
||||
self.db_path = Path(db_path)
|
||||
|
||||
async def init(self) -> None:
|
||||
self.db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
await db.execute(
|
||||
"""
|
||||
create table if not exists memories (
|
||||
id integer primary key autoincrement,
|
||||
memory_id text not null unique,
|
||||
text text not null,
|
||||
workspace text not null,
|
||||
conversation_id text,
|
||||
memory_type text not null,
|
||||
importance real not null,
|
||||
metadata_json text not null,
|
||||
created_at text not null,
|
||||
updated_at text not null
|
||||
)
|
||||
"""
|
||||
)
|
||||
await db.execute(
|
||||
"""
|
||||
create index if not exists idx_memories_workspace_created
|
||||
on memories(workspace, created_at)
|
||||
"""
|
||||
)
|
||||
await db.commit()
|
||||
|
||||
async def add(
|
||||
self,
|
||||
text: str,
|
||||
workspace: str,
|
||||
conversation_id: str | None = None,
|
||||
memory_type: str = "note",
|
||||
importance: float = 0.5,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> MemoryRecord:
|
||||
await self.init()
|
||||
now = utc_now()
|
||||
memory_id = f"mem_{uuid4().hex[:12]}"
|
||||
clean_text = " ".join(text.strip().split())
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
insert into memories(
|
||||
memory_id, text, workspace, conversation_id, memory_type,
|
||||
importance, metadata_json, created_at, updated_at
|
||||
) values (?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
memory_id,
|
||||
clean_text,
|
||||
workspace,
|
||||
conversation_id,
|
||||
memory_type or "note",
|
||||
max(0.0, min(float(importance), 1.0)),
|
||||
json.dumps(metadata or {}, ensure_ascii=False),
|
||||
now,
|
||||
now,
|
||||
),
|
||||
)
|
||||
await db.commit()
|
||||
row_id = cursor.lastrowid
|
||||
return MemoryRecord(
|
||||
id=row_id,
|
||||
memory_id=memory_id,
|
||||
text=clean_text,
|
||||
workspace=workspace,
|
||||
conversation_id=conversation_id,
|
||||
memory_type=memory_type or "note",
|
||||
importance=max(0.0, min(float(importance), 1.0)),
|
||||
metadata=metadata or {},
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
|
||||
async def list(
|
||||
self, workspace: str | None = None, limit: int = 50
|
||||
) -> list[MemoryRecord]:
|
||||
await self.init()
|
||||
bounded_limit = min(max(limit, 1), 200)
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
if workspace:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
select * from memories
|
||||
where workspace = ?
|
||||
order by importance desc, created_at desc
|
||||
limit ?
|
||||
""",
|
||||
(workspace, bounded_limit),
|
||||
)
|
||||
else:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
select * from memories
|
||||
order by importance desc, created_at desc
|
||||
limit ?
|
||||
""",
|
||||
(bounded_limit,),
|
||||
)
|
||||
rows = await cursor.fetchall()
|
||||
return [self._row_to_record(row) for row in rows]
|
||||
|
||||
async def search(
|
||||
self, query: str, workspace: str | None = None, limit: int = 20
|
||||
) -> list[MemoryRecord]:
|
||||
await self.init()
|
||||
bounded_limit = min(max(limit, 1), 100)
|
||||
pattern = f"%{query.strip()}%"
|
||||
async with aiosqlite.connect(self.db_path) as db:
|
||||
db.row_factory = aiosqlite.Row
|
||||
if workspace:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
select * from memories
|
||||
where workspace = ?
|
||||
and (text like ? or memory_type like ? or metadata_json like ?)
|
||||
order by importance desc, created_at desc
|
||||
limit ?
|
||||
""",
|
||||
(workspace, pattern, pattern, pattern, bounded_limit),
|
||||
)
|
||||
else:
|
||||
cursor = await db.execute(
|
||||
"""
|
||||
select * from memories
|
||||
where text like ? or memory_type like ? or metadata_json like ?
|
||||
order by importance desc, created_at desc
|
||||
limit ?
|
||||
""",
|
||||
(pattern, pattern, pattern, bounded_limit),
|
||||
)
|
||||
rows = await cursor.fetchall()
|
||||
return [self._row_to_record(row) for row in rows]
|
||||
|
||||
def _row_to_record(self, row: aiosqlite.Row) -> MemoryRecord:
|
||||
return MemoryRecord(
|
||||
id=row["id"],
|
||||
memory_id=row["memory_id"],
|
||||
text=row["text"],
|
||||
workspace=row["workspace"],
|
||||
conversation_id=row["conversation_id"],
|
||||
memory_type=row["memory_type"],
|
||||
importance=float(row["importance"]),
|
||||
metadata=json.loads(row["metadata_json"]),
|
||||
created_at=row["created_at"],
|
||||
updated_at=row["updated_at"],
|
||||
)
|
||||
Reference in New Issue
Block a user