Add gfa_core/mas/memory module
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
potaser
Executable
+181
@@ -0,0 +1,181 @@
|
||||
import json
|
||||
import sqlite3
|
||||
import os
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from gfa_core.mas.memory.base import BaseMemory, SearchResult
|
||||
from gfa_core.project_paths import MEMORY_ROOT
|
||||
|
||||
|
||||
@dataclass
|
||||
class AnalysisCase:
|
||||
case_id: str
|
||||
user_id: str | None = None
|
||||
gene_family: str | None = None
|
||||
species: str | None = None
|
||||
workflow: list[dict] = field(default_factory=list)
|
||||
results: dict[str, Any] = field(default_factory=dict)
|
||||
status: str = "pending"
|
||||
created_at: datetime = field(default_factory=datetime.now)
|
||||
updated_at: datetime = field(default_factory=datetime.now)
|
||||
completed_at: datetime | None = None
|
||||
feedback: dict[str, Any] = field(default_factory=dict)
|
||||
metadata: dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
class CaseMemory(BaseMemory):
|
||||
def __init__(self, db_path: str = None):
|
||||
if db_path is None:
|
||||
db_path = os.environ.get("BBAV5_CASE_DB") or MEMORY_ROOT / "cases.db"
|
||||
self._db_path = Path(db_path)
|
||||
self._db_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self._init_db()
|
||||
print(f"[CaseMemory] SQLite storage: {self._db_path}")
|
||||
|
||||
def _init_db(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS cases (
|
||||
case_id TEXT PRIMARY KEY,
|
||||
user_id TEXT,
|
||||
gene_family TEXT,
|
||||
species TEXT,
|
||||
workflow TEXT,
|
||||
results TEXT,
|
||||
status TEXT,
|
||||
created_at TEXT,
|
||||
updated_at TEXT,
|
||||
completed_at TEXT,
|
||||
feedback TEXT,
|
||||
metadata TEXT
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_user_id ON cases(user_id)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_gene_family ON cases(gene_family)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE INDEX IF NOT EXISTS idx_status ON cases(status)
|
||||
""")
|
||||
|
||||
async def store(self, key: str, value: Any, ttl: int | None = None) -> None:
|
||||
if isinstance(value, AnalysisCase):
|
||||
case = value
|
||||
else:
|
||||
case = AnalysisCase(case_id=key, **value)
|
||||
|
||||
case.updated_at = datetime.now()
|
||||
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("""
|
||||
INSERT OR REPLACE INTO cases
|
||||
(case_id, user_id, gene_family, species, workflow, results, status,
|
||||
created_at, updated_at, completed_at, feedback, metadata)
|
||||
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
||||
""", (
|
||||
case.case_id,
|
||||
case.user_id,
|
||||
case.gene_family,
|
||||
case.species,
|
||||
json.dumps(case.workflow),
|
||||
json.dumps(case.results),
|
||||
case.status,
|
||||
case.created_at.isoformat(),
|
||||
case.updated_at.isoformat(),
|
||||
case.completed_at.isoformat() if case.completed_at else None,
|
||||
json.dumps(case.feedback),
|
||||
json.dumps(case.metadata),
|
||||
))
|
||||
|
||||
async def retrieve(self, key: str) -> Any | None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
return self._row_to_case(row)
|
||||
|
||||
async def delete(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"DELETE FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.rowcount > 0
|
||||
|
||||
async def exists(self, key: str) -> bool:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT 1 FROM cases WHERE case_id = ?", (key,)
|
||||
)
|
||||
return cursor.fetchone() is not None
|
||||
|
||||
async def clear(self) -> None:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.execute("DELETE FROM cases")
|
||||
|
||||
def _row_to_case(self, row: sqlite3.Row) -> AnalysisCase:
|
||||
return AnalysisCase(
|
||||
case_id=row["case_id"],
|
||||
user_id=row["user_id"],
|
||||
gene_family=row["gene_family"],
|
||||
species=row["species"],
|
||||
workflow=json.loads(row["workflow"]) if row["workflow"] else [],
|
||||
results=json.loads(row["results"]) if row["results"] else {},
|
||||
status=row["status"],
|
||||
created_at=datetime.fromisoformat(row["created_at"]),
|
||||
updated_at=datetime.fromisoformat(row["updated_at"]),
|
||||
completed_at=datetime.fromisoformat(row["completed_at"]) if row["completed_at"] else None,
|
||||
feedback=json.loads(row["feedback"]) if row["feedback"] else {},
|
||||
metadata=json.loads(row["metadata"]) if row["metadata"] else {},
|
||||
)
|
||||
|
||||
async def list_by_user(self, user_id: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE user_id = ? ORDER BY created_at DESC",
|
||||
(user_id,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_gene_family(self, gene_family: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE gene_family = ? ORDER BY created_at DESC",
|
||||
(gene_family,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def list_by_status(self, status: str) -> list[AnalysisCase]:
|
||||
with sqlite3.connect(self._db_path) as conn:
|
||||
conn.row_factory = sqlite3.Row
|
||||
cursor = conn.execute(
|
||||
"SELECT * FROM cases WHERE status = ? ORDER BY created_at DESC",
|
||||
(status,)
|
||||
)
|
||||
return [self._row_to_case(row) for row in cursor.fetchall()]
|
||||
|
||||
async def update_status(self, case_id: str, status: str) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.status = status
|
||||
if status == "completed":
|
||||
case.completed_at = datetime.now()
|
||||
await self.store(case_id, case)
|
||||
|
||||
async def add_feedback(self, case_id: str, feedback: dict[str, Any]) -> None:
|
||||
case = await self.retrieve(case_id)
|
||||
if case:
|
||||
case.feedback.update(feedback)
|
||||
await self.store(case_id, case)
|
||||
Reference in New Issue
Block a user