Spaces:
Sleeping
Sleeping
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-
|
| 22 |
self._openai_model = OpenAIResponsesModel(
|
| 23 |
-
|
| 24 |
)
|
| 25 |
self._agent = Agent(
|
| 26 |
system_prompt=self._system_prompt,
|
| 27 |
model=self._openai_model,
|
| 28 |
output_type=SQLQueryExtractor,
|
| 29 |
-
|
| 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 |
-
|
|
|
|
| 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):
|