AGS - Test Model Component in UI, Compare Sessions

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

@@ -1,3 +1,4 @@
+import json
 import threading
 from datetime import datetime
 from pathlib import Path
@@ -12,6 +13,13 @@
 from .schema_manager import SchemaManager
 
 
+class CustomJSONEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if hasattr(obj, "get_secret_value") and callable(obj.get_secret_value):
+            return obj.get_secret_value()
+        return super().default(obj)
+
+
 class DatabaseManager:
     _init_lock = threading.Lock()
 
@@ -29,7 +37,9 @@ def __init__(self, engine_uri: str, base_dir: Optional[Union[str, Path]] = None)
         if base_dir is not None and isinstance(base_dir, str):
             base_dir = Path(base_dir)
 
-        self.engine = create_engine(engine_uri, connect_args=connection_args)
+        self.engine = create_engine(
+            engine_uri, connect_args=connection_args, json_serializer=lambda obj: json.dumps(obj, cls=CustomJSONEncoder)
+        )
         self.schema_manager = SchemaManager(
             engine=self.engine,
             base_dir=base_dir,
@@ -163,7 +173,7 @@ def upsert(self, model: SQLModel, return_json: bool = True) -> Response:
                     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
+                    model = existing_model
                     session.add(model)
                 else:
                     session.add(model)

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

@@ -5,7 +5,7 @@
 from typing import List, Optional, Union
 
 from autogen_core import ComponentModel
-from pydantic import ConfigDict
+from pydantic import ConfigDict, SecretStr
 from sqlalchemy import ForeignKey, Integer, String
 from sqlmodel import JSON, Column, DateTime, Field, SQLModel, func
 
@@ -141,7 +141,12 @@ class Gallery(SQLModel, table=True):
         sa_column=Column(JSON),
     )
 
-    model_config = ConfigDict(json_encoders={datetime: lambda v: v.isoformat()})  # type: ignore[call-arg]
+    model_config = ConfigDict(
+        json_encoders={
+            datetime: lambda v: v.isoformat(),
+            SecretStr: lambda v: v.get_secret_value(),  # Add this line
+        }
+    )  # type: ignore[call-arg]
 
 
 class Settings(SQLModel, table=True):

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

@@ -6,7 +6,7 @@
 from autogen_agentchat.messages import BaseChatMessage
 from autogen_core import ComponentModel
 from autogen_ext.models.openai import OpenAIChatCompletionClient
-from pydantic import BaseModel, ConfigDict
+from pydantic import BaseModel, ConfigDict, SecretStr
 
 
 class MessageConfig(BaseModel):
@@ -71,9 +71,7 @@ class GalleryConfig(BaseModel):
     components: GalleryComponents
 
     model_config = ConfigDict(
-        json_encoders={
-            datetime: lambda v: v.isoformat(),
-        }
+        json_encoders={datetime: lambda v: v.isoformat(), SecretStr: lambda v: v.get_secret_value()}
     )
 
 

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

@@ -1,13 +1,16 @@
+import os
 from datetime import datetime
 from typing import List, Optional
 
+import anthropic
 from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
 from autogen_agentchat.conditions import MaxMessageTermination, TextMentionTermination
 from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat
 from autogen_core import ComponentModel
 from autogen_core.models import ModelInfo
 from autogen_ext.agents.web_surfer import MultimodalWebSurfer
 from autogen_ext.code_executors.local import LocalCommandLineCodeExecutor
+from autogen_ext.models.anthropic import AnthropicChatCompletionClient
 from autogen_ext.models.openai import OpenAIChatCompletionClient
 from autogen_ext.models.openai._openai_client import AzureOpenAIChatCompletionClient
 from autogen_ext.tools.code_execution import PythonCodeExecutionTool
@@ -132,6 +135,12 @@ def build(self) -> GalleryConfig:
 def create_default_gallery() -> GalleryConfig:
     """Create a default gallery with all components including calculator and web surfer teams."""
 
+    # model clients require API keys to be set in the environment or passed in
+    # as arguments. For testing purposes, we set them to "test" if not already set.
+    for key in ["OPENAI_API_KEY", "AZURE_OPENAI_API_KEY", "ANTHROPIC_API_KEY"]:
+        if not os.environ.get(key):
+            os.environ[key] = "test"
+
     # url = "https://raw.githubusercontent.com/microsoft/autogen/refs/heads/main/python/packages/autogen-studio/autogenstudio/gallery/default.json"
     builder = GalleryBuilder(id="gallery_default", name="Default Component Gallery")
 
@@ -149,21 +158,30 @@ def create_default_gallery() -> GalleryConfig:
     mistral_vllm_model = OpenAIChatCompletionClient(
         model="TheBloke/Mistral-7B-Instruct-v0.2-GGUF",
         base_url="http://localhost:1234/v1",
-        model_info=ModelInfo(vision=False, function_calling=True, json_output=False, family="unknown"),
+        model_info=ModelInfo(
+            vision=False, function_calling=True, json_output=False, family="unknown", structured_output=False
+        ),
     )
     builder.add_model(
         mistral_vllm_model.dump_component(),
         label="Mistral-7B Local",
         description="Local Mistral-7B model client for instruction-based generation (Ollama, LMStudio).",
     )
 
+    anthropic_model = AnthropicChatCompletionClient(model="claude-3-7-sonnet-20250219")
+    builder.add_model(
+        anthropic_model.dump_component(),
+        label="Anthropic Claude-3-7",
+        description="Anthropic Claude-3 model client.",
+    )
+
     # create an azure mode
     az_model_client = AzureOpenAIChatCompletionClient(
         azure_deployment="{your-azure-deployment}",
         model="gpt-4o-mini",
         api_version="2024-06-01",
         azure_endpoint="https://{your-custom-endpoint}.openai.azure.com/",
-        api_key="sk-...",  # For key-based authentication.
+        api_key="test",
     )
     builder.add_model(
         az_model_client.dump_component(),

python/packages/autogen-studio/autogenstudio/validation/component_test_service.py

@@ -0,0 +1,188 @@
+# api/validator/test_service.py
+import asyncio
+from typing import Any, Dict, List, Optional
+
+from autogen_core import ComponentModel
+from autogen_core.models import ChatCompletionClient, UserMessage
+from pydantic import BaseModel
+
+
+class ComponentTestResult(BaseModel):
+    status: bool
+    message: str
+    data: Optional[Any] = None
+    logs: List[str] = []
+
+
+class ComponentTestRequest(BaseModel):
+    component: ComponentModel
+    model_client: Optional[Dict[str, Any]] = None
+    timeout: Optional[int] = 30
+
+
+class ComponentTestService:
+    @staticmethod
+    async def test_agent(
+        component: ComponentModel, model_client: Optional[ChatCompletionClient] = None
+    ) -> ComponentTestResult:
+        """Test an agent component with a simple message"""
+        try:
+            from autogen_agentchat.agents import AssistantAgent
+            from autogen_agentchat.messages import TextMessage
+            from autogen_core import CancellationToken
+
+            # If model_client is provided, use it; otherwise, use the component's model (if applicable)
+            agent_config = component.config or {}
+
+            # Try to load the agent
+            try:
+                # Construct the agent with the model client if provided
+                if model_client:
+                    agent_config["model_client"] = model_client
+
+                agent = AssistantAgent(name=agent_config.get("name", "assistant"), **agent_config)
+
+                logs = ["Agent component loaded successfully"]
+            except Exception as e:
+                return ComponentTestResult(
+                    status=False,
+                    message=f"Failed to initialize agent: {str(e)}",
+                    logs=[f"Agent initialization error: {str(e)}"],
+                )
+
+            # Test the agent with a simple message
+            test_question = "What is 2+2? Keep it brief."
+            try:
+                response = await agent.on_messages(
+                    [TextMessage(content=test_question, source="user")],
+                    cancellation_token=CancellationToken(),
+                )
+
+                # Check if we got a valid response
+                status = response and response.chat_message is not None
+
+                if status:
+                    logs.append(
+                        f"Agent responded with: {response.chat_message.content} to the question : {test_question}"
+                    )
+                else:
+                    logs.append("Agent did not return a valid response")
+
+                return ComponentTestResult(
+                    status=status,
+                    message="Agent test completed successfully" if status else "Agent test failed - no valid response",
+                    data=response.chat_message.model_dump() if status else None,
+                    logs=logs,
+                )
+            except Exception as e:
+                return ComponentTestResult(
+                    status=False,
+                    message=f"Error during agent response: {str(e)}",
+                    logs=logs + [f"Agent response error: {str(e)}"],
+                )
+
+        except Exception as e:
+            return ComponentTestResult(
+                status=False, message=f"Error testing agent component: {str(e)}", logs=[f"Exception: {str(e)}"]
+            )
+
+    @staticmethod
+    async def test_model(
+        component: ComponentModel, model_client: Optional[ChatCompletionClient] = None
+    ) -> ComponentTestResult:
+        """Test a model component with a simple prompt"""
+        try:
+            # Use the component itself as a model client
+            model = ChatCompletionClient.load_component(component)
+
+            # Prepare a simple test message
+            test_question = "What is 2+2? Give me only the answer."
+            messages = [UserMessage(content=test_question, source="user")]
+
+            # Try to get a response
+            response = await model.create(messages=messages)
+
+            # Test passes if we got a response with content
+            status = response and response.content is not None
+
+            logs = ["Model component loaded successfully"]
+            if status:
+                logs.append(f"Model responded with: {response.content} (Query:{test_question})")
+            else:
+                logs.append("Model did not return a valid response")
+
+            return ComponentTestResult(
+                status=status,
+                message="Model test completed successfully" if status else "Model test failed - no valid response",
+                data=response.model_dump() if status else None,
+                logs=logs,
+            )
+        except Exception as e:
+            return ComponentTestResult(
+                status=False, message=f"Error testing model component: {str(e)}", logs=[f"Exception: {str(e)}"]
+            )
+
+    @staticmethod
+    async def test_tool(component: ComponentModel) -> ComponentTestResult:
+        """Test a tool component with sample inputs"""
+        # Placeholder for tool test logic
+        return ComponentTestResult(
+            status=True, message="Tool test not yet implemented", logs=["Tool component loaded successfully"]
+        )
+
+    @staticmethod
+    async def test_team(
+        component: ComponentModel, model_client: Optional[ChatCompletionClient] = None
+    ) -> ComponentTestResult:
+        """Test a team component with a simple task"""
+        # Placeholder for team test logic
+        return ComponentTestResult(
+            status=True, message="Team test not yet implemented", logs=["Team component loaded successfully"]
+        )
+
+    @staticmethod
+    async def test_termination(component: ComponentModel) -> ComponentTestResult:
+        """Test a termination component with sample message history"""
+        # Placeholder for termination test logic
+        return ComponentTestResult(
+            status=True,
+            message="Termination test not yet implemented",
+            logs=["Termination component loaded successfully"],
+        )
+
+    @classmethod
+    async def test_component(
+        cls, component: ComponentModel, timeout: int = 60, model_client: Optional[ChatCompletionClient] = None
+    ) -> ComponentTestResult:
+        """Test a component based on its type with appropriate test inputs"""
+        try:
+            # Get component type
+            component_type = component.component_type
+
+            # Select test method based on component type
+            test_method = {
+                "agent": cls.test_agent,
+                "model": cls.test_model,
+                "tool": cls.test_tool,
+                "team": cls.test_team,
+                "termination": cls.test_termination,
+            }.get(component_type or "unknown")
+
+            if not test_method:
+                return ComponentTestResult(status=False, message=f"Unknown component type: {component_type}")
+
+            # Determine if the test method accepts a model_client parameter
+            accepts_model_client = component_type in ["agent", "model", "team"]
+
+            # Run test with timeout
+            try:
+                if accepts_model_client:
+                    result = await asyncio.wait_for(test_method(component, model_client), timeout=timeout)
+                else:
+                    result = await asyncio.wait_for(test_method(component), timeout=timeout)
+                return result
+            except asyncio.TimeoutError:
+                return ComponentTestResult(status=False, message=f"Component test exceeded the {timeout}s timeout")
+
+        except Exception as e:
+            return ComponentTestResult(status=False, message=f"Error testing component: {str(e)}")