Split agent mods and setup

Author: marktayl1

src/backend/common/config/config.py

@@ -15,7 +15,6 @@
 import os
 
 from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
-from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 
 
 class Config:
@@ -51,10 +50,6 @@ def __init__(self):
 
         self.__azure_credentials = DefaultAzureCredential()
 
-        self.ai_project_client = AzureAIAgent.create_client(
-            credential=self.get_azure_credentials()
-        )
-
     def get_azure_credentials(self):
         """Retrieve Azure credentials, either from environment variables or managed identity."""
         if all([self.azure_tenant_id, self.azure_client_id, self.azure_client_secret]):

src/backend/sql_agents/__init__.py

@@ -1,25 +1,30 @@
-"""This module initializes the agents and helpers for the"""
+# """This module initializes the agents and helpers for the"""
 
-from common.models.api import AgentType
-from sql_agents.fixer.agent import setup_fixer_agent
-from sql_agents.helpers.sk_utils import create_kernel_with_chat_completion
-from sql_agents.helpers.utils import get_prompt
-from sql_agents.migrator.agent import setup_migrator_agent
-from sql_agents.picker.agent import setup_picker_agent
-from sql_agents.semantic_verifier.agent import setup_semantic_verifier_agent
-from sql_agents.syntax_checker.agent import setup_syntax_checker_agent
+# from common.models.api import AgentType
+from sql_agents.fixer.agent import FixerAgent, setup_fixer_agent
+from sql_agents.migrator.agent import MigratorAgent, setup_migrator_agent
+from sql_agents.picker.agent import PickerAgent, setup_picker_agent
+from sql_agents.semantic_verifier.agent import (
+    SemanticVerifierAgent,
+    setup_semantic_verifier_agent,
+)
+from sql_agents.syntax_checker.agent import (
+    SyntaxCheckerAgent,
+    setup_syntax_checker_agent,
+)
 
-# Import the configuration function
-from .agent_config import AgentsConfigDialect, create_config
+# from sql_agents.agent_config import AgentBaseConfig
+# from sql_agents.agent_factory import SQLAgentFactory
 
 __all__ = [
-    "create_kernel_with_chat_completion",
     "setup_migrator_agent",
+    "MigratorAgent",
     "setup_fixer_agent",
+    "FixerAgent",
     "setup_picker_agent",
+    "PickerAgent",
     "setup_syntax_checker_agent",
+    "SyntaxCheckerAgent",
     "setup_semantic_verifier_agent",
-    "get_prompt",
-    "create_config",
-    "AgentType",
+    "SemanticVerifierAgent",
 ]

src/backend/sql_agents/agent_base.py

@@ -8,15 +8,14 @@
     ResponseFormatJsonSchema,
     ResponseFormatJsonSchemaType,
 )
-from common.config.config import app_config
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 from semantic_kernel.functions import KernelArguments
-from sql_agents.agent_config import AgentModelDeployment, AgentsConfigDialect
+from sql_agents.agent_config import AgentBaseConfig
 from sql_agents.helpers.utils import get_prompt
 
 # Type variable for response models
-T = TypeVar('T')
+T = TypeVar("T")
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -28,12 +27,12 @@ class BaseSQLAgent(Generic[T], ABC):
     def __init__(
         self,
         agent_type: AgentType,
-        config: AgentsConfigDialect,
-        deployment_name: AgentModelDeployment,
+        config: AgentBaseConfig,
+        deployment_name: None,
         temperature: float = 0.0,
     ):
         """Initialize the base SQL agent.
-        
+
         Args:
             agent_type: The type of agent to create.
             config: The dialect configuration for the agent.
@@ -55,7 +54,7 @@ def response_schema(self) -> type:
     @property
     def num_candidates(self) -> Optional[int]:
         """Get the number of candidates for this agent.
-        
+
         Returns:
             The number of candidates, or None if not applicable.
         """
@@ -64,32 +63,32 @@ def num_candidates(self) -> Optional[int]:
     @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,
+            "target": self.config.sql_to,
+            "source": self.config.sql_from,
         }
-        
+
         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."""
-        _deployment_name = self.deployment_name.value
         _name = self.agent_type.value
+        _deployment_name = self.config.model_type.get(self.agent_type)
 
         try:
             template_content = get_prompt(_name)
@@ -100,7 +99,7 @@ async def setup(self) -> AzureAIAgent:
         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(
+        agent_definition = await self.config.ai_project_client.agents.create_agent(
             model=_deployment_name,
             name=_name,
             instructions=template_content,
@@ -116,15 +115,15 @@ async def setup(self) -> AzureAIAgent:
 
         # Create a Semantic Kernel agent based on the agent definition
         agent_kwargs = {
-            "client": app_config.ai_project_client,
+            "client": self.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
@@ -139,4 +138,4 @@ async def execute(self, inputs: Any) -> T:
         """Execute the agent with the given inputs."""
         agent = await self.get_agent()
         response = await agent.invoke(inputs)
-        return response  # Type will be inferred from T
+        return response  # Type will be inferred from T

src/backend/sql_agents/agent_config.py

@@ -1,36 +1,37 @@
-"""Configuration for the agents module."""
+"""Configuration class for the agents.
+This class loads configuration values from environment variables and provides
+properties to access them. It also stores an Azure AI client and SQL dialect
+configuration for the agents, that will be set per batch.
+Access to .env variables requires adding the `python-dotenv` package to, or
+configuration of the env python path through the IDE. For example, in VSCode, the
+settings.json file in the .vscode folder should include the following:
+{
+    "python.envFile": "${workspaceFolder}/.env"
+}
+"""
 
-import json
 import os
 from enum import Enum
 
-from dotenv import load_dotenv
+from azure.ai.projects.aio import AIProjectClient
+from common.models.api import AgentType
 
-load_dotenv()
 
-
-class AgentModelDeployment(Enum):
+class AgentBaseConfig:
     """Agent model deployment names."""
 
-    MIGRATOR_AGENT_MODEL_DEPLOY = os.getenv("MIGRATOR_AGENT_MODEL_DEPLOY")
-    PICKER_AGENT_MODEL_DEPLOY = os.getenv("PICKER_AGENT_MODEL_DEPLOY")
-    FIXER_AGENT_MODEL_DEPLOY = os.getenv("FIXER_AGENT_MODEL_DEPLOY")
-    SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY = os.getenv(
-        "SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY"
-    )
-    SYNTAX_CHECKER_AGENT_MODEL_DEPLOY = os.getenv("SYNTAX_CHECKER_AGENT_MODEL_DEPLOY")
-    SELECTION_MODEL_DEPLOY = os.getenv("SELECTION_MODEL_DEPLOY")
-    TERMINATION_MODEL_DEPLOY = os.getenv("TERMINATION_MODEL_DEPLOY")
-
-
-class AgentsConfigDialect:
-    """Configuration for the agents module."""
-
-    def __init__(self, sql_dialect_in, sql_dialect_out):
-        self.sql_dialect_in = sql_dialect_in
-        self.sql_dialect_out = sql_dialect_out
-
-
-def create_config(sql_dialect_in, sql_dialect_out):
-    """Create and return a new AgentConfig object."""
-    return AgentsConfigDialect(sql_dialect_in, sql_dialect_out)
+    def __init__(self, project_client: AIProjectClient, sql_from: str, sql_to: str):
+
+        self.ai_project_client: AIProjectClient = project_client
+        self.sql_from = sql_from
+        self.sql_to = sql_to
+
+    model_type = {
+        AgentType.MIGRATOR: os.getenv("MIGRATOR_AGENT_MODEL_DEPLOY"),
+        AgentType.PICKER: os.getenv("PICKER_AGENT_MODEL_DEPLOY"),
+        AgentType.FIXER: os.getenv("FIXER_AGENT_MODEL_DEPLOY"),
+        AgentType.SEMANTIC_VERIFIER: os.getenv("SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY"),
+        AgentType.SYNTAX_CHECKER: os.getenv("SYNTAX_CHECKER_AGENT_MODEL_DEPLOY"),
+        AgentType.SELECTION: os.getenv("SELECTION_MODEL_DEPLOY"),
+        AgentType.TERMINATION: os.getenv("TERMINATION_MODEL_DEPLOY"),
+    }

src/backend/sql_agents/agent_factory.py

@@ -1,24 +1,25 @@
 """Factory for creating SQL migration agents."""
 
 import logging
-from typing import Type, TypeVar, Optional, Dict, Any
+from typing import Any, Dict, Optional, Type, TypeVar
 
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
+from sql_agents import (
+    FixerAgent,
+    MigratorAgent,
+    PickerAgent,
+    SemanticVerifierAgent,
+    SyntaxCheckerAgent,
+)
 from sql_agents.agent_base import BaseSQLAgent
-from sql_agents.agent_config import AgentModelDeployment, AgentsConfigDialect
-from sql_agents.migrator.agent import MigratorAgent
-from sql_agents.picker.agent import PickerAgent 
-from sql_agents.syntax_checker.agent import SyntaxCheckerAgent
-from sql_agents.fixer.agent import FixerAgent   
-from sql_agents.semantic_verifier.agent import SemanticVerifierAgent
-from sql_agents.helpers.utils import get_prompt
+from sql_agents.agent_config import AgentBaseConfig
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 # Type variable for agent response types
-T = TypeVar('T')
+T = TypeVar("T")
 
 
 class SQLAgentFactory:
@@ -36,20 +37,19 @@ class SQLAgentFactory:
     async def create_agent(
         cls,
         agent_type: AgentType,
-        config: AgentsConfigDialect,
-        deployment_name: AgentModelDeployment,
+        config: AgentBaseConfig,
         temperature: float = 0.0,
-        **kwargs
+        **kwargs,
     ) -> AzureAIAgent:
         """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.
             **kwargs: Additional parameters to pass to the agent constructor.
-            
+
         Returns:
             A configured AzureAIAgent instance.
         """
@@ -61,24 +61,29 @@ async def create_agent(
         params = {
             "agent_type": agent_type,
             "config": config,
-            "deployment_name": deployment_name,
             "temperature": temperature,
-            **kwargs
+            **kwargs,
         }
-            
+
         agent = agent_class(**params)
         return await agent.setup()
-    
+
     @classmethod
     def get_agent_class(cls, agent_type: AgentType) -> Type[BaseSQLAgent]:
         """Get the agent class for the specified type."""
         agent_class = cls._agent_classes.get(agent_type)
         if not agent_class:
             raise ValueError(f"Unknown agent type: {agent_type}")
         return agent_class
-    
+
     @classmethod
-    def register_agent_class(cls, agent_type: AgentType, agent_class: Type[BaseSQLAgent]) -> None:
+    def register_agent_class(
+        cls, agent_type: AgentType, agent_class: Type[BaseSQLAgent]
+    ) -> None:
         """Register a new agent class with the factory."""
         cls._agent_classes[agent_type] = agent_class
-        logger.info("Registered agent class %s for type %s", agent_class.__name__, agent_type.value)
+        logger.info(
+            "Registered agent class %s for type %s",
+            agent_class.__name__,
+            agent_type.value,
+        )