chore: import upstream snapshot with attribution

python/samples/05-end-to-end/chatkit-integration/store.py

@@ -0,0 +1,348 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""SQLite-based store implementation for ChatKit data persistence.
+
+This module provides a complete Store implementation using SQLite for data persistence.
+It includes proper thread safety, user isolation, and follows the ChatKit Store protocol.
+"""
+
+import sqlite3
+import uuid
+from typing import Any
+
+from chatkit.store import NotFoundError, Store
+from chatkit.types import (
+    Attachment,
+    Page,
+    ThreadItem,
+    ThreadMetadata,
+)
+from pydantic import BaseModel
+
+
+class ThreadData(BaseModel):
+    """Model for serializing thread data to SQLite."""
+
+    thread: ThreadMetadata
+
+
+class ItemData(BaseModel):
+    """Model for serializing thread item data to SQLite."""
+
+    item: ThreadItem
+
+
+class AttachmentData(BaseModel):
+    """Model for serializing attachment data to SQLite."""
+
+    attachment: Attachment
+
+
+class SQLiteStore(Store[dict[str, Any]]):
+    """SQLite-based store implementation for ChatKit data.
+
+    This implementation follows the pattern from the ChatKit Python tests
+    and provides persistent storage for threads, messages, and attachments.
+
+    Features:
+    - Thread-safe SQLite connections with WAL mode
+    - User isolation for multi-tenant support
+    - Proper error handling and transaction management
+    - Complete Store protocol implementation
+
+    Note: This is for demonstration purposes. In production, you should
+    implement proper error handling, connection pooling, and migration strategies.
+    """
+
+    def __init__(self, db_path: str | None = None):
+        self.db_path = db_path or "chatkit_demo.db"  # Use file-based DB for demo
+        self._create_tables()
+
+    def _create_connection(self):
+        # Enable thread safety and WAL mode for better concurrent access
+        conn = sqlite3.connect(self.db_path, check_same_thread=False)
+        conn.execute("PRAGMA journal_mode=WAL")
+        return conn
+
+    def _create_tables(self):
+        with self._create_connection() as conn:
+            # Create threads table
+            conn.execute(
+                """CREATE TABLE IF NOT EXISTS threads (
+                id TEXT PRIMARY KEY,
+                user_id TEXT NOT NULL,
+                created_at TEXT NOT NULL,
+                data TEXT NOT NULL
+                )"""
+            )
+
+            # Create items table
+            conn.execute(
+                """CREATE TABLE IF NOT EXISTS items (
+                id TEXT PRIMARY KEY,
+                thread_id TEXT NOT NULL,
+                user_id TEXT NOT NULL,
+                created_at TEXT NOT NULL,
+                data TEXT NOT NULL
+                )"""
+            )
+
+            # Create attachments table
+            conn.execute(
+                """CREATE TABLE IF NOT EXISTS attachments (
+                id TEXT PRIMARY KEY,
+                user_id TEXT NOT NULL,
+                data TEXT NOT NULL
+                )"""
+            )
+            conn.commit()
+
+    def generate_thread_id(self, context: dict[str, Any]) -> str:
+        return f"thr_{uuid.uuid4().hex[:8]}"
+
+    def generate_item_id(
+        self,
+        item_type: str,
+        thread: ThreadMetadata,
+        context: dict[str, Any],
+    ) -> str:
+        prefix_map = {
+            "message": "msg",
+            "tool_call": "tc",
+            "task": "tsk",
+            "workflow": "wf",
+            "attachment": "atc",
+        }
+        prefix = prefix_map.get(item_type, "itm")
+        return f"{prefix}_{uuid.uuid4().hex[:8]}"
+
+    async def load_thread(self, thread_id: str, context: dict[str, Any]) -> ThreadMetadata:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            cursor = conn.execute(
+                "SELECT data FROM threads WHERE id = ? AND user_id = ?",
+                (thread_id, user_id),
+            ).fetchone()
+
+            if cursor is None:
+                raise NotFoundError(f"Thread {thread_id} not found")
+
+            thread_data = ThreadData.model_validate_json(cursor[0])
+            return thread_data.thread
+
+    async def save_thread(self, thread: ThreadMetadata, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            thread_data = ThreadData(thread=thread)
+
+            # Replace existing thread data
+            conn.execute(
+                "DELETE FROM threads WHERE id = ? AND user_id = ?",
+                (thread.id, user_id),
+            )
+            conn.execute(
+                "INSERT INTO threads (id, user_id, created_at, data) VALUES (?, ?, ?, ?)",
+                (
+                    thread.id,
+                    user_id,
+                    thread.created_at.isoformat(),
+                    thread_data.model_dump_json(),
+                ),
+            )
+            conn.commit()
+
+    async def load_thread_items(
+        self,
+        thread_id: str,
+        after: str | None,
+        limit: int,
+        order: str,
+        context: dict[str, Any],
+    ) -> Page[ThreadItem]:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            created_after: str | None = None
+            if after:
+                after_cursor = conn.execute(
+                    "SELECT created_at FROM items WHERE id = ? AND user_id = ?",
+                    (after, user_id),
+                ).fetchone()
+                if after_cursor is None:
+                    raise NotFoundError(f"Item {after} not found")
+                created_after = after_cursor[0]
+
+            query = """
+                SELECT data FROM items
+                WHERE thread_id = ? AND user_id = ?
+            """
+            params: list[Any] = [thread_id, user_id]
+
+            if created_after:
+                query += " AND created_at > ?" if order == "asc" else " AND created_at < ?"
+                params.append(created_after)
+
+            query += f" ORDER BY created_at {order} LIMIT ?"
+            params.append(limit + 1)
+
+            items_cursor = conn.execute(query, params).fetchall()
+            items = [ItemData.model_validate_json(row[0]).item for row in items_cursor]
+
+            has_more = len(items) > limit
+            if has_more:
+                items = items[:limit]
+
+            return Page[ThreadItem](data=items, has_more=has_more, after=items[-1].id if items else None)
+
+    async def save_attachment(self, attachment: Attachment, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            attachment_data = AttachmentData(attachment=attachment)
+            conn.execute(
+                "INSERT OR REPLACE INTO attachments (id, user_id, data) VALUES (?, ?, ?)",
+                (
+                    attachment.id,
+                    user_id,
+                    attachment_data.model_dump_json(),
+                ),
+            )
+            conn.commit()
+
+    async def load_attachment(self, attachment_id: str, context: dict[str, Any]) -> Attachment:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            cursor = conn.execute(
+                "SELECT data FROM attachments WHERE id = ? AND user_id = ?",
+                (attachment_id, user_id),
+            ).fetchone()
+
+            if cursor is None:
+                raise NotFoundError(f"Attachment {attachment_id} not found")
+
+            attachment_data = AttachmentData.model_validate_json(cursor[0])
+            return attachment_data.attachment
+
+    async def delete_attachment(self, attachment_id: str, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            conn.execute(
+                "DELETE FROM attachments WHERE id = ? AND user_id = ?",
+                (attachment_id, user_id),
+            )
+            conn.commit()
+
+    async def load_threads(
+        self,
+        limit: int,
+        after: str | None,
+        order: str,
+        context: dict[str, Any],
+    ) -> Page[ThreadMetadata]:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            created_after: str | None = None
+            if after:
+                after_cursor = conn.execute(
+                    "SELECT created_at FROM threads WHERE id = ? AND user_id = ?",
+                    (after, user_id),
+                ).fetchone()
+                if after_cursor is None:
+                    raise NotFoundError(f"Thread {after} not found")
+                created_after = after_cursor[0]
+
+            query = "SELECT data FROM threads WHERE user_id = ?"
+            params: list[Any] = [user_id]
+
+            if created_after:
+                query += " AND created_at > ?" if order == "asc" else " AND created_at < ?"
+                params.append(created_after)
+
+            query += f" ORDER BY created_at {order} LIMIT ?"
+            params.append(limit + 1)
+
+            threads_cursor = conn.execute(query, params).fetchall()
+            threads = [ThreadData.model_validate_json(row[0]).thread for row in threads_cursor]
+
+            has_more = len(threads) > limit
+            if has_more:
+                threads = threads[:limit]
+
+            return Page[ThreadMetadata](data=threads, has_more=has_more, after=threads[-1].id if threads else None)
+
+    async def add_thread_item(self, thread_id: str, item: ThreadItem, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            item_data = ItemData(item=item)
+            conn.execute(
+                "INSERT INTO items (id, thread_id, user_id, created_at, data) VALUES (?, ?, ?, ?, ?)",
+                (
+                    item.id,
+                    thread_id,
+                    user_id,
+                    item.created_at.isoformat(),
+                    item_data.model_dump_json(),
+                ),
+            )
+            conn.commit()
+
+    async def save_item(self, thread_id: str, item: ThreadItem, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            item_data = ItemData(item=item)
+            conn.execute(
+                "UPDATE items SET data = ? WHERE id = ? AND thread_id = ? AND user_id = ?",
+                (
+                    item_data.model_dump_json(),
+                    item.id,
+                    thread_id,
+                    user_id,
+                ),
+            )
+            conn.commit()
+
+    async def load_item(self, thread_id: str, item_id: str, context: dict[str, Any]) -> ThreadItem:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            cursor = conn.execute(
+                "SELECT data FROM items WHERE id = ? AND thread_id = ? AND user_id = ?",
+                (item_id, thread_id, user_id),
+            ).fetchone()
+
+            if cursor is None:
+                raise NotFoundError(f"Item {item_id} not found in thread {thread_id}")
+
+            item_data = ItemData.model_validate_json(cursor[0])
+            return item_data.item
+
+    async def delete_thread(self, thread_id: str, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            conn.execute(
+                "DELETE FROM threads WHERE id = ? AND user_id = ?",
+                (thread_id, user_id),
+            )
+            conn.execute(
+                "DELETE FROM items WHERE thread_id = ? AND user_id = ?",
+                (thread_id, user_id),
+            )
+            conn.commit()
+
+    async def delete_thread_item(self, thread_id: str, item_id: str, context: dict[str, Any]) -> None:
+        user_id = context.get("user_id", "demo_user")
+
+        with self._create_connection() as conn:
+            conn.execute(
+                "DELETE FROM items WHERE id = ? AND thread_id = ? AND user_id = ?",
+                (item_id, thread_id, user_id),
+            )
+            conn.commit()