Fix warnings in AGS (#5320)

This PR does the following: 

- Fix warning messages in AGS on launch.
- Improve Cli message to include app URL on startup from command line
- Minor improvements default gallery generator. (add more default tools)
- Improve new session behaviour.



## Related issue number

Closes #5097

## Checks

python/packages/autogen-agentchat/src/autogen_agentchat/agents/_assistant_agent.py

@@ -64,7 +64,7 @@ class AssistantAgentConfig(BaseModel):
     model_context: ComponentModel | None = None
     description: str
     system_message: str | None = None
-    model_client_stream: bool
+    model_client_stream: bool = False
     reflect_on_tool_use: bool
     tool_call_summary_format: str
 

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

@@ -1,4 +1,6 @@
 import os
+import tempfile
+import warnings
 from typing import Optional
 
 import typer
@@ -9,6 +11,10 @@
 
 app = typer.Typer()
 
+# Ignore deprecation warnings from websockets
+warnings.filterwarnings("ignore", message="websockets.legacy is deprecated*")
+warnings.filterwarnings("ignore", message="websockets.server.WebSocketServerProtocol is deprecated*")
+
 
 @app.command()
 def ui(
@@ -33,14 +39,25 @@ def ui(
         appdir (str, optional): Path to the AutoGen Studio app directory. Defaults to None.
         database-uri (str, optional): Database URI to connect to. Defaults to None.
     """
-
-    os.environ["AUTOGENSTUDIO_API_DOCS"] = str(docs)
+    # Create temporary env file to share configuration with uvicorn workers
+    temp_env = tempfile.NamedTemporaryFile(mode="w", suffix=".env", delete=True)
+
+    # Write configuration
+    env_vars = {
+        "AUTOGENSTUDIO_HOST": host,
+        "AUTOGENSTUDIO_PORT": port,
+        "AUTOGENSTUDIO_API_DOCS": str(docs),
+    }
     if appdir:
-        os.environ["AUTOGENSTUDIO_APPDIR"] = appdir
+        env_vars["AUTOGENSTUDIO_APPDIR"] = appdir
     if database_uri:
-        os.environ["AUTOGENSTUDIO_DATABASE_URI"] = database_uri
+        env_vars["AUTOGENSTUDIO_DATABASE_URI"] = database_uri
     if upgrade_database:
-        os.environ["AUTOGENSTUDIO_UPGRADE_DATABASE"] = "1"
+        env_vars["AUTOGENSTUDIO_UPGRADE_DATABASE"] = "1"
+
+    for key, value in env_vars.items():
+        temp_env.write(f"{key}={value}\n")
+    temp_env.flush()
 
     uvicorn.run(
         "autogenstudio.web.app:app",
@@ -49,6 +66,7 @@ def ui(
         workers=workers,
         reload=reload,
         reload_excludes=["**/alembic/*", "**/alembic.ini", "**/versions/*"] if reload else None,
+        env_file=temp_env.name,
     )
 
 

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

@@ -1,5 +1,4 @@
 from .db_manager import DatabaseManager
-from .gallery_builder import GalleryBuilder, create_default_gallery
 
 __all__ = [
     "DatabaseManager",

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

@@ -1,5 +1,7 @@
+import io
 import os
 import shutil
+from contextlib import redirect_stdout
 from pathlib import Path
 from typing import List, Optional, Tuple
 
@@ -45,7 +47,7 @@ def __init__(
     def initialize_migrations(self, force: bool = False) -> bool:
         try:
             if force:
-                logger.info("Force reinitialization of migrations...")
+                # logger.info("Force reinitialization of migrations...")
                 self._cleanup_existing_alembic()
                 if not self._initialize_alembic():
                     return False
@@ -60,7 +62,7 @@ def initialize_migrations(self, force: bool = False) -> bool:
                         return False
 
             # Only generate initial revision if alembic is properly initialized
-            logger.info("Creating initial migration...")
+            # logger.info("Creating initial migration...")
             return self.generate_revision("Initial schema") is not None
 
         except Exception as e:
@@ -88,7 +90,7 @@ def _cleanup_existing_alembic(self) -> None:
         Completely remove existing Alembic configuration including versions.
         For fresh initialization, we don't need to preserve anything.
         """
-        logger.info("Cleaning up existing Alembic configuration...")
+        # logger.info("Cleaning up existing Alembic configuration...")
 
         # Remove entire alembic directory if it exists
         if self.alembic_dir.exists():
@@ -130,7 +132,7 @@ def _initialize_alembic(self) -> bool:
             self.alembic_dir.parent.mkdir(exist_ok=True)
 
             # Run alembic init to create fresh directory structure
-            logger.info("Initializing alembic directory structure...")
+            # logger.info("Initializing alembic directory structure...")
 
             # Create initial config file for alembic init
             config_content = self._generate_alembic_ini_content()
@@ -139,7 +141,9 @@ def _initialize_alembic(self) -> bool:
 
             # Use the config we just created
             config = Config(str(self.alembic_ini_path))
-            command.init(config, str(self.alembic_dir))
+
+            with redirect_stdout(io.StringIO()):
+                command.init(config, str(self.alembic_dir))
 
             # Update script template after initialization
             self.update_script_template()
@@ -265,7 +269,6 @@ def update_script_template(self):
             with open(template_path, "w") as f:
                 f.write(content)
 
-            logger.info("Updated script template")
             return True
 
         except Exception as e:
@@ -320,8 +323,6 @@ def _update_env_py(self, env_path: Path) -> None:
 
             with open(env_path, "w") as f:
                 f.write(content)
-
-            logger.info("Updated env.py with SQLModel metadata")
         except Exception as e:
             logger.error(f"Failed to update env.py: {e}")
             raise
@@ -481,8 +482,9 @@ def generate_revision(self, message: str = "auto") -> Optional[str]:
         """
         try:
             config = self.get_alembic_config()
-            command.revision(config, message=message, autogenerate=True)
-            return self.get_head_revision()
+            with redirect_stdout(io.StringIO()):
+                command.revision(config, message=message, autogenerate=True)
+                return self.get_head_revision()
 
         except Exception as e:
             logger.error(f"Failed to generate revision: {str(e)}")

python/packages/autogen-studio/autogenstudio/datamodel/db.py

@@ -6,6 +6,7 @@
 from uuid import UUID, uuid4
 
 from autogen_core import ComponentModel
+from pydantic import ConfigDict
 from sqlalchemy import ForeignKey, Integer
 from sqlmodel import JSON, Column, DateTime, Field, SQLModel, func
 
@@ -101,5 +102,4 @@ class Run(SQLModel, table=True):
     version: Optional[str] = "0.0.1"
     messages: Union[List[Message], List[dict]] = Field(default_factory=list, sa_column=Column(JSON))
 
-    class Config:
-        json_encoders = {UUID: str, datetime: lambda v: v.isoformat()}
+    model_config = ConfigDict(json_encoders={UUID: str, datetime: lambda v: v.isoformat()})

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

@@ -0,0 +1,3 @@
+from .builder import GalleryBuilder, create_default_gallery
+
+__all__ = ["GalleryBuilder", "create_default_gallery"]