Skip to content

Commit 9b6ec87

Browse files
committed
fixes and init run
1 parent 2dd1b90 commit 9b6ec87

File tree

9 files changed

+59
-83
lines changed

9 files changed

+59
-83
lines changed

src/backend/sql_agents/__init__.py

Lines changed: 26 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,30 +1,29 @@
1-
# """This module initializes the agents and helpers for the"""
1+
# # """This module initializes the agents and helpers for the"""
22

3-
# from common.models.api import AgentType
4-
from sql_agents.fixer.agent import FixerAgent, setup_fixer_agent
5-
from sql_agents.migrator.agent import MigratorAgent, setup_migrator_agent
6-
from sql_agents.picker.agent import PickerAgent, setup_picker_agent
7-
from sql_agents.semantic_verifier.agent import (
8-
SemanticVerifierAgent,
9-
setup_semantic_verifier_agent,
10-
)
11-
from sql_agents.syntax_checker.agent import (
12-
SyntaxCheckerAgent,
13-
setup_syntax_checker_agent,
14-
)
3+
# # from common.models.api import AgentType
4+
# from sql_agents.fixer.agent import FixerAgent
5+
# from sql_agents.fixer.setup import setup_fixer_agent
6+
# from sql_agents.migrator.agent import MigratorAgent
7+
# from sql_agents.migrator.setup import setup_migrator_agent
8+
# from sql_agents.picker.agent import PickerAgent
9+
# from sql_agents.picker.setup import setup_picker_agent
10+
# from sql_agents.semantic_verifier.agent import SemanticVerifierAgent
11+
# from sql_agents.semantic_verifier.setup import setup_semantic_verifier_agent
12+
# from sql_agents.syntax_checker.agent import SyntaxCheckerAgent
13+
# from sql_agents.syntax_checker.setup import setup_syntax_checker_agent
1514

16-
# from sql_agents.agent_config import AgentBaseConfig
17-
# from sql_agents.agent_factory import SQLAgentFactory
15+
# # from sql_agents.agent_config import AgentBaseConfig
16+
# # from sql_agents.agent_factory import SQLAgentFactory
1817

19-
__all__ = [
20-
"setup_migrator_agent",
21-
"MigratorAgent",
22-
"setup_fixer_agent",
23-
"FixerAgent",
24-
"setup_picker_agent",
25-
"PickerAgent",
26-
"setup_syntax_checker_agent",
27-
"SyntaxCheckerAgent",
28-
"setup_semantic_verifier_agent",
29-
"SemanticVerifierAgent",
30-
]
18+
# __all__ = [
19+
# "setup_migrator_agent",
20+
# "MigratorAgent",
21+
# "setup_fixer_agent",
22+
# "FixerAgent",
23+
# "setup_picker_agent",
24+
# "PickerAgent",
25+
# "setup_syntax_checker_agent",
26+
# "SyntaxCheckerAgent",
27+
# "setup_semantic_verifier_agent",
28+
# "SemanticVerifierAgent",
29+
# ]

src/backend/sql_agents/agent_factory.py

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,15 +5,13 @@
55

66
from common.models.api import AgentType
77
from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
8-
from sql_agents import (
9-
FixerAgent,
10-
MigratorAgent,
11-
PickerAgent,
12-
SemanticVerifierAgent,
13-
SyntaxCheckerAgent,
14-
)
158
from sql_agents.agent_base import BaseSQLAgent
169
from sql_agents.agent_config import AgentBaseConfig
10+
from sql_agents.fixer.agent import FixerAgent
11+
from sql_agents.migrator.agent import MigratorAgent
12+
from sql_agents.picker.agent import PickerAgent
13+
from sql_agents.semantic_verifier.agent import SemanticVerifierAgent
14+
from sql_agents.syntax_checker.agent import SyntaxCheckerAgent
1715

1816
logger = logging.getLogger(__name__)
1917
logger.setLevel(logging.DEBUG)

src/backend/sql_agents/migrator/agent.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,9 @@
22

33
import logging
44

5+
from sql_agents.agent_base import BaseSQLAgent
56
from sql_agents.migrator.response import MigratorResponse
67

7-
from backend.sql_agents.agent_base import BaseSQLAgent
8-
98
logger = logging.getLogger(__name__)
109
logger.setLevel(logging.DEBUG)
1110

src/backend/sql_agents/migrator/setup.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55
from common.models.api import AgentType
66
from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
77
from sql_agents.agent_config import AgentBaseConfig
8-
9-
from backend.sql_agents.agent_factory import SQLAgentFactory
8+
from sql_agents.agent_factory import SQLAgentFactory
109

1110
logger = logging.getLogger(__name__)
1211
logger.setLevel(logging.DEBUG)

src/backend/sql_agents/picker/agent.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,9 @@
22

33
import logging
44

5+
from sql_agents.agent_base import BaseSQLAgent
56
from sql_agents.picker.response import PickerResponse
67

7-
from backend.sql_agents.agent_base import BaseSQLAgent
8-
98
logger = logging.getLogger(__name__)
109
logger.setLevel(logging.DEBUG)
1110

src/backend/sql_agents/picker/setup.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55
from common.models.api import AgentType
66
from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
77
from sql_agents.agent_config import AgentBaseConfig
8-
9-
from backend.sql_agents.agent_factory import SQLAgentFactory
8+
from sql_agents.agent_factory import SQLAgentFactory
109

1110
logger = logging.getLogger(__name__)
1211
logger.setLevel(logging.DEBUG)

src/backend/sql_agents/semantic_verifier/agent.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,11 +4,10 @@
44
from typing import Any, Dict, Optional
55

66
from common.models.api import AgentType
7+
from sql_agents.agent_base import BaseSQLAgent
78
from sql_agents.agent_config import AgentBaseConfig
89
from sql_agents.semantic_verifier.response import SemanticVerifierResponse
910

10-
from backend.sql_agents.agent_base import BaseSQLAgent
11-
1211
logger = logging.getLogger(__name__)
1312
logger.setLevel(logging.DEBUG)
1413

src/backend/sql_agents/semantic_verifier/setup.py

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,7 @@
55
from common.models.api import AgentType
66
from semantic_kernel.agents.azure_ai.azure_ai_agent import AzureAIAgent
77
from sql_agents.agent_config import AgentBaseConfig
8-
9-
from backend.sql_agents.agent_factory import SQLAgentFactory
8+
from sql_agents.agent_factory import SQLAgentFactory
109

1110
logger = logging.getLogger(__name__)
1211
logger.setLevel(logging.DEBUG)

src/backend/sql_agents_start.py

Lines changed: 22 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -34,22 +34,20 @@
3434
ChatMessageContent,
3535
)
3636
from 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-
)
4537
from sql_agents.agent_config import AgentBaseConfig
4638
from sql_agents.fixer.response import FixerResponse
39+
from sql_agents.fixer.setup import setup_fixer_agent
4740
from sql_agents.helpers.selection_function import setup_selection_function
41+
from sql_agents.helpers.sk_utils import create_kernel_with_chat_completion
4842
from sql_agents.helpers.termination_function import setup_termination_function
4943
from sql_agents.helpers.utils import is_text
5044
from sql_agents.migrator.response import MigratorResponse
45+
from sql_agents.migrator.setup import setup_migrator_agent
5146
from sql_agents.picker.response import PickerResponse
47+
from sql_agents.picker.setup import setup_picker_agent
5248
from 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

@@ -84,28 +82,12 @@
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

135117
async 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

396381
async 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

Comments
 (0)