tinyagent-py 0.0.1__py3-none-any.whl → 0.0.4__py3-none-any.whl

storage/base.py

@@ -0,0 +1,49 @@
+from abc import ABC, abstractmethod
+from typing import Dict, Any, TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from tinyagent.tiny_agent import TinyAgent
+
+class Storage(ABC):
+    """
+    Abstract base class for TinyAgent session storage.
+    """
+
+    @abstractmethod
+    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
+        """
+        Persist the given agent state under `session_id`.
+        """
+        ...
+
+    @abstractmethod
+    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        """
+        Retrieve the agent state for `session_id`, or return {} if not found.
+        """
+        ...
+
+    @abstractmethod
+    async def close(self) -> None:
+        """
+        Clean up any resources (DB connections, file handles, etc.).
+        """
+        ...
+
+    def attach(self, agent: "TinyAgent") -> None:
+        """
+        Hook this storage to a TinyAgent so that on every `llm_end`
+        it will auto‐persist the agent's state.
+
+        Usage:
+            storage.attach(agent)
+        or in TinyAgent.__init__:
+            if storage: storage.attach(self)
+        """
+        async def _auto_save(event_name: str, agent: "TinyAgent", **kwargs):
+            if event_name != "llm_end":
+                return
+            state = agent.to_dict()
+            await self.save_session(agent.session_id, state)
+
+        agent.callbacks.append(_auto_save) 

storage/json_file_storage.py

@@ -0,0 +1,30 @@
+import json
+import asyncio
+from pathlib import Path
+from typing import Dict, Any, Union, Optional
+from tinyagent.storage import Storage
+
+class JsonFileStorage(Storage):
+    """
+    Persist TinyAgent sessions as individual JSON files.
+    """
+
+    def __init__(self, folder: Union[str, Path]):
+        self.folder = Path(folder)
+        self.folder.mkdir(parents=True, exist_ok=True)
+
+    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
+        path = self.folder / f"{session_id}_{user_id}.json"
+        # Write in a thread pool to avoid blocking the event loop
+        await asyncio.to_thread(path.write_text, json.dumps(data, indent=2), "utf-8")
+
+    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        path = self.folder / f"{session_id}_{user_id}.json"
+        if not path.exists():
+            return {}
+        text = await asyncio.to_thread(path.read_text, "utf-8")
+        return json.loads(text)
+
+    async def close(self) -> None:
+        # Nothing to clean up for file storage
+        return

storage/postgres_storage.py

@@ -0,0 +1,201 @@
+import asyncpg
+import json
+import logging
+from typing import Optional, Dict, Any
+from tinyagent.storage import Storage
+
+class PostgresStorage(Storage):
+    """
+    Persist TinyAgent sessions in a Postgres table with JSONB state.
+    """
+
+    def __init__(self, db_url: str, table_name: str = "tny_agent_sessions"):
+        self._dsn = db_url
+        self._table = table_name
+        self._pool: Optional[asyncpg.pool.Pool] = None
+        self.logger = logging.getLogger(__name__)
+
+    async def _ensure_table(self):
+        """Create the sessions table if it doesn't exist."""
+        self.logger.debug(f"Ensuring table {self._table} exists")
+        try:
+            async with self._pool.acquire() as conn:
+                await conn.execute(f"""
+                    CREATE TABLE IF NOT EXISTS {self._table} (
+                        agent_id TEXT PRIMARY KEY,
+                        session_id TEXT NOT NULL,
+                        user_id TEXT,
+                        memories JSONB,
+                        metadata JSONB,
+                        session_data JSONB,
+                        model_meta JSONB,
+                        created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
+                        updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
+                    );
+                    CREATE INDEX IF NOT EXISTS idx_{self._table}_session_id ON {self._table} (session_id);
+                    CREATE INDEX IF NOT EXISTS idx_{self._table}_user_id ON {self._table} (user_id);
+                """)
+                self.logger.info(f"Table {self._table} and indexes created/verified")
+        except Exception as e:
+            self.logger.error(f"Error creating table {self._table}: {str(e)}")
+            raise
+
+    async def _connect(self):
+        if not self._pool:
+            self.logger.debug(f"Connecting to PostgreSQL with DSN: {self._dsn[:10]}...")
+            try:
+                # Ensure statement_cache_size=0 to disable prepared statements for pgbouncer compatibility
+                self._pool = await asyncpg.create_pool(
+                    dsn=self._dsn, 
+                    statement_cache_size=0,
+                    min_size=1,
+                    max_size=10
+                )
+                self.logger.info("PostgreSQL connection pool created")
+                await self._ensure_table()
+            except Exception as e:
+                self.logger.error(f"Failed to connect to PostgreSQL: {str(e)}")
+                raise
+
+    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None):
+        self.logger.info(f"Saving session {session_id} for user {user_id}")
+        self.logger.debug(f"Save data: {json.dumps(data)[:200]}...")
+        
+        try:
+            await self._connect()
+            
+            # Extract data following the TinyAgent schema
+            metadata = data.get("metadata", {}) or {}
+            session_state = data.get("session_state", {}) or {}
+            
+            # Use session_id as agent_id if not provided
+            agent_id = metadata.get("agent_id", session_id)
+            self.logger.debug(f"Using agent_id: {agent_id}")
+            
+            # Extract specific components
+            memories = session_state.get("memory", {})
+            session_data = {"messages": session_state.get("messages", [])}
+            model_meta = metadata.get("model_meta", {})
+            
+            # Convert Python dictionaries to JSON strings for PostgreSQL
+            self.logger.debug("Converting Python dictionaries to JSON")
+            try:
+                memories_json = json.dumps(memories)
+                metadata_json = json.dumps(metadata)
+                session_data_json = json.dumps(session_data)
+                model_meta_json = json.dumps(model_meta)
+            except Exception as e:
+                self.logger.error(f"JSON serialization error: {str(e)}")
+                raise
+            
+            self.logger.debug("Executing PostgreSQL INSERT/UPDATE")
+            async with self._pool.acquire() as conn:
+                try:
+                    await conn.execute(f"""
+                        INSERT INTO {self._table} 
+                        (agent_id, session_id, user_id, memories, metadata, session_data, model_meta, updated_at)
+                        VALUES ($1, $2, $3, $4::jsonb, $5::jsonb, $6::jsonb, $7::jsonb, NOW())
+                        ON CONFLICT (agent_id) DO UPDATE
+                          SET session_id = EXCLUDED.session_id,
+                              user_id = EXCLUDED.user_id,
+                              memories = EXCLUDED.memories,
+                              metadata = EXCLUDED.metadata,
+                              session_data = EXCLUDED.session_data,
+                              model_meta = EXCLUDED.model_meta,
+                              updated_at = NOW();
+                    """, agent_id, session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json)
+                    self.logger.info(f"Session {session_id} saved successfully")
+                except Exception as e:
+                    self.logger.error(f"Database error during save: {str(e)}")
+                    raise
+        except Exception as e:
+            self.logger.error(f"Failed to save session {session_id}: {str(e)}")
+            raise
+
+    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        self.logger.info(f"Loading session {session_id} for user {user_id}")
+        
+        try:
+            await self._connect()
+            
+            async with self._pool.acquire() as conn:
+                # First try to find by session_id
+                query = f"""
+                    SELECT agent_id, session_id, user_id, memories, metadata, session_data, model_meta
+                      FROM {self._table}
+                     WHERE session_id = $1
+                """
+                params = [session_id]
+                
+                # Add user_id filter if provided
+                if user_id:
+                    query += " AND user_id = $2"
+                    params.append(user_id)
+                
+                self.logger.debug(f"Executing query: {query} with params: {params}")
+                row = await conn.fetchrow(query, *params)
+                
+                if not row:
+                    self.logger.warning(f"No session found for session_id={session_id}, user_id={user_id}")
+                    return {}
+                
+                self.logger.debug(f"Session found: {dict(row)}")
+                
+                # Parse JSON from PostgreSQL
+                try:
+                    # Check if values are already dictionaries or need parsing
+                    memories = row["memories"]
+                    if isinstance(memories, str):
+                        memories = json.loads(memories)
+                    
+                    metadata = row["metadata"]
+                    if isinstance(metadata, str):
+                        metadata = json.loads(metadata)
+                    
+                    session_data = row["session_data"]
+                    if isinstance(session_data, str):
+                        session_data = json.loads(session_data)
+                    
+                    model_meta = row["model_meta"]
+                    if isinstance(model_meta, str):
+                        model_meta = json.loads(model_meta)
+                except Exception as e:
+                    self.logger.error(f"Error parsing JSON from database: {str(e)}")
+                    raise
+                
+                # Update metadata with additional fields
+                metadata.update({
+                    "agent_id": row["agent_id"],
+                    "user_id": row["user_id"],
+                    "model_meta": model_meta
+                })
+                
+                # Construct session state
+                session_state = {
+                    "messages": session_data.get("messages", []),
+                    "memory": memories,
+                }
+                
+                result = {
+                    "session_id": row["session_id"],
+                    "metadata": metadata,
+                    "session_state": session_state
+                }
+                
+                self.logger.info(f"Session {session_id} loaded successfully")
+                self.logger.debug(f"Loaded data: {json.dumps(result)[:200]}...")
+                return result
+        except Exception as e:
+            self.logger.error(f"Failed to load session {session_id}: {str(e)}")
+            raise
+
+    async def close(self):
+        if self._pool:
+            self.logger.info("Closing PostgreSQL connection pool")
+            try:
+                await self._pool.close()
+                self._pool = None
+                self.logger.debug("PostgreSQL connection pool closed")
+            except Exception as e:
+                self.logger.error(f"Error closing PostgreSQL connection: {str(e)}") 
+

storage/redis_storage.py

@@ -0,0 +1,48 @@
+import json
+from typing import Dict, Any, Optional
+import redis.asyncio as aioredis
+from tinyagent.storage import Storage
+
+class RedisStorage(Storage):
+    """
+    Persist TinyAgent sessions in Redis.  Optionally expire them after `ttl` seconds.
+    """
+
+    def __init__(self, url: str = "redis://localhost", ttl: Optional[int] = None):
+        """
+        :param url: Redis connection URL, e.g. "redis://localhost:6379/0"
+        :param ttl: time‐to‐live in seconds (None ⇒ no expiry)
+        """
+        self.url = url
+        self.ttl = ttl
+        self._client: Optional[aioredis.Redis] = None
+
+    async def _connect(self):
+        if not self._client:
+            # from_url returns an asyncio‐enabled Redis client
+            self._client = aioredis.from_url(self.url)
+
+    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
+        await self._connect()
+        payload = json.dumps(data)
+        if self.ttl is not None:
+            # set with expiration
+            await self._client.set(f"{session_id}_{user_id}", payload, ex=self.ttl)
+        else:
+            # set without expiration
+            await self._client.set(f"{session_id}_{user_id}", payload)
+
+    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        await self._connect()
+        raw = await self._client.get(f"{session_id}_{user_id}")
+        if not raw:
+            return {}
+        # raw may be bytes or str
+        if isinstance(raw, bytes):
+            raw = raw.decode("utf-8")
+        return json.loads(raw)
+
+    async def close(self) -> None:
+        if self._client:
+            await self._client.close()
+            self._client = None

storage/sqlite_storage.py

@@ -0,0 +1,156 @@
+import aiosqlite
+import json
+import os
+from typing import Optional, Dict, Any
+from tinyagent.storage import Storage
+
+class SqliteStorage(Storage):
+    """
+    Persist TinyAgent sessions in a SQLite database with JSON state.
+    """
+
+    def __init__(self, db_path: str, table_name: str = "tny_agent_sessions"):
+        self._db_path = db_path
+        self._table = table_name
+        self._conn: Optional[aiosqlite.Connection] = None
+        
+        # Ensure the directory exists
+        os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)
+
+    async def _ensure_table(self):
+        """Create the sessions table if it doesn't exist."""
+        await self._conn.execute(f"""
+            CREATE TABLE IF NOT EXISTS {self._table} (
+                agent_id TEXT PRIMARY KEY,
+                session_id TEXT NOT NULL,
+                user_id TEXT,
+                memories TEXT,
+                metadata TEXT,
+                session_data TEXT,
+                model_meta TEXT,
+                created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
+            );
+        """)
+        
+        # Create indexes
+        await self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_session_id ON {self._table} (session_id);")
+        await self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_user_id ON {self._table} (user_id);")
+        await self._conn.commit()
+
+    async def _connect(self):
+        if not self._conn:
+            self._conn = await aiosqlite.connect(self._db_path)
+            self._conn.row_factory = aiosqlite.Row
+            await self._ensure_table()
+
+    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None):
+        await self._connect()
+        print(f"Saving session {session_id} for user {user_id} to sqlite {data}")
+        # Extract data following the TinyAgent schema
+        metadata = data.get("metadata", {}) or {}
+        session_state = data.get("session_state", {}) or {}
+        
+        # Use session_id as agent_id if not provided
+        agent_id = metadata.get("agent_id", session_id)
+        
+        # Extract specific components
+        memories = session_state.get("memory", {})
+        session_data = {"messages": session_state.get("messages", [])}
+        model_meta = metadata.get("model_meta", {})
+        
+        # Convert dictionaries to JSON strings
+        memories_json = json.dumps(memories)
+        metadata_json = json.dumps(metadata)
+        session_data_json = json.dumps(session_data)
+        model_meta_json = json.dumps(model_meta)
+        
+        # Check if record exists
+        cursor = await self._conn.execute(
+            f"SELECT 1 FROM {self._table} WHERE agent_id = ?", 
+            (agent_id,)
+        )
+        exists = await cursor.fetchone() is not None
+        
+        if exists:
+            # Update existing record
+            await self._conn.execute(f"""
+                UPDATE {self._table} SET
+                    session_id = ?,
+                    user_id = ?,
+                    memories = ?,
+                    metadata = ?,
+                    session_data = ?,
+                    model_meta = ?,
+                    updated_at = CURRENT_TIMESTAMP
+                WHERE agent_id = ?
+            """, (session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json, agent_id))
+        else:
+            # Insert new record
+            await self._conn.execute(f"""
+                INSERT INTO {self._table} 
+                (agent_id, session_id, user_id, memories, metadata, session_data, model_meta)
+                VALUES (?, ?, ?, ?, ?, ?, ?)
+            """, (agent_id, session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json))
+        
+        await self._conn.commit()
+
+    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
+        await self._connect()
+        
+        # Build query
+        query = f"""
+            SELECT agent_id, session_id, user_id, memories, metadata, session_data, model_meta
+            FROM {self._table}
+            WHERE session_id = ?
+        """
+        params = [session_id]
+        
+        # Add user_id filter if provided
+        if user_id:
+            query += " AND user_id = ?"
+            params.append(user_id)
+        
+        # Execute query
+        cursor = await self._conn.execute(query, params)
+        row = await cursor.fetchone()
+        
+        if not row:
+            return {}
+        
+        # Parse JSON strings
+        memories = json.loads(row["memories"]) if row["memories"] else {}
+        metadata = json.loads(row["metadata"]) if row["metadata"] else {}
+        session_data = json.loads(row["session_data"]) if row["session_data"] else {}
+        model_meta = json.loads(row["model_meta"]) if row["model_meta"] else {}
+        
+        # Update metadata with additional fields
+        metadata.update({
+            "agent_id": row["agent_id"],
+            "user_id": row["user_id"],
+            "model_meta": model_meta
+        })
+        
+        # Construct session state
+        session_state = {
+            "messages": session_data.get("messages", []),
+            "memory": memories,
+        }
+        
+        return {
+            "session_id": row["session_id"],
+            "metadata": metadata,
+            "session_state": session_state
+        }
+
+    async def close(self):
+        if self._conn:
+            await self._conn.close()
+            self._conn = None
+
+    async def __aenter__(self):
+        await self._connect()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.close()