3434 ChatMessageContent ,
3535)
3636from semantic_kernel .exceptions .service_exceptions import ServiceResponseException
37- from sql_agents import (
38- create_kernel_with_chat_completion ,
39- setup_fixer_agent ,
40- setup_migrator_agent ,
41- setup_picker_agent ,
42- setup_semantic_verifier_agent ,
43- setup_syntax_checker_agent ,
44- )
4537from sql_agents .agent_config import AgentBaseConfig
4638from sql_agents .fixer .response import FixerResponse
39+ from sql_agents .fixer .setup import setup_fixer_agent
4740from sql_agents .helpers .selection_function import setup_selection_function
41+ from sql_agents .helpers .sk_utils import create_kernel_with_chat_completion
4842from sql_agents .helpers .termination_function import setup_termination_function
4943from sql_agents .helpers .utils import is_text
5044from sql_agents .migrator .response import MigratorResponse
45+ from sql_agents .migrator .setup import setup_migrator_agent
5146from sql_agents .picker .response import PickerResponse
47+ from sql_agents .picker .setup import setup_picker_agent
5248from sql_agents .semantic_verifier .response import SemanticVerifierResponse
49+ from sql_agents .semantic_verifier .setup import setup_semantic_verifier_agent
50+ from sql_agents .syntax_checker .setup import setup_syntax_checker_agent
5351
5452# Loop through files from Cosmos DB.
5553
8482# return sub_str.split(":")[1].strip().strip('"')
8583
8684
87- async def configure_agents ():
85+ async def configure_agents (config : AgentBaseConfig ):
8886 try :
89- agent_fixer = await setup_fixer_agent (
90- AgentType .FIXER ,
91- agent_dialect_config ,
92- AgentModelDeployment .FIXER_AGENT_MODEL_DEPLOY ,
93- )
94- agent_migrator = await setup_migrator_agent (
95- AgentType .MIGRATOR ,
96- agent_dialect_config ,
97- AgentModelDeployment .MIGRATOR_AGENT_MODEL_DEPLOY ,
98- )
99- agent_picker = await setup_picker_agent (
100- AgentType .PICKER ,
101- agent_dialect_config ,
102- AgentModelDeployment .PICKER_AGENT_MODEL_DEPLOY ,
103- )
104- agent_syntax_checker = await setup_syntax_checker_agent (
105- AgentType .SYNTAX_CHECKER ,
106- agent_dialect_config ,
107- AgentModelDeployment .SYNTAX_CHECKER_AGENT_MODEL_DEPLOY ,
108- )
87+ agent_fixer = await setup_fixer_agent (config )
88+ agent_migrator = await setup_migrator_agent (config )
89+ agent_picker = await setup_picker_agent (config )
90+ agent_syntax_checker = await setup_syntax_checker_agent (config )
10991 selection_function = setup_selection_function (
11092 SELECTION_FUNCTION_NAME ,
11193 AgentType .MIGRATOR ,
@@ -133,7 +115,10 @@ async def configure_agents():
133115
134116
135117async def convert (
136- source_script , file : FileRecord , batch_service : BatchService , agent_config
118+ source_script ,
119+ file : FileRecord ,
120+ batch_service : BatchService ,
121+ agent_config : AgentBaseConfig ,
137122) -> str :
138123 """setup agents, selection and termination."""
139124 logger .info ("Migrating query: %s\n " , source_script )
@@ -149,7 +134,7 @@ async def convert(
149134 function = agent_config ["selection_function" ],
150135 kernel = create_kernel_with_chat_completion (
151136 AgentType .SELECTION .value ,
152- AgentModelDeployment . SELECTION_MODEL_DEPLOY . value ,
137+ agent_config . model_type [ AgentType . SELECTION ] ,
153138 ),
154139 result_parser = lambda result : (
155140 str (result .value [0 ]) if result .value is not None else AgentType .MIGRATOR
@@ -163,7 +148,7 @@ async def convert(
163148 function = agent_config ["termination_function" ],
164149 kernel = create_kernel_with_chat_completion (
165150 AgentType .TERMINATION .value ,
166- AgentModelDeployment . TERMINATION_MODEL_DEPLOY . value ,
151+ agent_config . model_type [ AgentType . TERMINATION ] ,
167152 ),
168153 result_parser = lambda result : TERMINATION_KEYWORD
169154 in str (result .value [0 ]).lower (),
@@ -297,7 +282,7 @@ async def convert(
297282
298283 # Invoke the semantic verifier agent to validate the migrated query
299284 semver_response = await invoke_semantic_verifier (
300- source_script , migrated_query , file , batch_service
285+ agent_config , source_script , migrated_query
301286 )
302287 semver_response = SemanticVerifierResponse .model_validate_json (
303288 semver_response or ""
@@ -394,7 +379,9 @@ async def convert(
394379
395380
396381async def invoke_semantic_verifier (
397- source_script , migrated_query , file : FileRecord , batch_service : BatchService
382+ config : AgentBaseConfig ,
383+ source_script : str ,
384+ migrated_query : str ,
398385):
399386 """Invoke the semantic verifier agent to validate the migrated query."""
400387 try :
@@ -410,9 +397,7 @@ async def invoke_semantic_verifier(
410397 )
411398
412399 agent_semantic_verifier = await setup_semantic_verifier_agent (
413- AgentType .SEMANTIC_VERIFIER ,
414- agent_dialect_config ,
415- AgentModelDeployment .SEMANTIC_VERIFIER_AGENT_MODEL_DEPLOY ,
400+ config ,
416401 source_script ,
417402 migrated_query ,
418403 )
0 commit comments