SQL查询子智能体，支持自动纠错循环（ReAct/Reflection 模式）
| 22 | |
| 23 | |
| 24 | class SQLQueryAgent: |
| 25 | """SQL查询子智能体，支持自动纠错循环（ReAct/Reflection 模式）""" |
| 26 | |
| 27 | def __init__(self, llm: BaseLLM, db_path: str, num_examples: int = 3): |
| 28 | """初始化SQL查询智能体 |
| 29 | |
| 30 | Args: |
| 31 | llm: 语言模型实例 |
| 32 | db_path: 数据库路径 |
| 33 | num_examples: Few-shot示例数量 |
| 34 | """ |
| 35 | self.llm = llm |
| 36 | self.db_path = db_path |
| 37 | self.num_examples = num_examples |
| 38 | |
| 39 | @staticmethod |
| 40 | def _llm_to_str(result) -> str: |
| 41 | """安全地从 LLM 返回值中提取文本，清理思考标签""" |
| 42 | import re |
| 43 | if isinstance(result, str): |
| 44 | text = result |
| 45 | elif hasattr(result, 'content'): |
| 46 | text = str(result.content) |
| 47 | elif hasattr(result, 'text'): |
| 48 | text = str(result.text) |
| 49 | else: |
| 50 | text = str(result) |
| 51 | text = re.sub(r'<think>[\s\S]*?</think>', '', text).strip() |
| 52 | text = re.sub(r'</think>', '', text).strip() |
| 53 | return text |
| 54 | |
| 55 | def _get_schema(self) -> str: |
| 56 | """获取数据库Schema""" |
| 57 | conn = sqlite3.connect(self.db_path) |
| 58 | cursor = conn.cursor() |
| 59 | |
| 60 | cursor.execute(""" |
| 61 | SELECT name FROM sqlite_master |
| 62 | WHERE type='table' AND name NOT LIKE 'sqlite_%' |
| 63 | ORDER BY name |
| 64 | """) |
| 65 | tables = cursor.fetchall() |
| 66 | |
| 67 | schema_text = "" |
| 68 | for table in tables: |
| 69 | table_name = table[0] |
| 70 | schema_text += f"\n表：{table_name}\n" |
| 71 | |
| 72 | cursor.execute(f"PRAGMA table_info({table_name})") |
| 73 | columns = cursor.fetchall() |
| 74 | |
| 75 | for col in columns: |
| 76 | cid, name, dtype, notnull, default, pk = col |
| 77 | pk_text = " (主键)" if pk else "" |
| 78 | notnull_text = " NOT NULL" if notnull else "" |
| 79 | schema_text += f" - {name}: {dtype}{notnull_text}{pk_text}\n" |
| 80 | |
| 81 | conn.close() |