PlantAI-dev/gfa_core/mas/memory/case.py

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)