ducklm/duck_core/memory/store.py

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"],
        )