feat: AI助手现在可以直接操作生成/清除中心化/标准化/离散化子变量

LeafYeeXYZ · LeafYeeXYZ · commit 99b0f066e90e · 2025-03-05T14:18:30.000+08:00
diff --git a/src/components/AI.tsx b/src/components/AI.tsx
@@ -1,6 +1,6 @@
 // 记得在 GREETTING 消息里说明可以使用的功能
 // TODO: 写好之后更新一下使用文档的 2.5
-import type { AIFunction } from '../types'
+import { type AIFunction, ALLOWED_DISCRETE_METHODS } from '../types'
 import { useAssistant } from '../lib/useAssistant'
 import { useZustand } from '../lib/useZustand'
 import { useState, useRef, useEffect } from 'react'
@@ -34,6 +34,7 @@ import { nav_to_plots_view } from '../lib/assistant/nav_to_plots_view'
 import { nav_to_statistics_view } from '../lib/assistant/nav_to_statistics_view'
 import { nav_to_tools_view } from '../lib/assistant/nav_to_tools_view'
 import { create_new_var } from '../lib/assistant/create_new_var'
+import { create_sub_var, clear_sub_var } from '../lib/assistant/create_sub_var'
 
 const md = markdownit({ html: true, breaks: true })
 const funcs: AIFunction[] = [
@@ -44,9 +45,11 @@ const funcs: AIFunction[] = [
   nav_to_statistics_view,
   nav_to_tools_view,
   create_new_var,
+  create_sub_var,
+  clear_sub_var,
 ]
 const GREETTING =
-  '你好, 我是 PsychPen 的 AI 助手, 可以帮你讲解 PsychPen 的使用方法、探索你的数据集、导出数据、跳转页面、生成新变量等. 请问有什么可以帮你的?'
+  '你好, 我是 PsychPen 的 AI 助手, 可以帮你讲解 PsychPen 的使用方法、探索你的数据集、导出数据、跳转页面、生成/清除子变量 (标准化/中心化/离散化)、生成新变量等. 请问有什么可以帮你的?'
 
 export function AI() {
   const { ai, model } = useAssistant()
@@ -147,6 +150,35 @@ export function AI() {
 
         try {
           switch (toolCall.function.name) {
+            case 'clear_sub_var': {
+              const { variable_names } = JSON.parse(toolCall.function.arguments)
+              if (
+                !Array.isArray(variable_names) ||
+                !variable_names.every((name) =>  (typeof name === 'string') && dataCols.some((col) => col.name === name))
+              ) {
+                throw new Error('变量名参数错误')
+              }
+              newMessages[1].content = `已请求清除变量 ${(variable_names as string[]).map((name) => `"${name}"`).join('、')} 的所有子变量, 等待用户手动确认`
+              break
+            }
+            case 'create_sub_var': {
+              const { variable_names, standardize, centralize, discretize } =
+                JSON.parse(toolCall.function.arguments)
+              if (
+                !Array.isArray(variable_names) ||
+                !variable_names.every((name) =>  (typeof name === 'string') && dataCols.some((col) => col.name === name))
+              ) {
+                throw new Error('变量名参数错误')
+              }
+              newMessages[1].content = `已请求生成变量 ${(variable_names as string[]).map((name) => `"${name}"`).join('、')} 的${[
+                standardize ? '标准化' : '',
+                centralize ? '中心化' : '',
+                discretize ? '离散化' : '',
+              ]
+                .filter((part) => part)
+                .join('、')}子变量, 等待用户手动确认`
+              break
+            }
             case 'create_new_var': {
               // eslint-disable-next-line @typescript-eslint/no-unused-vars
               const { variable_name, calc_expression: _ } = JSON.parse(
@@ -397,12 +429,150 @@ function ToolCall({ toolCall }: { toolCall: ChatCompletionMessageToolCall }) {
   const id = toolCall.id
   const name = toolCall.function.name
   const args = toolCall.function.arguments
-  const { dataRows, _VariableView_addNewVar, messageApi } = useZustand()
+  const { dataRows, _VariableView_addNewVar, messageApi, dataCols, _VariableView_updateData } = useZustand()
   const [done, setDone] = useState(false)
   const formerDone = sessionStorage.getItem(id) === 'done'
   let element: React.ReactElement | null = null
   let initDone = true
   switch (name) {
+    case 'clear_sub_var': {
+      const { variable_names } = JSON.parse(args) as { variable_names: string[] }
+      if (!formerDone) {
+        initDone = false
+      }
+      element = (
+        <>
+          <div>
+            执行函数{' '}
+            <Tag color='blue' style={{ margin: 0 }}>
+              {funcs.find(
+                (func) => func.tool.function.name === toolCall.function.name,
+              )?.label || `未知函数 (${toolCall.function.name})`}
+            </Tag>
+            {done ? ', 已' : ', 是否确认'}清除变量
+            {variable_names.map((name) => (
+              <Tag key={name} style={{ margin: 0, marginLeft: '0.3rem' }} color='blue'>
+                {name}
+              </Tag>
+            ))}{' '}
+            的所有子变量
+          </div>
+          <div>
+            <Button
+              block
+              disabled={done}
+              onClick={() => {
+                _VariableView_updateData(
+                  dataCols.map((col) => {
+                    if (variable_names.includes(col.name)) {
+                      return {
+                        ...col,
+                        subVars: undefined,
+                      }
+                    }
+                    return col
+                  }).filter((col) => col.derived !== true),
+                )
+                setDone(true)
+                sessionStorage.setItem(id, 'done')
+                messageApi?.success(`已成功清除变量 ${variable_names.map((name) => `"${name}"`).join('、')} 的所有子变量`)
+              }}
+            >
+              {done ? '已清除子变量' : '确认清除子变量'}
+            </Button>
+          </div>
+        </>
+      )
+      break
+    }
+    case 'create_sub_var': {
+      const { variable_names, standardize, centralize, discretize } = JSON.parse(args) as {
+        variable_names: string[]
+        standardize: boolean | undefined
+        centralize: boolean | undefined
+        discretize: {
+          method: ALLOWED_DISCRETE_METHODS
+          groups: number
+        } | undefined
+      }
+      if (!formerDone) {
+        initDone = false
+      }
+      const ALLOWED_METHOD = Object.values(ALLOWED_DISCRETE_METHODS)
+      const shouldDiscritize = Boolean(
+        typeof discretize === 'object' && 
+        discretize.method && 
+        discretize.groups && 
+        ALLOWED_METHOD.includes(discretize.method)
+      )
+      element = (
+        <>
+          <div>
+            执行函数{' '}
+            <Tag color='blue' style={{ margin: 0 }}>
+              {funcs.find(
+                (func) => func.tool.function.name === toolCall.function.name,
+              )?.label || `未知函数 (${toolCall.function.name})`}
+            </Tag>
+            {done ? ', 已' : ', 是否确认'}生成变量
+            {variable_names.map((name) => (
+              <Tag key={name} style={{ margin: 0, marginLeft: '0.3rem' }} color='blue'>
+                {name}
+              </Tag>
+            ))}{' '}
+            的
+            {[
+              standardize ? '标准化' : '',
+              centralize ? '中心化' : '',
+              shouldDiscritize ? `离散化 (${discretize!.method}, ${discretize!.groups} 组) ` : '',
+            ]
+              .filter((part) => part)
+              .join('、')}
+            子变量
+          </div>
+          <div>
+            <Button
+              block
+              disabled={done}
+              onClick={() => {
+                _VariableView_updateData(
+                  dataCols.map((col) => {
+                    if (variable_names.includes(col.name)) {
+                      return {
+                        ...col,
+                        subVars: {
+                          standard: Boolean(standardize) || col.subVars?.standard,
+                          center: Boolean(centralize) || col.subVars?.center,
+                          discrete: shouldDiscritize
+                            ? {
+                                method: discretize!.method,
+                                groups: discretize!.groups,
+                              }
+                            : col.subVars?.discrete,
+                        },
+                      }
+                    }
+                    return col
+                  }).filter((col) => col.derived !== true),
+                )
+                setDone(true)
+                sessionStorage.setItem(id, 'done')
+                messageApi?.success(`已成功生成变量 ${variable_names.map((name) => `"${name}"`).join('、')} 的${[
+                  standardize ? '标准化' : '',
+                  centralize ? '中心化' : '',
+                  shouldDiscritize ? '离散化' : '',
+                ]
+                  .filter((part) => part)
+                  .join('、')}子变量`)
+              }}
+            >
+              {done ? '已生成子变量' : '确认生成子变量'}
+            </Button>
+          </div>
+        </>
+      )
+      break
+    }
     case 'create_new_var': {
       const { variable_name, calc_expression } = JSON.parse(args)
       if (!formerDone) {
diff --git a/src/components/variable/DataFilter.tsx b/src/components/variable/DataFilter.tsx
@@ -209,7 +209,10 @@ export function DataFilter() {
       method === ALLOWED_FILTER_METHODS.BELOW_MEDIAN
     ) {
       return null
-    } else if (method === ALLOWED_FILTER_METHODS.EQUAL || method === ALLOWED_FILTER_METHODS.NOT_EQUAL) {
+    } else if (
+      method === ALLOWED_FILTER_METHODS.EQUAL ||
+      method === ALLOWED_FILTER_METHODS.NOT_EQUAL
+    ) {
       const variable = form.getFieldValue('variable')
       const options = Array.from(new Set(dataRows.map((row) => row[variable])))
         .sort((a, b) => Number(a) - Number(b))
diff --git a/src/components/variable/SubVariables.tsx b/src/components/variable/SubVariables.tsx
@@ -16,10 +16,12 @@ type Option = {
   discretizeGroups?: number
 }
 
-const DISCRETE_METHODS = Object.values(ALLOWED_DISCRETE_METHODS).map((method) => ({
-  label: method,
-  value: method,
-}))
+const DISCRETE_METHODS = Object.values(ALLOWED_DISCRETE_METHODS).map(
+  (method) => ({
+    label: method,
+    value: method,
+  }),
+)
 const ALLOW_SUBVARS: {
   en: string
   cn: string
diff --git a/src/lib/assistant/create_sub_var.ts b/src/lib/assistant/create_sub_var.ts
@@ -0,0 +1,83 @@
+import type { AIFunction } from '../../types'
+import { ALLOWED_DISCRETE_METHODS } from '../../types'
+
+export const create_sub_var: AIFunction = {
+  label: '生成子变量',
+  tool: {
+    type: 'function',
+    function: {
+      name: 'create_sub_var',
+      description:
+        '你可以调用这个函数来帮助用户把指定变量的标准化/中心化/离散化子变量',
+      parameters: {
+        type: 'object',
+        properties: {
+          variable_names: {
+            type: 'array',
+            description: '所有要生成子变量的变量名',
+            items: {
+              type: 'string',
+              description: '变量名',
+            },
+          },
+          standardize: {
+            type: 'boolean',
+            description: '是否生成标准化子变量',
+          },
+          centralize: {
+            type: 'boolean',
+            description: '是否生成中心化子变量',
+          },
+          discretize: {
+            type: 'object',
+            description: '如果要生成离散化子变量, 则需要指定离散化算法和分组数',
+            properties: {
+              method: {
+                type: 'string',
+                description: `离散化算法, 可选值为: ${Object.values(
+                  ALLOWED_DISCRETE_METHODS,
+                ).join(', ')}`,
+              },
+              groups: {
+                type: 'number',
+                description: '离散化分组数',
+              },
+            },
+            required: ['method', 'groups'],
+            additionalProperties: false,
+          },
+        },
+        required: ['variable_name'],
+        additionalProperties: false,
+      },
+      strict: true,
+    },
+  },
+}
+
+export const clear_sub_var: AIFunction = {
+  label: '清除子变量',
+  tool: {
+    type: 'function',
+    function: {
+      name: 'clear_sub_var',
+      description: '你可以调用这个函数来帮助用户清除指定变量的所有子变量',
+      parameters: {
+        type: 'object',
+        properties: {
+          variable_names: {
+            type: 'array',
+            description: '所有要清除子变量的变量名',
+            items: {
+              type: 'string',
+              description: '变量名',
+            },
+          },
+        },
+        required: ['variable_name'],
+        additionalProperties: false,
+      },
+      strict: true,
+    },
+  },
+}
diff --git a/src/lib/calculates/derive.ts b/src/lib/calculates/derive.ts
@@ -97,7 +97,11 @@ class Discrete {
    * @param groups 分组数
    * @param methed 离散化方法
    */
-  constructor(data: number[], groups: number, methed: ALLOWED_DISCRETE_METHODS) {
+  constructor(
+    data: number[],
+    groups: number,
+    methed: ALLOWED_DISCRETE_METHODS,
+  ) {
     this.method = methed
     this.groups = groups
     this.#data = data.toSorted((a, b) => a - b)
