初始化demo

agent_demo/persistence/__init__.py

@@ -0,0 +1,23 @@
+from .base import Session, SessionStore
+from .json_store import JSONSessionStore
+from .sqlite_store import SQLiteSessionStore
+
+
+__all__ = [
+    "Session",
+    "SessionStore",
+    "JSONSessionStore",
+    "SQLiteSessionStore",
+    "get_session_store",
+]
+
+
+def get_session_store() -> SessionStore:
+    from .. import config
+    backend = config.SESSION_BACKEND
+    if backend == "sqlite":
+        return SQLiteSessionStore(config.SQLITE_PATH)
+    if backend == "json":
+        return JSONSessionStore(config.SESSION_DIR)
+    raise ValueError(f"Unsupported SESSION_BACKEND: {backend}")
+

agent_demo/persistence/base.py

@@ -0,0 +1,42 @@
+import re
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import List
+
+SAFE_ID = re.compile(r"^[a-zA-Z0-9._-]")
+
+
+def sanitize_id(value: str, fallback:str) -> str:
+    """Sanitize a string to be used as an ID by removing unsafe characters."""
+    return SAFE_ID.sub("_", value) or fallback
+
+
+@dataclass
+class Session:
+    """Conversasion history for one (user_id, session_id) pair.
+    ``message`` is in Strands / Bedrock-converse format (provider-neutral::
+        {"role": "user"|"assistant"|"system",
+         "content": [{"text": ...} | {"toolUse": ...} | {"toolResult": ...} ]})
+    """
+    session_id: str
+    user_id: str
+    messages: List[dict] = field(default_factory=list)
+    
+
+class SessionStore(ABC):
+    """Abstract base class for session storage backends."""
+
+    @abstractmethod
+    def load(self, user_id: str, session_id: str) -> Session:
+        """Return prior history for the key (empty Session if none exists yet)."""
+        pass
+
+    @abstractmethod
+    def save(self, session: Session) -> None:
+        """Save a session to the store."""
+        pass
+    
+    @abstractmethod
+    def clear(self, user_id: str, session_id: str) -> None:
+        """Clear a session from the store."""
+        pass

agent_demo/persistence/json_store.py

@@ -0,0 +1,53 @@
+import json
+import os
+
+from .base import Session, SessionStore, sanitize_id
+
+
+class JSONSessionStore(SessionStore):
+    """A simple JSON file-based session store."""
+
+    def __init__(self, base_dir: str):
+        self.base_dir = base_dir
+        os.makedirs(base_dir, exist_ok=True)
+
+    def _path(self, user_id: str, session_id: str) -> str:
+        return os.path.join(
+            self.base_dir,
+            f"{sanitize_id(user_id, 'anon')}__{sanitize_id(session_id, 'default')}.json")
+
+    def load(self, user_id: str, session_id: str) -> Session:
+        path = self._path(user_id, session_id)
+        messages = []
+        if os.path.exists(path):
+            try:
+                with open(path, 'r', encoding='utf-8') as f:
+                    messages = json.load(f).get('messages', [])
+                    
+            except (json.JSONDecodeError, OSError):
+                messages = []
+        
+        return Session(session_id=session_id, user_id=user_id, messages=messages)
+
+
+    def save(self, session: Session) -> None:
+        path = self._path(session.user_id, session.session_id)
+        tmp = path + ".tmp"
+        
+        with open(path, 'w', encoding='utf-8') as f:
+            json.dump(
+                {
+                    "user_id": session.user_id,
+                    "session_id": session.session_id,
+                    "messages": session.messages,
+                },
+                f,
+                ensure_ascii=False,
+                indent=2,
+            )
+        os.replace(tmp, path)
+
+    def clear(self, user_id: str, session_id: str) -> None:
+        path = self._path(user_id, session_id)
+        if os.path.exists(path):
+            os.remove(path)

agent_demo/persistence/sqlite_store.py

@@ -0,0 +1,74 @@
+import json
+import sqlite3
+from typing import Optional
+
+from .base import Session, SessionStore
+
+_SCHEMA = """
+CREATE TABLE IF NOT EXISTS sessions (
+    user_id TEXT NOT NULL,
+    session_id TEXT NOT NULL,
+    messages TEXT NOT NULL,    -- JSON-encoded message list
+    updated_at REAL NOT NULL,
+    PRIMARY KEY (user_id, session_id)
+);
+"""
+
+
+class SQLiteSessionStore(SessionStore):
+    """A simple SQLite-based session store."""
+
+    def __init__(self, db_path: str):
+        self.db_path = db_path
+        self._init_schema()
+    
+    def _connect(self) -> sqlite3.Connection:
+        conn = sqlite3.connect(self.db_path, check_same_thread=False)
+        conn.row_factory = sqlite3.Row
+        return conn
+
+    def _init_schema(self):
+        with self._connect() as conn:
+            conn.execute(_SCHEMA)
+
+    def load(self, user_id: str, session_id: str) -> Session:
+        with self._connect() as conn:
+            row: Optional[sqlite3.Row] = conn.execute(
+                "SELECT messages FROM agent_session WHERE user_id = ? AND session_id = ?",
+                (user_id, session_id),
+            ).fetchone()
+            messages = []
+            if row is not None:
+                try:
+                    messages = json.loads(row["messages"])
+                except (json.JSONDecodeError, TypeError):
+                    messages = []
+            return Session(session_id=session_id, user_id=user_id, messages=messages)
+
+    def save(self, session: Session) -> None:
+        import time
+        
+        with self._connect() as conn:
+            conn.execute(
+                """
+                INSERT INTO agent_session (user_id, session_id, messages, updated_at)
+                VALUES (?, ?, ?, ?)
+                ON CONFLICT(user_id, session_id)
+                DO UPDATE SET messages=excluded.messages, updated_at=excluded.updated_at
+                """,
+                (
+                    session.user_id,
+                    session.session_id,
+                    json.dumps(session.messages, ensure_ascii=False),
+                    time.time()
+                ),
+            )
+            conn.commit()
+
+    def clear(self, user_id: str, session_id: str) -> None:
+        with self._connect() as conn:
+            conn.execute(
+                "DELETE FROM agent_session WHERE user_id = ? AND session_id = ?",
+                (user_id, session_id),
+            )
+            conn.commit()