tinyagent-py 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

tinyagent/storage/__init__.py

@@ -1,7 +1,5 @@
 from .base import Storage
-from .json_file_storage import JsonFileStorage
+
 # Import your concrete backends here if you like:
 # from .postgres_storage import PostgresStorage
-# from .json_file_storage import JsonFileStorage
-
-__all__ = ["Storage", "JsonFileStorage"]
+# from .json_file_storage import JsonFileStorage

tinyagent/storage/sqlite_storage.py

@@ -1,6 +1,7 @@
 import aiosqlite
 import json
 import os
+import logging
 from typing import Optional, Dict, Any
 from tinyagent.storage import Storage
 
@@ -13,6 +14,7 @@ class SqliteStorage(Storage):
         self._db_path = db_path
         self._table = table_name
         self._conn: Optional[aiosqlite.Connection] = None
+        self.logger = logging.getLogger(__name__)
         
         # Ensure the directory exists
         os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)
@@ -46,7 +48,8 @@ class SqliteStorage(Storage):
 
     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}")
+        self.logger.info(f"[sqlite] Saving session {session_id} for user {user_id}")
+        self.logger.debug(f"[sqlite] Save data: {json.dumps(data)[:200]}...")
         # Extract data following the TinyAgent schema
         metadata = data.get("metadata", {}) or {}
         session_state = data.get("session_state", {}) or {}

tinyagent/tiny_agent.py

@@ -186,8 +186,7 @@ class TinyAgent:
         self.temperature = temperature
         if model in ["o1", "o1-preview","o3","o4-mini"]:
             self.temperature = 1
-        if api_key:
-            litellm.api_key = api_key
+
         
         self.model_kwargs = model_kwargs
         self.encoder = tiktoken.get_encoding("o200k_base")
@@ -560,6 +559,7 @@ class TinyAgent:
             return str(result)
         except Exception as e:
             self.logger.error(f"Error executing custom tool {tool_name}: {str(e)}")
+            self.logger.error(f"Error: {traceback.format_exc()}")
             return f"Error executing tool {tool_name}: {str(e)}"
     
     async def run(self, user_input: str, max_turns: int = 10) -> str:
@@ -604,6 +604,7 @@ class TinyAgent:
                 
                 response = await litellm.acompletion(
                     model=self.model,
+                    api_key=self.api_key,
                     messages=self.messages,
                     tools=all_tools,
                     tool_choice="auto",
@@ -746,13 +747,58 @@ class TinyAgent:
 
     
     async def close(self):
-        """Clean up *all* MCP clients."""
+        """
+        Clean up all resources used by the agent including MCP clients and storage.
+        
+        This method should be called when the agent is no longer needed to ensure
+        proper resource cleanup, especially in web frameworks like FastAPI.
+        """
+        cleanup_errors = []
+        
+        # 1. First save any pending state if storage is configured
+        if self.storage:
+            try:
+                self.logger.debug(f"Saving final state before closing (session={self.session_id})")
+                await self.save_agent()
+            except Exception as e:
+                error_msg = f"Error saving final state: {str(e)}"
+                self.logger.error(error_msg)
+                cleanup_errors.append(error_msg)
+        
+        # 2. Close all MCP clients
         for client in self.mcp_clients:
             try:
+                self.logger.debug(f"Closing MCP client: {client}")
                 await client.close()
-            except RuntimeError as e:
-                self.logger.error(f"Error closing MCP client: {str(e)}")
-                # Continue closing other clients even if one fails
+            except Exception as e:
+                error_msg = f"Error closing MCP client: {str(e)}"
+                self.logger.error(error_msg)
+                cleanup_errors.append(error_msg)
+        
+        # 3. Close storage connection if available
+        if self.storage:
+            try:
+                self.logger.debug("Closing storage connection")
+                await self.storage.close()
+            except Exception as e:
+                error_msg = f"Error closing storage: {str(e)}"
+                self.logger.error(error_msg)
+                cleanup_errors.append(error_msg)
+        
+        # 4. Run any cleanup callbacks
+        try:
+            await self._run_callbacks("agent_cleanup")
+        except Exception as e:
+            error_msg = f"Error in cleanup callbacks: {str(e)}"
+            self.logger.error(error_msg)
+            cleanup_errors.append(error_msg)
+        
+        # Log summary of cleanup
+        if cleanup_errors:
+            self.logger.warning(f"TinyAgent cleanup completed with {len(cleanup_errors)} errors")
+        else:
+            self.logger.info(f"TinyAgent cleanup completed successfully (session={self.session_id})")
+
     def clear_conversation(self):
         """
         Clear the conversation history, preserving only the initial system prompt.

{tinyagent_py-0.0.5.dist-info → tinyagent_py-0.0.7.dist-info}/METADATA

@@ -1,7 +1,7 @@
 Metadata-Version: 2.4
 Name: tinyagent-py
-Version: 0.0.5
-Summary: Tiny Agent with MCP Client
+Version: 0.0.7
+Summary: Tiny Agent with MCP Client and Extendable Hooks, Tiny but powerful
 Author-email: Mahdi Golchin <golchin@askdev.ai>
 Project-URL: Homepage, https://github.com/askbudi/tinyagent
 Project-URL: Bug Tracker, https://github.com/askbudi/tinyagent/issues

tinyagent_py-0.0.7.dist-info/RECORD

@@ -0,0 +1,20 @@
+tinyagent/__init__.py,sha256=GrD21npMQGzl9ZYKYTP8VxHLzCfJCvA0oTKQZTkmnCw,117
+tinyagent/mcp_client.py,sha256=9dmLtJ8CTwKWKTH6K9z8CaCQuaicOH9ifAuNyX7Kdo0,6030
+tinyagent/memory_manager.py,sha256=tAaZZdxBJ235wJIyr04n3f2Damok4s2UXunTtur_p-4,44916
+tinyagent/tiny_agent.py,sha256=_qHQc2yroSI1Ihn5Ex3FH_ml7aaSNC9-yR540C16Cxs,41125
+tinyagent/hooks/__init__.py,sha256=UztCHjoqF5JyDolbWwkBsBZkWguDQg23l2GD_zMHt-s,178
+tinyagent/hooks/gradio_callback.py,sha256=jGsZlObAd6I5lN9cE53dDL_LfiB8I0tBsicuHwwmL-M,44833
+tinyagent/hooks/logging_manager.py,sha256=UpdmpQ7HRPyer-jrmQSXcBwi409tV9LnGvXSHjTcYTI,7935
+tinyagent/hooks/rich_code_ui_callback.py,sha256=PLcu5MOSoP4oZR3BtvcV9DquxcIT_d0WzSlkvaDcGOk,19820
+tinyagent/hooks/rich_ui_callback.py,sha256=5iCNOiJmhc1lOL7ZjaOt5Sk3rompko4zu_pAxfTVgJQ,22897
+tinyagent/storage/__init__.py,sha256=7qwfdD4smCl891xaRuiReSUgfOJFy7jJZsN0ul1iQdY,173
+tinyagent/storage/base.py,sha256=GGAMvOoslmm1INLFG_jtwOkRk2Qg39QXx-1LnN7fxDI,1474
+tinyagent/storage/json_file_storage.py,sha256=SYD8lvTHu2-FEHm1tZmsrcgEOirBrlUsUM186X-UPgI,1114
+tinyagent/storage/postgres_storage.py,sha256=IGwan8UXHNnTZFK1F8x4kvMDex3GAAGWUg9ePx_5IF4,9018
+tinyagent/storage/redis_storage.py,sha256=hu3y7wHi49HkpiR-AW7cWVQuTVOUk1WaB8TEPGUKVJ8,1742
+tinyagent/storage/sqlite_storage.py,sha256=ZyOYe0d_oHO1wOIT8FxKIbc67tP_0e_8FnM2Zq8Pwj8,5915
+tinyagent_py-0.0.7.dist-info/licenses/LICENSE,sha256=YIogcVQnknaaE4K-oaQylFWo8JGRBWnwmGb3fWB_Pww,1064
+tinyagent_py-0.0.7.dist-info/METADATA,sha256=C1qNkWSjwmrsaYVXkEyMrvsg-57nCZsIMmsUaZ4OOpo,9094
+tinyagent_py-0.0.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tinyagent_py-0.0.7.dist-info/top_level.txt,sha256=Ny8aJNchZpc2Vvhp3306L5vjceJakvFxBk-UjjVeA_I,10
+tinyagent_py-0.0.7.dist-info/RECORD,,

{tinyagent_py-0.0.5.dist-info → tinyagent_py-0.0.7.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.8.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

tinyagent/hooks/agno_storage_hook.py

@@ -1,128 +0,0 @@
-try:
-    import agno
-    from agno.storage.postgres import PostgresStorage
-    from agno.storage.sqlite import SqliteStorage
-    from agno.storage.session.agent import AgentSession
-except ImportError as e:
-    raise ImportError("agno is not installed. Please install it with `pip install agno`.", e)
-
-import asyncio
-from typing import Optional
-from agno.storage.postgres import PostgresStorage
-from agno.storage.sqlite import SqliteStorage
-from agno.storage.session.agent import AgentSession
-
-class PostgresStorageHook:
-    def __init__(
-        self,
-        table_name: str,
-        db_url: Optional[str] = None,
-        db_engine=None,
-        schema: Optional[str] = "ai",
-        schema_version: int = 1,
-        auto_upgrade_schema: bool = True,
-        mode: str = "agent",
-    ):
-        self.storage = PostgresStorage(
-            table_name=table_name,
-            db_url=db_url,
-            db_engine=db_engine,
-            schema=schema,
-            schema_version=schema_version,
-            auto_upgrade_schema=auto_upgrade_schema,
-            mode=mode,
-        )
-
-    async def __call__(self, event_name: str, agent, **kwargs):
-        if event_name == "agent_start":
-            # Load session from storage
-            session_id = getattr(agent, "session_id", None)
-            user_id = getattr(agent, "user_id", None)
-            if session_id:
-                session = self.storage.read(session_id=session_id, user_id=user_id)
-                if session:
-                    # Populate agent state from session
-                    agent.messages = session.session_data.get("messages", [])
-                    agent.memory = session.memory
-                    agent.metadata = session.extra_data
-                    # You may need to adapt this depending on tinyagent's state structure
-
-        elif event_name in ("llm_end", "agent_end"):
-            # Save session to storage
-            session_id = getattr(agent, "session_id", None)
-            user_id = getattr(agent, "user_id", None)
-            if session_id:
-                # Create AgentSession from agent state
-                session_data = {
-                    "messages": getattr(agent, "messages", []),
-                }
-                session = AgentSession(
-                    session_id=session_id,
-                    user_id=user_id,
-                    memory=getattr(agent, "memory", {}),
-                    session_data=session_data,
-                    extra_data=getattr(agent, "metadata", {}),
-                    agent_id=getattr(agent, "agent_id", None),
-                    team_session_id=None,
-                    agent_data=None,
-                )
-                await asyncio.to_thread(self.storage.upsert, session)
-
-class SqliteStorageHook:
-    def __init__(
-        self,
-        table_name: str,
-        db_url: Optional[str] = None,
-        db_file: Optional[str] = None,
-        db_engine=None,
-        schema_version: int = 1,
-        auto_upgrade_schema: bool = True,
-        mode: str = "agent",
-    ):
-        self.storage = SqliteStorage(
-            table_name=table_name,
-            db_url=db_url,
-            db_file=db_file,
-            db_engine=db_engine,
-            schema_version=schema_version,
-            auto_upgrade_schema=auto_upgrade_schema,
-            mode=mode,
-        )
-
-    async def __call__(self, event_name: str, agent, **kwargs):
-        if event_name == "agent_start":
-            # Load session from storage
-            session_id = getattr(agent, "session_id", None)
-            user_id = getattr(agent, "user_id", None)
-            print("Session ID",session_id)
-            print("User ID",user_id)
-            if session_id:
-                session = self.storage.read(session_id=session_id, user_id=user_id)
-                print(f"Session: {session}")
-                if session:
-                    # Populate agent state from session
-                    agent.messages = session.memory.get("messages", [])
-                    agent.memory = session.memory
-                    agent.metadata = session.extra_data
-
-        elif event_name in ("llm_end", "agent_end"):
-            # Save session to storage
-            print("Agent metadata",getattr(agent, "metadata", {}))
-            session_id = getattr(agent, "session_id", None)
-            user_id = getattr(agent, "user_id", None)
-            if session_id:
-                session_data = {
-                    "messages": getattr(agent, "messages", []),
-                }
-                session = AgentSession(
-                    session_id=session_id,
-                    user_id=user_id,
-                    memory=getattr(agent, "memory", {}),
-                    session_data=session_data,
-                    extra_data=getattr(agent, "metadata", {}),
-                    agent_id=getattr(agent, "agent_id", None),
-                    team_session_id=None,
-                    agent_data=None,
-                )
-                await asyncio.to_thread(self.storage.upsert, session)
-

tinyagent/storage/agno_storage.py

@@ -1,114 +0,0 @@
-import asyncio
-from typing import Dict, Any, Optional
-try:
-    from tinyagent.storage.base import Storage
-    from agno.storage.postgres import PostgresStorage as AgnoPG
-    from agno.storage.sqlite import SqliteStorage as AgnoSL
-    from agno.storage.session.agent import AgentSession
-except ImportError as e:
-    raise ImportError("agno is not installed. Please install it with `pip install agno`.", e)
-
-def _remap_agno_to_tiny(ag: Dict[str, Any]) -> Dict[str, Any]:
-    """Map a full Agno to_dict() into TinyAgent’s to_dict() shape."""
-    sess_id = ag.get("session_id", "")
-
-    # meta = everything except the session_state fields
-    metadata = {
-        k: v
-        for k, v in ag.items()
-        if k not in ("session_id", "session_data", "memory", "runs")
-    }
-
-    # Safe‐guard: use {} if any of these are None
-    _session_data = ag.get("session_data") or {}
-    _memory       = ag.get("memory") or {}
-   
-
-    session_state = {
-        "messages": _memory.get("messages", []),
-        "memory":   _memory,
-        **_session_data,
-    }
-
-    return {
-        "session_id": sess_id,
-        "metadata": metadata,
-        "session_state": session_state,
-    }
-
-def _remap_tiny_to_agno(tiny: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Given TinyAgent.to_dict() output:
-      { "session_id": str,
-        "metadata": {...},
-        "session_state": { "messages": [...], "memory": {...}, "runs": [...] }
-      }
-    produce a full AgnoSession.to_dict() shape:
-      { "session_id":..., "user_id":..., "memory":..., "runs":...,
-        "session_data": {"messages": [...]},
-        "extra_data":...,
-         ... (other Agno fields remain None/default) }
-    """
-    session_id = tiny["session_id"]
-    meta       = tiny.get("metadata", {}) or {}
-    state      = tiny.get("session_state", {}) or {}
-
-    return {
-        "session_id":      session_id,
-        "user_id":         meta.get("user_id"),
-        "team_session_id": meta.get("team_session_id"),
-        "memory":          state.get("memory", {}),
-        "runs":            state.get("runs", []),
-        "session_data":    {"messages": state.get("messages", [])},
-        "extra_data":      meta.get("extra_data"),
-        # created_at/updated_at/agent_id/agent_data will default in AgnoSession
-    }
-
-
-class AgnoPostgresStorage(Storage):
-    def __init__(self, table_name: str, db_url: str, schema: str = "ai", mode: str = "agent"):
-        super().__init__()
-        self.backend = AgnoPG(table_name=table_name, schema=schema, db_url=db_url, mode=mode)
-        self.backend.create()
-
-    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
-        # Pack TinyAgent dict into AgnoSession record
-        agno_dict = _remap_tiny_to_agno(data)
-        session_obj = AgentSession.from_dict(agno_dict)
-        loop = asyncio.get_event_loop()
-        await loop.run_in_executor(None, self.backend.upsert, session_obj)
-
-    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
-        loop = asyncio.get_event_loop()
-        agno_obj = await loop.run_in_executor(None, self.backend.read, session_id, user_id)
-        if not agno_obj:
-            return {}
-        ag = agno_obj.to_dict()
-        return _remap_agno_to_tiny(ag)
-
-    async def close(self) -> None:
-        pass
-
-
-class AgnoSqliteStorage(Storage):
-    def __init__(self, table_name: str, db_url: str, mode: str = "agent"):
-        super().__init__()
-        self.backend = AgnoSL(table_name=table_name, db_url=db_url, mode=mode)
-        self.backend.create()
-
-    async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
-        agno_dict = _remap_tiny_to_agno(data)
-        session_obj = AgentSession.from_dict(agno_dict)
-        loop = asyncio.get_event_loop()
-        await loop.run_in_executor(None, self.backend.upsert, session_obj)
-
-    async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
-        loop = asyncio.get_event_loop()
-        agno_obj = await loop.run_in_executor(None, self.backend.read, session_id, user_id)
-        if not agno_obj:
-            return {}
-        ag = agno_obj.to_dict()
-        return _remap_agno_to_tiny(ag)
-
-    async def close(self) -> None:
-        pass

tinyagent_py-0.0.5.dist-info/RECORD

@@ -1,20 +0,0 @@
-tinyagent/__init__.py,sha256=GrD21npMQGzl9ZYKYTP8VxHLzCfJCvA0oTKQZTkmnCw,117
-tinyagent/mcp_client.py,sha256=ZyhIWFQvOveNVhMKkMkUMYdbqRSuqjK_m-ApP2llFVE,6215
-tinyagent/tiny_agent.py,sha256=89OLYj95VUW7uoHpkWXZjmxsQISUMjHffUmDT-3-Nd4,39250
-tinyagent/hooks/__init__.py,sha256=UztCHjoqF5JyDolbWwkBsBZkWguDQg23l2GD_zMHt-s,178
-tinyagent/hooks/agno_storage_hook.py,sha256=5qvvjmtraanPa-A46Zstrqq3s1e-sC7Ly0o3zifuw_4,5003
-tinyagent/hooks/gradio_callback.py,sha256=jGsZlObAd6I5lN9cE53dDL_LfiB8I0tBsicuHwwmL-M,44833
-tinyagent/hooks/logging_manager.py,sha256=UpdmpQ7HRPyer-jrmQSXcBwi409tV9LnGvXSHjTcYTI,7935
-tinyagent/hooks/rich_ui_callback.py,sha256=5iCNOiJmhc1lOL7ZjaOt5Sk3rompko4zu_pAxfTVgJQ,22897
-tinyagent/storage/__init__.py,sha256=NebvYxwEGJtvPnRO9dGa-bgOwA7cPkLjFHnMWDxMg5I,261
-tinyagent/storage/agno_storage.py,sha256=ol4qwdH-9jYjBjDvsYkHh7I-vu8uHArPtQylUpoEaCc,4322
-tinyagent/storage/base.py,sha256=GGAMvOoslmm1INLFG_jtwOkRk2Qg39QXx-1LnN7fxDI,1474
-tinyagent/storage/json_file_storage.py,sha256=SYD8lvTHu2-FEHm1tZmsrcgEOirBrlUsUM186X-UPgI,1114
-tinyagent/storage/postgres_storage.py,sha256=IGwan8UXHNnTZFK1F8x4kvMDex3GAAGWUg9ePx_5IF4,9018
-tinyagent/storage/redis_storage.py,sha256=hu3y7wHi49HkpiR-AW7cWVQuTVOUk1WaB8TEPGUKVJ8,1742
-tinyagent/storage/sqlite_storage.py,sha256=7lk1XZpr2t4s2bjVr9-AqrI74w4hwkuK3taWtyJZhBc,5769
-tinyagent_py-0.0.5.dist-info/licenses/LICENSE,sha256=YIogcVQnknaaE4K-oaQylFWo8JGRBWnwmGb3fWB_Pww,1064
-tinyagent_py-0.0.5.dist-info/METADATA,sha256=EW9xVFf0vYbe1DQYxQAlNgOXUzuP9g2pqc__yx4V8NQ,9054
-tinyagent_py-0.0.5.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
-tinyagent_py-0.0.5.dist-info/top_level.txt,sha256=Ny8aJNchZpc2Vvhp3306L5vjceJakvFxBk-UjjVeA_I,10
-tinyagent_py-0.0.5.dist-info/RECORD,,