summaryrefslogtreecommitdiff
path: root/memory_store.py
blob: 2c69113605660cd2123c2208b8e5362c1d160b67 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
from datetime import datetime
from uuid import UUID

from database.database import Database
from memory import LongTermMemory, MemoryCategory


class LongTermMemoryStore:
    def __init__(self, database: Database) -> None:
        self.database = database

    def add(self, memory: LongTermMemory) -> None:
        created_at = memory.created_at.isoformat().replace("+00:00", "Z")

        with self.database.connect() as connection:
            connection.execute(
                """
                INSERT INTO memories (id, content, category, created_at)
                VALUES (?, ?, ?, ?)
                """,
                (str(memory.id), memory.content, memory.category.value, created_at),
            )

    def get(self, memory_id: UUID) -> LongTermMemory | None:
        with self.database.connect() as connection:
            row = connection.execute(
                """
                SELECT id, content, category, created_at
                FROM memories
                WHERE id = ?
                """,
                (str(memory_id),),
            ).fetchone()

        if row is None:
            return None

        return self._from_row(row)

    def list(self) -> list[LongTermMemory]:
        with self.database.connect() as connection:
            rows = connection.execute(
                """
                SELECT id, content, category, created_at
                FROM memories
                ORDER BY id
                """
            ).fetchall()

        return [self._from_row(row) for row in rows]

    def delete(self, memory_id: UUID) -> None:
        with self.database.connect() as connection:
            connection.execute(
                "DELETE FROM memories WHERE id = ?",
                (str(memory_id),),
            )

    def search(
        self,
        query: str,
        category: MemoryCategory | None = None,
        limit: int = 10,
    ) -> list[LongTermMemory]:
        if limit <= 0:
            raise ValueError("Search limit must be positive")

        if not query.strip():
            raise ValueError("Search query must not be empty")

        sql = """
            SELECT id, content, category, created_at
            FROM memories
            WHERE content LIKE ?
        """
        parameters: list[str | int] = [f"%{query}%"]

        if category is not None:
            sql += " AND category = ?"
            parameters.append(category.value)

        sql += " ORDER BY id LIMIT ?"
        parameters.append(limit)

        with self.database.connect() as connection:
            rows = connection.execute(sql, parameters).fetchall()

        return [self._from_row(row) for row in rows]

    @staticmethod
    def _from_row(row: tuple[str, str, str, str]) -> LongTermMemory:
        memory_id, content, category, created_at = row

        return LongTermMemory(
            id=UUID(memory_id),
            content=content,
            category=MemoryCategory(category),
            created_at=datetime.fromisoformat(created_at),
        )