update agent base and agent factory

Author: Fr4nc3

src/backend/sql_agents/agent_base.py

@@ -2,7 +2,7 @@
 
 import logging
 from abc import ABC, abstractmethod
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, List, Optional, TypeVar, Union
 
 from azure.ai.projects.models import (
     ResponseFormatJsonSchema,
@@ -30,26 +30,61 @@ def __init__(
         agent_type: AgentType,
         config: AgentsConfigDialect,
         deployment_name: AgentModelDeployment,
-
+        temperature: float = 0.0,
     ):
-        """Initialize the base SQL agent."""
+        """Initialize the base SQL agent.
+        
+        Args:
+            agent_type: The type of agent to create.
+            config: The dialect configuration for the agent.
+            deployment_name: The model deployment to use.
+            temperature: The temperature parameter for the model.
+        """
         self.agent_type = agent_type
         self.config = config
         self.deployment_name = deployment_name
+        self.temperature = temperature
         self.agent: AzureAIAgent = None
 
-
     @property
     @abstractmethod
     def response_schema(self) -> type:
         """Get the response schema for this agent."""
         pass
 
     @property
-    @abstractmethod
-    def num_candidates(self) -> int:
-        """Get the number of candidates for this agent."""
-        pass
+    def num_candidates(self) -> Optional[int]:
+        """Get the number of candidates for this agent.
+        
+        Returns:
+            The number of candidates, or None if not applicable.
+        """
+        return None
+
+    @property
+    def plugins(self) -> Optional[List[Union[str, Any]]]:
+        """Get the plugins for this agent.
+        
+        Returns:
+            A list of plugins, or None if not applicable.
+        """
+        return None
+
+    def get_kernel_arguments(self) -> KernelArguments:
+        """Get the kernel arguments for this agent.
+        
+        Returns:
+            A KernelArguments object with the necessary arguments.
+        """
+        args = {
+            "target": self.config.sql_dialect_out,
+            "source": self.config.sql_dialect_in,
+        }
+        
+        if self.num_candidates is not None:
+            args["numCandidates"] = str(self.num_candidates)
+            
+        return KernelArguments(**args)
 
     async def setup(self) -> AzureAIAgent:
         """Setup the agent with Azure AI."""
@@ -62,11 +97,7 @@ async def setup(self) -> AzureAIAgent:
             logger.error("Prompt file for %s not found.", _name)
             raise ValueError(f"Prompt file for {_name} not found.") from exc
 
-        kernel_args = KernelArguments(
-            target=self.config.sql_dialect_out,
-            numCandidates=str(self.num_candidates),
-            source=self.config.sql_dialect_in,
-        )
+        kernel_args = self.get_kernel_arguments()
 
         # Define an agent on the Azure AI agent service
         agent_definition = await app_config.ai_project_client.agents.create_agent(
@@ -84,11 +115,17 @@ async def setup(self) -> AzureAIAgent:
         )
 
         # Create a Semantic Kernel agent based on the agent definition
-        self.agent = AzureAIAgent(
-            client=app_config.ai_project_client,
-            definition=agent_definition,
-            arguments=kernel_args,
-        )
+        agent_kwargs = {
+            "client": app_config.ai_project_client,
+            "definition": agent_definition,
+            "arguments": kernel_args,
+        }
+        
+        # Add plugins if specified
+        if self.plugins:
+            agent_kwargs["plugins"] = self.plugins
+            
+        self.agent = AzureAIAgent(**agent_kwargs)
 
         return self.agent
 

src/backend/sql_agents/agent_factory.py

@@ -1,13 +1,13 @@
 """Factory for creating SQL migration agents."""
 
 import logging
-from typing import Type, TypeVar
+from typing import Type, TypeVar, Optional, Dict, Any
 
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 from sql_agents.agent_base import BaseSQLAgent
 from sql_agents.agent_config import AgentModelDeployment, AgentsConfigDialect
-from sql_agents.specific_agents import MigratorAgent, PickerAgent
+from sql_agents.specific_agents import MigratorAgent, PickerAgent, SyntaxCheckerAgent
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -22,6 +22,7 @@ class SQLAgentFactory:
     _agent_classes = {
         AgentType.PICKER: PickerAgent,
         AgentType.MIGRATOR: MigratorAgent,
+        AgentType.SYNTAX_CHECKER: SyntaxCheckerAgent,
     }
 
     @classmethod
@@ -30,13 +31,38 @@ async def create_agent(
         agent_type: AgentType,
         config: AgentsConfigDialect,
         deployment_name: AgentModelDeployment,
+        temperature: float = 0.0,
+        extra_params: Optional[Dict[str, Any]] = None,
     ) -> AzureAIAgent:
-        """Create and setup an agent of the specified type."""
+        """Create and setup an agent of the specified type.
+        
+        Args:
+            agent_type: The type of agent to create.
+            config: The dialect configuration for the agent.
+            deployment_name: The model deployment to use.
+            temperature: The temperature parameter for the model.
+            extra_params: Additional parameters to pass to the agent constructor.
+            
+        Returns:
+            A configured AzureAIAgent instance.
+        """
         agent_class = cls._agent_classes.get(agent_type)
         if not agent_class:
             raise ValueError(f"Unknown agent type: {agent_type}")
 
-        agent = agent_class(agent_type, config, deployment_name)
+        # Prepare constructor parameters
+        params = {
+            "agent_type": agent_type,
+            "config": config,
+            "deployment_name": deployment_name,
+            "temperature": temperature,
+        }
+        
+        # Add any extra parameters provided
+        if extra_params:
+            params.update(extra_params)
+            
+        agent = agent_class(**params)
         return await agent.setup()
 
     @classmethod