gopidesupavan commented on code in PR #67438:
URL: https://github.com/apache/airflow/pull/67438#discussion_r3299426936
##########
providers/common/ai/src/airflow/providers/common/ai/toolsets/sql.py:
##########
@@ -164,50 +150,36 @@ def _get_db_hook(self) -> DbApiHook:
return self._hook
# ------------------------------------------------------------------
- # AbstractToolset interface
+ # BaseToolset interface
# ------------------------------------------------------------------
- async def get_tools(self, ctx: RunContext[Any]) -> dict[str,
ToolsetTool[Any]]:
- tools: dict[str, ToolsetTool[Any]] = {}
-
- for name, description, schema in (
- ("list_tables", "List available table names in the database.",
_LIST_TABLES_SCHEMA),
- ("get_schema", "Get column names and types for a table.",
_GET_SCHEMA_SCHEMA),
- ("query", "Execute a SQL query and return rows as JSON.",
_QUERY_SCHEMA),
- ("check_query", "Validate SQL syntax without executing it.",
_CHECK_QUERY_SCHEMA),
- ):
- # sequential=True because all tools use a shared DbApiHook with
- # synchronous I/O — they must not run concurrently.
- tool_def = ToolDefinition(
- name=name,
- description=description,
- parameters_json_schema=schema,
- sequential=True,
- )
- tools[name] = ToolsetTool(
- toolset=self,
- tool_def=tool_def,
- max_retries=1,
- args_validator=_PASSTHROUGH_VALIDATOR,
- )
- return tools
-
- async def call_tool(
- self,
- name: str,
- tool_args: dict[str, Any],
- ctx: RunContext[Any],
- tool: ToolsetTool[Any],
- ) -> Any:
- if name == "list_tables":
- return self._list_tables()
- if name == "get_schema":
- return self._get_schema(tool_args["table_name"])
- if name == "query":
- return self._query(tool_args["sql"])
- if name == "check_query":
- return self._check_query(tool_args["sql"])
- raise ValueError(f"Unknown tool: {name!r}")
+ def as_tools(self) -> list[ToolSpec]:
Review Comment:
Yeah good call :) updated it now thank you :)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
