Aryan Jain commited on
Commit
71dcc32
·
1 Parent(s): f961e0d

fix: correct method signature for _model_schema_to_text; update OpenAI model name and fix parameter naming in PydanticAgent; enhance SQL query verification with logging and error handling

Browse files
src/services/_data_extractor.py CHANGED
@@ -34,7 +34,7 @@ class DataExtractorService:
34
  async def __aexit__(self, exc_type, exc_value, traceback):
35
  pass
36
 
37
- async def _model_schema_to_text(model_class) -> str:
38
  """Convert SQLAlchemy model schema to LLM-friendly text."""
39
  mapper = inspect(model_class)
40
  table_name = model_class.__tablename__
 
34
  async def __aexit__(self, exc_type, exc_value, traceback):
35
  pass
36
 
37
+ async def _model_schema_to_text(self, model_class) -> str:
38
  """Convert SQLAlchemy model schema to LLM-friendly text."""
39
  mapper = inspect(model_class)
40
  table_name = model_class.__tablename__
src/utils/_pydantic_agent.py CHANGED
@@ -18,15 +18,15 @@ class PydanticAgent:
18
  self
19
  ):
20
  self._system_prompt = SQL_QUERY_EXTRACTOR_PROMPT
21
- self._openai_model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-5.2")
22
  self._openai_model = OpenAIResponsesModel(
23
- model=self._openai_model_name,
24
  )
25
  self._agent = Agent(
26
  system_prompt=self._system_prompt,
27
  model=self._openai_model,
28
  output_type=SQLQueryExtractor,
29
- responses_model_settings=OpenAIResponsesModelSettings(temperature=0.0),
30
  tools=[self._verify_sql_query],
31
  retries=5,
32
  )
@@ -37,7 +37,8 @@ class PydanticAgent:
37
  async def __aexit__(self, exc_type, exc_value, traceback):
38
  pass
39
 
40
- async def _verify_sql_query(self, sql_query):
 
41
  try:
42
  words_shoould_not_present_in_sql_query = [
43
  "DELETE",
@@ -58,7 +59,8 @@ class PydanticAgent:
58
  async with DatabaseConfig.async_session() as session:
59
  await session.execute(text(sql_query))
60
  except Exception as e:
61
- raise e
 
62
  return True
63
 
64
  async def _run_with_backoff(self, agent, *args, retries=5, **kwargs):
 
18
  self
19
  ):
20
  self._system_prompt = SQL_QUERY_EXTRACTOR_PROMPT
21
+ self._openai_model_name = os.getenv("OPENAI_MODEL_NAME", "gpt-4.1-mini")
22
  self._openai_model = OpenAIResponsesModel(
23
+ model_name=self._openai_model_name,
24
  )
25
  self._agent = Agent(
26
  system_prompt=self._system_prompt,
27
  model=self._openai_model,
28
  output_type=SQLQueryExtractor,
29
+ model_settings=OpenAIResponsesModelSettings(temperature=0.0),
30
  tools=[self._verify_sql_query],
31
  retries=5,
32
  )
 
37
  async def __aexit__(self, exc_type, exc_value, traceback):
38
  pass
39
 
40
+ async def _verify_sql_query(self, sql_query: str) -> bool|str:
41
+ logger.info(f"Verifying SQL query: {sql_query}")
42
  try:
43
  words_shoould_not_present_in_sql_query = [
44
  "DELETE",
 
59
  async with DatabaseConfig.async_session() as session:
60
  await session.execute(text(sql_query))
61
  except Exception as e:
62
+ logger.error(e)
63
+ return str(e) + "\nPlease generate SQL Query again"
64
  return True
65
 
66
  async def _run_with_backoff(self, agent, *args, retries=5, **kwargs):