most changes

Author: marktayl1

src/backend/common/config/config.py

@@ -1,12 +1,26 @@
+"""Configuration class for the application.
+This class loads configuration values from environment variables and provides
+methods to access them. It also initializes an Azure AI client using the
+provided credentials.
+It uses the `azure.identity` library to handle authentication and
+authorization with Azure services.
+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 os
 
 from azure.identity.aio import ClientSecretCredential, DefaultAzureCredential
-from dotenv import load_dotenv
-
-load_dotenv()
+from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 
 
 class Config:
+    """Configuration class for the application."""
+
     def __init__(self):
         self.azure_tenant_id = os.getenv("AZURE_TENANT_ID", "")
         self.azure_client_id = os.getenv("AZURE_CLIENT_ID", "")
@@ -37,6 +51,10 @@ 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/requirements.txt

@@ -34,7 +34,7 @@ structlog
 typing-extensions
 python-jose[cryptography]
 passlib[bcrypt]
-semantic-kernel==1.23.1
+semantic-kernel[azure]==1.26.1
 openai
 sqlparse
 sqlglot

src/backend/sql_agents/__init__.py

@@ -13,13 +13,13 @@
 from .agent_config import AgentsConfigDialect, create_config
 
 __all__ = [
+    "create_kernel_with_chat_completion",
     "setup_migrator_agent",
     "setup_fixer_agent",
     "setup_picker_agent",
     "setup_syntax_checker_agent",
     "setup_semantic_verifier_agent",
     "get_prompt",
-    "create_kernel_with_chat_completion",
     "create_config",
     "AgentType",
 ]

src/backend/sql_agents/fixer/agent.py

@@ -2,49 +2,57 @@
 
 import logging
 
+from azure.ai.projects.models import (
+    ResponseFormatJsonSchema,
+    ResponseFormatJsonSchemaType,
+)
+from common.config.config import app_config
 from common.models.api import AgentType
-from sql_agents.helpers.sk_utils import create_kernel_with_chat_completion
-from sql_agents.helpers.utils import get_prompt
-from semantic_kernel.agents import ChatCompletionAgent
+from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 from semantic_kernel.kernel import KernelArguments
-from semantic_kernel.prompt_template import PromptTemplateConfig
 from sql_agents.agent_config import AgentModelDeployment, AgentsConfigDialect
 from sql_agents.fixer.response import FixerResponse
+from sql_agents.helpers.utils import get_prompt
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.DEBUG)
-
-
-def setup_fixer_agent(
+async def setup_fixer_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
-) -> ChatCompletionAgent:
+) -> AzureAIAgent:
     """Setup the fixer agent."""
     _deployment_name = deployment_name.value
     _name = name.value
-    kernel = create_kernel_with_chat_completion(_name, _deployment_name)
 
     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
 
-    # prompt = replace_tags(template_content, {"target": config.sql_dialect_out})
-
-    settings = kernel.get_prompt_execution_settings_from_service_id(service_id=_name)
-    settings.response_format = FixerResponse
-    settings.temperature = 0.0
-
-    kernel_args = KernelArguments(target=config.sql_dialect_out, settings=settings)
+    kernel_args = KernelArguments(target=config.sql_dialect_out)
 
-    fixer_agent = ChatCompletionAgent(
-        kernel=kernel,
+    # 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,
     )
 

src/backend/sql_agents/fixer/response.py

@@ -1,7 +1,7 @@
-from pydantic import BaseModel
+from semantic_kernel.kernel_pydantic import KernelBaseModel
 
 
-class FixerResponse(BaseModel):
+class FixerResponse(KernelBaseModel):
     """
     Model for the response of the fixer
     """

src/backend/sql_agents/migrator/agent.py

@@ -2,51 +2,62 @@
 
 import logging
 
+from azure.ai.projects.models import (
+    ResponseFormatJsonSchema,
+    ResponseFormatJsonSchemaType,
+)
+from common.config.config import app_config
 from common.models.api import AgentType
-from sql_agents.helpers.sk_utils import create_kernel_with_chat_completion
-from sql_agents.helpers.utils import get_prompt
-from semantic_kernel.agents import ChatCompletionAgent
+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.helpers.utils import get_prompt
 from sql_agents.migrator.response import MigratorResponse
 
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
 
-def setup_migrator_agent(
+async def setup_migrator_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
-) -> ChatCompletionAgent:
+) -> AzureAIAgent:
     """Setup the migrator agent."""
     _deployment_name = deployment_name.value
     _name = name.value
-    NUM_CANDIDATES = 3
-
-    kernel = create_kernel_with_chat_completion(_name, _deployment_name)
+    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
 
-    settings = kernel.get_prompt_execution_settings_from_service_id(
-        service_id="migrator"
-    )
-    settings.response_format = MigratorResponse
-    settings.temperature = 0.0
-
     kernel_args = KernelArguments(
         target=config.sql_dialect_out,
-        numCandidates=str(NUM_CANDIDATES),
+        numCandidates=str(num_candidates),
         source=config.sql_dialect_in,
-        settings=settings,
     )
 
-    migrator_agent = ChatCompletionAgent(
-        kernel=kernel,
-        name=name,
+    # 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,
     )
 

src/backend/sql_agents/migrator/response.py

@@ -1,7 +1,7 @@
-from pydantic import BaseModel
+from semantic_kernel.kernel_pydantic import KernelBaseModel
 
 
-class MigratorCandidate(BaseModel):
+class MigratorCandidate(KernelBaseModel):
     """
     Model for a single candidate for migration
     """
@@ -10,7 +10,7 @@ class MigratorCandidate(BaseModel):
     candidate_query: str
 
 
-class MigratorResponse(BaseModel):
+class MigratorResponse(KernelBaseModel):
     """
     Model for the response of the migrator
     """
@@ -19,4 +19,4 @@ class MigratorResponse(BaseModel):
     candidates: list[MigratorCandidate]
     input_error: str | None = None
     summary: str | None = None
-    rai_error: str | None = None
+    rai_error: str | None = None

src/backend/sql_agents/picker/agent.py

@@ -2,12 +2,16 @@
 
 import logging
 
+from azure.ai.projects.models import (
+    ResponseFormatJsonSchema,
+    ResponseFormatJsonSchemaType,
+)
+from common.config.config import app_config
 from common.models.api import AgentType
-from sql_agents.helpers.sk_utils import create_kernel_with_chat_completion
-from sql_agents.helpers.utils import get_prompt
-from semantic_kernel.agents import ChatCompletionAgent
+from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
 from semantic_kernel.kernel import KernelArguments
 from sql_agents.agent_config import AgentModelDeployment, AgentsConfigDialect
+from sql_agents.helpers.utils import get_prompt
 from sql_agents.picker.response import PickerResponse
 
 logger = logging.getLogger(__name__)
@@ -16,35 +20,45 @@
 NUM_CANDIDATES = 3
 
 
-def setup_picker_agent(
+async def setup_picker_agent(
     name: AgentType, config: AgentsConfigDialect, deployment_name: AgentModelDeployment
-) -> ChatCompletionAgent:
+) -> AzureAIAgent:
     """Setup the picker agent."""
     _deployment_name = deployment_name.value
     _name = name.value
-    kernel = create_kernel_with_chat_completion(_name, _deployment_name)
 
     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
 
-    settings = kernel.get_prompt_execution_settings_from_service_id(service_id="picker")
-    settings.response_format = PickerResponse
-    settings.temperature = 0.0
-
     kernel_args = KernelArguments(
         target=config.sql_dialect_out,
         numCandidates=str(NUM_CANDIDATES),
         source=config.sql_dialect_in,
-        settings=settings,
     )
 
-    picker_agent = ChatCompletionAgent(
-        kernel=kernel,
+    # 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,
     )
 

src/backend/sql_agents/picker/response.py

@@ -1,12 +1,12 @@
-from pydantic import BaseModel
+from semantic_kernel.kernel_pydantic import KernelBaseModel
 
 
-class PickerCandidateSummary(BaseModel):
+class PickerCandidateSummary(KernelBaseModel):
     candidate_index: int
     candidate_summary: str
 
 
-class PickerResponse(BaseModel):
+class PickerResponse(KernelBaseModel):
     """
     The response of the picker agent.
     """