add agent class and agent factory update

Author: Fr4nc3

src/backend/sql_agents/agent_factory.py

@@ -7,7 +7,12 @@
 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, SyntaxCheckerAgent
+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
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
@@ -23,6 +28,8 @@ class SQLAgentFactory:
         AgentType.PICKER: PickerAgent,
         AgentType.MIGRATOR: MigratorAgent,
         AgentType.SYNTAX_CHECKER: SyntaxCheckerAgent,
+        AgentType.FIXER: FixerAgent,
+        AgentType.SEMANTIC_VERIFIER: SemanticVerifierAgent,
     }
 
     @classmethod
@@ -32,7 +39,7 @@ async def create_agent(
         config: AgentsConfigDialect,
         deployment_name: AgentModelDeployment,
         temperature: float = 0.0,
-        extra_params: Optional[Dict[str, Any]] = None,
+        **kwargs
     ) -> AzureAIAgent:
         """Create and setup an agent of the specified type.
         
@@ -41,7 +48,7 @@ async def create_agent(
             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.
+            **kwargs: Additional parameters to pass to the agent constructor.
             
         Returns:
             A configured AzureAIAgent instance.
@@ -56,11 +63,8 @@ async def create_agent(
             "config": config,
             "deployment_name": deployment_name,
             "temperature": temperature,
+            **kwargs
         }
-        
-        # Add any extra parameters provided
-        if extra_params:
-            params.update(extra_params)
 
         agent = agent_class(**params)
         return await agent.setup()

src/backend/sql_agents/fixer/agent.py

@@ -6,6 +6,10 @@
     ResponseFormatJsonSchema,
     ResponseFormatJsonSchemaType,
 )
+from backend.sql_agents.agent_base import (
+    BaseSQLAgent,
+)  # Ensure this import is correct and the module exists
+from backend.sql_agents.agent_factory import SQLAgentFactory
 from common.config.config import app_config
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
@@ -18,42 +22,58 @@
 logger.setLevel(logging.DEBUG)
 
 
+class FixerAgent(BaseSQLAgent[FixerResponse]):
+    """Fixer agent for correcting SQL syntax errors."""
+
+    @property
+    def response_schema(self) -> type:
+        """Get the response schema for the fixer agent."""
+        return FixerResponse
+
+
+# async def setup_fixer_agent(
+#     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
+# ) -> AzureAIAgent:
+#     """Setup the fixer agent."""
+#     _deployment_name = deployment_name.value
+#     _name = name.value
+
+#     try:
+#         template_content = get_prompt(_name)
+#     except FileNotFoundError as exc:
+#         logger.error("Prompt file for %s not found.", _name)
+#         raise ValueError(f"Prompt file for {_name} not found.") from exc
+
+#     kernel_args = KernelArguments(target=config.sql_dialect_out)
+
+#     # Define an agent on the Azure AI agent service
+#     agent_definition = await app_config.ai_project_client.agents.create_agent(
+#         model=_deployment_name,
+#         name=_name,
+#         instructions=template_content,
+#         temperature=0.0,
+#         response_format=ResponseFormatJsonSchemaType(
+#             json_schema=ResponseFormatJsonSchema(
+#                 name="FixerResponse",
+#                 description="respond with fixer response",
+#                 schema=FixerResponse.model_json_schema(),
+#             )
+#         ),
+#     )
+
+#     # Create a Semantic Kernel agent based on the agent definition.
+#     # Add RAG with docs programmatically for this one
+#     fixer_agent = AzureAIAgent(
+#         client=app_config.ai_project_client,
+#         definition=agent_definition,
+#         arguments=kernel_args,
+#     )
+
+#     return fixer_agent
+
+
 async def setup_fixer_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
 ) -> AzureAIAgent:
-    """Setup the fixer agent."""
-    _deployment_name = deployment_name.value
-    _name = name.value
-
-    try:
-        template_content = get_prompt(_name)
-    except FileNotFoundError as exc:
-        logger.error("Prompt file for %s not found.", _name)
-        raise ValueError(f"Prompt file for {_name} not found.") from exc
-
-    kernel_args = KernelArguments(target=config.sql_dialect_out)
-
-    # Define an agent on the Azure AI agent service
-    agent_definition = await app_config.ai_project_client.agents.create_agent(
-        model=_deployment_name,
-        name=_name,
-        instructions=template_content,
-        temperature=0.0,
-        response_format=ResponseFormatJsonSchemaType(
-            json_schema=ResponseFormatJsonSchema(
-                name="FixerResponse",
-                description="respond with fixer response",
-                schema=FixerResponse.model_json_schema(),
-            )
-        ),
-    )
-
-    # Create a Semantic Kernel agent based on the agent definition.
-    # Add RAG with docs programmatically for this one
-    fixer_agent = AzureAIAgent(
-        client=app_config.ai_project_client,
-        definition=agent_definition,
-        arguments=kernel_args,
-    )
-
-    return fixer_agent
+    """Setup the fixer agent using the factory."""
+    return await SQLAgentFactory.create_agent(name, config, deployment_name)

src/backend/sql_agents/migrator/agent.py

@@ -6,6 +6,8 @@
     ResponseFormatJsonSchema,
     ResponseFormatJsonSchemaType,
 )
+from backend.sql_agents.agent_base import BaseSQLAgent
+from backend.sql_agents.agent_factory import SQLAgentFactory
 from common.config.config import app_config
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
@@ -18,47 +20,68 @@
 logger.setLevel(logging.DEBUG)
 
 
+class MigratorAgent(BaseSQLAgent[MigratorResponse]):
+    """Migrator agent for translating SQL from one dialect to another."""
+
+    @property
+    def response_schema(self) -> type:
+        """Get the response schema for the migrator agent."""
+        return MigratorResponse
+
+    @property
+    def num_candidates(self) -> int:
+        """Get the number of candidates for the migrator agent."""
+        return 3
+
+
 async def setup_migrator_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
 ) -> AzureAIAgent:
-    """Setup the migrator agent."""
-    _deployment_name = deployment_name.value
-    _name = name.value
-    num_candidates = 3
-
-    try:
-        template_content = get_prompt(_name)
-    except FileNotFoundError as exc:
-        logger.error("Prompt file for %s not found.", _name)
-        raise ValueError(f"Prompt file for {_name} not found.") from exc
-
-    kernel_args = KernelArguments(
-        target=config.sql_dialect_out,
-        numCandidates=str(num_candidates),
-        source=config.sql_dialect_in,
-    )
-
-    # Define an agent on the Azure AI agent service
-    agent_definition = await app_config.ai_project_client.agents.create_agent(
-        model=_deployment_name,
-        name=_name,
-        instructions=template_content,
-        temperature=0.0,
-        response_format=ResponseFormatJsonSchemaType(
-            json_schema=ResponseFormatJsonSchema(
-                name="MigratorResponse",
-                description="respond with migrator response",
-                schema=MigratorResponse.model_json_schema(),
-            )
-        ),
-    )
-
-    # Create a Semantic Kernel agent based on the agent definition.
-    # Add RAG with docs programmatically for this one
-    migrator_agent = AzureAIAgent(
-        client=app_config.ai_project_client,
-        definition=agent_definition,
-        arguments=kernel_args,
-    )
-
-    return migrator_agent
+    """Setup the migrator agent using the factory."""
+    return await SQLAgentFactory.create_agent(name, config, deployment_name)
+
+
+# async def setup_migrator_agent(
+#     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
+# ) -> AzureAIAgent:
+#     """Setup the migrator agent."""
+#     _deployment_name = deployment_name.value
+#     _name = name.value
+#     num_candidates = 3
+
+#     try:
+#         template_content = get_prompt(_name)
+#     except FileNotFoundError as exc:
+#         logger.error("Prompt file for %s not found.", _name)
+#         raise ValueError(f"Prompt file for {_name} not found.") from exc
+
+#     kernel_args = KernelArguments(
+#         target=config.sql_dialect_out,
+#         numCandidates=str(num_candidates),
+#         source=config.sql_dialect_in,
+#     )
+
+#     # Define an agent on the Azure AI agent service
+#     agent_definition = await app_config.ai_project_client.agents.create_agent(
+#         model=_deployment_name,
+#         name=_name,
+#         instructions=template_content,
+#         temperature=0.0,
+#         response_format=ResponseFormatJsonSchemaType(
+#             json_schema=ResponseFormatJsonSchema(
+#                 name="MigratorResponse",
+#                 description="respond with migrator response",
+#                 schema=MigratorResponse.model_json_schema(),
+#             )
+#         ),
+#     )
+
+#     # Create a Semantic Kernel agent based on the agent definition.
+#     # Add RAG with docs programmatically for this one
+#     migrator_agent = AzureAIAgent(
+#         client=app_config.ai_project_client,
+#         definition=agent_definition,
+#         arguments=kernel_args,
+#     )
+
+#     return migrator_agent

src/backend/sql_agents/picker/agent.py

@@ -6,6 +6,8 @@
     ResponseFormatJsonSchema,
     ResponseFormatJsonSchemaType,
 )
+from backend.sql_agents.agent_base import BaseSQLAgent
+from backend.sql_agents.agent_factory import SQLAgentFactory
 from common.config.config import app_config
 from common.models.api import AgentType
 from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
@@ -17,49 +19,70 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
-NUM_CANDIDATES = 3
+# NUM_CANDIDATES = 3
+
+
+class PickerAgent(BaseSQLAgent[PickerResponse]):
+    """Picker agent for selecting the best SQL translation candidate."""
+
+    @property
+    def response_schema(self) -> type:
+        """Get the response schema for the picker agent."""
+        return PickerResponse
+
+    @property
+    def num_candidates(self) -> int:
+        """Get the number of candidates for the picker agent."""
+        return 3
 
 
 async def setup_picker_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
 ) -> AzureAIAgent:
-    """Setup the picker agent."""
-    _deployment_name = deployment_name.value
-    _name = name.value
-
-    try:
-        template_content = get_prompt(_name)
-    except FileNotFoundError as exc:
-        logger.error("Prompt file for %s not found.", _name)
-        raise ValueError(f"Prompt file for {_name} not found.") from exc
-
-    kernel_args = KernelArguments(
-        target=config.sql_dialect_out,
-        numCandidates=str(NUM_CANDIDATES),
-        source=config.sql_dialect_in,
-    )
-
-    # Define an agent on the Azure AI agent service
-    agent_definition = await app_config.ai_project_client.agents.create_agent(
-        model=_deployment_name,
-        name=_name,
-        instructions=template_content,
-        temperature=0.0,
-        response_format=ResponseFormatJsonSchemaType(
-            json_schema=ResponseFormatJsonSchema(
-                name="PickerResponse",
-                description="respond with picker response",
-                schema=PickerResponse.model_json_schema(),
-            )
-        ),
-    )
-
-    # Create a Semantic Kernel agent based on the agent definition.
-    # Add RAG with docs programmatically for this one
-    picker_agent = AzureAIAgent(
-        client=app_config.ai_project_client,
-        definition=agent_definition,
-        arguments=kernel_args,
-    )
-
-    return picker_agent
+    """Setup the picker agent using the factory."""
+    return await SQLAgentFactory.create_agent(name, config, deployment_name)
+
+
+# async def setup_picker_agent(
+#     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
+# ) -> AzureAIAgent:
+#     """Setup the picker agent."""
+#     _deployment_name = deployment_name.value
+#     _name = name.value
+
+#     try:
+#         template_content = get_prompt(_name)
+#     except FileNotFoundError as exc:
+#         logger.error("Prompt file for %s not found.", _name)
+#         raise ValueError(f"Prompt file for {_name} not found.") from exc
+
+#     kernel_args = KernelArguments(
+#         target=config.sql_dialect_out,
+#         numCandidates=str(NUM_CANDIDATES),
+#         source=config.sql_dialect_in,
+#     )
+
+#     # Define an agent on the Azure AI agent service
+#     agent_definition = await app_config.ai_project_client.agents.create_agent(
+#         model=_deployment_name,
+#         name=_name,
+#         instructions=template_content,
+#         temperature=0.0,
+#         response_format=ResponseFormatJsonSchemaType(
+#             json_schema=ResponseFormatJsonSchema(
+#                 name="PickerResponse",
+#                 description="respond with picker response",
+#                 schema=PickerResponse.model_json_schema(),
+#             )
+#         ),
+#     )
+
+#     # Create a Semantic Kernel agent based on the agent definition.
+#     # Add RAG with docs programmatically for this one
+#     picker_agent = AzureAIAgent(
+#         client=app_config.ai_project_client,
+#         definition=agent_definition,
+#         arguments=kernel_args,
+#     )
+
+#     return picker_agent