tinyagent-py 0.0.1__tar.gz → 0.0.3__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 (30) hide show
  1. tinyagent_py-0.0.3/PKG-INFO +207 -0
  2. tinyagent_py-0.0.3/README.md +189 -0
  3. {tinyagent_py-0.0.1 → tinyagent_py-0.0.3}/pyproject.toml +6 -3
  4. tinyagent_py-0.0.3/tinyagent/hooks/__init__.py +4 -0
  5. tinyagent_py-0.0.3/tinyagent/hooks/agno_storage_hook.py +128 -0
  6. tinyagent_py-0.0.3/tinyagent/hooks/gradio_callback.py +966 -0
  7. tinyagent_py-0.0.3/tinyagent/hooks/logging_manager.py +213 -0
  8. tinyagent_py-0.0.3/tinyagent/hooks/rich_ui_callback.py +559 -0
  9. tinyagent_py-0.0.3/tinyagent/storage/__init__.py +7 -0
  10. tinyagent_py-0.0.3/tinyagent/storage/agno_storage.py +114 -0
  11. tinyagent_py-0.0.3/tinyagent/storage/base.py +49 -0
  12. tinyagent_py-0.0.3/tinyagent/storage/json_file_storage.py +30 -0
  13. tinyagent_py-0.0.3/tinyagent/storage/postgres_storage.py +201 -0
  14. tinyagent_py-0.0.3/tinyagent/storage/redis_storage.py +48 -0
  15. tinyagent_py-0.0.3/tinyagent/storage/sqlite_storage.py +156 -0
  16. tinyagent_py-0.0.3/tinyagent/tinyagent_py.egg-info/PKG-INFO +207 -0
  17. tinyagent_py-0.0.3/tinyagent/tinyagent_py.egg-info/SOURCES.txt +20 -0
  18. tinyagent_py-0.0.3/tinyagent/tinyagent_py.egg-info/top_level.txt +2 -0
  19. tinyagent_py-0.0.1/PKG-INFO +0 -79
  20. tinyagent_py-0.0.1/README.md +0 -60
  21. tinyagent_py-0.0.1/tinyagent/__init__.py +0 -4
  22. tinyagent_py-0.0.1/tinyagent/mcp_client.py +0 -52
  23. tinyagent_py-0.0.1/tinyagent/tiny_agent.py +0 -247
  24. tinyagent_py-0.0.1/tinyagent_py.egg-info/PKG-INFO +0 -79
  25. tinyagent_py-0.0.1/tinyagent_py.egg-info/SOURCES.txt +0 -11
  26. tinyagent_py-0.0.1/tinyagent_py.egg-info/top_level.txt +0 -1
  27. {tinyagent_py-0.0.1 → tinyagent_py-0.0.3}/LICENSE +0 -0
  28. {tinyagent_py-0.0.1 → tinyagent_py-0.0.3}/setup.cfg +0 -0
  29. {tinyagent_py-0.0.1 → tinyagent_py-0.0.3/tinyagent}/tinyagent_py.egg-info/dependency_links.txt +0 -0
  30. {tinyagent_py-0.0.1 → tinyagent_py-0.0.3/tinyagent}/tinyagent_py.egg-info/requires.txt +0 -0
@@ -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).
@@ -0,0 +1,189 @@
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
+ pip install tinyagent-py
23
+ ```
24
+
25
+ ### Using uv
26
+ ```bash
27
+ uv pip install tinyagent-py
28
+ ```
29
+
30
+ ## Usage
31
+
32
+ ```python
33
+ from tinyagent import TinyAgent
34
+ from textwrap import dedent
35
+ import asyncio
36
+ import os
37
+
38
+ async def test_agent(task, model="o4-mini", api_key=None):
39
+ # Initialize the agent with model and API key
40
+ agent = TinyAgent(
41
+ model=model, # Or any model supported by LiteLLM
42
+ api_key=os.environ.get("OPENAI_API_KEY") if not api_key else api_key # Set your API key as an env variable
43
+ )
44
+
45
+ try:
46
+ # Connect to an MCP server
47
+ # Replace with your actual server command and args
48
+ await agent.connect_to_server("npx", ["@openbnb/mcp-server-airbnb", "--ignore-robots-txt"])
49
+
50
+ # Run the agent with a user query
51
+ result = await agent.run(task)
52
+ print("\nFinal result:", result)
53
+ return result
54
+ finally:
55
+ # Clean up resources
56
+ await agent.close()
57
+
58
+ # Example usage
59
+ task = dedent("""
60
+ I need accommodation in Toronto between 15th to 20th of May. Give me 5 options for 2 adults.
61
+ """)
62
+ await test_agent(task, model="gpt-4.1-mini")
63
+ ```
64
+
65
+ ## How the TinyAgent Hook System Works
66
+
67
+ 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.
68
+
69
+ ### How Hooks Work
70
+
71
+ - **Hooks** are just callables (functions or classes with `__call__`) that receive events from the agent.
72
+ - You register hooks using `agent.add_callback(hook)`.
73
+ - Hooks are called with:
74
+ `event_name, agent, **kwargs`
75
+ - Events include:
76
+ - `"agent_start"`: Agent is starting a new run
77
+ - `"message_add"`: A new message is added to the conversation
78
+ - `"llm_start"`: LLM is about to be called
79
+ - `"llm_end"`: LLM call finished
80
+ - `"agent_end"`: Agent is done (final result)
81
+ - (MCPClient also emits `"tool_start"` and `"tool_end"` for tool calls)
82
+
83
+ Hooks can be **async** or regular functions. If a hook is a class with an async `__call__`, it will be awaited.
84
+
85
+ #### Example: Adding a Custom Hook
86
+
87
+ ```python
88
+ def my_logger_hook(event_name, agent, **kwargs):
89
+ print(f"[{event_name}] {kwargs}")
90
+
91
+ agent.add_callback(my_logger_hook)
92
+ ```
93
+
94
+ #### Example: Async Hook
95
+
96
+ ```python
97
+ async def my_async_hook(event_name, agent, **kwargs):
98
+ if event_name == "agent_end":
99
+ print("Agent finished with result:", kwargs.get("result"))
100
+
101
+ agent.add_callback(my_async_hook)
102
+ ```
103
+
104
+ #### Example: Class-based Hook
105
+
106
+ ```python
107
+ class MyHook:
108
+ async def __call__(self, event_name, agent, **kwargs):
109
+ if event_name == "llm_start":
110
+ print("LLM is starting...")
111
+
112
+ agent.add_callback(MyHook())
113
+ ```
114
+
115
+ ### How to Extend the Hook System
116
+
117
+ - **Create your own hook**: Write a function or class as above.
118
+ - **Register it**: Use `agent.add_callback(your_hook)`.
119
+ - **Listen for events**: Check `event_name` and use `**kwargs` for event data.
120
+ - **See examples**: Each official hook (see below) includes a `run_example()` in its file.
121
+
122
+ ---
123
+
124
+ ## List of Available Hooks
125
+
126
+ You can import and use these hooks from `tinyagent.hooks`:
127
+
128
+ | Hook Name | Description | Example Import |
129
+ |--------------------------|--------------------------------------------------|-------------------------------------------------|
130
+ | `LoggingManager` | Granular logging control for all modules | `from tinyagent.hooks.logging_manager import LoggingManager` |
131
+ | `RichUICallback` | Rich terminal UI (with [rich](https://github.com/Textualize/rich)) | `from tinyagent.hooks.rich_ui_callback import RichUICallback` |
132
+ | `GradioCallback` | Interactive browser-based chat UI: file uploads, live thinking, tool calls, token stats | `from tinyagent.hooks.gradio_callback import GradioCallback` |
133
+
134
+ To see more details and usage, check the docstrings and `run_example()` in each hook file.
135
+
136
+ ## Using the GradioCallback Hook
137
+
138
+ The `GradioCallback` hook lets you spin up a full-featured web chat interface for your agent in just a few lines. You get:
139
+
140
+ Features:
141
+ - **Browser-based chat** with streaming updates
142
+ - **File uploads** (\*.pdf, \*.docx, \*.txt) that the agent can reference
143
+ - **Live “thinking” view** so you see intermediate thoughts
144
+ - **Collapsible tool-call sections** showing inputs & outputs
145
+ - **Real-time token usage** (prompt, completion, total)
146
+ - **Toggleable display options** for thinking & tool calls
147
+ - **Non-blocking launch** for asyncio apps (`prevent_thread_lock=True`)
148
+
149
+ ```python
150
+ import asyncio
151
+ from tinyagent import TinyAgent
152
+ from tinyagent.hooks.gradio_callback import GradioCallback
153
+ async def main():
154
+ # 1. Initialize your agent
155
+ agent = TinyAgent(model="gpt-4.1-mini", api_key="YOUR_API_KEY")
156
+ # 2. (Optional) Add tools or connect to MCP servers
157
+ # await agent.connect_to_server("npx", ["-y","@openbnb/mcp-server-airbnb","--ignore-robots-txt"])
158
+ # 3. Instantiate the Gradio UI callback
159
+ gradio_ui = GradioCallback(
160
+ file_upload_folder="uploads/",
161
+ show_thinking=True,
162
+ show_tool_calls=True
163
+ )
164
+ # 4. Register the callback with the agent
165
+ agent.add_callback(gradio_ui)
166
+ # 5. Launch the web interface (non-blocking)
167
+ gradio_ui.launch(
168
+ agent,
169
+ title="TinyAgent Chat",
170
+ description="Ask me to plan a trip or fetch data!",
171
+ share=False,
172
+ prevent_thread_lock=True
173
+ )
174
+ if __name__ == "__main__":
175
+ asyncio.run(main())
176
+ ```
177
+ ---
178
+
179
+ ## Contributing Hooks
180
+
181
+ - Place new hooks in the `tinyagent/hooks/` directory.
182
+ - Add an example usage as `async def run_example()` in the same file.
183
+ - Use `"gpt-4.1-mini"` as the default model in examples.
184
+
185
+ ---
186
+
187
+ ## License
188
+
189
+ 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.3"
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
  ]
@@ -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
+