| 108 | |
| 109 | class SqlGPT: |
| 110 | def __init__(self, sqlUrl: Optional[str], userName: Optional[str], password: Optional[str]): |
| 111 | self.sqlUrl = sqlUrl |
| 112 | self.userName = userName |
| 113 | self.password = password |
| 114 | |
| 115 | self.llm = LLM() |
| 116 | |
| 117 | self.sqlHelper = SQLHelper() |
| 118 | self.orm = orm() |
| 119 | |
| 120 | self.sql_prompt_template = PromptTemplate.from_template( |
| 121 | "You now need to act as an SQL command intelligence generator, my current requirement is {question}. " |
| 122 | "You only need to give the specific SQL command, you don't need to give any explanation." |
| 123 | ) |
| 124 | |
| 125 | self.optimize_prompt_template = PromptTemplate.from_template( |
| 126 | "You now need to act as an SQL command optimizer and optimize the following sql statement {sql}, " |
| 127 | "the structure of the associated table is {structure}, and the index of the associated table is {index}. " |
| 128 | "Index optimization: Add appropriate indexes to your QL statements to improve query efficiency. " |
| 129 | "SQL statement refactoring: Optimize the structure of SQL statements to reduce query time" |
| 130 | ) |
| 131 | self.error_prompt_template = PromptTemplate.from_template( |
| 132 | "The input sql statement is {sql}, the error is {error}, " |
| 133 | "and the structure of the associated database is {structure}." |
| 134 | ) |
| 135 | |
| 136 | def generateSQL(self, question: Optional[str], dbID: Optional[str], need_operate: Optional[bool] = False, only_sql: Optional[bool] = True): |
| 137 | question = self.sql_prompt_template.format(question=question) |