Add Session Saving to AGS

python/packages/autogen-studio/autogenstudio/__init__.py

@@ -1,3 +1,18 @@
-from .datamodel import *
+from .database.db_manager import DatabaseManager
+from .datamodel import Agent, AgentConfig, Model, ModelConfig, Team, TeamConfig, Tool, ToolConfig
+from .teammanager import TeamManager
 from .version import __version__
-from .teammanager import *
+
+__all__ = [
+    "Tool",
+    "Model",
+    "DatabaseManager",
+    "Team",
+    "Agent",
+    "ToolConfig",
+    "ModelConfig",
+    "TeamConfig",
+    "AgentConfig",
+    "TeamManager",
+    "__version__",
+]

python/packages/autogen-studio/autogenstudio/cli.py

@@ -15,7 +15,7 @@ def ui(
     host: str = "127.0.0.1",
     port: int = 8081,
     workers: int = 1,
-    reload: Annotated[bool, typer.Option("--reload")] = True,
+    reload: Annotated[bool, typer.Option("--reload")] = False,
     docs: bool = True,
     appdir: str = None,
     database_uri: Optional[str] = None,
@@ -48,11 +48,7 @@ def ui(
         port=port,
         workers=workers,
         reload=reload,
-        reload_excludes=[
-            "**/alembic/*",
-            "**/alembic.ini",
-            "**/versions/*"
-        ] if reload else None
+        reload_excludes=["**/alembic/*", "**/alembic.ini", "**/versions/*"] if reload else None,
     )
 
 

python/packages/autogen-studio/autogenstudio/components/agents/userproxy.py

@@ -1,11 +1,11 @@
-from typing import Callable, List, Optional, Sequence, Union, Awaitable
+import asyncio
 from inspect import iscoroutinefunction
+from typing import Awaitable, Callable, List, Optional, Sequence, Union
 
 from autogen_agentchat.agents import BaseChatAgent
 from autogen_agentchat.base import Response
 from autogen_agentchat.messages import ChatMessage, TextMessage
 from autogen_core.base import CancellationToken
-import asyncio
 
 
 class UserProxyAgent(BaseChatAgent):
@@ -15,13 +15,11 @@ def __init__(
         self,
         name: str,
         description: Optional[str] = "a",
-        input_func: Optional[Union[Callable[..., str],
-                                   Callable[..., Awaitable[str]]]] = None
+        input_func: Optional[Union[Callable[..., str], Callable[..., Awaitable[str]]]] = None,
     ) -> None:
         super().__init__(name, description=description)
         self.input_func = input_func or input
-        self._is_async = iscoroutinefunction(
-            input_func) if input_func else False
+        self._is_async = iscoroutinefunction(input_func) if input_func else False
 
     @property
     def produced_message_types(self) -> List[type[ChatMessage]]:
@@ -35,7 +33,6 @@ async def _get_input(self, prompt: str) -> str:
             return await asyncio.get_event_loop().run_in_executor(None, self.input_func, prompt)
 
     async def on_messages(self, messages: Sequence[ChatMessage], cancellation_token: CancellationToken) -> Response:
-
         try:
             user_input = await self._get_input("Enter your response: ")
             return Response(chat_message=TextMessage(content=user_input, source=self.name))

python/packages/autogen-studio/autogenstudio/database/__init__.py

@@ -1,3 +1,3 @@
-from .db_manager import DatabaseManager
-from .component_factory import ComponentFactory, Component
+from .component_factory import Component, ComponentFactory
 from .config_manager import ConfigurationManager
+from .db_manager import DatabaseManager