182 lines
6.8 KiB
Python
182 lines
6.8 KiB
Python
|
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)
|