From 0e985d4b40443a802455433cbf0140be4c4ab9db Mon Sep 17 00:00:00 2001 From: Victor Dibia Date: Sat, 9 Nov 2024 14:32:24 -0800 Subject: [PATCH] v1 of AutoGen Studio on AgentChat (#4097) * add skeleton worflow manager * add test notebook * update test nb * add sample team spec * refactor requirements to agentchat and ext * add base provider to return agentchat agents from json spec * initial api refactor, update dbmanager * api refactor * refactor tests * ags api tutorial update * ui refactor * general refactor * minor refactor updates * backend api refaactor * ui refactor and update * implement v1 for streaming connection with ui updates * backend refactor * ui refactor * minor ui tweak * minor refactor and tweaks * general refactor * update tests * sync uv.lock with main * uv lock update --- python/packages/autogen-studio/.gitignore | 4 + .../autogen-studio/autogenstudio/__init__.py | 3 +- .../autogenstudio/chatmanager.py | 179 - .../autogen-studio/autogenstudio/cli.py | 12 +- .../autogenstudio/database/__init__.py | 6 +- .../autogenstudio/database/alembic.ini | 116 - .../database/component_factory.py | 355 + .../autogenstudio/database/config_manager.py | 322 + .../autogenstudio/database/db_manager.py | 424 + .../autogenstudio/database/dbmanager.py | 491 - .../autogenstudio/database/migrations/README | 1 - .../autogenstudio/database/migrations/env.py | 80 - .../database/migrations/script.py.mako | 27 - .../autogenstudio/database/schema_manager.py | 505 + .../autogenstudio/database/utils.py | 361 - .../autogen-studio/autogenstudio/datamodel.py | 297 - .../autogenstudio/datamodel/__init__.py | 2 + .../autogenstudio/datamodel/db.py | 282 + .../autogenstudio/datamodel/types.py | 136 + .../autogenstudio/teammanager.py | 67 + .../autogenstudio/utils/dbdefaults.json | 242 - .../autogenstudio/utils/utils.py | 346 +- .../autogen-studio/autogenstudio/web/app.py | 576 +- .../autogenstudio/web/config.py | 18 + .../autogen-studio/autogenstudio/web/deps.py | 201 + .../autogenstudio/web/initialization.py | 110 + .../migrations => web/managers}/__init__.py | 0 .../autogenstudio/web/managers/connection.py | 247 + .../autogenstudio/web/routes/__init__.py | 0 .../autogenstudio/web/routes/agents.py | 181 + .../autogenstudio/web/routes/models.py | 95 + .../autogenstudio/web/routes/runs.py | 76 + .../autogenstudio/web/routes/sessions.py | 114 + .../autogenstudio/web/routes/teams.py | 146 + .../autogenstudio/web/routes/tools.py | 103 + .../autogenstudio/web/routes/ws.py | 74 + .../websocket_connection_manager.py | 135 - .../autogenstudio/workflowmanager.py | 1066 -- .../autogen-studio/frontend/.env.default | 2 +- .../autogen-studio/frontend/.gitignore | 8 +- .../packages/autogen-studio/frontend/LICENSE | 21 - .../autogen-studio/frontend/README.md | 4 +- .../autogen-studio/frontend/gatsby-config.ts | 14 +- .../autogen-studio/frontend/package.json | 70 +- .../autogen-studio/frontend/postcss.config.js | 10 +- .../frontend/src/components/atoms.tsx | 873 - .../frontend/src/components/contentheader.tsx | 157 + .../frontend/src/components/footer.tsx | 2 +- .../frontend/src/components/icons.tsx | 56 +- .../frontend/src/components/layout.tsx | 120 +- .../frontend/src/components/sidebar.tsx | 148 + .../frontend/src/components/types.ts | 127 - .../frontend/src/components/types/app.ts | 5 + .../src/components/types/datamodel.ts | 86 + .../frontend/src/components/utils.ts | 568 +- .../src/components/views/builder/agents.tsx | 385 - .../src/components/views/builder/build.tsx | 81 - .../src/components/views/builder/models.tsx | 403 - .../src/components/views/builder/skills.tsx | 380 - .../views/builder/utils/agentconfig.tsx | 517 - .../components/views/builder/utils/export.tsx | 207 - .../views/builder/utils/modelconfig.tsx | 388 - .../views/builder/utils/selectors.tsx | 1359 -- .../views/builder/utils/skillconfig.tsx | 295 - .../views/builder/utils/workflowconfig.tsx | 279 - .../src/components/views/builder/workflow.tsx | 428 - .../src/components/views/gallery/gallery.tsx | 207 - .../components/views/playground/chat/chat.tsx | 449 + .../views/playground/chat/chatinput.tsx | 128 + .../views/playground/chat/messagelist.tsx | 139 + .../views/playground/chat/thread.tsx | 212 + .../components/views/playground/chat/types.ts | 47 + .../components/views/playground/chatbox.tsx | 824 - .../components/views/playground/metadata.tsx | 242 - .../src/components/views/playground/ra.tsx | 94 - .../components/views/playground/sessions.tsx | 460 - .../components/views/playground/sidebar.tsx | 49 - .../views/playground/utils/charts/bar.tsx | 58 - .../views/playground/utils/profiler.tsx | 125 - .../views/playground/utils/selectors.tsx | 122 - .../src/components/views/shared/markdown.tsx | 20 + .../src/components/views/shared/monaco.tsx | 48 + .../components/views/shared/session/api.ts | 115 + .../views/shared/session/editor.tsx | 157 + .../components/views/shared/session/list.tsx | 76 + .../views/shared/session/manager.tsx | 190 + .../components/views/shared/session/types.ts | 24 + .../src/components/views/shared/team/api.ts | 127 + .../components/views/shared/team/editor.tsx | 171 + .../src/components/views/shared/team/list.tsx | 76 + .../components/views/shared/team/manager.tsx | 153 + .../src/components/views/shared/team/types.ts | 17 + .../frontend/src/hooks/store.tsx | 67 +- .../frontend/src/images/icon.png | 4 +- .../frontend/src/images/landing/nodata.svg | 1 + .../svgs => src/images/landing}/welcome.svg | 2 +- .../autogen-studio/frontend/src/index.d.ts | 3 + .../autogen-studio/frontend/src/pages/404.tsx | 15 +- .../frontend/src/pages/build.tsx | 3 +- .../frontend/src/pages/gallery/index.tsx | 28 - .../frontend/src/pages/index.tsx | 4 +- .../frontend/src/pages/settings.tsx | 35 + .../frontend/src/styles/global.css | 79 +- .../frontend/tailwind.config.js | 7 +- .../autogen-studio/frontend/tsconfig.json | 114 +- .../autogen-studio/frontend/yarn.lock | 13455 ++++++++++++++++ .../autogen-studio/notebooks/team.json | 32 + .../notebooks/travel_groupchat.json | 273 - .../autogen-studio/notebooks/tutorial.ipynb | 225 +- .../autogen-studio/notebooks/two_agent.json | 112 - python/packages/autogen-studio/pyproject.toml | 4 +- .../test/test_save_skills_to_file.py | 56 - .../autogen-studio/test/test_skills_prompt.py | 47 - .../packages/autogen-studio/tests/__init__.py | 0 .../tests/test_component_factory.py | 246 + .../autogen-studio/tests/test_db_manager.py | 183 + python/uv.lock | 586 - 117 files changed, 20720 insertions(+), 13584 deletions(-) delete mode 100644 python/packages/autogen-studio/autogenstudio/chatmanager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/database/alembic.ini create mode 100644 python/packages/autogen-studio/autogenstudio/database/component_factory.py create mode 100644 python/packages/autogen-studio/autogenstudio/database/config_manager.py create mode 100644 python/packages/autogen-studio/autogenstudio/database/db_manager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/database/dbmanager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/database/migrations/README delete mode 100644 python/packages/autogen-studio/autogenstudio/database/migrations/env.py delete mode 100644 python/packages/autogen-studio/autogenstudio/database/migrations/script.py.mako create mode 100644 python/packages/autogen-studio/autogenstudio/database/schema_manager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/database/utils.py delete mode 100644 python/packages/autogen-studio/autogenstudio/datamodel.py create mode 100644 python/packages/autogen-studio/autogenstudio/datamodel/__init__.py create mode 100644 python/packages/autogen-studio/autogenstudio/datamodel/db.py create mode 100644 python/packages/autogen-studio/autogenstudio/datamodel/types.py create mode 100644 python/packages/autogen-studio/autogenstudio/teammanager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/utils/dbdefaults.json create mode 100644 python/packages/autogen-studio/autogenstudio/web/config.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/deps.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/initialization.py rename python/packages/autogen-studio/autogenstudio/{database/migrations => web/managers}/__init__.py (100%) create mode 100644 python/packages/autogen-studio/autogenstudio/web/managers/connection.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/__init__.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/agents.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/models.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/runs.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/sessions.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/teams.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/tools.py create mode 100644 python/packages/autogen-studio/autogenstudio/web/routes/ws.py delete mode 100644 python/packages/autogen-studio/autogenstudio/websocket_connection_manager.py delete mode 100644 python/packages/autogen-studio/autogenstudio/workflowmanager.py delete mode 100644 python/packages/autogen-studio/frontend/LICENSE delete mode 100644 python/packages/autogen-studio/frontend/src/components/atoms.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/contentheader.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/sidebar.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/types.ts create mode 100644 python/packages/autogen-studio/frontend/src/components/types/app.ts create mode 100644 python/packages/autogen-studio/frontend/src/components/types/datamodel.ts delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/agents.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/build.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/models.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/skills.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/agentconfig.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/export.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/modelconfig.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/selectors.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/skillconfig.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/utils/workflowconfig.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/builder/workflow.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/gallery/gallery.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chat/chat.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chat/chatinput.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chat/messagelist.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chat/thread.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chat/types.ts delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/chatbox.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/metadata.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/ra.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/sessions.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/sidebar.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/utils/charts/bar.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/utils/profiler.tsx delete mode 100644 python/packages/autogen-studio/frontend/src/components/views/playground/utils/selectors.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/markdown.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/monaco.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/session/api.ts create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/session/editor.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/session/list.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/session/manager.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/session/types.ts create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/team/api.ts create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/team/editor.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/team/list.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/team/manager.tsx create mode 100644 python/packages/autogen-studio/frontend/src/components/views/shared/team/types.ts create mode 100644 python/packages/autogen-studio/frontend/src/images/landing/nodata.svg rename python/packages/autogen-studio/frontend/{static/images/svgs => src/images/landing}/welcome.svg (98%) create mode 100644 python/packages/autogen-studio/frontend/src/index.d.ts delete mode 100644 python/packages/autogen-studio/frontend/src/pages/gallery/index.tsx create mode 100644 python/packages/autogen-studio/frontend/src/pages/settings.tsx create mode 100644 python/packages/autogen-studio/frontend/yarn.lock create mode 100644 python/packages/autogen-studio/notebooks/team.json delete mode 100644 python/packages/autogen-studio/notebooks/travel_groupchat.json delete mode 100644 python/packages/autogen-studio/notebooks/two_agent.json delete mode 100644 python/packages/autogen-studio/test/test_save_skills_to_file.py delete mode 100644 python/packages/autogen-studio/test/test_skills_prompt.py create mode 100644 python/packages/autogen-studio/tests/__init__.py create mode 100644 python/packages/autogen-studio/tests/test_component_factory.py create mode 100644 python/packages/autogen-studio/tests/test_db_manager.py diff --git a/python/packages/autogen-studio/.gitignore b/python/packages/autogen-studio/.gitignore index 549ce16b6db9..cf5c0a525432 100644 --- a/python/packages/autogen-studio/.gitignore +++ b/python/packages/autogen-studio/.gitignore @@ -2,6 +2,8 @@ database.sqlite .cache/* autogenstudio/web/files/user/* autogenstudio/test +autogenstudio/database/alembic.ini +autogenstudio/database/alembic/* autogenstudio/web/files/ui/* OAI_CONFIG_LIST scratch/ @@ -10,8 +12,10 @@ autogenstudio/web/ui/* autogenstudio/web/skills/user/* .release.sh .nightly.sh +notebooks/test notebooks/work_dir/* +notebooks/test.db # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/python/packages/autogen-studio/autogenstudio/__init__.py b/python/packages/autogen-studio/autogenstudio/__init__.py index acc477c5cd88..833950913331 100644 --- a/python/packages/autogen-studio/autogenstudio/__init__.py +++ b/python/packages/autogen-studio/autogenstudio/__init__.py @@ -1,4 +1,3 @@ -from .chatmanager import * from .datamodel import * from .version import __version__ -from .workflowmanager import * +from .teammanager import * diff --git a/python/packages/autogen-studio/autogenstudio/chatmanager.py b/python/packages/autogen-studio/autogenstudio/chatmanager.py deleted file mode 100644 index e8ed3abfd627..000000000000 --- a/python/packages/autogen-studio/autogenstudio/chatmanager.py +++ /dev/null @@ -1,179 +0,0 @@ -import os -from datetime import datetime -from queue import Queue -from typing import Any, Dict, List, Optional, Tuple, Union - -from loguru import logger - -from .datamodel import Message -from .websocket_connection_manager import WebSocketConnectionManager -from .workflowmanager import WorkflowManager - - -class AutoGenChatManager: - """ - This class handles the automated generation and management of chat interactions - using an automated workflow configuration and message queue. - """ - - def __init__( - self, message_queue: Queue, websocket_manager: WebSocketConnectionManager = None, human_input_timeout: int = 180 - ) -> None: - """ - Initializes the AutoGenChatManager with a message queue. - - :param message_queue: A queue to use for sending messages asynchronously. - """ - self.message_queue = message_queue - self.websocket_manager = websocket_manager - self.a_human_input_timeout = human_input_timeout - - def send(self, message: dict) -> None: - """ - Sends a message by putting it into the message queue. - - :param message: The message string to be sent. - """ - if self.message_queue is not None: - self.message_queue.put_nowait(message) - - async def a_send(self, message: dict) -> None: - """ - Asynchronously sends a message via the WebSocketManager class - - :param message: The message string to be sent. - """ - for connection, socket_client_id in self.websocket_manager.active_connections: - if message["connection_id"] == socket_client_id: - logger.info( - f"Sending message to connection_id: {message['connection_id']}. Connection ID: {socket_client_id}" - ) - await self.websocket_manager.send_message(message, connection) - else: - logger.info( - f"Skipping message for connection_id: {message['connection_id']}. Connection ID: {socket_client_id}" - ) - - async def a_prompt_for_input(self, prompt: dict, timeout: int = 60) -> str: - """ - Sends the user a prompt and waits for a response asynchronously via the WebSocketManager class - - :param message: The message string to be sent. - """ - - for connection, socket_client_id in self.websocket_manager.active_connections: - if prompt["connection_id"] == socket_client_id: - logger.info( - f"Sending message to connection_id: {prompt['connection_id']}. Connection ID: {socket_client_id}" - ) - try: - result = await self.websocket_manager.get_input(prompt, connection, timeout) - return result - except Exception as e: - return f"Error: {e}\nTERMINATE" - else: - logger.info( - f"Skipping message for connection_id: {prompt['connection_id']}. Connection ID: {socket_client_id}" - ) - - def chat( - self, - message: Message, - history: List[Dict[str, Any]], - workflow: Any = None, - connection_id: Optional[str] = None, - user_dir: Optional[str] = None, - **kwargs, - ) -> Message: - """ - Processes an incoming message according to the agent's workflow configuration - and generates a response. - - :param message: An instance of `Message` representing an incoming message. - :param history: A list of dictionaries, each representing a past interaction. - :param flow_config: An instance of `AgentWorkFlowConfig`. If None, defaults to a standard configuration. - :param connection_id: An optional connection identifier. - :param kwargs: Additional keyword arguments. - :return: An instance of `Message` representing a response. - """ - - # create a working director for workflow based on user_dir/session_id/time_hash - work_dir = os.path.join( - user_dir, - str(message.session_id), - datetime.now().strftime("%Y%m%d_%H-%M-%S"), - ) - os.makedirs(work_dir, exist_ok=True) - - # if no flow config is provided, use the default - if workflow is None: - raise ValueError("Workflow must be specified") - - workflow_manager = WorkflowManager( - workflow=workflow, - history=history, - work_dir=work_dir, - send_message_function=self.send, - a_send_message_function=self.a_send, - connection_id=connection_id, - ) - - message_text = message.content.strip() - result_message: Message = workflow_manager.run(message=f"{message_text}", clear_history=False, history=history) - - result_message.user_id = message.user_id - result_message.session_id = message.session_id - return result_message - - async def a_chat( - self, - message: Message, - history: List[Dict[str, Any]], - workflow: Any = None, - connection_id: Optional[str] = None, - user_dir: Optional[str] = None, - **kwargs, - ) -> Message: - """ - Processes an incoming message according to the agent's workflow configuration - and generates a response. - - :param message: An instance of `Message` representing an incoming message. - :param history: A list of dictionaries, each representing a past interaction. - :param flow_config: An instance of `AgentWorkFlowConfig`. If None, defaults to a standard configuration. - :param connection_id: An optional connection identifier. - :param kwargs: Additional keyword arguments. - :return: An instance of `Message` representing a response. - """ - - # create a working director for workflow based on user_dir/session_id/time_hash - work_dir = os.path.join( - user_dir, - str(message.session_id), - datetime.now().strftime("%Y%m%d_%H-%M-%S"), - ) - os.makedirs(work_dir, exist_ok=True) - - # if no flow config is provided, use the default - if workflow is None: - raise ValueError("Workflow must be specified") - - workflow_manager = WorkflowManager( - workflow=workflow, - history=history, - work_dir=work_dir, - send_message_function=self.send, - a_send_message_function=self.a_send, - a_human_input_function=self.a_prompt_for_input, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=connection_id, - ) - - message_text = message.content.strip() - result_message: Message = await workflow_manager.a_run( - message=f"{message_text}", clear_history=False, history=history - ) - - result_message.user_id = message.user_id - result_message.session_id = message.session_id - return result_message diff --git a/python/packages/autogen-studio/autogenstudio/cli.py b/python/packages/autogen-studio/autogenstudio/cli.py index 81fee7991455..b8612f4ad97b 100644 --- a/python/packages/autogen-studio/autogenstudio/cli.py +++ b/python/packages/autogen-studio/autogenstudio/cli.py @@ -15,10 +15,11 @@ def ui( host: str = "127.0.0.1", port: int = 8081, workers: int = 1, - reload: Annotated[bool, typer.Option("--reload")] = False, + reload: Annotated[bool, typer.Option("--reload")] = True, docs: bool = True, appdir: str = None, database_uri: Optional[str] = None, + upgrade_database: bool = False, ): """ Run the AutoGen Studio UI. @@ -30,7 +31,7 @@ def ui( reload (bool, optional): Whether to reload the UI on code changes. Defaults to False. docs (bool, optional): Whether to generate API docs. Defaults to False. 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. Examples include sqlite:///autogenstudio.db, postgresql://user:password@localhost/autogenstudio. + database-uri (str, optional): Database URI to connect to. Defaults to None. """ os.environ["AUTOGENSTUDIO_API_DOCS"] = str(docs) @@ -38,6 +39,8 @@ def ui( os.environ["AUTOGENSTUDIO_APPDIR"] = appdir if database_uri: os.environ["AUTOGENSTUDIO_DATABASE_URI"] = database_uri + if upgrade_database: + os.environ["AUTOGENSTUDIO_UPGRADE_DATABASE"] = "1" uvicorn.run( "autogenstudio.web.app:app", @@ -45,6 +48,11 @@ def ui( port=port, workers=workers, reload=reload, + reload_excludes=[ + "**/alembic/*", + "**/alembic.ini", + "**/versions/*" + ] if reload else None ) diff --git a/python/packages/autogen-studio/autogenstudio/database/__init__.py b/python/packages/autogen-studio/autogenstudio/database/__init__.py index 0518c24ba4fa..ac87c41f0bd7 100644 --- a/python/packages/autogen-studio/autogenstudio/database/__init__.py +++ b/python/packages/autogen-studio/autogenstudio/database/__init__.py @@ -1,3 +1,3 @@ -# from .dbmanager import * -from .dbmanager import * -from .utils import * +from .db_manager import DatabaseManager +from .component_factory import ComponentFactory +from .config_manager import ConfigurationManager diff --git a/python/packages/autogen-studio/autogenstudio/database/alembic.ini b/python/packages/autogen-studio/autogenstudio/database/alembic.ini deleted file mode 100644 index cd413a26066c..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/alembic.ini +++ /dev/null @@ -1,116 +0,0 @@ -# A generic, single database configuration. - -[alembic] -# path to migration scripts -script_location = migrations - -# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s -# Uncomment the line below if you want the files to be prepended with date and time -# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file -# for all available tokens -# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s - -# sys.path path, will be prepended to sys.path if present. -# defaults to the current working directory. -prepend_sys_path = . - -# timezone to use when rendering the date within the migration file -# as well as the filename. -# If specified, requires the python>=3.9 or backports.zoneinfo library. -# Any required deps can installed by adding `alembic[tz]` to the pip requirements -# string value is passed to ZoneInfo() -# leave blank for localtime -# timezone = - -# max length of characters to apply to the -# "slug" field -# truncate_slug_length = 40 - -# set to 'true' to run the environment during -# the 'revision' command, regardless of autogenerate -# revision_environment = false - -# set to 'true' to allow .pyc and .pyo files without -# a source .py file to be detected as revisions in the -# versions/ directory -# sourceless = false - -# version location specification; This defaults -# to migrations/versions. When using multiple version -# directories, initial revisions must be specified with --version-path. -# The path separator used here should be the separator specified by "version_path_separator" below. -# version_locations = %(here)s/bar:%(here)s/bat:migrations/versions - -# version path separator; As mentioned above, this is the character used to split -# version_locations. The default within new alembic.ini files is "os", which uses os.pathsep. -# If this key is omitted entirely, it falls back to the legacy behavior of splitting on spaces and/or commas. -# Valid values for version_path_separator are: -# -# version_path_separator = : -# version_path_separator = ; -# version_path_separator = space -version_path_separator = os # Use os.pathsep. Default configuration used for new projects. - -# set to 'true' to search source files recursively -# in each "version_locations" directory -# new in Alembic version 1.10 -# recursive_version_locations = false - -# the output encoding used when revision files -# are written from script.py.mako -# output_encoding = utf-8 - -sqlalchemy.url = driver://user:pass@localhost/dbname - - -[post_write_hooks] -# post_write_hooks defines scripts or Python functions that are run -# on newly generated revision scripts. See the documentation for further -# detail and examples - -# format using "black" - use the console_scripts runner, against the "black" entrypoint -# hooks = black -# black.type = console_scripts -# black.entrypoint = black -# black.options = -l 79 REVISION_SCRIPT_FILENAME - -# lint with attempts to fix using "ruff" - use the exec runner, execute a binary -# hooks = ruff -# ruff.type = exec -# ruff.executable = %(here)s/.venv/bin/ruff -# ruff.options = --fix REVISION_SCRIPT_FILENAME - -# Logging configuration -[loggers] -keys = root,sqlalchemy,alembic - -[handlers] -keys = console - -[formatters] -keys = generic - -[logger_root] -level = WARN -handlers = console -qualname = - -[logger_sqlalchemy] -level = WARN -handlers = -qualname = sqlalchemy.engine - -[logger_alembic] -level = INFO -handlers = -qualname = alembic - -[handler_console] -class = StreamHandler -args = (sys.stderr,) -level = NOTSET -formatter = generic - -[formatter_generic] -format = %(levelname)-5.5s [%(name)s] %(message)s -datefmt = %H:%M:%S diff --git a/python/packages/autogen-studio/autogenstudio/database/component_factory.py b/python/packages/autogen-studio/autogenstudio/database/component_factory.py new file mode 100644 index 000000000000..708a292d6da4 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/database/component_factory.py @@ -0,0 +1,355 @@ +import os +from pathlib import Path +from typing import List, Literal, Union, Optional, Dict, Any, Type +from datetime import datetime +import json +from autogen_agentchat.task import MaxMessageTermination, TextMentionTermination, StopMessageTermination +import yaml +import logging +from packaging import version + +from ..datamodel import ( + TeamConfig, AgentConfig, ModelConfig, ToolConfig, + TeamTypes, AgentTypes, ModelTypes, ToolTypes, + ComponentType, ComponentConfig, ComponentConfigInput, TerminationConfig, TerminationTypes, Response +) +from autogen_agentchat.agents import AssistantAgent +from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat +from autogen_ext.models import OpenAIChatCompletionClient +from autogen_core.components.tools import FunctionTool + +logger = logging.getLogger(__name__) + +# Type definitions for supported components +TeamComponent = Union[RoundRobinGroupChat, SelectorGroupChat] +AgentComponent = Union[AssistantAgent] # Will grow with more agent types +# Will grow with more model types +ModelComponent = Union[OpenAIChatCompletionClient] +ToolComponent = Union[FunctionTool] # Will grow with more tool types +TerminationComponent = Union[MaxMessageTermination, + StopMessageTermination, TextMentionTermination] + +# Config type definitions + +Component = Union[TeamComponent, AgentComponent, ModelComponent, ToolComponent] + + +ReturnType = Literal['object', 'dict', 'config'] +Component = Union[RoundRobinGroupChat, SelectorGroupChat, + AssistantAgent, OpenAIChatCompletionClient, FunctionTool] + + +class ComponentFactory: + """Creates and manages agent components with versioned configuration loading""" + + SUPPORTED_VERSIONS = { + ComponentType.TEAM: ["1.0.0"], + ComponentType.AGENT: ["1.0.0"], + ComponentType.MODEL: ["1.0.0"], + ComponentType.TOOL: ["1.0.0"], + ComponentType.TERMINATION: ["1.0.0"] + } + + def __init__(self): + self._model_cache: Dict[str, OpenAIChatCompletionClient] = {} + self._tool_cache: Dict[str, FunctionTool] = {} + self._last_cache_clear = datetime.now() + + async def load(self, component: ComponentConfigInput, return_type: ReturnType = 'object') -> Union[Component, dict, ComponentConfig]: + """ + Universal loader for any component type + + Args: + component: Component configuration (file path, dict, or ComponentConfig) + return_type: Type of return value ('object', 'dict', or 'config') + + Returns: + Component instance, config dict, or ComponentConfig based on return_type + + Raises: + ValueError: If component type is unknown or version unsupported + """ + try: + # Load and validate config + if isinstance(component, (str, Path)): + component_dict = await self._load_from_file(component) + config = self._dict_to_config(component_dict) + elif isinstance(component, dict): + config = self._dict_to_config(component) + else: + config = component + + # Validate version + if not self._is_version_supported(config.component_type, config.version): + raise ValueError( + f"Unsupported version {config.version} for " + f"component type {config.component_type}. " + f"Supported versions: {self.SUPPORTED_VERSIONS[config.component_type]}" + ) + + # Return early if dict or config requested + if return_type == 'dict': + return config.model_dump() + elif return_type == 'config': + return config + + # Otherwise create and return component instance + handlers = { + ComponentType.TEAM: self.load_team, + ComponentType.AGENT: self.load_agent, + ComponentType.MODEL: self.load_model, + ComponentType.TOOL: self.load_tool, + ComponentType.TERMINATION: self.load_termination + } + + handler = handlers.get(config.component_type) + if not handler: + raise ValueError( + f"Unknown component type: {config.component_type}") + + return await handler(config) + + except Exception as e: + logger.error(f"Failed to load component: {str(e)}") + raise + + async def load_directory(self, directory: Union[str, Path], check_exists: bool = False, return_type: ReturnType = 'object') -> List[Union[Component, dict, ComponentConfig]]: + """ + Import all component configurations from a directory. + """ + components = [] + try: + directory = Path(directory) + # Using Path.iterdir() instead of os.listdir + for path in list(directory.glob("*")): + if path.suffix.lower().endswith(('.json', '.yaml', '.yml')): + try: + component = await self.load(path, return_type) + components.append(component) + except Exception as e: + logger.info( + f"Failed to load component: {str(e)}, {path}") + + return components + except Exception as e: + logger.info(f"Failed to load directory: {str(e)}") + return components + + def _dict_to_config(self, config_dict: dict) -> ComponentConfig: + """Convert dictionary to appropriate config type based on component_type""" + if "component_type" not in config_dict: + raise ValueError("component_type is required in configuration") + + config_types = { + ComponentType.TEAM: TeamConfig, + ComponentType.AGENT: AgentConfig, + ComponentType.MODEL: ModelConfig, + ComponentType.TOOL: ToolConfig, + ComponentType.TERMINATION: TerminationConfig # Add mapping for termination + } + + component_type = ComponentType(config_dict["component_type"]) + config_class = config_types.get(component_type) + + if not config_class: + raise ValueError(f"Unknown component type: {component_type}") + + return config_class(**config_dict) + + async def load_termination(self, config: TerminationConfig) -> TerminationComponent: + """Create termination condition instance from configuration.""" + try: + if config.termination_type == TerminationTypes.MAX_MESSAGES: + return MaxMessageTermination(max_messages=config.max_messages) + elif config.termination_type == TerminationTypes.STOP_MESSAGE: + return StopMessageTermination() + elif config.termination_type == TerminationTypes.TEXT_MENTION: + if not config.text: + raise ValueError( + "text parameter required for TextMentionTermination") + return TextMentionTermination(text=config.text) + else: + raise ValueError( + f"Unsupported termination type: {config.termination_type}") + except Exception as e: + logger.error(f"Failed to create termination condition: {str(e)}") + raise ValueError( + f"Termination condition creation failed: {str(e)}") + + async def load_team(self, config: TeamConfig) -> TeamComponent: + """Create team instance from configuration.""" + try: + # Load participants (agents) + participants = [] + for participant in config.participants: + agent = await self.load(participant) + participants.append(agent) + + # Load model client if specified + model_client = None + if config.model_client: + model_client = await self.load(config.model_client) + + # Load termination condition if specified + termination = None + if config.termination_condition: + # Now we can use the universal load() method since termination is a proper component + termination = await self.load(config.termination_condition) + + # Create team based on type + if config.team_type == TeamTypes.ROUND_ROBIN: + return RoundRobinGroupChat( + participants=participants, + termination_condition=termination + ) + elif config.team_type == TeamTypes.SELECTOR: + if not model_client: + raise ValueError( + "SelectorGroupChat requires a model_client") + return SelectorGroupChat( + participants=participants, + model_client=model_client, + termination_condition=termination + ) + else: + raise ValueError(f"Unsupported team type: {config.team_type}") + + except Exception as e: + logger.error(f"Failed to create team {config.name}: {str(e)}") + raise ValueError(f"Team creation failed: {str(e)}") + + async def load_agent(self, config: AgentConfig) -> AgentComponent: + """Create agent instance from configuration.""" + try: + # Load model client if specified + model_client = None + if config.model_client: + model_client = await self.load(config.model_client) + system_message = config.system_message if config.system_message else "You are a helpful assistant" + # Load tools if specified + tools = [] + if config.tools: + for tool_config in config.tools: + tool = await self.load(tool_config) + tools.append(tool) + + if config.agent_type == AgentTypes.ASSISTANT: + return AssistantAgent( + name=config.name, + model_client=model_client, + tools=tools, + system_message=system_message + ) + else: + raise ValueError( + f"Unsupported agent type: {config.agent_type}") + + except Exception as e: + logger.error(f"Failed to create agent {config.name}: {str(e)}") + raise ValueError(f"Agent creation failed: {str(e)}") + + async def load_model(self, config: ModelConfig) -> ModelComponent: + """Create model instance from configuration.""" + try: + # Check cache first + cache_key = str(config.model_dump()) + if cache_key in self._model_cache: + logger.debug(f"Using cached model for {config.model}") + return self._model_cache[cache_key] + + if config.model_type == ModelTypes.OPENAI: + model = OpenAIChatCompletionClient( + model=config.model, + api_key=config.api_key, + base_url=config.base_url + ) + self._model_cache[cache_key] = model + return model + else: + raise ValueError( + f"Unsupported model type: {config.model_type}") + + except Exception as e: + logger.error(f"Failed to create model {config.model}: {str(e)}") + raise ValueError(f"Model creation failed: {str(e)}") + + async def load_tool(self, config: ToolConfig) -> ToolComponent: + """Create tool instance from configuration.""" + try: + # Validate required fields + if not all([config.name, config.description, config.content, config.tool_type]): + raise ValueError("Tool configuration missing required fields") + + # Check cache first + cache_key = str(config.model_dump()) + if cache_key in self._tool_cache: + logger.debug(f"Using cached tool '{config.name}'") + return self._tool_cache[cache_key] + + if config.tool_type == ToolTypes.PYTHON_FUNCTION: + tool = FunctionTool( + name=config.name, + description=config.description, + func=self._func_from_string(config.content) + ) + self._tool_cache[cache_key] = tool + return tool + else: + raise ValueError(f"Unsupported tool type: {config.tool_type}") + + except Exception as e: + logger.error(f"Failed to create tool '{config.name}': {str(e)}") + raise + + # Helper methods remain largely the same + async def _load_from_file(self, path: Union[str, Path]) -> dict: + """Load configuration from JSON or YAML file.""" + path = Path(path) + if not path.exists(): + raise FileNotFoundError(f"Config file not found: {path}") + + try: + with open(path) as f: + if path.suffix == '.json': + return json.load(f) + elif path.suffix in ('.yml', '.yaml'): + return yaml.safe_load(f) + else: + raise ValueError(f"Unsupported file format: {path.suffix}") + except Exception as e: + raise ValueError(f"Failed to load file {path}: {str(e)}") + + def _func_from_string(self, content: str) -> callable: + """Convert function string to callable.""" + try: + namespace = {} + exec(content, namespace) + for item in namespace.values(): + if callable(item) and not isinstance(item, type): + return item + raise ValueError("No function found in provided code") + except Exception as e: + raise ValueError(f"Failed to create function: {str(e)}") + + def _is_version_supported(self, component_type: ComponentType, ver: str) -> bool: + """Check if version is supported for component type.""" + try: + v = version.parse(ver) + return ver in self.SUPPORTED_VERSIONS[component_type] + except version.InvalidVersion: + return False + + async def cleanup(self) -> None: + """Cleanup resources and clear caches.""" + for model in self._model_cache.values(): + if hasattr(model, 'cleanup'): + await model.cleanup() + + for tool in self._tool_cache.values(): + if hasattr(tool, 'cleanup'): + await tool.cleanup() + + self._model_cache.clear() + self._tool_cache.clear() + self._last_cache_clear = datetime.now() + logger.info("Cleared all component caches") diff --git a/python/packages/autogen-studio/autogenstudio/database/config_manager.py b/python/packages/autogen-studio/autogenstudio/database/config_manager.py new file mode 100644 index 000000000000..be4fa8ad6912 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/database/config_manager.py @@ -0,0 +1,322 @@ +import logging +from typing import Optional, Union, Dict, Any, List +from pathlib import Path +from loguru import logger +from ..datamodel import ( + Model, Team, Agent, Tool, + Response, ComponentTypes, LinkTypes, + ComponentConfigInput +) + +from .component_factory import ComponentFactory +from .db_manager import DatabaseManager + + +class ConfigurationManager: + """Manages persistence and relationships of components using ComponentFactory for validation""" + + DEFAULT_UNIQUENESS_FIELDS = { + ComponentTypes.MODEL: ['model_type', 'model'], + ComponentTypes.TOOL: ['name'], + ComponentTypes.AGENT: ['agent_type', 'name'], + ComponentTypes.TEAM: ['team_type', 'name'] + } + + def __init__(self, db_manager: DatabaseManager, uniqueness_fields: Dict[ComponentTypes, List[str]] = None): + self.db_manager = db_manager + self.component_factory = ComponentFactory() + self.uniqueness_fields = uniqueness_fields or self.DEFAULT_UNIQUENESS_FIELDS + + async def import_component(self, component_config: ComponentConfigInput, user_id: str, check_exists: bool = False) -> Response: + """ + Import a component configuration, validate it, and store the resulting component. + + Args: + component_config: Configuration for the component (file path, dict, or ComponentConfig) + user_id: User ID to associate with imported component + check_exists: Whether to check for existing components before storing (default: False) + + Returns: + Response containing import results or error + """ + try: + # Get validated config as dict + config = await self.component_factory.load(component_config, return_type='dict') + + # Get component type + component_type = self._determine_component_type(config) + if not component_type: + raise ValueError( + f"Unable to determine component type from config") + + # Check existence if requested + if check_exists: + existing = self._check_exists(component_type, config, user_id) + if existing: + return Response( + message=self._format_exists_message( + component_type, config), + status=True, + data={"id": existing.id} + ) + + # Route to appropriate storage method + if component_type == ComponentTypes.TEAM: + return await self._store_team(config, user_id, check_exists) + elif component_type == ComponentTypes.AGENT: + return await self._store_agent(config, user_id, check_exists) + elif component_type == ComponentTypes.MODEL: + return await self._store_model(config, user_id) + elif component_type == ComponentTypes.TOOL: + return await self._store_tool(config, user_id) + else: + raise ValueError( + f"Unsupported component type: {component_type}") + + except Exception as e: + logger.error(f"Failed to import component: {str(e)}") + return Response(message=str(e), status=False) + + async def import_directory(self, directory: Union[str, Path], user_id: str, check_exists: bool = False) -> Response: + """ + Import all component configurations from a directory. + + Args: + directory: Path to directory containing configuration files + user_id: User ID to associate with imported components + check_exists: Whether to check for existing components before storing (default: False) + + Returns: + Response containing import results for all files + """ + try: + configs = await self.component_factory.load_directory(directory, return_type='dict') + + results = [] + for config in configs: + result = await self.import_component(config, user_id, check_exists) + results.append({ + "component": self._get_component_type(config), + "status": result.status, + "message": result.message, + "id": result.data.get("id") if result.status else None + }) + + return Response( + message="Directory import complete", + status=True, + data=results + ) + + except Exception as e: + logger.error(f"Failed to import directory: {str(e)}") + return Response(message=str(e), status=False) + + async def _store_team(self, config: dict, user_id: str, check_exists: bool = False) -> Response: + """Store team component and manage its relationships with agents""" + try: + # Store the team + team_db = Team( + user_id=user_id, + config=config + ) + team_result = self.db_manager.upsert(team_db) + if not team_result.status: + return team_result + + team_id = team_result.data["id"] + + # Handle participants (agents) + for participant in config.get("participants", []): + if check_exists: + # Check for existing agent + agent_type = self._determine_component_type(participant) + existing_agent = self._check_exists( + agent_type, participant, user_id) + if existing_agent: + # Link existing agent + self.db_manager.link( + LinkTypes.TEAM_AGENT, + team_id, + existing_agent.id + ) + logger.info( + f"Linked existing agent to team: {existing_agent}") + continue + + # Store and link new agent + agent_result = await self._store_agent(participant, user_id, check_exists) + if agent_result.status: + self.db_manager.link( + LinkTypes.TEAM_AGENT, + team_id, + agent_result.data["id"] + ) + + return team_result + + except Exception as e: + logger.error(f"Failed to store team: {str(e)}") + return Response(message=str(e), status=False) + + async def _store_agent(self, config: dict, user_id: str, check_exists: bool = False) -> Response: + """Store agent component and manage its relationships with tools and model""" + try: + # Store the agent + agent_db = Agent( + user_id=user_id, + config=config + ) + agent_result = self.db_manager.upsert(agent_db) + if not agent_result.status: + return agent_result + + agent_id = agent_result.data["id"] + + # Handle model client + if "model_client" in config: + if check_exists: + # Check for existing model + model_type = self._determine_component_type( + config["model_client"]) + existing_model = self._check_exists( + model_type, config["model_client"], user_id) + if existing_model: + # Link existing model + self.db_manager.link( + LinkTypes.AGENT_MODEL, + agent_id, + existing_model.id + ) + logger.info( + f"Linked existing model to agent: {existing_model.config.model_type}") + else: + # Store and link new model + model_result = await self._store_model(config["model_client"], user_id) + if model_result.status: + self.db_manager.link( + LinkTypes.AGENT_MODEL, + agent_id, + model_result.data["id"] + ) + else: + # Store and link new model without checking + model_result = await self._store_model(config["model_client"], user_id) + if model_result.status: + self.db_manager.link( + LinkTypes.AGENT_MODEL, + agent_id, + model_result.data["id"] + ) + + # Handle tools + for tool_config in config.get("tools", []): + if check_exists: + # Check for existing tool + tool_type = self._determine_component_type(tool_config) + existing_tool = self._check_exists( + tool_type, tool_config, user_id) + if existing_tool: + # Link existing tool + self.db_manager.link( + LinkTypes.AGENT_TOOL, + agent_id, + existing_tool.id + ) + logger.info( + f"Linked existing tool to agent: {existing_tool.config.name}") + continue + + # Store and link new tool + tool_result = await self._store_tool(tool_config, user_id) + if tool_result.status: + self.db_manager.link( + LinkTypes.AGENT_TOOL, + agent_id, + tool_result.data["id"] + ) + + return agent_result + + except Exception as e: + logger.error(f"Failed to store agent: {str(e)}") + return Response(message=str(e), status=False) + + async def _store_model(self, config: dict, user_id: str) -> Response: + """Store model component (leaf node - no relationships)""" + try: + model_db = Model( + user_id=user_id, + config=config + ) + return self.db_manager.upsert(model_db) + + except Exception as e: + logger.error(f"Failed to store model: {str(e)}") + return Response(message=str(e), status=False) + + async def _store_tool(self, config: dict, user_id: str) -> Response: + """Store tool component (leaf node - no relationships)""" + try: + tool_db = Tool( + user_id=user_id, + config=config + ) + return self.db_manager.upsert(tool_db) + + except Exception as e: + logger.error(f"Failed to store tool: {str(e)}") + return Response(message=str(e), status=False) + + def _check_exists(self, component_type: ComponentTypes, config: dict, user_id: str) -> Optional[Union[Model, Tool, Agent, Team]]: + """Check if component exists based on configured uniqueness fields.""" + fields = self.uniqueness_fields.get(component_type, []) + if not fields: + return None + + component_class = { + ComponentTypes.MODEL: Model, + ComponentTypes.TOOL: Tool, + ComponentTypes.AGENT: Agent, + ComponentTypes.TEAM: Team + }.get(component_type) + + components = self.db_manager.get( + component_class, {"user_id": user_id}).data + + for component in components: + matches = all( + component.config.get(field) == config.get(field) + for field in fields + ) + if matches: + return component + + return None + + def _format_exists_message(self, component_type: ComponentTypes, config: dict) -> str: + """Format existence message with identifying fields.""" + fields = self.uniqueness_fields.get(component_type, []) + field_values = [f"{field}='{config.get(field)}'" for field in fields] + return f"{component_type.value} with {' and '.join(field_values)} already exists" + + def _determine_component_type(self, config: dict) -> Optional[ComponentTypes]: + """Determine component type from configuration dictionary""" + if "team_type" in config: + return ComponentTypes.TEAM + elif "agent_type" in config: + return ComponentTypes.AGENT + elif "model_type" in config: + return ComponentTypes.MODEL + elif "tool_type" in config: + return ComponentTypes.TOOL + return None + + def _get_component_type(self, config: dict) -> str: + """Helper to get component type string from config""" + component_type = self._determine_component_type(config) + return component_type.value if component_type else "unknown" + + async def cleanup(self): + """Cleanup resources""" + await self.component_factory.cleanup() diff --git a/python/packages/autogen-studio/autogenstudio/database/db_manager.py b/python/packages/autogen-studio/autogenstudio/database/db_manager.py new file mode 100644 index 000000000000..b1808f245e2f --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/database/db_manager.py @@ -0,0 +1,424 @@ +import threading +from datetime import datetime +from typing import Optional + +from loguru import logger +from sqlalchemy import exc, text, func +from sqlmodel import Session, SQLModel, and_, create_engine, select +from .schema_manager import SchemaManager + +from ..datamodel import ( + Response, + LinkTypes +) +# from .dbutils import init_db_samples + + +class DatabaseManager: + """A class to manage database operations""" + + _init_lock = threading.Lock() + + def __init__(self, engine_uri: str, auto_upgrade: bool = True): + connection_args = { + "check_same_thread": True} if "sqlite" in engine_uri else {} + self.engine = create_engine(engine_uri, connect_args=connection_args) + self.schema_manager = SchemaManager( + engine=self.engine, + auto_upgrade=auto_upgrade, + ) + + # Check and upgrade on startup + upgraded, status = self.schema_manager.check_and_upgrade() + if upgraded: + logger.info("Database schema was upgraded automatically") + else: + logger.info(f"Schema status: {status}") + + def reset_db(self, recreate_tables: bool = True): + """ + Reset the database by dropping all tables and optionally recreating them. + + Args: + recreate_tables (bool): If True, recreates the tables after dropping them. + Set to False if you want to call create_db_and_tables() separately. + """ + if not self._init_lock.acquire(blocking=False): + logger.warning("Database reset already in progress") + return Response( + message="Database reset already in progress", + status=False, + data=None + ) + + try: + # Dispose existing connections + self.engine.dispose() + with Session(self.engine) as session: + try: + # Disable foreign key checks for SQLite + if 'sqlite' in str(self.engine.url): + session.exec(text('PRAGMA foreign_keys=OFF')) + + # Drop all tables + SQLModel.metadata.drop_all(self.engine) + logger.info("All tables dropped successfully") + + # Re-enable foreign key checks for SQLite + if 'sqlite' in str(self.engine.url): + session.exec(text('PRAGMA foreign_keys=ON')) + + session.commit() + + except Exception as e: + session.rollback() + raise e + finally: + session.close() + self._init_lock.release() + + if recreate_tables: + logger.info("Recreating tables...") + self.create_db_and_tables() + + return Response( + message="Database reset successfully" if recreate_tables else "Database tables dropped successfully", + status=True, + data=None + ) + + except Exception as e: + error_msg = f"Error while resetting database: {str(e)}" + logger.error(error_msg) + return Response( + message=error_msg, + status=False, + data=None + ) + finally: + if self._init_lock.locked(): + self._init_lock.release() + logger.info("Database reset lock released") + + def create_db_and_tables(self): + """Create a new database and tables""" + with self._init_lock: + try: + SQLModel.metadata.create_all(self.engine) + logger.info("Database tables created successfully") + try: + # init_db_samples(self) + pass + except Exception as e: + logger.info( + "Error while initializing database samples: " + str(e)) + except Exception as e: + logger.info("Error while creating database tables:" + str(e)) + + def upsert(self, model: SQLModel, return_json: bool = True): + """Create or update an entity + + Args: + model (SQLModel): The model instance to create or update + return_json (bool, optional): If True, returns the model as a dictionary. + If False, returns the SQLModel instance. Defaults to True. + + Returns: + Response: Contains status, message and data (either dict or SQLModel based on return_json) + """ + status = True + model_class = type(model) + existing_model = None + + with Session(self.engine) as session: + try: + existing_model = session.exec( + select(model_class).where(model_class.id == model.id)).first() + if existing_model: + model.updated_at = datetime.now() + for key, value in model.model_dump().items(): + setattr(existing_model, key, value) + model = existing_model # Use the updated existing model + session.add(model) + else: + session.add(model) + session.commit() + session.refresh(model) + except Exception as e: + session.rollback() + logger.error("Error while updating/creating " + + str(model_class.__name__) + ": " + str(e)) + status = False + + return Response( + message=( + f"{model_class.__name__} Updated Successfully" + if existing_model + else f"{model_class.__name__} Created Successfully" + ), + status=status, + data=model.model_dump() if return_json else model, + ) + + def _model_to_dict(self, model_obj): + return {col.name: getattr(model_obj, col.name) for col in model_obj.__table__.columns} + + def get( + self, + model_class: SQLModel, + filters: dict = None, + return_json: bool = False, + order: str = "desc", + ): + """List entities""" + with Session(self.engine) as session: + result = [] + status = True + status_message = "" + + try: + statement = select(model_class) + if filters: + conditions = [getattr(model_class, col) == + value for col, value in filters.items()] + statement = statement.where(and_(*conditions)) + + if hasattr(model_class, "created_at") and order: + order_by_clause = getattr( + model_class.created_at, order)() # Dynamically apply asc/desc + statement = statement.order_by(order_by_clause) + + items = session.exec(statement).all() + result = [self._model_to_dict( + item) if return_json else item for item in items] + status_message = f"{model_class.__name__} Retrieved Successfully" + except Exception as e: + session.rollback() + status = False + status_message = f"Error while fetching {model_class.__name__}" + logger.error("Error while getting items: " + + str(model_class.__name__) + " " + str(e)) + + return Response(message=status_message, status=status, data=result) + + def delete(self, model_class: SQLModel, filters: dict = None): + """Delete an entity""" + status_message = "" + status = True + + with Session(self.engine) as session: + try: + statement = select(model_class) + if filters: + conditions = [ + getattr(model_class, col) == value for col, value in filters.items()] + statement = statement.where(and_(*conditions)) + + rows = session.exec(statement).all() + + if rows: + for row in rows: + session.delete(row) + session.commit() + status_message = f"{model_class.__name__} Deleted Successfully" + else: + status_message = "Row not found" + logger.info(f"Row with filters {filters} not found") + + except exc.IntegrityError as e: + session.rollback() + status = False + status_message = f"Integrity error: The {model_class.__name__} is linked to another entity and cannot be deleted. {e}" + # Log the specific integrity error + logger.error(status_message) + except Exception as e: + session.rollback() + status = False + status_message = f"Error while deleting: {e}" + logger.error(status_message) + + return Response(message=status_message, status=status, data=None) + + def link( + self, + link_type: LinkTypes, + primary_id: int, + secondary_id: int, + sequence: Optional[int] = None, + ): + """Link two entities with automatic sequence handling.""" + with Session(self.engine) as session: + try: + # Get classes from LinkTypes + primary_class = link_type.primary_class + secondary_class = link_type.secondary_class + link_table = link_type.link_table + + # Get entities + primary_entity = session.get(primary_class, primary_id) + secondary_entity = session.get(secondary_class, secondary_id) + + if not primary_entity or not secondary_entity: + return Response(message="One or both entities do not exist", status=False) + + # Get field names + primary_id_field = f"{primary_class.__name__.lower()}_id" + secondary_id_field = f"{secondary_class.__name__.lower()}_id" + + # Check for existing link + existing_link = session.exec( + select(link_table).where( + and_( + getattr(link_table, primary_id_field) == primary_id, + getattr( + link_table, secondary_id_field) == secondary_id + ) + ) + ).first() + + if existing_link: + return Response(message="Link already exists", status=False) + + # Get the next sequence number if not provided + if sequence is None: + max_seq_result = session.exec( + select(func.max(link_table.sequence)).where( + getattr(link_table, primary_id_field) == primary_id + ) + ).first() + sequence = 0 if max_seq_result is None else max_seq_result + 1 + + # Create new link + new_link = link_table(**{ + primary_id_field: primary_id, + secondary_id_field: secondary_id, + 'sequence': sequence + }) + session.add(new_link) + session.commit() + + return Response( + message=f"Entities linked successfully with sequence {sequence}", + status=True + ) + + except Exception as e: + session.rollback() + return Response(message=f"Error linking entities: {str(e)}", status=False) + + def unlink( + self, + link_type: LinkTypes, + primary_id: int, + secondary_id: int, + sequence: Optional[int] = None + ): + """Unlink two entities and reorder sequences if needed.""" + with Session(self.engine) as session: + try: + # Get classes from LinkTypes + primary_class = link_type.primary_class + secondary_class = link_type.secondary_class + link_table = link_type.link_table + + # Get field names + primary_id_field = f"{primary_class.__name__.lower()}_id" + secondary_id_field = f"{secondary_class.__name__.lower()}_id" + + # Find existing link + statement = select(link_table).where( + and_( + getattr(link_table, primary_id_field) == primary_id, + getattr(link_table, secondary_id_field) == secondary_id + ) + ) + + if sequence is not None: + statement = statement.where( + link_table.sequence == sequence) + + existing_link = session.exec(statement).first() + + if not existing_link: + return Response(message="Link does not exist", status=False) + + deleted_sequence = existing_link.sequence + session.delete(existing_link) + + # Reorder sequences for remaining links + remaining_links = session.exec( + select(link_table) + .where(getattr(link_table, primary_id_field) == primary_id) + .where(link_table.sequence > deleted_sequence) + .order_by(link_table.sequence) + ).all() + + # Decrease sequence numbers to fill the gap + for link in remaining_links: + link.sequence -= 1 + + session.commit() + + return Response( + message="Entities unlinked successfully and sequences reordered", + status=True + ) + + except Exception as e: + session.rollback() + return Response(message=f"Error unlinking entities: {str(e)}", status=False) + + def get_linked_entities( + self, + link_type: LinkTypes, + primary_id: int, + return_json: bool = False, + ): + """Get linked entities based on link type and primary ID, ordered by sequence.""" + with Session(self.engine) as session: + try: + # Get classes from LinkTypes + primary_class = link_type.primary_class + secondary_class = link_type.secondary_class + link_table = link_type.link_table + + # Get field names + primary_id_field = f"{primary_class.__name__.lower()}_id" + secondary_id_field = f"{secondary_class.__name__.lower()}_id" + + # Query both link and entity, ordered by sequence + items = session.exec( + select(secondary_class) + .join(link_table, getattr(link_table, secondary_id_field) == secondary_class.id) + .where(getattr(link_table, primary_id_field) == primary_id) + .order_by(link_table.sequence) + ).all() + + result = [ + item.model_dump() if return_json else item for item in items] + + return Response( + message="Linked entities retrieved successfully", + status=True, + data=result + ) + + except Exception as e: + logger.error(f"Error getting linked entities: {str(e)}") + return Response( + message=f"Error getting linked entities: {str(e)}", + status=False, + data=[] + ) + # Add new close method + + async def close(self): + """Close database connections and cleanup resources""" + logger.info("Closing database connections...") + try: + # Dispose of the SQLAlchemy engine + self.engine.dispose() + logger.info("Database connections closed successfully") + except Exception as e: + logger.error(f"Error closing database connections: {str(e)}") + raise diff --git a/python/packages/autogen-studio/autogenstudio/database/dbmanager.py b/python/packages/autogen-studio/autogenstudio/database/dbmanager.py deleted file mode 100644 index 6a02a0a7038c..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/dbmanager.py +++ /dev/null @@ -1,491 +0,0 @@ -import threading -from datetime import datetime -from typing import Optional - -from loguru import logger -from sqlalchemy import exc -from sqlmodel import Session, SQLModel, and_, create_engine, select - -from ..datamodel import ( - Agent, - AgentLink, - AgentModelLink, - AgentSkillLink, - Model, - Response, - Skill, - Workflow, - WorkflowAgentLink, - WorkflowAgentType, -) -from .utils import init_db_samples - -valid_link_types = ["agent_model", "agent_skill", "agent_agent", "workflow_agent"] - - -class WorkflowAgentMap(SQLModel): - agent: Agent - link: WorkflowAgentLink - - -class DBManager: - """A class to manage database operations""" - - _init_lock = threading.Lock() # Class-level lock - - def __init__(self, engine_uri: str): - connection_args = {"check_same_thread": True} if "sqlite" in engine_uri else {} - self.engine = create_engine(engine_uri, connect_args=connection_args) - # run_migration(engine_uri=engine_uri) - - def create_db_and_tables(self): - """Create a new database and tables""" - with self._init_lock: # Use the lock - try: - SQLModel.metadata.create_all(self.engine) - try: - init_db_samples(self) - except Exception as e: - logger.info("Error while initializing database samples: " + str(e)) - except Exception as e: - logger.info("Error while creating database tables:" + str(e)) - - def upsert(self, model: SQLModel): - """Create a new entity""" - # check if the model exists, update else add - status = True - model_class = type(model) - existing_model = None - - with Session(self.engine) as session: - try: - existing_model = session.exec(select(model_class).where(model_class.id == model.id)).first() - if existing_model: - model.updated_at = datetime.now() - for key, value in model.model_dump().items(): - setattr(existing_model, key, value) - model = existing_model - session.add(model) - else: - session.add(model) - session.commit() - session.refresh(model) - except Exception as e: - session.rollback() - logger.error("Error while updating " + str(model_class.__name__) + ": " + str(e)) - status = False - - response = Response( - message=( - f"{model_class.__name__} Updated Successfully " - if existing_model - else f"{model_class.__name__} Created Successfully" - ), - status=status, - data=model.model_dump(), - ) - - return response - - def _model_to_dict(self, model_obj): - return {col.name: getattr(model_obj, col.name) for col in model_obj.__table__.columns} - - def get_items( - self, - model_class: SQLModel, - session: Session, - filters: dict = None, - return_json: bool = False, - order: str = "desc", - ): - """List all entities""" - result = [] - status = True - status_message = "" - - try: - if filters: - conditions = [getattr(model_class, col) == value for col, value in filters.items()] - statement = select(model_class).where(and_(*conditions)) - - if hasattr(model_class, "created_at") and order: - if order == "desc": - statement = statement.order_by(model_class.created_at.desc()) - else: - statement = statement.order_by(model_class.created_at.asc()) - else: - statement = select(model_class) - - if return_json: - result = [self._model_to_dict(row) for row in session.exec(statement).all()] - else: - result = session.exec(statement).all() - status_message = f"{model_class.__name__} Retrieved Successfully" - except Exception as e: - session.rollback() - status = False - status_message = f"Error while fetching {model_class.__name__}" - logger.error("Error while getting items: " + str(model_class.__name__) + " " + str(e)) - - response: Response = Response( - message=status_message, - status=status, - data=result, - ) - return response - - def get( - self, - model_class: SQLModel, - filters: dict = None, - return_json: bool = False, - order: str = "desc", - ): - """List all entities""" - - with Session(self.engine) as session: - response = self.get_items(model_class, session, filters, return_json, order) - return response - - def delete(self, model_class: SQLModel, filters: dict = None): - """Delete an entity""" - row = None - status_message = "" - status = True - - with Session(self.engine) as session: - try: - if filters: - conditions = [getattr(model_class, col) == value for col, value in filters.items()] - row = session.exec(select(model_class).where(and_(*conditions))).all() - else: - row = session.exec(select(model_class)).all() - if row: - for row in row: - session.delete(row) - session.commit() - status_message = f"{model_class.__name__} Deleted Successfully" - else: - print(f"Row with filters {filters} not found") - logger.info("Row with filters + filters + not found") - status_message = "Row not found" - except exc.IntegrityError as e: - session.rollback() - logger.error("Integrity ... Error while deleting: " + str(e)) - status_message = f"The {model_class.__name__} is linked to another entity and cannot be deleted." - status = False - except Exception as e: - session.rollback() - logger.error("Error while deleting: " + str(e)) - status_message = f"Error while deleting: {e}" - status = False - response = Response( - message=status_message, - status=status, - data=None, - ) - return response - - def get_linked_entities( - self, - link_type: str, - primary_id: int, - return_json: bool = False, - agent_type: Optional[str] = None, - sequence_id: Optional[int] = None, - ): - """ - Get all entities linked to the primary entity. - - Args: - link_type (str): The type of link to retrieve, e.g., "agent_model". - primary_id (int): The identifier for the primary model. - return_json (bool): Whether to return the result as a JSON object. - - Returns: - List[SQLModel]: A list of linked entities. - """ - - linked_entities = [] - - if link_type not in valid_link_types: - return [] - - status = True - status_message = "" - - with Session(self.engine) as session: - try: - if link_type == "agent_model": - # get the agent - agent = self.get_items(Agent, filters={"id": primary_id}, session=session).data[0] - linked_entities = agent.models - elif link_type == "agent_skill": - agent = self.get_items(Agent, filters={"id": primary_id}, session=session).data[0] - linked_entities = agent.skills - elif link_type == "agent_agent": - agent = self.get_items(Agent, filters={"id": primary_id}, session=session).data[0] - linked_entities = agent.agents - elif link_type == "workflow_agent": - linked_entities = session.exec( - select(WorkflowAgentLink, Agent) - .join(Agent, WorkflowAgentLink.agent_id == Agent.id) - .where( - WorkflowAgentLink.workflow_id == primary_id, - ) - ).all() - - linked_entities = [WorkflowAgentMap(agent=agent, link=link) for link, agent in linked_entities] - linked_entities = sorted(linked_entities, key=lambda x: x.link.sequence_id) # type: ignore - except Exception as e: - logger.error("Error while getting linked entities: " + str(e)) - status_message = f"Error while getting linked entities: {e}" - status = False - if return_json: - linked_entities = [row.model_dump() for row in linked_entities] - - response = Response( - message=status_message, - status=status, - data=linked_entities, - ) - - return response - - def link( - self, - link_type: str, - primary_id: int, - secondary_id: int, - agent_type: Optional[str] = None, - sequence_id: Optional[int] = None, - ) -> Response: - """ - Link two entities together. - - Args: - link_type (str): The type of link to create, e.g., "agent_model". - primary_id (int): The identifier for the primary model. - secondary_id (int): The identifier for the secondary model. - agent_type (Optional[str]): The type of agent, e.g., "sender" or receiver. - - Returns: - Response: The response of the linking operation, including success status and message. - """ - - # TBD verify that is creator of the primary entity being linked - status = True - status_message = "" - primary_model = None - secondary_model = None - - if link_type not in valid_link_types: - status = False - status_message = f"Invalid link type: {link_type}. Valid link types are: {valid_link_types}" - else: - with Session(self.engine) as session: - try: - if link_type == "agent_model": - primary_model = session.exec(select(Agent).where(Agent.id == primary_id)).first() - secondary_model = session.exec(select(Model).where(Model.id == secondary_id)).first() - if primary_model is None or secondary_model is None: - status = False - status_message = "One or both entity records do not exist." - else: - # check if the link already exists - existing_link = session.exec( - select(AgentModelLink).where( - AgentModelLink.agent_id == primary_id, - AgentModelLink.model_id == secondary_id, - ) - ).first() - if existing_link: # link already exists - return Response( - message=( - f"{secondary_model.__class__.__name__} already linked " - f"to {primary_model.__class__.__name__}" - ), - status=False, - ) - else: - primary_model.models.append(secondary_model) - elif link_type == "agent_agent": - primary_model = session.exec(select(Agent).where(Agent.id == primary_id)).first() - secondary_model = session.exec(select(Agent).where(Agent.id == secondary_id)).first() - if primary_model is None or secondary_model is None: - status = False - status_message = "One or both entity records do not exist." - else: - # check if the link already exists - existing_link = session.exec( - select(AgentLink).where( - AgentLink.parent_id == primary_id, - AgentLink.agent_id == secondary_id, - ) - ).first() - if existing_link: - return Response( - message=( - f"{secondary_model.__class__.__name__} already linked " - f"to {primary_model.__class__.__name__}" - ), - status=False, - ) - else: - primary_model.agents.append(secondary_model) - - elif link_type == "agent_skill": - primary_model = session.exec(select(Agent).where(Agent.id == primary_id)).first() - secondary_model = session.exec(select(Skill).where(Skill.id == secondary_id)).first() - if primary_model is None or secondary_model is None: - status = False - status_message = "One or both entity records do not exist." - else: - # check if the link already exists - existing_link = session.exec( - select(AgentSkillLink).where( - AgentSkillLink.agent_id == primary_id, - AgentSkillLink.skill_id == secondary_id, - ) - ).first() - if existing_link: - return Response( - message=( - f"{secondary_model.__class__.__name__} already linked " - f"to {primary_model.__class__.__name__}" - ), - status=False, - ) - else: - primary_model.skills.append(secondary_model) - elif link_type == "workflow_agent": - primary_model = session.exec(select(Workflow).where(Workflow.id == primary_id)).first() - secondary_model = session.exec(select(Agent).where(Agent.id == secondary_id)).first() - if primary_model is None or secondary_model is None: - status = False - status_message = "One or both entity records do not exist." - else: - # check if the link already exists - existing_link = session.exec( - select(WorkflowAgentLink).where( - WorkflowAgentLink.workflow_id == primary_id, - WorkflowAgentLink.agent_id == secondary_id, - WorkflowAgentLink.agent_type == agent_type, - WorkflowAgentLink.sequence_id == sequence_id, - ) - ).first() - if existing_link: - return Response( - message=( - f"{secondary_model.__class__.__name__} already linked " - f"to {primary_model.__class__.__name__}" - ), - status=False, - ) - else: - # primary_model.agents.append(secondary_model) - workflow_agent_link = WorkflowAgentLink( - workflow_id=primary_id, - agent_id=secondary_id, - agent_type=agent_type, - sequence_id=sequence_id, - ) - session.add(workflow_agent_link) - # add and commit the link - session.add(primary_model) - session.commit() - status_message = ( - f"{secondary_model.__class__.__name__} successfully linked " - f"to {primary_model.__class__.__name__}" - ) - - except Exception as e: - session.rollback() - logger.error("Error while linking: " + str(e)) - status = False - status_message = f"Error while linking due to an exception: {e}" - - response = Response( - message=status_message, - status=status, - ) - - return response - - def unlink( - self, - link_type: str, - primary_id: int, - secondary_id: int, - agent_type: Optional[str] = None, - sequence_id: Optional[int] = 0, - ) -> Response: - """ - Unlink two entities. - - Args: - link_type (str): The type of link to remove, e.g., "agent_model". - primary_id (int): The identifier for the primary model. - secondary_id (int): The identifier for the secondary model. - agent_type (Optional[str]): The type of agent, e.g., "sender" or receiver. - - Returns: - Response: The response of the unlinking operation, including success status and message. - """ - status = True - status_message = "" - print("primary", primary_id, "secondary", secondary_id, "sequence", sequence_id, "agent_type", agent_type) - - if link_type not in valid_link_types: - status = False - status_message = f"Invalid link type: {link_type}. Valid link types are: {valid_link_types}" - return Response(message=status_message, status=status) - - with Session(self.engine) as session: - try: - if link_type == "agent_model": - existing_link = session.exec( - select(AgentModelLink).where( - AgentModelLink.agent_id == primary_id, - AgentModelLink.model_id == secondary_id, - ) - ).first() - elif link_type == "agent_skill": - existing_link = session.exec( - select(AgentSkillLink).where( - AgentSkillLink.agent_id == primary_id, - AgentSkillLink.skill_id == secondary_id, - ) - ).first() - elif link_type == "agent_agent": - existing_link = session.exec( - select(AgentLink).where( - AgentLink.parent_id == primary_id, - AgentLink.agent_id == secondary_id, - ) - ).first() - elif link_type == "workflow_agent": - existing_link = session.exec( - select(WorkflowAgentLink).where( - WorkflowAgentLink.workflow_id == primary_id, - WorkflowAgentLink.agent_id == secondary_id, - WorkflowAgentLink.agent_type == agent_type, - WorkflowAgentLink.sequence_id == sequence_id, - ) - ).first() - - if existing_link: - session.delete(existing_link) - session.commit() - status_message = "Link removed successfully." - else: - status = False - status_message = "Link does not exist." - - except Exception as e: - session.rollback() - logger.error("Error while unlinking: " + str(e)) - status = False - status_message = f"Error while unlinking due to an exception: {e}" - - return Response(message=status_message, status=status) diff --git a/python/packages/autogen-studio/autogenstudio/database/migrations/README b/python/packages/autogen-studio/autogenstudio/database/migrations/README deleted file mode 100644 index 2500aa1bcf72..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/migrations/README +++ /dev/null @@ -1 +0,0 @@ -Generic single-database configuration. diff --git a/python/packages/autogen-studio/autogenstudio/database/migrations/env.py b/python/packages/autogen-studio/autogenstudio/database/migrations/env.py deleted file mode 100644 index 1431492ad910..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/migrations/env.py +++ /dev/null @@ -1,80 +0,0 @@ -import os -from logging.config import fileConfig - -from alembic import context -from sqlalchemy import engine_from_config, pool -from sqlmodel import SQLModel - -from autogenstudio.datamodel import * -from autogenstudio.utils import get_db_uri - -# this is the Alembic Config object, which provides -# access to the values within the .ini file in use. -config = context.config -config.set_main_option("sqlalchemy.url", get_db_uri()) - -# Interpret the config file for Python logging. -# This line sets up loggers basically. -if config.config_file_name is not None: - fileConfig(config.config_file_name) - -# add your model's MetaData object here -# for 'autogenerate' support -# from myapp import mymodel -# target_metadata = mymodel.Base.metadata -target_metadata = SQLModel.metadata - -# other values from the config, defined by the needs of env.py, -# can be acquired: -# my_important_option = config.get_main_option("my_important_option") -# ... etc. - - -def run_migrations_offline() -> None: - """Run migrations in 'offline' mode. - - This configures the context with just a URL - and not an Engine, though an Engine is acceptable - here as well. By skipping the Engine creation - we don't even need a DBAPI to be available. - - Calls to context.execute() here emit the given string to the - script output. - - """ - url = config.get_main_option("sqlalchemy.url") - context.configure( - url=url, - target_metadata=target_metadata, - literal_binds=True, - dialect_opts={"paramstyle": "named"}, - ) - - with context.begin_transaction(): - context.run_migrations() - - -def run_migrations_online() -> None: - """Run migrations in 'online' mode. - - In this scenario we need to create an Engine - and associate a connection with the context. - - """ - connectable = engine_from_config( - config.get_section(config.config_ini_section, {}), - prefix="sqlalchemy.", - poolclass=pool.NullPool, - ) - - with connectable.connect() as connection: - context.configure(connection=connection, target_metadata=target_metadata) - - with context.begin_transaction(): - context.run_migrations() - - -if context.is_offline_mode(): - run_migrations_offline() -else: - run_migrations_online() diff --git a/python/packages/autogen-studio/autogenstudio/database/migrations/script.py.mako b/python/packages/autogen-studio/autogenstudio/database/migrations/script.py.mako deleted file mode 100644 index 6ce3351093cf..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/migrations/script.py.mako +++ /dev/null @@ -1,27 +0,0 @@ -"""${message} - -Revision ID: ${up_revision} -Revises: ${down_revision | comma,n} -Create Date: ${create_date} - -""" -from typing import Sequence, Union - -from alembic import op -import sqlalchemy as sa -import sqlmodel -${imports if imports else ""} - -# revision identifiers, used by Alembic. -revision: str = ${repr(up_revision)} -down_revision: Union[str, None] = ${repr(down_revision)} -branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)} -depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)} - - -def upgrade() -> None: - ${upgrades if upgrades else "pass"} - - -def downgrade() -> None: - ${downgrades if downgrades else "pass"} diff --git a/python/packages/autogen-studio/autogenstudio/database/schema_manager.py b/python/packages/autogen-studio/autogenstudio/database/schema_manager.py new file mode 100644 index 000000000000..450e0a5d76a3 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/database/schema_manager.py @@ -0,0 +1,505 @@ +import os +from pathlib import Path +import shutil +from typing import Optional, Tuple, List +from loguru import logger +from alembic import command +from alembic.config import Config +from alembic.runtime.migration import MigrationContext +from alembic.script import ScriptDirectory +from alembic.autogenerate import compare_metadata +from sqlalchemy import Engine +from sqlmodel import SQLModel + + +class SchemaManager: + """ + Manages database schema validation and migrations using Alembic. + Provides automatic schema validation, migrations, and safe upgrades. + + Args: + engine: SQLAlchemy engine instance + auto_upgrade: Whether to automatically upgrade schema when differences found + init_mode: Controls initialization behavior: + - "none": No automatic initialization (raises error if not set up) + - "auto": Initialize if not present (default) + - "force": Always reinitialize, removing existing configuration + """ + + def __init__( + self, + engine: Engine, + auto_upgrade: bool = True, + init_mode: str = "auto" + ): + if init_mode not in ["none", "auto", "force"]: + raise ValueError("init_mode must be one of: none, auto, force") + + self.engine = engine + self.auto_upgrade = auto_upgrade + + # Set up paths relative to this file + self.base_dir = Path(__file__).parent + self.alembic_dir = self.base_dir / 'alembic' + self.alembic_ini_path = self.base_dir / 'alembic.ini' + + # Handle initialization based on mode + if init_mode == "none": + self._validate_alembic_setup() + else: + self._ensure_alembic_setup(force=init_mode == "force") + + def _cleanup_existing_alembic(self) -> None: + """ + Safely removes existing Alembic configuration while preserving versions directory. + """ + logger.info( + "Cleaning up existing Alembic configuration while preserving versions...") + + # Create a backup of versions directory if it exists + if self.alembic_dir.exists() and (self.alembic_dir / 'versions').exists(): + logger.info("Preserving existing versions directory") + + # Remove alembic directory contents EXCEPT versions + if self.alembic_dir.exists(): + for item in self.alembic_dir.iterdir(): + if item.name != 'versions': + try: + if item.is_dir(): + shutil.rmtree(item) + logger.info(f"Removed directory: {item}") + else: + item.unlink() + logger.info(f"Removed file: {item}") + except Exception as e: + logger.error(f"Failed to remove {item}: {e}") + + # Remove alembic.ini if it exists + if self.alembic_ini_path.exists(): + try: + self.alembic_ini_path.unlink() + logger.info( + f"Removed existing alembic.ini: {self.alembic_ini_path}") + except Exception as e: + logger.error(f"Failed to remove alembic.ini: {e}") + + def _ensure_alembic_setup(self, *, force: bool = False) -> None: + """ + Ensures Alembic is properly set up, initializing if necessary. + + Args: + force: If True, removes existing configuration and reinitializes + """ + try: + self._validate_alembic_setup() + if force: + logger.info( + "Force initialization requested. Cleaning up existing configuration...") + self._cleanup_existing_alembic() + self._initialize_alembic() + except FileNotFoundError: + logger.info("Alembic configuration not found. Initializing...") + if self.alembic_dir.exists(): + logger.warning( + "Found existing alembic directory but missing configuration") + self._cleanup_existing_alembic() + self._initialize_alembic() + logger.info("Alembic initialization complete") + + def _initialize_alembic(self) -> str: + """Initializes Alembic configuration in the local directory.""" + logger.info("Initializing Alembic configuration...") + + # Check if versions exists + has_versions = (self.alembic_dir / 'versions').exists() + logger.info(f"Existing versions directory found: {has_versions}") + + # Create base directories + self.alembic_dir.mkdir(exist_ok=True) + if not has_versions: + (self.alembic_dir / 'versions').mkdir(exist_ok=True) + + # Write alembic.ini + ini_content = self._generate_alembic_ini_content() + with open(self.alembic_ini_path, 'w') as f: + f.write(ini_content) + logger.info("Created alembic.ini") + + if not has_versions: + # Only run init if no versions directory + config = self.get_alembic_config() + command.init(config, str(self.alembic_dir)) + logger.info("Initialized new Alembic directory structure") + else: + # Create minimal env.py if it doesn't exist + env_path = self.alembic_dir / 'env.py' + if not env_path.exists(): + self._create_minimal_env_py(env_path) + logger.info("Created minimal env.py") + else: + # Update existing env.py + self._update_env_py(env_path) + logger.info("Updated existing env.py") + + logger.info(f"Alembic setup completed at {self.base_dir}") + return str(self.alembic_ini_path) + + def _create_minimal_env_py(self, env_path: Path) -> None: + """Creates a minimal env.py file for Alembic.""" + content = ''' +from logging.config import fileConfig +from sqlalchemy import engine_from_config +from sqlalchemy import pool +from alembic import context +from sqlmodel import SQLModel + +config = context.config +if config.config_file_name is not None: + fileConfig(config.config_file_name) + +target_metadata = SQLModel.metadata + +def run_migrations_offline() -> None: + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + compare_type=True + ) + with context.begin_transaction(): + context.run_migrations() + +def run_migrations_online() -> None: + connectable = engine_from_config( + config.get_section(config.config_ini_section), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + with connectable.connect() as connection: + context.configure( + connection=connection, + target_metadata=target_metadata, + compare_type=True + ) + with context.begin_transaction(): + context.run_migrations() + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online()''' + + with open(env_path, 'w') as f: + f.write(content) + + def _generate_alembic_ini_content(self) -> str: + """ + Generates content for alembic.ini file. + """ + return f""" +[alembic] +script_location = {self.alembic_dir} +sqlalchemy.url = {self.engine.url} + +[loggers] +keys = root,sqlalchemy,alembic + +[handlers] +keys = console + +[formatters] +keys = generic + +[logger_root] +level = WARN +handlers = console +qualname = + +[logger_sqlalchemy] +level = WARN +handlers = +qualname = sqlalchemy.engine + +[logger_alembic] +level = INFO +handlers = +qualname = alembic + +[handler_console] +class = StreamHandler +args = (sys.stderr,) +level = NOTSET +formatter = generic + +[formatter_generic] +format = %(levelname)-5.5s [%(name)s] %(message)s +datefmt = %H:%M:%S +""".strip() + + def _update_env_py(self, env_path: Path) -> None: + """ + Updates the env.py file to use SQLModel metadata. + """ + try: + with open(env_path, 'r') as f: + content = f.read() + + # Add SQLModel import + if "from sqlmodel import SQLModel" not in content: + content = "from sqlmodel import SQLModel\n" + content + + # Replace target_metadata + content = content.replace( + "target_metadata = None", + "target_metadata = SQLModel.metadata" + ) + + # Add compare_type=True to context.configure + if "context.configure(" in content and "compare_type=True" not in content: + content = content.replace( + "context.configure(", + "context.configure(compare_type=True," + ) + + 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 + + # Fixed: use keyword-only argument + def _ensure_alembic_setup(self, *, force: bool = False) -> None: + """ + Ensures Alembic is properly set up, initializing if necessary. + + Args: + force: If True, removes existing configuration and reinitializes + """ + try: + self._validate_alembic_setup() + if force: + logger.info( + "Force initialization requested. Cleaning up existing configuration...") + self._cleanup_existing_alembic() + self._initialize_alembic() + except FileNotFoundError: + logger.info("Alembic configuration not found. Initializing...") + if self.alembic_dir.exists(): + logger.warning( + "Found existing alembic directory but missing configuration") + self._cleanup_existing_alembic() + self._initialize_alembic() + logger.info("Alembic initialization complete") + + def _validate_alembic_setup(self) -> None: + """Validates that Alembic is properly configured.""" + if not self.alembic_ini_path.exists(): + raise FileNotFoundError("Alembic configuration not found") + + def get_alembic_config(self) -> Config: + """ + Gets Alembic configuration. + + Returns: + Config: Alembic Config object + + Raises: + FileNotFoundError: If alembic.ini cannot be found + """ + if not self.alembic_ini_path.exists(): + raise FileNotFoundError("Could not find alembic.ini") + + return Config(str(self.alembic_ini_path)) + + def get_current_revision(self) -> Optional[str]: + """ + Gets the current database revision. + + Returns: + str: Current revision string or None if no revision + """ + with self.engine.connect() as conn: + context = MigrationContext.configure(conn) + return context.get_current_revision() + + def get_head_revision(self) -> str: + """ + Gets the latest available revision. + + Returns: + str: Head revision string + """ + config = self.get_alembic_config() + script = ScriptDirectory.from_config(config) + return script.get_current_head() + + def get_schema_differences(self) -> List[tuple]: + """ + Detects differences between current database and models. + + Returns: + List[tuple]: List of differences found + """ + with self.engine.connect() as conn: + context = MigrationContext.configure(conn) + diff = compare_metadata(context, SQLModel.metadata) + return list(diff) + + def check_schema_status(self) -> Tuple[bool, str]: + """ + Checks if database schema matches current models and migrations. + + Returns: + Tuple[bool, str]: (needs_upgrade, status_message) + """ + try: + current_rev = self.get_current_revision() + head_rev = self.get_head_revision() + + if current_rev != head_rev: + return True, f"Database needs upgrade: {current_rev} -> {head_rev}" + + differences = self.get_schema_differences() + if differences: + changes_desc = "\n".join(str(diff) for diff in differences) + return True, f"Unmigrated changes detected:\n{changes_desc}" + + return False, "Database schema is up to date" + + except Exception as e: + logger.error(f"Error checking schema status: {str(e)}") + return True, f"Error checking schema: {str(e)}" + + def upgrade_schema(self, revision: str = "head") -> bool: + """ + Upgrades database schema to specified revision. + + Args: + revision: Target revision (default: "head") + + Returns: + bool: True if upgrade successful + """ + try: + config = self.get_alembic_config() + command.upgrade(config, revision) + logger.info(f"Schema upgraded successfully to {revision}") + return True + + except Exception as e: + logger.error(f"Schema upgrade failed: {str(e)}") + return False + + def check_and_upgrade(self) -> Tuple[bool, str]: + """ + Checks schema status and upgrades if necessary (and auto_upgrade is True). + + Returns: + Tuple[bool, str]: (action_taken, status_message) + """ + needs_upgrade, status = self.check_schema_status() + + if needs_upgrade: + if self.auto_upgrade: + if self.upgrade_schema(): + return True, "Schema was automatically upgraded" + else: + return False, "Automatic schema upgrade failed" + else: + return False, f"Schema needs upgrade but auto_upgrade is disabled. Status: {status}" + + return False, status + + def generate_revision(self, message: str = "auto") -> Optional[str]: + """ + Generates new migration revision for current schema changes. + + Args: + message: Revision message + + Returns: + str: Revision ID if successful, None otherwise + """ + try: + config = self.get_alembic_config() + 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)}") + return None + + def get_pending_migrations(self) -> List[str]: + """ + Gets list of pending migrations that need to be applied. + + Returns: + List[str]: List of pending migration revision IDs + """ + config = self.get_alembic_config() + script = ScriptDirectory.from_config(config) + + current = self.get_current_revision() + head = self.get_head_revision() + + if current == head: + return [] + + pending = [] + for rev in script.iterate_revisions(current, head): + pending.append(rev.revision) + + return pending + + def print_status(self) -> None: + """Prints current migration status information to logger.""" + current = self.get_current_revision() + head = self.get_head_revision() + differences = self.get_schema_differences() + pending = self.get_pending_migrations() + + logger.info("=== Database Schema Status ===") + logger.info(f"Current revision: {current}") + logger.info(f"Head revision: {head}") + logger.info(f"Pending migrations: {len(pending)}") + for rev in pending: + logger.info(f" - {rev}") + logger.info(f"Unmigrated changes: {len(differences)}") + for diff in differences: + logger.info(f" - {diff}") + + def ensure_schema_up_to_date(self) -> bool: + """ + Ensures the database schema is up to date, generating and applying migrations if needed. + + Returns: + bool: True if schema is up to date or was successfully updated + """ + try: + # Check for unmigrated changes + differences = self.get_schema_differences() + if differences: + # Generate new migration + revision = self.generate_revision("auto-generated") + if not revision: + return False + logger.info(f"Generated new migration: {revision}") + + # Apply any pending migrations + upgraded, status = self.check_and_upgrade() + if not upgraded and "needs upgrade" in status.lower(): + return False + + return True + + except Exception as e: + logger.error(f"Failed to ensure schema is up to date: {e}") + return False diff --git a/python/packages/autogen-studio/autogenstudio/database/utils.py b/python/packages/autogen-studio/autogenstudio/database/utils.py deleted file mode 100644 index ac77a9161498..000000000000 --- a/python/packages/autogen-studio/autogenstudio/database/utils.py +++ /dev/null @@ -1,361 +0,0 @@ -# from .util import get_app_root -import os -import time -from datetime import datetime -from pathlib import Path -from typing import Any - -from alembic import command, util -from alembic.config import Config -from loguru import logger - -# from ..utils.db_utils import get_db_uri -from sqlmodel import Session, create_engine, text - -from autogen.agentchat import AssistantAgent - -from ..datamodel import ( - Agent, - AgentConfig, - AgentType, - CodeExecutionConfigTypes, - Model, - Skill, - Workflow, - WorkflowAgentLink, - WorkFlowType, -) - - -def workflow_from_id(workflow_id: int, dbmanager: Any): - workflow = dbmanager.get(Workflow, filters={"id": workflow_id}).data - if not workflow or len(workflow) == 0: - raise ValueError("The specified workflow does not exist.") - workflow = workflow[0].model_dump(mode="json") - workflow_agent_links = dbmanager.get(WorkflowAgentLink, filters={"workflow_id": workflow_id}).data - - def dump_agent(agent: Agent): - exclude = [] - if agent.type != AgentType.groupchat: - exclude = [ - "admin_name", - "messages", - "max_round", - "admin_name", - "speaker_selection_method", - "allow_repeat_speaker", - ] - return agent.model_dump(warnings=False, mode="json", exclude=exclude) - - def get_agent(agent_id): - with Session(dbmanager.engine) as session: - agent: Agent = dbmanager.get_items(Agent, filters={"id": agent_id}, session=session).data[0] - agent_dict = dump_agent(agent) - agent_dict["skills"] = [Skill.model_validate(skill.model_dump(mode="json")) for skill in agent.skills] - model_exclude = [ - "id", - "agent_id", - "created_at", - "updated_at", - "user_id", - "description", - ] - models = [model.model_dump(mode="json", exclude=model_exclude) for model in agent.models] - agent_dict["models"] = [model.model_dump(mode="json") for model in agent.models] - - if len(models) > 0: - agent_dict["config"]["llm_config"] = agent_dict.get("config", {}).get("llm_config", {}) - llm_config = agent_dict["config"]["llm_config"] - if llm_config: - llm_config["config_list"] = models - agent_dict["config"]["llm_config"] = llm_config - agent_dict["agents"] = [get_agent(agent.id) for agent in agent.agents] - return agent_dict - - agents = [] - for link in workflow_agent_links: - agent_dict = get_agent(link.agent_id) - agents.append({"agent": agent_dict, "link": link.model_dump(mode="json")}) - # workflow[str(link.agent_type.value)] = agent_dict - if workflow["type"] == WorkFlowType.sequential.value: - # sort agents by sequence_id in link - agents = sorted(agents, key=lambda x: x["link"]["sequence_id"]) - workflow["agents"] = agents - return workflow - - -def run_migration(engine_uri: str): - database_dir = Path(__file__).parent - script_location = database_dir / "migrations" - - engine = create_engine(engine_uri) - buffer = open(script_location / "alembic.log", "w") - alembic_cfg = Config(stdout=buffer) - alembic_cfg.set_main_option("script_location", str(script_location)) - alembic_cfg.set_main_option("sqlalchemy.url", engine_uri) - - print(f"Running migrations with engine_uri: {engine_uri}") - - should_initialize_alembic = False - with Session(engine) as session: - try: - session.exec(text("SELECT * FROM alembic_version")) - except Exception: - logger.info("Alembic not initialized") - should_initialize_alembic = True - else: - logger.info("Alembic already initialized") - - if should_initialize_alembic: - try: - logger.info("Initializing alembic") - command.ensure_version(alembic_cfg) - command.upgrade(alembic_cfg, "head") - logger.info("Alembic initialized") - except Exception as exc: - logger.error(f"Error initializing alembic: {exc}") - raise RuntimeError("Error initializing alembic") from exc - - logger.info(f"Running DB migrations in {script_location}") - - try: - buffer.write(f"{datetime.now().isoformat()}: Checking migrations\n") - command.check(alembic_cfg) - except Exception as exc: - if isinstance(exc, (util.exc.CommandError, util.exc.AutogenerateDiffsDetected)): - try: - command.upgrade(alembic_cfg, "head") - time.sleep(3) - except Exception as exc: - logger.error(f"Error running migrations: {exc}") - - try: - buffer.write(f"{datetime.now().isoformat()}: Checking migrations\n") - command.check(alembic_cfg) - except util.exc.AutogenerateDiffsDetected as exc: - logger.info(f"AutogenerateDiffsDetected: {exc}") - # raise RuntimeError( - # f"There's a mismatch between the models and the database.\n{exc}") - except util.exc.CommandError as exc: - logger.error(f"CommandError: {exc}") - # raise RuntimeError(f"Error running migrations: {exc}") - - -def init_db_samples(dbmanager: Any): - workflows = dbmanager.get(Workflow).data - workflow_names = [w.name for w in workflows] - if "Default Workflow" in workflow_names and "Travel Planning Workflow" in workflow_names: - logger.info("Database already initialized with Default and Travel Planning Workflows") - return - logger.info("Initializing database with Default and Travel Planning Workflows") - - # models - google_gemini_model = Model( - model="gemini-1.5-pro-latest", - description="Google's Gemini model", - user_id="guestuser@gmail.com", - api_type="google", - ) - azure_model = Model( - model="gpt4-turbo", - description="Azure OpenAI model", - user_id="guestuser@gmail.com", - api_type="azure", - base_url="https://api.your azureendpoint.com/v1", - ) - zephyr_model = Model( - model="zephyr", - description="Local Huggingface Zephyr model via vLLM, LMStudio or Ollama", - base_url="http://localhost:1234/v1", - user_id="guestuser@gmail.com", - api_type="open_ai", - ) - - gpt_4_model = Model( - model="gpt-4-1106-preview", description="OpenAI GPT-4 model", user_id="guestuser@gmail.com", api_type="open_ai" - ) - - anthropic_sonnet_model = Model( - model="claude-3-5-sonnet-20240620", - description="Anthropic's Claude 3.5 Sonnet model", - api_type="anthropic", - user_id="guestuser@gmail.com", - ) - - # skills - generate_pdf_skill = Skill( - name="generate_and_save_pdf", - description="Generate and save a pdf file based on the provided input sections.", - user_id="guestuser@gmail.com", - libraries=["requests", "fpdf", "PIL"], - content='import uuid\nimport requests\nfrom fpdf import FPDF\nfrom typing import List, Dict, Optional\nfrom pathlib import Path\nfrom PIL import Image, ImageDraw, ImageOps\nfrom io import BytesIO\n\ndef generate_and_save_pdf(\n sections: List[Dict[str, Optional[str]]], \n output_file: str = "report.pdf", \n report_title: str = "PDF Report"\n) -> None:\n """\n Function to generate a beautiful PDF report in A4 paper format. \n\n :param sections: A list of sections where each section is represented by a dictionary containing:\n - title: The title of the section.\n - level: The heading level (e.g., "title", "h1", "h2").\n - content: The content or body text of the section.\n - image: (Optional) The URL or local path to the image.\n :param output_file: The name of the output PDF file. (default is "report.pdf")\n :param report_title: The title of the report. (default is "PDF Report")\n :return: None\n """\n\n def get_image(image_url_or_path):\n if image_url_or_path.startswith("http://") or image_url_or_path.startswith("https://"):\n response = requests.get(image_url_or_path)\n if response.status_code == 200:\n return BytesIO(response.content)\n elif Path(image_url_or_path).is_file():\n return open(image_url_or_path, \'rb\')\n return None\n\n def add_rounded_corners(img, radius=6):\n mask = Image.new(\'L\', img.size, 0)\n draw = ImageDraw.Draw(mask)\n draw.rounded_rectangle([(0, 0), img.size], radius, fill=255)\n img = ImageOps.fit(img, mask.size, centering=(0.5, 0.5))\n img.putalpha(mask)\n return img\n\n class PDF(FPDF):\n def header(self):\n self.set_font("Arial", "B", 12)\n self.cell(0, 10, report_title, 0, 1, "C")\n \n def chapter_title(self, txt): \n self.set_font("Arial", "B", 12)\n self.cell(0, 10, txt, 0, 1, "L")\n self.ln(2)\n \n def chapter_body(self, body):\n self.set_font("Arial", "", 12)\n self.multi_cell(0, 10, body)\n self.ln()\n\n def add_image(self, img_data):\n img = Image.open(img_data)\n img = add_rounded_corners(img)\n img_path = Path(f"temp_{uuid.uuid4().hex}.png")\n img.save(img_path, format="PNG")\n self.image(str(img_path), x=None, y=None, w=190 if img.width > 190 else img.width)\n self.ln(10)\n img_path.unlink()\n\n pdf = PDF()\n pdf.add_page()\n font_size = {"title": 16, "h1": 14, "h2": 12, "body": 12}\n\n for section in sections:\n title, level, content, image = section.get("title", ""), section.get("level", "h1"), section.get("content", ""), section.get("image")\n pdf.set_font("Arial", "B" if level in font_size else "", font_size.get(level, font_size["body"]))\n pdf.chapter_title(title)\n\n if content: pdf.chapter_body(content)\n if image:\n img_data = get_image(image)\n if img_data:\n pdf.add_image(img_data)\n if isinstance(img_data, BytesIO):\n img_data.close()\n\n pdf.output(output_file)\n print(f"PDF report saved as {output_file}")\n\n# # Example usage\n# sections = [\n# {\n# "title": "Introduction - Early Life",\n# "level": "h1",\n# "image": "https://picsum.photos/536/354",\n# "content": ("Marie Curie was born on 7 November 1867 in Warsaw, Poland. "\n# "She was the youngest of five children. Both of her parents were teachers. "\n# "Her father was a math and physics instructor, and her mother was the head of a private school. "\n# "Marie\'s curiosity and brilliance were evident from an early age."),\n# },\n# {\n# "title": "Academic Accomplishments",\n# "level": "h2",\n# "content": ("Despite many obstacles, Marie Curie earned degrees in physics and mathematics from the University of Paris. "\n# "She conducted groundbreaking research on radioactivity, becoming the first woman to win a Nobel Prize. "\n# "Her achievements paved the way for future generations of scientists, particularly women in STEM fields."),\n# },\n# {\n# "title": "Major Discoveries",\n# "level": "h2",\n# "image": "https://picsum.photos/536/354",\n# "content": ("One of Marie Curie\'s most notable discoveries was that of radium and polonium, two radioactive elements. "\n# "Her meticulous work not only advanced scientific understanding but also had practical applications in medicine and industry."),\n# },\n# {\n# "title": "Conclusion - Legacy",\n# "level": "h1",\n# "content": ("Marie Curie\'s legacy lives on through her contributions to science, her role as a trailblazer for women in STEM, "\n# "and the ongoing impact of her discoveries on modern medicine and technology. "\n# "Her life and work remain an inspiration to many, demonstrating the power of perseverance and intellectual curiosity."),\n# },\n# ]\n\n# generate_and_save_pdf_report(sections, "my_report.pdf", "The Life of Marie Curie")', - ) - generate_image_skill = Skill( - name="generate_and_save_images", - secrets=[{"secret": "OPENAI_API_KEY", "value": None}], - libraries=["openai"], - description="Generate and save images based on a user's query.", - content='\nfrom typing import List\nimport uuid\nimport requests # to perform HTTP requests\nfrom pathlib import Path\n\nfrom openai import OpenAI\n\n\ndef generate_and_save_images(query: str, image_size: str = "1024x1024") -> List[str]:\n """\n Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI\'s DALL-E model and saves them to disk. Use the code below anytime there is a request to create an image.\n\n :param query: A natural language description of the image to be generated.\n :param image_size: The size of the image to be generated. (default is "1024x1024")\n :return: A list of filenames for the saved images.\n """\n\n client = OpenAI() # Initialize the OpenAI client\n response = client.images.generate(model="dall-e-3", prompt=query, n=1, size=image_size) # Generate images\n\n # List to store the file names of saved images\n saved_files = []\n\n # Check if the response is successful\n if response.data:\n for image_data in response.data:\n # Generate a random UUID as the file name\n file_name = str(uuid.uuid4()) + ".png" # Assuming the image is a PNG\n file_path = Path(file_name)\n\n img_url = image_data.url\n img_response = requests.get(img_url)\n if img_response.status_code == 200:\n # Write the binary content to a file\n with open(file_path, "wb") as img_file:\n img_file.write(img_response.content)\n print(f"Image saved to {file_path}")\n saved_files.append(str(file_path))\n else:\n print(f"Failed to download the image from {img_url}")\n else:\n print("No image data found in the response!")\n\n # Return the list of saved files\n return saved_files\n\n\n# Example usage of the function:\n# generate_and_save_images("A cute baby sea otter")\n', - user_id="guestuser@gmail.com", - ) - - # agents - - planner_assistant_config = AgentConfig( - name="planner_assistant", - description="Assistant Agent", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message="You are a helpful assistant that can suggest a travel plan for a user and utilize any context information provided. You are the primary cordinator who will receive suggestions or advice from other agents (local_assistant, language_assistant). You must ensure that the finally plan integrates the suggestions from other agents or team members. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE.", - code_execution_config=CodeExecutionConfigTypes.none, - llm_config={}, - ) - planner_assistant = Agent( - user_id="guestuser@gmail.com", - type=AgentType.assistant, - config=planner_assistant_config.model_dump(mode="json"), - ) - - local_assistant_config = AgentConfig( - name="local_assistant", - description="Local Assistant Agent", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message="You are a local assistant that can suggest local activities or places to visit for a user and can utilize any context information provided. You can suggest local activities, places to visit, restaurants to eat at, etc. You can also provide information about the weather, local events, etc. You can provide information about the local area, but you cannot suggest a complete travel plan. You can only provide information about the local area.", - code_execution_config=CodeExecutionConfigTypes.none, - llm_config={}, - ) - local_assistant = Agent( - user_id="guestuser@gmail.com", type=AgentType.assistant, config=local_assistant_config.model_dump(mode="json") - ) - - language_assistant_config = AgentConfig( - name="language_assistant", - description="Language Assistant Agent", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message="You are a helpful assistant that can review travel plans, providing feedback on important/critical tips about how best to address language or communication challenges for the given destination. If the plan already includes language tips, you can mention that the plan is satisfactory, with rationale.", - code_execution_config=CodeExecutionConfigTypes.none, - llm_config={}, - ) - language_assistant = Agent( - user_id="guestuser@gmail.com", - type=AgentType.assistant, - config=language_assistant_config.model_dump(mode="json"), - ) - - # group chat agent - travel_groupchat_config = AgentConfig( - name="travel_groupchat", - admin_name="groupchat", - description="Group Chat Agent Configuration", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message="You are a group chat manager", - code_execution_config=CodeExecutionConfigTypes.none, - default_auto_reply="TERMINATE", - llm_config={}, - speaker_selection_method="auto", - ) - travel_groupchat_agent = Agent( - user_id="guestuser@gmail.com", type=AgentType.groupchat, config=travel_groupchat_config.model_dump(mode="json") - ) - - user_proxy_config = AgentConfig( - name="user_proxy", - description="User Proxy Agent Configuration", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message="You are a helpful assistant", - code_execution_config=CodeExecutionConfigTypes.local, - default_auto_reply="TERMINATE", - llm_config=False, - ) - user_proxy = Agent( - user_id="guestuser@gmail.com", type=AgentType.userproxy, config=user_proxy_config.model_dump(mode="json") - ) - - default_assistant_config = AgentConfig( - name="default_assistant", - description="Assistant Agent", - human_input_mode="NEVER", - max_consecutive_auto_reply=25, - system_message=AssistantAgent.DEFAULT_SYSTEM_MESSAGE, - code_execution_config=CodeExecutionConfigTypes.none, - llm_config={}, - ) - default_assistant = Agent( - user_id="guestuser@gmail.com", type=AgentType.assistant, config=default_assistant_config.model_dump(mode="json") - ) - - # workflows - travel_workflow = Workflow( - name="Travel Planning Workflow", - description="Travel workflow", - user_id="guestuser@gmail.com", - sample_tasks=["Plan a 3 day trip to Hawaii Islands.", "Plan an eventful and exciting trip to Uzbeksitan."], - ) - default_workflow = Workflow( - name="Default Workflow", - description="Default workflow", - user_id="guestuser@gmail.com", - sample_tasks=[ - "paint a picture of a glass of ethiopian coffee, freshly brewed in a tall glass cup, on a table right in front of a lush green forest scenery", - "Plot the stock price of NVIDIA YTD.", - ], - ) - - with Session(dbmanager.engine) as session: - session.add(zephyr_model) - session.add(google_gemini_model) - session.add(azure_model) - session.add(gpt_4_model) - session.add(anthropic_sonnet_model) - session.add(generate_image_skill) - session.add(generate_pdf_skill) - session.add(user_proxy) - session.add(default_assistant) - session.add(travel_groupchat_agent) - session.add(planner_assistant) - session.add(local_assistant) - session.add(language_assistant) - - session.add(travel_workflow) - session.add(default_workflow) - session.commit() - - dbmanager.link(link_type="agent_model", primary_id=default_assistant.id, secondary_id=gpt_4_model.id) - dbmanager.link(link_type="agent_skill", primary_id=default_assistant.id, secondary_id=generate_image_skill.id) - dbmanager.link( - link_type="workflow_agent", primary_id=default_workflow.id, secondary_id=user_proxy.id, agent_type="sender" - ) - dbmanager.link( - link_type="workflow_agent", - primary_id=default_workflow.id, - secondary_id=default_assistant.id, - agent_type="receiver", - ) - - # link agents to travel groupchat agent - - dbmanager.link(link_type="agent_agent", primary_id=travel_groupchat_agent.id, secondary_id=planner_assistant.id) - dbmanager.link(link_type="agent_agent", primary_id=travel_groupchat_agent.id, secondary_id=local_assistant.id) - dbmanager.link( - link_type="agent_agent", primary_id=travel_groupchat_agent.id, secondary_id=language_assistant.id - ) - dbmanager.link(link_type="agent_agent", primary_id=travel_groupchat_agent.id, secondary_id=user_proxy.id) - dbmanager.link(link_type="agent_model", primary_id=travel_groupchat_agent.id, secondary_id=gpt_4_model.id) - dbmanager.link(link_type="agent_model", primary_id=planner_assistant.id, secondary_id=gpt_4_model.id) - dbmanager.link(link_type="agent_model", primary_id=local_assistant.id, secondary_id=gpt_4_model.id) - dbmanager.link(link_type="agent_model", primary_id=language_assistant.id, secondary_id=gpt_4_model.id) - - dbmanager.link( - link_type="workflow_agent", primary_id=travel_workflow.id, secondary_id=user_proxy.id, agent_type="sender" - ) - dbmanager.link( - link_type="workflow_agent", - primary_id=travel_workflow.id, - secondary_id=travel_groupchat_agent.id, - agent_type="receiver", - ) - logger.info("Successfully initialized database with Default and Travel Planning Workflows") diff --git a/python/packages/autogen-studio/autogenstudio/datamodel.py b/python/packages/autogen-studio/autogenstudio/datamodel.py deleted file mode 100644 index ee48818d599f..000000000000 --- a/python/packages/autogen-studio/autogenstudio/datamodel.py +++ /dev/null @@ -1,297 +0,0 @@ -from datetime import datetime -from enum import Enum -from typing import Any, Callable, Dict, List, Literal, Optional, Union - -from sqlalchemy import ForeignKey, Integer, orm -from sqlmodel import ( - JSON, - Column, - DateTime, - Field, - Relationship, - SQLModel, - func, -) -from sqlmodel import ( - Enum as SqlEnum, -) - -# added for python3.11 and sqlmodel 0.0.22 incompatibility -if hasattr(SQLModel, "model_config"): - SQLModel.model_config["protected_namespaces"] = () -elif hasattr(SQLModel, "Config"): - - class CustomSQLModel(SQLModel): - class Config: - protected_namespaces = () - - SQLModel = CustomSQLModel -else: - print("Warning: Unable to set protected_namespaces.") - -# pylint: disable=protected-access - - -class MessageMeta(SQLModel, table=False): - task: Optional[str] = None - messages: Optional[List[Dict[str, Any]]] = None - summary_method: Optional[str] = "last" - files: Optional[List[dict]] = None - time: Optional[datetime] = None - log: Optional[List[dict]] = None - usage: Optional[List[dict]] = None - - -class Message(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - role: str - content: str - session_id: Optional[int] = Field( - default=None, sa_column=Column(Integer, ForeignKey("session.id", ondelete="CASCADE")) - ) - connection_id: Optional[str] = None - meta: Optional[Union[MessageMeta, dict]] = Field(default={}, sa_column=Column(JSON)) - - -class Session(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - workflow_id: Optional[int] = Field(default=None, foreign_key="workflow.id") - name: Optional[str] = None - description: Optional[str] = None - - -class AgentSkillLink(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - agent_id: int = Field(default=None, primary_key=True, foreign_key="agent.id") - skill_id: int = Field(default=None, primary_key=True, foreign_key="skill.id") - - -class AgentModelLink(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - agent_id: int = Field(default=None, primary_key=True, foreign_key="agent.id") - model_id: int = Field(default=None, primary_key=True, foreign_key="model.id") - - -class Skill(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" - name: str - content: str - description: Optional[str] = None - secrets: Optional[List[dict]] = Field(default_factory=list, sa_column=Column(JSON)) - libraries: Optional[List[str]] = Field(default_factory=list, sa_column=Column(JSON)) - agents: List["Agent"] = Relationship(back_populates="skills", link_model=AgentSkillLink) - - -class LLMConfig(SQLModel, table=False): - """Data model for LLM Config for AutoGen""" - - config_list: List[Any] = Field(default_factory=list) - temperature: float = 0 - cache_seed: Optional[Union[int, None]] = None - timeout: Optional[int] = None - max_tokens: Optional[int] = 2048 - extra_body: Optional[dict] = None - - -class ModelTypes(str, Enum): - openai = "open_ai" - cerebras = "cerebras" - google = "google" - azure = "azure" - anthropic = "anthropic" - mistral = "mistral" - together = "together" - groq = "groq" - - -class Model(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" - model: str - api_key: Optional[str] = None - base_url: Optional[str] = None - api_type: ModelTypes = Field(default=ModelTypes.openai, sa_column=Column(SqlEnum(ModelTypes))) - api_version: Optional[str] = None - description: Optional[str] = None - agents: List["Agent"] = Relationship(back_populates="models", link_model=AgentModelLink) - - -class CodeExecutionConfigTypes(str, Enum): - local = "local" - docker = "docker" - none = "none" - - -class AgentConfig(SQLModel, table=False): - name: Optional[str] = None - human_input_mode: str = "NEVER" - max_consecutive_auto_reply: int = 10 - system_message: Optional[str] = None - is_termination_msg: Optional[Union[bool, str, Callable]] = None - code_execution_config: CodeExecutionConfigTypes = Field( - default=CodeExecutionConfigTypes.local, sa_column=Column(SqlEnum(CodeExecutionConfigTypes)) - ) - default_auto_reply: Optional[str] = "" - description: Optional[str] = None - llm_config: Optional[Union[LLMConfig, bool]] = Field(default=False, sa_column=Column(JSON)) - - admin_name: Optional[str] = "Admin" - messages: Optional[List[Dict]] = Field(default_factory=list) - max_round: Optional[int] = 100 - speaker_selection_method: Optional[str] = "auto" - allow_repeat_speaker: Optional[Union[bool, List["AgentConfig"]]] = True - - -class AgentType(str, Enum): - assistant = "assistant" - userproxy = "userproxy" - groupchat = "groupchat" - - -class WorkflowAgentType(str, Enum): - sender = "sender" - receiver = "receiver" - planner = "planner" - sequential = "sequential" - - -class WorkflowAgentLink(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - workflow_id: int = Field(default=None, primary_key=True, foreign_key="workflow.id") - agent_id: int = Field(default=None, primary_key=True, foreign_key="agent.id") - agent_type: WorkflowAgentType = Field( - default=WorkflowAgentType.sender, - sa_column=Column(SqlEnum(WorkflowAgentType), primary_key=True), - ) - sequence_id: Optional[int] = Field(default=0, primary_key=True) - - -class AgentLink(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - parent_id: Optional[int] = Field(default=None, foreign_key="agent.id", primary_key=True) - agent_id: Optional[int] = Field(default=None, foreign_key="agent.id", primary_key=True) - - -class Agent(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" - type: AgentType = Field(default=AgentType.assistant, sa_column=Column(SqlEnum(AgentType))) - config: Union[AgentConfig, dict] = Field(default_factory=AgentConfig, sa_column=Column(JSON)) - skills: List[Skill] = Relationship(back_populates="agents", link_model=AgentSkillLink) - models: List[Model] = Relationship(back_populates="agents", link_model=AgentModelLink) - workflows: List["Workflow"] = Relationship(link_model=WorkflowAgentLink, back_populates="agents") - parents: List["Agent"] = Relationship( - back_populates="agents", - link_model=AgentLink, - sa_relationship_kwargs=dict( - primaryjoin="Agent.id==AgentLink.agent_id", - secondaryjoin="Agent.id==AgentLink.parent_id", - ), - ) - agents: List["Agent"] = Relationship( - back_populates="parents", - link_model=AgentLink, - sa_relationship_kwargs=dict( - primaryjoin="Agent.id==AgentLink.parent_id", - secondaryjoin="Agent.id==AgentLink.agent_id", - ), - ) - task_instruction: Optional[str] = None - - -class WorkFlowType(str, Enum): - autonomous = "autonomous" - sequential = "sequential" - - -class WorkFlowSummaryMethod(str, Enum): - last = "last" - none = "none" - llm = "llm" - - -class Workflow(SQLModel, table=True): - __table_args__ = {"sqlite_autoincrement": True} - id: Optional[int] = Field(default=None, primary_key=True) - created_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), server_default=func.now()), - ) # pylint: disable=not-callable - updated_at: datetime = Field( - default_factory=datetime.now, - sa_column=Column(DateTime(timezone=True), onupdate=func.now()), - ) # pylint: disable=not-callable - user_id: Optional[str] = None - version: Optional[str] = "0.0.1" - name: str - description: str - agents: List[Agent] = Relationship(back_populates="workflows", link_model=WorkflowAgentLink) - type: WorkFlowType = Field(default=WorkFlowType.autonomous, sa_column=Column(SqlEnum(WorkFlowType))) - summary_method: Optional[WorkFlowSummaryMethod] = Field( - default=WorkFlowSummaryMethod.last, - sa_column=Column(SqlEnum(WorkFlowSummaryMethod)), - ) - sample_tasks: Optional[List[str]] = Field(default_factory=list, sa_column=Column(JSON)) - - -class Response(SQLModel): - message: str - status: bool - data: Optional[Any] = None - - -class SocketMessage(SQLModel, table=False): - connection_id: str - data: Dict[str, Any] - type: str diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py b/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py new file mode 100644 index 000000000000..6b7b4098df4b --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/datamodel/__init__.py @@ -0,0 +1,2 @@ +from .db import * +from .types import * diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/db.py b/python/packages/autogen-studio/autogenstudio/datamodel/db.py new file mode 100644 index 000000000000..2f8210029a9a --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/datamodel/db.py @@ -0,0 +1,282 @@ +# defines how core data types in autogenstudio are serialized and stored in the database + +from datetime import datetime +from enum import Enum +from typing import List, Optional, Union, Tuple, Type +from sqlalchemy import ForeignKey, Integer, UniqueConstraint +from sqlmodel import JSON, Column, DateTime, Field, SQLModel, func, Relationship, SQLModel +from uuid import UUID, uuid4 + +from .types import ToolConfig, ModelConfig, AgentConfig, TeamConfig, MessageConfig, MessageMeta + +# added for python3.11 and sqlmodel 0.0.22 incompatibility +if hasattr(SQLModel, "model_config"): + SQLModel.model_config["protected_namespaces"] = () +elif hasattr(SQLModel, "Config"): + class CustomSQLModel(SQLModel): + class Config: + protected_namespaces = () + + SQLModel = CustomSQLModel +else: + print("Warning: Unable to set protected_namespaces.") + +# pylint: disable=protected-access + + +class ComponentTypes(Enum): + TEAM = "team" + AGENT = "agent" + MODEL = "model" + TOOL = "tool" + + @property + def model_class(self) -> Type[SQLModel]: + return { + ComponentTypes.TEAM: Team, + ComponentTypes.AGENT: Agent, + ComponentTypes.MODEL: Model, + ComponentTypes.TOOL: Tool + }[self] + + +class LinkTypes(Enum): + AGENT_MODEL = "agent_model" + AGENT_TOOL = "agent_tool" + TEAM_AGENT = "team_agent" + + @property + # type: ignore + def link_config(self) -> Tuple[Type[SQLModel], Type[SQLModel], Type[SQLModel]]: + return { + LinkTypes.AGENT_MODEL: (Agent, Model, AgentModelLink), + LinkTypes.AGENT_TOOL: (Agent, Tool, AgentToolLink), + LinkTypes.TEAM_AGENT: (Team, Agent, TeamAgentLink) + }[self] + + @property + def primary_class(self) -> Type[SQLModel]: # type: ignore + return self.link_config[0] + + @property + def secondary_class(self) -> Type[SQLModel]: # type: ignore + return self.link_config[1] + + @property + def link_table(self) -> Type[SQLModel]: # type: ignore + return self.link_config[2] + + +# link models +class AgentToolLink(SQLModel, table=True): + __table_args__ = ( + UniqueConstraint('agent_id', 'sequence', + name='unique_agent_tool_sequence'), + {'sqlite_autoincrement': True} + ) + agent_id: int = Field(default=None, primary_key=True, + foreign_key="agent.id") + tool_id: int = Field(default=None, primary_key=True, foreign_key="tool.id") + sequence: Optional[int] = Field(default=0, primary_key=True) + + +class AgentModelLink(SQLModel, table=True): + __table_args__ = ( + UniqueConstraint('agent_id', 'sequence', + name='unique_agent_tool_sequence'), + {'sqlite_autoincrement': True} + ) + agent_id: int = Field(default=None, primary_key=True, + foreign_key="agent.id") + model_id: int = Field(default=None, primary_key=True, + foreign_key="model.id") + sequence: Optional[int] = Field(default=0, primary_key=True) + + +class TeamAgentLink(SQLModel, table=True): + __table_args__ = ( + UniqueConstraint('agent_id', 'sequence', + name='unique_agent_tool_sequence'), + {'sqlite_autoincrement': True} + ) + team_id: int = Field(default=None, primary_key=True, foreign_key="team.id") + agent_id: int = Field(default=None, primary_key=True, + foreign_key="agent.id") + sequence: Optional[int] = Field(default=0, primary_key=True) + +# database models + + +class Tool(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + config: Union[ToolConfig, dict] = Field( + default_factory=ToolConfig, sa_column=Column(JSON)) + agents: List["Agent"] = Relationship( + back_populates="tools", link_model=AgentToolLink) + + +class Model(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + config: Union[ModelConfig, dict] = Field( + default_factory=ModelConfig, sa_column=Column(JSON)) + agents: List["Agent"] = Relationship( + back_populates="models", link_model=AgentModelLink) + + +class Team(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + config: Union[TeamConfig, dict] = Field( + default_factory=TeamConfig, sa_column=Column(JSON)) + agents: List["Agent"] = Relationship( + back_populates="teams", link_model=TeamAgentLink) + + +class Agent(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + config: Union[AgentConfig, dict] = Field( + default_factory=AgentConfig, sa_column=Column(JSON)) + tools: List[Tool] = Relationship( + back_populates="agents", link_model=AgentToolLink) + models: List[Model] = Relationship( + back_populates="agents", link_model=AgentModelLink) + teams: List[Team] = Relationship( + back_populates="agents", link_model=TeamAgentLink) + + +class Message(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + config: Union[MessageConfig, dict] = Field( + default_factory=MessageConfig, sa_column=Column(JSON)) + session_id: Optional[int] = Field( + default=None, sa_column=Column(Integer, ForeignKey("session.id", ondelete="CASCADE")) + ) + run_id: Optional[UUID] = Field( + default=None, foreign_key="run.id" + ) + + message_meta: Optional[Union[MessageMeta, dict]] = Field( + default={}, sa_column=Column(JSON)) + + +class Session(SQLModel, table=True): + __table_args__ = {"sqlite_autoincrement": True} + id: Optional[int] = Field(default=None, primary_key=True) + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()), + ) # pylint: disable=not-callable + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()), + ) # pylint: disable=not-callable + user_id: Optional[str] = None + version: Optional[str] = "0.0.1" + team_id: Optional[int] = Field( + default=None, sa_column=Column(Integer, ForeignKey("team.id", ondelete="CASCADE")) + ) + name: Optional[str] = None + + +class RunStatus(str, Enum): + CREATED = "created" + ACTIVE = "active" + COMPLETE = "complete" + ERROR = "error" + STOPPED = "stopped" + + +class Run(SQLModel, table=True): + """Represents a single execution run within a session""" + __table_args__ = {"sqlite_autoincrement": True} + + # Primary key using UUID + id: UUID = Field( + default_factory=uuid4, + primary_key=True, + index=True + ) + + # Timestamps using the same pattern as other models + created_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), server_default=func.now()) + ) + updated_at: datetime = Field( + default_factory=datetime.now, + sa_column=Column(DateTime(timezone=True), onupdate=func.now()) + ) + + # Foreign key to Session + session_id: Optional[int] = Field( + default=None, + sa_column=Column( + Integer, + ForeignKey("session.id", ondelete="CASCADE"), + nullable=False + ) + ) + + # Run status and metadata + status: RunStatus = Field(default=RunStatus.CREATED) + error_message: Optional[str] = None + + # Metadata storage following pattern from Message model + run_meta: dict = Field(default={}, sa_column=Column(JSON)) + + # Version tracking like other models + version: Optional[str] = "0.0.1" diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/types.py b/python/packages/autogen-studio/autogenstudio/datamodel/types.py new file mode 100644 index 000000000000..3f059eba1903 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/datamodel/types.py @@ -0,0 +1,136 @@ +from datetime import datetime +from enum import Enum +from pathlib import Path +from typing import Any, Dict, List, Optional, Union + +from pydantic import BaseModel +from autogen_agentchat.base._task import TaskResult + + +class ModelTypes(str, Enum): + OPENAI = "OpenAIChatCompletionClient" + + +class ToolTypes(str, Enum): + PYTHON_FUNCTION = "PythonFunction" + + +class AgentTypes(str, Enum): + ASSISTANT = "AssistantAgent" + CODING = "CodingAssistantAgent" + + +class TeamTypes(str, Enum): + ROUND_ROBIN = "RoundRobinGroupChat" + SELECTOR = "SelectorGroupChat" + + +class TerminationTypes(str, Enum): + MAX_MESSAGES = "MaxMessageTermination" + STOP_MESSAGE = "StopMessageTermination" + TEXT_MENTION = "TextMentionTermination" + + +class ComponentType(str, Enum): + TEAM = "team" + AGENT = "agent" + MODEL = "model" + TOOL = "tool" + TERMINATION = "termination" + + +class BaseConfig(BaseModel): + model_config = { + "protected_namespaces": () + } + version: str = "1.0.0" + component_type: ComponentType + + +class MessageConfig(BaseModel): + source: str + content: str + message_type: Optional[str] = "text" + + +class ModelConfig(BaseConfig): + model: str + model_type: ModelTypes + api_key: Optional[str] = None + base_url: Optional[str] = None + component_type: ComponentType = ComponentType.MODEL + + +class ToolConfig(BaseConfig): + name: str + description: str + content: str + tool_type: ToolTypes + component_type: ComponentType = ComponentType.TOOL + + +class AgentConfig(BaseConfig): + name: str + agent_type: AgentTypes + system_message: Optional[str] = None + model_client: Optional[ModelConfig] = None + tools: Optional[List[ToolConfig]] = None + description: Optional[str] = None + component_type: ComponentType = ComponentType.AGENT + + +class TerminationConfig(BaseConfig): + termination_type: TerminationTypes + max_messages: Optional[int] = None + text: Optional[str] = None + component_type: ComponentType = ComponentType.TERMINATION + + +class TeamConfig(BaseConfig): + name: str + participants: List[AgentConfig] + team_type: TeamTypes + model_client: Optional[ModelConfig] = None + termination_condition: Optional[TerminationConfig] = None + component_type: ComponentType = ComponentType.TEAM + + +class TeamResult(BaseModel): + task_result: TaskResult + usage: str + duration: float + + +class MessageMeta(BaseModel): + task: Optional[str] = None + task_result: Optional[TaskResult] = None + summary_method: Optional[str] = "last" + files: Optional[List[dict]] = None + time: Optional[datetime] = None + log: Optional[List[dict]] = None + usage: Optional[List[dict]] = None + +# web request/response data models + + +class Response(BaseModel): + message: str + status: bool + data: Optional[Any] = None + + +class SocketMessage(BaseModel): + connection_id: str + data: Dict[str, Any] + type: str + + +ComponentConfig = Union[ + TeamConfig, + AgentConfig, + ModelConfig, + ToolConfig, + TerminationConfig +] + +ComponentConfigInput = Union[str, Path, dict, ComponentConfig] diff --git a/python/packages/autogen-studio/autogenstudio/teammanager.py b/python/packages/autogen-studio/autogenstudio/teammanager.py new file mode 100644 index 000000000000..e50f740472d6 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/teammanager.py @@ -0,0 +1,67 @@ +from typing import AsyncGenerator, Union, Optional +import time +from .database import ComponentFactory +from .datamodel import TeamResult, TaskResult, ComponentConfigInput +from autogen_agentchat.messages import InnerMessage, ChatMessage +from autogen_core.base import CancellationToken + + +class TeamManager: + def __init__(self) -> None: + self.component_factory = ComponentFactory() + + async def run_stream( + self, + task: str, + team_config: ComponentConfigInput, + cancellation_token: Optional[CancellationToken] = None + ) -> AsyncGenerator[Union[InnerMessage, ChatMessage, TaskResult], None]: + """Stream the team's execution results""" + start_time = time.time() + + try: + # Let factory handle all config processing + team = await self.component_factory.load(team_config) + + stream = team.run_stream( + task=task, + cancellation_token=cancellation_token + ) + + async for message in stream: + if cancellation_token and cancellation_token.is_cancelled(): + break + + if isinstance(message, TaskResult): + yield TeamResult( + task_result=message, + usage="", + duration=time.time() - start_time + ) + else: + yield message + + except Exception as e: + raise e + + async def run( + self, + task: str, + team_config: ComponentConfigInput, + cancellation_token: Optional[CancellationToken] = None + ) -> TeamResult: + """Original non-streaming run method with optional cancellation""" + start_time = time.time() + + # Let factory handle all config processing + team = await self.component_factory.load(team_config) + result = await team.run( + task=task, + cancellation_token=cancellation_token + ) + + return TeamResult( + task_result=result, + usage="", + duration=time.time() - start_time + ) diff --git a/python/packages/autogen-studio/autogenstudio/utils/dbdefaults.json b/python/packages/autogen-studio/autogenstudio/utils/dbdefaults.json deleted file mode 100644 index 7f36325266ea..000000000000 --- a/python/packages/autogen-studio/autogenstudio/utils/dbdefaults.json +++ /dev/null @@ -1,242 +0,0 @@ -{ - "models": [ - { - "model": "gpt-4", - "api_key": "Your Azure API key here", - "base_url": "Your Azure base URL here", - "api_type": "azure", - "api_version": "Your Azure API version here", - "description": "Azure Open AI model configuration" - }, - { - "model": "gpt-4-1106-preview", - "description": "OpenAI model configuration" - }, - { - "model": "TheBloke/zephyr-7B-alpha-AWQ", - "api_key": "EMPTY", - "base_url": "http://localhost:8000/v1", - "description": "Local model example with vLLM server endpoint" - } - ], - "agents": [ - { - "type": "userproxy", - - "config": { - "name": "userproxy", - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 5, - "system_message": "You are a helpful assistant.", - "default_auto_reply": "TERMINATE", - "llm_config": false, - "code_execution_config": { - "work_dir": null, - "use_docker": false - }, - "description": "A user proxy agent that executes code." - } - }, - { - "type": "assistant", - "skills": [ - { - "title": "find_papers_arxiv", - "description": "This skill finds relevant papers on arXiv given a query.", - "content": "import os\nimport re\nimport json\nimport hashlib\n\n\ndef search_arxiv(query, max_results=10):\n \"\"\"\n Searches arXiv for the given query using the arXiv API, then returns the search results. This is a helper function. In most cases, callers will want to use 'find_relevant_papers( query, max_results )' instead.\n\n Args:\n query (str): The search query.\n max_results (int, optional): The maximum number of search results to return. Defaults to 10.\n\n Returns:\n jresults (list): A list of dictionaries. Each dictionary contains fields such as 'title', 'authors', 'summary', and 'pdf_url'\n\n Example:\n >>> results = search_arxiv(\"attention is all you need\")\n >>> print(results)\n \"\"\"\n\n import arxiv\n\n key = hashlib.md5((\"search_arxiv(\" + str(max_results) + \")\" + query).encode(\"utf-8\")).hexdigest()\n # Create the cache if it doesn't exist\n cache_dir = \".cache\"\n if not os.path.isdir(cache_dir):\n os.mkdir(cache_dir)\n\n fname = os.path.join(cache_dir, key + \".cache\")\n\n # Cache hit\n if os.path.isfile(fname):\n fh = open(fname, \"r\", encoding=\"utf-8\")\n data = json.loads(fh.read())\n fh.close()\n return data\n\n # Normalize the query, removing operator keywords\n query = re.sub(r\"[^\\s\\w]\", \" \", query.lower())\n query = re.sub(r\"\\s(and|or|not)\\s\", \" \", \" \" + query + \" \")\n query = re.sub(r\"[^\\s\\w]\", \" \", query.lower())\n query = re.sub(r\"\\s+\", \" \", query).strip()\n\n search = arxiv.Search(query=query, max_results=max_results, sort_by=arxiv.SortCriterion.Relevance)\n\n jresults = list()\n for result in search.results():\n r = dict()\n r[\"entry_id\"] = result.entry_id\n r[\"updated\"] = str(result.updated)\n r[\"published\"] = str(result.published)\n r[\"title\"] = result.title\n r[\"authors\"] = [str(a) for a in result.authors]\n r[\"summary\"] = result.summary\n r[\"comment\"] = result.comment\n r[\"journal_ref\"] = result.journal_ref\n r[\"doi\"] = result.doi\n r[\"primary_category\"] = result.primary_category\n r[\"categories\"] = result.categories\n r[\"links\"] = [str(link) for link in result.links]\n r[\"pdf_url\"] = result.pdf_url\n jresults.append(r)\n\n if len(jresults) > max_results:\n jresults = jresults[0:max_results]\n\n # Save to cache\n fh = open(fname, \"w\")\n fh.write(json.dumps(jresults))\n fh.close()\n return jresults\n", - "file_name": "find_papers_arxiv" - }, - { - "title": "generate_images", - "description": "This skill generates images from a given query using OpenAI's DALL-E model and saves them to disk.", - "content": "from typing import List\nimport uuid\nimport requests # to perform HTTP requests\nfrom pathlib import Path\n\nfrom openai import OpenAI\n\n\ndef generate_and_save_images(query: str, image_size: str = \"1024x1024\") -> List[str]:\n \"\"\"\n Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI's DALL-E model and saves them to disk. Use the code below anytime there is a request to create an image.\n\n :param query: A natural language description of the image to be generated.\n :param image_size: The size of the image to be generated. (default is \"1024x1024\")\n :return: A list of filenames for the saved images.\n \"\"\"\n\n client = OpenAI() # Initialize the OpenAI client\n response = client.images.generate(model=\"dall-e-3\", prompt=query, n=1, size=image_size) # Generate images\n\n # List to store the file names of saved images\n saved_files = []\n\n # Check if the response is successful\n if response.data:\n for image_data in response.data:\n # Generate a random UUID as the file name\n file_name = str(uuid.uuid4()) + \".png\" # Assuming the image is a PNG\n file_path = Path(file_name)\n\n img_url = image_data.url\n img_response = requests.get(img_url)\n if img_response.status_code == 200:\n # Write the binary content to a file\n with open(file_path, \"wb\") as img_file:\n img_file.write(img_response.content)\n print(f\"Image saved to {file_path}\")\n saved_files.append(str(file_path))\n else:\n print(f\"Failed to download the image from {img_url}\")\n else:\n print(\"No image data found in the response!\")\n\n # Return the list of saved files\n return saved_files\n\n\n# Example usage of the function:\n# generate_and_save_images(\"A cute baby sea otter\")\n" - } - ], - "config": { - "name": "primary_assistant", - "description": "A primary assistant agent that writes plans and code to solve tasks.", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": null - }, - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 8, - "system_message": "You are a helpful AI assistant. Solve tasks using your coding and language skills. In the following cases, suggest python code (in a python coding block) or shell script (in a sh coding block) for the user to execute. 1. When you need to collect info, use the code to output the info you need, for example, browse or search the web, download/read a file, print the content of a webpage or a file, get the current date/time, check the operating system. After sufficient info is printed and the task is ready to be solved based on your language skill, you can solve the task by yourself. 2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly. Solve the task step by step if you need to. If a plan is not provided, explain your plan first. Be clear which step uses code, and which step uses your language skill. When using code, you must indicate the script type in the code block. The user cannot provide any other feedback or perform any other action beyond executing the code you suggest. The user can't modify your code. So do not suggest incomplete code which requires users to modify. Don't use a code block if it's not intended to be executed by the user. If you want the user to save the code in a file before executing it, put # filename: inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use 'print' function for the output when relevant. Check the execution result returned by the user. If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try. When you find an answer, verify the answer carefully. Include verifiable evidence in your response if possible. Reply 'TERMINATE' in the end when everything is done." - } - } - ], - "skills": [ - { - "title": "fetch_profile", - "description": "This skill fetches the text content from a personal website.", - "content": "from typing import Optional\nimport requests\nfrom bs4 import BeautifulSoup\n\n\ndef fetch_user_profile(url: str) -> Optional[str]:\n \"\"\"\n Fetches the text content from a personal website.\n\n Given a URL of a person's personal website, this function scrapes\n the content of the page and returns the text found within the .\n\n Args:\n url (str): The URL of the person's personal website.\n\n Returns:\n Optional[str]: The text content of the website's body, or None if any error occurs.\n \"\"\"\n try:\n # Send a GET request to the URL\n response = requests.get(url)\n # Check for successful access to the webpage\n if response.status_code == 200:\n # Parse the HTML content of the page using BeautifulSoup\n soup = BeautifulSoup(response.text, \"html.parser\")\n # Extract the content of the tag\n body_content = soup.find(\"body\")\n # Return all the text in the body tag, stripping leading/trailing whitespaces\n return \" \".join(body_content.stripped_strings) if body_content else None\n else:\n # Return None if the status code isn't 200 (success)\n return None\n except requests.RequestException:\n # Return None if any request-related exception is caught\n return None\n" - }, - { - "title": "generate_images", - "description": "This skill generates images from a given query using OpenAI's DALL-E model and saves them to disk.", - "content": "from typing import List\nimport uuid\nimport requests # to perform HTTP requests\nfrom pathlib import Path\n\nfrom openai import OpenAI\n\n\ndef generate_and_save_images(query: str, image_size: str = \"1024x1024\") -> List[str]:\n \"\"\"\n Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI's DALL-E model and saves them to disk. Use the code below anytime there is a request to create an image.\n\n :param query: A natural language description of the image to be generated.\n :param image_size: The size of the image to be generated. (default is \"1024x1024\")\n :return: A list of filenames for the saved images.\n \"\"\"\n\n client = OpenAI() # Initialize the OpenAI client\n response = client.images.generate(model=\"dall-e-3\", prompt=query, n=1, size=image_size) # Generate images\n\n # List to store the file names of saved images\n saved_files = []\n\n # Check if the response is successful\n if response.data:\n for image_data in response.data:\n # Generate a random UUID as the file name\n file_name = str(uuid.uuid4()) + \".png\" # Assuming the image is a PNG\n file_path = Path(file_name)\n\n img_url = image_data.url\n img_response = requests.get(img_url)\n if img_response.status_code == 200:\n # Write the binary content to a file\n with open(file_path, \"wb\") as img_file:\n img_file.write(img_response.content)\n print(f\"Image saved to {file_path}\")\n saved_files.append(str(file_path))\n else:\n print(f\"Failed to download the image from {img_url}\")\n else:\n print(\"No image data found in the response!\")\n\n # Return the list of saved files\n return saved_files\n\n\n# Example usage of the function:\n# generate_and_save_images(\"A cute baby sea otter\")\n" - } - ], - "workflows": [ - { - "name": "Travel Agent Group Chat Workflow", - "description": "A group chat workflow", - "type": "groupchat", - "sender": { - "type": "userproxy", - "config": { - "name": "userproxy", - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 5, - "system_message": "You are a helpful assistant.", - "code_execution_config": { - "work_dir": null, - "use_docker": false - } - } - }, - "receiver": { - "type": "groupchat", - "config": { - "name": "group_chat_manager", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": 42 - }, - "human_input_mode": "NEVER", - "system_message": "Group chat manager" - }, - "groupchat_config": { - "admin_name": "Admin", - "max_round": 10, - "speaker_selection_method": "auto", - "agents": [ - { - "type": "assistant", - "config": { - "name": "travel_planner", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": 42 - }, - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 8, - "system_message": "You are a helpful assistant that can suggest a travel plan for a user. You are the primary cordinator who will receive suggestions or advice from other agents (local_assistant, language_assistant). You must ensure that the finally plan integrates the suggestions from other agents or team members. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE." - } - }, - { - "type": "assistant", - "config": { - "name": "local_assistant", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": 42 - }, - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 8, - "system_message": "You are a helpful assistant that can review travel plans, providing critical feedback on how the trip can be enriched for enjoyment of the local culture. If the plan already includes local experiences, you can mention that the plan is satisfactory, with rationale." - } - }, - { - "type": "assistant", - "config": { - "name": "language_assistant", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": 42 - }, - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 8, - "system_message": "You are a helpful assistant that can review travel plans, providing feedback on important/critical tips about how best to address language or communication challenges for the given destination. If the plan already includes language tips, you can mention that the plan is satisfactory, with rationale." - } - } - ] - } - } - }, - { - "name": "General Agent Workflow", - "description": "This workflow is used for general purpose tasks.", - "sender": { - "type": "userproxy", - "config": { - "name": "userproxy", - "description": "A user proxy agent that executes code.", - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 10, - "system_message": "You are a helpful assistant.", - "default_auto_reply": "TERMINATE", - "llm_config": false, - "code_execution_config": { - "work_dir": null, - "use_docker": false - } - } - }, - "receiver": { - "type": "assistant", - - "skills": [ - { - "title": "find_papers_arxiv", - "description": "This skill finds relevant papers on arXiv given a query.", - "content": "import os\nimport re\nimport json\nimport hashlib\n\n\ndef search_arxiv(query, max_results=10):\n \"\"\"\n Searches arXiv for the given query using the arXiv API, then returns the search results. This is a helper function. In most cases, callers will want to use 'find_relevant_papers( query, max_results )' instead.\n\n Args:\n query (str): The search query.\n max_results (int, optional): The maximum number of search results to return. Defaults to 10.\n\n Returns:\n jresults (list): A list of dictionaries. Each dictionary contains fields such as 'title', 'authors', 'summary', and 'pdf_url'\n\n Example:\n >>> results = search_arxiv(\"attention is all you need\")\n >>> print(results)\n \"\"\"\n\n import arxiv\n\n key = hashlib.md5((\"search_arxiv(\" + str(max_results) + \")\" + query).encode(\"utf-8\")).hexdigest()\n # Create the cache if it doesn't exist\n cache_dir = \".cache\"\n if not os.path.isdir(cache_dir):\n os.mkdir(cache_dir)\n\n fname = os.path.join(cache_dir, key + \".cache\")\n\n # Cache hit\n if os.path.isfile(fname):\n fh = open(fname, \"r\", encoding=\"utf-8\")\n data = json.loads(fh.read())\n fh.close()\n return data\n\n # Normalize the query, removing operator keywords\n query = re.sub(r\"[^\\s\\w]\", \" \", query.lower())\n query = re.sub(r\"\\s(and|or|not)\\s\", \" \", \" \" + query + \" \")\n query = re.sub(r\"[^\\s\\w]\", \" \", query.lower())\n query = re.sub(r\"\\s+\", \" \", query).strip()\n\n search = arxiv.Search(query=query, max_results=max_results, sort_by=arxiv.SortCriterion.Relevance)\n\n jresults = list()\n for result in search.results():\n r = dict()\n r[\"entry_id\"] = result.entry_id\n r[\"updated\"] = str(result.updated)\n r[\"published\"] = str(result.published)\n r[\"title\"] = result.title\n r[\"authors\"] = [str(a) for a in result.authors]\n r[\"summary\"] = result.summary\n r[\"comment\"] = result.comment\n r[\"journal_ref\"] = result.journal_ref\n r[\"doi\"] = result.doi\n r[\"primary_category\"] = result.primary_category\n r[\"categories\"] = result.categories\n r[\"links\"] = [str(link) for link in result.links]\n r[\"pdf_url\"] = result.pdf_url\n jresults.append(r)\n\n if len(jresults) > max_results:\n jresults = jresults[0:max_results]\n\n # Save to cache\n fh = open(fname, \"w\")\n fh.write(json.dumps(jresults))\n fh.close()\n return jresults\n" - }, - { - "title": "generate_images", - "description": "This skill generates images from a given query using OpenAI's DALL-E model and saves them to disk.", - "content": "from typing import List\nimport uuid\nimport requests # to perform HTTP requests\nfrom pathlib import Path\n\nfrom openai import OpenAI\n\n\ndef generate_and_save_images(query: str, image_size: str = \"1024x1024\") -> List[str]:\n \"\"\"\n Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI's DALL-E model and saves them to disk. Use the code below anytime there is a request to create an image.\n\n :param query: A natural language description of the image to be generated.\n :param image_size: The size of the image to be generated. (default is \"1024x1024\")\n :return: A list of filenames for the saved images.\n \"\"\"\n\n client = OpenAI() # Initialize the OpenAI client\n response = client.images.generate(model=\"dall-e-3\", prompt=query, n=1, size=image_size) # Generate images\n\n # List to store the file names of saved images\n saved_files = []\n\n # Check if the response is successful\n if response.data:\n for image_data in response.data:\n # Generate a random UUID as the file name\n file_name = str(uuid.uuid4()) + \".png\" # Assuming the image is a PNG\n file_path = Path(file_name)\n\n img_url = image_data.url\n img_response = requests.get(img_url)\n if img_response.status_code == 200:\n # Write the binary content to a file\n with open(file_path, \"wb\") as img_file:\n img_file.write(img_response.content)\n print(f\"Image saved to {file_path}\")\n saved_files.append(str(file_path))\n else:\n print(f\"Failed to download the image from {img_url}\")\n else:\n print(\"No image data found in the response!\")\n\n # Return the list of saved files\n return saved_files\n\n\n# Example usage of the function:\n# generate_and_save_images(\"A cute baby sea otter\")\n" - } - ], - "config": { - "description": "Default assistant to generate plans and write code to solve tasks.", - "name": "primary_assistant", - "llm_config": { - "config_list": [ - { - "model": "gpt-4-1106-preview" - } - ], - "temperature": 0.1, - "timeout": 600, - "cache_seed": null - }, - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 15, - "system_message": "You are a helpful AI assistant. Solve tasks using your coding and language skills. In the following cases, suggest python code (in a python coding block) or shell script (in a sh coding block) for the user to execute. 1. When you need to collect info, use the code to output the info you need, for example, browse or search the web, download/read a file, print the content of a webpage or a file, get the current date/time, check the operating system. After sufficient info is printed and the task is ready to be solved based on your language skill, you can solve the task by yourself. 2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly. Solve the task step by step if you need to. If a plan is not provided, explain your plan first. Be clear which step uses code, and which step uses your language skill. When using code, you must indicate the script type in the code block. The user cannot provide any other feedback or perform any other action beyond executing the code you suggest. The user can't modify your code. So do not suggest incomplete code which requires users to modify. Don't use a code block if it's not intended to be executed by the user. If you want the user to save the code in a file before executing it, put # filename: inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use 'print' function for the output when relevant. Check the execution result returned by the user. If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try. When you find an answer, verify the answer carefully. Include verifiable evidence in your response if possible. Reply 'TERMINATE' in the end when everything is done." - } - }, - "type": "twoagents" - } - ] -} diff --git a/python/packages/autogen-studio/autogenstudio/utils/utils.py b/python/packages/autogen-studio/autogenstudio/utils/utils.py index 88f310d6ffc3..419a6e4a66d2 100644 --- a/python/packages/autogen-studio/autogenstudio/utils/utils.py +++ b/python/packages/autogen-studio/autogenstudio/utils/utils.py @@ -10,10 +10,8 @@ from dotenv import load_dotenv from loguru import logger -from autogen.coding import DockerCommandLineCodeExecutor, LocalCommandLineCodeExecutor -from autogen.oai.client import ModelClient, OpenAIWrapper -from ..datamodel import CodeExecutionConfigTypes, Model, Skill +from ..datamodel import Model from ..version import APP_NAME @@ -44,26 +42,6 @@ def str_to_datetime(dt_str: str) -> datetime: return datetime.fromisoformat(dt_str) -def clear_folder(folder_path: str) -> None: - """ - Clear the contents of a folder. - - :param folder_path: The path to the folder to clear. - """ - # exit if the folder does not exist - if not os.path.exists(folder_path): - return - # exit if the folder does not exist - if not os.path.exists(folder_path): - return - for file in os.listdir(folder_path): - file_path = os.path.join(folder_path, file) - if os.path.isfile(file_path): - os.remove(file_path) - elif os.path.isdir(file_path): - shutil.rmtree(file_path) - - def get_file_type(file_path: str) -> str: """ @@ -153,31 +131,6 @@ def get_file_type(file_path: str) -> str: return file_type -def serialize_file(file_path: str) -> Tuple[str, str]: - """ - Reads a file from a given file path, base64 encodes its content, - and returns the base64 encoded string along with the file type. - - The file type is determined by the file extension. If the file extension is not - recognized, 'unknown' will be used as the file type. - - :param file_path: The path to the file to be serialized. - :return: A tuple containing the base64 encoded string of the file and the file type. - """ - - file_type = get_file_type(file_path) - - # Read the file and encode its contents - try: - with open(file_path, "rb") as file: - file_content = file.read() - base64_encoded_content = base64.b64encode(file_content).decode("utf-8") - except Exception as e: - raise IOError(f"An error occurred while reading the file: {e}") from e - - return base64_encoded_content, file_type - - def get_modified_files(start_timestamp: float, end_timestamp: float, source_dir: str) -> List[Dict[str, str]]: """ Identify files from source_dir that were modified within a specified timestamp range. @@ -200,7 +153,8 @@ def get_modified_files(start_timestamp: float, end_timestamp: float, source_dir: for root, dirs, files in os.walk(source_dir): # Update directories and files to exclude those to be ignored dirs[:] = [d for d in dirs if d not in ignore_files] - files[:] = [f for f in files if f not in ignore_files and os.path.splitext(f)[1] not in ignore_extensions] + files[:] = [f for f in files if f not in ignore_files and os.path.splitext(f)[ + 1] not in ignore_extensions] for file in files: file_path = os.path.join(root, file) @@ -209,7 +163,9 @@ def get_modified_files(start_timestamp: float, end_timestamp: float, source_dir: # Verify if the file was modified within the given timestamp range if start_timestamp <= file_mtime <= end_timestamp: file_relative_path = ( - "files/user" + file_path.split("files/user", 1)[1] if "files/user" in file_path else "" + "files/user" + + file_path.split( + "files/user", 1)[1] if "files/user" in file_path else "" ) file_type = get_file_type(file_path) @@ -289,138 +245,6 @@ def init_app_folders(app_file_path: str) -> Dict[str, str]: return folders -def get_skills_prompt(skills: List[Skill], work_dir: str) -> str: - """ - Create a prompt with the content of all skills and write the skills to a file named skills.py in the work_dir. - - :param skills: A dictionary skills - :return: A string containing the content of all skills - """ - - instruction = """ - -While solving the task you may use functions below which will be available in a file called skills.py . -To use a function skill.py in code, IMPORT THE FUNCTION FROM skills.py and then use the function. -If you need to install python packages, write shell code to -install via pip and use --quiet option. - - """ - prompt = "" # filename: skills.py - - for skill in skills: - if not isinstance(skill, Skill): - skill = Skill(**skill) - if skill.secrets: - for secret in skill.secrets: - if secret.get("value") is not None: - os.environ[secret["secret"]] = secret["value"] - prompt += f""" - -##### Begin of {skill.name} ##### -from skills import {skill.name} # Import the function from skills.py - -{skill.content} - -#### End of {skill.name} #### - - """ - - return instruction + prompt - - -def save_skills_to_file(skills: List[Skill], work_dir: str) -> None: - """ - Write the skills to a file named skills.py in the work_dir. - - :param skills: A dictionary skills - """ - - # TBD: Double check for duplicate skills? - - # check if work_dir exists - if not os.path.exists(work_dir): - os.makedirs(work_dir) - - skills_content = "" - for skill in skills: - if not isinstance(skill, Skill): - skill = Skill(**skill) - - skills_content += f""" - -##### Begin of {skill.name} ##### - -{skill.content} - -#### End of {skill.name} #### - - """ - - # overwrite skills.py in work_dir - with open(os.path.join(work_dir, "skills.py"), "w", encoding="utf-8") as f: - f.write(skills_content) - - -def delete_files_in_folder(folders: Union[str, List[str]]) -> None: - """ - Delete all files and directories in the specified folders. - - :param folders: A list of folders or a single folder string - """ - - if isinstance(folders, str): - folders = [folders] - - for folder in folders: - # Check if the folder exists - if not os.path.isdir(folder): - continue - - # List all the entries in the directory - for entry in os.listdir(folder): - # Get the full path - path = os.path.join(folder, entry) - try: - if os.path.isfile(path) or os.path.islink(path): - # Remove the file or link - os.remove(path) - elif os.path.isdir(path): - # Remove the directory and all its content - shutil.rmtree(path) - except Exception as e: - # Print the error message and skip - logger.info(f"Failed to delete {path}. Reason: {e}") - - -def extract_successful_code_blocks(messages: List[Dict[str, str]]) -> List[str]: - """ - Parses through a list of messages containing code blocks and execution statuses, - returning the array of code blocks that executed successfully and retains - the backticks for Markdown rendering. - - Parameters: - messages (List[Dict[str, str]]): A list of message dictionaries containing 'content' and 'role' keys. - - Returns: - List[str]: A list containing the code blocks that were successfully executed, including backticks. - """ - successful_code_blocks = [] - # Regex pattern to capture code blocks enclosed in triple backticks. - code_block_regex = r"```[\s\S]*?```" - - for i, row in enumerate(messages): - message = row["message"] - if message["role"] == "user" and "execution succeeded" in message["content"]: - if i > 0 and messages[i - 1]["message"]["role"] == "assistant": - prev_content = messages[i - 1]["message"]["content"] - # Find all matches for code blocks - code_blocks = re.findall(code_block_regex, prev_content) - # Add the code blocks with backticks - successful_code_blocks.extend(code_blocks) - - return successful_code_blocks - - def sanitize_model(model: Model): """ Sanitize model dictionary to remove None values and empty strings and only keep valid keys. @@ -429,7 +253,8 @@ def sanitize_model(model: Model): model = model.model_dump() valid_keys = ["model", "base_url", "api_key", "api_type", "api_version"] # only add key if value is not None - sanitized_model = {k: v for k, v in model.items() if (v is not None and v != "") and k in valid_keys} + sanitized_model = {k: v for k, v in model.items() if ( + v is not None and v != "") and k in valid_keys} return sanitized_model @@ -440,134 +265,29 @@ def test_model(model: Model): print("Testing model", model) - sanitized_model = sanitize_model(model) - client = OpenAIWrapper(config_list=[sanitized_model]) - response = client.create( - messages=[ - { - "role": "system", - "content": "You are a helpful assistant that can add numbers. ONLY RETURN THE RESULT.", - }, - { - "role": "user", - "content": "2+2=", - }, - ], - cache_seed=None, - ) - return response.choices[0].message.content - - -def load_code_execution_config(code_execution_type: CodeExecutionConfigTypes, work_dir: str): - """ - Load the code execution configuration based on the code execution type. - - :param code_execution_type: The code execution type. - :param work_dir: The working directory to store code execution files. - :return: The code execution configuration. - - """ - work_dir = Path(work_dir) - work_dir.mkdir(exist_ok=True) - executor = None - if code_execution_type == CodeExecutionConfigTypes.local: - executor = LocalCommandLineCodeExecutor(work_dir=work_dir) - elif code_execution_type == CodeExecutionConfigTypes.docker: - try: - executor = DockerCommandLineCodeExecutor(work_dir=work_dir) - except Exception as e: - logger.error(f"Error initializing Docker executor: {e}") - return False - elif code_execution_type == CodeExecutionConfigTypes.none: - return False - else: - raise ValueError(f"Invalid code execution type: {code_execution_type}") - code_execution_config = { - "executor": executor, - } - return code_execution_config - -def summarize_chat_history(task: str, messages: List[Dict[str, str]], client: ModelClient): - """ - Summarize the chat history using the model endpoint and returning the response. - """ - summarization_system_prompt = f""" - You are a helpful assistant that is able to review the chat history between a set of agents (userproxy agents, assistants etc) as they try to address a given TASK and provide a summary. Be SUCCINCT but also comprehensive enough to allow others (who cannot see the chat history) understand and recreate the solution. - - The task requested by the user is: - === - {task} - === - The summary should focus on extracting the actual solution to the task from the chat history (assuming the task was addressed) such that any other agent reading the summary will understand what the actual solution is. Use a neutral tone and DO NOT directly mention the agents. Instead only focus on the actions that were carried out (e.g. do not say 'assistant agent generated some code visualization code ..' instead say say 'visualization code was generated ..'. The answer should be framed as a response to the user task. E.g. if the task is "What is the height of the Eiffel tower", the summary should be "The height of the Eiffel Tower is ..."). - """ - summarization_prompt = [ - { - "role": "system", - "content": summarization_system_prompt, - }, - { - "role": "user", - "content": f"Summarize the following chat history. {str(messages)}", - }, - ] - response = client.create(messages=summarization_prompt, cache_seed=None) - return response.choices[0].message.content - - -def get_autogen_log(db_path="logs.db"): - """ - Fetches data the autogen logs database. - Args: - dbname (str): Name of the database file. Defaults to "logs.db". - table (str): Name of the table to query. Defaults to "chat_completions". - - Returns: - list: A list of dictionaries, where each dictionary represents a row from the table. - """ - import json - import sqlite3 - - con = sqlite3.connect(db_path) - query = """ - SELECT - chat_completions.*, - agents.name AS agent_name - FROM - chat_completions - JOIN - agents ON chat_completions.wrapper_id = agents.wrapper_id - """ - cursor = con.execute(query) - rows = cursor.fetchall() - column_names = [description[0] for description in cursor.description] - data = [dict(zip(column_names, row)) for row in rows] - for row in data: - response = json.loads(row["response"]) - print(response) - total_tokens = response.get("usage", {}).get("total_tokens", 0) - row["total_tokens"] = total_tokens - con.close() - return data - - -def find_key_value(d, target_key): - """ - Recursively search for a key in a nested dictionary and return its value. - """ - if d is None: - return None - - if isinstance(d, dict): - if target_key in d: - return d[target_key] - for k in d: - item = find_key_value(d[k], target_key) - if item is not None: - return item - elif isinstance(d, list): - for i in d: - item = find_key_value(i, target_key) - if item is not None: - return item - return None +# def summarize_chat_history(task: str, messages: List[Dict[str, str]], client: ModelClient): +# """ +# Summarize the chat history using the model endpoint and returning the response. +# """ +# summarization_system_prompt = f""" +# You are a helpful assistant that is able to review the chat history between a set of agents (userproxy agents, assistants etc) as they try to address a given TASK and provide a summary. Be SUCCINCT but also comprehensive enough to allow others (who cannot see the chat history) understand and recreate the solution. + +# The task requested by the user is: +# === +# {task} +# === +# The summary should focus on extracting the actual solution to the task from the chat history (assuming the task was addressed) such that any other agent reading the summary will understand what the actual solution is. Use a neutral tone and DO NOT directly mention the agents. Instead only focus on the actions that were carried out (e.g. do not say 'assistant agent generated some code visualization code ..' instead say say 'visualization code was generated ..'. The answer should be framed as a response to the user task. E.g. if the task is "What is the height of the Eiffel tower", the summary should be "The height of the Eiffel Tower is ..."). +# """ +# summarization_prompt = [ +# { +# "role": "system", +# "content": summarization_system_prompt, +# }, +# { +# "role": "user", +# "content": f"Summarize the following chat history. {str(messages)}", +# }, +# ] +# response = client.create(messages=summarization_prompt, cache_seed=None) +# return response.choices[0].message.content diff --git a/python/packages/autogen-studio/autogenstudio/web/app.py b/python/packages/autogen-studio/autogenstudio/web/app.py index d86e2dc439fd..8a0c2ce19ab4 100644 --- a/python/packages/autogen-studio/autogenstudio/web/app.py +++ b/python/packages/autogen-studio/autogenstudio/web/app.py @@ -1,93 +1,63 @@ -import asyncio +# api/app.py import os -import queue -import threading -import traceback -from contextlib import asynccontextmanager -from typing import Any, Union - -from fastapi import FastAPI, WebSocket, WebSocketDisconnect +# import logging +from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles +from contextlib import asynccontextmanager +from typing import AsyncGenerator from loguru import logger -from openai import OpenAIError - -from ..chatmanager import AutoGenChatManager -from ..database import workflow_from_id -from ..database.dbmanager import DBManager -from ..datamodel import Agent, Message, Model, Response, Session, Skill, Workflow -from ..profiler import Profiler -from ..utils import check_and_cast_datetime_fields, init_app_folders, sha256_hash, test_model -from ..version import VERSION -from ..websocket_connection_manager import WebSocketConnectionManager - -profiler = Profiler() -managers = {"chat": None} # manage calls to autogen -# Create thread-safe queue for messages between api thread and autogen threads -message_queue = queue.Queue() -active_connections = [] -active_connections_lock = asyncio.Lock() -websocket_manager = WebSocketConnectionManager( - active_connections=active_connections, - active_connections_lock=active_connections_lock, -) - - -def message_handler(): - while True: - message = message_queue.get() - logger.info( - "** Processing Agent Message on Queue: Active Connections: " - + str([client_id for _, client_id in websocket_manager.active_connections]) - + " **" - ) - for connection, socket_client_id in websocket_manager.active_connections: - if message["connection_id"] == socket_client_id: - logger.info( - f"Sending message to connection_id: {message['connection_id']}. Connection ID: {socket_client_id}" - ) - asyncio.run(websocket_manager.send_message(message, connection)) - else: - logger.info( - f"Skipping message for connection_id: {message['connection_id']}. Connection ID: {socket_client_id}" - ) - message_queue.task_done() +from .routes import sessions, runs, teams, agents, models, tools, ws +from .deps import init_managers, cleanup_managers +from .config import settings +from .initialization import AppInitializer +from ..version import VERSION -message_handler_thread = threading.Thread(target=message_handler, daemon=True) -message_handler_thread.start() +# Configure logging +# logger = logging.getLogger(__name__) +# logging.basicConfig(level=logging.INFO) +# Initialize application app_file_path = os.path.dirname(os.path.abspath(__file__)) -folders = init_app_folders(app_file_path) -ui_folder_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "ui") - -database_engine_uri = folders["database_engine_uri"] -dbmanager = DBManager(engine_uri=database_engine_uri) - -HUMAN_INPUT_TIMEOUT_SECONDS = 180 +initializer = AppInitializer(settings, app_file_path) @asynccontextmanager -async def lifespan(app: FastAPI): - print("***** App started *****") - managers["chat"] = AutoGenChatManager( - message_queue=message_queue, - websocket_manager=websocket_manager, - human_input_timeout=HUMAN_INPUT_TIMEOUT_SECONDS, - ) - dbmanager.create_db_and_tables() +async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]: + """ + Lifecycle manager for the FastAPI application. + Handles initialization and cleanup of application resources. + """ + # Startup + logger.info("Initializing application...") + try: + # Initialize managers (DB, Connection, Team) + await init_managers(initializer.database_uri, initializer.config_dir) + logger.info("Managers initialized successfully") + + # Any other initialization code + logger.info("Application startup complete") - yield - # Close all active connections - await websocket_manager.disconnect_all() - print("***** App stopped *****") + except Exception as e: + logger.error(f"Failed to initialize application: {str(e)}") + raise + yield # Application runs here -app = FastAPI(lifespan=lifespan) + # Shutdown + try: + logger.info("Cleaning up application resources...") + await cleanup_managers() + logger.info("Application shutdown complete") + except Exception as e: + logger.error(f"Error during shutdown: {str(e)}") +# Create FastAPI application +app = FastAPI(lifespan=lifespan, debug=True) -# allow cross origin requests for testing on localhost:800* ports only +# CORS middleware configuration app.add_middleware( CORSMiddleware, allow_origins=[ @@ -101,412 +71,114 @@ async def lifespan(app: FastAPI): allow_headers=["*"], ) -show_docs = os.environ.get("AUTOGENSTUDIO_API_DOCS", "False").lower() == "true" -docs_url = "/docs" if show_docs else None +# Create API router with version and documentation api = FastAPI( root_path="/api", title="AutoGen Studio API", version=VERSION, - docs_url=docs_url, - description="AutoGen Studio is a low-code tool for building and testing multi-agent workflows using AutoGen.", + description="AutoGen Studio is a low-code tool for building and testing multi-agent workflows.", + docs_url="/docs" if settings.API_DOCS else None, ) -# mount an api route such that the main route serves the ui and the /api -app.mount("/api", api) -app.mount("/", StaticFiles(directory=ui_folder_path, html=True), name="ui") -api.mount( - "/files", - StaticFiles(directory=folders["files_static_root"], html=True), - name="files", +# Include all routers with their prefixes +api.include_router( + sessions.router, + prefix="/sessions", + tags=["sessions"], + responses={404: {"description": "Not found"}}, ) +api.include_router( + runs.router, + prefix="/runs", + tags=["runs"], + responses={404: {"description": "Not found"}}, +) -# manage websocket connections - - -def create_entity(model: Any, model_class: Any, filters: dict = None): - """Create a new entity""" - model = check_and_cast_datetime_fields(model) - try: - response: Response = dbmanager.upsert(model) - return response.model_dump(mode="json") - - except Exception as ex_error: - print(ex_error) - return { - "status": False, - "message": f"Error occurred while creating {model_class.__name__}: " + str(ex_error), - } - - -def list_entity( - model_class: Any, - filters: dict = None, - return_json: bool = True, - order: str = "desc", -): - """List all entities for a user""" - return dbmanager.get(model_class, filters=filters, return_json=return_json, order=order) - - -def delete_entity(model_class: Any, filters: dict = None): - """Delete an entity""" - - return dbmanager.delete(filters=filters, model_class=model_class) - - -@api.get("/skills") -async def list_skills(user_id: str): - """List all skills for a user""" - filters = {"user_id": user_id} - return list_entity(Skill, filters=filters) - - -@api.post("/skills") -async def create_skill(skill: Skill): - """Create a new skill""" - filters = {"user_id": skill.user_id} - return create_entity(skill, Skill, filters=filters) - - -@api.delete("/skills/delete") -async def delete_skill(skill_id: int, user_id: str): - """Delete a skill""" - filters = {"id": skill_id, "user_id": user_id} - return delete_entity(Skill, filters=filters) - - -@api.get("/models") -async def list_models(user_id: str): - """List all models for a user""" - filters = {"user_id": user_id} - return list_entity(Model, filters=filters) - - -@api.post("/models") -async def create_model(model: Model): - """Create a new model""" - return create_entity(model, Model) - - -@api.post("/models/test") -async def test_model_endpoint(model: Model): - """Test a model""" - try: - response = test_model(model) - return { - "status": True, - "message": "Model tested successfully", - "data": response, - } - except (OpenAIError, Exception) as ex_error: - return { - "status": False, - "message": "Error occurred while testing model: " + str(ex_error), - } - - -@api.delete("/models/delete") -async def delete_model(model_id: int, user_id: str): - """Delete a model""" - filters = {"id": model_id, "user_id": user_id} - return delete_entity(Model, filters=filters) - - -@api.get("/agents") -async def list_agents(user_id: str): - """List all agents for a user""" - filters = {"user_id": user_id} - return list_entity(Agent, filters=filters) - - -@api.post("/agents") -async def create_agent(agent: Agent): - """Create a new agent""" - return create_entity(agent, Agent) - - -@api.delete("/agents/delete") -async def delete_agent(agent_id: int, user_id: str): - """Delete an agent""" - filters = {"id": agent_id, "user_id": user_id} - return delete_entity(Agent, filters=filters) - - -@api.post("/agents/link/model/{agent_id}/{model_id}") -async def link_agent_model(agent_id: int, model_id: int): - """Link a model to an agent""" - return dbmanager.link(link_type="agent_model", primary_id=agent_id, secondary_id=model_id) - - -@api.delete("/agents/link/model/{agent_id}/{model_id}") -async def unlink_agent_model(agent_id: int, model_id: int): - """Unlink a model from an agent""" - return dbmanager.unlink(link_type="agent_model", primary_id=agent_id, secondary_id=model_id) - - -@api.get("/agents/link/model/{agent_id}") -async def get_agent_models(agent_id: int): - """Get all models linked to an agent""" - return dbmanager.get_linked_entities("agent_model", agent_id, return_json=True) - - -@api.post("/agents/link/skill/{agent_id}/{skill_id}") -async def link_agent_skill(agent_id: int, skill_id: int): - """Link an a skill to an agent""" - return dbmanager.link(link_type="agent_skill", primary_id=agent_id, secondary_id=skill_id) - - -@api.delete("/agents/link/skill/{agent_id}/{skill_id}") -async def unlink_agent_skill(agent_id: int, skill_id: int): - """Unlink an a skill from an agent""" - return dbmanager.unlink(link_type="agent_skill", primary_id=agent_id, secondary_id=skill_id) - - -@api.get("/agents/link/skill/{agent_id}") -async def get_agent_skills(agent_id: int): - """Get all skills linked to an agent""" - return dbmanager.get_linked_entities("agent_skill", agent_id, return_json=True) - - -@api.post("/agents/link/agent/{primary_agent_id}/{secondary_agent_id}") -async def link_agent_agent(primary_agent_id: int, secondary_agent_id: int): - """Link an agent to another agent""" - return dbmanager.link( - link_type="agent_agent", - primary_id=primary_agent_id, - secondary_id=secondary_agent_id, - ) - - -@api.delete("/agents/link/agent/{primary_agent_id}/{secondary_agent_id}") -async def unlink_agent_agent(primary_agent_id: int, secondary_agent_id: int): - """Unlink an agent from another agent""" - return dbmanager.unlink( - link_type="agent_agent", - primary_id=primary_agent_id, - secondary_id=secondary_agent_id, - ) - - -@api.get("/agents/link/agent/{agent_id}") -async def get_linked_agents(agent_id: int): - """Get all agents linked to an agent""" - return dbmanager.get_linked_entities("agent_agent", agent_id, return_json=True) - - -@api.get("/workflows") -async def list_workflows(user_id: str): - """List all workflows for a user""" - filters = {"user_id": user_id} - return list_entity(Workflow, filters=filters) - - -@api.get("/workflows/{workflow_id}") -async def get_workflow(workflow_id: int, user_id: str): - """Get a workflow""" - filters = {"id": workflow_id, "user_id": user_id} - return list_entity(Workflow, filters=filters) - - -@api.get("/workflows/export/{workflow_id}") -async def export_workflow(workflow_id: int, user_id: str): - """Export a user workflow""" - response = Response(message="Workflow exported successfully", status=True, data=None) - try: - workflow_details = workflow_from_id(workflow_id, dbmanager=dbmanager) - response.data = workflow_details - except Exception as ex_error: - response.message = "Error occurred while exporting workflow: " + str(ex_error) - response.status = False - return response.model_dump(mode="json") - - -@api.post("/workflows") -async def create_workflow(workflow: Workflow): - """Create a new workflow""" - return create_entity(workflow, Workflow) - - -@api.delete("/workflows/delete") -async def delete_workflow(workflow_id: int, user_id: str): - """Delete a workflow""" - filters = {"id": workflow_id, "user_id": user_id} - return delete_entity(Workflow, filters=filters) - - -@api.post("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}") -async def link_workflow_agent(workflow_id: int, agent_id: int, agent_type: str): - """Link an agent to a workflow""" - return dbmanager.link( - link_type="workflow_agent", - primary_id=workflow_id, - secondary_id=agent_id, - agent_type=agent_type, - ) - - -@api.post("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}/{sequence_id}") -async def link_workflow_agent_sequence(workflow_id: int, agent_id: int, agent_type: str, sequence_id: int): - """Link an agent to a workflow""" - print("Sequence ID: ", sequence_id) - return dbmanager.link( - link_type="workflow_agent", - primary_id=workflow_id, - secondary_id=agent_id, - agent_type=agent_type, - sequence_id=sequence_id, - ) - - -@api.delete("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}") -async def unlink_workflow_agent(workflow_id: int, agent_id: int, agent_type: str): - """Unlink an agent from a workflow""" - return dbmanager.unlink( - link_type="workflow_agent", - primary_id=workflow_id, - secondary_id=agent_id, - agent_type=agent_type, - ) - - -@api.delete("/workflows/link/agent/{workflow_id}/{agent_id}/{agent_type}/{sequence_id}") -async def unlink_workflow_agent_sequence(workflow_id: int, agent_id: int, agent_type: str, sequence_id: int): - """Unlink an agent from a workflow sequence""" - return dbmanager.unlink( - link_type="workflow_agent", - primary_id=workflow_id, - secondary_id=agent_id, - agent_type=agent_type, - sequence_id=sequence_id, - ) - - -@api.get("/workflows/link/agent/{workflow_id}") -async def get_linked_workflow_agents(workflow_id: int): - """Get all agents linked to a workflow""" - return dbmanager.get_linked_entities( - link_type="workflow_agent", - primary_id=workflow_id, - return_json=True, - ) - - -@api.get("/profiler/{message_id}") -async def profile_agent_task_run(message_id: int): - """Profile an agent task run""" - try: - agent_message = dbmanager.get(Message, filters={"id": message_id}).data[0] - - profile = profiler.profile(agent_message) - return { - "status": True, - "message": "Agent task run profiled successfully", - "data": profile, - } - except Exception as ex_error: - return { - "status": False, - "message": "Error occurred while profiling agent task run: " + str(ex_error), - } - - -@api.get("/sessions") -async def list_sessions(user_id: str): - """List all sessions for a user""" - filters = {"user_id": user_id} - return list_entity(Session, filters=filters) - - -@api.post("/sessions") -async def create_session(session: Session): - """Create a new session""" - return create_entity(session, Session) +api.include_router( + teams.router, + prefix="/teams", + tags=["teams"], + responses={404: {"description": "Not found"}}, +) +api.include_router( + agents.router, + prefix="/agents", + tags=["agents"], + responses={404: {"description": "Not found"}}, +) -@api.delete("/sessions/delete") -async def delete_session(session_id: int, user_id: str): - """Delete a session""" - filters = {"id": session_id, "user_id": user_id} - return delete_entity(Session, filters=filters) +api.include_router( + models.router, + prefix="/models", + tags=["models"], + responses={404: {"description": "Not found"}}, +) +api.include_router( + tools.router, + prefix="/tools", + tags=["tools"], + responses={404: {"description": "Not found"}}, +) -@api.get("/sessions/{session_id}/messages") -async def list_messages(user_id: str, session_id: int): - """List all messages for a use session""" - filters = {"user_id": user_id, "session_id": session_id} - return list_entity(Message, filters=filters, order="asc", return_json=True) +api.include_router( + ws.router, + prefix="/ws", + tags=["websocket"], + responses={404: {"description": "Not found"}}, +) -@api.post("/sessions/{session_id}/workflow/{workflow_id}/run") -async def run_session_workflow(message: Message, session_id: int, workflow_id: int): - """Runs a workflow on provided message""" - try: - user_message_history = ( - dbmanager.get( - Message, - filters={"user_id": message.user_id, "session_id": message.session_id}, - return_json=True, - ).data - if session_id is not None - else [] - ) - # save incoming message - dbmanager.upsert(message) - user_dir = os.path.join(folders["files_static_root"], "user", sha256_hash(message.user_id)) - os.makedirs(user_dir, exist_ok=True) - workflow = workflow_from_id(workflow_id, dbmanager=dbmanager) - agent_response: Message = await managers["chat"].a_chat( - message=message, - history=user_message_history, - user_dir=user_dir, - workflow=workflow, - connection_id=message.connection_id, - ) - - response: Response = dbmanager.upsert(agent_response) - return response.model_dump(mode="json") - except Exception as ex_error: - return { - "status": False, - "message": "Error occurred while processing message: " + str(ex_error), - } +# Version endpoint @api.get("/version") async def get_version(): + """Get API version""" return { "status": True, "message": "Version retrieved successfully", "data": {"version": VERSION}, } +# Health check endpoint + + +@api.get("/health") +async def health_check(): + """API health check endpoint""" + return { + "status": True, + "message": "Service is healthy", + } + +# Mount static file directories +app.mount("/api", api) +app.mount( + "/files", + StaticFiles(directory=initializer.static_root, html=True), + name="files", +) +app.mount("/", StaticFiles(directory=initializer.ui_root, html=True), name="ui") -# websockets +# Error handlers -async def process_socket_message(data: dict, websocket: WebSocket, client_id: str): - print(f"Client says: {data['type']}") - if data["type"] == "user_message": - user_message = Message(**data["data"]) - session_id = data["data"].get("session_id", None) - workflow_id = data["data"].get("workflow_id", None) - response = await run_session_workflow(message=user_message, session_id=session_id, workflow_id=workflow_id) - response_socket_message = { - "type": "agent_response", - "data": response, - "connection_id": client_id, - } - await websocket_manager.send_message(response_socket_message, websocket) +@app.exception_handler(500) +async def internal_error_handler(request, exc): + logger.error(f"Internal error: {str(exc)}") + return { + "status": False, + "message": "Internal server error", + "detail": str(exc) if settings.API_DOCS else "Internal server error" + } -@api.websocket("/ws/{client_id}") -async def websocket_endpoint(websocket: WebSocket, client_id: str): - await websocket_manager.connect(websocket, client_id) - try: - while True: - data = await websocket.receive_json() - await process_socket_message(data, websocket, client_id) - except WebSocketDisconnect: - print(f"Client #{client_id} is disconnected") - await websocket_manager.disconnect(websocket) +def create_app() -> FastAPI: + """ + Factory function to create and configure the FastAPI application. + Useful for testing and different deployment scenarios. + """ + return app diff --git a/python/packages/autogen-studio/autogenstudio/web/config.py b/python/packages/autogen-studio/autogenstudio/web/config.py new file mode 100644 index 000000000000..128edada9bf3 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/config.py @@ -0,0 +1,18 @@ +# api/config.py +from pydantic_settings import BaseSettings + + +class Settings(BaseSettings): + DATABASE_URI: str = "sqlite:///./autogen.db" + API_DOCS: bool = False + CLEANUP_INTERVAL: int = 300 # 5 minutes + SESSION_TIMEOUT: int = 3600 # 1 hour + CONFIG_DIR: str = "configs" # Default config directory relative to app_root + DEFAULT_USER_ID: str = "guestuser@gmail.com" + UPGRADE_DATABASE: bool = False + + class Config: + env_prefix = "AUTOGENSTUDIO_" + + +settings = Settings() diff --git a/python/packages/autogen-studio/autogenstudio/web/deps.py b/python/packages/autogen-studio/autogenstudio/web/deps.py new file mode 100644 index 000000000000..b4c08e952aeb --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/deps.py @@ -0,0 +1,201 @@ +# api/deps.py +from typing import Optional +from fastapi import Depends, HTTPException, status +import logging +from contextlib import contextmanager + +from ..database import DatabaseManager +from .managers.connection import WebSocketManager +from ..teammanager import TeamManager +from .config import settings +from ..database import ConfigurationManager + +logger = logging.getLogger(__name__) + +# Global manager instances +_db_manager: Optional[DatabaseManager] = None +_websocket_manager: Optional[WebSocketManager] = None +_team_manager: Optional[TeamManager] = None + +# Context manager for database sessions + + +@contextmanager +def get_db_context(): + """Provide a transactional scope around a series of operations.""" + if not _db_manager: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Database manager not initialized" + ) + try: + yield _db_manager + except Exception as e: + logger.error(f"Database operation failed: {str(e)}") + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Database operation failed" + ) + +# Dependency providers + + +async def get_db() -> DatabaseManager: + """Dependency provider for database manager""" + if not _db_manager: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Database manager not initialized" + ) + return _db_manager + + +async def get_websocket_manager() -> WebSocketManager: + """Dependency provider for connection manager""" + if not _websocket_manager: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Connection manager not initialized" + ) + return _websocket_manager + + +async def get_team_manager() -> TeamManager: + """Dependency provider for team manager""" + if not _team_manager: + raise HTTPException( + status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, + detail="Team manager not initialized" + ) + return _team_manager + +# Authentication dependency + + +async def get_current_user( + # Add your authentication logic here + # For example: token: str = Depends(oauth2_scheme) +) -> str: + """ + Dependency for getting the current authenticated user. + Replace with your actual authentication logic. + """ + # Implement your user authentication here + return "user_id" # Replace with actual user identification + +# Manager initialization and cleanup + + +async def init_managers(database_uri: str, config_dir: str) -> None: + """Initialize all manager instances""" + global _db_manager, _websocket_manager, _team_manager + + logger.info("Initializing managers...") + + try: + # Initialize database manager + _db_manager = DatabaseManager( + engine_uri=database_uri, auto_upgrade=settings.UPGRADE_DATABASE) + _db_manager.create_db_and_tables() + + # init default team config + + _team_config_manager = ConfigurationManager(db_manager=_db_manager) + import_result = await _team_config_manager.import_directory( + config_dir, settings.DEFAULT_USER_ID, check_exists=True) + + # Initialize connection manager + _websocket_manager = WebSocketManager( + db_manager=_db_manager + ) + logger.info("Connection manager initialized") + + # Initialize team manager + _team_manager = TeamManager() + logger.info("Team manager initialized") + + except Exception as e: + logger.error(f"Failed to initialize managers: {str(e)}") + await cleanup_managers() # Cleanup any partially initialized managers + raise + + +async def cleanup_managers() -> None: + """Cleanup and shutdown all manager instances""" + global _db_manager, _websocket_manager, _team_manager + + logger.info("Cleaning up managers...") + + # Cleanup connection manager first to ensure all active connections are closed + if _websocket_manager: + try: + await _websocket_manager.cleanup() + except Exception as e: + logger.error(f"Error cleaning up connection manager: {str(e)}") + finally: + _websocket_manager = None + + # TeamManager doesn't need explicit cleanup since WebSocketManager handles it + _team_manager = None + + # Cleanup database manager last + if _db_manager: + try: + await _db_manager.close() + except Exception as e: + logger.error(f"Error cleaning up database manager: {str(e)}") + finally: + _db_manager = None + + logger.info("All managers cleaned up") + +# Utility functions for dependency management + + +def get_manager_status() -> dict: + """Get the initialization status of all managers""" + return { + "database_manager": _db_manager is not None, + "websocket_manager": _websocket_manager is not None, + "team_manager": _team_manager is not None + } + +# Combined dependencies + + +async def get_managers(): + """Get all managers in one dependency""" + return { + "db": await get_db(), + "connection": await get_websocket_manager(), + "team": await get_team_manager() + } + +# Error handling for manager operations + + +class ManagerOperationError(Exception): + """Custom exception for manager operation errors""" + + def __init__(self, manager_name: str, operation: str, detail: str): + self.manager_name = manager_name + self.operation = operation + self.detail = detail + super().__init__(f"{manager_name} failed during {operation}: {detail}") + +# Dependency for requiring specific managers + + +def require_managers(*manager_names: str): + """Decorator to require specific managers for a route""" + async def dependency(): + status = get_manager_status() + missing = [name for name in manager_names if not status.get( + f"{name}_manager")] + if missing: + raise HTTPException( + status_code=status.HTTP_503_SERVICE_UNAVAILABLE, + detail=f"Required managers not available: {', '.join(missing)}" + ) + return True + return Depends(dependency) diff --git a/python/packages/autogen-studio/autogenstudio/web/initialization.py b/python/packages/autogen-studio/autogenstudio/web/initialization.py new file mode 100644 index 000000000000..e938d222062d --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/initialization.py @@ -0,0 +1,110 @@ +# api/initialization.py +import os +from pathlib import Path +from typing import Dict +from pydantic import BaseModel +from loguru import logger +from dotenv import load_dotenv + +from .config import Settings + + +class _AppPaths(BaseModel): + """Internal model representing all application paths""" + app_root: Path + static_root: Path + user_files: Path + ui_root: Path + config_dir: Path + database_uri: str + + +class AppInitializer: + """Handles application initialization including paths and environment setup""" + + def __init__(self, settings: Settings, app_path: str): + """ + Initialize the application structure. + + Args: + settings: Application settings + app_path: Path to the application code directory + """ + self.settings = settings + self._app_path = Path(app_path) + self._paths = self._init_paths() + self._create_directories() + self._load_environment() + logger.info(f"Initialized application data folder: {self.app_root}") + + def _get_app_root(self) -> Path: + """Determine application root directory""" + if app_dir := os.getenv("AUTOGENSTUDIO_APPDIR"): + return Path(app_dir) + return Path.home() / ".autogenstudio" + + def _get_database_uri(self, app_root: Path) -> str: + """Generate database URI based on settings or environment""" + if db_uri := os.getenv("AUTOGENSTUDIO_DATABASE_URI"): + return db_uri + return self.settings.DATABASE_URI.replace( + "./", str(app_root) + "/" + ) + + def _init_paths(self) -> _AppPaths: + """Initialize and return AppPaths instance""" + app_root = self._get_app_root() + return _AppPaths( + app_root=app_root, + static_root=app_root / "files", + user_files=app_root / "files" / "user", + ui_root=self._app_path / "ui", + config_dir=app_root / self.settings.CONFIG_DIR, + database_uri=self._get_database_uri(app_root) + ) + + def _create_directories(self) -> None: + """Create all required directories""" + self.app_root.mkdir(parents=True, exist_ok=True) + dirs = [self.static_root, self.user_files, + self.ui_root, self.config_dir] + for path in dirs: + path.mkdir(parents=True, exist_ok=True) + + def _load_environment(self) -> None: + """Load environment variables from .env file if it exists""" + env_file = self.app_root / ".env" + if env_file.exists(): + logger.info(f"Loading environment variables from {env_file}") + load_dotenv(str(env_file)) + + # Properties for accessing paths + @property + def app_root(self) -> Path: + """Root directory for the application""" + return self._paths.app_root + + @property + def static_root(self) -> Path: + """Directory for static files""" + return self._paths.static_root + + @property + def user_files(self) -> Path: + """Directory for user files""" + return self._paths.user_files + + @property + def ui_root(self) -> Path: + """Directory for UI files""" + return self._paths.ui_root + + @property + def config_dir(self) -> Path: + """Directory for configuration files""" + return self._paths.config_dir + + @property + def database_uri(self) -> str: + """Database connection URI""" + return self._paths.database_uri diff --git a/python/packages/autogen-studio/autogenstudio/database/migrations/__init__.py b/python/packages/autogen-studio/autogenstudio/web/managers/__init__.py similarity index 100% rename from python/packages/autogen-studio/autogenstudio/database/migrations/__init__.py rename to python/packages/autogen-studio/autogenstudio/web/managers/__init__.py diff --git a/python/packages/autogen-studio/autogenstudio/web/managers/connection.py b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py new file mode 100644 index 000000000000..b46887b7ac1a --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py @@ -0,0 +1,247 @@ +# managers/connection.py +from fastapi import WebSocket, WebSocketDisconnect +from typing import Dict, Optional, Any +from uuid import UUID +import logging +from datetime import datetime, timezone + +from ...datamodel import Run, RunStatus, TeamResult +from ...database import DatabaseManager +from autogen_agentchat.messages import InnerMessage, ChatMessage +from autogen_core.base import CancellationToken + +logger = logging.getLogger(__name__) + + +class WebSocketManager: + """Manages WebSocket connections and message streaming for team task execution""" + + def __init__(self, db_manager: DatabaseManager): + self.db_manager = db_manager + self._connections: Dict[UUID, WebSocket] = {} + self._cancellation_tokens: Dict[UUID, CancellationToken] = {} + + async def connect(self, websocket: WebSocket, run_id: UUID) -> bool: + """Initialize WebSocket connection for a run + + Args: + websocket: The WebSocket connection to initialize + run_id: UUID of the run to associate with this connection + + Returns: + bool: True if connection was successful, False otherwise + """ + try: + await websocket.accept() + self._connections[run_id] = websocket + + run = await self._get_run(run_id) + if run: + run.status = RunStatus.ACTIVE + self.db_manager.upsert(run) + + await self._send_message(run_id, { + "type": "system", + "status": "connected", + "timestamp": datetime.now(timezone.utc).isoformat() + }) + + return True + + except Exception as e: + logger.error(f"Connection error for run {run_id}: {e}") + return False + + async def start_stream( + self, + run_id: UUID, + team_manager: Any, + task: str, + team_config: dict + ) -> None: + """Start streaming task execution + + Args: + run_id: UUID of the run + team_manager: Instance of the team manager + task: Task string to execute + team_config: Team configuration dictionary + """ + if run_id not in self._connections: + raise ValueError(f"No active connection for run {run_id}") + + cancellation_token = CancellationToken() + self._cancellation_tokens[run_id] = cancellation_token + + try: + async for message in team_manager.run_stream( + task=task, + team_config=team_config, + cancellation_token=cancellation_token + ): + + if cancellation_token.is_cancelled(): + logger.info(f"Stream cancelled for run {run_id}") + break + + formatted_message = self._format_message(message) + if formatted_message: + await self._send_message(run_id, formatted_message) + + # Only send completion if not cancelled + if not cancellation_token.is_cancelled(): + # await self._send_message(run_id, { + # "type": "completion", + # "status": "complete", + # "timestamp": datetime.now(timezone.utc).isoformat() + # }) + await self._update_run_status(run_id, RunStatus.COMPLETE) + else: + await self._send_message(run_id, { + "type": "completion", + "status": "cancelled", + "timestamp": datetime.now(timezone.utc).isoformat() + }) + await self._update_run_status(run_id, RunStatus.STOPPED) + + except Exception as e: + logger.error(f"Stream error for run {run_id}: {e}") + await self._handle_stream_error(run_id, e) + + finally: + self._cancellation_tokens.pop(run_id, None) + + async def stop_run(self, run_id: UUID) -> None: + """Stop a running task""" + if run_id in self._cancellation_tokens: + logger.info(f"Stopping run {run_id}") + self._cancellation_tokens[run_id].cancel() + + # Send final message if connection still exists + if run_id in self._connections: + try: + await self._send_message(run_id, { + "type": "completion", + "status": "cancelled", + "timestamp": datetime.now(timezone.utc).isoformat() + }) + except Exception: + pass + + async def disconnect(self, run_id: UUID) -> None: + """Clean up connection and associated resources""" + logger.info(f"Disconnecting run {run_id}") + + # First cancel any running tasks + await self.stop_run(run_id) + + # Then clean up resources without trying to close the socket again + if run_id in self._connections: + self._connections.pop(run_id, None) + self._cancellation_tokens.pop(run_id, None) + + async def _send_message(self, run_id: UUID, message: dict) -> None: + """Send a message through the WebSocket + + Args: + run_id: UUID of the run + message: Message dictionary to send + """ + try: + if run_id in self._connections: + await self._connections[run_id].send_json(message) + except WebSocketDisconnect: + logger.warning( + f"WebSocket disconnected while sending message for run {run_id}") + await self.disconnect(run_id) + except Exception as e: + logger.error(f"Error sending message for run {run_id}: {e}") + await self._handle_stream_error(run_id, e) + + async def _handle_stream_error(self, run_id: UUID, error: Exception) -> None: + """Handle stream errors consistently + + Args: + run_id: UUID of the run + error: Exception that occurred + """ + try: + await self._send_message(run_id, { + "type": "completion", + "status": "error", + "error": str(error), + "timestamp": datetime.now(timezone.utc).isoformat() + }) + except Exception as send_error: + logger.error( + f"Failed to send error message for run {run_id}: {send_error}") + + await self._update_run_status(run_id, RunStatus.ERROR, str(error)) + + def _format_message(self, message: Any) -> Optional[dict]: + """Format message for WebSocket transmission + + Args: + message: Message to format + + Returns: + Optional[dict]: Formatted message or None if formatting fails + """ + try: + if isinstance(message, (InnerMessage, ChatMessage)): + return { + "type": "message", + "data": message.model_dump() + } + elif isinstance(message, TeamResult): + return { + "type": "result", + "data": message.model_dump(), + "status": "complete", + } + return None + except Exception as e: + logger.error(f"Message formatting error: {e}") + return None + + async def _get_run(self, run_id: UUID) -> Optional[Run]: + """Get run from database + + Args: + run_id: UUID of the run to retrieve + + Returns: + Optional[Run]: Run object if found, None otherwise + """ + response = self.db_manager.get( + Run, filters={"id": run_id}, return_json=False) + return response.data[0] if response.status and response.data else None + + async def _update_run_status( + self, + run_id: UUID, + status: RunStatus, + error: Optional[str] = None + ) -> None: + """Update run status in database + + Args: + run_id: UUID of the run to update + status: New status to set + error: Optional error message + """ + run = await self._get_run(run_id) + if run: + run.status = status + run.error_message = error + self.db_manager.upsert(run) + + @property + def active_connections(self) -> set[UUID]: + """Get set of active run IDs""" + return set(self._connections.keys()) + + @property + def active_runs(self) -> set[UUID]: + """Get set of runs with active cancellation tokens""" + return set(self._cancellation_tokens.keys()) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/__init__.py b/python/packages/autogen-studio/autogenstudio/web/routes/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/agents.py b/python/packages/autogen-studio/autogenstudio/web/routes/agents.py new file mode 100644 index 000000000000..183dbf2a5bee --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/agents.py @@ -0,0 +1,181 @@ +# api/routes/agents.py +from fastapi import APIRouter, Depends, HTTPException +from typing import Dict +from ..deps import get_db +from ...datamodel import Agent, Model, Tool + +router = APIRouter() + + +@router.get("/") +async def list_agents( + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all agents for a user""" + response = db.get(Agent, filters={"user_id": user_id}) + return { + "status": True, + "data": response.data + } + + +@router.get("/{agent_id}") +async def get_agent( + agent_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Get a specific agent""" + response = db.get( + Agent, + filters={"id": agent_id, "user_id": user_id} + ) + if not response.status or not response.data: + raise HTTPException(status_code=404, detail="Agent not found") + return { + "status": True, + "data": response.data[0] + } + + +@router.post("/") +async def create_agent( + agent: Agent, + db=Depends(get_db) +) -> Dict: + """Create a new agent""" + response = db.upsert(agent) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + return { + "status": True, + "data": response.data + } + + +@router.delete("/{agent_id}") +async def delete_agent( + agent_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Delete an agent""" + response = db.delete( + filters={"id": agent_id, "user_id": user_id}, + model_class=Agent + ) + return { + "status": True, + "message": "Agent deleted successfully" + } + +# Agent-Model link endpoints + + +@router.post("/{agent_id}/models/{model_id}") +async def link_agent_model( + agent_id: int, + model_id: int, + db=Depends(get_db) +) -> Dict: + """Link a model to an agent""" + response = db.link( + link_type="agent_model", + primary_id=agent_id, + secondary_id=model_id + ) + return { + "status": True, + "message": "Model linked to agent successfully" + } + + +@router.delete("/{agent_id}/models/{model_id}") +async def unlink_agent_model( + agent_id: int, + model_id: int, + db=Depends(get_db) +) -> Dict: + """Unlink a model from an agent""" + response = db.unlink( + link_type="agent_model", + primary_id=agent_id, + secondary_id=model_id + ) + return { + "status": True, + "message": "Model unlinked from agent successfully" + } + + +@router.get("/{agent_id}/models") +async def get_agent_models( + agent_id: int, + db=Depends(get_db) +) -> Dict: + """Get all models linked to an agent""" + response = db.get_linked_entities( + link_type="agent_model", + primary_id=agent_id, + return_json=True + ) + return { + "status": True, + "data": response.data + } + +# Agent-Tool link endpoints + + +@router.post("/{agent_id}/tools/{tool_id}") +async def link_agent_tool( + agent_id: int, + tool_id: int, + db=Depends(get_db) +) -> Dict: + """Link a tool to an agent""" + response = db.link( + link_type="agent_tool", + primary_id=agent_id, + secondary_id=tool_id + ) + return { + "status": True, + "message": "Tool linked to agent successfully" + } + + +@router.delete("/{agent_id}/tools/{tool_id}") +async def unlink_agent_tool( + agent_id: int, + tool_id: int, + db=Depends(get_db) +) -> Dict: + """Unlink a tool from an agent""" + response = db.unlink( + link_type="agent_tool", + primary_id=agent_id, + secondary_id=tool_id + ) + return { + "status": True, + "message": "Tool unlinked from agent successfully" + } + + +@router.get("/{agent_id}/tools") +async def get_agent_tools( + agent_id: int, + db=Depends(get_db) +) -> Dict: + """Get all tools linked to an agent""" + response = db.get_linked_entities( + link_type="agent_tool", + primary_id=agent_id, + return_json=True + ) + return { + "status": True, + "data": response.data + } diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/models.py b/python/packages/autogen-studio/autogenstudio/web/routes/models.py new file mode 100644 index 000000000000..9b57e6255458 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/models.py @@ -0,0 +1,95 @@ +# api/routes/models.py +from fastapi import APIRouter, Depends, HTTPException +from typing import Dict +from openai import OpenAIError +from ..deps import get_db +from ...datamodel import Model +from ...utils import test_model + +router = APIRouter() + + +@router.get("/") +async def list_models( + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all models for a user""" + response = db.get(Model, filters={"user_id": user_id}) + return { + "status": True, + "data": response.data + } + + +@router.get("/{model_id}") +async def get_model( + model_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Get a specific model""" + response = db.get( + Model, + filters={"id": model_id, "user_id": user_id} + ) + if not response.status or not response.data: + raise HTTPException(status_code=404, detail="Model not found") + return { + "status": True, + "data": response.data[0] + } + + +@router.post("/") +async def create_model( + model: Model, + db=Depends(get_db) +) -> Dict: + """Create a new model""" + response = db.upsert(model) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + return { + "status": True, + "data": response.data + } + + +@router.delete("/{model_id}") +async def delete_model( + model_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Delete a model""" + response = db.delete( + filters={"id": model_id, "user_id": user_id}, + model_class=Model + ) + return { + "status": True, + "message": "Model deleted successfully" + } + + +@router.post("/test") +async def test_model_endpoint(model: Model) -> Dict: + """Test a model configuration""" + try: + response = test_model(model) + return { + "status": True, + "message": "Model tested successfully", + "data": response + } + except OpenAIError as e: + raise HTTPException( + status_code=400, + detail=f"OpenAI API error: {str(e)}" + ) + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Error testing model: {str(e)}" + ) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/runs.py b/python/packages/autogen-studio/autogenstudio/web/routes/runs.py new file mode 100644 index 000000000000..7fb5e7475ac0 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/runs.py @@ -0,0 +1,76 @@ +# /api/runs routes +from fastapi import APIRouter, Body, Depends, HTTPException +from uuid import UUID +from typing import Dict + +from pydantic import BaseModel +from ..deps import get_db, get_websocket_manager, get_team_manager +from ...datamodel import Run, Session, Message, Team, RunStatus, MessageConfig + +from ...teammanager import TeamManager +from autogen_core.base import CancellationToken + +router = APIRouter() + + +class CreateRunRequest(BaseModel): + session_id: int + user_id: str + + +@router.post("/") +async def create_run( + request: CreateRunRequest, + db=Depends(get_db), +) -> Dict: + """Create a new run""" + session_response = db.get( + Session, + filters={"id": request.session_id, "user_id": request.user_id}, + return_json=False + ) + if not session_response.status or not session_response.data: + raise HTTPException(status_code=404, detail="Session not found") + + try: + + run = db.upsert(Run(session_id=request.session_id), return_json=False) + return { + "status": run.status, + "data": {"run_id": str(run.data.id)} + } + + # } + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post("/{run_id}/start") +async def start_run( + run_id: UUID, + message: Message = Body(...), + ws_manager=Depends(get_websocket_manager), + team_manager=Depends(get_team_manager), + db=Depends(get_db), +) -> Dict: + """Start streaming task execution""" + + if isinstance(message.config, dict): + message.config = MessageConfig(**message.config) + + session = db.get(Session, filters={ + "id": message.session_id}, return_json=False) + + team = db.get( + Team, filters={"id": session.data[0].team_id}, return_json=False) + + try: + await ws_manager.start_stream(run_id, team_manager, message.config.content, team.data[0].config) + return { + "status": True, + "message": "Stream started successfully", + "data": {"run_id": str(run_id)} + } + + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py b/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py new file mode 100644 index 000000000000..f74ee6288154 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/sessions.py @@ -0,0 +1,114 @@ +# api/routes/sessions.py +from fastapi import APIRouter, Depends, HTTPException +from typing import Dict +from ..deps import get_db +from ...datamodel import Session, Message + +router = APIRouter() + + +@router.get("/") +async def list_sessions( + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all sessions for a user""" + response = db.get(Session, filters={"user_id": user_id}) + return { + "status": True, + "data": response.data + } + + +@router.get("/{session_id}") +async def get_session( + session_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Get a specific session""" + response = db.get( + Session, + filters={"id": session_id, "user_id": user_id} + ) + if not response.status or not response.data: + raise HTTPException(status_code=404, detail="Session not found") + return { + "status": True, + "data": response.data[0] + } + + +@router.post("/") +async def create_session( + session: Session, + db=Depends(get_db) +) -> Dict: + """Create a new session""" + response = db.upsert(session) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + return { + "status": True, + "data": response.data + } + + +@router.put("/{session_id}") +async def update_session( + session_id: int, + user_id: str, + session: Session, + db=Depends(get_db) +) -> Dict: + """Update an existing session""" + # First verify the session belongs to user + existing = db.get( + Session, + filters={"id": session_id, "user_id": user_id} + ) + if not existing.status or not existing.data: + raise HTTPException(status_code=404, detail="Session not found") + + # Update the session + response = db.upsert(session) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + + return { + "status": True, + "data": response.data, + "message": "Session updated successfully" + } + + +@router.delete("/{session_id}") +async def delete_session( + session_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Delete a session""" + response = db.delete( + filters={"id": session_id, "user_id": user_id}, + model_class=Session + ) + return { + "status": True, + "message": "Session deleted successfully" + } + + +@router.get("/{session_id}/messages") +async def list_messages( + session_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all messages for a session""" + filters = {"session_id": session_id, "user_id": user_id} + response = db.get(Message, filters=filters, order="asc") + return { + "status": True, + "data": response.data + } diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/teams.py b/python/packages/autogen-studio/autogenstudio/web/routes/teams.py new file mode 100644 index 000000000000..854c195d3c71 --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/teams.py @@ -0,0 +1,146 @@ +# api/routes/teams.py +from fastapi import APIRouter, Depends, HTTPException +from typing import Dict +from ..deps import get_db +from ...datamodel import Team + +router = APIRouter() + + +@router.get("/") +async def list_teams( + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all teams for a user""" + response = db.get(Team, filters={"user_id": user_id}) + return { + "status": True, + "data": response.data + } + + +@router.get("/{team_id}") +async def get_team( + team_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Get a specific team""" + response = db.get( + Team, + filters={"id": team_id, "user_id": user_id} + ) + if not response.status or not response.data: + raise HTTPException(status_code=404, detail="Team not found") + return { + "status": True, + "data": response.data[0] + } + + +@router.post("/") +async def create_team( + team: Team, + db=Depends(get_db) +) -> Dict: + """Create a new team""" + response = db.upsert(team) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + return { + "status": True, + "data": response.data + } + + +@router.delete("/{team_id}") +async def delete_team( + team_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Delete a team""" + response = db.delete( + filters={"id": team_id, "user_id": user_id}, + model_class=Team + ) + return { + "status": True, + "message": "Team deleted successfully" + } + +# Team-Agent link endpoints + + +@router.post("/{team_id}/agents/{agent_id}") +async def link_team_agent( + team_id: int, + agent_id: int, + db=Depends(get_db) +) -> Dict: + """Link an agent to a team""" + response = db.link( + link_type="team_agent", + primary_id=team_id, + secondary_id=agent_id + ) + return { + "status": True, + "message": "Agent linked to team successfully" + } + + +@router.post("/{team_id}/agents/{agent_id}/{sequence_id}") +async def link_team_agent_sequence( + team_id: int, + agent_id: int, + sequence_id: int, + db=Depends(get_db) +) -> Dict: + """Link an agent to a team with sequence""" + response = db.link( + link_type="team_agent", + primary_id=team_id, + secondary_id=agent_id, + sequence_id=sequence_id + ) + return { + "status": True, + "message": "Agent linked to team with sequence successfully" + } + + +@router.delete("/{team_id}/agents/{agent_id}") +async def unlink_team_agent( + team_id: int, + agent_id: int, + db=Depends(get_db) +) -> Dict: + """Unlink an agent from a team""" + response = db.unlink( + link_type="team_agent", + primary_id=team_id, + secondary_id=agent_id + ) + return { + "status": True, + "message": "Agent unlinked from team successfully" + } + + +@router.get("/{team_id}/agents") +async def get_team_agents( + team_id: int, + db=Depends(get_db) +) -> Dict: + """Get all agents linked to a team""" + response = db.get_linked_entities( + link_type="team_agent", + primary_id=team_id, + return_json=True + ) + return { + "status": True, + "data": response.data + } diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/tools.py b/python/packages/autogen-studio/autogenstudio/web/routes/tools.py new file mode 100644 index 000000000000..d73b626038ad --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/tools.py @@ -0,0 +1,103 @@ +# api/routes/tools.py +from fastapi import APIRouter, Depends, HTTPException +from typing import Dict +from ..deps import get_db +from ...datamodel import Tool + +router = APIRouter() + + +@router.get("/") +async def list_tools( + user_id: str, + db=Depends(get_db) +) -> Dict: + """List all tools for a user""" + response = db.get(Tool, filters={"user_id": user_id}) + return { + "status": True, + "data": response.data + } + + +@router.get("/{tool_id}") +async def get_tool( + tool_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Get a specific tool""" + response = db.get( + Tool, + filters={"id": tool_id, "user_id": user_id} + ) + if not response.status or not response.data: + raise HTTPException(status_code=404, detail="Tool not found") + return { + "status": True, + "data": response.data[0] + } + + +@router.post("/") +async def create_tool( + tool: Tool, + db=Depends(get_db) +) -> Dict: + """Create a new tool""" + response = db.upsert(tool) + if not response.status: + raise HTTPException(status_code=400, detail=response.message) + return { + "status": True, + "data": response.data + } + + +@router.delete("/{tool_id}") +async def delete_tool( + tool_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Delete a tool""" + response = db.delete( + filters={"id": tool_id, "user_id": user_id}, + model_class=Tool + ) + return { + "status": True, + "message": "Tool deleted successfully" + } + + +@router.post("/{tool_id}/test") +async def test_tool( + tool_id: int, + user_id: str, + db=Depends(get_db) +) -> Dict: + """Test a tool configuration""" + # Get tool + tool_response = db.get( + Tool, + filters={"id": tool_id, "user_id": user_id} + ) + if not tool_response.status or not tool_response.data: + raise HTTPException(status_code=404, detail="Tool not found") + + tool = tool_response.data[0] + + try: + # Implement tool testing logic here + # This would depend on the tool type and configuration + return { + "status": True, + "message": "Tool tested successfully", + "data": {"tool_id": tool_id} + } + except Exception as e: + raise HTTPException( + status_code=500, + detail=f"Error testing tool: {str(e)}" + ) diff --git a/python/packages/autogen-studio/autogenstudio/web/routes/ws.py b/python/packages/autogen-studio/autogenstudio/web/routes/ws.py new file mode 100644 index 000000000000..8fd6844ff3ad --- /dev/null +++ b/python/packages/autogen-studio/autogenstudio/web/routes/ws.py @@ -0,0 +1,74 @@ +# api/ws.py +from fastapi import APIRouter, WebSocket, WebSocketDisconnect, Depends, HTTPException +from typing import Dict +from uuid import UUID +import logging +import json +from datetime import datetime + +from ..deps import get_websocket_manager, get_db, get_team_manager +from ...datamodel import Run, RunStatus + +router = APIRouter() +logger = logging.getLogger(__name__) + + +@router.websocket("/runs/{run_id}") +async def run_websocket( + websocket: WebSocket, + run_id: UUID, + ws_manager=Depends(get_websocket_manager), + db=Depends(get_db), + team_manager=Depends(get_team_manager) +): + """WebSocket endpoint for run communication""" + # Verify run exists and is in valid state + run_response = db.get(Run, filters={"id": run_id}, return_json=False) + if not run_response.status or not run_response.data: + await websocket.close(code=4004, reason="Run not found") + return + + run = run_response.data[0] + if run.status not in [RunStatus.CREATED, RunStatus.ACTIVE]: + await websocket.close(code=4003, reason="Run not in valid state") + return + + # Connect websocket + connected = await ws_manager.connect(websocket, run_id) + if not connected: + await websocket.close(code=4002, reason="Failed to establish connection") + return + + try: + logger.info(f"WebSocket connection established for run {run_id}") + + while True: + try: + raw_message = await websocket.receive_text() + message = json.loads(raw_message) + + if message.get("type") == "stop": + logger.info(f"Received stop request for run {run_id}") + await ws_manager.stop_run(run_id) + break + + elif message.get("type") == "ping": + await websocket.send_json({ + "type": "pong", + "timestamp": datetime.utcnow().isoformat() + }) + + except json.JSONDecodeError: + logger.warning(f"Invalid JSON received: {raw_message}") + await websocket.send_json({ + "type": "error", + "error": "Invalid message format", + "timestamp": datetime.utcnow().isoformat() + }) + + except WebSocketDisconnect: + logger.info(f"WebSocket disconnected for run {run_id}") + except Exception as e: + logger.error(f"WebSocket error: {str(e)}") + finally: + await ws_manager.disconnect(run_id) diff --git a/python/packages/autogen-studio/autogenstudio/websocket_connection_manager.py b/python/packages/autogen-studio/autogenstudio/websocket_connection_manager.py deleted file mode 100644 index 73f7ef896811..000000000000 --- a/python/packages/autogen-studio/autogenstudio/websocket_connection_manager.py +++ /dev/null @@ -1,135 +0,0 @@ -import asyncio -from typing import Any, Dict, List, Optional, Tuple, Union - -import websockets -from fastapi import WebSocket, WebSocketDisconnect - - -class WebSocketConnectionManager: - """ - Manages WebSocket connections including sending, broadcasting, and managing the lifecycle of connections. - """ - - def __init__( - self, - active_connections: List[Tuple[WebSocket, str]] = None, - active_connections_lock: asyncio.Lock = None, - ) -> None: - """ - Initializes WebSocketConnectionManager with an optional list of active WebSocket connections. - - :param active_connections: A list of tuples, each containing a WebSocket object and its corresponding client_id. - """ - if active_connections is None: - active_connections = [] - self.active_connections_lock = active_connections_lock - self.active_connections: List[Tuple[WebSocket, str]] = active_connections - - async def connect(self, websocket: WebSocket, client_id: str) -> None: - """ - Accepts a new WebSocket connection and appends it to the active connections list. - - :param websocket: The WebSocket instance representing a client connection. - :param client_id: A string representing the unique identifier of the client. - """ - await websocket.accept() - async with self.active_connections_lock: - self.active_connections.append((websocket, client_id)) - print(f"New Connection: {client_id}, Total: {len(self.active_connections)}") - - async def disconnect(self, websocket: WebSocket) -> None: - """ - Disconnects and removes a WebSocket connection from the active connections list. - - :param websocket: The WebSocket instance to remove. - """ - async with self.active_connections_lock: - try: - self.active_connections = [conn for conn in self.active_connections if conn[0] != websocket] - print(f"Connection Closed. Total: {len(self.active_connections)}") - except ValueError: - print("Error: WebSocket connection not found") - - async def disconnect_all(self) -> None: - """ - Disconnects all active WebSocket connections. - """ - for connection, _ in self.active_connections[:]: - await self.disconnect(connection) - - async def send_message(self, message: Union[Dict, str], websocket: WebSocket) -> None: - """ - Sends a JSON message to a single WebSocket connection. - - :param message: A JSON serializable dictionary containing the message to send. - :param websocket: The WebSocket instance through which to send the message. - """ - try: - async with self.active_connections_lock: - await websocket.send_json(message) - except WebSocketDisconnect: - print("Error: Tried to send a message to a closed WebSocket") - await self.disconnect(websocket) - except websockets.exceptions.ConnectionClosedOK: - print("Error: WebSocket connection closed normally") - await self.disconnect(websocket) - except Exception as e: - print(f"Error in sending message: {str(e)}", message) - await self.disconnect(websocket) - - async def get_input(self, prompt: Union[Dict, str], websocket: WebSocket, timeout: int = 60) -> str: - """ - Sends a JSON message to a single WebSocket connection as a prompt for user input. - Waits on a user response or until the given timeout elapses. - - :param prompt: A JSON serializable dictionary containing the message to send. - :param websocket: The WebSocket instance through which to send the message. - """ - response = "Error: Unexpected response.\nTERMINATE" - try: - async with self.active_connections_lock: - await websocket.send_json(prompt) - result = await asyncio.wait_for(websocket.receive_json(), timeout=timeout) - data = result.get("data") - if data: - response = data.get("content", "Error: Unexpected response format\nTERMINATE") - else: - response = "Error: Unexpected response format\nTERMINATE" - - except asyncio.TimeoutError: - response = f"The user was timed out after {timeout} seconds of inactivity.\nTERMINATE" - except WebSocketDisconnect: - print("Error: Tried to send a message to a closed WebSocket") - await self.disconnect(websocket) - response = "The user was disconnected\nTERMINATE" - except websockets.exceptions.ConnectionClosedOK: - print("Error: WebSocket connection closed normally") - await self.disconnect(websocket) - response = "The user was disconnected\nTERMINATE" - except Exception as e: - print(f"Error in sending message: {str(e)}", prompt) - await self.disconnect(websocket) - response = f"Error: {e}\nTERMINATE" - - return response - - async def broadcast(self, message: Dict) -> None: - """ - Broadcasts a JSON message to all active WebSocket connections. - - :param message: A JSON serializable dictionary containing the message to broadcast. - """ - # Create a message dictionary with the desired format - message_dict = {"message": message} - - for connection, _ in self.active_connections[:]: - try: - if connection.client_state == websockets.protocol.State.OPEN: - # Call send_message method with the message dictionary and current WebSocket connection - await self.send_message(message_dict, connection) - else: - print("Error: WebSocket connection is closed") - await self.disconnect(connection) - except (WebSocketDisconnect, websockets.exceptions.ConnectionClosedOK) as e: - print(f"Error: WebSocket disconnected or closed({str(e)})") - await self.disconnect(connection) diff --git a/python/packages/autogen-studio/autogenstudio/workflowmanager.py b/python/packages/autogen-studio/autogenstudio/workflowmanager.py deleted file mode 100644 index 2da3b58b7cec..000000000000 --- a/python/packages/autogen-studio/autogenstudio/workflowmanager.py +++ /dev/null @@ -1,1066 +0,0 @@ -import json -import os -import time -from datetime import datetime -from typing import Any, Coroutine, Dict, List, Optional, Union - -import autogen - -from .datamodel import ( - Agent, - AgentType, - CodeExecutionConfigTypes, - Message, - SocketMessage, - Workflow, - WorkFlowSummaryMethod, - WorkFlowType, -) -from .utils import ( - clear_folder, - find_key_value, - get_modified_files, - get_skills_prompt, - load_code_execution_config, - sanitize_model, - save_skills_to_file, - summarize_chat_history, -) - - -class AutoWorkflowManager: - """ - WorkflowManager class to load agents from a provided configuration and run a chat between them. - """ - - def __init__( - self, - workflow: Union[Dict, str], - history: Optional[List[Message]] = None, - work_dir: str = None, - clear_work_dir: bool = True, - send_message_function: Optional[callable] = None, - a_send_message_function: Optional[Coroutine] = None, - a_human_input_function: Optional[callable] = None, - a_human_input_timeout: Optional[int] = 60, - connection_id: Optional[str] = None, - ) -> None: - """ - Initializes the WorkflowManager with agents specified in the config and optional message history. - - Args: - workflow (Union[Dict, str]): The workflow configuration. This can be a dictionary or a string which is a path to a JSON file. - history (Optional[List[Message]]): The message history. - work_dir (str): The working directory. - clear_work_dir (bool): If set to True, clears the working directory. - send_message_function (Optional[callable]): The function to send messages. - a_send_message_function (Optional[Coroutine]): Async coroutine to send messages. - a_human_input_function (Optional[callable]): Async coroutine to prompt the user for input. - a_human_input_timeout (Optional[int]): A time (in seconds) to wait for user input. After this time, the a_human_input_function will timeout and end the conversation. - connection_id (Optional[str]): The connection identifier. - """ - if isinstance(workflow, str): - if os.path.isfile(workflow): - with open(workflow, "r") as file: - self.workflow = json.load(file) - else: - raise FileNotFoundError(f"The file {workflow} does not exist.") - elif isinstance(workflow, dict): - self.workflow = workflow - else: - raise ValueError("The 'workflow' parameter should be either a dictionary or a valid JSON file path") - - # TODO - improved typing for workflow - self.workflow_skills = [] - self.send_message_function = send_message_function - self.a_send_message_function = a_send_message_function - self.a_human_input_function = a_human_input_function - self.a_human_input_timeout = a_human_input_timeout - self.connection_id = connection_id - self.work_dir = work_dir or "work_dir" - self.code_executor_pool = { - CodeExecutionConfigTypes.local: load_code_execution_config( - CodeExecutionConfigTypes.local, work_dir=self.work_dir - ), - CodeExecutionConfigTypes.docker: load_code_execution_config( - CodeExecutionConfigTypes.docker, work_dir=self.work_dir - ), - } - if clear_work_dir: - clear_folder(self.work_dir) - self.agent_history = [] - self.history = history or [] - self.sender = None - self.receiver = None - - def _run_workflow(self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False) -> None: - """ - Runs the workflow based on the provided configuration. - - Args: - message: The initial message to start the chat. - history: A list of messages to populate the agents' history. - clear_history: If set to True, clears the chat history before initiating. - - """ - for agent in self.workflow.get("agents", []): - if agent.get("link").get("agent_type") == "sender": - self.sender = self.load(agent.get("agent")) - elif agent.get("link").get("agent_type") == "receiver": - self.receiver = self.load(agent.get("agent")) - if self.sender and self.receiver: - # save all agent skills to skills.py - save_skills_to_file(self.workflow_skills, self.work_dir) - if history: - self._populate_history(history) - self.sender.initiate_chat( - self.receiver, - message=message, - clear_history=clear_history, - ) - else: - raise ValueError("Sender and receiver agents are not defined in the workflow configuration.") - - async def _a_run_workflow( - self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False - ) -> None: - """ - Asynchronously runs the workflow based on the provided configuration. - - Args: - message: The initial message to start the chat. - history: A list of messages to populate the agents' history. - clear_history: If set to True, clears the chat history before initiating. - - """ - for agent in self.workflow.get("agents", []): - if agent.get("link").get("agent_type") == "sender": - self.sender = self.load(agent.get("agent")) - elif agent.get("link").get("agent_type") == "receiver": - self.receiver = self.load(agent.get("agent")) - if self.sender and self.receiver: - # save all agent skills to skills.py - save_skills_to_file(self.workflow_skills, self.work_dir) - if history: - self._populate_history(history) - await self.sender.a_initiate_chat( - self.receiver, - message=message, - clear_history=clear_history, - ) - else: - raise ValueError("Sender and receiver agents are not defined in the workflow configuration.") - - def _serialize_agent( - self, - agent: Agent, - mode: str = "python", - include: Optional[List[str]] = {"config"}, - exclude: Optional[List[str]] = None, - ) -> Dict: - """ """ - # exclude = ["id","created_at", "updated_at","user_id","type"] - exclude = exclude or {} - include = include or {} - if agent.type != AgentType.groupchat: - exclude.update( - { - "config": { - "admin_name", - "messages", - "max_round", - "admin_name", - "speaker_selection_method", - "allow_repeat_speaker", - } - } - ) - else: - include = { - "config": { - "admin_name", - "messages", - "max_round", - "admin_name", - "speaker_selection_method", - "allow_repeat_speaker", - } - } - result = agent.model_dump(warnings=False, exclude=exclude, include=include, mode=mode) - return result["config"] - - def process_message( - self, - sender: autogen.Agent, - receiver: autogen.Agent, - message: Dict, - request_reply: bool = False, - silent: bool = False, - sender_type: str = "agent", - ) -> None: - """ - Processes the message and adds it to the agent history. - - Args: - - sender: The sender of the message. - receiver: The receiver of the message. - message: The message content. - request_reply: If set to True, the message will be added to agent history. - silent: determining verbosity. - sender_type: The type of the sender of the message. - """ - - message = message if isinstance(message, dict) else {"content": message, "role": "user"} - message_payload = { - "recipient": receiver.name, - "sender": sender.name, - "message": message, - "timestamp": datetime.now().isoformat(), - "sender_type": sender_type, - "connection_id": self.connection_id, - "message_type": "agent_message", - } - # if the agent will respond to the message, or the message is sent by a groupchat agent. - # This avoids adding groupchat broadcast messages to the history (which are sent with request_reply=False), - # or when agent populated from history - if request_reply is not False or sender_type == "groupchat": - self.agent_history.append(message_payload) # add to history - if self.send_message_function: # send over the message queue - socket_msg = SocketMessage( - type="agent_message", - data=message_payload, - connection_id=self.connection_id, - ) - self.send_message_function(socket_msg.dict()) - - async def a_process_message( - self, - sender: autogen.Agent, - receiver: autogen.Agent, - message: Dict, - request_reply: bool = False, - silent: bool = False, - sender_type: str = "agent", - ) -> None: - """ - Asynchronously processes the message and adds it to the agent history. - - Args: - - sender: The sender of the message. - receiver: The receiver of the message. - message: The message content. - request_reply: If set to True, the message will be added to agent history. - silent: determining verbosity. - sender_type: The type of the sender of the message. - """ - - message = message if isinstance(message, dict) else {"content": message, "role": "user"} - message_payload = { - "recipient": receiver.name, - "sender": sender.name, - "message": message, - "timestamp": datetime.now().isoformat(), - "sender_type": sender_type, - "connection_id": self.connection_id, - "message_type": "agent_message", - } - # if the agent will respond to the message, or the message is sent by a groupchat agent. - # This avoids adding groupchat broadcast messages to the history (which are sent with request_reply=False), - # or when agent populated from history - if request_reply is not False or sender_type == "groupchat": - self.agent_history.append(message_payload) # add to history - socket_msg = SocketMessage( - type="agent_message", - data=message_payload, - connection_id=self.connection_id, - ) - if self.a_send_message_function: # send over the message queue - await self.a_send_message_function(socket_msg.dict()) - elif self.send_message_function: # send over the message queue - self.send_message_function(socket_msg.dict()) - - def _populate_history(self, history: List[Message]) -> None: - """ - Populates the agent message history from the provided list of messages. - - Args: - history: A list of messages to populate the agents' history. - """ - for msg in history: - if isinstance(msg, dict): - msg = Message(**msg) - if msg.role == "user": - self.sender.send( - msg.content, - self.receiver, - request_reply=False, - silent=True, - ) - elif msg.role == "assistant": - self.receiver.send( - msg.content, - self.sender, - request_reply=False, - silent=True, - ) - - def sanitize_agent(self, agent: Dict) -> Agent: - """ """ - - skills = agent.get("skills", []) - - # When human input mode is not NEVER and no model is attached, the ui is passing bogus llm_config. - configured_models = agent.get("models") - if not configured_models or len(configured_models) == 0: - agent["config"]["llm_config"] = False - - agent = Agent.model_validate(agent) - agent.config.is_termination_msg = agent.config.is_termination_msg or ( - lambda x: "TERMINATE" in x.get("content", "").rstrip()[-20:] - ) - - def get_default_system_message(agent_type: str) -> str: - if agent_type == "assistant": - return autogen.AssistantAgent.DEFAULT_SYSTEM_MESSAGE - else: - return "You are a helpful AI Assistant." - - if agent.config.llm_config is not False: - config_list = [] - for llm in agent.config.llm_config.config_list: - # check if api_key is present either in llm or env variable - if "api_key" not in llm and "OPENAI_API_KEY" not in os.environ: - error_message = f"api_key is not present in llm_config or OPENAI_API_KEY env variable for agent ** {agent.config.name}**. Update your workflow to provide an api_key to use the LLM." - raise ValueError(error_message) - - # only add key if value is not None - sanitized_llm = sanitize_model(llm) - config_list.append(sanitized_llm) - agent.config.llm_config.config_list = config_list - - agent.config.code_execution_config = self.code_executor_pool.get(agent.config.code_execution_config, False) - - if skills: - for skill in skills: - self.workflow_skills.append(skill) - skills_prompt = "" - skills_prompt = get_skills_prompt(skills, self.work_dir) - if agent.config.system_message: - agent.config.system_message = agent.config.system_message + "\n\n" + skills_prompt - else: - agent.config.system_message = get_default_system_message(agent.type) + "\n\n" + skills_prompt - return agent - - def load(self, agent: Any) -> autogen.Agent: - """ - Loads an agent based on the provided agent specification. - - Args: - agent_spec: The specification of the agent to be loaded. - - Returns: - An instance of the loaded agent. - """ - if not agent: - raise ValueError( - "An agent configuration in this workflow is empty. Please provide a valid agent configuration." - ) - - linked_agents = agent.get("agents", []) - agent = self.sanitize_agent(agent) - if agent.type == "groupchat": - groupchat_agents = [self.load(agent) for agent in linked_agents] - group_chat_config = self._serialize_agent(agent) - group_chat_config["agents"] = groupchat_agents - groupchat = autogen.GroupChat(**group_chat_config) - agent = ExtendedGroupChatManager( - groupchat=groupchat, - message_processor=self.process_message, - a_message_processor=self.a_process_message, - a_human_input_function=self.a_human_input_function, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=self.connection_id, - llm_config=agent.config.llm_config.model_dump(), - ) - return agent - - else: - if agent.type == "assistant": - agent = ExtendedConversableAgent( - **self._serialize_agent(agent), - message_processor=self.process_message, - a_message_processor=self.a_process_message, - a_human_input_function=self.a_human_input_function, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=self.connection_id, - ) - elif agent.type == "userproxy": - agent = ExtendedConversableAgent( - **self._serialize_agent(agent), - message_processor=self.process_message, - a_message_processor=self.a_process_message, - a_human_input_function=self.a_human_input_function, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=self.connection_id, - ) - else: - raise ValueError(f"Unknown agent type: {agent.type}") - return agent - - def _generate_output( - self, - message_text: str, - summary_method: str, - ) -> str: - """ - Generates the output response based on the workflow configuration and agent history. - - :param message_text: The text of the incoming message. - :param flow: An instance of `WorkflowManager`. - :param flow_config: An instance of `AgentWorkFlowConfig`. - :return: The output response as a string. - """ - - output = "" - if summary_method == WorkFlowSummaryMethod.last: - (self.agent_history) - last_message = self.agent_history[-1]["message"]["content"] if self.agent_history else "" - output = last_message - elif summary_method == WorkFlowSummaryMethod.llm: - client = self.receiver.client - if self.connection_id: - status_message = SocketMessage( - type="agent_status", - data={ - "status": "summarizing", - "message": "Summarizing agent dialogue", - }, - connection_id=self.connection_id, - ) - self.send_message_function(status_message.model_dump(mode="json")) - output = summarize_chat_history( - task=message_text, - messages=self.agent_history, - client=client, - ) - - elif summary_method == "none": - output = "" - return output - - def _get_agent_usage(self, agent: autogen.Agent): - final_usage = [] - default_usage = {"total_cost": 0, "total_tokens": 0} - agent_usage = agent.client.total_usage_summary if agent.client else default_usage - agent_usage = { - "agent": agent.name, - "total_cost": find_key_value(agent_usage, "total_cost") or 0, - "total_tokens": find_key_value(agent_usage, "total_tokens") or 0, - } - final_usage.append(agent_usage) - - if type(agent) == ExtendedGroupChatManager: - print("groupchat found, processing", len(agent.groupchat.agents)) - for agent in agent.groupchat.agents: - agent_usage = agent.client.total_usage_summary if agent.client else default_usage or default_usage - agent_usage = { - "agent": agent.name, - "total_cost": find_key_value(agent_usage, "total_cost") or 0, - "total_tokens": find_key_value(agent_usage, "total_tokens") or 0, - } - final_usage.append(agent_usage) - return final_usage - - def _get_usage_summary(self): - sender_usage = self._get_agent_usage(self.sender) - receiver_usage = self._get_agent_usage(self.receiver) - - all_usage = [] - all_usage.extend(sender_usage) - all_usage.extend(receiver_usage) - # all_usage = [sender_usage, receiver_usage] - return all_usage - - def run(self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False) -> Message: - """ - Initiates a chat between the sender and receiver agents with an initial message - and an option to clear the history. - - Args: - message: The initial message to start the chat. - clear_history: If set to True, clears the chat history before initiating. - """ - - start_time = time.time() - self._run_workflow(message=message, history=history, clear_history=clear_history) - end_time = time.time() - - output = self._generate_output(message, self.workflow.get("summary_method", "last")) - - usage = self._get_usage_summary() - # print("usage", usage) - - result_message = Message( - content=output, - role="assistant", - meta={ - "messages": self.agent_history, - "summary_method": self.workflow.get("summary_method", "last"), - "time": end_time - start_time, - "files": get_modified_files(start_time, end_time, source_dir=self.work_dir), - "usage": usage, - }, - ) - return result_message - - async def a_run( - self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False - ) -> Message: - """ - Asynchronously initiates a chat between the sender and receiver agents with an initial message - and an option to clear the history. - - Args: - message: The initial message to start the chat. - clear_history: If set to True, clears the chat history before initiating. - """ - - start_time = time.time() - await self._a_run_workflow(message=message, history=history, clear_history=clear_history) - end_time = time.time() - - output = self._generate_output(message, self.workflow.get("summary_method", "last")) - - usage = self._get_usage_summary() - # print("usage", usage) - - result_message = Message( - content=output, - role="assistant", - meta={ - "messages": self.agent_history, - "summary_method": self.workflow.get("summary_method", "last"), - "time": end_time - start_time, - "files": get_modified_files(start_time, end_time, source_dir=self.work_dir), - "usage": usage, - }, - ) - return result_message - - -class SequentialWorkflowManager: - """ - WorkflowManager class to load agents from a provided configuration and run a chat between them sequentially. - """ - - def __init__( - self, - workflow: Union[Dict, str], - history: Optional[List[Message]] = None, - work_dir: str = None, - clear_work_dir: bool = True, - send_message_function: Optional[callable] = None, - a_send_message_function: Optional[Coroutine] = None, - a_human_input_function: Optional[callable] = None, - a_human_input_timeout: Optional[int] = 60, - connection_id: Optional[str] = None, - ) -> None: - """ - Initializes the WorkflowManager with agents specified in the config and optional message history. - - Args: - workflow (Union[Dict, str]): The workflow configuration. This can be a dictionary or a string which is a path to a JSON file. - history (Optional[List[Message]]): The message history. - work_dir (str): The working directory. - clear_work_dir (bool): If set to True, clears the working directory. - send_message_function (Optional[callable]): The function to send messages. - a_send_message_function (Optional[Coroutine]): Async coroutine to send messages. - a_human_input_function (Optional[callable]): Async coroutine to prompt for human input. - a_human_input_timeout (Optional[int]): A time (in seconds) to wait for user input. After this time, the a_human_input_function will timeout and end the conversation. - connection_id (Optional[str]): The connection identifier. - """ - if isinstance(workflow, str): - if os.path.isfile(workflow): - with open(workflow, "r") as file: - self.workflow = json.load(file) - else: - raise FileNotFoundError(f"The file {workflow} does not exist.") - elif isinstance(workflow, dict): - self.workflow = workflow - else: - raise ValueError("The 'workflow' parameter should be either a dictionary or a valid JSON file path") - - # TODO - improved typing for workflow - self.send_message_function = send_message_function - self.a_send_message_function = a_send_message_function - self.a_human_input_function = a_human_input_function - self.a_human_input_timeout = a_human_input_timeout - self.connection_id = connection_id - self.work_dir = work_dir or "work_dir" - if clear_work_dir: - clear_folder(self.work_dir) - self.agent_history = [] - self.history = history or [] - self.sender = None - self.receiver = None - self.model_client = None - - def _run_workflow(self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False) -> None: - """ - Runs the workflow based on the provided configuration. - - Args: - message: The initial message to start the chat. - history: A list of messages to populate the agents' history. - clear_history: If set to True, clears the chat history before initiating. - - """ - user_proxy = { - "config": { - "name": "user_proxy", - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 25, - "code_execution_config": "local", - "default_auto_reply": "TERMINATE", - "description": "User Proxy Agent Configuration", - "llm_config": False, - "type": "userproxy", - } - } - sequential_history = [] - for i, agent in enumerate(self.workflow.get("agents", [])): - workflow = Workflow( - name="agent workflow", type=WorkFlowType.autonomous, summary_method=WorkFlowSummaryMethod.llm - ) - workflow = workflow.model_dump(mode="json") - agent = agent.get("agent") - workflow["agents"] = [ - {"agent": user_proxy, "link": {"agent_type": "sender"}}, - {"agent": agent, "link": {"agent_type": "receiver"}}, - ] - - auto_workflow = AutoWorkflowManager( - workflow=workflow, - history=history, - work_dir=self.work_dir, - clear_work_dir=True, - send_message_function=self.send_message_function, - a_send_message_function=self.a_send_message_function, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=self.connection_id, - ) - task_prompt = ( - f""" - Your primary instructions are as follows: - {agent.get("task_instruction")} - Context for addressing your task is below: - ======= - {str(sequential_history)} - ======= - Now address your task: - """ - if i > 0 - else message - ) - result = auto_workflow.run(message=task_prompt, clear_history=clear_history) - sequential_history.append(result.content) - self.model_client = auto_workflow.receiver.client - print(f"======== end of sequence === {i}============") - self.agent_history.extend(result.meta.get("messages", [])) - - async def _a_run_workflow( - self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False - ) -> None: - """ - Asynchronously runs the workflow based on the provided configuration. - - Args: - message: The initial message to start the chat. - history: A list of messages to populate the agents' history. - clear_history: If set to True, clears the chat history before initiating. - - """ - user_proxy = { - "config": { - "name": "user_proxy", - "human_input_mode": "NEVER", - "max_consecutive_auto_reply": 25, - "code_execution_config": "local", - "default_auto_reply": "TERMINATE", - "description": "User Proxy Agent Configuration", - "llm_config": False, - "type": "userproxy", - } - } - sequential_history = [] - for i, agent in enumerate(self.workflow.get("agents", [])): - workflow = Workflow( - name="agent workflow", type=WorkFlowType.autonomous, summary_method=WorkFlowSummaryMethod.llm - ) - workflow = workflow.model_dump(mode="json") - agent = agent.get("agent") - workflow["agents"] = [ - {"agent": user_proxy, "link": {"agent_type": "sender"}}, - {"agent": agent, "link": {"agent_type": "receiver"}}, - ] - - auto_workflow = AutoWorkflowManager( - workflow=workflow, - history=history, - work_dir=self.work_dir, - clear_work_dir=True, - send_message_function=self.send_message_function, - a_send_message_function=self.a_send_message_function, - a_human_input_function=self.a_human_input_function, - a_human_input_timeout=self.a_human_input_timeout, - connection_id=self.connection_id, - ) - task_prompt = ( - f""" - Your primary instructions are as follows: - {agent.get("task_instruction")} - Context for addressing your task is below: - ======= - {str(sequential_history)} - ======= - Now address your task: - """ - if i > 0 - else message - ) - result = await auto_workflow.a_run(message=task_prompt, clear_history=clear_history) - sequential_history.append(result.content) - self.model_client = auto_workflow.receiver.client - print(f"======== end of sequence === {i}============") - self.agent_history.extend(result.meta.get("messages", [])) - - def _generate_output( - self, - message_text: str, - summary_method: str, - ) -> str: - """ - Generates the output response based on the workflow configuration and agent history. - - :param message_text: The text of the incoming message. - :param flow: An instance of `WorkflowManager`. - :param flow_config: An instance of `AgentWorkFlowConfig`. - :return: The output response as a string. - """ - - output = "" - if summary_method == WorkFlowSummaryMethod.last: - (self.agent_history) - last_message = self.agent_history[-1]["message"]["content"] if self.agent_history else "" - output = last_message - elif summary_method == WorkFlowSummaryMethod.llm: - if self.connection_id: - status_message = SocketMessage( - type="agent_status", - data={ - "status": "summarizing", - "message": "Summarizing agent dialogue", - }, - connection_id=self.connection_id, - ) - self.send_message_function(status_message.model_dump(mode="json")) - output = summarize_chat_history( - task=message_text, - messages=self.agent_history, - client=self.model_client, - ) - - elif summary_method == "none": - output = "" - return output - - def run(self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False) -> Message: - """ - Initiates a chat between the sender and receiver agents with an initial message - and an option to clear the history. - - Args: - message: The initial message to start the chat. - clear_history: If set to True, clears the chat history before initiating. - """ - - start_time = time.time() - self._run_workflow(message=message, history=history, clear_history=clear_history) - end_time = time.time() - output = self._generate_output(message, self.workflow.get("summary_method", "last")) - - result_message = Message( - content=output, - role="assistant", - meta={ - "messages": self.agent_history, - "summary_method": self.workflow.get("summary_method", "last"), - "time": end_time - start_time, - "files": get_modified_files(start_time, end_time, source_dir=self.work_dir), - "task": message, - }, - ) - return result_message - - async def a_run( - self, message: str, history: Optional[List[Message]] = None, clear_history: bool = False - ) -> Message: - """ - Asynchronously initiates a chat between the sender and receiver agents with an initial message - and an option to clear the history. - - Args: - message: The initial message to start the chat. - clear_history: If set to True, clears the chat history before initiating. - """ - - start_time = time.time() - await self._a_run_workflow(message=message, history=history, clear_history=clear_history) - end_time = time.time() - output = self._generate_output(message, self.workflow.get("summary_method", "last")) - - result_message = Message( - content=output, - role="assistant", - meta={ - "messages": self.agent_history, - "summary_method": self.workflow.get("summary_method", "last"), - "time": end_time - start_time, - "files": get_modified_files(start_time, end_time, source_dir=self.work_dir), - "task": message, - }, - ) - return result_message - - -class WorkflowManager: - """ - WorkflowManager class to load agents from a provided configuration and run a chat between them. - """ - - def __new__( - self, - workflow: Union[Dict, str], - history: Optional[List[Message]] = None, - work_dir: str = None, - clear_work_dir: bool = True, - send_message_function: Optional[callable] = None, - a_send_message_function: Optional[Coroutine] = None, - a_human_input_function: Optional[callable] = None, - a_human_input_timeout: Optional[int] = 60, - connection_id: Optional[str] = None, - ) -> None: - """ - Initializes the WorkflowManager with agents specified in the config and optional message history. - - Args: - workflow (Union[Dict, str]): The workflow configuration. This can be a dictionary or a string which is a path to a JSON file. - history (Optional[List[Message]]): The message history. - work_dir (str): The working directory. - clear_work_dir (bool): If set to True, clears the working directory. - send_message_function (Optional[callable]): The function to send messages. - a_send_message_function (Optional[Coroutine]): Async coroutine to send messages. - a_human_input_function (Optional[callable]): Async coroutine to prompt for user input. - a_human_input_timeout (Optional[int]): A time (in seconds) to wait for user input. After this time, the a_human_input_function will timeout and end the conversation. - connection_id (Optional[str]): The connection identifier. - """ - if isinstance(workflow, str): - if os.path.isfile(workflow): - with open(workflow, "r") as file: - self.workflow = json.load(file) - else: - raise FileNotFoundError(f"The file {workflow} does not exist.") - elif isinstance(workflow, dict): - self.workflow = workflow - else: - raise ValueError("The 'workflow' parameter should be either a dictionary or a valid JSON file path") - - if self.workflow.get("type") == WorkFlowType.autonomous.value: - return AutoWorkflowManager( - workflow=workflow, - history=history, - work_dir=work_dir, - clear_work_dir=clear_work_dir, - send_message_function=send_message_function, - a_send_message_function=a_send_message_function, - a_human_input_function=a_human_input_function, - a_human_input_timeout=a_human_input_timeout, - connection_id=connection_id, - ) - elif self.workflow.get("type") == WorkFlowType.sequential.value: - return SequentialWorkflowManager( - workflow=workflow, - history=history, - work_dir=work_dir, - clear_work_dir=clear_work_dir, - send_message_function=send_message_function, - a_send_message_function=a_send_message_function, - a_human_input_function=a_human_input_function, - a_human_input_timeout=a_human_input_timeout, - connection_id=connection_id, - ) - - -class ExtendedConversableAgent(autogen.ConversableAgent): - def __init__( - self, - message_processor=None, - a_message_processor=None, - a_human_input_function=None, - a_human_input_timeout: Optional[int] = 60, - connection_id=None, - *args, - **kwargs, - ): - - super().__init__(*args, **kwargs) - self.message_processor = message_processor - self.a_message_processor = a_message_processor - self.a_human_input_function = a_human_input_function - self.a_human_input_response = None - self.a_human_input_timeout = a_human_input_timeout - self.connection_id = connection_id - - def receive( - self, - message: Union[Dict, str], - sender: autogen.Agent, - request_reply: Optional[bool] = None, - silent: Optional[bool] = False, - ): - if self.message_processor: - self.message_processor(sender, self, message, request_reply, silent, sender_type="agent") - super().receive(message, sender, request_reply, silent) - - async def a_receive( - self, - message: Union[Dict, str], - sender: autogen.Agent, - request_reply: Optional[bool] = None, - silent: Optional[bool] = False, - ) -> None: - if self.a_message_processor: - await self.a_message_processor(sender, self, message, request_reply, silent, sender_type="agent") - elif self.message_processor: - self.message_processor(sender, self, message, request_reply, silent, sender_type="agent") - await super().a_receive(message, sender, request_reply, silent) - - # Strangely, when the response from a_get_human_input == "" (empty string) the libs call into the - # sync version. I guess that's "just in case", but it's odd because replying with an empty string - # is the intended way for the user to signal the underlying libs that they want to system to go forward - # with whatever function call, tool call or AI generated response the request calls for. Oh well, - # Que Sera Sera. - def get_human_input(self, prompt: str) -> str: - if self.a_human_input_response is None: - return super().get_human_input(prompt) - else: - response = self.a_human_input_response - self.a_human_input_response = None - return response - - async def a_get_human_input(self, prompt: str) -> str: - if self.message_processor and self.a_human_input_function: - message_dict = {"content": prompt, "role": "system", "type": "user-input-request"} - - message_payload = { - "recipient": self.name, - "sender": "system", - "message": message_dict, - "timestamp": datetime.now().isoformat(), - "sender_type": "system", - "connection_id": self.connection_id, - "message_type": "agent_message", - } - - socket_msg = SocketMessage( - type="user_input_request", - data=message_payload, - connection_id=self.connection_id, - ) - self.a_human_input_response = await self.a_human_input_function( - socket_msg.dict(), self.a_human_input_timeout - ) - return self.a_human_input_response - - else: - result = await super().a_get_human_input(prompt) - return result - - -class ExtendedGroupChatManager(autogen.GroupChatManager): - def __init__( - self, - message_processor=None, - a_message_processor=None, - a_human_input_function=None, - a_human_input_timeout: Optional[int] = 60, - connection_id=None, - *args, - **kwargs, - ): - super().__init__(*args, **kwargs) - self.message_processor = message_processor - self.a_message_processor = a_message_processor - self.a_human_input_function = a_human_input_function - self.a_human_input_response = None - self.a_human_input_timeout = a_human_input_timeout - self.connection_id = connection_id - - def receive( - self, - message: Union[Dict, str], - sender: autogen.Agent, - request_reply: Optional[bool] = None, - silent: Optional[bool] = False, - ): - if self.message_processor: - self.message_processor(sender, self, message, request_reply, silent, sender_type="groupchat") - super().receive(message, sender, request_reply, silent) - - async def a_receive( - self, - message: Union[Dict, str], - sender: autogen.Agent, - request_reply: Optional[bool] = None, - silent: Optional[bool] = False, - ) -> None: - if self.a_message_processor: - await self.a_message_processor(sender, self, message, request_reply, silent, sender_type="agent") - elif self.message_processor: - self.message_processor(sender, self, message, request_reply, silent, sender_type="agent") - await super().a_receive(message, sender, request_reply, silent) - - def get_human_input(self, prompt: str) -> str: - if self.a_human_input_response is None: - return super().get_human_input(prompt) - else: - response = self.a_human_input_response - self.a_human_input_response = None - return response - - async def a_get_human_input(self, prompt: str) -> str: - if self.message_processor and self.a_human_input_function: - message_dict = {"content": prompt, "role": "system", "type": "user-input-request"} - - message_payload = { - "recipient": self.name, - "sender": "system", - "message": message_dict, - "timestamp": datetime.now().isoformat(), - "sender_type": "system", - "connection_id": self.connection_id, - "message_type": "agent_message", - } - socket_msg = SocketMessage( - type="user_input_request", - data=message_payload, - connection_id=self.connection_id, - ) - result = await self.a_human_input_function(socket_msg.dict(), self.a_human_input_timeout) - return result - - else: - result = await super().a_get_human_input(prompt) - return result diff --git a/python/packages/autogen-studio/frontend/.env.default b/python/packages/autogen-studio/frontend/.env.default index 7f0839b275d2..9bd224b3320c 100644 --- a/python/packages/autogen-studio/frontend/.env.default +++ b/python/packages/autogen-studio/frontend/.env.default @@ -1 +1 @@ -GATSBY_API_URL=http://127.0.0.1:8081/api +GATSBY_API_URL=http://127.0.0.1:8081/api \ No newline at end of file diff --git a/python/packages/autogen-studio/frontend/.gitignore b/python/packages/autogen-studio/frontend/.gitignore index 8a0ea868f24b..f48f83a425b5 100644 --- a/python/packages/autogen-studio/frontend/.gitignore +++ b/python/packages/autogen-studio/frontend/.gitignore @@ -1,8 +1,6 @@ node_modules/ .cache/ -public/ - +public +src/gatsby-types.d.ts .env.development -.env.production - -yarn.lock +.env.production \ No newline at end of file diff --git a/python/packages/autogen-studio/frontend/LICENSE b/python/packages/autogen-studio/frontend/LICENSE deleted file mode 100644 index 16ab6489c8ed..000000000000 --- a/python/packages/autogen-studio/frontend/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -MIT License - -Copyright (c) 2022 Victor Dibia - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/python/packages/autogen-studio/frontend/README.md b/python/packages/autogen-studio/frontend/README.md index b707495cf42a..768541b0d06a 100644 --- a/python/packages/autogen-studio/frontend/README.md +++ b/python/packages/autogen-studio/frontend/README.md @@ -2,8 +2,8 @@ Run the UI in dev mode (make changes and see them reflected in the browser with hotreloading): -- npm install -- npm run start +- yarn install +- yarn start This should start the server on port 8000. diff --git a/python/packages/autogen-studio/frontend/gatsby-config.ts b/python/packages/autogen-studio/frontend/gatsby-config.ts index f66761c24be8..6497a72fe558 100644 --- a/python/packages/autogen-studio/frontend/gatsby-config.ts +++ b/python/packages/autogen-studio/frontend/gatsby-config.ts @@ -14,22 +14,20 @@ require("dotenv").config({ }); const config: GatsbyConfig = { - pathPrefix: process.env.PREFIX_PATH_VALUE || '', + pathPrefix: process.env.PREFIX_PATH_VALUE || "", siteMetadata: { title: `AutoGen Studio [Beta]`, description: `Build Multi-Agent Apps`, siteUrl: `http://tbd.place`, }, - flags: { - LAZY_IMAGES: true, - FAST_DEV: true, - DEV_SSR: false, - }, + // More easily incorporate content into your pages through automatic TypeScript type generation and better GraphQL IntelliSense. + // If you use VSCode you can also use the GraphQL plugin + // Learn more at: https://gatsby.dev/graphql-typegen + graphqlTypegen: true, plugins: [ - "gatsby-plugin-sass", + "gatsby-plugin-postcss", "gatsby-plugin-image", "gatsby-plugin-sitemap", - "gatsby-plugin-postcss", { resolve: "gatsby-plugin-manifest", options: { diff --git a/python/packages/autogen-studio/frontend/package.json b/python/packages/autogen-studio/frontend/package.json index 7a06f09dac03..24ba03e190f3 100644 --- a/python/packages/autogen-studio/frontend/package.json +++ b/python/packages/autogen-studio/frontend/package.json @@ -1,9 +1,9 @@ { - "name": "AutoGen_Studio", + "name": "autogentstudio", "version": "1.0.0", "private": true, "description": "AutoGen Studio - Build LLM Enabled Agents", - "author": "SPIRAL Team", + "author": "Microsoft", "keywords": [ "gatsby" ], @@ -17,55 +17,41 @@ "typecheck": "tsc --noEmit" }, "dependencies": { - "@ant-design/charts": "^1.3.6", + "@ant-design/charts": "^2.2.1", "@ant-design/plots": "^2.2.2", "@headlessui/react": "^1.7.16", "@heroicons/react": "^2.0.18", - "@mdx-js/mdx": "^1.6.22", - "@mdx-js/react": "^1.6.22", + "@mdx-js/react": "^3.1.0", "@monaco-editor/react": "^4.6.0", - "@tailwindcss/line-clamp": "^0.4.0", "@tailwindcss/typography": "^0.5.9", - "@types/lodash.debounce": "^4.0.9", - "@types/react-syntax-highlighter": "^15.5.10", "antd": "^5.1.0", - "autoprefixer": "^10.4.7", - "gatsby": "^4.14.0", - "gatsby-plugin-image": "^2.14.1", - "gatsby-plugin-manifest": "^4.14.0", - "gatsby-plugin-mdx": "^3.14.0", - "gatsby-plugin-postcss": "^5.14.0", - "gatsby-plugin-sass": "^5.14.0", - "gatsby-plugin-sharp": "^4.14.1", - "gatsby-plugin-sitemap": "^5.14.0", - "gatsby-source-filesystem": "^4.14.0", - "gatsby-transformer-sharp": "^4.14.0", - "jszip": "^3.10.1", - "lodash.debounce": "^4.0.8", - "papaparse": "^5.4.1", - "postcss": "^8.4.13", + "autoprefixer": "^10.4.20", + "gatsby": "^5.13.7", + "gatsby-plugin-image": "^3.13.1", + "gatsby-plugin-manifest": "^5.13.1", + "gatsby-plugin-mdx": "^5.13.1", + "gatsby-plugin-postcss": "^6.13.1", + "gatsby-plugin-sharp": "^5.13.1", + "gatsby-plugin-sitemap": "^6.13.1", + "gatsby-source-filesystem": "^5.13.1", + "gatsby-transformer-sharp": "^5.13.1", + "install": "^0.13.0", + "lucide-react": "^0.454.0", + "postcss": "^8.4.47", "react": "^18.2.0", - "react-contenteditable": "^3.3.6", "react-dom": "^18.2.0", - "react-inner-image-zoom": "^3.0.2", - "react-markdown": "^8.0.7", - "react-resizable": "^3.0.5", - "react-router-dom": "^6.3.0", - "react-syntax-highlighter": "^15.5.0", - "remark-gfm": "^3.0.1", - "sass": "^1.51.0", - "tailwindcss": "^3.0.24", - "uuid": "^9.0.1", - "zustand": "^4.4.6" + "react-markdown": "^9.0.1", + "tailwindcss": "^3.4.14", + "yarn": "^1.22.22", + "zustand": "^5.0.1" }, "devDependencies": { - "@types/node": "^18.7.13", - "@types/papaparse": "^5.3.14", - "@types/react": "^18.2.48", - "@types/react-dom": "^18.2.15", - "@types/react-inner-image-zoom": "^3.0.0", - "@types/react-resizable": "^3.0.2", - "@types/uuid": "^9.0.8", - "typescript": "^4.6.4" + "@types/lodash.debounce": "^4.0.9", + "@types/node": "^20.11.19", + "@types/react": "^18.2.55", + "@types/react-dom": "^18.2.19", + "@types/react-syntax-highlighter": "^15.5.10", + "@types/uuid": "^10.0.0", + "typescript": "^5.3.3" } } diff --git a/python/packages/autogen-studio/frontend/postcss.config.js b/python/packages/autogen-studio/frontend/postcss.config.js index e8351e2a4d51..33ad091d26d8 100644 --- a/python/packages/autogen-studio/frontend/postcss.config.js +++ b/python/packages/autogen-studio/frontend/postcss.config.js @@ -1,4 +1,6 @@ - -module.exports = () => ({ - plugins: [require("tailwindcss")], -}) +module.exports = { + plugins: { + tailwindcss: {}, + autoprefixer: {}, + }, +} diff --git a/python/packages/autogen-studio/frontend/src/components/atoms.tsx b/python/packages/autogen-studio/frontend/src/components/atoms.tsx deleted file mode 100644 index 8f52e60281b7..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/atoms.tsx +++ /dev/null @@ -1,873 +0,0 @@ -import { - ChevronDownIcon, - ChevronUpIcon, - Cog8ToothIcon, - XMarkIcon, - ClipboardIcon, - InformationCircleIcon, -} from "@heroicons/react/24/outline"; -import React, { ReactNode, useEffect, useRef, useState } from "react"; -import Icon from "./icons"; -import { Modal, Table, Tooltip, theme } from "antd"; -import Editor from "@monaco-editor/react"; -import Papa from "papaparse"; -import remarkGfm from "remark-gfm"; -import ReactMarkdown from "react-markdown"; -import { atomDark } from "react-syntax-highlighter/dist/esm/styles/prism"; -import { Prism as SyntaxHighlighter } from "react-syntax-highlighter"; -import { truncateText } from "./utils"; - -const { useToken } = theme; -interface CodeProps { - node?: any; - inline?: any; - className?: any; - children?: React.ReactNode; -} - -interface IProps { - children?: ReactNode; - title?: string | ReactNode; - subtitle?: string | ReactNode; - count?: number; - active?: boolean; - cursor?: string; - icon?: ReactNode; - padding?: string; - className?: string; - open?: boolean; - hoverable?: boolean; - onClick?: () => void; - loading?: boolean; -} - -export const SectionHeader = ({ - children, - title, - subtitle, - count, - icon, -}: IProps) => { - return ( -
-

- {/* {count !== null && {count}} */} - {icon && <>{icon}} - {title} - {count !== null && ( - {count} - )} -

- {subtitle && {subtitle}} - {children} -
- ); -}; - -export const IconButton = ({ - onClick, - icon, - className, - active = false, -}: IProps) => { - return ( - - {icon} - - ); -}; - -export const LaunchButton = ({ - children, - onClick, - className = "p-3 px-5 ", -}: any) => { - return ( - - ); -}; - -export const SecondaryButton = ({ children, onClick, className }: any) => { - return ( - - ); -}; - -export const Card = ({ - children, - title, - subtitle, - hoverable = true, - active, - cursor = "cursor-pointer", - className = "p-3", - onClick, -}: IProps) => { - let border = active - ? "border-accent" - : "border-secondary hover:border-accent "; - border = hoverable ? border : "border-secondary"; - - return ( - - ); -}; - -export const CollapseBox = ({ - title, - subtitle, - children, - className = " p-3", - open = false, -}: IProps) => { - const [isOpen, setIsOpen] = React.useState(open); - const chevronClass = "h-4 cursor-pointer inline-block mr-1"; - return ( -
{ - if (e.detail > 1) { - e.preventDefault(); - } - }} - className="bordper border-secondary rounded" - > -
{ - setIsOpen(!isOpen); - }} - className={`cursor-pointer bg-secondary p-2 rounded ${ - isOpen ? "rounded-b-none " : " " - }"}`} - > - {isOpen && } - {!isOpen && } - - - {" "} - {/* {isOpen ? "hide" : "show"} section | */} - {title} - -
- - {isOpen && ( -
- {children} -
- )} -
- ); -}; - -export const HighLight = ({ children }: IProps) => { - return {children}; -}; - -export const LoadBox = ({ - subtitle, - className = "my-2 text-accent ", -}: IProps) => { - return ( -
- - {" "} - - {" "} - {subtitle} -
- ); -}; - -export const LoadingBar = ({ children }: IProps) => { - return ( - <> -
- - - - - {children} -
-
-
-
- - ); -}; - -export const MessageBox = ({ title, children, className }: IProps) => { - const messageBox = useRef(null); - - const closeMessage = () => { - if (messageBox.current) { - messageBox.current.remove(); - } - }; - - return ( -
- {" "} -
-
- {/* - - {" "} */} - {title} -
-
- { - closeMessage(); - }} - className=" border border-secondary bg-secondary brightness-125 hover:brightness-100 cursor-pointer transition duration-200 inline-block px-1 pb-1 rounded text-primary" - > - - -
-
- {children} -
- ); -}; - -export const GroupView = ({ - children, - title, - className = "text-primary bg-primary ", -}: any) => { - return ( -
-
-
- {title} -
-
{children}
-
-
- ); -}; - -export const ExpandView = ({ - children, - icon = null, - className = "", - title = "Detail View", -}: any) => { - const [isOpen, setIsOpen] = React.useState(false); - let windowAspect = 1; - if (typeof window !== "undefined") { - windowAspect = window.innerWidth / window.innerHeight; - } - const minImageWidth = 400; - return ( -
-
{ - setIsOpen(true); - }} - className="text-xs mb-2 h-full w-full break-words" - > - {icon ? icon : children} -
- {isOpen && ( - setIsOpen(false)} - footer={null} - > - {/* resize} - lockAspectRatio={false} - handle={ -
- -
- } - width={800} - height={minImageWidth * windowAspect} - minConstraints={[minImageWidth, minImageWidth * windowAspect]} - maxConstraints={[900, 900 * windowAspect]} - className="overflow-auto w-full rounded select-none " - > */} - {children} - {/* */} - - )} -
- ); -}; - -export const LoadingOverlay = ({ children, loading }: IProps) => { - return ( - <> - {loading && ( - <> -
- {/* Overlay background */} -
-
- {/* Center BounceLoader without inheriting the opacity */} - -
- - )} -
{children}
- - ); -}; - -export const MarkdownView = ({ - data, - className = "", - showCode = true, -}: { - data: string; - className?: string; - showCode?: boolean; -}) => { - function processString(inputString: string): string { - // TODO: Had to add this temp measure while debugging. Why is it null? - if (!inputString) { - console.log("inputString is null!") - } - inputString = inputString && inputString.replace(/\n/g, " \n"); - const markdownPattern = /```markdown\s+([\s\S]*?)\s+```/g; - return inputString?.replace(markdownPattern, (match, content) => content); - } - const [showCopied, setShowCopied] = React.useState(false); - - const CodeView = ({ props, children, language }: any) => { - const [codeVisible, setCodeVisible] = React.useState(showCode); - return ( -
-
-
{ - setCodeVisible(!codeVisible); - }} - className=" flex-1 mr-4 " - > - {!codeVisible && ( -
- - show -
- )} - - {codeVisible && ( -
- {" "} - - hide -
- )} -
- {/*
*/} -
- {showCopied && ( -
- {" "} - 🎉 Copied!{" "} -
- )} - { - navigator.clipboard.writeText(data); - // message.success("Code copied to clipboard"); - setShowCopied(true); - setTimeout(() => { - setShowCopied(false); - }, 3000); - }} - className=" inline-block duration-300 text-white hover:text-accent w-5 h-5" - /> -
-
- {codeVisible && ( - - {String(children).replace(/\n$/, "")} - - )} -
- ); - }; - - return ( -
- - ) : ( - - {children} - - ); - }, - }} - > - {processString(data)} - -
- ); -}; - -interface ICodeProps { - code: string; - language: string; - title?: string; - showLineNumbers?: boolean; - className?: string | undefined; - wrapLines?: boolean; - maxWidth?: string; - maxHeight?: string; - minHeight?: string; -} - -export const CodeBlock = ({ - code, - language = "python", - showLineNumbers = false, - className = " ", - wrapLines = false, - maxHeight = "400px", - minHeight = "auto", -}: ICodeProps) => { - const codeString = code; - - const [showCopied, setShowCopied] = React.useState(false); - return ( -
-
-
-
-
- {showCopied && ( -
- {" "} - 🎉 Copied!{" "} -
- )} - { - navigator.clipboard.writeText(codeString); - // message.success("Code copied to clipboard"); - setShowCopied(true); - setTimeout(() => { - setShowCopied(false); - }, 6000); - }} - className="m-2 inline-block duration-300 text-white hover:text-accent w-5 h-5" - /> -
-
-
-
- - {codeString} - -
-
- ); -}; - -// Controls Row -export const ControlRowView = ({ - title, - description, - value, - control, - className, - truncateLength = 20, -}: { - title: string; - description: string; - value: string | number | boolean; - control: any; - className?: string; - truncateLength?: number; -}) => { - return ( -
-
- {title} - - {truncateText(value + "", truncateLength)} - {" "} - - - -
- {control} -
-
- ); -}; - -export const BounceLoader = ({ - className, - title = "", -}: { - className?: string; - title?: string; -}) => { - return ( -
-
- - - -
- {title} -
- ); -}; - -export const ImageLoader = ({ - src, - className = "", -}: { - src: string; - className?: string; -}) => { - const [isLoading, setIsLoading] = useState(true); - - return ( -
- {isLoading && ( -
- -
- )} - Dynamic content setIsLoading(false)} - /> -
- ); -}; - -type DataRow = { [key: string]: any }; -export const CsvLoader = ({ - csvUrl, - className, -}: { - csvUrl: string; - className?: string; -}) => { - const [data, setData] = useState([]); - const [columns, setColumns] = useState([]); - const [isLoading, setIsLoading] = useState(true); - const [pageSize, setPageSize] = useState(50); - - useEffect(() => { - const fetchData = async () => { - try { - const response = await fetch(csvUrl); - const csvString = await response.text(); - const parsedData = Papa.parse(csvString, { - header: true, - dynamicTyping: true, - skipEmptyLines: true, - }); - setData(parsedData.data as DataRow[]); - - // Use the keys of the first object for column headers - const firstRow = parsedData.data[0] as DataRow; // Type assertion - const columnHeaders: any[] = Object.keys(firstRow).map((key) => { - const val = { - title: key.charAt(0).toUpperCase() + key.slice(1), // Capitalize the key for the title - dataIndex: key, - key: key, - }; - if (typeof firstRow[key] === "number") { - return { - ...val, - sorter: (a: DataRow, b: DataRow) => a[key] - b[key], - }; - } - return val; - }); - setColumns(columnHeaders); - setIsLoading(false); - } catch (error) { - console.error("Error fetching CSV data:", error); - setIsLoading(false); - } - }; - - fetchData(); - }, [csvUrl]); - - // calculate x scroll, based on number of columns - const scrollX = columns.length * 150; - - return ( -
- { - setPageSize(pagination.pageSize || 50); - }} - /> - - ); -}; - -export const CodeLoader = ({ - url, - className, -}: { - url: string; - className?: string; -}) => { - const [isLoading, setIsLoading] = useState(true); - const [code, setCode] = useState(null); - - React.useEffect(() => { - fetch(url) - .then((response) => response.text()) - .then((data) => { - setCode(data); - setIsLoading(false); - }); - }, [url]); - - return ( -
- {isLoading && ( -
- -
- )} - - {!isLoading && } -
- ); -}; - -export const PdfViewer = ({ url }: { url: string }) => { - const [loading, setLoading] = useState(true); - - React.useEffect(() => { - // Assuming the URL is directly usable as the source for the tag - setLoading(false); - // Note: No need to handle the creation and cleanup of a blob URL or converting file content as it's not provided anymore. - }, [url]); - - // Render the PDF viewer - return ( -
- {loading &&

Loading PDF...

} - {!loading && ( - -

PDF cannot be displayed.

- - )} -
- ); -}; - -export const MonacoEditor = ({ - value, - editorRef, - language, - onChange, - minimap = true, -}: { - value: string; - onChange?: (value: string) => void; - editorRef: any; - language: string; - minimap?: boolean; -}) => { - const [isEditorReady, setIsEditorReady] = useState(false); - const onEditorDidMount = (editor: any, monaco: any) => { - editorRef.current = editor; - setIsEditorReady(true); - }; - return ( -
- { - if (onChange && value) { - onChange(value); - } - }} - onMount={onEditorDidMount} - theme="vs-dark" - options={{ - wordWrap: "on", - wrappingIndent: "indent", - wrappingStrategy: "advanced", - minimap: { - enabled: minimap, - }, - }} - /> -
- ); -}; - -export const CardHoverBar = ({ - items, -}: { - items: { - title: string; - icon: any; - hoverText: string; - onClick: (e: any) => void; - }[]; -}) => { - const itemRows = items.map((item, i) => { - return ( -
- - - -
- ); - }); - return ( -
{ - e.stopPropagation(); - }} - className=" mt-2 text-right opacity-0 group-hover:opacity-100 " - > - {itemRows} -
- ); -}; - -export const AgentRow = ({ message }: { message: any }) => { - return ( - - {message.sender} ( to{" "} - {message.recipient} ) - - } - className="m" - > - - - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/contentheader.tsx b/python/packages/autogen-studio/frontend/src/components/contentheader.tsx new file mode 100644 index 000000000000..b66bc0d7da27 --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/contentheader.tsx @@ -0,0 +1,157 @@ +import React from "react"; +import { Menu } from "@headlessui/react"; +import { + BellIcon, + MoonIcon, + SunIcon, + MagnifyingGlassIcon, +} from "@heroicons/react/24/outline"; +import { + ChevronDown, + PanelLeftClose, + PanelLeftOpen, + Menu as MenuIcon, +} from "lucide-react"; +import { Tooltip } from "antd"; +import { appContext } from "../hooks/provider"; +import { useConfigStore } from "../hooks/store"; + +type ContentHeaderProps = { + title?: string; + onMobileMenuToggle: () => void; + isMobileMenuOpen: boolean; +}; + +const classNames = (...classes: (string | undefined | boolean)[]) => { + return classes.filter(Boolean).join(" "); +}; + +const ContentHeader = ({ + title, + onMobileMenuToggle, + isMobileMenuOpen, +}: ContentHeaderProps) => { + const { darkMode, setDarkMode, user, logout } = React.useContext(appContext); + const { sidebar, setSidebarState } = useConfigStore(); + const { isExpanded } = sidebar; + + return ( +
+
+ {/* Mobile Menu Button */} + + + {/* Desktop Sidebar Toggle - Hidden on Mobile */} +
+ + + +
+ +
+ {/* Search */} +
+
+ + + + +
+ + {/* Right side header items */} +
+ {/* Dark Mode Toggle */} + + + {/* Notifications */} + + + {/* Separator */} +
+ + {/* User Menu */} + {user && ( + + + {user.avatar_url ? ( + {user.name} + ) : ( +
+ {user.name?.[0]} +
+ )} + + + {user.name} + + + + + + + {({ active }) => ( + logout()} + className={`${ + active ? "bg-secondary" : "" + } block px-4 py-2 text-sm text-primary`} + > + Sign out + + )} + + + + )} +
+
+
+
+ ); +}; + +export default ContentHeader; diff --git a/python/packages/autogen-studio/frontend/src/components/footer.tsx b/python/packages/autogen-studio/frontend/src/components/footer.tsx index 104547b443d3..6a8828d5ea28 100644 --- a/python/packages/autogen-studio/frontend/src/components/footer.tsx +++ b/python/packages/autogen-studio/frontend/src/components/footer.tsx @@ -17,7 +17,7 @@ const Footer = () => { } }, []); return ( -
+
Maintained by the AutoGen{" "} { className={` ${sizeClass} inline-block `} xmlns="http://www.w3.org/2000/svg" fill="currentColor" - viewBox="0 0 93 90" + viewBox="0 0 290 264" > + + + + + + + + + + + + + + + + + + + + + ); } diff --git a/python/packages/autogen-studio/frontend/src/components/layout.tsx b/python/packages/autogen-studio/frontend/src/components/layout.tsx index 9142470534eb..42182b07684f 100644 --- a/python/packages/autogen-studio/frontend/src/components/layout.tsx +++ b/python/packages/autogen-studio/frontend/src/components/layout.tsx @@ -1,10 +1,16 @@ import * as React from "react"; -import Header from "./header"; +import { Dialog } from "@headlessui/react"; +import { X } from "lucide-react"; import { appContext } from "../hooks/provider"; +import { useConfigStore } from "../hooks/store"; import Footer from "./footer"; - -/// import ant css import "antd/dist/reset.css"; +import SideBar from "./sidebar"; +import ContentHeader from "./contentheader"; + +const classNames = (...classes: (string | undefined | boolean)[]) => { + return classes.filter(Boolean).join(" "); +}; type Props = { title: string; @@ -23,38 +29,96 @@ const Layout = ({ showHeader = true, restricted = false, }: Props) => { - const layoutContent = ( -
- {showHeader &&
} -
- {meta?.title + " | " + title} -
{children}
-
-
-
- ); - const { darkMode } = React.useContext(appContext); + const { sidebar } = useConfigStore(); + const { isExpanded } = sidebar; + const [isMobileMenuOpen, setIsMobileMenuOpen] = React.useState(false); + + // Close mobile menu on route change + React.useEffect(() => { + setIsMobileMenuOpen(false); + }, [link]); + React.useEffect(() => { document.getElementsByTagName("html")[0].className = `${ darkMode === "dark" ? "dark bg-primary" : "light bg-primary" - } `; + }`; }, [darkMode]); - return ( - - {(context: any) => { - if (restricted) { - return
{context.user && layoutContent}
; - } else { - return layoutContent; - } - }} - + const layoutContent = ( +
+ {/* Mobile menu */} + setIsMobileMenuOpen(false)} + className="relative z-50 md:hidden" + > + {/* Backdrop */} + + + {/* Desktop sidebar */} +
+ +
+ + {/* Content area */} +
+ {showHeader && ( + setIsMobileMenuOpen(!isMobileMenuOpen)} + /> + )} + +
{children}
+ +
+
+
); + + // Handle restricted content + if (restricted) { + return ( + + {(context: any) => { + if (context.user) { + return layoutContent; + } + return null; + }} + + ); + } + + return layoutContent; }; export default Layout; diff --git a/python/packages/autogen-studio/frontend/src/components/sidebar.tsx b/python/packages/autogen-studio/frontend/src/components/sidebar.tsx new file mode 100644 index 000000000000..b9decb10170d --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/sidebar.tsx @@ -0,0 +1,148 @@ +import React from "react"; +import { Link } from "gatsby"; +import { useConfigStore } from "../hooks/store"; +import { Tooltip } from "antd"; +import { Blocks, Settings, MessagesSquare } from "lucide-react"; +import Icon from "./icons"; + +const navigation = [ + // { name: "Build", href: "/build", icon: Blocks }, + { name: "Playground", href: "/", icon: MessagesSquare }, +]; + +const classNames = (...classes: (string | undefined | boolean)[]) => { + return classes.filter(Boolean).join(" "); +}; + +type SidebarProps = { + link: string; + meta?: { + title: string; + description: string; + }; + isMobile: boolean; +}; + +const Sidebar = ({ link, meta, isMobile }: SidebarProps) => { + const { sidebar } = useConfigStore(); + const { isExpanded } = sidebar; + + // Always show full sidebar in mobile view + const showFull = isMobile || isExpanded; + + return ( +
+ {/* App Logo/Title */} +
+
+ +
+ {showFull && ( +
+ + {meta?.title} + + {meta?.description} +
+ )} +
+ + {/* Navigation */} + +
+ ); +}; + +export default Sidebar; diff --git a/python/packages/autogen-studio/frontend/src/components/types.ts b/python/packages/autogen-studio/frontend/src/components/types.ts deleted file mode 100644 index ca51003e7ed0..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/types.ts +++ /dev/null @@ -1,127 +0,0 @@ -export type NotificationType = "success" | "info" | "warning" | "error"; - -export interface IMessage { - user_id: string; - role: string; - content: string; - created_at?: string; - updated_at?: string; - session_id?: number; - connection_id?: string; - workflow_id?: number; - meta?: any; - id?: number; -} - -export interface IStatus { - message: string; - status: boolean; - data?: any; -} - -export interface IChatMessage { - text: string; - sender: "user" | "bot"; - meta?: any; - id?: number; -} - -export interface ILLMConfig { - config_list: Array; - timeout?: number; - cache_seed?: number | null; - temperature: number; - max_tokens: number; -} - -export interface IAgentConfig { - name: string; - llm_config?: ILLMConfig | false; - human_input_mode: string; - max_consecutive_auto_reply: number; - system_message: string | ""; - is_termination_msg?: boolean | string; - default_auto_reply?: string | null; - code_execution_config?: "none" | "local" | "docker"; - description?: string; - - admin_name?: string; - messages?: Array; - max_round?: number; - speaker_selection_method?: string; - allow_repeat_speaker?: boolean; -} - -export interface IAgent { - type?: "assistant" | "userproxy" | "groupchat"; - config: IAgentConfig; - created_at?: string; - updated_at?: string; - id?: number; - skills?: Array; - user_id?: string; -} - -export interface IWorkflow { - name: string; - description: string; - sender?: IAgent; - receiver?: IAgent; - type?: "autonomous" | "sequential"; - created_at?: string; - updated_at?: string; - summary_method?: "none" | "last" | "llm"; - id?: number; - user_id?: string; -} - -export interface IModelConfig { - model: string; - api_key?: string; - api_version?: string; - base_url?: string; - api_type?: "open_ai" | "azure" | "google" | "anthropic" | "mistral"; - user_id?: string; - created_at?: string; - updated_at?: string; - description?: string; - id?: number; -} - -export interface IMetadataFile { - name: string; - path: string; - extension: string; - content: string; - type: string; -} - -export interface IChatSession { - id?: number; - user_id: string; - workflow_id?: number; - created_at?: string; - updated_at?: string; - name: string; -} - -export interface IGalleryItem { - id: number; - messages: Array; - session: IChatSession; - tags: Array; - created_at: string; - updated_at: string; -} - -export interface ISkill { - name: string; - content: string; - secrets?: any[]; - libraries?: string[]; - id?: number; - description?: string; - user_id?: string; - created_at?: string; - updated_at?: string; -} diff --git a/python/packages/autogen-studio/frontend/src/components/types/app.ts b/python/packages/autogen-studio/frontend/src/components/types/app.ts new file mode 100644 index 000000000000..e71635902323 --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/types/app.ts @@ -0,0 +1,5 @@ +export interface IStatus { + message: string; + status: boolean; + data?: any; +} diff --git a/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts b/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts new file mode 100644 index 000000000000..da4999d376a2 --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts @@ -0,0 +1,86 @@ +export interface RequestUsage { + prompt_tokens: number; + completion_tokens: number; +} + +export interface MessageConfig { + source: string; + content: string; + models_usage?: RequestUsage; +} + +export interface DBModel { + id?: number; + user_id?: string; + created_at?: string; + updated_at?: string; +} + +export interface Message extends DBModel { + config: MessageConfig; + session_id: number; + run_id: string; +} + +export interface Session extends DBModel { + name: string; + team_id?: string; +} + +export interface TeamConfig { + name: string; + participants: AgentConfig[]; + team_type: TeamTypes; + model_client?: ModelConfig; + termination_condition?: TerminationConfig; +} + +export interface Team extends DBModel { + config: TeamConfig; +} + +export type ModelTypes = "OpenAIChatCompletionClient"; + +export type AgentTypes = "AssistantAgent" | "CodingAssistantAgent"; + +export type TeamTypes = "RoundRobinGroupChat" | "SelectorGroupChat"; + +export type TerminationTypes = + | "MaxMessageTermination" + | "StopMessageTermination" + | "TextMentionTermination"; + +export interface ModelConfig { + model: string; + model_type: ModelTypes; + api_key?: string; + base_url?: string; +} + +export interface ToolConfig { + name: string; + description: string; + content: string; +} + +export interface AgentConfig { + name: string; + agent_type: AgentTypes; + system_message?: string; + model_client?: ModelConfig; + tools?: ToolConfig[]; + description?: string; +} + +export interface TerminationConfig { + termination_type: TerminationTypes; + max_messages?: number; + text?: string; +} + +export interface TaskResult { + messages: MessageConfig[]; + usage: string; + duration: number; + stop_reason?: string; +} diff --git a/python/packages/autogen-studio/frontend/src/components/utils.ts b/python/packages/autogen-studio/frontend/src/components/utils.ts index e70590153a88..2f3d21230da7 100644 --- a/python/packages/autogen-studio/frontend/src/components/utils.ts +++ b/python/packages/autogen-studio/frontend/src/components/utils.ts @@ -1,12 +1,4 @@ -import { - IAgent, - IAgentConfig, - ILLMConfig, - IModelConfig, - ISkill, - IStatus, - IWorkflow, -} from "./types"; +import { IStatus } from "./types"; export const getServerUrl = () => { return process.env.GATSBY_API_URL || "/api"; @@ -100,10 +92,6 @@ export function fetchJSON( onFinal(); }); } -export const capitalize = (s: string) => { - if (typeof s !== "string") return ""; - return s.charAt(0).toUpperCase() + s.slice(1); -}; export function eraseCookie(name: string) { document.cookie = name + "=; Path=/; Expires=Thu, 01 Jan 1970 00:00:01 GMT;"; @@ -116,435 +104,6 @@ export function truncateText(text: string, length = 50) { return text; } -export const getCaretCoordinates = () => { - let caretX, caretY; - const selection = window.getSelection(); - if (selection && selection?.rangeCount !== 0) { - const range = selection.getRangeAt(0).cloneRange(); - range.collapse(false); - const rect = range.getClientRects()[0]; - if (rect) { - caretX = rect.left; - caretY = rect.top; - } - } - return { caretX, caretY }; -}; - -export const getPrefixSuffix = (container: any) => { - let prefix = ""; - let suffix = ""; - if (window.getSelection) { - const sel = window.getSelection(); - if (sel && sel.rangeCount > 0) { - let range = sel.getRangeAt(0).cloneRange(); - range.collapse(true); - range.setStart(container!, 0); - prefix = range.toString(); - - range = sel.getRangeAt(0).cloneRange(); - range.collapse(true); - range.setEnd(container, container.childNodes.length); - - suffix = range.toString(); - console.log("prefix", prefix); - console.log("suffix", suffix); - } - } - return { prefix, suffix }; -}; - -export const uid = () => { - return Date.now().toString(36) + Math.random().toString(36).substr(2); -}; - -export const setCaretToEnd = (element: HTMLElement) => { - const range = document.createRange(); - const selection = window.getSelection(); - range.selectNodeContents(element); - range.collapse(false); - selection?.removeAllRanges(); - selection?.addRange(range); - element.focus(); -}; - -// return a color between a start and end color using a percentage -export const ColorTween = ( - startColor: string, - endColor: string, - percent: number -) => { - // example startColor = "#ff0000" endColor = "#0000ff" percent = 0.5 - const start = { - r: parseInt(startColor.substring(1, 3), 16), - g: parseInt(startColor.substring(3, 5), 16), - b: parseInt(startColor.substring(5, 7), 16), - }; - const end = { - r: parseInt(endColor.substring(1, 3), 16), - g: parseInt(endColor.substring(3, 5), 16), - b: parseInt(endColor.substring(5, 7), 16), - }; - const r = Math.floor(start.r + (end.r - start.r) * percent); - const g = Math.floor(start.g + (end.g - start.g) * percent); - const b = Math.floor(start.b + (end.b - start.b) * percent); - return `rgb(${r}, ${g}, ${b})`; -}; - -export const guid = () => { - var w = () => { - return Math.floor((1 + Math.random()) * 0x10000) - .toString(16) - .substring(1); - }; - return `${w()}${w()}-${w()}-${w()}-${w()}-${w()}${w()}${w()}`; -}; - -/** - * Takes a string and returns the first n characters followed by asterisks. - * @param {string} str - The string to obscure - * @param {number} n - Number of characters to show before obscuring - * @returns {string} The obscured string with first n characters in clear text - */ -export const obscureString = (str: string, n: number = 3) => { - if (n < 0 || n > str.length) { - console.log("n cannot be less than 0 or greater than the string length."); - return str; - } - // First n characters in clear text - var clearText = str.substring(0, n); - // Remaining characters replaced with asterisks - var obscured = clearText + "*".repeat(str.length - n); - - return obscured; -}; - -/** - * Converts a number of seconds into a human-readable string representing the duration in days, hours, minutes, and seconds. - * @param {number} seconds - The number of seconds to convert. - * @returns {string} A well-formatted duration string. - */ -export const formatDuration = (seconds: number) => { - const units = [ - { label: " day", seconds: 86400 }, - { label: " hr", seconds: 3600 }, - { label: " min", seconds: 60 }, - { label: " sec", seconds: 1 }, - ]; - - let remainingSeconds = seconds; - const parts = []; - - for (const { label, seconds: unitSeconds } of units) { - const count = Math.floor(remainingSeconds / unitSeconds); - if (count > 0) { - parts.push(count + (count > 1 ? label + "s" : label)); - remainingSeconds -= count * unitSeconds; - } - } - - return parts.length > 0 ? parts.join(" ") : "0 sec"; -}; - -export const sampleModelConfig = (modelType: string = "open_ai") => { - const openaiConfig: IModelConfig = { - model: "gpt-4-1106-preview", - api_type: "open_ai", - description: "OpenAI GPT-4 model", - }; - const azureConfig: IModelConfig = { - model: "gpt-4", - api_type: "azure", - api_version: "v1", - base_url: "https://youazureendpoint.azure.com/", - description: "Azure model", - }; - - const googleConfig: IModelConfig = { - model: "gemini-1.0-pro", - api_type: "google", - description: "Google Gemini Model model", - }; - - const anthropicConfig: IModelConfig = { - model: "claude-3-5-sonnet-20240620", - api_type: "anthropic", - description: "Claude 3.5 Sonnet model", - }; - - const mistralConfig: IModelConfig = { - model: "mistral", - api_type: "mistral", - description: "Mistral model", - }; - - switch (modelType) { - case "open_ai": - return openaiConfig; - case "azure": - return azureConfig; - case "google": - return googleConfig; - case "anthropic": - return anthropicConfig; - case "mistral": - return mistralConfig; - default: - return openaiConfig; - } -}; - -export const getRandomIntFromDateAndSalt = (salt: number = 43444) => { - const currentDate = new Date(); - const seed = currentDate.getTime() + salt; - const randomValue = Math.sin(seed) * 10000; - const randomInt = Math.floor(randomValue) % 100; - return randomInt; -}; - -export const getSampleWorkflow = (workflow_type: string = "autonomous") => { - const autonomousWorkflow: IWorkflow = { - name: "Default Chat Workflow", - description: "Autonomous Workflow", - type: "autonomous", - summary_method: "llm", - }; - const sequentialWorkflow: IWorkflow = { - name: "Default Sequential Workflow", - description: "Sequential Workflow", - type: "sequential", - summary_method: "llm", - }; - - if (workflow_type === "autonomous") { - return autonomousWorkflow; - } else if (workflow_type === "sequential") { - return sequentialWorkflow; - } else { - return autonomousWorkflow; - } -}; - -export const sampleAgentConfig = (agent_type: string = "assistant") => { - const llm_config: ILLMConfig = { - config_list: [], - temperature: 0.1, - timeout: 600, - cache_seed: null, - max_tokens: 4000, - }; - - const userProxyConfig: IAgentConfig = { - name: "userproxy", - human_input_mode: "NEVER", - description: "User Proxy", - max_consecutive_auto_reply: 25, - system_message: "You are a helpful assistant.", - default_auto_reply: "TERMINATE", - llm_config: false, - code_execution_config: "local", - }; - const userProxyFlowSpec: IAgent = { - type: "userproxy", - config: userProxyConfig, - }; - - const assistantConfig: IAgentConfig = { - name: "primary_assistant", - description: "Primary Assistant", - llm_config: llm_config, - human_input_mode: "NEVER", - max_consecutive_auto_reply: 25, - code_execution_config: "none", - system_message: - "You are a helpful AI assistant. Solve tasks using your coding and language skills. In the following cases, suggest python code (in a python coding block) or shell script (in a sh coding block) for the user to execute. 1. When you need to collect info, use the code to output the info you need, for example, browse or search the web, download/read a file, print the content of a webpage or a file, get the current date/time, check the operating system. After sufficient info is printed and the task is ready to be solved based on your language skill, you can solve the task by yourself. 2. When you need to perform some task with code, use the code to perform the task and output the result. Finish the task smartly. Solve the task step by step if you need to. If a plan is not provided, explain your plan first. Be clear which step uses code, and which step uses your language skill. When using code, you must indicate the script type in the code block. The user cannot provide any other feedback or perform any other action beyond executing the code you suggest. The user can't modify your code. So do not suggest incomplete code which requires users to modify. Don't use a code block if it's not intended to be executed by the user. If you want the user to save the code in a file before executing it, put # filename: inside the code block as the first line. Don't include multiple code blocks in one response. Do not ask users to copy and paste the result. Instead, use 'print' function for the output when relevant. Check the execution result returned by the user. If the result indicates there is an error, fix the error and output the code again. Suggest the full code instead of partial code or code changes. If the error can't be fixed or if the task is not solved even after the code is executed successfully, analyze the problem, revisit your assumption, collect additional info you need, and think of a different approach to try. When you find an answer, verify the answer carefully. Include verifiable evidence in your response if possible. Reply 'TERMINATE' in the end when everything is done.", - }; - - const assistantFlowSpec: IAgent = { - type: "assistant", - config: assistantConfig, - }; - - const groupChatAssistantConfig = Object.assign( - { - admin_name: "groupchat_assistant", - messages: [], - max_round: 10, - speaker_selection_method: "auto", - allow_repeat_speaker: false, - }, - assistantConfig - ); - groupChatAssistantConfig.name = "groupchat_assistant"; - groupChatAssistantConfig.system_message = - "You are a helpful assistant skilled at cordinating a group of other assistants to solve a task. "; - groupChatAssistantConfig.description = "Group Chat Assistant"; - - const groupChatFlowSpec: IAgent = { - type: "groupchat", - config: groupChatAssistantConfig, - }; - - if (agent_type === "userproxy") { - return userProxyFlowSpec; - } else if (agent_type === "assistant") { - return assistantFlowSpec; - } else if (agent_type === "groupchat") { - return groupChatFlowSpec; - } else { - return assistantFlowSpec; - } -}; - -export const getSampleSkill = () => { - const content = ` -from typing import List -import uuid -import requests # to perform HTTP requests -from pathlib import Path - -from openai import OpenAI - - -def generate_and_save_images(query: str, image_size: str = "1024x1024") -> List[str]: - """ - Function to paint, draw or illustrate images based on the users query or request. Generates images from a given query using OpenAI's DALL-E model and saves them to disk. Use the code below anytime there is a request to create an image. - - :param query: A natural language description of the image to be generated. - :param image_size: The size of the image to be generated. (default is "1024x1024") - :return: A list of filenames for the saved images. - """ - - client = OpenAI() # Initialize the OpenAI client - response = client.images.generate(model="dall-e-3", prompt=query, n=1, size=image_size) # Generate images - - # List to store the file names of saved images - saved_files = [] - - # Check if the response is successful - if response.data: - for image_data in response.data: - # Generate a random UUID as the file name - file_name = str(uuid.uuid4()) + ".png" # Assuming the image is a PNG - file_path = Path(file_name) - - img_url = image_data.url - img_response = requests.get(img_url) - if img_response.status_code == 200: - # Write the binary content to a file - with open(file_path, "wb") as img_file: - img_file.write(img_response.content) - print(f"Image saved to {file_path}") - saved_files.append(str(file_path)) - else: - print(f"Failed to download the image from {img_url}") - else: - print("No image data found in the response!") - - # Return the list of saved files - return saved_files - - -# Example usage of the function: -# generate_and_save_images("A cute baby sea otter") - `; - - const skill: ISkill = { - name: "generate_and_save_images", - description: "Generate and save images based on a user's query.", - content: content, - }; - - return skill; -}; - -export const timeAgo = ( - dateString: string, - returnFormatted: boolean = false -): string => { - // if dateStr is empty, return empty string - if (!dateString) { - return ""; - } - // Parse the date string into a Date object - const timestamp = new Date(dateString); - - // Check for invalid date - if (isNaN(timestamp.getTime())) { - throw new Error("Invalid date string provided."); - } - - // Get the current time - const now = new Date(); - - // Calculate the difference in milliseconds - const timeDifference = now.getTime() - timestamp.getTime(); - - // Convert time difference to minutes and hours - const minutesAgo = Math.floor(timeDifference / (1000 * 60)); - const hoursAgo = Math.floor(minutesAgo / 60); - - // Format the date into a readable format e.g. "November 27, 2021, 3:45 PM" - const options: Intl.DateTimeFormatOptions = { - month: "long", - day: "numeric", - year: "numeric", - hour: "numeric", - minute: "numeric", - }; - const formattedDate = timestamp.toLocaleDateString(undefined, options); - - if (returnFormatted) { - return formattedDate; - } - - // Determine the time difference string - let timeAgoStr: string; - if (minutesAgo < 1) { - timeAgoStr = "just now"; - } else if (minutesAgo < 60) { - // Less than an hour ago, display minutes - timeAgoStr = `${minutesAgo} ${minutesAgo === 1 ? "minute" : "minutes"} ago`; - } else if (hoursAgo < 24) { - // Less than a day ago, display hours - timeAgoStr = `${hoursAgo} ${hoursAgo === 1 ? "hour" : "hours"} ago`; - } else { - // More than a day ago, display the formatted date - timeAgoStr = formattedDate; - } - - // Return the final readable string - return timeAgoStr; -}; - -export const examplePrompts = [ - { - title: "Stock Price", - prompt: - "Plot a chart of NVDA and TESLA stock price for 2023. Save the result to a file named nvda_tesla.png", - }, - { - title: "Sine Wave", - prompt: - "Write a python script to plot a sine wave and save it to disc as a png file sine_wave.png", - }, - { - title: "Markdown", - prompt: - "List out the top 5 rivers in africa and their length and return that as a markdown table. Do not try to write any code, just write the table", - }, - { - title: "Paint", - prompt: - "paint a picture of a glass of ethiopian coffee, freshly brewed in a tall glass cup, on a table right in front of a lush green forest scenery", - }, - { - title: "Travel", - prompt: - "Plan a 2 day trip to hawaii. Limit to 3 activities per day, be as brief as possible!", - }, -]; - export const fetchVersion = () => { const versionUrl = getServerUrl() + "/version"; return fetch(versionUrl) @@ -557,128 +116,3 @@ export const fetchVersion = () => { return null; }); }; - -/** - * Recursively sanitizes JSON objects by replacing specific keys with a given value. - * @param {JsonValue} data - The JSON data to be sanitized. - * @param {string[]} keys - An array of keys to be replaced in the JSON object. - * @param {string} replacement - The value to use as replacement for the specified keys. - * @returns {JsonValue} - The sanitized JSON data. - */ -export const sanitizeConfig = ( - data: any, - keys: string[] = ["api_key", "id", "created_at", "updated_at", "secrets"] -): any => { - if (Array.isArray(data)) { - return data.map((item) => sanitizeConfig(item, keys)); - } else if (typeof data === "object" && data !== null) { - Object.keys(data).forEach((key) => { - if (keys.includes(key)) { - delete data[key]; - } else { - data[key] = sanitizeConfig(data[key], keys); - } - }); - } - return data; -}; - -/** - * Checks the input text against the regex '^[a-zA-Z0-9_-]{1,64}$' and returns an object with - * status, message, and sanitizedText. Status is boolean indicating whether input text is valid, - * message provides information about the outcome, and sanitizedText contains a valid version - * of the input text or the original text if it was already valid. - * - * @param text - The input string to be checked and sanitized. - * @returns An object containing a status, a message, and sanitizedText. - */ -export const checkAndSanitizeInput = ( - text: string -): { status: boolean; message: string; sanitizedText: string } => { - // Create a regular expression pattern to match valid characters - const regexPattern: RegExp = /^[a-zA-Z0-9_-]{1,64}$/; - let status: boolean = true; - let message: string; - let sanitizedText: string; - - // Check if the input text matches the pattern - if (regexPattern.test(text)) { - // Text already adheres to the pattern - message = `The text '${text}' is valid.`; - sanitizedText = text; - } else { - // The text does not match; sanitize the input - status = false; - sanitizedText = text.replace(/[^a-zA-Z0-9_-]/g, "_").slice(0, 64); - message = `'${text}' is invalid. Consider using '${sanitizedText}' instead.`; - } - - return { status, message, sanitizedText }; -}; - -export const isValidConfig = ( - jsonObj: any, - templateObj: any, - diffThreshold: number = 4 -): { - status: boolean; - message: string; -} => { - // Check if both parameters are indeed objects and not null - if ( - typeof jsonObj !== "object" || - jsonObj === null || - Array.isArray(jsonObj) || - typeof templateObj !== "object" || - templateObj === null || - Array.isArray(templateObj) - ) { - return { - status: false, - message: - "Invalid input: One or both parameters are not objects, or are null or arrays.", - }; - } - - const jsonKeys = new Set(Object.keys(jsonObj)); - const templateKeys = new Set(Object.keys(templateObj)); - - if (jsonKeys.size !== templateKeys.size) { - if (Math.abs(jsonKeys.size - templateKeys.size) > diffThreshold) { - return { - status: false, - message: - "Configuration does not match template: Number of keys differ.", - }; - } - } - - for (const key of templateKeys) { - if (!jsonKeys.has(key)) { - return { - status: false, - message: `Configuration does not match template: Missing key '${key}' in configuration.`, - }; - } - - // If the value is an object, recursively validate - if ( - typeof templateObj[key] === "object" && - templateObj[key] !== null && - !Array.isArray(templateObj[key]) - ) { - const result = isValidConfig(jsonObj[key], templateObj[key]); - if (!result.status) { - return { - status: false, - message: `Configuration error in nested key '${key}': ${result.message}`, - }; - } - } - } - - return { - status: true, - message: "Configuration is valid.", - }; -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/agents.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/agents.tsx deleted file mode 100644 index 6fcb505cc7e6..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/agents.tsx +++ /dev/null @@ -1,385 +0,0 @@ -import { - ArrowDownTrayIcon, - ArrowUpTrayIcon, - DocumentDuplicateIcon, - InformationCircleIcon, - PlusIcon, - TrashIcon, -} from "@heroicons/react/24/outline"; -import { Dropdown, MenuProps, Modal, message } from "antd"; -import * as React from "react"; -import { IAgent, IStatus } from "../../types"; -import { appContext } from "../../../hooks/provider"; -import { - fetchJSON, - getServerUrl, - sanitizeConfig, - timeAgo, - truncateText, -} from "../../utils"; -import { BounceLoader, Card, CardHoverBar, LoadingOverlay } from "../../atoms"; -import { AgentViewer } from "./utils/agentconfig"; - -const AgentsView = ({}: any) => { - const [loading, setLoading] = React.useState(false); - const [error, setError] = React.useState({ - status: true, - message: "All good", - }); - - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const listAgentsUrl = `${serverUrl}/agents?user_id=${user?.email}`; - - const [agents, setAgents] = React.useState([]); - const [selectedAgent, setSelectedAgent] = React.useState(null); - - const [showNewAgentModal, setShowNewAgentModal] = React.useState(false); - - const [showAgentModal, setShowAgentModal] = React.useState(false); - - const sampleAgent = { - config: { - name: "sample_agent", - description: "Sample agent description", - human_input_mode: "NEVER", - max_consecutive_auto_reply: 3, - system_message: "", - }, - }; - const [newAgent, setNewAgent] = React.useState(sampleAgent); - - const deleteAgent = (agent: IAgent) => { - setError(null); - setLoading(true); - - const deleteAgentUrl = `${serverUrl}/agents/delete?user_id=${user?.email}&agent_id=${agent.id}`; - // const fetch; - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - user_id: user?.email, - agent: agent, - }), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchAgents(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(deleteAgentUrl, payLoad, onSuccess, onError); - }; - - const fetchAgents = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setAgents(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listAgentsUrl, payLoad, onSuccess, onError); - }; - - React.useEffect(() => { - if (user) { - // console.log("fetching messages", messages); - fetchAgents(); - } - }, []); - - const agentRows = (agents || []).map((agent: IAgent, i: number) => { - const cardItems = [ - { - title: "Download", - icon: ArrowDownTrayIcon, - onClick: (e: any) => { - e.stopPropagation(); - // download workflow as workflow.name.json - const element = document.createElement("a"); - const sanitizedAgent = sanitizeConfig(agent); - const file = new Blob([JSON.stringify(sanitizedAgent)], { - type: "application/json", - }); - element.href = URL.createObjectURL(file); - element.download = `agent_${agent.config.name}.json`; - document.body.appendChild(element); // Required for this to work in FireFox - element.click(); - }, - hoverText: "Download", - }, - { - title: "Make a Copy", - icon: DocumentDuplicateIcon, - onClick: (e: any) => { - e.stopPropagation(); - let newAgent = { ...sanitizeConfig(agent) }; - newAgent.config.name = `${agent.config.name}_copy`; - console.log("newAgent", newAgent); - setNewAgent(newAgent); - setShowNewAgentModal(true); - }, - hoverText: "Make a Copy", - }, - { - title: "Delete", - icon: TrashIcon, - onClick: (e: any) => { - e.stopPropagation(); - deleteAgent(agent); - }, - hoverText: "Delete", - }, - ]; - return ( -
  • - - {truncateText(agent.config.name || "", 25)} -
    • - } - onClick={() => { - setSelectedAgent(agent); - setShowAgentModal(true); - }} - > - -
    - {timeAgo(agent.updated_at || "")} -
    - - - - ); - }); - - const AgentModal = ({ - agent, - setAgent, - showAgentModal, - setShowAgentModal, - handler, - }: { - agent: IAgent | null; - setAgent: (agent: IAgent | null) => void; - showAgentModal: boolean; - setShowAgentModal: (show: boolean) => void; - handler?: (agent: IAgent | null) => void; - }) => { - const [localAgent, setLocalAgent] = React.useState(agent); - - const closeModal = () => { - setShowAgentModal(false); - if (handler) { - handler(localAgent); - } - }; - - return ( - Agent Configuration} - width={800} - open={showAgentModal} - onOk={() => { - closeModal(); - }} - onCancel={() => { - closeModal(); - }} - footer={[]} - > - {agent && ( - - )} - {/* {JSON.stringify(localAgent)} */} - - ); - }; - - const uploadAgent = () => { - const input = document.createElement("input"); - input.type = "file"; - input.accept = ".json"; - input.onchange = (e: any) => { - const file = e.target.files[0]; - const reader = new FileReader(); - reader.onload = (e: any) => { - const contents = e.target.result; - if (contents) { - try { - const agent = JSON.parse(contents); - // TBD validate that it is a valid agent - if (!agent.config) { - throw new Error( - "Invalid agent file. An agent must have a config" - ); - } - setNewAgent(agent); - setShowNewAgentModal(true); - } catch (err) { - message.error( - "Invalid agent file. Please upload a valid agent file." - ); - } - } - }; - reader.readAsText(file); - }; - input.click(); - }; - - const agentsMenuItems: MenuProps["items"] = [ - // { - // type: "divider", - // }, - { - key: "uploadagent", - label: ( -
    - - Upload Agent -
    - ), - }, - ]; - - const agentsMenuItemOnClick: MenuProps["onClick"] = ({ key }) => { - if (key === "uploadagent") { - uploadAgent(); - return; - } - }; - - return ( -
    - { - fetchAgents(); - }} - /> - - { - fetchAgents(); - }} - /> - -
    -
    -
    -
    - {" "} - Agents ({agentRows.length}){" "} -
    -
    - { - setShowNewAgentModal(true); - }} - > - - New Agent - -
    -
    - -
    - {" "} - Configure an agent that can reused in your agent workflow . -
    - Tip: You can also create a Group of Agents ( New Agent - - GroupChat) which can have multiple agents in it. -
    -
    - {agents && agents.length > 0 && ( -
    - -
      {agentRows}
    -
    - )} - - {agents && agents.length === 0 && !loading && ( -
    - - No agents found. Please create a new agent. -
    - )} - - {loading && ( -
    - {" "} - {" "} - loading .. -
    - )} -
    -
    -
    - ); -}; - -export default AgentsView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/build.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/build.tsx deleted file mode 100644 index bf8128fe0f9c..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/build.tsx +++ /dev/null @@ -1,81 +0,0 @@ -import * as React from "react"; -import SkillsView from "./skills"; -import AgentsView from "./agents"; -import WorkflowView from "./workflow"; -import { Tabs } from "antd"; -import { - BugAntIcon, - CpuChipIcon, - Square2StackIcon, - Square3Stack3DIcon, -} from "@heroicons/react/24/outline"; -import ModelsView from "./models"; - -const BuildView = () => { - return ( -
    - {/*
    Build
    */} -
    - {" "} - Create skills, agents and workflows for building multiagent capabilities{" "} -
    - -
    - {" "} - - {" "} - - Skills -
    - ), - key: "1", - children: , - }, - { - label: ( -
    - {" "} - - Models -
    - ), - key: "2", - children: , - }, - { - label: ( - <> - - Agents - - ), - key: "3", - children: , - }, - { - label: ( - <> - - Workflows - - ), - key: "4", - children: , - }, - ]} - /> -
    - -
    -
    - ); -}; - -export default BuildView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/models.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/models.tsx deleted file mode 100644 index 87ae739b62e7..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/models.tsx +++ /dev/null @@ -1,403 +0,0 @@ -import { - ArrowDownTrayIcon, - ArrowUpTrayIcon, - DocumentDuplicateIcon, - InformationCircleIcon, - PlusIcon, - TrashIcon, -} from "@heroicons/react/24/outline"; -import { Dropdown, MenuProps, Modal, message } from "antd"; -import * as React from "react"; -import { IModelConfig, IStatus } from "../../types"; -import { appContext } from "../../../hooks/provider"; -import { - fetchJSON, - getServerUrl, - sanitizeConfig, - timeAgo, - truncateText, -} from "../../utils"; -import { BounceLoader, Card, CardHoverBar, LoadingOverlay } from "../../atoms"; -import { ModelConfigView } from "./utils/modelconfig"; - -const ModelsView = ({}: any) => { - const [loading, setLoading] = React.useState(false); - const [error, setError] = React.useState({ - status: true, - message: "All good", - }); - - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const listModelsUrl = `${serverUrl}/models?user_id=${user?.email}`; - const createModelUrl = `${serverUrl}/models`; - const testModelUrl = `${serverUrl}/models/test`; - - const defaultModel: IModelConfig = { - model: "gpt-4-1106-preview", - description: "Sample OpenAI GPT-4 model", - user_id: user?.email, - }; - - const [models, setModels] = React.useState([]); - const [selectedModel, setSelectedModel] = React.useState( - null - ); - const [newModel, setNewModel] = React.useState( - defaultModel - ); - - const [showNewModelModal, setShowNewModelModal] = React.useState(false); - const [showModelModal, setShowModelModal] = React.useState(false); - - const deleteModel = (model: IModelConfig) => { - setError(null); - setLoading(true); - const deleteModelUrl = `${serverUrl}/models/delete?user_id=${user?.email}&model_id=${model.id}`; - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchModels(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(deleteModelUrl, payLoad, onSuccess, onError); - }; - - const fetchModels = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setModels(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listModelsUrl, payLoad, onSuccess, onError); - }; - - const createModel = (model: IModelConfig) => { - setError(null); - setLoading(true); - model.user_id = user?.email; - - const payLoad = { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, - body: JSON.stringify(model), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - const updatedModels = [data.data].concat(models || []); - setModels(updatedModels); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(createModelUrl, payLoad, onSuccess, onError); - }; - - React.useEffect(() => { - if (user) { - // console.log("fetching messages", messages); - fetchModels(); - } - }, []); - - const modelRows = (models || []).map((model: IModelConfig, i: number) => { - const cardItems = [ - { - title: "Download", - icon: ArrowDownTrayIcon, - onClick: (e: any) => { - e.stopPropagation(); - // download workflow as workflow.name.json - const element = document.createElement("a"); - const sanitizedSkill = sanitizeConfig(model); - const file = new Blob([JSON.stringify(sanitizedSkill)], { - type: "application/json", - }); - element.href = URL.createObjectURL(file); - element.download = `model_${model.model}.json`; - document.body.appendChild(element); // Required for this to work in FireFox - element.click(); - }, - hoverText: "Download", - }, - { - title: "Make a Copy", - icon: DocumentDuplicateIcon, - onClick: (e: any) => { - e.stopPropagation(); - let newModel = { ...sanitizeConfig(model) }; - newModel.model = `${model.model}_copy`; - setNewModel(newModel); - setShowNewModelModal(true); - }, - hoverText: "Make a Copy", - }, - { - title: "Delete", - icon: TrashIcon, - onClick: (e: any) => { - e.stopPropagation(); - deleteModel(model); - }, - hoverText: "Delete", - }, - ]; - return ( -
  • - {truncateText(model.model || "", 20)}
    • - } - onClick={() => { - setSelectedModel(model); - setShowModelModal(true); - }} - > -
    - {" "} - {truncateText(model.description || model.model || "", 70)} -
    -
    - {timeAgo(model.updated_at || "")} -
    - - - - ); - }); - - const ModelModal = ({ - model, - setModel, - showModelModal, - setShowModelModal, - handler, - }: { - model: IModelConfig; - setModel: (model: IModelConfig | null) => void; - showModelModal: boolean; - setShowModelModal: (show: boolean) => void; - handler?: (agent: IModelConfig) => void; - }) => { - const [localModel, setLocalModel] = React.useState(model); - - const closeModal = () => { - setModel(null); - setShowModelModal(false); - if (handler) { - handler(model); - } - }; - - return ( - - Model Specification{" "} - {model?.model}{" "} - - } - width={800} - open={showModelModal} - footer={[]} - onOk={() => { - closeModal(); - }} - onCancel={() => { - closeModal(); - }} - > - {model && ( - - )} - - ); - }; - - const uploadModel = () => { - const input = document.createElement("input"); - input.type = "file"; - input.accept = ".json"; - input.onchange = (e: any) => { - const file = e.target.files[0]; - const reader = new FileReader(); - reader.onload = (e: any) => { - const contents = e.target.result; - if (contents) { - try { - const model = JSON.parse(contents); - if (model) { - setNewModel(model); - setShowNewModelModal(true); - } - } catch (e) { - message.error("Invalid model file"); - } - } - }; - reader.readAsText(file); - }; - input.click(); - }; - - const modelsMenuItems: MenuProps["items"] = [ - // { - // type: "divider", - // }, - { - key: "uploadmodel", - label: ( -
    - - Upload Model -
    - ), - }, - ]; - - const modelsMenuItemOnClick: MenuProps["onClick"] = ({ key }) => { - if (key === "uploadmodel") { - uploadModel(); - return; - } - }; - - return ( -
    - {selectedModel && ( - { - fetchModels(); - }} - /> - )} - { - fetchModels(); - }} - /> - -
    -
    -
    -
    - {" "} - Models ({modelRows.length}){" "} -
    -
    - { - setShowNewModelModal(true); - }} - > - - New Model - -
    -
    - -
    - {" "} - Create model configurations that can be reused in your agents and - workflows. {selectedModel?.model} -
    - {models && models.length > 0 && ( -
    - -
      {modelRows}
    -
    - )} - - {models && models.length === 0 && !loading && ( -
    - - No models found. Please create a new model which can be reused - with agents. -
    - )} - - {loading && ( -
    - {" "} - {" "} - loading .. -
    - )} -
    -
    -
    - ); -}; - -export default ModelsView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/skills.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/skills.tsx deleted file mode 100644 index 7d3dfe75611f..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/skills.tsx +++ /dev/null @@ -1,380 +0,0 @@ -import { - ArrowDownTrayIcon, - ArrowUpTrayIcon, - CodeBracketIcon, - CodeBracketSquareIcon, - DocumentDuplicateIcon, - InformationCircleIcon, - KeyIcon, - PlusIcon, - TrashIcon, -} from "@heroicons/react/24/outline"; -import { Button, Input, Modal, message, MenuProps, Dropdown, Tabs } from "antd"; -import * as React from "react"; -import { ISkill, IStatus } from "../../types"; -import { appContext } from "../../../hooks/provider"; -import { - fetchJSON, - getSampleSkill, - getServerUrl, - sanitizeConfig, - timeAgo, - truncateText, -} from "../../utils"; -import { - BounceLoader, - Card, - CardHoverBar, - LoadingOverlay, - MonacoEditor, -} from "../../atoms"; -import { SkillSelector } from "./utils/selectors"; -import { SkillConfigView } from "./utils/skillconfig"; - -const SkillsView = ({}: any) => { - const [loading, setLoading] = React.useState(false); - const [error, setError] = React.useState({ - status: true, - message: "All good", - }); - - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const listSkillsUrl = `${serverUrl}/skills?user_id=${user?.email}`; - const saveSkillsUrl = `${serverUrl}/skills`; - - const [skills, setSkills] = React.useState([]); - const [selectedSkill, setSelectedSkill] = React.useState(null); - - const [showSkillModal, setShowSkillModal] = React.useState(false); - const [showNewSkillModal, setShowNewSkillModal] = React.useState(false); - - const sampleSkill = getSampleSkill(); - const [newSkill, setNewSkill] = React.useState(sampleSkill); - - const deleteSkill = (skill: ISkill) => { - setError(null); - setLoading(true); - // const fetch; - const deleteSkillUrl = `${serverUrl}/skills/delete?user_id=${user?.email}&skill_id=${skill.id}`; - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - user_id: user?.email, - skill: skill, - }), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchSkills(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(deleteSkillUrl, payLoad, onSuccess, onError); - }; - - const fetchSkills = () => { - setError(null); - setLoading(true); - // const fetch; - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - // message.success(data.message); - console.log("skills", data.data); - setSkills(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listSkillsUrl, payLoad, onSuccess, onError); - }; - - React.useEffect(() => { - if (user) { - // console.log("fetching messages", messages); - fetchSkills(); - } - }, []); - - const skillRows = (skills || []).map((skill: ISkill, i: number) => { - const cardItems = [ - { - title: "Download", - icon: ArrowDownTrayIcon, - onClick: (e: any) => { - e.stopPropagation(); - // download workflow as workflow.name.json - const element = document.createElement("a"); - const sanitizedSkill = sanitizeConfig(skill); - const file = new Blob([JSON.stringify(sanitizedSkill)], { - type: "application/json", - }); - element.href = URL.createObjectURL(file); - element.download = `skill_${skill.name}.json`; - document.body.appendChild(element); // Required for this to work in FireFox - element.click(); - }, - hoverText: "Download", - }, - { - title: "Make a Copy", - icon: DocumentDuplicateIcon, - onClick: (e: any) => { - e.stopPropagation(); - let newSkill = { ...sanitizeConfig(skill) }; - newSkill.name = `${skill.name}_copy`; - setNewSkill(newSkill); - setShowNewSkillModal(true); - }, - hoverText: "Make a Copy", - }, - { - title: "Delete", - icon: TrashIcon, - onClick: (e: any) => { - e.stopPropagation(); - deleteSkill(skill); - }, - hoverText: "Delete", - }, - ]; - return ( -
  • -
    - {" "} - { - setSelectedSkill(skill); - setShowSkillModal(true); - }} - > - -
    - {timeAgo(skill.updated_at || "")} -
    - -     -
    -
    -
    • - ); - }); - - const SkillModal = ({ - skill, - setSkill, - showSkillModal, - setShowSkillModal, - handler, - }: { - skill: ISkill | null; - setSkill: any; - showSkillModal: boolean; - setShowSkillModal: any; - handler: any; - }) => { - const editorRef = React.useRef(null); - const [localSkill, setLocalSkill] = React.useState(skill); - - const closeModal = () => { - setSkill(null); - setShowSkillModal(false); - if (handler) { - handler(skill); - } - }; - - return ( - - Skill Specification{" "} - {localSkill?.name}{" "} - - } - width={800} - open={showSkillModal} - onCancel={() => { - setShowSkillModal(false); - }} - footer={[]} - > - {localSkill && ( - - )} - - ); - }; - - const uploadSkill = () => { - const fileInput = document.createElement("input"); - fileInput.type = "file"; - fileInput.accept = ".json"; - fileInput.onchange = (e: any) => { - const file = e.target.files[0]; - const reader = new FileReader(); - reader.onload = (e) => { - const content = e.target?.result; - if (content) { - try { - const skill = JSON.parse(content as string); - if (skill) { - setNewSkill(skill); - setShowNewSkillModal(true); - } - } catch (e) { - message.error("Invalid skill file"); - } - } - }; - reader.readAsText(file); - }; - fileInput.click(); - }; - - const skillsMenuItems: MenuProps["items"] = [ - // { - // type: "divider", - // }, - { - key: "uploadskill", - label: ( -
    - - Upload Skill -
    - ), - }, - ]; - - const skillsMenuItemOnClick: MenuProps["onClick"] = ({ key }) => { - if (key === "uploadskill") { - uploadSkill(); - return; - } - }; - - return ( -
    - { - fetchSkills(); - }} - /> - - { - fetchSkills(); - }} - /> - -
    -
    -
    -
      - {" "} - Skills ({skillRows.length}){" "} -
    -
    - { - setShowNewSkillModal(true); - }} - > - - New Skill - -
    -
    -
    - {" "} - Skills are python functions that agents can use to solve tasks.{" "} -
    - {skills && skills.length > 0 && ( -
    - -
    {skillRows}
    -
    - )} - - {skills && skills.length === 0 && !loading && ( -
    - - No skills found. Please create a new skill. -
    - )} - {loading && ( -
    - {" "} - {" "} - loading .. -
    - )} -
    -
    -
    - ); -}; - -export default SkillsView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/agentconfig.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/agentconfig.tsx deleted file mode 100644 index a62e6fc6ef14..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/agentconfig.tsx +++ /dev/null @@ -1,517 +0,0 @@ -import React from "react"; -import { CollapseBox, ControlRowView } from "../../../atoms"; -import { checkAndSanitizeInput, fetchJSON, getServerUrl } from "../../../utils"; -import { - Button, - Form, - Input, - Select, - Slider, - Tabs, - message, - theme, -} from "antd"; -import { - BugAntIcon, - CpuChipIcon, - UserGroupIcon, -} from "@heroicons/react/24/outline"; -import { appContext } from "../../../../hooks/provider"; -import { - AgentSelector, - AgentTypeSelector, - ModelSelector, - SkillSelector, -} from "./selectors"; -import { IAgent, ILLMConfig } from "../../../types"; -import TextArea from "antd/es/input/TextArea"; - -const { useToken } = theme; - -export const AgentConfigView = ({ - agent, - setAgent, - close, -}: { - agent: IAgent; - setAgent: (agent: IAgent) => void; - close: () => void; -}) => { - const nameValidation = checkAndSanitizeInput(agent?.config?.name); - const [error, setError] = React.useState(null); - const [loading, setLoading] = React.useState(false); - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const createAgentUrl = `${serverUrl}/agents`; - const [controlChanged, setControlChanged] = React.useState(false); - - const onControlChange = (value: any, key: string) => { - // if (key === "llm_config") { - // if (value.config_list.length === 0) { - // value = false; - // } - // } - const updatedAgent = { - ...agent, - config: { ...agent.config, [key]: value }, - }; - - setAgent(updatedAgent); - setControlChanged(true); - }; - - const llm_config: ILLMConfig = agent?.config?.llm_config || { - config_list: [], - temperature: 0.1, - max_tokens: 4000, - }; - - const createAgent = (agent: IAgent) => { - setError(null); - setLoading(true); - // const fetch; - - console.log("agent", agent); - agent.user_id = user?.email; - const payLoad = { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, - body: JSON.stringify(agent), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - console.log("agents", data.data); - const newAgent = data.data; - setAgent(newAgent); - } else { - message.error(data.message); - } - setLoading(false); - // setNewAgent(sampleAgent); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - const onFinal = () => { - setLoading(false); - setControlChanged(false); - }; - - fetchJSON(createAgentUrl, payLoad, onSuccess, onError, onFinal); - }; - - const hasChanged = - (!controlChanged || !nameValidation.status) && agent?.id !== undefined; - - return ( -
    -
    -
    -
    - - { - onControlChange(e.target.value, "name"); - }} - /> - {!nameValidation.status && ( -
    - {nameValidation.message} -
    - )} - - } - /> - - { - onControlChange(e.target.value, "description"); - }} - /> - } - /> - - { - onControlChange(value, "max_consecutive_auto_reply"); - }} - /> - } - /> - - { - onControlChange(value, "human_input_mode"); - }} - options={ - [ - { label: "NEVER", value: "NEVER" }, - { label: "TERMINATE", value: "TERMINATE" }, - { label: "ALWAYS", value: "ALWAYS" }, - ] as any - } - /> - } - /> - - { - onControlChange(e.target.value, "system_message"); - }} - /> - } - /> - -
    - {" "} - - { - const llm_config = { - ...agent.config.llm_config, - temperature: value, - }; - onControlChange(llm_config, "llm_config"); - }} - /> - } - /> - - { - onControlChange(e.target.value, "default_auto_reply"); - }} - /> - } - /> - - { - const llm_config = { - ...agent.config.llm_config, - max_tokens: value, - }; - onControlChange(llm_config, "llm_config"); - }} - /> - } - /> - { - onControlChange(value, "code_execution_config"); - }} - options={ - [ - { label: "None", value: "none" }, - { label: "Local", value: "local" }, - { label: "Docker", value: "docker" }, - ] as any - } - /> - } - /> - -
    -
    - {/* ====================== Group Chat Config ======================= */} - {agent.type === "groupchat" && ( -
    - { - if (agent?.config) { - onControlChange(value, "speaker_selection_method"); - } - }} - options={ - [ - { label: "Auto", value: "auto" }, - { label: "Round Robin", value: "round_robin" }, - { label: "Random", value: "random" }, - ] as any - } - /> - } - /> - - { - onControlChange(e.target.value, "admin_name"); - }} - /> - } - /> - - { - onControlChange(value, "max_round"); - }} - /> - } - /> - - { - onControlChange(value, "allow_repeat_speaker"); - }} - options={ - [ - { label: "True", value: true }, - { label: "False", value: false }, - ] as any - } - /> - } - /> -
    - )} -
    - - -
    - {" "} - {!hasChanged && ( - - )} - -
    -
    - ); -}; - -export const AgentViewer = ({ - agent, - setAgent, - close, -}: { - agent: IAgent | null; - setAgent: (newAgent: IAgent) => void; - close: () => void; -}) => { - let items = [ - { - label: ( -
    - {" "} - - Agent Configuration -
    - ), - key: "1", - children: ( -
    - {!agent?.type && ( - - )} - - {agent?.type && agent && ( - - )} -
    - ), - }, - ]; - if (agent) { - if (agent?.id) { - if (agent.type && agent.type === "groupchat") { - items.push({ - label: ( -
    - {" "} - - Agents -
    - ), - key: "2", - children: , - }); - } - - items.push({ - label: ( -
    - {" "} - - Models -
    - ), - key: "3", - children: , - }); - - items.push({ - label: ( - <> - - Skills - - ), - key: "4", - children: , - }); - } - } - - return ( -
    - {/* */} - -
    - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/export.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/export.tsx deleted file mode 100644 index bb74bd0e2e37..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/export.tsx +++ /dev/null @@ -1,207 +0,0 @@ -import { Button, Modal, message } from "antd"; -import * as React from "react"; -import { IWorkflow } from "../../../types"; -import { ArrowDownTrayIcon } from "@heroicons/react/24/outline"; -import { - checkAndSanitizeInput, - fetchJSON, - getServerUrl, - sanitizeConfig, -} from "../../../utils"; -import { appContext } from "../../../../hooks/provider"; -import { CodeBlock } from "../../../atoms"; - -export const ExportWorkflowModal = ({ - workflow, - show, - setShow, -}: { - workflow: IWorkflow | null; - show: boolean; - setShow: (show: boolean) => void; -}) => { - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - - const [error, setError] = React.useState(null); - const [loading, setLoading] = React.useState(false); - const [workflowDetails, setWorkflowDetails] = React.useState(null); - - const getWorkflowCode = (workflow: IWorkflow) => { - const workflowCode = `from autogenstudio import WorkflowManager -# load workflow from exported json workflow file. -workflow_manager = WorkflowManager(workflow="path/to/your/workflow_.json") - -# run the workflow on a task -task_query = "What is the height of the Eiffel Tower?. Dont write code, just respond to the question." -workflow_manager.run(message=task_query)`; - return workflowCode; - }; - - const getCliWorkflowCode = (workflow: IWorkflow) => { - const workflowCode = `autogenstudio serve --workflow=workflow.json --port=5000 - `; - return workflowCode; - }; - - const getGunicornWorkflowCode = (workflow: IWorkflow) => { - const workflowCode = `gunicorn -w $((2 * $(getconf _NPROCESSORS_ONLN) + 1)) --timeout 12600 -k uvicorn.workers.UvicornWorker autogenstudio.web.app:app --bind `; - - return workflowCode; - }; - - const fetchWorkFlow = (workflow: IWorkflow) => { - setError(null); - setLoading(true); - // const fetch; - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - const downloadWorkflowUrl = `${serverUrl}/workflows/export/${workflow.id}?user_id=${user?.email}`; - - const onSuccess = (data: any) => { - if (data && data.status) { - setWorkflowDetails(data.data); - console.log("workflow details", data.data); - - const sanitized_name = - checkAndSanitizeInput(workflow.name).sanitizedText || workflow.name; - const file_name = `workflow_${sanitized_name}.json`; - const workflowData = sanitizeConfig(data.data); - const file = new Blob([JSON.stringify(workflowData)], { - type: "application/json", - }); - const downloadUrl = URL.createObjectURL(file); - const a = document.createElement("a"); - a.href = downloadUrl; - a.download = file_name; - a.click(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(downloadWorkflowUrl, payLoad, onSuccess, onError); - }; - - React.useEffect(() => { - if (workflow && workflow.id && show) { - // fetchWorkFlow(workflow.id); - console.log("workflow modal ... component loaded", workflow); - } - }, [show]); - - return ( - - Export Workflow - - {workflow?.name} - {" "} - - } - width={800} - open={show} - onOk={() => { - setShow(false); - }} - onCancel={() => { - setShow(false); - }} - footer={[]} - > -
    -
    - {" "} - You can use the following steps to start integrating your workflow - into your application.{" "} -
    - {workflow && workflow.id && ( - <> -
    -
    -
    Step 1
    -
    - Download your workflow as a JSON file by clicking the button - below. -
    - -
    - -
    -
    - -
    -
    Step 2
    -
    - Copy the following code snippet and paste it into your - application to run your workflow on a task. -
    -
    - -
    -
    -
    - -
    -
    - Step 3 (Deploy) -
    -
    - You can also deploy your workflow as an API endpoint using the - autogenstudio python CLI. -
    - -
    - - -
    - Note: this will start a endpoint on port 5000. You can change - the port by changing the port number. You can also scale this - using multiple workers (e.g., via an application server like - gunicorn) or wrap it in a docker container and deploy on a - cloud provider like Azure. -
    - - -
    -
    - - )} -
    -     - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/modelconfig.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/modelconfig.tsx deleted file mode 100644 index c4a39956ba0f..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/modelconfig.tsx +++ /dev/null @@ -1,388 +0,0 @@ -import React from "react"; -import { fetchJSON, getServerUrl, sampleModelConfig } from "../../../utils"; -import { Button, Input, message, theme } from "antd"; -import { - CpuChipIcon, - InformationCircleIcon, -} from "@heroicons/react/24/outline"; -import { IModelConfig, IStatus } from "../../../types"; -import { Card, ControlRowView } from "../../../atoms"; -import TextArea from "antd/es/input/TextArea"; -import { appContext } from "../../../../hooks/provider"; - -const ModelTypeSelector = ({ - model, - setModel, -}: { - model: IModelConfig; - setModel: (newModel: IModelConfig) => void; -}) => { - const modelTypes = [ - { - label: "OpenAI", - value: "open_ai", - description: "OpenAI or other endpoints that implement the OpenAI API", - icon: , - hint: "In addition to OpenAI models, You can also use OSS models via tools like Ollama, vLLM, LMStudio etc. that provide OpenAI compatible endpoint.", - }, - { - label: "Azure OpenAI", - value: "azure", - description: "Azure OpenAI endpoint", - icon: , - hint: "Azure OpenAI endpoint", - }, - { - label: "Gemini", - value: "google", - description: "Gemini", - icon: , - hint: "Gemini", - }, - { - label: "Claude", - value: "anthropic", - description: "Anthropic Claude", - icon: , - hint: "Anthropic Claude models", - }, - { - label: "Mistral", - value: "mistral", - description: "Mistral", - icon: , - hint: "Mistral models", - }, - ]; - - const [selectedType, setSelectedType] = React.useState( - model?.api_type - ); - - const modelTypeRows = modelTypes.map((modelType: any, i: number) => { - return ( -
  • { - setSelectedHint(modelType.hint); - }} - role="listitem" - key={"modeltype" + i} - className="w-36" - > - {modelType.label}
    • } - onClick={() => { - setSelectedType(modelType.value); - if (model) { - const sampleModel = sampleModelConfig(modelType.value); - setModel(sampleModel); - // setAgent(sampleAgent); - } - }} - > -
    - {" "} -
    {modelType.icon}
    - - {" "} - {modelType.description} - -
    - - - ); - }); - - const [selectedHint, setSelectedHint] = React.useState("open_ai"); - - return ( - <> -
    Select Model Type
    -
      {modelTypeRows}
    - -
    - - {selectedHint} -
    - - ); -}; - -const ModelConfigMainView = ({ - model, - setModel, - close, -}: { - model: IModelConfig; - setModel: (newModel: IModelConfig) => void; - close: () => void; -}) => { - const [loading, setLoading] = React.useState(false); - const [modelStatus, setModelStatus] = React.useState(null); - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - const testModelUrl = `${serverUrl}/models/test`; - const createModelUrl = `${serverUrl}/models`; - - // const [model, setmodel] = React.useState( - // model - // ); - const testModel = (model: IModelConfig) => { - setModelStatus(null); - setLoading(true); - model.user_id = user?.email; - const payLoad = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify(model), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - setModelStatus(data.data); - } else { - message.error(data.message); - } - setLoading(false); - setModelStatus(data); - }; - const onError = (err: any) => { - message.error(err.message); - setLoading(false); - }; - fetchJSON(testModelUrl, payLoad, onSuccess, onError); - }; - const createModel = (model: IModelConfig) => { - setLoading(true); - model.user_id = user?.email; - const payLoad = { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, - body: JSON.stringify(model), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - setModel(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - message.error(err.message); - setLoading(false); - }; - const onFinal = () => { - setLoading(false); - setControlChanged(false); - }; - fetchJSON(createModelUrl, payLoad, onSuccess, onError, onFinal); - }; - - const [controlChanged, setControlChanged] = React.useState(false); - - const updateModelConfig = (key: string, value: string) => { - if (model) { - const updatedModelConfig = { ...model, [key]: value }; - // setmodel(updatedModelConfig); - setModel(updatedModelConfig); - } - setControlChanged(true); - }; - - const hasChanged = !controlChanged && model.id !== undefined; - - return ( -
    -
    - Enter parameters for your{" "} - {model.api_type} model. -
    -
    -
    - { - updateModelConfig("model", e.target.value); - }} - /> - } - /> - - { - updateModelConfig("base_url", e.target.value); - }} - /> - } - /> -
    -
    - { - updateModelConfig("api_key", e.target.value); - }} - /> - } - /> - {model?.api_type == "azure" && ( - { - updateModelConfig("api_version", e.target.value); - }} - /> - } - /> - )} -
    -
    - - { - updateModelConfig("description", e.target.value); - }} - /> - } - /> - - {model?.api_type === "azure" && ( -
    - Note: For Azure OAI models, you will need to specify all fields. -
    - )} - - {modelStatus && ( -
    - - {modelStatus.message} - - {/* Note */} -
    - )} - -
    - - - {!hasChanged && ( - - )} - - -
    -
    - ); -}; - -export const ModelConfigView = ({ - model, - setModel, - close, -}: { - model: IModelConfig; - setModel: (newModel: IModelConfig) => void; - close: () => void; -}) => { - return ( -
    -
    - {!model?.api_type && ( - - )} - - {model?.api_type && model && ( - - )} -
    -
    - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/selectors.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/selectors.tsx deleted file mode 100644 index 79275fe4ba2f..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/selectors.tsx +++ /dev/null @@ -1,1359 +0,0 @@ -import React, { useEffect, useState } from "react"; -import { IAgent, IModelConfig, ISkill, IWorkflow } from "../../../types"; -import { Card } from "../../../atoms"; -import { - fetchJSON, - getSampleWorkflow, - getServerUrl, - obscureString, - sampleAgentConfig, - truncateText, -} from "../../../utils"; -import { - Divider, - Dropdown, - MenuProps, - Space, - Tooltip, - message, - theme, -} from "antd"; -import { - ArrowLongRightIcon, - ChatBubbleLeftRightIcon, - CodeBracketSquareIcon, - ExclamationTriangleIcon, - InformationCircleIcon, - PlusIcon, - RectangleGroupIcon, - UserCircleIcon, - XMarkIcon, -} from "@heroicons/react/24/outline"; -import { appContext } from "../../../../hooks/provider"; - -const { useToken } = theme; - -export const SkillSelector = ({ agentId }: { agentId: number }) => { - const [error, setError] = useState(null); - const [loading, setLoading] = useState(false); - const [skills, setSkills] = useState([]); - const [agentSkills, setAgentSkills] = useState([]); - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - const listSkillsUrl = `${serverUrl}/skills?user_id=${user?.email}`; - const listAgentSkillsUrl = `${serverUrl}/agents/link/skill/${agentId}`; - - const fetchSkills = () => { - setError(null); - setLoading(true); - - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setSkills(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(listSkillsUrl, payLoad, onSuccess, onError); - }; - - const fetchAgentSkills = () => { - setError(null); - setLoading(true); - - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setAgentSkills(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(listAgentSkillsUrl, payLoad, onSuccess, onError); - }; - - const linkAgentSkill = (agentId: number, skillId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - }; - const linkSkillUrl = `${serverUrl}/agents/link/skill/${agentId}/${skillId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchAgentSkills(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(linkSkillUrl, payLoad, onSuccess, onError); - }; - - const unLinkAgentSkill = (agentId: number, skillId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - }; - const linkSkillUrl = `${serverUrl}/agents/link/skill/${agentId}/${skillId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchAgentSkills(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(linkSkillUrl, payLoad, onSuccess, onError); - }; - - useEffect(() => { - fetchSkills(); - fetchAgentSkills(); - }, [agentId]); - - const skillItems: MenuProps["items"] = skills.map((skill, index) => ({ - key: index, - label: ( - <> -
    {skill.name}
    -
    - {truncateText(skill.description || "", 20)} -
    - - ), - value: index, - })); - - const skillOnClick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedSkill = skills[selectedIndex]; - - if (selectedSkill && selectedSkill.id) { - linkAgentSkill(agentId, selectedSkill.id); - } - }; - - const { token } = useToken(); - const contentStyle: React.CSSProperties = { - backgroundColor: token.colorBgElevated, - borderRadius: token.borderRadiusLG, - boxShadow: token.boxShadowSecondary, - }; - - const handleRemoveSkill = (index: number) => { - const skill = agentSkills[index]; - if (skill && skill.id) { - unLinkAgentSkill(agentId, skill.id); - } - }; - - const AddSkillsDropDown = () => { - return ( - ( -
    - {React.cloneElement(menu as React.ReactElement, { - style: { boxShadow: "none" }, - })} - {skills.length === 0 && ( - <> - - -
    - {" "} - - {" "} - Please create skills in the Skills tab - -
    - - )} -
    - )} - > -
    - add -
    -     - ); - }; - - const agentSkillButtons = agentSkills.map((skill, i) => { - const tooltipText = ( - <> -
    {skill.name}
    -
    - {truncateText(skill.description || "", 90)} -
    - - ); - return ( -
    showModal(config, i)} - > -
    - {" "} - -
    {skill.name}
    {" "} -     -
    { - e.stopPropagation(); // Prevent opening the modal to edit - handleRemoveSkill(i); - }} - className="ml-1 text-primary hover:text-accent duration-300" - > - -
    -
    -
    - ); - }); - - return ( -
    - {agentSkills && agentSkills.length > 0 && ( -
    - {agentSkills.length} Skills - linked to this agent -
    - )} - - {(!agentSkills || agentSkills.length === 0) && ( -
    - No skills - currently linked to this agent. Please add a skill using the button - below. -
    - )} - -
    - {agentSkillButtons} - -
    -
    - ); -}; - -export const AgentTypeSelector = ({ - agent, - setAgent, -}: { - agent: IAgent | null; - setAgent: (agent: IAgent) => void; -}) => { - const iconClass = "h-6 w-6 inline-block "; - const agentTypes = [ - { - label: "User Proxy Agent", - value: "userproxy", - description: <>Typically represents the user and executes code. , - icon: , - }, - { - label: "Assistant Agent", - value: "assistant", - description: <>Plan and generate code to solve user tasks, - icon: , - }, - { - label: "GroupChat ", - value: "groupchat", - description: <>Manage group chat interactions, - icon: , - }, - ]; - const [selectedAgentType, setSelectedAgentType] = React.useState< - string | null - >(null); - - const agentTypeRows = agentTypes.map((agentType: any, i: number) => { - return ( -
  • - {agentType.label}} - onClick={() => { - setSelectedAgentType(agentType.value); - if (agent) { - const sampleAgent = sampleAgentConfig(agentType.value); - setAgent(sampleAgent); - } - }} - > -
    - {" "} -
    {agentType.icon}
    - - {" "} - {agentType.description} - -
    -     -
    • - ); - }); - - return ( - <> -
    Select Agent Type
    -
      {agentTypeRows}
    - - ); -}; - -export const WorkflowTypeSelector = ({ - workflow, - setWorkflow, -}: { - workflow: IWorkflow; - setWorkflow: (workflow: IWorkflow) => void; -}) => { - const iconClass = "h-6 w-6 inline-block "; - const workflowTypes = [ - { - label: "Autonomous (Chat)", - value: "autonomous", - description: - "Includes an initiator and receiver. The initiator is typically a user proxy agent, while the receiver could be any agent type (assistant or groupchat", - icon: , - }, - { - label: "Sequential", - value: "sequential", - description: - " Includes a list of agents in a given order. Each agent should have an nstruction and will summarize and pass on the results of their work to the next agent", - icon: , - }, - ]; - const [seletectedWorkflowType, setSelectedWorkflowType] = React.useState< - string | null - >(null); - - const workflowTypeRows = workflowTypes.map((workflowType: any, i: number) => { - return ( -
  • - {workflowType.label}} - onClick={() => { - setSelectedWorkflowType(workflowType.value); - if (workflow) { - const sampleWorkflow = getSampleWorkflow(workflowType.value); - setWorkflow(sampleWorkflow); - } - }} - > -
    - {" "} -
    {workflowType.icon}
    - - {" "} - {truncateText(workflowType.description, 60)} - -
    -     -
    • - ); - }); - - return ( - <> -
    Select Workflow Type
    -
      {workflowTypeRows}
    - - ); -}; - -export const AgentSelector = ({ agentId }: { agentId: number }) => { - const [error, setError] = useState(null); - const [loading, setLoading] = useState(false); - const [agents, setAgents] = useState([]); - const [targetAgents, setTargetAgents] = useState([]); - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - - const listAgentsUrl = `${serverUrl}/agents?user_id=${user?.email}`; - const listTargetAgentsUrl = `${serverUrl}/agents/link/agent/${agentId}`; - - const fetchAgents = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setAgents(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(listAgentsUrl, payLoad, onSuccess, onError); - }; - - const fetchTargetAgents = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setTargetAgents(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(listTargetAgentsUrl, payLoad, onSuccess, onError); - }; - - const linkAgentAgent = (agentId: number, targetAgentId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - }; - const linkAgentUrl = `${serverUrl}/agents/link/agent/${agentId}/${targetAgentId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchTargetAgents(); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(linkAgentUrl, payLoad, onSuccess, onError); - }; - - const unLinkAgentAgent = (agentId: number, targetAgentId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - }; - const linkAgentUrl = `${serverUrl}/agents/link/agent/${agentId}/${targetAgentId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchTargetAgents(); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(linkAgentUrl, payLoad, onSuccess, onError); - }; - - useEffect(() => { - fetchAgents(); - fetchTargetAgents(); - }, []); - - const agentItems: MenuProps["items"] = - agents.length > 0 - ? agents.map((agent, index) => ({ - key: index, - label: ( - <> -
    {agent.config.name}
    -
    - {truncateText(agent.config.description || "", 20)} -
    - - ), - value: index, - })) - : [ - { - key: -1, - label: <>No agents found, - value: 0, - }, - ]; - - const agentOnClick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedAgent = agents[selectedIndex]; - - if (selectedAgent && selectedAgent.id) { - linkAgentAgent(agentId, selectedAgent.id); - } - }; - - const handleRemoveAgent = (index: number) => { - const agent = targetAgents[index]; - if (agent && agent.id) { - unLinkAgentAgent(agentId, agent.id); - } - }; - - const { token } = useToken(); - const contentStyle: React.CSSProperties = { - backgroundColor: token.colorBgElevated, - borderRadius: token.borderRadiusLG, - boxShadow: token.boxShadowSecondary, - }; - - const AddAgentDropDown = () => { - return ( - ( -
    - {React.cloneElement(menu as React.ReactElement, { - style: { boxShadow: "none" }, - })} - {agents.length === 0 && ( - <> - - -
    - {" "} - - {" "} - Please create agents in the Agents tab - -
    - - )} -
    - )} - > -
    - add -
    -     - ); - }; - - const agentButtons = targetAgents.map((agent, i) => { - const tooltipText = ( - <> -
    {agent.config.name}
    -
    - {truncateText(agent.config.description || "", 90)} -
    - - ); - return ( -
    -
    - {" "} - -
    {agent.config.name}
    {" "} -     -
    { - e.stopPropagation(); // Prevent opening the modal to edit - handleRemoveAgent(i); - }} - className="ml-1 text-primary hover:text-accent duration-300" - > - -
    -
    -
    - ); - }); - - return ( -
    - {targetAgents && targetAgents.length > 0 && ( -
    - {targetAgents.length} Agents - linked to this agent -
    - )} - - {(!targetAgents || targetAgents.length === 0) && ( -
    - No agents - currently linked to this agent. Please add an agent using the button - below. -
    - )} - -
    - {agentButtons} - -
    -
    - ); -}; - -export const ModelSelector = ({ agentId }: { agentId: number }) => { - const [error, setError] = useState(null); - const [loading, setLoading] = useState(false); - const [models, setModels] = useState([]); - const [agentModels, setAgentModels] = useState([]); - const serverUrl = getServerUrl(); - - const { user } = React.useContext(appContext); - const listModelsUrl = `${serverUrl}/models?user_id=${user?.email}`; - const listAgentModelsUrl = `${serverUrl}/agents/link/model/${agentId}`; - - const fetchModels = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - // message.success(data.message); - setModels(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listModelsUrl, payLoad, onSuccess, onError); - }; - - const fetchAgentModels = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - // message.success(data.message); - setAgentModels(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listAgentModelsUrl, payLoad, onSuccess, onError); - }; - - const linkAgentModel = (agentId: number, modelId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - }; - const linkModelUrl = `${serverUrl}/agents/link/model/${agentId}/${modelId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - console.log("linked model", data); - fetchAgentModels(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(linkModelUrl, payLoad, onSuccess, onError); - }; - - const unLinkAgentModel = (agentId: number, modelId: number) => { - setError(null); - setLoading(true); - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - }; - const linkModelUrl = `${serverUrl}/agents/link/model/${agentId}/${modelId}`; - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - console.log("unlinked model", data); - fetchAgentModels(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(linkModelUrl, payLoad, onSuccess, onError); - }; - - useEffect(() => { - fetchModels(); - fetchAgentModels(); - }, []); - - const modelItems: MenuProps["items"] = - models.length > 0 - ? models.map((model: IModelConfig, index: number) => ({ - key: index, - label: ( - <> -
    {model.model}
    -
    - {truncateText(model.description || "", 20)} -
    - - ), - value: index, - })) - : [ - { - key: -1, - label: <>No models found, - value: 0, - }, - ]; - - const modelOnClick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedModel = models[selectedIndex]; - - console.log("selected model", selectedModel); - if (selectedModel && selectedModel.id) { - linkAgentModel(agentId, selectedModel.id); - } - }; - - const menuStyle: React.CSSProperties = { - boxShadow: "none", - }; - - const { token } = useToken(); - const contentStyle: React.CSSProperties = { - backgroundColor: token.colorBgElevated, - borderRadius: token.borderRadiusLG, - boxShadow: token.boxShadowSecondary, - }; - - const AddModelsDropDown = () => { - return ( - ( -
    - {React.cloneElement(menu as React.ReactElement, { - style: menuStyle, - })} - {models.length === 0 && ( - <> - - -
    - - {" "} - {" "} - Please create models in the Model tab - -
    - - )} -
    - )} - > -
    - add -
    -     - ); - }; - - const handleRemoveModel = (index: number) => { - const model = agentModels[index]; - if (model && model.id) { - unLinkAgentModel(agentId, model.id); - } - }; - - const agentModelButtons = agentModels.map((model, i) => { - const tooltipText = ( - <> -
    {model.model}
    - {model.base_url &&
    {model.base_url}
    } - {model.api_key &&
    {obscureString(model.api_key, 3)}
    } -
    - {truncateText(model.description || "", 90)} -
    - - ); - return ( -
    showModal(config, i)} - > -
    - {" "} - -
    {model.model}
    {" "} -     -
    { - e.stopPropagation(); // Prevent opening the modal to edit - handleRemoveModel(i); - }} - className="ml-1 text-primary hover:text-accent duration-300" - > - -
    -
    -
    - ); - }); - - return ( -
    - {agentModels && agentModels.length > 0 && ( - <> -
    - {" "} - {agentModels.length} Models - linked to this agent{" "} -
    - - )} - - {(!agentModels || agentModels.length == 0) && ( -
    - - No models currently linked to this agent. Please add a model using the - button below. -
    - )} - -
    - {agentModelButtons} - -
    -
    - ); -}; - -export const WorkflowAgentSelector = ({ - workflow, -}: { - workflow: IWorkflow; -}) => { - const [error, setError] = useState(null); - const [loading, setLoading] = useState(false); - const [agents, setAgents] = useState([]); - const [linkedAgents, setLinkedAgents] = useState([]); - - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - - const listAgentsUrl = `${serverUrl}/agents?user_id=${user?.email}`; - - const fetchAgents = () => { - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setAgents(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(listAgentsUrl, payLoad, onSuccess, onError); - }; - - const fetchLinkedAgents = () => { - const listTargetAgentsUrl = `${serverUrl}/workflows/link/agent/${workflow.id}`; - setError(null); - setLoading(true); - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setLinkedAgents(data.data); - console.log("linked agents", data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listTargetAgentsUrl, payLoad, onSuccess, onError); - }; - - const linkWorkflowAgent = ( - workflowId: number, - targetAgentId: number, - agentType: string, - sequenceId?: number - ) => { - setError(null); - setLoading(true); - const payLoad = { - method: "POST", - headers: { - "Content-Type": "application/json", - }, - }; - let linkAgentUrl; - linkAgentUrl = `${serverUrl}/workflows/link/agent/${workflowId}/${targetAgentId}/${agentType}`; - if (agentType === "sequential") { - linkAgentUrl = `${serverUrl}/workflows/link/agent/${workflowId}/${targetAgentId}/${agentType}/${sequenceId}`; - } - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchLinkedAgents(); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(linkAgentUrl, payLoad, onSuccess, onError); - }; - - const unlinkWorkflowAgent = (agent: IAgent, link: any) => { - setError(null); - setLoading(true); - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - }; - - let unlinkAgentUrl; - unlinkAgentUrl = `${serverUrl}/workflows/link/agent/${workflow.id}/${agent.id}/${link.agent_type}`; - if (link.agent_type === "sequential") { - unlinkAgentUrl = `${serverUrl}/workflows/link/agent/${workflow.id}/${agent.id}/${link.agent_type}/${link.sequence_id}`; - } - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchLinkedAgents(); - } else { - message.error(data.message); - } - setLoading(false); - }; - - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - - fetchJSON(unlinkAgentUrl, payLoad, onSuccess, onError); - }; - - useEffect(() => { - fetchAgents(); - fetchLinkedAgents(); - }, []); - - const agentItems: MenuProps["items"] = - agents.length > 0 - ? agents.map((agent, index) => ({ - key: index, - label: ( - <> -
    {agent.config.name}
    -
    - {truncateText(agent.config.description || "", 20)} -
    - - ), - value: index, - })) - : [ - { - key: -1, - label: <>No agents found, - value: 0, - }, - ]; - - const receiverOnclick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedAgent = agents[selectedIndex]; - if (selectedAgent && selectedAgent.id && workflow.id) { - linkWorkflowAgent(workflow.id, selectedAgent.id, "receiver"); - } - }; - - const sequenceOnclick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedAgent = agents[selectedIndex]; - - if (selectedAgent && selectedAgent.id && workflow.id) { - const sequenceId = - linkedAgents.length > 0 - ? linkedAgents[linkedAgents.length - 1].link.sequence_id + 1 - : 0; - linkWorkflowAgent( - workflow.id, - selectedAgent.id, - "sequential", - sequenceId - ); - } - }; - - const senderOnClick: MenuProps["onClick"] = ({ key }) => { - const selectedIndex = parseInt(key.toString()); - let selectedAgent = agents[selectedIndex]; - - if (selectedAgent && selectedAgent.id && workflow.id) { - linkWorkflowAgent(workflow.id, selectedAgent.id, "sender"); - } - }; - - const handleRemoveAgent = (agent: IAgent, link: any) => { - if (agent && agent.id && workflow.id) { - unlinkWorkflowAgent(agent, link); - } - console.log(link); - }; - - const { token } = useToken(); - const contentStyle: React.CSSProperties = { - backgroundColor: token.colorBgElevated, - borderRadius: token.borderRadiusLG, - boxShadow: token.boxShadowSecondary, - }; - - const AddAgentDropDown = ({ - title, - onClick, - agentType, - }: { - title?: string; - onClick: MenuProps["onClick"]; - agentType: string; - }) => { - const targetAgents = linkedAgents.filter( - (row) => row.link.agent_type === agentType - ); - - const agentButtons = targetAgents.map(({ agent, link }, i) => { - const tooltipText = ( - <> -
    {agent.config.name}
    -
    - {truncateText(agent.config.description || "", 90)} -
    - - ); - return ( -
    -
    - {" "} -
    - {" "} - -
    {agent.config.name}
    {" "} -     -
    { - e.stopPropagation(); // Prevent opening the modal to edit - handleRemoveAgent(agent, link); - }} - className="ml-1 text-primary hover:text-accent duration-300" - > - -
    -
    -
    - {link.agent_type === "sequential" && - i !== targetAgents.length - 1 && ( -
    - {" "} -
    - )} -
    - ); - }); - - return ( -
    -
    - {(!targetAgents || targetAgents.length === 0) && ( -
    - No{" "} - {title} agent linked to this workflow. -
    - )} -
    {agentButtons}
    -
    - - {targetAgents && targetAgents.length == 1 && ( -
    - you can - remove current agents and add new ones. -
    - )} - {((targetAgents.length < 1 && agentType !== "sequential") || - agentType === "sequential") && ( - ( -
    - {React.cloneElement(menu as React.ReactElement, { - style: { boxShadow: "none" }, - })} - {agents.length === 0 && ( - <> - - -
    - {" "} - - {" "} - Please create agents in the Agents tab - -
    - - )} -
    - )} - > -
    - {" "} -
    - Add {title} -
    -
    -     - )} -
    - ); - }; - - return ( -
    - {workflow.type === "autonomous" && ( -
    -
    -

    - Initiator{" "} - - - -

    -
      - -
    -
    -
    -

    Receiver

    -
      - -
    -
    -
    - )} - - {workflow.type === "sequential" && ( -
    -
    Agents
    -
      - -
    -
    - )} -
    - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/skillconfig.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/skillconfig.tsx deleted file mode 100644 index 8a7a2f24c7f0..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/skillconfig.tsx +++ /dev/null @@ -1,295 +0,0 @@ -import React from "react"; -import { fetchJSON, getServerUrl, sampleModelConfig } from "../../../utils"; -import { Button, Input, message, theme } from "antd"; -import { - CpuChipIcon, - EyeIcon, - EyeSlashIcon, - InformationCircleIcon, - PlusIcon, - TrashIcon, -} from "@heroicons/react/24/outline"; -import { ISkill, IStatus } from "../../../types"; -import { Card, ControlRowView, MonacoEditor } from "../../../atoms"; -import TextArea from "antd/es/input/TextArea"; -import { appContext } from "../../../../hooks/provider"; - -const SecretsEditor = ({ - secrets = [], - updateSkillConfig, -}: { - secrets: { secret: string; value: string }[]; - updateSkillConfig: (key: string, value: any) => void; -}) => { - const [editingIndex, setEditingIndex] = React.useState(null); - const [newSecret, setNewSecret] = React.useState(""); - const [newValue, setNewValue] = React.useState(""); - - const toggleEditing = (index: number) => { - setEditingIndex(editingIndex === index ? null : index); - }; - - const handleAddSecret = () => { - if (newSecret && newValue) { - const updatedSecrets = [ - ...secrets, - { secret: newSecret, value: newValue }, - ]; - updateSkillConfig("secrets", updatedSecrets); - setNewSecret(""); - setNewValue(""); - } - }; - - const handleRemoveSecret = (index: number) => { - const updatedSecrets = secrets.filter((_, i) => i !== index); - updateSkillConfig("secrets", updatedSecrets); - }; - - const handleSecretChange = (index: number, key: string, value: string) => { - const updatedSecrets = secrets.map((item, i) => - i === index ? { ...item, [key]: value } : item - ); - updateSkillConfig("secrets", updatedSecrets); - }; - - return ( -
    - {secrets && ( -
    - {secrets.map((secret, index) => ( -
    - - handleSecretChange(index, "secret", e.target.value) - } - className="flex-1" - /> - - handleSecretChange(index, "value", e.target.value) - } - className="flex-1" - /> -
    - ))} -
    - )} -
    - setNewSecret(e.target.value)} - className="flex-1" - /> - setNewValue(e.target.value)} - className="flex-1" - /> -
    -
    - ); -}; - -export const SkillConfigView = ({ - skill, - setSkill, - close, -}: { - skill: ISkill; - setSkill: (newModel: ISkill) => void; - close: () => void; -}) => { - const [loading, setLoading] = React.useState(false); - - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - const testModelUrl = `${serverUrl}/skills/test`; - const createSkillUrl = `${serverUrl}/skills`; - - const createSkill = (skill: ISkill) => { - setLoading(true); - skill.user_id = user?.email; - const payLoad = { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, - body: JSON.stringify(skill), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - setSkill(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - message.error(err.message); - setLoading(false); - }; - const onFinal = () => { - setLoading(false); - setControlChanged(false); - }; - fetchJSON(createSkillUrl, payLoad, onSuccess, onError, onFinal); - }; - - const [controlChanged, setControlChanged] = React.useState(false); - - const updateSkillConfig = (key: string, value: string) => { - if (skill) { - const updatedSkill = { ...skill, [key]: value }; - // setSkill(updatedModelConfig); - setSkill(updatedSkill); - } - setControlChanged(true); - }; - - const hasChanged = !controlChanged && skill.id !== undefined; - const editorRef = React.useRef(null); - - return ( -
    - {skill && ( -
    -
    -
    -
    -
    - { - updateSkillConfig("content", value); - }} - /> -
    -
    -
    -
    -
    - { - updateSkillConfig("name", e.target.value); - }} - /> - } - /> - - { - updateSkillConfig("description", e.target.value); - }} - /> - } - /> - - - } - /> -
    -
    -
    -
    - )} - -
    - {/* */} - - {!hasChanged && ( - - )} - - -
    -
    - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/workflowconfig.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/utils/workflowconfig.tsx deleted file mode 100644 index c42c2e9be302..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/utils/workflowconfig.tsx +++ /dev/null @@ -1,279 +0,0 @@ -import React from "react"; -import { IWorkflow, IStatus, IChatSession } from "../../../types"; -import { ControlRowView } from "../../../atoms"; -import { - fetchJSON, - getRandomIntFromDateAndSalt, - getServerUrl, -} from "../../../utils"; -import { Button, Drawer, Input, Select, Tabs, message, theme } from "antd"; -import { appContext } from "../../../../hooks/provider"; -import { BugAntIcon, UserGroupIcon } from "@heroicons/react/24/outline"; -import { WorkflowAgentSelector, WorkflowTypeSelector } from "./selectors"; -import ChatBox from "../../playground/chatbox"; - -export const WorkflowViewConfig = ({ - workflow, - setWorkflow, - close, -}: { - workflow: IWorkflow; - setWorkflow: (newFlowConfig: IWorkflow) => void; - close: () => void; -}) => { - const [loading, setLoading] = React.useState(false); - const [error, setError] = React.useState(null); - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const createWorkflowUrl = `${serverUrl}/workflows`; - - const [controlChanged, setControlChanged] = React.useState(false); - const [localWorkflow, setLocalWorkflow] = React.useState(workflow); - - const updateFlowConfig = (key: string, value: string) => { - // When an updatedFlowConfig is created using localWorkflow, if the contents of FlowConfigViewer Modal are changed after the Agent Specification Modal is updated, the updated contents of the Agent Specification Modal are not saved. Fixed to localWorkflow->flowConfig. Fixed a bug. - const updatedFlowConfig = { ...workflow, [key]: value }; - - setLocalWorkflow(updatedFlowConfig); - setWorkflow(updatedFlowConfig); - setControlChanged(true); - }; - - const createWorkflow = (workflow: IWorkflow) => { - setError(null); - setLoading(true); - // const fetch; - workflow.user_id = user?.email; - - const payLoad = { - method: "POST", - headers: { - Accept: "application/json", - "Content-Type": "application/json", - }, - body: JSON.stringify(workflow), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - const newWorkflow = data.data; - setWorkflow(newWorkflow); - } else { - message.error(data.message); - } - setLoading(false); - // setNewAgent(sampleAgent); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - const onFinal = () => { - setLoading(false); - setControlChanged(false); - }; - - fetchJSON(createWorkflowUrl, payLoad, onSuccess, onError, onFinal); - }; - - const hasChanged = !controlChanged && workflow.id !== undefined; - const [drawerOpen, setDrawerOpen] = React.useState(false); - - const openDrawer = () => { - setDrawerOpen(true); - }; - - const closeDrawer = () => { - setDrawerOpen(false); - }; - - const dummySession: IChatSession = { - user_id: user?.email || "test_session_user_id", - workflow_id: workflow?.id, - name: "test_session", - }; - - return ( - <> - {/*
    {flowConfig.name}
    */} -
    - updateFlowConfig("name", e.target.value)} - /> - } - /> - - updateFlowConfig("description", e.target.value)} - /> - } - /> - - - updateFlowConfig("summary_method", value) - } - options={ - [ - { label: "last", value: "last" }, - { label: "none", value: "none" }, - { label: "llm", value: "llm" }, - ] as any - } - /> - } - /> -
    - -
    - {" "} - {!hasChanged && ( - - )} - {workflow?.id && ( - - )} - -
    - - {workflow?.name || "Test Workflow"}} - size="large" - onClose={closeDrawer} - open={drawerOpen} - > -
    - {drawerOpen && ( - - )} -
    -     - - ); -}; - -export const WorflowViewer = ({ - workflow, - setWorkflow, - close, -}: { - workflow: IWorkflow; - setWorkflow: (workflow: IWorkflow) => void; - close: () => void; -}) => { - let items = [ - { - label: ( -
    - {" "} - - Workflow Configuration -
    - ), - key: "1", - children: ( -
    - {!workflow?.type && ( - - )} - - {workflow?.type && workflow && ( - - )} -
    - ), - }, - ]; - if (workflow) { - if (workflow?.id) { - items.push({ - label: ( -
    - {" "} - - Agents -
    - ), - key: "2", - children: ( - <> - {" "} - - ), - }); - } - } - - const { user } = React.useContext(appContext); - - return ( -
    - -
    - ); -}; diff --git a/python/packages/autogen-studio/frontend/src/components/views/builder/workflow.tsx b/python/packages/autogen-studio/frontend/src/components/views/builder/workflow.tsx deleted file mode 100644 index 025ad77c7dd2..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/builder/workflow.tsx +++ /dev/null @@ -1,428 +0,0 @@ -import { - ArrowDownTrayIcon, - ArrowUpTrayIcon, - CodeBracketSquareIcon, - DocumentDuplicateIcon, - InformationCircleIcon, - PlusIcon, - TrashIcon, - UserGroupIcon, - UsersIcon, -} from "@heroicons/react/24/outline"; -import { Dropdown, MenuProps, Modal, message } from "antd"; -import * as React from "react"; -import { IWorkflow, IStatus } from "../../types"; -import { appContext } from "../../../hooks/provider"; -import { - fetchJSON, - getServerUrl, - sanitizeConfig, - timeAgo, - truncateText, -} from "../../utils"; -import { BounceLoader, Card, CardHoverBar, LoadingOverlay } from "../../atoms"; -import { WorflowViewer } from "./utils/workflowconfig"; -import { ExportWorkflowModal } from "./utils/export"; - -const WorkflowView = ({}: any) => { - const [loading, setLoading] = React.useState(false); - const [error, setError] = React.useState({ - status: true, - message: "All good", - }); - const { user } = React.useContext(appContext); - const serverUrl = getServerUrl(); - const listWorkflowsUrl = `${serverUrl}/workflows?user_id=${user?.email}`; - const saveWorkflowsUrl = `${serverUrl}/workflows`; - - const [workflows, setWorkflows] = React.useState([]); - const [selectedWorkflow, setSelectedWorkflow] = - React.useState(null); - const [selectedExportWorkflow, setSelectedExportWorkflow] = - React.useState(null); - - const sampleWorkflow: IWorkflow = { - name: "Sample Agent Workflow", - description: "Sample Agent Workflow", - }; - const [newWorkflow, setNewWorkflow] = React.useState( - sampleWorkflow - ); - - const [showWorkflowModal, setShowWorkflowModal] = React.useState(false); - const [showNewWorkflowModal, setShowNewWorkflowModal] = React.useState(false); - - const fetchWorkFlow = () => { - setError(null); - setLoading(true); - // const fetch; - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - setWorkflows(data.data); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(listWorkflowsUrl, payLoad, onSuccess, onError); - }; - - const deleteWorkFlow = (workflow: IWorkflow) => { - setError(null); - setLoading(true); - // const fetch; - const deleteWorkflowsUrl = `${serverUrl}/workflows/delete?user_id=${user?.email}&workflow_id=${workflow.id}`; - const payLoad = { - method: "DELETE", - headers: { - "Content-Type": "application/json", - }, - body: JSON.stringify({ - user_id: user?.email, - workflow: workflow, - }), - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - message.success(data.message); - fetchWorkFlow(); - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(deleteWorkflowsUrl, payLoad, onSuccess, onError); - }; - - React.useEffect(() => { - if (user) { - // console.log("fetching messages", messages); - fetchWorkFlow(); - } - }, []); - - React.useEffect(() => { - if (selectedWorkflow) { - setShowWorkflowModal(true); - } - }, [selectedWorkflow]); - - const [showExportModal, setShowExportModal] = React.useState(false); - - const workflowRows = (workflows || []).map( - (workflow: IWorkflow, i: number) => { - const cardItems = [ - { - title: "Export", - icon: CodeBracketSquareIcon, - onClick: (e: any) => { - e.stopPropagation(); - setSelectedExportWorkflow(workflow); - setShowExportModal(true); - }, - hoverText: "Export", - }, - { - title: "Download", - icon: ArrowDownTrayIcon, - onClick: (e: any) => { - e.stopPropagation(); - // download workflow as workflow.name.json - const element = document.createElement("a"); - const sanitizedWorkflow = sanitizeConfig(workflow); - const file = new Blob([JSON.stringify(sanitizedWorkflow)], { - type: "application/json", - }); - element.href = URL.createObjectURL(file); - element.download = `workflow_${workflow.name}.json`; - document.body.appendChild(element); // Required for this to work in FireFox - element.click(); - }, - hoverText: "Download", - }, - { - title: "Make a Copy", - icon: DocumentDuplicateIcon, - onClick: (e: any) => { - e.stopPropagation(); - let newWorkflow = { ...sanitizeConfig(workflow) }; - newWorkflow.name = `${workflow.name}_copy`; - setNewWorkflow(newWorkflow); - setShowNewWorkflowModal(true); - }, - hoverText: "Make a Copy", - }, - { - title: "Delete", - icon: TrashIcon, - onClick: (e: any) => { - e.stopPropagation(); - deleteWorkFlow(workflow); - }, - hoverText: "Delete", - }, - ]; - return ( -
  • - {truncateText(workflow.name, 25)}} - onClick={() => { - setSelectedWorkflow(workflow); - }} - > - -
    - {timeAgo(workflow.updated_at || "")} -
    - - -     -
    • - ); - } - ); - - const WorkflowModal = ({ - workflow, - setWorkflow, - showModal, - setShowModal, - handler, - }: { - workflow: IWorkflow | null; - setWorkflow?: (workflow: IWorkflow | null) => void; - showModal: boolean; - setShowModal: (show: boolean) => void; - handler?: (workflow: IWorkflow) => void; - }) => { - const [localWorkflow, setLocalWorkflow] = React.useState( - workflow - ); - - const closeModal = () => { - setShowModal(false); - if (handler) { - handler(localWorkflow as IWorkflow); - } - }; - - return ( - - Workflow Specification{" "} - - {localWorkflow?.name} - {" "} - - } - width={800} - open={showModal} - onOk={() => { - closeModal(); - }} - onCancel={() => { - closeModal(); - }} - footer={[]} - > - <> - {localWorkflow && ( - - )} - - - ); - }; - - const uploadWorkflow = () => { - const input = document.createElement("input"); - input.type = "file"; - input.accept = ".json"; - input.onchange = (e: any) => { - const file = e.target.files[0]; - const reader = new FileReader(); - reader.onload = (e: any) => { - const contents = e.target.result; - if (contents) { - try { - const workflow = JSON.parse(contents); - // TBD validate that it is a valid workflow - setNewWorkflow(workflow); - setShowNewWorkflowModal(true); - } catch (err) { - message.error("Invalid workflow file"); - } - } - }; - reader.readAsText(file); - }; - input.click(); - }; - - const workflowTypes: MenuProps["items"] = [ - // { - // key: "twoagents", - // label: ( - //
    - // {" "} - // - // Two Agents - //
    - // ), - // }, - // { - // key: "groupchat", - // label: ( - //
    - // - // Group Chat - //
    - // ), - // }, - // { - // type: "divider", - // }, - { - key: "uploadworkflow", - label: ( -
    - - Upload Workflow -
    - ), - }, - ]; - - const showWorkflow = (config: IWorkflow) => { - setSelectedWorkflow(config); - setShowWorkflowModal(true); - }; - - const workflowTypesOnClick: MenuProps["onClick"] = ({ key }) => { - if (key === "uploadworkflow") { - uploadWorkflow(); - return; - } - showWorkflow(sampleWorkflow); - }; - - return ( -
    - { - fetchWorkFlow(); - }} - /> - - { - fetchWorkFlow(); - }} - /> - - - -
    -
    -
    -
    - {" "} - Workflows ({workflowRows.length}){" "} -
    -
    - { - showWorkflow(sampleWorkflow); - }} - > - - New Workflow - -
    -
    -
    - {" "} - Configure an agent workflow that can be used to handle tasks. -
    - {workflows && workflows.length > 0 && ( -
    - -
      {workflowRows}
    -
    - )} - {workflows && workflows.length === 0 && !loading && ( -
    - - No workflows found. Please create a new workflow. -
    - )} - {loading && ( -
    - {" "} - {" "} - loading .. -
    - )} -
    -
    -
    - ); -}; - -export default WorkflowView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/gallery/gallery.tsx b/python/packages/autogen-studio/frontend/src/components/views/gallery/gallery.tsx deleted file mode 100644 index 53f1be444938..000000000000 --- a/python/packages/autogen-studio/frontend/src/components/views/gallery/gallery.tsx +++ /dev/null @@ -1,207 +0,0 @@ -import * as React from "react"; -import { appContext } from "../../../hooks/provider"; -import { fetchJSON, getServerUrl, timeAgo, truncateText } from "../../utils"; -import { IGalleryItem, IStatus } from "../../types"; -import { Button, message } from "antd"; -import { BounceLoader, Card } from "../../atoms"; -import { - ChevronLeftIcon, - InformationCircleIcon, -} from "@heroicons/react/24/outline"; -import { navigate } from "gatsby"; -import ChatBox from "../playground/chatbox"; - -const GalleryView = ({ location }: any) => { - const serverUrl = getServerUrl(); - const { user } = React.useContext(appContext); - const [loading, setLoading] = React.useState(false); - const [gallery, setGallery] = React.useState(null); - const [currentGallery, setCurrentGallery] = - React.useState(null); - const listGalleryUrl = `${serverUrl}/gallery?user_id=${user?.email}`; - const [error, setError] = React.useState({ - status: true, - message: "All good", - }); - const [currentGalleryId, setCurrentGalleryId] = React.useState( - null - ); - - React.useEffect(() => { - // get gallery id from url - const urlParams = new URLSearchParams(location.search); - const galleryId = urlParams.get("id"); - - if (galleryId) { - // Fetch gallery details using the galleryId - fetchGallery(galleryId); - setCurrentGalleryId(galleryId); - } else { - // Redirect to an error page or home page if the id is not found - // navigate("/"); - fetchGallery(null); - } - }, []); - - const fetchGallery = (galleryId: string | null) => { - const fetchGalleryUrl = galleryId - ? `${serverUrl}/gallery?gallery_id=${galleryId}` - : listGalleryUrl; - setError(null); - setLoading(true); - // const fetch; - const payLoad = { - method: "GET", - headers: { - "Content-Type": "application/json", - }, - }; - - const onSuccess = (data: any) => { - if (data && data.status) { - // message.success(data.message); - console.log("gallery", data); - if (galleryId) { - // Set the currently viewed gallery item - setCurrentGallery(data.data[0]); - } else { - setGallery(data.data); - } - // Set the list of gallery items - } else { - message.error(data.message); - } - setLoading(false); - }; - const onError = (err: any) => { - setError(err); - message.error(err.message); - setLoading(false); - }; - fetchJSON(fetchGalleryUrl, payLoad, onSuccess, onError); - }; - - const GalleryContent = ({ item }: { item: IGalleryItem }) => { - return ( -
    -
    - This session contains {item.messages.length} messages and was created{" "} - {timeAgo(item.timestamp)} -
    -
    - -
    -
    - ); - }; - - const TagsView = ({ tags }: { tags: string[] }) => { - const tagsView = tags.map((tag: string, index: number) => { - return ( -
    - - {tag} - -
    - ); - }); - return
    {tagsView}
    ; - }; - - const galleryRows = gallery?.map((item: IGalleryItem, index: number) => { - const isSelected = currentGallery?.id === item.id; - return ( -
    - { - setCurrentGallery(item); - // add to history - navigate(`/gallery?id=${item.id}`); - }} - className="h-full p-2 cursor-pointer" - title={truncateText(item.messages[0]?.content || "", 20)} - > -
    - {" "} - {truncateText(item.messages[0]?.content || "", 80)} -
    -
    - {" "} - {item.messages.length} message{item.messages.length > 1 && "s"} -
    -
    - {" "} -
    -
    {timeAgo(item.timestamp)}
    -     -
    - ); - }); - - return ( -
    -
    Gallery
    - - {/* back to gallery button */} - - {currentGallery && ( -
    - -
    - )} - - {!currentGallery && ( - <> -
    - View a collection of AutoGen agent specifications and sessions{" "} -
    -
    - {galleryRows} -
    - - )} - - {gallery && gallery.length === 0 && ( -
    - - No gallery items found. Please create a chat session and publish to - gallery. -
    - )} - - {currentGallery && ( -
    - -
    - )} - - {loading && ( -
    -
    - {" "} - -
    - loading gallery -
    - )} -
    - ); -}; - -export default GalleryView; diff --git a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chat.tsx b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chat.tsx new file mode 100644 index 000000000000..55a837d52c10 --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chat.tsx @@ -0,0 +1,449 @@ +import * as React from "react"; +import { message } from "antd"; +import { getServerUrl } from "../../../utils"; +import { SessionManager } from "../../shared/session/manager"; +import { IStatus } from "../../../types/app"; +import { Message } from "../../../types/datamodel"; +import { useConfigStore } from "../../../../hooks/store"; +import { appContext } from "../../../../hooks/provider"; +import ChatInput from "./chatinput"; +import { ModelUsage, SocketMessage, ThreadState, ThreadStatus } from "./types"; +import { MessageList } from "./messagelist"; +import TeamManager from "../../shared/team/manager"; + +const logo = require("../../../../images/landing/welcome.svg").default; + +export default function ChatView({ + initMessages, +}: { + initMessages: Message[]; +}) { + const serverUrl = getServerUrl(); + const [loading, setLoading] = React.useState(false); + const [error, setError] = React.useState({ + status: true, + message: "All good", + }); + const [messages, setMessages] = React.useState(initMessages); + const [threadMessages, setThreadMessages] = React.useState< + Record + >({}); + const chatContainerRef = React.useRef(null); + + const { user } = React.useContext(appContext); + const { session, sessions } = useConfigStore(); + const [activeSockets, setActiveSockets] = React.useState< + Record + >({}); + + React.useEffect(() => { + if (chatContainerRef.current) { + chatContainerRef.current.scrollTo({ + top: chatContainerRef.current.scrollHeight, + behavior: "smooth", + }); + } + }, [messages, threadMessages]); + + React.useEffect(() => { + return () => { + Object.values(activeSockets).forEach((socket) => socket.close()); + }; + }, [activeSockets]); + + const getBaseUrl = (url: string): string => { + try { + // Remove protocol (http:// or https://) + let baseUrl = url.replace(/(^\w+:|^)\/\//, ""); + + // Handle both localhost and production cases + if (baseUrl.startsWith("localhost")) { + // For localhost, keep the port if it exists + baseUrl = baseUrl.replace("/api", ""); + } else if (baseUrl === "/api") { + // For production where url is just '/api' + baseUrl = window.location.host; + } else { + // For other cases, remove '/api' and trailing slash + baseUrl = baseUrl.replace("/api", "").replace(/\/$/, ""); + } + + return baseUrl; + } catch (error) { + console.error("Error processing server URL:", error); + throw new Error("Invalid server URL configuration"); + } + }; + + const createRun = async (sessionId: number): Promise => { + const payload = { session_id: sessionId, user_id: user?.email || "" }; + + const response = await fetch(`${serverUrl}/runs`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(payload), + }); + + if (!response.ok) { + throw new Error("Failed to create run"); + } + + const data = await response.json(); + return data.data.run_id; + }; + + const startRun = async (runId: string, query: string) => { + const messagePayload = { + user_id: user?.email, + session_id: session?.id, + config: { + content: query, + source: "user", + }, + }; + + const response = await fetch(`${serverUrl}/runs/${runId}/start`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(messagePayload), + }); + + if (!response.ok) { + throw new Error("Failed to start run"); + } + + return await response.json(); + }; + + interface RequestUsage { + prompt_tokens: number; + completion_tokens: number; + } + + const connectWebSocket = (runId: string, query: string) => { + const baseUrl = getBaseUrl(serverUrl); + // Determine if we should use ws:// or wss:// based on current protocol + const wsProtocol = window.location.protocol === "https:" ? "wss:" : "ws:"; + const wsUrl = `${wsProtocol}//${baseUrl}/api/ws/runs/${runId}`; + + console.log("Connecting to WebSocket URL:", wsUrl); // For debugging + + const socket = new WebSocket(wsUrl); + let isClosing = false; + + const closeSocket = () => { + if (!isClosing && socket.readyState !== WebSocket.CLOSED) { + isClosing = true; + socket.close(); + setActiveSockets((prev) => { + const newSockets = { ...prev }; + delete newSockets[runId]; + return newSockets; + }); + } + }; + + socket.onopen = async () => { + try { + setActiveSockets((prev) => ({ + ...prev, + [runId]: socket, + })); + + setThreadMessages((prev) => ({ + ...prev, + [runId]: { + messages: [], + status: "streaming", + isExpanded: true, + }, + })); + + setMessages((prev: Message[]) => + prev.map((msg: Message) => { + if (msg.run_id === runId && msg.config.source === "bot") { + return { + ...msg, + config: { + ...msg.config, + content: "Starting...", + }, + }; + } + return msg; + }) + ); + + // Start the run only after socket is connected + await startRun(runId, query); + } catch (error) { + console.error("Error starting run:", error); + message.error("Failed to start run"); + closeSocket(); + + setThreadMessages((prev) => ({ + ...prev, + [runId]: { + ...prev[runId], + status: "error", + isExpanded: true, + }, + })); + } + }; + + socket.onmessage = (event) => { + const message: SocketMessage = JSON.parse(event.data); + + switch (message.type) { + case "message": + setThreadMessages((prev) => { + const currentThread = prev[runId] || { + messages: [], + status: "streaming", + isExpanded: true, + }; + + const models_usage: ModelUsage | undefined = message.data + ?.models_usage + ? { + prompt_tokens: message.data.models_usage.prompt_tokens, + completion_tokens: + message.data.models_usage.completion_tokens, + } + : undefined; + + const newMessage = { + source: message.data?.source || "", + content: message.data?.content || "", + models_usage, + }; + + return { + ...prev, + [runId]: { + ...currentThread, + messages: [...currentThread.messages, newMessage], + status: "streaming", + }, + }; + }); + break; + + case "result": + case "completion": + setThreadMessages((prev) => { + const currentThread = prev[runId]; + if (!currentThread) return prev; + + const finalMessage = message.data?.task_result?.messages + ?.filter((msg: any) => msg.content !== "TERMINATE") + .pop(); + + const status: ThreadStatus = message.status || "complete"; + // Capture completion reason from task_result + const reason = + message.data?.task_result?.stop_reason || + (message.error ? `Error: ${message.error}` : undefined); + + return { + ...prev, + [runId]: { + ...currentThread, + status: status, + reason: reason, + isExpanded: true, + finalResult: finalMessage, + messages: currentThread.messages, + }, + }; + }); + closeSocket(); + break; + } + }; + + socket.onclose = (event) => { + console.log( + `WebSocket closed for run ${runId}. Code: ${event.code}, Reason: ${event.reason}` + ); + + if (!isClosing) { + setActiveSockets((prev) => { + const newSockets = { ...prev }; + delete newSockets[runId]; + return newSockets; + }); + + setThreadMessages((prev) => { + const thread = prev[runId]; + if (thread && thread.status === "streaming") { + return { + ...prev, + [runId]: { + ...thread, + status: "complete", + reason: event.reason || "Connection closed", + }, + }; + } + return prev; + }); + } + }; + + socket.onerror = (error) => { + console.error("WebSocket error:", error); + message.error("WebSocket connection error"); + + setThreadMessages((prev) => { + const thread = prev[runId]; + if (!thread) return prev; + + return { + ...prev, + [runId]: { + ...thread, + status: "error", + reason: "WebSocket connection error occurred", + isExpanded: true, + }, + }; + }); + + closeSocket(); + }; + + return socket; + }; + + const cancelRun = async (runId: string) => { + const socket = activeSockets[runId]; + if (socket && socket.readyState === WebSocket.OPEN) { + socket.send(JSON.stringify({ type: "stop" })); + + setThreadMessages((prev) => ({ + ...prev, + [runId]: { + ...prev[runId], + status: "cancelled", + reason: "Cancelled by user", + isExpanded: true, + }, + })); + } + }; + + const runTask = async (query: string) => { + setError(null); + setLoading(true); + + if (!session?.id) { + setLoading(false); + return; + } + + let runId: string | null = null; + + try { + runId = (await createRun(session.id)) + ""; + + const userMessage: Message = { + config: { + content: query, + source: "user", + }, + session_id: session.id, + run_id: runId, + }; + + const botMessage: Message = { + config: { + content: "Thinking...", + source: "bot", + }, + session_id: session.id, + run_id: runId, + }; + + setMessages((prev) => [...prev, userMessage, botMessage]); + connectWebSocket(runId, query); // Now passing query to connectWebSocket + } catch (err) { + console.error("Error:", err); + message.error("Error during request processing"); + + if (runId) { + if (activeSockets[runId]) { + activeSockets[runId].close(); + } + + setThreadMessages((prev) => ({ + ...prev, + [runId!]: { + ...prev[runId!], + status: "error", + isExpanded: true, + }, + })); + } + + setError({ + status: false, + message: err instanceof Error ? err.message : "Unknown error occurred", + }); + } finally { + setLoading(false); + } + }; + + React.useEffect(() => { + // session changed + if (session) { + setMessages([]); + setThreadMessages({}); + } + }, [session]); + + return ( +
    +
    +
    + +
    + +
    +
    +
    + +
    + + {sessions?.length === 0 ? ( +
    +
    + Welcome + Welcome! Create a session to get started! +
    +
    + ) : ( + <> + {session && ( +
    + +
    + )} + + )} +
    +
    + ); +} diff --git a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chatinput.tsx b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chatinput.tsx new file mode 100644 index 000000000000..b170f4bcc4c8 --- /dev/null +++ b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/chatinput.tsx @@ -0,0 +1,128 @@ +"use client"; + +import { + PaperAirplaneIcon, + Cog6ToothIcon, + ExclamationTriangleIcon, +} from "@heroicons/react/24/outline"; +import * as React from "react"; +import { IStatus } from "../../../types/app"; + +interface ChatInputProps { + onSubmit: (text: string) => void; + loading: boolean; + error: IStatus | null; +} +export default function ChatInput({ + onSubmit, + loading, + error, +}: ChatInputProps) { + const textAreaRef = React.useRef(null); + const [previousLoading, setPreviousLoading] = React.useState(loading); + const [text, setText] = React.useState(""); + + const textAreaDefaultHeight = "64px"; + + // Handle textarea auto-resize + React.useEffect(() => { + if (textAreaRef.current) { + textAreaRef.current.style.height = textAreaDefaultHeight; + const scrollHeight = textAreaRef.current.scrollHeight; + textAreaRef.current.style.height = `${scrollHeight}px`; + } + }, [text]); + + // Clear input when loading changes from true to false (meaning the response is complete) + React.useEffect(() => { + if (previousLoading && !loading && !error) { + resetInput(); + } + setPreviousLoading(loading); + }, [loading, error, previousLoading]); + + const resetInput = () => { + if (textAreaRef.current) { + textAreaRef.current.value = ""; + textAreaRef.current.style.height = textAreaDefaultHeight; + setText(""); + } + }; + + const handleTextChange = (event: React.ChangeEvent) => { + setText(event.target.value); + }; + + const handleSubmit = () => { + if (textAreaRef.current?.value && !loading) { + const query = textAreaRef.current.value; + onSubmit(query); + // Don't reset immediately - wait for response to complete + } + }; + + const handleKeyDown = (event: React.KeyboardEvent) => { + if (event.key === "Enter" && !event.shiftKey) { + event.preventDefault(); + handleSubmit(); + } + }; + + return ( +
    +
    +
    { + e.preventDefault(); + handleSubmit(); + }} + > +