diff --git a/python/packages/autogen-studio/autogenstudio/database/component_factory.py b/python/packages/autogen-studio/autogenstudio/database/component_factory.py index 5f3a48aa0fa5..595a11bbbfe0 100644 --- a/python/packages/autogen-studio/autogenstudio/database/component_factory.py +++ b/python/packages/autogen-studio/autogenstudio/database/component_factory.py @@ -10,9 +10,8 @@ from autogen_agentchat.task import MaxMessageTermination, StopMessageTermination, TextMentionTermination from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat from autogen_core.components.tools import FunctionTool -from autogen_ext.models import OpenAIChatCompletionClient from autogen_ext.agents import MultimodalWebSurfer - +from autogen_ext.models import OpenAIChatCompletionClient from ..datamodel.types import ( AgentConfig, @@ -37,11 +36,9 @@ AgentComponent = Union[AssistantAgent, MultimodalWebSurfer] ModelComponent = Union[OpenAIChatCompletionClient] ToolComponent = Union[FunctionTool] # Will grow with more tool types -TerminationComponent = Union[MaxMessageTermination, - StopMessageTermination, TextMentionTermination] +TerminationComponent = Union[MaxMessageTermination, StopMessageTermination, TextMentionTermination] -Component = Union[TeamComponent, AgentComponent, - ModelComponent, ToolComponent, TerminationComponent] +Component = Union[TeamComponent, AgentComponent, ModelComponent, ToolComponent, TerminationComponent] ReturnType = Literal["object", "dict", "config"] @@ -122,8 +119,7 @@ async def load( handler = handlers.get(config.component_type) if not handler: - raise ValueError( - f"Unknown component type: {config.component_type}") + raise ValueError(f"Unknown component type: {config.component_type}") return await handler(config) @@ -147,8 +143,7 @@ async def load_directory( component = await self.load(path, return_type=return_type) components.append(component) except Exception as e: - logger.info( - f"Failed to load component: {str(e)}, {path}") + logger.info(f"Failed to load component: {str(e)}, {path}") return components except Exception as e: @@ -181,11 +176,9 @@ async def load_termination(self, config: TerminationConfig) -> TerminationCompon try: if config.termination_type == TerminationTypes.COMBINATION: if not config.conditions or len(config.conditions) < 2: - raise ValueError( - "Combination termination requires at least 2 conditions") + raise ValueError("Combination termination requires at least 2 conditions") if not config.operator: - raise ValueError( - "Combination termination requires an operator (and/or)") + raise ValueError("Combination termination requires an operator (and/or)") # Load first two conditions conditions = [await self.load_termination(cond) for cond in config.conditions[:2]] @@ -200,8 +193,7 @@ async def load_termination(self, config: TerminationConfig) -> TerminationCompon elif config.termination_type == TerminationTypes.MAX_MESSAGES: if config.max_messages is None: - raise ValueError( - "max_messages parameter required for MaxMessageTermination") + raise ValueError("max_messages parameter required for MaxMessageTermination") return MaxMessageTermination(max_messages=config.max_messages) elif config.termination_type == TerminationTypes.STOP_MESSAGE: @@ -209,18 +201,15 @@ async def load_termination(self, config: TerminationConfig) -> TerminationCompon elif config.termination_type == TerminationTypes.TEXT_MENTION: if not config.text: - raise ValueError( - "text parameter required for TextMentionTermination") + raise ValueError("text parameter required for TextMentionTermination") return TextMentionTermination(text=config.text) else: - raise ValueError( - f"Unsupported termination type: {config.termination_type}") + 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)}") from e + raise ValueError(f"Termination condition creation failed: {str(e)}") from e async def load_team(self, config: TeamConfig, input_func: Optional[Callable] = None) -> TeamComponent: """Create team instance from configuration.""" @@ -246,8 +235,7 @@ async def load_team(self, config: TeamConfig, input_func: Optional[Callable] = N return RoundRobinGroupChat(participants=participants, termination_condition=termination) elif config.team_type == TeamTypes.SELECTOR: if not model_client: - raise ValueError( - "SelectorGroupChat requires a model_client") + raise ValueError("SelectorGroupChat requires a model_client") selector_prompt = config.selector_prompt if config.selector_prompt else DEFAULT_SELECTOR_PROMPT return SelectorGroupChat( participants=participants, @@ -306,8 +294,7 @@ async def load_agent(self, config: AgentConfig, input_func: Optional[Callable] = ) else: - raise ValueError( - f"Unsupported agent type: {config.agent_type}") + raise ValueError(f"Unsupported agent type: {config.agent_type}") except Exception as e: logger.error(f"Failed to create agent {config.name}: {str(e)}") @@ -323,13 +310,11 @@ async def load_model(self, config: ModelConfig) -> ModelComponent: 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) + 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}") + raise ValueError(f"Unsupported model type: {config.model_type}") except Exception as e: logger.error(f"Failed to create model {config.model}: {str(e)}") @@ -350,8 +335,7 @@ async def load_tool(self, config: ToolConfig) -> ToolComponent: if config.tool_type == ToolTypes.PYTHON_FUNCTION: tool = FunctionTool( - name=config.name, description=config.description, func=self._func_from_string( - config.content) + name=config.name, description=config.description, func=self._func_from_string(config.content) ) self._tool_cache[cache_key] = tool return tool @@ -396,8 +380,7 @@ def _is_version_supported(self, component_type: ComponentTypes, ver: str) -> boo """Check if version is supported for component type.""" try: version = Version(ver) - supported = [Version(v) - for v in self.SUPPORTED_VERSIONS[component_type]] + supported = [Version(v) for v in self.SUPPORTED_VERSIONS[component_type]] return any(version == v for v in supported) except ValueError: return False diff --git a/python/packages/autogen-studio/autogenstudio/datamodel/types.py b/python/packages/autogen-studio/autogenstudio/datamodel/types.py index e2c3bb9a3d96..795a10b56419 100644 --- a/python/packages/autogen-studio/autogenstudio/datamodel/types.py +++ b/python/packages/autogen-studio/autogenstudio/datamodel/types.py @@ -136,7 +136,6 @@ class SocketMessage(BaseModel): type: str -ComponentConfig = Union[TeamConfig, AgentConfig, - ModelConfig, ToolConfig, TerminationConfig] +ComponentConfig = Union[TeamConfig, AgentConfig, ModelConfig, ToolConfig, TerminationConfig] ComponentConfigInput = Union[str, Path, dict, ComponentConfig] diff --git a/python/packages/autogen-studio/autogenstudio/web/managers/connection.py b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py index 5322086a5b35..cc83995be6b1 100644 --- a/python/packages/autogen-studio/autogenstudio/web/managers/connection.py +++ b/python/packages/autogen-studio/autogenstudio/web/managers/connection.py @@ -5,18 +5,14 @@ from uuid import UUID from autogen_agentchat.base._task import TaskResult -from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage -from ...datamodel import Run, RunStatus, TeamResult -from ...database import DatabaseManager -from ...teammanager import TeamManager -from autogen_agentchat.messages import AgentMessage, ChatMessage, TextMessage, MultiModalMessage +from autogen_agentchat.messages import AgentMessage, ChatMessage, MultiModalMessage, TextMessage from autogen_core.base import CancellationToken +from autogen_core.components import Image as AGImage from fastapi import WebSocket, WebSocketDisconnect from ...database import DatabaseManager from ...datamodel import Message, MessageConfig, Run, RunStatus, TeamResult from ...teammanager import TeamManager -from autogen_core.components import Image as AGImage logger = logging.getLogger(__name__) @@ -42,8 +38,7 @@ def __init__(self, db_manager: DatabaseManager): def _get_stop_message(self, reason: str) -> dict: return TeamResult( - task_result=TaskResult(messages=[TextMessage( - source="user", content=reason)], stop_reason=reason), + task_result=TaskResult(messages=[TextMessage(source="user", content=reason)], stop_reason=reason), usage="", duration=0, ).model_dump() @@ -57,8 +52,7 @@ async def connect(self, websocket: WebSocket, run_id: UUID) -> bool: self._input_responses[run_id] = asyncio.Queue() await self._send_message( - run_id, {"type": "system", "status": "connected", - "timestamp": datetime.now(timezone.utc).isoformat()} + run_id, {"type": "system", "status": "connected", "timestamp": datetime.now(timezone.utc).isoformat()} ) return True @@ -80,8 +74,7 @@ async def start_stream(self, run_id: UUID, task: str, team_config: dict) -> None # Update run with task and status run = await self._get_run(run_id) if run: - run.task = MessageConfig( - content=task, source="user").model_dump() + run.task = MessageConfig(content=task, source="user").model_dump() run.status = RunStatus.ACTIVE self.db_manager.upsert(run) @@ -91,8 +84,7 @@ async def start_stream(self, run_id: UUID, task: str, team_config: dict) -> None task=task, team_config=team_config, input_func=input_func, cancellation_token=cancellation_token ): if cancellation_token.is_cancelled() or run_id in self._closed_connections: - logger.info( - f"Stream cancelled or connection closed for run {run_id}") + logger.info(f"Stream cancelled or connection closed for run {run_id}") break formatted_message = self._format_message(message) @@ -110,8 +102,7 @@ async def start_stream(self, run_id: UUID, task: str, team_config: dict) -> None if final_result: await self._update_run(run_id, RunStatus.COMPLETE, team_result=final_result) else: - logger.warning( - f"No final result captured for completed run {run_id}") + logger.warning(f"No final result captured for completed run {run_id}") await self._update_run_status(run_id, RunStatus.COMPLETE) else: await self._send_message( @@ -191,8 +182,7 @@ async def handle_input_response(self, run_id: UUID, response: str) -> None: if run_id in self._input_responses: await self._input_responses[run_id].put(response) else: - logger.warning( - f"Received input response for inactive run {run_id}") + logger.warning(f"Received input response for inactive run {run_id}") async def stop_run(self, run_id: UUID, reason: str) -> None: if run_id in self._cancellation_tokens: @@ -247,8 +237,7 @@ async def _send_message(self, run_id: UUID, message: dict) -> None: message: Message dictionary to send """ if run_id in self._closed_connections: - logger.warning( - f"Attempted to send message to closed connection for run {run_id}") + logger.warning(f"Attempted to send message to closed connection for run {run_id}") return try: @@ -256,12 +245,10 @@ async def _send_message(self, run_id: UUID, message: dict) -> None: websocket = self._connections[run_id] await websocket.send_json(message) except WebSocketDisconnect: - logger.warning( - f"WebSocket disconnected while sending message for run {run_id}") + 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}, {message}") + logger.error(f"Error sending message for run {run_id}: {e}, {message}") # Don't try to send error message here to avoid potential recursive loop await self._update_run_status(run_id, RunStatus.ERROR, str(e)) await self.disconnect(run_id) @@ -303,7 +290,10 @@ def _format_message(self, message: Any) -> Optional[dict]: message_dump = message.model_dump() message_dump["content"] = [ message_dump["content"][0], - {"url": f"data:image/png;base64,{message_dump["content"][1]['data']}", "alt": "WebSurfer Screenshot"}, + { + "url": f"data:image/png;base64,{message_dump['content'][1]['data']}", + "alt": "WebSurfer Screenshot", + }, ] return {"type": "message", "data": message_dump} elif isinstance(message, (AgentMessage, ChatMessage)): @@ -329,8 +319,7 @@ async def _get_run(self, run_id: UUID) -> Optional[Run]: Returns: Optional[Run]: Run object if found, None otherwise """ - response = self.db_manager.get( - Run, filters={"id": run_id}, return_json=False) + 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: @@ -349,8 +338,7 @@ async def _update_run_status(self, run_id: UUID, status: RunStatus, error: Optio async def cleanup(self) -> None: """Clean up all active connections and resources when server is shutting down""" - logger.info( - f"Cleaning up {len(self.active_connections)} active connections") + logger.info(f"Cleaning up {len(self.active_connections)} active connections") try: # First cancel all running tasks @@ -361,8 +349,7 @@ async def cleanup(self) -> None: if run and run.status == RunStatus.ACTIVE: interrupted_result = TeamResult( task_result=TaskResult( - messages=[TextMessage( - source="system", content="Run interrupted by server shutdown")], + messages=[TextMessage(source="system", content="Run interrupted by server shutdown")], stop_reason="server_shutdown", ), usage="", diff --git a/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts b/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts index 26de84ae41b1..e408a1d84562 100644 --- a/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts +++ b/python/packages/autogen-studio/frontend/src/components/types/datamodel.ts @@ -6,6 +6,7 @@ export interface RequestUsage { export interface ImageContent { url: string; alt?: string; + data?: string; } export interface FunctionCall { @@ -95,15 +96,6 @@ export interface BaseConfig { version?: string; } -// WebSocket message types -export type ThreadStatus = - | "streaming" - | "complete" - | "error" - | "cancelled" - | "awaiting_input" - | "timeout"; - export interface WebSocketMessage { type: "message" | "result" | "completion" | "input_request" | "error"; data?: AgentMessageConfig | TaskResult; @@ -119,7 +111,10 @@ export interface TaskResult { export type ModelTypes = "OpenAIChatCompletionClient"; -export type AgentTypes = "AssistantAgent" | "CodingAssistantAgent" | "MultimodalWebSurfer"; +export type AgentTypes = + | "AssistantAgent" + | "CodingAssistantAgent" + | "MultimodalWebSurfer"; export type TeamTypes = "RoundRobinGroupChat" | "SelectorGroupChat"; @@ -197,14 +192,6 @@ export interface Run { error_message?: string; } -// Separate transient state -interface TransientRunState { - pendingInput?: { - prompt: string; - isPending: boolean; - }; -} - export type RunStatus = | "created" | "active" // covers 'streaming' diff --git a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/agentflow/agentnode.tsx b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/agentflow/agentnode.tsx index e9a66fd5b9ba..626a31499dbb 100644 --- a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/agentflow/agentnode.tsx +++ b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/agentflow/agentnode.tsx @@ -8,7 +8,7 @@ import { Bot, Flag, } from "lucide-react"; -import { ThreadStatus } from "../../../../types/datamodel"; +import { RunStatus } from "../../../../types/datamodel"; export type NodeType = "agent" | "user" | "end"; @@ -18,7 +18,7 @@ export interface AgentNodeData { agentType?: string; description?: string; isActive?: boolean; - status?: ThreadStatus | null; + status?: RunStatus | null; reason?: string | null; draggable: boolean; } @@ -54,7 +54,7 @@ function AgentNode({ data, isConnectable }: AgentNodeProps) { return ; case "error": return ; - case "cancelled": + case "stopped": return ; default: return null; diff --git a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/rendermessage.tsx b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/rendermessage.tsx index 66d4ed7e0403..d9f9346fce2e 100644 --- a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/rendermessage.tsx +++ b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/rendermessage.tsx @@ -1,11 +1,77 @@ -import React from "react"; -import { User, Bot } from "lucide-react"; +import React, { useState, memo } from "react"; +import { User, Bot, Maximize2, Minimize2 } from "lucide-react"; import { AgentMessageConfig, FunctionCall, FunctionExecutionResult, ImageContent, } from "../../../types/datamodel"; +import { ClickableImage, TruncatableText } from "../../shared/atoms"; + +const TEXT_THRESHOLD = 400; +const JSON_THRESHOLD = 800; + +// Helper function to get image source from either format +const getImageSource = (item: ImageContent): string => { + if (item.url) { + return item.url; + } + if (item.data) { + // Assume PNG if no type specified - we can enhance this later if needed + return `data:image/png;base64,${item.data}`; + } + // Fallback placeholder if neither url nor data is present + return "/api/placeholder/400/320"; +}; + +const RenderMultiModal: React.FC<{ content: (string | ImageContent)[] }> = ({ + content, +}) => ( +
+ {content.map((item, index) => + typeof item === "string" ? ( + + ) : ( + + ) + )} +
+); +const RenderToolCall: React.FC<{ content: FunctionCall[] }> = ({ content }) => ( +
+ {content.map((call) => ( +
+
Function: {call.name}
+ +
+ ))} +
+); + +const RenderToolResult: React.FC<{ content: FunctionExecutionResult[] }> = ({ + content, +}) => ( +
+ {content.map((result) => ( +
+
Result ID: {result.call_id}
+ +
+ ))} +
+); export const messageUtils = { isToolCallContent(content: unknown): content is FunctionCall[] { @@ -25,7 +91,9 @@ export const messageUtils = { return content.every( (item) => typeof item === "string" || - (typeof item === "object" && item !== null && "url" in item) + (typeof item === "object" && + item !== null && + ("url" in item || "data" in item)) ); }, @@ -53,53 +121,6 @@ interface MessageProps { className?: string; } -const RenderToolCall: React.FC<{ content: FunctionCall[] }> = ({ content }) => ( -
- {content.map((call) => ( -
-
Function: {call.name}
-
- {JSON.stringify(JSON.parse(call.arguments), null, 2)} -
-
- ))} -
-); - -const RenderMultiModal: React.FC<{ content: (string | ImageContent)[] }> = ({ - content, -}) => ( -
- {content.map((item, index) => - typeof item === "string" ? ( -

{item}

- ) : ( - {item.alt - ) - )} -
-); - -const RenderToolResult: React.FC<{ content: FunctionExecutionResult[] }> = ({ - content, -}) => ( -
- {content.map((result) => ( -
-
Result ID: {result.call_id}
- 
-          {result.content}
-
-
- ))} -
-); - export const RenderMessage: React.FC = ({ message, isLast = false, @@ -143,7 +164,10 @@ export const RenderMessage: React.FC = ({ ) : messageUtils.isFunctionExecutionResult(content) ? ( ) : ( -
{content}
+ )} diff --git a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/runview.tsx b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/runview.tsx index ab1a58c543df..5dd9ba8c5506 100644 --- a/python/packages/autogen-studio/frontend/src/components/views/playground/chat/runview.tsx +++ b/python/packages/autogen-studio/frontend/src/components/views/playground/chat/runview.tsx @@ -11,9 +11,9 @@ import { import { Run, Message, TeamConfig } from "../../../types/datamodel"; import AgentFlow from "./agentflow/agentflow"; import { RenderMessage } from "./rendermessage"; -import LoadingDots from "../../shared/atoms"; import InputRequestView from "./inputrequest"; import { Tooltip } from "antd"; +import { LoadingDots } from "../../shared/atoms"; interface RunViewProps { run: Run; diff --git a/python/packages/autogen-studio/frontend/src/components/views/shared/atoms.tsx b/python/packages/autogen-studio/frontend/src/components/views/shared/atoms.tsx index 15dbc8c6f5a7..fcd95ca76093 100644 --- a/python/packages/autogen-studio/frontend/src/components/views/shared/atoms.tsx +++ b/python/packages/autogen-studio/frontend/src/components/views/shared/atoms.tsx @@ -1,5 +1,5 @@ -import React from "react"; -import { Loader2 } from "lucide-react"; +import React, { memo, useState } from "react"; +import { Loader2, Maximize2, Minimize2, X } from "lucide-react"; export const LoadingIndicator = ({ size = 16 }: { size: number }) => (
@@ -40,4 +40,127 @@ export const LoadingDots = ({ size = 8 }) => { ); }; -export default LoadingDots; +export const TruncatableText = memo( + ({ + content, + isJson = false, + className = "", + jsonThreshold = 1000, + textThreshold = 500, + }: { + content: string; + isJson?: boolean; + className?: string; + jsonThreshold?: number; + textThreshold?: number; + }) => { + const [isExpanded, setIsExpanded] = useState(false); + const threshold = isJson ? jsonThreshold : textThreshold; + const shouldTruncate = content.length > threshold; + + const toggleExpand = () => { + setIsExpanded(!isExpanded); + }; + + const displayContent = + shouldTruncate && !isExpanded + ? content.slice(0, threshold) + "..." + : content; + + return ( +
+
+ {displayContent} + {shouldTruncate && !isExpanded && ( +
+ )} +
+ + {shouldTruncate && ( +
+ +
+ )} +
+ ); + } +); + +const FullScreenImage: React.FC<{ + src: string; + alt: string; + onClose: () => void; +}> = ({ src, alt, onClose }) => { + return ( +
+ + {alt} e.stopPropagation()} + /> +
+ ); +}; + +export const ClickableImage: React.FC<{ + src: string; + alt: string; + className?: string; +}> = ({ src, alt, className = "" }) => { + const [isFullScreen, setIsFullScreen] = useState(false); + + return ( + <> + {alt} setIsFullScreen(true)} + /> + {isFullScreen && ( + setIsFullScreen(false)} + /> + )} + + ); +};