Coverage for mindsdb / interfaces / agents / mindsdb_database_agent.py: 0%

1""" 

2Wrapper around MindsDB's executor and integration controller following the implementation of the original 

3langchain.sql_database.SQLDatabase class to partly replicate its behavior. 

4""" 

5 

6from typing import Any, Iterable, List, Optional 

7 

8from langchain_community.utilities import SQLDatabase 

9 

10from mindsdb.utilities import log 

11from mindsdb.interfaces.skills.sql_agent import SQLAgent 

12 

13logger = log.getLogger(__name__) 

14 

15 

16def extract_essential(input: str) -> str: 

17 """Sometimes LLM include to input unnecessary data. We can't control stochastic nature of LLM, so we need to 

18 'clean' input somehow. LLM prompt contains instruction to enclose input between '$START$' and '$STOP$'. 

19 """ 

20 if "$START$" in input: 

21 input = input.partition("$START$")[-1] 

22 if "$STOP$" in input: 

23 input = input.partition("$STOP$")[0] 

24 return input.strip(" ") 

25 

26 

27class MindsDBSQL(SQLDatabase): 

28 @staticmethod 

29 def custom_init(sql_agent: "SQLAgent") -> "MindsDBSQL": 

30 instance = MindsDBSQL() 

31 instance._sql_agent = sql_agent 

32 return instance 

33 

34 """ Can't modify signature, as LangChain does a Pydantic check.""" 

35 

36 def __init__( 

37 self, 

38 engine: Optional[Any] = None, 

39 schema: Optional[str] = None, 

40 metadata: Optional[Any] = None, 

41 ignore_tables: Optional[List[str]] = None, 

42 include_tables: Optional[List[str]] = None, 

43 sample_rows_in_table_info: int = 3, 

44 indexes_in_table_info: bool = False, 

45 custom_table_info: Optional[dict] = None, 

46 view_support: bool = True, 

47 max_string_length: int = 300, 

48 lazy_table_reflection: bool = False, 

49 ): 

50 pass 

51 

52 @property 

53 def dialect(self) -> str: 

54 return "mindsdb" 

55 

56 @property 

57 def table_info(self) -> str: 

58 """Information about all tables in the database.""" 

59 return self._sql_agent.get_table_info() 

60 

61 @property 

62 def kb_info(self) -> str: 

63 """Information about all knowledge bases in the database.""" 

64 return self._sql_agent.get_knowledge_base_info() 

65 

66 def get_usable_table_names(self) -> Iterable[str]: 

67 """Information about all tables in the database.""" 

68 try: 

69 return self._sql_agent.get_usable_table_names() 

70 except Exception as e: 

71 # If there's an error accessing tables through the integration, handle it gracefully 

72 logger.warning(f"Error getting usable table names: {str(e)}") 

73 # Return an empty list instead of raising an exception 

74 return [] 

75 

76 def get_table_info_no_throw(self, table_names: Optional[List[str]] = None) -> str: 

77 for i in range(len(table_names)): 

78 table_names[i] = extract_essential(table_names[i]) 

79 return self._sql_agent.get_table_info_safe(table_names) 

80 

81 def run_no_throw(self, command: str, fetch: str = "all") -> str: 

82 """Execute a SQL command and return the result as a string. 

83 

84 This method catches any exceptions and returns an error message instead of raising an exception. 

85 

86 Args: 

87 command: The SQL command to execute 

88 fetch: Whether to fetch 'all' results or just 'one' 

89 

90 Returns: 

91 A string representation of the result or an error message 

92 """ 

93 command = extract_essential(command) 

94 

95 try: 

96 # Log the query for debugging 

97 logger.info(f"Executing SQL query: {command}") 

98 

99 return self._sql_agent.query(command) 

100 

101 except Exception as e: 

102 logger.exception("Error executing SQL command:") 

103 # If this is a knowledge base query, provide a more helpful error message 

104 if "knowledge_base" in command.lower() or any( 

105 kb in command for kb in self._sql_agent.get_usable_knowledge_base_names() 

106 ): 

107 return f"Error executing knowledge base query: {str(e)}. Please check that the knowledge base exists and your query syntax is correct." 

108 return f"Error: {str(e)}" 

109 

110 def get_usable_knowledge_base_names(self) -> List[str]: 

111 """Get a list of usable knowledge base names. 

112 

113 Returns: 

114 A list of knowledge base names that can be used in queries 

115 """ 

116 try: 

117 return self._sql_agent.get_usable_knowledge_base_names() 

118 except Exception: 

119 logger.exception("Error getting usable knowledge base names:") 

120 return [] 

121 

122 def check_knowledge_base_permission(self, name): 

123 """Get a list of usable knowledge base names. 

124 

125 Returns: 

126 A list of knowledge base names that can be used in queries 

127 """ 

128 

129 return self._sql_agent.check_knowledge_base_permission(name)