310 lines
13 KiB
Python
310 lines
13 KiB
Python
import aiosqlite
|
|
import json
|
|
import logging
|
|
from datetime import datetime
|
|
from typing import List, Dict, Any, Optional
|
|
|
|
DB_PATH = "data/context_hub.db"
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
async def init_db():
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS rules (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
scope TEXT UNIQUE NOT NULL,
|
|
content TEXT NOT NULL
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS api_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
agent_name TEXT UNIQUE NOT NULL,
|
|
api_key TEXT UNIQUE NOT NULL,
|
|
last_used TIMESTAMP
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS audit_log (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
agent_name TEXT,
|
|
scope TEXT,
|
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
ip_address TEXT
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS ports (
|
|
port INTEGER PRIMARY KEY,
|
|
service_name TEXT NOT NULL,
|
|
description TEXT
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS agents (
|
|
agent_name TEXT PRIMARY KEY,
|
|
agent_type TEXT,
|
|
scopes TEXT NOT NULL,
|
|
repo_url TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TABLE IF NOT EXISTS memory_entries (
|
|
id TEXT PRIMARY KEY,
|
|
scope TEXT NOT NULL,
|
|
project TEXT,
|
|
type TEXT NOT NULL,
|
|
title TEXT NOT NULL,
|
|
body TEXT NOT NULL,
|
|
tags TEXT,
|
|
status TEXT NOT NULL DEFAULT 'active',
|
|
superseded_by TEXT,
|
|
created_by TEXT NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE VIRTUAL TABLE IF NOT EXISTS memory_entries_fts USING fts5(
|
|
title, body, tags, content='memory_entries', content_rowid='rowid'
|
|
)
|
|
""")
|
|
await db.execute("""
|
|
CREATE TRIGGER IF NOT EXISTS memory_entries_ai AFTER INSERT ON memory_entries BEGIN
|
|
INSERT INTO memory_entries_fts(rowid, title, body, tags) VALUES (new.rowid, new.title, new.body, new.tags);
|
|
END
|
|
""")
|
|
await db.execute("""
|
|
CREATE TRIGGER IF NOT EXISTS memory_entries_ad AFTER DELETE ON memory_entries BEGIN
|
|
INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, body, tags) VALUES('delete', old.rowid, old.title, old.body, old.tags);
|
|
END
|
|
""")
|
|
await db.execute("""
|
|
CREATE TRIGGER IF NOT EXISTS memory_entries_au AFTER UPDATE ON memory_entries BEGIN
|
|
INSERT INTO memory_entries_fts(memory_entries_fts, rowid, title, body, tags) VALUES('delete', old.rowid, old.title, old.body, old.tags);
|
|
INSERT INTO memory_entries_fts(rowid, title, body, tags) VALUES (new.rowid, new.title, new.body, new.tags);
|
|
END
|
|
""")
|
|
await db.commit()
|
|
|
|
async def get_rule(scope: str) -> Optional[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT content FROM rules WHERE scope = ?", (scope,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return json.loads(row[0])
|
|
return None
|
|
|
|
async def set_rule(scope: str, content: Dict[str, Any]):
|
|
content_str = json.dumps(content)
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO rules (scope, content) VALUES (?, ?) ON CONFLICT(scope) DO UPDATE SET content=?",
|
|
(scope, content_str, content_str)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_all_rules() -> Dict[str, Any]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT scope, content FROM rules") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return {row[0]: json.loads(row[1]) for row in rows}
|
|
|
|
async def get_all_ports() -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT port, service_name, description FROM ports") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def add_port(port: int, service_name: str, description: str = ""):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO ports (port, service_name, description) VALUES (?, ?, ?) ON CONFLICT(port) DO UPDATE SET service_name=?, description=?",
|
|
(port, service_name, description, service_name, description)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_port(port: int) -> Optional[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT port, service_name, description FROM ports WHERE port = ?", (port,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return dict(row)
|
|
return None
|
|
|
|
async def add_api_key(agent_name: str, api_key: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO api_keys (agent_name, api_key) VALUES (?, ?) ON CONFLICT(agent_name) DO UPDATE SET api_key=?",
|
|
(agent_name, api_key, api_key)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_agent_by_key(api_key: str) -> Optional[str]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
async with db.execute("SELECT agent_name FROM api_keys WHERE api_key = ?", (api_key,)) as cursor:
|
|
row = await cursor.fetchone()
|
|
if row:
|
|
return row[0]
|
|
return None
|
|
|
|
async def update_key_last_used(api_key: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute("UPDATE api_keys SET last_used = CURRENT_TIMESTAMP WHERE api_key = ?", (api_key,))
|
|
await db.commit()
|
|
|
|
async def get_all_keys() -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT id, agent_name, api_key, last_used FROM api_keys") as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
async def add_audit_log(agent_name: str, scope: str, ip_address: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"INSERT INTO audit_log (agent_name, scope, ip_address) VALUES (?, ?, ?)",
|
|
(agent_name, scope, ip_address)
|
|
)
|
|
await db.commit()
|
|
|
|
async def get_audit_logs(limit: int = 50) -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT id, agent_name, scope, timestamp, ip_address FROM audit_log ORDER BY timestamp DESC LIMIT ?", (limit,)) as cursor:
|
|
rows = await cursor.fetchall()
|
|
return [dict(row) for row in rows]
|
|
|
|
|
|
_TYPE_ABBR = {
|
|
"decision": "dec",
|
|
"constraint": "cst",
|
|
"best-practice": "bp",
|
|
"common-error": "err",
|
|
"do-not-use": "dnu",
|
|
}
|
|
|
|
async def _next_entry_id(db, scope: str, type_: str) -> str:
|
|
abbr = _TYPE_ABBR.get(type_, type_)
|
|
prefix = f"{scope}-{abbr}-"
|
|
async with db.execute(
|
|
"SELECT id FROM memory_entries WHERE scope=? AND type=? ORDER BY created_at DESC, id DESC LIMIT 1",
|
|
(scope, type_)
|
|
) as cur:
|
|
row = await cur.fetchone()
|
|
n = 1
|
|
if row:
|
|
try:
|
|
n = int(row[0].rsplit("-", 1)[-1]) + 1
|
|
except ValueError:
|
|
n = 1
|
|
return f"{prefix}{n:04d}"
|
|
|
|
async def add_memory_entry(scope: str, type_: str, title: str, body: str, created_by: str,
|
|
project: Optional[str] = None, tags: Optional[List[str]] = None) -> Dict[str, Any]:
|
|
tags_json = json.dumps(tags or [])
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
entry_id = await _next_entry_id(db, scope, type_)
|
|
await db.execute(
|
|
"""INSERT INTO memory_entries (id, scope, project, type, title, body, tags, created_by)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(entry_id, scope, project, type_, title, body, tags_json, created_by)
|
|
)
|
|
await db.commit()
|
|
return {"id": entry_id, "scope": scope, "project": project, "type": type_,
|
|
"title": title, "body": body, "tags": tags or [], "status": "active", "created_by": created_by}
|
|
|
|
async def list_memory_entries(scope: str, project: Optional[str] = None,
|
|
type_: Optional[str] = None, status: str = "active") -> List[Dict[str, Any]]:
|
|
query = ("SELECT id, scope, project, type, title, body, tags, status, superseded_by, "
|
|
"created_by, created_at, updated_at FROM memory_entries WHERE scope=?")
|
|
params: List[Any] = [scope]
|
|
if project is not None:
|
|
query += " AND project=?"
|
|
params.append(project)
|
|
if type_ is not None:
|
|
query += " AND type=?"
|
|
params.append(type_)
|
|
if status:
|
|
query += " AND status=?"
|
|
params.append(status)
|
|
query += " ORDER BY created_at DESC"
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(query, params) as cursor:
|
|
rows = await cursor.fetchall()
|
|
result = []
|
|
for row in rows:
|
|
d = dict(row)
|
|
d["tags"] = json.loads(d["tags"]) if d["tags"] else []
|
|
result.append(d)
|
|
return result
|
|
|
|
async def search_memory_entries(query_text: str, scope: Optional[str] = None) -> List[Dict[str, Any]]:
|
|
sql = ("SELECT me.id, me.scope, me.project, me.type, me.title, me.body, me.tags, "
|
|
"me.status, me.created_by, me.created_at "
|
|
"FROM memory_entries_fts JOIN memory_entries me ON me.rowid = memory_entries_fts.rowid "
|
|
"WHERE memory_entries_fts MATCH ? AND me.status = 'active'")
|
|
params: List[Any] = [query_text]
|
|
if scope:
|
|
sql += " AND me.scope = ?"
|
|
params.append(scope)
|
|
sql += " ORDER BY rank"
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute(sql, params) as cursor:
|
|
rows = await cursor.fetchall()
|
|
result = []
|
|
for row in rows:
|
|
d = dict(row)
|
|
d["tags"] = json.loads(d["tags"]) if d["tags"] else []
|
|
result.append(d)
|
|
return result
|
|
|
|
async def supersede_memory_entry(entry_id: str, superseded_by: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"UPDATE memory_entries SET status='superseded', superseded_by=?, updated_at=CURRENT_TIMESTAMP WHERE id=?",
|
|
(superseded_by, entry_id)
|
|
)
|
|
await db.commit()
|
|
|
|
async def archive_memory_entry(entry_id: str):
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"UPDATE memory_entries SET status='archived', updated_at=CURRENT_TIMESTAMP WHERE id=?",
|
|
(entry_id,)
|
|
)
|
|
await db.commit()
|
|
|
|
|
|
async def add_agent(agent_name: str, scopes: List[str], agent_type: Optional[str] = None,
|
|
repo_url: Optional[str] = None) -> Dict[str, Any]:
|
|
agent_name = agent_name.upper()
|
|
scopes_json = json.dumps(scopes)
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
await db.execute(
|
|
"""INSERT INTO agents (agent_name, agent_type, scopes, repo_url) VALUES (?, ?, ?, ?)
|
|
ON CONFLICT(agent_name) DO UPDATE SET agent_type=excluded.agent_type,
|
|
scopes=excluded.scopes, repo_url=excluded.repo_url""",
|
|
(agent_name, agent_type, scopes_json, repo_url)
|
|
)
|
|
await db.commit()
|
|
return {"agent_name": agent_name, "agent_type": agent_type, "scopes": scopes, "repo_url": repo_url}
|
|
|
|
async def list_agents() -> List[Dict[str, Any]]:
|
|
async with aiosqlite.connect(DB_PATH) as db:
|
|
db.row_factory = aiosqlite.Row
|
|
async with db.execute("SELECT agent_name, agent_type, scopes, repo_url, created_at FROM agents ORDER BY agent_name") as cursor:
|
|
rows = await cursor.fetchall()
|
|
result = []
|
|
for row in rows:
|
|
d = dict(row)
|
|
d["scopes"] = json.loads(d["scopes"])
|
|
result.append(d)
|
|
return result
|