tinyagent-py 0.0.1__tar.gz → 0.0.4__tar.gz

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.
Files changed (31) hide show
  1. tinyagent_py-0.0.4/PKG-INFO +252 -0
  2. tinyagent_py-0.0.4/README.md +219 -0
  3. {tinyagent_py-0.0.1 → tinyagent_py-0.0.4}/pyproject.toml +26 -6
  4. tinyagent_py-0.0.4/tinyagent/hooks/__init__.py +4 -0
  5. tinyagent_py-0.0.4/tinyagent/hooks/agno_storage_hook.py +128 -0
  6. tinyagent_py-0.0.4/tinyagent/hooks/gradio_callback.py +966 -0
  7. tinyagent_py-0.0.4/tinyagent/hooks/logging_manager.py +213 -0
  8. tinyagent_py-0.0.4/tinyagent/hooks/rich_ui_callback.py +559 -0
  9. tinyagent_py-0.0.4/tinyagent/storage/__init__.py +7 -0
  10. tinyagent_py-0.0.4/tinyagent/storage/agno_storage.py +114 -0
  11. tinyagent_py-0.0.4/tinyagent/storage/base.py +49 -0
  12. tinyagent_py-0.0.4/tinyagent/storage/json_file_storage.py +30 -0
  13. tinyagent_py-0.0.4/tinyagent/storage/postgres_storage.py +201 -0
  14. tinyagent_py-0.0.4/tinyagent/storage/redis_storage.py +48 -0
  15. tinyagent_py-0.0.4/tinyagent/storage/sqlite_storage.py +156 -0
  16. tinyagent_py-0.0.4/tinyagent/tinyagent_py.egg-info/PKG-INFO +252 -0
  17. tinyagent_py-0.0.4/tinyagent/tinyagent_py.egg-info/SOURCES.txt +20 -0
  18. tinyagent_py-0.0.4/tinyagent/tinyagent_py.egg-info/requires.txt +25 -0
  19. tinyagent_py-0.0.4/tinyagent/tinyagent_py.egg-info/top_level.txt +2 -0
  20. tinyagent_py-0.0.1/PKG-INFO +0 -79
  21. tinyagent_py-0.0.1/README.md +0 -60
  22. tinyagent_py-0.0.1/tinyagent/__init__.py +0 -4
  23. tinyagent_py-0.0.1/tinyagent/mcp_client.py +0 -52
  24. tinyagent_py-0.0.1/tinyagent/tiny_agent.py +0 -247
  25. tinyagent_py-0.0.1/tinyagent_py.egg-info/PKG-INFO +0 -79
  26. tinyagent_py-0.0.1/tinyagent_py.egg-info/SOURCES.txt +0 -11
  27. tinyagent_py-0.0.1/tinyagent_py.egg-info/requires.txt +0 -6
  28. tinyagent_py-0.0.1/tinyagent_py.egg-info/top_level.txt +0 -1
  29. {tinyagent_py-0.0.1 → tinyagent_py-0.0.4}/LICENSE +0 -0
  30. {tinyagent_py-0.0.1 → tinyagent_py-0.0.4}/setup.cfg +0 -0
  31. {tinyagent_py-0.0.1 → tinyagent_py-0.0.4/tinyagent}/tinyagent_py.egg-info/dependency_links.txt +0 -0
@@ -0,0 +1,252 @@
1
+ Metadata-Version: 2.4
2
+ Name: tinyagent-py
3
+ Version: 0.0.4
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
+ Requires-Dist: tiktoken
16
+ Requires-Dist: uuid
17
+ Provides-Extra: dev
18
+ Requires-Dist: pytest; extra == "dev"
19
+ Requires-Dist: black; extra == "dev"
20
+ Requires-Dist: isort; extra == "dev"
21
+ Requires-Dist: mypy; extra == "dev"
22
+ Provides-Extra: postgres
23
+ Requires-Dist: asyncpg>=0.27.0; extra == "postgres"
24
+ Provides-Extra: sqlite
25
+ Requires-Dist: aiosqlite>=0.18.0; extra == "sqlite"
26
+ Provides-Extra: gradio
27
+ Requires-Dist: gradio>=3.50.0; extra == "gradio"
28
+ Provides-Extra: all
29
+ Requires-Dist: asyncpg>=0.27.0; extra == "all"
30
+ Requires-Dist: aiosqlite>=0.18.0; extra == "all"
31
+ Requires-Dist: gradio>=3.50.0; extra == "all"
32
+ Dynamic: license-file
33
+
34
+ # tinyagent
35
+ Tiny Agent: 100 lines Agent with MCP
36
+ ![TinyAgent Logo](https://raw.githubusercontent.com/askbudi/tinyagent/main/public/logo.png)
37
+
38
+
39
+
40
+ Inspired by:
41
+ - [Tiny Agents blog post](https://huggingface.co/blog/tiny-agents)
42
+ - [12-factor-agents repository](https://github.com/humanlayer/12-factor-agents)
43
+ - Created by chatting to the source code of JS Tiny Agent using [AskDev.ai](https://askdev.ai/search)
44
+
45
+ ## Quick Links
46
+ - [Build your own Tiny Agent](https://askdev.ai/github/askbudi/tinyagent)
47
+
48
+ ## Overview
49
+ 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.
50
+
51
+ ## Installation
52
+
53
+ ### Using pip
54
+ ```bash
55
+ # Basic installation
56
+ pip install tinyagent-py
57
+
58
+ # Install with all optional dependencies
59
+ pip install tinyagent-py[all]
60
+
61
+ # Install with PostgreSQL support
62
+ pip install tinyagent-py[postgres]
63
+
64
+ # Install with SQLite support
65
+ pip install tinyagent-py[sqlite]
66
+
67
+ # Install with Gradio UI support
68
+ pip install tinyagent-py[gradio]
69
+
70
+ ```
71
+
72
+ ### Using uv
73
+ ```bash
74
+ # Basic installation
75
+ uv pip install tinyagent-py
76
+
77
+ # Install with PostgreSQL support
78
+ uv pip install tinyagent-py[postgres]
79
+
80
+ # Install with SQLite support
81
+ uv pip install tinyagent-py[sqlite]
82
+
83
+ # Install with Gradio UI support
84
+ uv pip install tinyagent-py[gradio]
85
+
86
+ # Install with all optional dependencies
87
+ uv pip install tinyagent-py[all]
88
+
89
+ # Install with development tools
90
+ uv pip install tinyagent-py[dev]
91
+ ```
92
+
93
+ ## Usage
94
+
95
+ ```python
96
+ from tinyagent import TinyAgent
97
+ from textwrap import dedent
98
+ import asyncio
99
+ import os
100
+
101
+ async def test_agent(task, model="o4-mini", api_key=None):
102
+ # Initialize the agent with model and API key
103
+ agent = TinyAgent(
104
+ model=model, # Or any model supported by LiteLLM
105
+ api_key=os.environ.get("OPENAI_API_KEY") if not api_key else api_key # Set your API key as an env variable
106
+ )
107
+
108
+ try:
109
+ # Connect to an MCP server
110
+ # Replace with your actual server command and args
111
+ await agent.connect_to_server("npx", ["@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
112
+
113
+ # Run the agent with a user query
114
+ result = await agent.run(task)
115
+ print("\nFinal result:", result)
116
+ return result
117
+ finally:
118
+ # Clean up resources
119
+ await agent.close()
120
+
121
+ # Example usage
122
+ task = dedent("""
123
+ I need accommodation in Toronto between 15th to 20th of May. Give me 5 options for 2 adults.
124
+ """)
125
+ await test_agent(task, model="gpt-4.1-mini")
126
+ ```
127
+
128
+ ## How the TinyAgent Hook System Works
129
+
130
+ 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.
131
+
132
+ ### How Hooks Work
133
+
134
+ - **Hooks** are just callables (functions or classes with `__call__`) that receive events from the agent.
135
+ - You register hooks using `agent.add_callback(hook)`.
136
+ - Hooks are called with:
137
+ `event_name, agent, **kwargs`
138
+ - Events include:
139
+ - `"agent_start"`: Agent is starting a new run
140
+ - `"message_add"`: A new message is added to the conversation
141
+ - `"llm_start"`: LLM is about to be called
142
+ - `"llm_end"`: LLM call finished
143
+ - `"agent_end"`: Agent is done (final result)
144
+ - (MCPClient also emits `"tool_start"` and `"tool_end"` for tool calls)
145
+
146
+ Hooks can be **async** or regular functions. If a hook is a class with an async `__call__`, it will be awaited.
147
+
148
+ #### Example: Adding a Custom Hook
149
+
150
+ ```python
151
+ def my_logger_hook(event_name, agent, **kwargs):
152
+ print(f"[{event_name}] {kwargs}")
153
+
154
+ agent.add_callback(my_logger_hook)
155
+ ```
156
+
157
+ #### Example: Async Hook
158
+
159
+ ```python
160
+ async def my_async_hook(event_name, agent, **kwargs):
161
+ if event_name == "agent_end":
162
+ print("Agent finished with result:", kwargs.get("result"))
163
+
164
+ agent.add_callback(my_async_hook)
165
+ ```
166
+
167
+ #### Example: Class-based Hook
168
+
169
+ ```python
170
+ class MyHook:
171
+ async def __call__(self, event_name, agent, **kwargs):
172
+ if event_name == "llm_start":
173
+ print("LLM is starting...")
174
+
175
+ agent.add_callback(MyHook())
176
+ ```
177
+
178
+ ### How to Extend the Hook System
179
+
180
+ - **Create your own hook**: Write a function or class as above.
181
+ - **Register it**: Use `agent.add_callback(your_hook)`.
182
+ - **Listen for events**: Check `event_name` and use `**kwargs` for event data.
183
+ - **See examples**: Each official hook (see below) includes a `run_example()` in its file.
184
+
185
+ ---
186
+
187
+ ## List of Available Hooks
188
+
189
+ You can import and use these hooks from `tinyagent.hooks`:
190
+
191
+ | Hook Name | Description | Example Import |
192
+ |--------------------------|--------------------------------------------------|-------------------------------------------------|
193
+ | `LoggingManager` | Granular logging control for all modules | `from tinyagent.hooks.logging_manager import LoggingManager` |
194
+ | `RichUICallback` | Rich terminal UI (with [rich](https://github.com/Textualize/rich)) | `from tinyagent.hooks.rich_ui_callback import RichUICallback` |
195
+ | `GradioCallback` | Interactive browser-based chat UI: file uploads, live thinking, tool calls, token stats | `from tinyagent.hooks.gradio_callback import GradioCallback` |
196
+
197
+ To see more details and usage, check the docstrings and `run_example()` in each hook file.
198
+
199
+ ## Using the GradioCallback Hook
200
+
201
+ The `GradioCallback` hook lets you spin up a full-featured web chat interface for your agent in just a few lines. You get:
202
+
203
+ Features:
204
+ - **Browser-based chat** with streaming updates
205
+ - **File uploads** (\*.pdf, \*.docx, \*.txt) that the agent can reference
206
+ - **Live "thinking" view** so you see intermediate thoughts
207
+ - **Collapsible tool-call sections** showing inputs & outputs
208
+ - **Real-time token usage** (prompt, completion, total)
209
+ - **Toggleable display options** for thinking & tool calls
210
+ - **Non-blocking launch** for asyncio apps (`prevent_thread_lock=True`)
211
+
212
+ ```python
213
+ import asyncio
214
+ from tinyagent import TinyAgent
215
+ from tinyagent.hooks.gradio_callback import GradioCallback
216
+ async def main():
217
+ # 1. Initialize your agent
218
+ agent = TinyAgent(model="gpt-4.1-mini", api_key="YOUR_API_KEY")
219
+ # 2. (Optional) Add tools or connect to MCP servers
220
+ # await agent.connect_to_server("npx", ["-y","@openbnb/mcp-server-airbnb","--ignore-robots-txt"])
221
+ # 3. Instantiate the Gradio UI callback
222
+ gradio_ui = GradioCallback(
223
+ file_upload_folder="uploads/",
224
+ show_thinking=True,
225
+ show_tool_calls=True
226
+ )
227
+ # 4. Register the callback with the agent
228
+ agent.add_callback(gradio_ui)
229
+ # 5. Launch the web interface (non-blocking)
230
+ gradio_ui.launch(
231
+ agent,
232
+ title="TinyAgent Chat",
233
+ description="Ask me to plan a trip or fetch data!",
234
+ share=False,
235
+ prevent_thread_lock=True
236
+ )
237
+ if __name__ == "__main__":
238
+ asyncio.run(main())
239
+ ```
240
+ ---
241
+
242
+ ## Contributing Hooks
243
+
244
+ - Place new hooks in the `tinyagent/hooks/` directory.
245
+ - Add an example usage as `async def run_example()` in the same file.
246
+ - Use `"gpt-4.1-mini"` as the default model in examples.
247
+
248
+ ---
249
+
250
+ ## License
251
+
252
+ MIT License. See [LICENSE](LICENSE).
@@ -0,0 +1,219 @@
1
+ # tinyagent
2
+ Tiny Agent: 100 lines Agent with MCP
3
+ ![TinyAgent Logo](https://raw.githubusercontent.com/askbudi/tinyagent/main/public/logo.png)
4
+
5
+
6
+
7
+ Inspired by:
8
+ - [Tiny Agents blog post](https://huggingface.co/blog/tiny-agents)
9
+ - [12-factor-agents repository](https://github.com/humanlayer/12-factor-agents)
10
+ - Created by chatting to the source code of JS Tiny Agent using [AskDev.ai](https://askdev.ai/search)
11
+
12
+ ## Quick Links
13
+ - [Build your own Tiny Agent](https://askdev.ai/github/askbudi/tinyagent)
14
+
15
+ ## Overview
16
+ 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.
17
+
18
+ ## Installation
19
+
20
+ ### Using pip
21
+ ```bash
22
+ # Basic installation
23
+ pip install tinyagent-py
24
+
25
+ # Install with all optional dependencies
26
+ pip install tinyagent-py[all]
27
+
28
+ # Install with PostgreSQL support
29
+ pip install tinyagent-py[postgres]
30
+
31
+ # Install with SQLite support
32
+ pip install tinyagent-py[sqlite]
33
+
34
+ # Install with Gradio UI support
35
+ pip install tinyagent-py[gradio]
36
+
37
+ ```
38
+
39
+ ### Using uv
40
+ ```bash
41
+ # Basic installation
42
+ uv pip install tinyagent-py
43
+
44
+ # Install with PostgreSQL support
45
+ uv pip install tinyagent-py[postgres]
46
+
47
+ # Install with SQLite support
48
+ uv pip install tinyagent-py[sqlite]
49
+
50
+ # Install with Gradio UI support
51
+ uv pip install tinyagent-py[gradio]
52
+
53
+ # Install with all optional dependencies
54
+ uv pip install tinyagent-py[all]
55
+
56
+ # Install with development tools
57
+ uv pip install tinyagent-py[dev]
58
+ ```
59
+
60
+ ## Usage
61
+
62
+ ```python
63
+ from tinyagent import TinyAgent
64
+ from textwrap import dedent
65
+ import asyncio
66
+ import os
67
+
68
+ async def test_agent(task, model="o4-mini", api_key=None):
69
+ # Initialize the agent with model and API key
70
+ agent = TinyAgent(
71
+ model=model, # Or any model supported by LiteLLM
72
+ api_key=os.environ.get("OPENAI_API_KEY") if not api_key else api_key # Set your API key as an env variable
73
+ )
74
+
75
+ try:
76
+ # Connect to an MCP server
77
+ # Replace with your actual server command and args
78
+ await agent.connect_to_server("npx", ["@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
79
+
80
+ # Run the agent with a user query
81
+ result = await agent.run(task)
82
+ print("\nFinal result:", result)
83
+ return result
84
+ finally:
85
+ # Clean up resources
86
+ await agent.close()
87
+
88
+ # Example usage
89
+ task = dedent("""
90
+ I need accommodation in Toronto between 15th to 20th of May. Give me 5 options for 2 adults.
91
+ """)
92
+ await test_agent(task, model="gpt-4.1-mini")
93
+ ```
94
+
95
+ ## How the TinyAgent Hook System Works
96
+
97
+ 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.
98
+
99
+ ### How Hooks Work
100
+
101
+ - **Hooks** are just callables (functions or classes with `__call__`) that receive events from the agent.
102
+ - You register hooks using `agent.add_callback(hook)`.
103
+ - Hooks are called with:
104
+ `event_name, agent, **kwargs`
105
+ - Events include:
106
+ - `"agent_start"`: Agent is starting a new run
107
+ - `"message_add"`: A new message is added to the conversation
108
+ - `"llm_start"`: LLM is about to be called
109
+ - `"llm_end"`: LLM call finished
110
+ - `"agent_end"`: Agent is done (final result)
111
+ - (MCPClient also emits `"tool_start"` and `"tool_end"` for tool calls)
112
+
113
+ Hooks can be **async** or regular functions. If a hook is a class with an async `__call__`, it will be awaited.
114
+
115
+ #### Example: Adding a Custom Hook
116
+
117
+ ```python
118
+ def my_logger_hook(event_name, agent, **kwargs):
119
+ print(f"[{event_name}] {kwargs}")
120
+
121
+ agent.add_callback(my_logger_hook)
122
+ ```
123
+
124
+ #### Example: Async Hook
125
+
126
+ ```python
127
+ async def my_async_hook(event_name, agent, **kwargs):
128
+ if event_name == "agent_end":
129
+ print("Agent finished with result:", kwargs.get("result"))
130
+
131
+ agent.add_callback(my_async_hook)
132
+ ```
133
+
134
+ #### Example: Class-based Hook
135
+
136
+ ```python
137
+ class MyHook:
138
+ async def __call__(self, event_name, agent, **kwargs):
139
+ if event_name == "llm_start":
140
+ print("LLM is starting...")
141
+
142
+ agent.add_callback(MyHook())
143
+ ```
144
+
145
+ ### How to Extend the Hook System
146
+
147
+ - **Create your own hook**: Write a function or class as above.
148
+ - **Register it**: Use `agent.add_callback(your_hook)`.
149
+ - **Listen for events**: Check `event_name` and use `**kwargs` for event data.
150
+ - **See examples**: Each official hook (see below) includes a `run_example()` in its file.
151
+
152
+ ---
153
+
154
+ ## List of Available Hooks
155
+
156
+ You can import and use these hooks from `tinyagent.hooks`:
157
+
158
+ | Hook Name | Description | Example Import |
159
+ |--------------------------|--------------------------------------------------|-------------------------------------------------|
160
+ | `LoggingManager` | Granular logging control for all modules | `from tinyagent.hooks.logging_manager import LoggingManager` |
161
+ | `RichUICallback` | Rich terminal UI (with [rich](https://github.com/Textualize/rich)) | `from tinyagent.hooks.rich_ui_callback import RichUICallback` |
162
+ | `GradioCallback` | Interactive browser-based chat UI: file uploads, live thinking, tool calls, token stats | `from tinyagent.hooks.gradio_callback import GradioCallback` |
163
+
164
+ To see more details and usage, check the docstrings and `run_example()` in each hook file.
165
+
166
+ ## Using the GradioCallback Hook
167
+
168
+ The `GradioCallback` hook lets you spin up a full-featured web chat interface for your agent in just a few lines. You get:
169
+
170
+ Features:
171
+ - **Browser-based chat** with streaming updates
172
+ - **File uploads** (\*.pdf, \*.docx, \*.txt) that the agent can reference
173
+ - **Live "thinking" view** so you see intermediate thoughts
174
+ - **Collapsible tool-call sections** showing inputs & outputs
175
+ - **Real-time token usage** (prompt, completion, total)
176
+ - **Toggleable display options** for thinking & tool calls
177
+ - **Non-blocking launch** for asyncio apps (`prevent_thread_lock=True`)
178
+
179
+ ```python
180
+ import asyncio
181
+ from tinyagent import TinyAgent
182
+ from tinyagent.hooks.gradio_callback import GradioCallback
183
+ async def main():
184
+ # 1. Initialize your agent
185
+ agent = TinyAgent(model="gpt-4.1-mini", api_key="YOUR_API_KEY")
186
+ # 2. (Optional) Add tools or connect to MCP servers
187
+ # await agent.connect_to_server("npx", ["-y","@openbnb/mcp-server-airbnb","--ignore-robots-txt"])
188
+ # 3. Instantiate the Gradio UI callback
189
+ gradio_ui = GradioCallback(
190
+ file_upload_folder="uploads/",
191
+ show_thinking=True,
192
+ show_tool_calls=True
193
+ )
194
+ # 4. Register the callback with the agent
195
+ agent.add_callback(gradio_ui)
196
+ # 5. Launch the web interface (non-blocking)
197
+ gradio_ui.launch(
198
+ agent,
199
+ title="TinyAgent Chat",
200
+ description="Ask me to plan a trip or fetch data!",
201
+ share=False,
202
+ prevent_thread_lock=True
203
+ )
204
+ if __name__ == "__main__":
205
+ asyncio.run(main())
206
+ ```
207
+ ---
208
+
209
+ ## Contributing Hooks
210
+
211
+ - Place new hooks in the `tinyagent/hooks/` directory.
212
+ - Add an example usage as `async def run_example()` in the same file.
213
+ - Use `"gpt-4.1-mini"` as the default model in examples.
214
+
215
+ ---
216
+
217
+ ## License
218
+
219
+ MIT License. See [LICENSE](LICENSE).
@@ -1,13 +1,16 @@
1
1
  [build-system]
2
- requires = ["setuptools>=61.0", "wheel"]
2
+ requires = ["setuptools>=77.0", "wheel"]
3
3
  build-backend = "setuptools.build_meta"
4
4
 
5
+ [tool.setuptools.packages.find]
6
+ where = ["tinyagent"]
7
+ exclude = ["public", "public.*"]
8
+
5
9
  [project]
6
10
  name = "tinyagent-py"
7
- version = "0.0.1"
11
+ version = "0.0.4"
8
12
  description = "Tiny Agent with MCP Client"
9
13
  readme = "README.md"
10
- license = {text = "MIT"}
11
14
  authors = [
12
15
  {name="Mahdi Golchin", email="golchin@askdev.ai"}
13
16
  ]
@@ -15,15 +18,32 @@ requires-python = ">=3.8"
15
18
  dependencies = [
16
19
  "mcp",
17
20
  "litellm",
18
- "openai"
21
+ "openai",
22
+ "tiktoken",
23
+ "uuid",
19
24
  # add whatever else you need…
20
25
  ]
21
26
 
22
27
  [project.optional-dependencies]
23
28
  dev = [
24
29
  "pytest",
25
-
26
-
30
+ "black",
31
+ "isort",
32
+ "mypy"
33
+ ]
34
+ postgres = [
35
+ "asyncpg>=0.27.0"
36
+ ]
37
+ sqlite = [
38
+ "aiosqlite>=0.18.0"
39
+ ]
40
+ gradio = [
41
+ "gradio>=3.50.0"
42
+ ]
43
+ all = [
44
+ "asyncpg>=0.27.0",
45
+ "aiosqlite>=0.18.0",
46
+ "gradio>=3.50.0"
27
47
  ]
28
48
  [project.urls]
29
49
  "Homepage" = "https://github.com/askbudi/tinyagent"
@@ -0,0 +1,4 @@
1
+ #from .rich_ui_agent import RichUICallback
2
+ from .rich_ui_callback import RichUICallback
3
+ from .logging_manager import LoggingManager
4
+ __all__ = ["RichUICallback", "LoggingManager"]
@@ -0,0 +1,128 @@
1
+ try:
2
+ import agno
3
+ from agno.storage.postgres import PostgresStorage
4
+ from agno.storage.sqlite import SqliteStorage
5
+ from agno.storage.session.agent import AgentSession
6
+ except ImportError as e:
7
+ raise ImportError("agno is not installed. Please install it with `pip install agno`.", e)
8
+
9
+ import asyncio
10
+ from typing import Optional
11
+ from agno.storage.postgres import PostgresStorage
12
+ from agno.storage.sqlite import SqliteStorage
13
+ from agno.storage.session.agent import AgentSession
14
+
15
+ class PostgresStorageHook:
16
+ def __init__(
17
+ self,
18
+ table_name: str,
19
+ db_url: Optional[str] = None,
20
+ db_engine=None,
21
+ schema: Optional[str] = "ai",
22
+ schema_version: int = 1,
23
+ auto_upgrade_schema: bool = True,
24
+ mode: str = "agent",
25
+ ):
26
+ self.storage = PostgresStorage(
27
+ table_name=table_name,
28
+ db_url=db_url,
29
+ db_engine=db_engine,
30
+ schema=schema,
31
+ schema_version=schema_version,
32
+ auto_upgrade_schema=auto_upgrade_schema,
33
+ mode=mode,
34
+ )
35
+
36
+ async def __call__(self, event_name: str, agent, **kwargs):
37
+ if event_name == "agent_start":
38
+ # Load session from storage
39
+ session_id = getattr(agent, "session_id", None)
40
+ user_id = getattr(agent, "user_id", None)
41
+ if session_id:
42
+ session = self.storage.read(session_id=session_id, user_id=user_id)
43
+ if session:
44
+ # Populate agent state from session
45
+ agent.messages = session.session_data.get("messages", [])
46
+ agent.memory = session.memory
47
+ agent.metadata = session.extra_data
48
+ # You may need to adapt this depending on tinyagent's state structure
49
+
50
+ elif event_name in ("llm_end", "agent_end"):
51
+ # Save session to storage
52
+ session_id = getattr(agent, "session_id", None)
53
+ user_id = getattr(agent, "user_id", None)
54
+ if session_id:
55
+ # Create AgentSession from agent state
56
+ session_data = {
57
+ "messages": getattr(agent, "messages", []),
58
+ }
59
+ session = AgentSession(
60
+ session_id=session_id,
61
+ user_id=user_id,
62
+ memory=getattr(agent, "memory", {}),
63
+ session_data=session_data,
64
+ extra_data=getattr(agent, "metadata", {}),
65
+ agent_id=getattr(agent, "agent_id", None),
66
+ team_session_id=None,
67
+ agent_data=None,
68
+ )
69
+ await asyncio.to_thread(self.storage.upsert, session)
70
+
71
+ class SqliteStorageHook:
72
+ def __init__(
73
+ self,
74
+ table_name: str,
75
+ db_url: Optional[str] = None,
76
+ db_file: Optional[str] = None,
77
+ db_engine=None,
78
+ schema_version: int = 1,
79
+ auto_upgrade_schema: bool = True,
80
+ mode: str = "agent",
81
+ ):
82
+ self.storage = SqliteStorage(
83
+ table_name=table_name,
84
+ db_url=db_url,
85
+ db_file=db_file,
86
+ db_engine=db_engine,
87
+ schema_version=schema_version,
88
+ auto_upgrade_schema=auto_upgrade_schema,
89
+ mode=mode,
90
+ )
91
+
92
+ async def __call__(self, event_name: str, agent, **kwargs):
93
+ if event_name == "agent_start":
94
+ # Load session from storage
95
+ session_id = getattr(agent, "session_id", None)
96
+ user_id = getattr(agent, "user_id", None)
97
+ print("Session ID",session_id)
98
+ print("User ID",user_id)
99
+ if session_id:
100
+ session = self.storage.read(session_id=session_id, user_id=user_id)
101
+ print(f"Session: {session}")
102
+ if session:
103
+ # Populate agent state from session
104
+ agent.messages = session.memory.get("messages", [])
105
+ agent.memory = session.memory
106
+ agent.metadata = session.extra_data
107
+
108
+ elif event_name in ("llm_end", "agent_end"):
109
+ # Save session to storage
110
+ print("Agent metadata",getattr(agent, "metadata", {}))
111
+ session_id = getattr(agent, "session_id", None)
112
+ user_id = getattr(agent, "user_id", None)
113
+ if session_id:
114
+ session_data = {
115
+ "messages": getattr(agent, "messages", []),
116
+ }
117
+ session = AgentSession(
118
+ session_id=session_id,
119
+ user_id=user_id,
120
+ memory=getattr(agent, "memory", {}),
121
+ session_data=session_data,
122
+ extra_data=getattr(agent, "metadata", {}),
123
+ agent_id=getattr(agent, "agent_id", None),
124
+ team_session_id=None,
125
+ agent_data=None,
126
+ )
127
+ await asyncio.to_thread(self.storage.upsert, session)
128
+