重构: ai chat 系统提示词支持传入更多的数据库信息,减少交互次数

sjjian · sjjian · commit 801a5ad14ba7 · 2026-02-03T12:43:20.000+08:00
diff --git a/client/lib/models/ai.dart b/client/lib/models/ai.dart
@@ -124,9 +124,15 @@ abstract class AIChatUserMessageModel with _$AIChatUserMessageModel {
   const factory AIChatUserMessageModel({
     required AIChatMessageId id,
     required String content,
+    String? ref,
   }) = _AIChatUserMessageModel;
 
-  String toMessage() => content;
+  String toMessage() {
+    final refText = ref?.trim() ?? '';
+    if (refText.isEmpty) return content;
+    // ref 作为“额外上下文”，拼到 user message 里供 LLM 使用，但 UI 仍可只展示 content。
+    return '$content\n\nref:\n$refText';
+  }
 }
 
 @freezed
diff --git a/client/lib/models/instances.dart b/client/lib/models/instances.dart
@@ -87,5 +87,6 @@ abstract class PaginationInstanceListModel with _$PaginationInstanceListModel {
 abstract class InstanceMetadataModel with _$InstanceMetadataModel {
   const factory InstanceMetadataModel({
     required List<MetaDataNode> metadata,
+    required String? version,
   }) = _InstanceMetadataModel;
 }
diff --git a/client/lib/models/sessions.dart b/client/lib/models/sessions.dart
@@ -288,7 +288,7 @@ abstract class SessionAIChatModel with _$SessionAIChatModel {
     required SessionId sessionId,
     required String? currentSchema,
     required DatabaseType? dbType,
-    required List<MetaDataNode>? metadata,
+    required InstanceMetadataModel? metadata,
     required ConnId? connId,
     required SQLConnectState? state,
     required AIChatModel chatModel,
diff --git a/client/lib/repositories/instances/instances.dart b/client/lib/repositories/instances/instances.dart
@@ -257,7 +257,8 @@ class InstanceRepoImpl extends InstanceRepo {
       conn = SessionConn(model: instance);
       await conn.connect();
       final metadataNode = await conn.metadata();
-      return InstanceMetadataModel(metadata: metadataNode);
+      final version = await conn.version();
+      return InstanceMetadataModel(metadata: metadataNode, version: version);
     } catch (e) {
       rethrow;
     } finally {
diff --git a/client/lib/repositories/instances/session_conn.dart b/client/lib/repositories/instances/session_conn.dart
@@ -233,6 +233,10 @@ class SessionConn {
   Future<List<MetaDataNode>> metadata() async {
     return await conn2!.metadata();
   }
+
+  Future<String?> version() async {
+    return await conn2!.version();
+  }
 }
 
 @Riverpod(keepAlive: true)
diff --git a/client/lib/screens/sessions/ai_chat/input_user.dart b/client/lib/screens/sessions/ai_chat/input_user.dart
@@ -30,18 +30,79 @@ class SessionChatInputCard extends ConsumerStatefulWidget {
 }
 
 class _SessionChatInputCardState extends ConsumerState<SessionChatInputCard> {
+  Set<String> _extractMentionedTables(String encoded) {
+    final segments = MentionSegmentSerializer.decode(encoded);
+    final tableNames = <String>{};
+
+    // 只取“真正的 mention token”
+    for (final s in segments) {
+      if (s is MentionSegment) {
+        tableNames.add(s.label);
+      }
+    }
+
+    return tableNames;
+  }
+
+  MetaDataNode? _findSchemaNode(SessionAIChatModel chatModel) {
+    if (chatModel.metadata == null || chatModel.currentSchema == null) return null;
+    final root = MetaDataNode(MetaType.instance, "", items: chatModel.metadata!.metadata);
+    MetaDataNode? schemaNode;
+    root.visitor((node, _) {
+      if (schemaNode != null) return false;
+      if (node.type == MetaType.schema && node.value == chatModel.currentSchema) {
+        schemaNode = node;
+        return false;
+      }
+      return true;
+    });
+    return schemaNode;
+  }
+
+  String _buildTableRef(SessionAIChatModel chatModel, Iterable<String> mentionedTables) {
+    final schemaNode = _findSchemaNode(chatModel);
+    if (schemaNode == null) return '';
+
+    final tables = mentionedTables.toSet().toList()..sort();
+    if (tables.isEmpty) return '';
+
+    final b = StringBuffer();
+    for (final tableName in tables) {
+      MetaDataNode? tableNode;
+      for (final n in (schemaNode.items ?? const <MetaDataNode>[])) {
+        if (n.type == MetaType.table && n.value == tableName) {
+          tableNode = n;
+          break;
+        }
+      }
+
+      // 直接复用 MetaDataNode.toString() 的 JSON 序列化（见 db_driver_metadata.dart）
+      if (tableNode != null) {
+        b.writeln(tableNode.toString());
+      }
+    }
+
+    return b.toString().trimRight();
+  }
+
   Future<void> _sendMessage(AIChatId chatId, SessionAIChatModel chatModel) async {
     final chatInputController = SessionController.sessionController(chatModel.sessionId).chatInputController;
     final encoded = chatInputController.text.trim();
-    final text = MentionSegmentSerializer.decode(encoded).map((s) => s.toDisplayText()).join().trim();
+    final segments = MentionSegmentSerializer.decode(encoded);
+    final text = segments.map((s) => s.toDisplayText()).join().trim();
     chatInputController.clear();
 
+    // 如果用户通过 @ 提及了表，则把表结构信息放到 ref 里
+    final mentionedTables = _extractMentionedTables(encoded);
+    final refText = _buildTableRef(chatModel, mentionedTables);
+
     // 调用AIChatService的chat方法
     await ref.read(aIChatServiceProvider.notifier).chat(
           chatId,
           chatModel.llmAgents.lastUsedLLMAgent!.id,
           genChatSystemPrompt(chatModel),
           message: text,
+          refText: refText.isEmpty ? null : refText,
         );
 
     final scrollController = SessionController.sessionController(chatModel.sessionId).aiChatScrollController;
@@ -80,9 +141,8 @@ class _SessionChatInputCardState extends ConsumerState<SessionChatInputCard> {
             // 输入框
             ChatInputFieldWidget(
               model: widget.model,
-              onSubmitted: widget.model.canSendMessage()
-                  ? () => _sendMessage(widget.model.chatModel.id, widget.model)
-                  : null,
+              onSubmitted:
+                  widget.model.canSendMessage() ? () => _sendMessage(widget.model.chatModel.id, widget.model) : null,
             ),
 
             const SizedBox(height: kSpacingSmall),
@@ -375,12 +435,9 @@ class _ChatInputFieldWidgetState extends ConsumerState<ChatInputFieldWidget> {
     if (widget.model.metadata == null || widget.model.currentSchema == null) {
       return [];
     }
-    final schema = MetaDataNode(MetaType.instance, "", items: widget.model.metadata!);
+    final schema = MetaDataNode(MetaType.instance, "", items: widget.model.metadata!.metadata);
     final schemaNodes = schema.getChildren(MetaType.schema, widget.model.currentSchema!);
-    return schemaNodes
-        .where((e) => e.type == MetaType.table)
-        .map((e) => e.value)
-        .toList();
+    return schemaNodes.where((e) => e.type == MetaType.table).map((e) => e.value).toList();
   }
 
   List<String> _filterAndSortTables(List<String> allTables, String query) {
diff --git a/client/lib/services/ai/chat.dart b/client/lib/services/ai/chat.dart
@@ -131,7 +131,13 @@ class AIChatService extends _$AIChatService {
   }
 
   /// 进行AI对话，请求接口，存储消息并刷新使用 provider 来动态刷新页面
-  Future<void> chat(AIChatId id, LLMAgentId agentId, String systemPrompt, {String? message}) async {
+  Future<void> chat(
+    AIChatId id,
+    LLMAgentId agentId,
+    String systemPrompt, {
+    String? message,
+    String? refText,
+  }) async {
     final repo = ref.read(aiChatRepoProvider);
     final model = repo.getAIChatById(id);
     if (model == null) {
@@ -142,7 +148,13 @@ class AIChatService extends _$AIChatService {
     if (message != null) {
       repo.addMessage(
         id,
-        AIChatMessageItem.userMessage(AIChatUserMessageModel(id: AIChatMessageId.generate(), content: message)),
+        AIChatMessageItem.userMessage(
+          AIChatUserMessageModel(
+            id: AIChatMessageId.generate(),
+            content: message,
+            ref: refText,
+          ),
+        ),
       );
     }
 
diff --git a/client/lib/services/ai/prompt.dart b/client/lib/services/ai/prompt.dart
@@ -8,38 +8,32 @@ To confirm that you are available, please only return a number 1 to me.
 
 const chatTemplate = """
 你是一个智能SQL客户端助手. 你正在与一个使用数据库工具的用户对话. 你正在帮助用户回答关于数据库的问题.
-## 数据库信息:
+## 当前数据库connection的一些基本信息:
 db type: {dbType}
+db version: {dbVersion}
+current schema: {currentSchema}
 
 ## 用户输入的格式:
 用户会通过@符号指定表名并在当前对话里将表信息传递给你, 给你辅助回答问题.
-@table_name 表示表名, 例如: @users.
-ref:
-table_name信息:
-- 表名: users
-- 表描述: 用户表
-- 表字段:
-  - id: 用户ID
-  - name: 用户名
-  - email: 用户邮箱
-  - created_at: 创建时间
-  - updated_at: 更新时间
+@table_name 表示表名, 例如: @users. 在`ref:`后面会传递表信息给你, 你需要根据表信息来辅助回答问题.
 
 ## 注意点: 
 - 你只能回答或解决与数据库相关的问题;
 - 如果回复包含SQL, 每个SQL应该被包裹在一个 ```sql``` 块中;
+- 信任用户传递的表信息，除非用户显式的表达你需要重新查询它;
 - 数据库的query查询是非常重要的工具, 你除了使用它进行数据库信息获取外，还可以用它来进行任务逻辑计算, 例如：`SELECT 100 * 30 as result`;
 - 在使用query工具时尽可能一次获取更多想要的信息, 避免多次调用query工具;
 - 在使用query工具时要保持返回必要信息, 不要返回无关信息，例如：只返回需要的列和行;
 - 在使用query工具时要注意性能问题,例如: 可使用limit等限制返回数据量, 避免返回过多数据导致性能问题;
 """;
 
 String genChatSystemPrompt(SessionAIChatModel model) {
-  String prompt = chatTemplate;
-  if (model.dbType != null) {
-    prompt = prompt.replaceAll("{dbType}", model.dbType!.name);
-  }
-  return prompt;
+  final dbVersion = (model.metadata?.version ?? "").trim();
+  final currentSchema = (model.currentSchema ?? "").trim();
+  return chatTemplate
+      .replaceAll("{dbType}", model.dbType?.name ?? "-")
+      .replaceAll("{dbVersion}", dbVersion.isEmpty ? "-" : dbVersion)
+      .replaceAll("{currentSchema}", currentSchema.isEmpty ? "-" : currentSchema);
 }
 
 // 导入任务的文件命名
diff --git a/client/lib/services/sessions/session_chat.dart b/client/lib/services/sessions/session_chat.dart
@@ -45,7 +45,7 @@ class SessionAIChatNotifier extends _$SessionAIChatNotifier {
       sessionId: session.sessionId,
       currentSchema: session.currentSchema,
       dbType: session.dbType,
-      metadata: metadata?.value?.metadata,
+      metadata: metadata?.value,
       connId: session.connId,
       state: session.connState,
       llmAgents: llmAgents,
diff --git a/pkg/db_driver/lib/src/db_driver_interface.dart b/pkg/db_driver/lib/src/db_driver_interface.dart
@@ -106,6 +106,7 @@ abstract class BaseConnection {
   Future<List<String>> schemas();
   Future<String?> getCurrentSchema();
   Future<void> setCurrentSchema(String schema);
+  Future<String> version();
 
   void listen(
       {Function()? onCloseCallback,
diff --git a/pkg/db_driver/lib/src/db_driver_mysql.dart b/pkg/db_driver/lib/src/db_driver_mysql.dart
@@ -258,6 +258,13 @@ class MySQLConnection extends BaseConnection {
     }
   }
 
+  @override
+  Future<String> version() async {
+    final results = await query("SELECT VERSION() AS version");
+    final rows = results.rows;
+    return rows.first.getString("version") ?? "";
+  }
+
   @override
   Future<List<MetaDataNode>> metadata() async {
     // ref: https://dev.mysql.com/doc/refman/8.4/en/information-schema-columns-table.html
diff --git a/pkg/db_driver/lib/src/db_driver_pg.dart b/pkg/db_driver/lib/src/db_driver_pg.dart
@@ -109,6 +109,13 @@ class PGConnection extends BaseConnection {
     await query("SELECT 1");
   }
 
+  @override
+  Future<String> version() async {
+    final results = await query("SELECT version() AS version");
+    final rows = results.rows;
+    return rows.first.getString("version") ?? "";
+  }
+
   @override
   Future<void> killQuery() async {
     return;

Original file line number	Diff line number	Diff line change
`@@ -87,5 +87,6 @@ abstract class PaginationInstanceListModel with _$PaginationInstanceListModel {`
`87`	`87`	`abstract class InstanceMetadataModel with _$InstanceMetadataModel {`
`88`	`88`	`const factory InstanceMetadataModel({`
`89`	`89`	`required List<MetaDataNode> metadata,`
	`90`	`+ required String? version,`
`90`	`91`	`}) = _InstanceMetadataModel;`
`91`	`92`	`}`
Original file line number	Diff line number	Diff line change
`@@ -233,6 +233,10 @@ class SessionConn {`
`233`	`233`	`Future<List<MetaDataNode>> metadata() async {`
`234`	`234`	`return await conn2!.metadata();`
`235`	`235`	`}`
	`236`	`+`
	`237`	`+ Future<String?> version() async {`
	`238`	`+ return await conn2!.version();`
	`239`	`+ }`
`236`	`240`	`}`
`237`	`241`
`238`	`242`	`@Riverpod(keepAlive: true)`