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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
storage/base.py ADDED
@@ -0,0 +1,49 @@
1
+ from abc import ABC, abstractmethod
2
+ from typing import Dict, Any, TYPE_CHECKING, Optional
3
+
4
+ if TYPE_CHECKING:
5
+ from tinyagent.tiny_agent import TinyAgent
6
+
7
+ class Storage(ABC):
8
+ """
9
+ Abstract base class for TinyAgent session storage.
10
+ """
11
+
12
+ @abstractmethod
13
+ async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
14
+ """
15
+ Persist the given agent state under `session_id`.
16
+ """
17
+ ...
18
+
19
+ @abstractmethod
20
+ async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
21
+ """
22
+ Retrieve the agent state for `session_id`, or return {} if not found.
23
+ """
24
+ ...
25
+
26
+ @abstractmethod
27
+ async def close(self) -> None:
28
+ """
29
+ Clean up any resources (DB connections, file handles, etc.).
30
+ """
31
+ ...
32
+
33
+ def attach(self, agent: "TinyAgent") -> None:
34
+ """
35
+ Hook this storage to a TinyAgent so that on every `llm_end`
36
+ it will auto‐persist the agent's state.
37
+
38
+ Usage:
39
+ storage.attach(agent)
40
+ or in TinyAgent.__init__:
41
+ if storage: storage.attach(self)
42
+ """
43
+ async def _auto_save(event_name: str, agent: "TinyAgent", **kwargs):
44
+ if event_name != "llm_end":
45
+ return
46
+ state = agent.to_dict()
47
+ await self.save_session(agent.session_id, state)
48
+
49
+ agent.callbacks.append(_auto_save)
@@ -0,0 +1,30 @@
1
+ import json
2
+ import asyncio
3
+ from pathlib import Path
4
+ from typing import Dict, Any, Union, Optional
5
+ from tinyagent.storage import Storage
6
+
7
+ class JsonFileStorage(Storage):
8
+ """
9
+ Persist TinyAgent sessions as individual JSON files.
10
+ """
11
+
12
+ def __init__(self, folder: Union[str, Path]):
13
+ self.folder = Path(folder)
14
+ self.folder.mkdir(parents=True, exist_ok=True)
15
+
16
+ async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
17
+ path = self.folder / f"{session_id}_{user_id}.json"
18
+ # Write in a thread pool to avoid blocking the event loop
19
+ await asyncio.to_thread(path.write_text, json.dumps(data, indent=2), "utf-8")
20
+
21
+ async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
22
+ path = self.folder / f"{session_id}_{user_id}.json"
23
+ if not path.exists():
24
+ return {}
25
+ text = await asyncio.to_thread(path.read_text, "utf-8")
26
+ return json.loads(text)
27
+
28
+ async def close(self) -> None:
29
+ # Nothing to clean up for file storage
30
+ return
@@ -0,0 +1,201 @@
1
+ import asyncpg
2
+ import json
3
+ import logging
4
+ from typing import Optional, Dict, Any
5
+ from tinyagent.storage import Storage
6
+
7
+ class PostgresStorage(Storage):
8
+ """
9
+ Persist TinyAgent sessions in a Postgres table with JSONB state.
10
+ """
11
+
12
+ def __init__(self, db_url: str, table_name: str = "tny_agent_sessions"):
13
+ self._dsn = db_url
14
+ self._table = table_name
15
+ self._pool: Optional[asyncpg.pool.Pool] = None
16
+ self.logger = logging.getLogger(__name__)
17
+
18
+ async def _ensure_table(self):
19
+ """Create the sessions table if it doesn't exist."""
20
+ self.logger.debug(f"Ensuring table {self._table} exists")
21
+ try:
22
+ async with self._pool.acquire() as conn:
23
+ await conn.execute(f"""
24
+ CREATE TABLE IF NOT EXISTS {self._table} (
25
+ agent_id TEXT PRIMARY KEY,
26
+ session_id TEXT NOT NULL,
27
+ user_id TEXT,
28
+ memories JSONB,
29
+ metadata JSONB,
30
+ session_data JSONB,
31
+ model_meta JSONB,
32
+ created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
33
+ updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
34
+ );
35
+ CREATE INDEX IF NOT EXISTS idx_{self._table}_session_id ON {self._table} (session_id);
36
+ CREATE INDEX IF NOT EXISTS idx_{self._table}_user_id ON {self._table} (user_id);
37
+ """)
38
+ self.logger.info(f"Table {self._table} and indexes created/verified")
39
+ except Exception as e:
40
+ self.logger.error(f"Error creating table {self._table}: {str(e)}")
41
+ raise
42
+
43
+ async def _connect(self):
44
+ if not self._pool:
45
+ self.logger.debug(f"Connecting to PostgreSQL with DSN: {self._dsn[:10]}...")
46
+ try:
47
+ # Ensure statement_cache_size=0 to disable prepared statements for pgbouncer compatibility
48
+ self._pool = await asyncpg.create_pool(
49
+ dsn=self._dsn,
50
+ statement_cache_size=0,
51
+ min_size=1,
52
+ max_size=10
53
+ )
54
+ self.logger.info("PostgreSQL connection pool created")
55
+ await self._ensure_table()
56
+ except Exception as e:
57
+ self.logger.error(f"Failed to connect to PostgreSQL: {str(e)}")
58
+ raise
59
+
60
+ async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None):
61
+ self.logger.info(f"Saving session {session_id} for user {user_id}")
62
+ self.logger.debug(f"Save data: {json.dumps(data)[:200]}...")
63
+
64
+ try:
65
+ await self._connect()
66
+
67
+ # Extract data following the TinyAgent schema
68
+ metadata = data.get("metadata", {}) or {}
69
+ session_state = data.get("session_state", {}) or {}
70
+
71
+ # Use session_id as agent_id if not provided
72
+ agent_id = metadata.get("agent_id", session_id)
73
+ self.logger.debug(f"Using agent_id: {agent_id}")
74
+
75
+ # Extract specific components
76
+ memories = session_state.get("memory", {})
77
+ session_data = {"messages": session_state.get("messages", [])}
78
+ model_meta = metadata.get("model_meta", {})
79
+
80
+ # Convert Python dictionaries to JSON strings for PostgreSQL
81
+ self.logger.debug("Converting Python dictionaries to JSON")
82
+ try:
83
+ memories_json = json.dumps(memories)
84
+ metadata_json = json.dumps(metadata)
85
+ session_data_json = json.dumps(session_data)
86
+ model_meta_json = json.dumps(model_meta)
87
+ except Exception as e:
88
+ self.logger.error(f"JSON serialization error: {str(e)}")
89
+ raise
90
+
91
+ self.logger.debug("Executing PostgreSQL INSERT/UPDATE")
92
+ async with self._pool.acquire() as conn:
93
+ try:
94
+ await conn.execute(f"""
95
+ INSERT INTO {self._table}
96
+ (agent_id, session_id, user_id, memories, metadata, session_data, model_meta, updated_at)
97
+ VALUES ($1, $2, $3, $4::jsonb, $5::jsonb, $6::jsonb, $7::jsonb, NOW())
98
+ ON CONFLICT (agent_id) DO UPDATE
99
+ SET session_id = EXCLUDED.session_id,
100
+ user_id = EXCLUDED.user_id,
101
+ memories = EXCLUDED.memories,
102
+ metadata = EXCLUDED.metadata,
103
+ session_data = EXCLUDED.session_data,
104
+ model_meta = EXCLUDED.model_meta,
105
+ updated_at = NOW();
106
+ """, agent_id, session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json)
107
+ self.logger.info(f"Session {session_id} saved successfully")
108
+ except Exception as e:
109
+ self.logger.error(f"Database error during save: {str(e)}")
110
+ raise
111
+ except Exception as e:
112
+ self.logger.error(f"Failed to save session {session_id}: {str(e)}")
113
+ raise
114
+
115
+ async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
116
+ self.logger.info(f"Loading session {session_id} for user {user_id}")
117
+
118
+ try:
119
+ await self._connect()
120
+
121
+ async with self._pool.acquire() as conn:
122
+ # First try to find by session_id
123
+ query = f"""
124
+ SELECT agent_id, session_id, user_id, memories, metadata, session_data, model_meta
125
+ FROM {self._table}
126
+ WHERE session_id = $1
127
+ """
128
+ params = [session_id]
129
+
130
+ # Add user_id filter if provided
131
+ if user_id:
132
+ query += " AND user_id = $2"
133
+ params.append(user_id)
134
+
135
+ self.logger.debug(f"Executing query: {query} with params: {params}")
136
+ row = await conn.fetchrow(query, *params)
137
+
138
+ if not row:
139
+ self.logger.warning(f"No session found for session_id={session_id}, user_id={user_id}")
140
+ return {}
141
+
142
+ self.logger.debug(f"Session found: {dict(row)}")
143
+
144
+ # Parse JSON from PostgreSQL
145
+ try:
146
+ # Check if values are already dictionaries or need parsing
147
+ memories = row["memories"]
148
+ if isinstance(memories, str):
149
+ memories = json.loads(memories)
150
+
151
+ metadata = row["metadata"]
152
+ if isinstance(metadata, str):
153
+ metadata = json.loads(metadata)
154
+
155
+ session_data = row["session_data"]
156
+ if isinstance(session_data, str):
157
+ session_data = json.loads(session_data)
158
+
159
+ model_meta = row["model_meta"]
160
+ if isinstance(model_meta, str):
161
+ model_meta = json.loads(model_meta)
162
+ except Exception as e:
163
+ self.logger.error(f"Error parsing JSON from database: {str(e)}")
164
+ raise
165
+
166
+ # Update metadata with additional fields
167
+ metadata.update({
168
+ "agent_id": row["agent_id"],
169
+ "user_id": row["user_id"],
170
+ "model_meta": model_meta
171
+ })
172
+
173
+ # Construct session state
174
+ session_state = {
175
+ "messages": session_data.get("messages", []),
176
+ "memory": memories,
177
+ }
178
+
179
+ result = {
180
+ "session_id": row["session_id"],
181
+ "metadata": metadata,
182
+ "session_state": session_state
183
+ }
184
+
185
+ self.logger.info(f"Session {session_id} loaded successfully")
186
+ self.logger.debug(f"Loaded data: {json.dumps(result)[:200]}...")
187
+ return result
188
+ except Exception as e:
189
+ self.logger.error(f"Failed to load session {session_id}: {str(e)}")
190
+ raise
191
+
192
+ async def close(self):
193
+ if self._pool:
194
+ self.logger.info("Closing PostgreSQL connection pool")
195
+ try:
196
+ await self._pool.close()
197
+ self._pool = None
198
+ self.logger.debug("PostgreSQL connection pool closed")
199
+ except Exception as e:
200
+ self.logger.error(f"Error closing PostgreSQL connection: {str(e)}")
201
+
@@ -0,0 +1,48 @@
1
+ import json
2
+ from typing import Dict, Any, Optional
3
+ import redis.asyncio as aioredis
4
+ from tinyagent.storage import Storage
5
+
6
+ class RedisStorage(Storage):
7
+ """
8
+ Persist TinyAgent sessions in Redis. Optionally expire them after `ttl` seconds.
9
+ """
10
+
11
+ def __init__(self, url: str = "redis://localhost", ttl: Optional[int] = None):
12
+ """
13
+ :param url: Redis connection URL, e.g. "redis://localhost:6379/0"
14
+ :param ttl: time‐to‐live in seconds (None ⇒ no expiry)
15
+ """
16
+ self.url = url
17
+ self.ttl = ttl
18
+ self._client: Optional[aioredis.Redis] = None
19
+
20
+ async def _connect(self):
21
+ if not self._client:
22
+ # from_url returns an asyncio‐enabled Redis client
23
+ self._client = aioredis.from_url(self.url)
24
+
25
+ async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None) -> None:
26
+ await self._connect()
27
+ payload = json.dumps(data)
28
+ if self.ttl is not None:
29
+ # set with expiration
30
+ await self._client.set(f"{session_id}_{user_id}", payload, ex=self.ttl)
31
+ else:
32
+ # set without expiration
33
+ await self._client.set(f"{session_id}_{user_id}", payload)
34
+
35
+ async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
36
+ await self._connect()
37
+ raw = await self._client.get(f"{session_id}_{user_id}")
38
+ if not raw:
39
+ return {}
40
+ # raw may be bytes or str
41
+ if isinstance(raw, bytes):
42
+ raw = raw.decode("utf-8")
43
+ return json.loads(raw)
44
+
45
+ async def close(self) -> None:
46
+ if self._client:
47
+ await self._client.close()
48
+ self._client = None
@@ -0,0 +1,156 @@
1
+ import aiosqlite
2
+ import json
3
+ import os
4
+ from typing import Optional, Dict, Any
5
+ from tinyagent.storage import Storage
6
+
7
+ class SqliteStorage(Storage):
8
+ """
9
+ Persist TinyAgent sessions in a SQLite database with JSON state.
10
+ """
11
+
12
+ def __init__(self, db_path: str, table_name: str = "tny_agent_sessions"):
13
+ self._db_path = db_path
14
+ self._table = table_name
15
+ self._conn: Optional[aiosqlite.Connection] = None
16
+
17
+ # Ensure the directory exists
18
+ os.makedirs(os.path.dirname(os.path.abspath(db_path)), exist_ok=True)
19
+
20
+ async def _ensure_table(self):
21
+ """Create the sessions table if it doesn't exist."""
22
+ await self._conn.execute(f"""
23
+ CREATE TABLE IF NOT EXISTS {self._table} (
24
+ agent_id TEXT PRIMARY KEY,
25
+ session_id TEXT NOT NULL,
26
+ user_id TEXT,
27
+ memories TEXT,
28
+ metadata TEXT,
29
+ session_data TEXT,
30
+ model_meta TEXT,
31
+ created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
32
+ updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP
33
+ );
34
+ """)
35
+
36
+ # Create indexes
37
+ await self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_session_id ON {self._table} (session_id);")
38
+ await self._conn.execute(f"CREATE INDEX IF NOT EXISTS idx_{self._table}_user_id ON {self._table} (user_id);")
39
+ await self._conn.commit()
40
+
41
+ async def _connect(self):
42
+ if not self._conn:
43
+ self._conn = await aiosqlite.connect(self._db_path)
44
+ self._conn.row_factory = aiosqlite.Row
45
+ await self._ensure_table()
46
+
47
+ async def save_session(self, session_id: str, data: Dict[str, Any], user_id: Optional[str] = None):
48
+ await self._connect()
49
+ print(f"Saving session {session_id} for user {user_id} to sqlite {data}")
50
+ # Extract data following the TinyAgent schema
51
+ metadata = data.get("metadata", {}) or {}
52
+ session_state = data.get("session_state", {}) or {}
53
+
54
+ # Use session_id as agent_id if not provided
55
+ agent_id = metadata.get("agent_id", session_id)
56
+
57
+ # Extract specific components
58
+ memories = session_state.get("memory", {})
59
+ session_data = {"messages": session_state.get("messages", [])}
60
+ model_meta = metadata.get("model_meta", {})
61
+
62
+ # Convert dictionaries to JSON strings
63
+ memories_json = json.dumps(memories)
64
+ metadata_json = json.dumps(metadata)
65
+ session_data_json = json.dumps(session_data)
66
+ model_meta_json = json.dumps(model_meta)
67
+
68
+ # Check if record exists
69
+ cursor = await self._conn.execute(
70
+ f"SELECT 1 FROM {self._table} WHERE agent_id = ?",
71
+ (agent_id,)
72
+ )
73
+ exists = await cursor.fetchone() is not None
74
+
75
+ if exists:
76
+ # Update existing record
77
+ await self._conn.execute(f"""
78
+ UPDATE {self._table} SET
79
+ session_id = ?,
80
+ user_id = ?,
81
+ memories = ?,
82
+ metadata = ?,
83
+ session_data = ?,
84
+ model_meta = ?,
85
+ updated_at = CURRENT_TIMESTAMP
86
+ WHERE agent_id = ?
87
+ """, (session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json, agent_id))
88
+ else:
89
+ # Insert new record
90
+ await self._conn.execute(f"""
91
+ INSERT INTO {self._table}
92
+ (agent_id, session_id, user_id, memories, metadata, session_data, model_meta)
93
+ VALUES (?, ?, ?, ?, ?, ?, ?)
94
+ """, (agent_id, session_id, user_id, memories_json, metadata_json, session_data_json, model_meta_json))
95
+
96
+ await self._conn.commit()
97
+
98
+ async def load_session(self, session_id: str, user_id: Optional[str] = None) -> Dict[str, Any]:
99
+ await self._connect()
100
+
101
+ # Build query
102
+ query = f"""
103
+ SELECT agent_id, session_id, user_id, memories, metadata, session_data, model_meta
104
+ FROM {self._table}
105
+ WHERE session_id = ?
106
+ """
107
+ params = [session_id]
108
+
109
+ # Add user_id filter if provided
110
+ if user_id:
111
+ query += " AND user_id = ?"
112
+ params.append(user_id)
113
+
114
+ # Execute query
115
+ cursor = await self._conn.execute(query, params)
116
+ row = await cursor.fetchone()
117
+
118
+ if not row:
119
+ return {}
120
+
121
+ # Parse JSON strings
122
+ memories = json.loads(row["memories"]) if row["memories"] else {}
123
+ metadata = json.loads(row["metadata"]) if row["metadata"] else {}
124
+ session_data = json.loads(row["session_data"]) if row["session_data"] else {}
125
+ model_meta = json.loads(row["model_meta"]) if row["model_meta"] else {}
126
+
127
+ # Update metadata with additional fields
128
+ metadata.update({
129
+ "agent_id": row["agent_id"],
130
+ "user_id": row["user_id"],
131
+ "model_meta": model_meta
132
+ })
133
+
134
+ # Construct session state
135
+ session_state = {
136
+ "messages": session_data.get("messages", []),
137
+ "memory": memories,
138
+ }
139
+
140
+ return {
141
+ "session_id": row["session_id"],
142
+ "metadata": metadata,
143
+ "session_state": session_state
144
+ }
145
+
146
+ async def close(self):
147
+ if self._conn:
148
+ await self._conn.close()
149
+ self._conn = None
150
+
151
+ async def __aenter__(self):
152
+ await self._connect()
153
+ return self
154
+
155
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
156
+ await self.close()
@@ -0,0 +1,207 @@
1
+ Metadata-Version: 2.4
2
+ Name: tinyagent-py
3
+ Version: 0.0.3
4
+ Summary: Tiny Agent with MCP Client
5
+ Author-email: Mahdi Golchin <golchin@askdev.ai>
6
+ Project-URL: Homepage, https://github.com/askbudi/tinyagent
7
+ Project-URL: Bug Tracker, https://github.com/askbudi/tinyagent/issues
8
+ Project-URL: Chat, https://askdev.ai/github/askbudi/tinyagent
9
+ Requires-Python: >=3.8
10
+ Description-Content-Type: text/markdown
11
+ License-File: LICENSE
12
+ Requires-Dist: mcp
13
+ Requires-Dist: litellm
14
+ Requires-Dist: openai
15
+ Provides-Extra: dev
16
+ Requires-Dist: pytest; extra == "dev"
17
+ Dynamic: license-file
18
+
19
+ # tinyagent
20
+ Tiny Agent: 100 lines Agent with MCP
21
+ ![TinyAgent Logo](https://raw.githubusercontent.com/askbudi/tinyagent/main/public/logo.png)
22
+
23
+
24
+
25
+ Inspired by:
26
+ - [Tiny Agents blog post](https://huggingface.co/blog/tiny-agents)
27
+ - [12-factor-agents repository](https://github.com/humanlayer/12-factor-agents)
28
+ - Created by chatting to the source code of JS Tiny Agent using [AskDev.ai](https://askdev.ai/search)
29
+
30
+ ## Quick Links
31
+ - [Build your own Tiny Agent](https://askdev.ai/github/askbudi/tinyagent)
32
+
33
+ ## Overview
34
+ This is a tiny agent that uses MCP and LiteLLM to interact with a model. You have full control over the agent, you can add any tools you like from MCP and extend the agent using its event system.
35
+
36
+ ## Installation
37
+
38
+ ### Using pip
39
+ ```bash
40
+ pip install tinyagent-py
41
+ ```
42
+
43
+ ### Using uv
44
+ ```bash
45
+ uv pip install tinyagent-py
46
+ ```
47
+
48
+ ## Usage
49
+
50
+ ```python
51
+ from tinyagent import TinyAgent
52
+ from textwrap import dedent
53
+ import asyncio
54
+ import os
55
+
56
+ async def test_agent(task, model="o4-mini", api_key=None):
57
+ # Initialize the agent with model and API key
58
+ agent = TinyAgent(
59
+ model=model, # Or any model supported by LiteLLM
60
+ api_key=os.environ.get("OPENAI_API_KEY") if not api_key else api_key # Set your API key as an env variable
61
+ )
62
+
63
+ try:
64
+ # Connect to an MCP server
65
+ # Replace with your actual server command and args
66
+ await agent.connect_to_server("npx", ["@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
67
+
68
+ # Run the agent with a user query
69
+ result = await agent.run(task)
70
+ print("\nFinal result:", result)
71
+ return result
72
+ finally:
73
+ # Clean up resources
74
+ await agent.close()
75
+
76
+ # Example usage
77
+ task = dedent("""
78
+ I need accommodation in Toronto between 15th to 20th of May. Give me 5 options for 2 adults.
79
+ """)
80
+ await test_agent(task, model="gpt-4.1-mini")
81
+ ```
82
+
83
+ ## How the TinyAgent Hook System Works
84
+
85
+ TinyAgent is designed to be **extensible** via a simple, event-driven hook (callback) system. This allows you to add custom logic, logging, UI, memory, or any other behavior at key points in the agent's lifecycle.
86
+
87
+ ### How Hooks Work
88
+
89
+ - **Hooks** are just callables (functions or classes with `__call__`) that receive events from the agent.
90
+ - You register hooks using `agent.add_callback(hook)`.
91
+ - Hooks are called with:
92
+ `event_name, agent, **kwargs`
93
+ - Events include:
94
+ - `"agent_start"`: Agent is starting a new run
95
+ - `"message_add"`: A new message is added to the conversation
96
+ - `"llm_start"`: LLM is about to be called
97
+ - `"llm_end"`: LLM call finished
98
+ - `"agent_end"`: Agent is done (final result)
99
+ - (MCPClient also emits `"tool_start"` and `"tool_end"` for tool calls)
100
+
101
+ Hooks can be **async** or regular functions. If a hook is a class with an async `__call__`, it will be awaited.
102
+
103
+ #### Example: Adding a Custom Hook
104
+
105
+ ```python
106
+ def my_logger_hook(event_name, agent, **kwargs):
107
+ print(f"[{event_name}] {kwargs}")
108
+
109
+ agent.add_callback(my_logger_hook)
110
+ ```
111
+
112
+ #### Example: Async Hook
113
+
114
+ ```python
115
+ async def my_async_hook(event_name, agent, **kwargs):
116
+ if event_name == "agent_end":
117
+ print("Agent finished with result:", kwargs.get("result"))
118
+
119
+ agent.add_callback(my_async_hook)
120
+ ```
121
+
122
+ #### Example: Class-based Hook
123
+
124
+ ```python
125
+ class MyHook:
126
+ async def __call__(self, event_name, agent, **kwargs):
127
+ if event_name == "llm_start":
128
+ print("LLM is starting...")
129
+
130
+ agent.add_callback(MyHook())
131
+ ```
132
+
133
+ ### How to Extend the Hook System
134
+
135
+ - **Create your own hook**: Write a function or class as above.
136
+ - **Register it**: Use `agent.add_callback(your_hook)`.
137
+ - **Listen for events**: Check `event_name` and use `**kwargs` for event data.
138
+ - **See examples**: Each official hook (see below) includes a `run_example()` in its file.
139
+
140
+ ---
141
+
142
+ ## List of Available Hooks
143
+
144
+ You can import and use these hooks from `tinyagent.hooks`:
145
+
146
+ | Hook Name | Description | Example Import |
147
+ |--------------------------|--------------------------------------------------|-------------------------------------------------|
148
+ | `LoggingManager` | Granular logging control for all modules | `from tinyagent.hooks.logging_manager import LoggingManager` |
149
+ | `RichUICallback` | Rich terminal UI (with [rich](https://github.com/Textualize/rich)) | `from tinyagent.hooks.rich_ui_callback import RichUICallback` |
150
+ | `GradioCallback` | Interactive browser-based chat UI: file uploads, live thinking, tool calls, token stats | `from tinyagent.hooks.gradio_callback import GradioCallback` |
151
+
152
+ To see more details and usage, check the docstrings and `run_example()` in each hook file.
153
+
154
+ ## Using the GradioCallback Hook
155
+
156
+ The `GradioCallback` hook lets you spin up a full-featured web chat interface for your agent in just a few lines. You get:
157
+
158
+ Features:
159
+ - **Browser-based chat** with streaming updates
160
+ - **File uploads** (\*.pdf, \*.docx, \*.txt) that the agent can reference
161
+ - **Live “thinking” view** so you see intermediate thoughts
162
+ - **Collapsible tool-call sections** showing inputs & outputs
163
+ - **Real-time token usage** (prompt, completion, total)
164
+ - **Toggleable display options** for thinking & tool calls
165
+ - **Non-blocking launch** for asyncio apps (`prevent_thread_lock=True`)
166
+
167
+ ```python
168
+ import asyncio
169
+ from tinyagent import TinyAgent
170
+ from tinyagent.hooks.gradio_callback import GradioCallback
171
+ async def main():
172
+ # 1. Initialize your agent
173
+ agent = TinyAgent(model="gpt-4.1-mini", api_key="YOUR_API_KEY")
174
+ # 2. (Optional) Add tools or connect to MCP servers
175
+ # await agent.connect_to_server("npx", ["-y","@openbnb/mcp-server-airbnb","--ignore-robots-txt"])
176
+ # 3. Instantiate the Gradio UI callback
177
+ gradio_ui = GradioCallback(
178
+ file_upload_folder="uploads/",
179
+ show_thinking=True,
180
+ show_tool_calls=True
181
+ )
182
+ # 4. Register the callback with the agent
183
+ agent.add_callback(gradio_ui)
184
+ # 5. Launch the web interface (non-blocking)
185
+ gradio_ui.launch(
186
+ agent,
187
+ title="TinyAgent Chat",
188
+ description="Ask me to plan a trip or fetch data!",
189
+ share=False,
190
+ prevent_thread_lock=True
191
+ )
192
+ if __name__ == "__main__":
193
+ asyncio.run(main())
194
+ ```
195
+ ---
196
+
197
+ ## Contributing Hooks
198
+
199
+ - Place new hooks in the `tinyagent/hooks/` directory.
200
+ - Add an example usage as `async def run_example()` in the same file.
201
+ - Use `"gpt-4.1-mini"` as the default model in examples.
202
+
203
+ ---
204
+
205
+ ## License
206
+
207
+ MIT License. See [LICENSE](LICENSE).