Coverage for mindsdb / interfaces / skills / custom / text2sql / mindsdb_sql_tool.py: 0%

1from typing import Type 

2import re 

3from mindsdb_sql_parser import parse_sql 

4from pydantic import BaseModel, Field 

5from langchain_core.tools import BaseTool 

6 

7 

8from mindsdb.interfaces.agents.mindsdb_database_agent import extract_essential 

9 

10 

11class _MindsDBSQLParserToolInput(BaseModel): 

12 tool_input: str = Field("", description="A SQL query to validate.") 

13 

14 

15class MindsDBSQLParserTool(BaseTool): 

16 name: str = "mindsdb_sql_parser_tool" 

17 description: str = "Parse a SQL query to check it is valid MindsDB SQL." 

18 args_schema: Type[BaseModel] = _MindsDBSQLParserToolInput 

19 

20 def _clean_query(self, query: str) -> str: 

21 # Sometimes LLM can input markdown into query tools. 

22 cmd = re.sub(r'```(sql)?', '', query) 

23 return cmd 

24 

25 def _query_options(self, query): 

26 yield query 

27 if '\\_' in query: 

28 yield query.replace('\\_', '_') 

29 

30 def _run(self, query: str): 

31 """Validate the SQL query.""" 

32 query = extract_essential(query) 

33 clean_query = self._clean_query(query) 

34 for query in self._query_options(clean_query): 

35 try: 

36 ast_query = parse_sql(query) 

37 return "".join(f"valid query: {ast_query.to_string()}") 

38 except Exception as e: 

39 error = "".join(f"invalid query, with error: {e}") 

40 continue 

41 return error