feat: support binding context var (#1227)

Co-authored-by: Joel <iamjoel007@gmail.com>

api/commands.py

@@ -647,6 +647,76 @@ def update_app_model_configs(batch_size):
 
             pbar.update(len(data_batch))
 
+@click.command('migrate_default_input_to_dataset_query_variable')
+@click.option("--batch-size", default=500, help="Number of records to migrate in each batch.")
+def migrate_default_input_to_dataset_query_variable(batch_size):
+
+    click.secho("Starting...", fg='green')
+
+    total_records = db.session.query(AppModelConfig) \
+        .join(App, App.app_model_config_id == AppModelConfig.id) \
+        .filter(App.mode == 'completion') \
+        .filter(AppModelConfig.dataset_query_variable == None) \
+        .count()
+    
+    if total_records == 0:
+        click.secho("No data to migrate.", fg='green')
+        return
+
+    num_batches = (total_records + batch_size - 1) // batch_size
+    
+    with tqdm(total=total_records, desc="Migrating Data") as pbar:
+        for i in range(num_batches):
+            offset = i * batch_size
+            limit = min(batch_size, total_records - offset)
+
+            click.secho(f"Fetching batch {i + 1}/{num_batches} from source database...", fg='green')
+
+            data_batch = db.session.query(AppModelConfig) \
+                .join(App, App.app_model_config_id == AppModelConfig.id) \
+                .filter(App.mode == 'completion') \
+                .filter(AppModelConfig.dataset_query_variable == None) \
+                .order_by(App.created_at) \
+                .offset(offset).limit(limit).all()
+
+            if not data_batch:
+                click.secho("No more data to migrate.", fg='green')
+                break
+
+            try:
+                click.secho(f"Migrating {len(data_batch)} records...", fg='green')
+                for data in data_batch:
+                    config = AppModelConfig.to_dict(data)
+
+                    tools = config["agent_mode"]["tools"]
+                    dataset_exists = "dataset" in str(tools)
+                    if not dataset_exists:
+                        continue
+
+                    user_input_form = config.get("user_input_form", [])
+                    for form in user_input_form:
+                        paragraph = form.get('paragraph')
+                        if paragraph \
+                            and paragraph.get('variable') == 'query':
+                                data.dataset_query_variable = 'query'
+                                break
+                        
+                        if paragraph \
+                            and paragraph.get('variable') == 'default_input':
+                                data.dataset_query_variable = 'default_input'
+                                break
+
+                db.session.commit()
+
+            except Exception as e:
+                click.secho(f"Error while migrating data: {e}, app_id: {data.app_id}, app_model_config_id: {data.id}",
+                            fg='red')
+                continue
+            
+            click.secho(f"Successfully migrated batch {i + 1}/{num_batches}.", fg='green')
+
+            pbar.update(len(data_batch))
+
 
 def register_commands(app):
     app.cli.add_command(reset_password)
@@ -660,3 +730,4 @@ def register_commands(app):
     app.cli.add_command(update_qdrant_indexes)
     app.cli.add_command(update_app_model_configs)
     app.cli.add_command(normalization_collections)
+    app.cli.add_command(migrate_default_input_to_dataset_query_variable)

api/controllers/console/app/app.py

@@ -34,6 +34,7 @@ model_config_fields = {
     'sensitive_word_avoidance': fields.Raw(attribute='sensitive_word_avoidance_dict'),
     'model': fields.Raw(attribute='model_dict'),
     'user_input_form': fields.Raw(attribute='user_input_form_list'),
+    'dataset_query_variable': fields.String,
     'pre_prompt': fields.String,
     'agent_mode': fields.Raw(attribute='agent_mode_dict'),
 }
@@ -162,7 +163,8 @@ class AppListApi(Resource):
             model_configuration = AppModelConfigService.validate_configuration(
                 tenant_id=current_user.current_tenant_id,
                 account=current_user,
-                config=model_config_dict
+                config=model_config_dict,
+                mode=args['mode']
             )
 
             app = App(

api/controllers/console/app/model_config.py

@@ -31,7 +31,8 @@ class ModelConfigResource(Resource):
         model_configuration = AppModelConfigService.validate_configuration(
             tenant_id=current_user.current_tenant_id,
             account=current_user,
-            config=request.json
+            config=request.json,
+            mode=app_model.mode
         )
 
         new_app_model_config = AppModelConfig(

api/core/completion.py

@@ -108,12 +108,14 @@ class Completion:
                 retriever_from=retriever_from
             )
 
+            query_for_agent = cls.get_query_for_agent(app, app_model_config, query, inputs)
+
             # run agent executor
             agent_execute_result = None
-            if agent_executor:
-                should_use_agent = agent_executor.should_use_agent(query)
+            if query_for_agent and agent_executor:
+                should_use_agent = agent_executor.should_use_agent(query_for_agent)
                 if should_use_agent:
-                    agent_execute_result = agent_executor.run(query)
+                    agent_execute_result = agent_executor.run(query_for_agent)
 
             # When no extra pre prompt is specified,
             # the output of the agent can be used directly as the main output content without calling LLM again
@@ -142,6 +144,13 @@ class Completion:
             logging.warning(f'ChunkedEncodingError: {e}')
             conversation_message_task.end()
             return
+        
+    @classmethod
+    def get_query_for_agent(cls, app: App, app_model_config: AppModelConfig, query: str, inputs: dict) -> str:
+        if app.mode != 'completion':
+            return query
+        
+        return inputs.get(app_model_config.dataset_query_variable, "")
 
     @classmethod
     def run_final_llm(cls, model_instance: BaseLLM, mode: str, app_model_config: AppModelConfig, query: str,

api/migrations/versions/ab23c11305d4_add_dataset_query_variable_at_app_model_.py

@@ -0,0 +1,31 @@
+"""add dataset query variable at app model configs.
+
+Revision ID: ab23c11305d4
+Revises: 6e2cfb077b04
+Create Date: 2023-09-26 12:22:59.044088
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision = 'ab23c11305d4'
+down_revision = '6e2cfb077b04'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('app_model_configs', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('dataset_query_variable', sa.String(length=255), nullable=True))
+
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    with op.batch_alter_table('app_model_configs', schema=None) as batch_op:
+        batch_op.drop_column('dataset_query_variable')
+
+    # ### end Alembic commands ###

api/models/model.py

@@ -88,6 +88,7 @@ class AppModelConfig(db.Model):
     more_like_this = db.Column(db.Text)
     model = db.Column(db.Text)
     user_input_form = db.Column(db.Text)
+    dataset_query_variable = db.Column(db.String(255))
     pre_prompt = db.Column(db.Text)
     agent_mode = db.Column(db.Text)
     sensitive_word_avoidance = db.Column(db.Text)
@@ -152,6 +153,7 @@ class AppModelConfig(db.Model):
             "sensitive_word_avoidance": self.sensitive_word_avoidance_dict,
             "model": self.model_dict,
             "user_input_form": self.user_input_form_list,
+            "dataset_query_variable": self.dataset_query_variable,
             "pre_prompt": self.pre_prompt,
             "agent_mode": self.agent_mode_dict
         }
@@ -170,6 +172,7 @@ class AppModelConfig(db.Model):
             if model_config.get('sensitive_word_avoidance') else None
         self.model = json.dumps(model_config['model'])
         self.user_input_form = json.dumps(model_config['user_input_form'])
+        self.dataset_query_variable = model_config['dataset_query_variable']
         self.pre_prompt = model_config['pre_prompt']
         self.agent_mode = json.dumps(model_config['agent_mode'])
         self.retriever_resource = json.dumps(model_config['retriever_resource']) \
@@ -191,6 +194,7 @@ class AppModelConfig(db.Model):
             sensitive_word_avoidance=self.sensitive_word_avoidance,
             model=self.model,
             user_input_form=self.user_input_form,
+            dataset_query_variable=self.dataset_query_variable,
             pre_prompt=self.pre_prompt,
             agent_mode=self.agent_mode
         )

api/services/app_model_config_service.py

@@ -81,7 +81,7 @@ class AppModelConfigService:
         return filtered_cp
 
     @staticmethod
-    def validate_configuration(tenant_id: str, account: Account, config: dict) -> dict:
+    def validate_configuration(tenant_id: str, account: Account, config: dict, mode: str) -> dict:
         # opening_statement
         if 'opening_statement' not in config or not config["opening_statement"]:
             config["opening_statement"] = ""
@@ -335,6 +335,9 @@ class AppModelConfigService:
 
                 if not AppModelConfigService.is_dataset_exists(account, tool_item["id"]):
                     raise ValueError("Dataset ID does not exist, please check your permission.")
+        
+        # dataset_query_variable
+        AppModelConfigService.is_dataset_query_variable_valid(config, mode)
 
         # Filter out extra parameters
         filtered_config = {
@@ -351,8 +354,25 @@ class AppModelConfigService:
                 "completion_params": config["model"]["completion_params"]
             },
             "user_input_form": config["user_input_form"],
+            "dataset_query_variable": config["dataset_query_variable"],
             "pre_prompt": config["pre_prompt"],
             "agent_mode": config["agent_mode"]
         }
 
         return filtered_config
+    
+    @staticmethod
+    def is_dataset_query_variable_valid(config: dict, mode: str) -> None:
+        # Only check when mode is completion
+        if mode != 'completion':
+            return
+        
+        agent_mode = config.get("agent_mode", {})
+        tools = agent_mode.get("tools", [])
+        dataset_exists = "dataset" in str(tools)
+        
+        dataset_query_variable = config.get("dataset_query_variable")
+
+        if dataset_exists and not dataset_query_variable:
+            raise ValueError("Dataset query variable is required when dataset is exist")
+

api/services/completion_service.py

@@ -117,7 +117,8 @@ class CompletionService:
                 model_config = AppModelConfigService.validate_configuration(
                     tenant_id=app_model.tenant_id,
                     account=user,
-                    config=args['model_config']
+                    config=args['model_config'],
+                    mode=app_model.mode
                 )
 
                 app_model_config = AppModelConfig(

web/app/components/app/configuration/base/warning-mask/cannot-query-dataset.tsx

@@ -0,0 +1,31 @@
+'use client'
+import type { FC } from 'react'
+import React from 'react'
+import { useTranslation } from 'react-i18next'
+import WarningMask from '.'
+import Button from '@/app/components/base/button'
+
+export type IFormattingChangedProps = {
+  onConfirm: () => void
+}
+
+const FormattingChanged: FC<IFormattingChangedProps> = ({
+  onConfirm,
+}) => {
+  const { t } = useTranslation()
+
+  return (
+    <WarningMask
+      title={t('appDebug.feature.dataSet.queryVariable.unableToQueryDataSet')}
+      description={t('appDebug.feature.dataSet.queryVariable.unableToQueryDataSetTip')}
+      footer={
+        <div className='flex space-x-2'>
+          <Button type='primary' className='flex items-center justify-start !h-8 !w-[96px]' onClick={onConfirm}>
+            <span className='text-[13px] font-medium'>{t('appDebug.feature.dataSet.queryVariable.ok')}</span>
+          </Button>
+        </div>
+      }
+    />
+  )
+}
+export default React.memo(FormattingChanged)

web/app/components/app/configuration/base/warning-mask/formatting-changed.tsx

@@ -1,10 +1,11 @@
 'use client'
-import React, { FC } from 'react'
+import type { FC } from 'react'
+import React from 'react'
 import { useTranslation } from 'react-i18next'
 import WarningMask from '.'
 import Button from '@/app/components/base/button'
 
-export interface IFormattingChangedProps {
+export type IFormattingChangedProps = {
   onConfirm: () => void
   onCancel: () => void
 }
@@ -17,7 +18,7 @@ const icon = (
 
 const FormattingChanged: FC<IFormattingChangedProps> = ({
   onConfirm,
-  onCancel
+  onCancel,
 }) => {
   const { t } = useTranslation()
 

web/app/components/app/configuration/dataset-config/context-var/index.tsx

@@ -0,0 +1,39 @@
+'use client'
+import type { FC } from 'react'
+import React from 'react'
+import { useTranslation } from 'react-i18next'
+import cn from 'classnames'
+import type { Props } from './var-picker'
+import VarPicker from './var-picker'
+import { BracketsX } from '@/app/components/base/icons/src/vender/line/development'
+import Tooltip from '@/app/components/base/tooltip'
+import { HelpCircle } from '@/app/components/base/icons/src/vender/line/general'
+
+const ContextVar: FC<Props> = (props) => {
+  const { t } = useTranslation()
+  const { value, options } = props
+  const currItem = options.find(item => item.value === value)
+  const notSetVar = !currItem
+  return (
+    <div className={cn(notSetVar ? 'rounded-bl-xl rounded-br-xl bg-[#FEF0C7] border-[#FEF0C7]' : 'border-gray-200', 'flex justify-between items-center h-12 px-3 border-t ')}>
+      <div className='flex items-center space-x-1 shrink-0'>
+        <div className='p-1'>
+          <BracketsX className='w-4 h-4 text-primary-500'/>
+        </div>
+        <div className='mr-1 text-sm font-medium text-gray-800'>{t('appDebug.feature.dataSet.queryVariable.title')}</div>
+        <Tooltip
+          htmlContent={<div className='w-[180px]'>
+            {t('appDebug.feature.dataSet.queryVariable.tip')}
+          </div>}
+          selector='context-var-tooltip'
+        >
+          <HelpCircle className='w-3.5 h-3.5 text-gray-400'/>
+        </Tooltip>
+      </div>
+
+      <VarPicker {...props} />
+    </div>
+  )
+}
+
+export default React.memo(ContextVar)

web/app/components/app/configuration/dataset-config/context-var/style.module.css

@@ -0,0 +1,3 @@
+.trigger:hover .dropdownIcon {
+  color: #98A2B3;
+}

web/app/components/app/configuration/dataset-config/context-var/var-picker.tsx

@@ -0,0 +1,99 @@
+'use client'
+import type { FC } from 'react'
+import React, { useState } from 'react'
+import { useTranslation } from 'react-i18next'
+import { ChevronDownIcon } from '@heroicons/react/24/outline'
+import cn from 'classnames'
+import s from './style.module.css'
+import {
+  PortalToFollowElem,
+  PortalToFollowElemContent,
+  PortalToFollowElemTrigger,
+} from '@/app/components/base/portal-to-follow-elem'
+import type { IInputTypeIconProps } from '@/app/components/app/configuration/config-var/input-type-icon'
+import IconTypeIcon from '@/app/components/app/configuration/config-var/input-type-icon'
+
+type Option = { name: string; value: string; type: string }
+export type Props = {
+  value: string | undefined
+  options: Option[]
+  onChange: (value: string) => void
+}
+
+const VarItem: FC<{ item: Option }> = ({ item }) => (
+  <div className='flex items-center h-[18px] px-1 bg-[#EFF8FF] rounded space-x-1'>
+    <IconTypeIcon type={item.type as IInputTypeIconProps['type']} className='text-[#1570EF]' />
+    <div className='flex text-xs font-medium text-[#1570EF]'>
+      <span className='opacity-60'>{'{{'}</span>
+      <span className='max-w-[150px] truncate'>{item.value}</span>
+      <span className='opacity-60'>{'}}'}</span>
+    </div>
+  </div>
+)
+const VarPicker: FC<Props> = ({
+  value,
+  options,
+  onChange,
+}) => {
+  const { t } = useTranslation()
+  const [open, setOpen] = useState(false)
+  const currItem = options.find(item => item.value === value)
+  const notSetVar = !currItem
+  return (
+    <PortalToFollowElem
+      open={open}
+      onOpenChange={setOpen}
+      placement='bottom-end'
+      offset={{
+        mainAxis: 8,
+      }}
+    >
+      <PortalToFollowElemTrigger onClick={() => setOpen(v => !v)}>
+        <div className={cn(
+          s.trigger,
+          notSetVar ? 'bg-[#FFFCF5] border-[#FEDF89] text-[#DC6803]' : ' hover:bg-gray-50 border-gray-200 text-primary-600',
+          open ? 'bg-gray-50' : 'bg-white',
+          `
+          flex items-center h-8 justify-center px-2 space-x-1 rounded-lg border  shadow-xs cursor-pointer
+          text-[13px]  font-medium
+          `)}>
+          <div>
+            {value
+              ? (
+                <VarItem item={currItem as Option} />
+              )
+              : (<div>
+                {t('appDebug.feature.dataSet.queryVariable.choosePlaceholder')}
+              </div>)}
+          </div>
+          <ChevronDownIcon className={cn(s.dropdownIcon, open && 'rotate-180 text-[#98A2B3]', 'w-3.5 h-3.5')} />
+        </div>
+      </PortalToFollowElemTrigger>
+      <PortalToFollowElemContent style={{ zIndex: 1000 }}>
+        {options.length > 0
+          ? (<div className='w-[240px] max-h-[50vh] overflow-y-auto p-1  border bg-white border-gray-200 rounded-lg shadow-lg'>
+            {options.map(({ name, value, type }, index) => (
+              <div
+                key={index}
+                className='px-3 py-1 flex rounded-lg hover:bg-gray-50 cursor-pointer'
+                onClick={() => {
+                  onChange(value)
+                  setOpen(false)
+                }}
+              >
+                <VarItem item={{ name, value, type }} />
+              </div>
+            ))}
+          </div>)
+          : (
+            <div className='w-[240px] p-6 bg-white border border-gray-200 rounded-lg shadow-lg'>
+              <div className='mb-1 text-sm font-medium text-gray-700'>{t('appDebug.feature.dataSet.queryVariable.noVar')}</div>
+              <div className='text-xs leading-normal text-gray-500'>{t('appDebug.feature.dataSet.queryVariable.noVarTip')}</div>
+            </div>
+          )}
+
+      </PortalToFollowElemContent>
+    </PortalToFollowElem>
+  )
+}
+export default React.memo(VarPicker)

web/app/components/app/configuration/dataset-config/index.tsx

@@ -10,8 +10,10 @@ import FeaturePanel from '../base/feature-panel'
 import OperationBtn from '../base/operation-btn'
 import CardItem from './card-item'
 import SelectDataSet from './select-dataset'
+import ContextVar from './context-var'
 import ConfigContext from '@/context/debug-configuration'
 import type { DataSet } from '@/models/datasets'
+import { AppType } from '@/types/app'
 
 const Icon = (
   <svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
@@ -23,9 +25,12 @@ const Icon = (
 const DatasetConfig: FC = () => {
   const { t } = useTranslation()
   const {
+    mode,
     dataSets: dataSet,
     setDataSets: setDataSet,
     setFormattingChanged,
+    modelConfig,
+    setModelConfig,
   } = useContext(ConfigContext)
   const selectedIds = dataSet.map(item => item.id)
 
@@ -60,6 +65,25 @@ const DatasetConfig: FC = () => {
     setFormattingChanged(true)
   }
 
+  const promptVariables = modelConfig.configs.prompt_variables
+  const promptVariablesToSelect = promptVariables.map(item => ({
+    name: item.name,
+    type: item.type,
+    value: item.key,
+  }))
+  const selectedContextVar = promptVariables?.find(item => item.is_context_var)
+  const handleSelectContextVar = (selectedValue: string) => {
+    const newModelConfig = produce(modelConfig, (draft) => {
+      draft.configs.prompt_variables = modelConfig.configs.prompt_variables.map((item) => {
+        return ({
+          ...item,
+          is_context_var: item.key === selectedValue,
+        })
+      })
+    })
+    setModelConfig(newModelConfig)
+  }
+
   return (
     <FeaturePanel
       className='mt-3'
@@ -67,10 +91,11 @@ const DatasetConfig: FC = () => {
       title={t('appDebug.feature.dataSet.title')}
       headerRight={<OperationBtn type="add" onClick={showSelectDataSet} />}
       hasHeaderBottomBorder={!hasData}
+      noBodySpacing
     >
       {hasData
         ? (
-          <div className='flex flex-wrap justify-between'>
+          <div className='flex flex-wrap mt-1 px-3 justify-between'>
             {dataSet.map(item => (
               <CardItem
                 className="mb-2"
@@ -82,9 +107,19 @@ const DatasetConfig: FC = () => {
           </div>
         )
         : (
-          <div className='pt-2 pb-1 text-xs text-gray-500'>{t('appDebug.feature.dataSet.noData')}</div>
+          <div className='mt-1 px-3 pb-3'>
+            <div className='pt-2 pb-1 text-xs text-gray-500'>{t('appDebug.feature.dataSet.noData')}</div>
+          </div>
         )}
 
+      {mode === AppType.completion && dataSet.length > 0 && (
+        <ContextVar
+          value={selectedContextVar?.key}
+          options={promptVariablesToSelect}
+          onChange={handleSelectContextVar}
+        />
+      )}
+
       {isShowSelectDataSet && (
         <SelectDataSet
           isShow={isShowSelectDataSet}

web/app/components/app/configuration/debug/index.tsx

@@ -10,6 +10,7 @@ import dayjs from 'dayjs'
 import HasNotSetAPIKEY from '../base/warning-mask/has-not-set-api'
 import FormattingChanged from '../base/warning-mask/formatting-changed'
 import GroupName from '../base/group-name'
+import CannotQueryDataset from '../base/warning-mask/cannot-query-dataset'
 import { AppType } from '@/types/app'
 import PromptValuePanel, { replaceStringWithValues } from '@/app/components/app/configuration/prompt-value-panel'
 import type { IChatItem } from '@/app/components/app/chat/type'
@@ -23,7 +24,6 @@ import { promptVariablesToUserInputsForm } from '@/utils/model-config'
 import TextGeneration from '@/app/components/app/text-generate/item'
 import { IS_CE_EDITION } from '@/config'
 import { useProviderContext } from '@/context/provider-context'
-
 type IDebug = {
   hasSetAPIKEY: boolean
   onSetting: () => void
@@ -52,6 +52,7 @@ const Debug: FC<IDebug> = ({
     dataSets,
     modelConfig,
     completionParams,
+    hasSetContextVar,
   } = useContext(ConfigContext)
   const { speech2textDefaultModel } = useProviderContext()
   const [chatList, setChatList, getChatList] = useGetState<IChatItem[]>([])
@@ -77,6 +78,7 @@ const Debug: FC<IDebug> = ({
   const [isResponsing, { setTrue: setResponsingTrue, setFalse: setResponsingFalse }] = useBoolean(false)
   const [abortController, setAbortController] = useState<AbortController | null>(null)
   const [isShowFormattingChangeConfirm, setIsShowFormattingChangeConfirm] = useState(false)
+  const [isShowCannotQueryDataset, setShowCannotQueryDataset] = useState(false)
   const [isShowSuggestion, setIsShowSuggestion] = useState(false)
   const [messageTaskId, setMessageTaskId] = useState('')
   const [hasStopResponded, setHasStopResponded, getHasStopResponded] = useGetState(false)
@@ -157,6 +159,7 @@ const Debug: FC<IDebug> = ({
     const postModelConfig: BackendModelConfig = {
       pre_prompt: modelConfig.configs.prompt_template,
       user_input_form: promptVariablesToUserInputsForm(modelConfig.configs.prompt_variables),
+      dataset_query_variable: '',
       opening_statement: introduction,
       more_like_this: {
         enabled: false,
@@ -298,6 +301,7 @@ const Debug: FC<IDebug> = ({
   }, [controlClearChatMessage])
 
   const [completionRes, setCompletionRes] = useState('')
+  const [messageId, setMessageId] = useState<string | null>(null)
 
   const sendTextCompletion = async () => {
     if (isResponsing) {
@@ -305,6 +309,11 @@ const Debug: FC<IDebug> = ({
       return false
     }
 
+    if (dataSets.length > 0 && !hasSetContextVar) {
+      setShowCannotQueryDataset(true)
+      return true
+    }
+
     if (!checkCanSend())
       return
 
@@ -314,10 +323,12 @@ const Debug: FC<IDebug> = ({
         id,
       },
     }))
+    const contextVar = modelConfig.configs.prompt_variables.find(item => item.is_context_var)?.key
 
     const postModelConfig: BackendModelConfig = {
       pre_prompt: modelConfig.configs.prompt_template,
       user_input_form: promptVariablesToUserInputsForm(modelConfig.configs.prompt_variables),
+      dataset_query_variable: contextVar || '',
       opening_statement: introduction,
       suggested_questions_after_answer: suggestedQuestionsAfterAnswerConfig,
       speech_to_text: speechToTextConfig,
@@ -340,13 +351,15 @@ const Debug: FC<IDebug> = ({
     }
 
     setCompletionRes('')
+    setMessageId('')
     const res: string[] = []
 
     setResponsingTrue()
     sendCompletionMessage(appId, data, {
-      onData: (data: string) => {
+      onData: (data: string, _isFirstMessage: boolean, { messageId }) => {
         res.push(data)
         setCompletionRes(res.join(''))
+        setMessageId(messageId)
       },
       onCompleted() {
         setResponsingFalse()
@@ -415,6 +428,7 @@ const Debug: FC<IDebug> = ({
                 content={completionRes}
                 isLoading={!completionRes && isResponsing}
                 isInstalledApp={false}
+                messageId={messageId}
               />
             )}
           </div>
@@ -425,6 +439,11 @@ const Debug: FC<IDebug> = ({
             onCancel={handleCancel}
           />
         )}
+        {isShowCannotQueryDataset && (
+          <CannotQueryDataset
+            onConfirm={() => setShowCannotQueryDataset(false)}
+          />
+        )}
       </div>
 
       {!hasSetAPIKEY && (<HasNotSetAPIKEY isTrailFinished={!IS_CE_EDITION} onSetting={onSetting} />)}

web/app/components/app/configuration/index.tsx

@@ -100,7 +100,8 @@ const Configuration: FC = () => {
   }
 
   const [dataSets, setDataSets] = useState<DataSet[]>([])
-
+  const contextVar = modelConfig.configs.prompt_variables.find(item => item.is_context_var)?.key
+  const hasSetContextVar = !!contextVar
   const syncToPublishedConfig = (_publishedConfig: PublichConfig) => {
     const modelConfig = _publishedConfig.modelConfig
     setModelConfig(_publishedConfig.modelConfig)
@@ -178,7 +179,7 @@ const Configuration: FC = () => {
           model_id: model.name,
           configs: {
             prompt_template: modelConfig.pre_prompt,
-            prompt_variables: userInputsFormToPromptVariables(modelConfig.user_input_form),
+            prompt_variables: userInputsFormToPromptVariables(modelConfig.user_input_form, modelConfig.dataset_query_variable),
           },
           opening_statement: modelConfig.opening_statement,
           more_like_this: modelConfig.more_like_this,
@@ -196,16 +197,22 @@ const Configuration: FC = () => {
     })
   }, [appId])
 
-  const cannotPublish = mode === AppType.completion && !modelConfig.configs.prompt_template
+  const promptEmpty = mode === AppType.completion && !modelConfig.configs.prompt_template
+  const contextVarEmpty = mode === AppType.completion && dataSets.length > 0 && !hasSetContextVar
+  const cannotPublish = promptEmpty || contextVarEmpty
   const saveAppConfig = async () => {
     const modelId = modelConfig.model_id
     const promptTemplate = modelConfig.configs.prompt_template
     const promptVariables = modelConfig.configs.prompt_variables
 
-    if (cannotPublish) {
+    if (promptEmpty) {
       notify({ type: 'error', message: t('appDebug.otherError.promptNoBeEmpty'), duration: 3000 })
       return
     }
+    if (contextVarEmpty) {
+      notify({ type: 'error', message: t('appDebug.feature.dataSet.queryVariable.contextVarNotEmpty'), duration: 3000 })
+      return
+    }
     const postDatasets = dataSets.map(({ id }) => ({
       dataset: {
         enabled: true,
@@ -217,6 +224,7 @@ const Configuration: FC = () => {
     const data: BackendModelConfig = {
       pre_prompt: promptTemplate,
       user_input_form: promptVariablesToUserInputsForm(promptVariables),
+      dataset_query_variable: contextVar || '',
       opening_statement: introduction || '',
       more_like_this: moreLikeThisConfig,
       suggested_questions_after_answer: suggestedQuestionsAfterAnswerConfig,
@@ -298,6 +306,7 @@ const Configuration: FC = () => {
       setModelConfig,
       dataSets,
       setDataSets,
+      hasSetContextVar,
     }}
     >
       <>

web/app/components/app/configuration/prompt-value-panel/index.tsx

@@ -6,7 +6,7 @@ import { useContext } from 'use-context-selector'
 import {
   PlayIcon,
 } from '@heroicons/react/24/solid'
-import VarIcon from '../base/icons/var-icon'
+import { BracketsX as VarIcon } from '@/app/components/base/icons/src/vender/line/development'
 import ConfigContext from '@/context/debug-configuration'
 import type { PromptVariable } from '@/models/debug'
 import { AppType } from '@/types/app'
@@ -131,7 +131,7 @@ const PromptValuePanel: FC<IPromptValuePanelProps> = ({
           `${!userInputFieldCollapse && 'mb-2'}`
         }>
           <div className='flex items-center space-x-1 cursor-pointer' onClick={() => setUserInputFieldCollapse(!userInputFieldCollapse)}>
-            <div className='flex items-center justify-center w-4 h-4'><VarIcon /></div>
+            <div className='flex items-center justify-center w-4 h-4'><VarIcon className='w-4 h-4 text-primary-500'/></div>
             <div className='text-xs font-medium text-gray-800'>{t('appDebug.inputs.userInputField')}</div>
             {
               userInputFieldCollapse

web/app/components/base/confirm/common.tsx

@@ -17,6 +17,7 @@ export type ConfirmCommonProps = {
   onConfirm?: () => void
   showOperate?: boolean
   showOperateCancel?: boolean
+  confirmBtnClassName?: string
   confirmText?: string
 }
 
@@ -29,6 +30,7 @@ const ConfirmCommon: FC<ConfirmCommonProps> = ({
   onConfirm,
   showOperate = true,
   showOperateCancel = true,
+  confirmBtnClassName,
   confirmText,
 }) => {
   const { t } = useTranslation()
@@ -72,7 +74,7 @@ const ConfirmCommon: FC<ConfirmCommonProps> = ({
               }
               <Button
                 type='primary'
-                className=''
+                className={confirmBtnClassName || ''}
                 onClick={onConfirm}
               >
                 {confirmText || CONFIRM_MAP[type].confirmText}

web/app/components/base/icons/assets/vender/line/development/brackets-x.svg

@@ -0,0 +1,3 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M18.5708 20C19.8328 20 20.8568 18.977 20.8568 17.714V13.143L21.9998 12L20.8568 10.857V6.286C20.8568 5.023 19.8338 4 18.5708 4M5.429 4C4.166 4 3.143 5.023 3.143 6.286V10.857L2 12L3.143 13.143V17.714C3.143 18.977 4.166 20 5.429 20M15 9L9 15M9 9L15 15" stroke="black" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"/>
+</svg>

web/app/components/base/icons/assets/vender/solid/editor/paragraph.svg

@@ -0,0 +1,5 @@
+<svg width="24" height="24" viewBox="0 0 24 24" fill="none" xmlns="http://www.w3.org/2000/svg">
+<path d="M2 6.5C2 5.67157 2.67157 5 3.5 5H20.5C21.3284 5 22 5.67157 22 6.5C22 7.32843 21.3284 8 20.5 8H3.5C2.67157 8 2 7.32843 2 6.5Z" fill="black"/>
+<path d="M2 12.5C2 11.6716 2.67157 11 3.5 11H20.5C21.3284 11 22 11.6716 22 12.5C22 13.3284 21.3284 14 20.5 14H3.5C2.67157 14 2 13.3284 2 12.5Z" fill="black"/>
+<path d="M2 18.5C2 17.6716 2.67157 17 3.5 17H12.5C13.3284 17 14 17.6716 14 18.5C14 19.3284 13.3284 20 12.5 20H3.5C2.67157 20 2 19.3284 2 18.5Z" fill="black"/>
+</svg>

web/app/components/base/icons/src/vender/line/development/BracketsX.json

@@ -0,0 +1,29 @@
+{
+	"icon": {
+		"type": "element",
+		"isRootNode": true,
+		"name": "svg",
+		"attributes": {
+			"width": "24",
+			"height": "24",
+			"viewBox": "0 0 24 24",
+			"fill": "none",
+			"xmlns": "http://www.w3.org/2000/svg"
+		},
+		"children": [
+			{
+				"type": "element",
+				"name": "path",
+				"attributes": {
+					"d": "M18.5708 20C19.8328 20 20.8568 18.977 20.8568 17.714V13.143L21.9998 12L20.8568 10.857V6.286C20.8568 5.023 19.8338 4 18.5708 4M5.429 4C4.166 4 3.143 5.023 3.143 6.286V10.857L2 12L3.143 13.143V17.714C3.143 18.977 4.166 20 5.429 20M15 9L9 15M9 9L15 15",
+					"stroke": "currentColor",
+					"stroke-width": "2",
+					"stroke-linecap": "round",
+					"stroke-linejoin": "round"
+				},
+				"children": []
+			}
+		]
+	},
+	"name": "BracketsX"
+}

web/app/components/base/icons/src/vender/line/development/BracketsX.tsx

@@ -0,0 +1,16 @@
+// GENERATE BY script
+// DON NOT EDIT IT MANUALLY
+
+import * as React from 'react'
+import data from './BracketsX.json'
+import IconBase from '@/app/components/base/icons/IconBase'
+import type { IconBaseProps, IconData } from '@/app/components/base/icons/IconBase'
+
+const Icon = React.forwardRef<React.MutableRefObject<SVGElement>, Omit<IconBaseProps, 'data'>>((
+  props,
+  ref,
+) => <IconBase {...props} ref={ref} data={data as IconData} />)
+
+Icon.displayName = 'BracketsX'
+
+export default Icon

web/app/components/base/icons/src/vender/line/development/index.ts

@@ -1,3 +1,4 @@
+export { default as BracketsX } from './BracketsX'
 export { default as Container } from './Container'
 export { default as Database01 } from './Database01'
 export { default as Database03 } from './Database03'

web/app/components/base/icons/src/vender/solid/editor/Paragraph.json

@@ -0,0 +1,44 @@
+{
+	"icon": {
+		"type": "element",
+		"isRootNode": true,
+		"name": "svg",
+		"attributes": {
+			"width": "24",
+			"height": "24",
+			"viewBox": "0 0 24 24",
+			"fill": "none",
+			"xmlns": "http://www.w3.org/2000/svg"
+		},
+		"children": [
+			{
+				"type": "element",
+				"name": "path",
+				"attributes": {
+					"d": "M2 6.5C2 5.67157 2.67157 5 3.5 5H20.5C21.3284 5 22 5.67157 22 6.5C22 7.32843 21.3284 8 20.5 8H3.5C2.67157 8 2 7.32843 2 6.5Z",
+					"fill": "currentColor"
+				},
+				"children": []
+			},
+			{
+				"type": "element",
+				"name": "path",
+				"attributes": {
+					"d": "M2 12.5C2 11.6716 2.67157 11 3.5 11H20.5C21.3284 11 22 11.6716 22 12.5C22 13.3284 21.3284 14 20.5 14H3.5C2.67157 14 2 13.3284 2 12.5Z",
+					"fill": "currentColor"
+				},
+				"children": []
+			},
+			{
+				"type": "element",
+				"name": "path",
+				"attributes": {
+					"d": "M2 18.5C2 17.6716 2.67157 17 3.5 17H12.5C13.3284 17 14 17.6716 14 18.5C14 19.3284 13.3284 20 12.5 20H3.5C2.67157 20 2 19.3284 2 18.5Z",
+					"fill": "currentColor"
+				},
+				"children": []
+			}
+		]
+	},
+	"name": "Paragraph"
+}

web/app/components/base/icons/src/vender/solid/editor/Paragraph.tsx

@@ -0,0 +1,16 @@
+// GENERATE BY script
+// DON NOT EDIT IT MANUALLY
+
+import * as React from 'react'
+import data from './Paragraph.json'
+import IconBase from '@/app/components/base/icons/IconBase'
+import type { IconBaseProps, IconData } from '@/app/components/base/icons/IconBase'
+
+const Icon = React.forwardRef<React.MutableRefObject<SVGElement>, Omit<IconBaseProps, 'data'>>((
+  props,
+  ref,
+) => <IconBase {...props} ref={ref} data={data as IconData} />)
+
+Icon.displayName = 'Paragraph'
+
+export default Icon

web/app/components/base/icons/src/vender/solid/editor/TypeSquare.tsx

@@ -0,0 +1,16 @@
+// GENERATE BY script
+// DON NOT EDIT IT MANUALLY
+
+import * as React from 'react'
+import data from './TypeSquare.json'
+import IconBase from '@/app/components/base/icons/IconBase'
+import type { IconBaseProps, IconData } from '@/app/components/base/icons/IconBase'
+
+const Icon = React.forwardRef<React.MutableRefObject<SVGElement>, Omit<IconBaseProps, 'data'>>((
+  props,
+  ref,
+) => <IconBase {...props} ref={ref} data={data as IconData} />)
+
+Icon.displayName = 'TypeSquare'
+
+export default Icon

web/app/components/base/icons/src/vender/solid/editor/index.ts

@@ -1,2 +1,4 @@
 export { default as Brush01 } from './Brush01'
 export { default as Citations } from './Citations'
+export { default as Paragraph } from './Paragraph'
+export { default as TypeSquare } from './TypeSquare'

web/app/components/base/icons/src/vender/solid/general/CheckDone01.tsx

@@ -0,0 +1,16 @@
+// GENERATE BY script
+// DON NOT EDIT IT MANUALLY
+
+import * as React from 'react'
+import data from './CheckDone01.json'
+import IconBase from '@/app/components/base/icons/IconBase'
+import type { IconBaseProps, IconData } from '@/app/components/base/icons/IconBase'
+
+const Icon = React.forwardRef<React.MutableRefObject<SVGElement>, Omit<IconBaseProps, 'data'>>((
+  props,
+  ref,
+) => <IconBase {...props} ref={ref} data={data as IconData} />)
+
+Icon.displayName = 'CheckDone01'
+
+export default Icon

web/app/components/base/icons/src/vender/solid/general/index.ts

@@ -1,4 +1,5 @@
 export { default as CheckCircle } from './CheckCircle'
+export { default as CheckDone01 } from './CheckDone01'
 export { default as Download02 } from './Download02'
 export { default as Target04 } from './Target04'
 export { default as XCircle } from './XCircle'

web/context/debug-configuration.ts

@@ -37,6 +37,7 @@ type IDebugConfiguration = {
   setModelConfig: (modelConfig: ModelConfig) => void
   dataSets: DataSet[]
   setDataSets: (dataSet: DataSet[]) => void
+  hasSetContextVar: boolean
 }
 
 const DebugConfigurationContext = createContext<IDebugConfiguration>({
@@ -102,6 +103,7 @@ const DebugConfigurationContext = createContext<IDebugConfiguration>({
   setModelConfig: () => { },
   dataSets: [],
   setDataSets: () => { },
+  hasSetContextVar: false,
 })
 
 export default DebugConfigurationContext

web/i18n/lang/app-debug.en.ts

@@ -66,6 +66,19 @@ const translation = {
       noDataSet: 'No dataset found',
       toCreate: 'Go to create',
       notSupportSelectMulti: 'Currently only support one dataset',
+      queryVariable: {
+        title: 'Query variable',
+        tip: 'This variable will be used as the query input for context retrieval, obtaining context information related to the input of this variable.',
+        choosePlaceholder: 'Choose query variable',
+        noVar: 'No variables',
+        noVarTip: 'please create a variable under the Variables section',
+        unableToQueryDataSet: 'Unable to query the dataset',
+        unableToQueryDataSetTip: 'Unable to query the dataset successfully, please choose a context query variable in the context section.',
+        ok: 'OK',
+        contextVarNotEmpty: 'context query variable can not be empty',
+        deleteContextVarTitle: 'Delete variable “{{varName}}”?',
+        deleteContextVarTip: 'This variable has been set as a context query variable, and removing it will impact the normal use of the dataset. If you still need to delete it, please reselect it in the context section.',
+      },
     },
   },
   automatic: {

web/i18n/lang/app-debug.zh.ts

@@ -66,6 +66,19 @@ const translation = {
       noDataSet: '未找到数据集',
       toCreate: '去创建',
       notSupportSelectMulti: '目前只支持引用一个数据集',
+      queryVariable: {
+        title: '查询变量',
+        tip: '该变量将用作上下文检索的查询输入，获取与该变量的输入相关的上下文信息。',
+        choosePlaceholder: '请选择变量',
+        noVar: '没有变量',
+        noVarTip: '请创建变量',
+        unableToQueryDataSet: '无法查询数据集',
+        unableToQueryDataSetTip: '无法成功查询数据集，请在上下文部分选择一个上下文查询变量。',
+        ok: '好的',
+        contextVarNotEmpty: '上下文查询变量不能为空',
+        deleteContextVarTitle: '删除变量“{{varName}}”？',
+        deleteContextVarTip: '该变量已被设置为上下文查询变量，删除该变量将影响数据集的正常使用。 如果您仍需要删除它，请在上下文部分中重新选择它。',
+      },
     },
   },
   automatic: {

web/models/debug.ts

@@ -8,6 +8,7 @@ export type PromptVariable = {
   required: boolean
   options?: string[]
   max_length?: number
+  is_context_var?: boolean
 }
 
 export type CompletionParams = {

web/tailwind.config.js

@@ -27,6 +27,7 @@ module.exports = {
           200: '#C3DDFD',
           300: '#A4CAFE',
           400: '#528BFF',
+          500: '#2970FF',
           600: '#1C64F2',
           700: '#1A56DB',
         },

web/types/app.ts

@@ -90,6 +90,7 @@ export type ModelConfig = {
   opening_statement: string
   pre_prompt: string
   user_input_form: UserInputFormItem[]
+  dataset_query_variable?: string
   more_like_this: {
     enabled: boolean
   }

web/utils/model-config.ts

@@ -1,12 +1,13 @@
 import type { UserInputFormItem } from '@/types/app'
 import type { PromptVariable } from '@/models/debug'
 
-export const userInputsFormToPromptVariables = (useInputs: UserInputFormItem[] | null) => {
+export const userInputsFormToPromptVariables = (useInputs: UserInputFormItem[] | null, dataset_query_variable?: string) => {
   if (!useInputs)
     return []
   const promptVariables: PromptVariable[] = []
   useInputs.forEach((item: any) => {
     const isParagraph = !!item.paragraph
+
     const [type, content] = (() => {
       if (isParagraph)
         return ['paragraph', item.paragraph]
@@ -16,6 +17,8 @@ export const userInputsFormToPromptVariables = (useInputs: UserInputFormItem[] |
 
       return ['select', item.select]
     })()
+    const is_context_var = dataset_query_variable === content.variable
+
     if (type === 'string' || type === 'paragraph') {
       promptVariables.push({
         key: content.variable,
@@ -24,6 +27,7 @@ export const userInputsFormToPromptVariables = (useInputs: UserInputFormItem[] |
         type,
         max_length: content.max_length,
         options: [],
+        is_context_var,
       })
     }
     else {
@@ -33,6 +37,7 @@ export const userInputsFormToPromptVariables = (useInputs: UserInputFormItem[] |
         required: content.required,
         type: 'select',
         options: content.options,
+        is_context_var,
       })
     }
   })

web/yarn.lock

@@ -3752,11 +3752,7 @@ mdast-util-from-markdown@^0.8.5:
     parse-entities "^2.0.0"
     unist-util-stringify-position "^2.0.0"
 
-<<<<<<< Updated upstream
-mdast-util-from-markdown@^1.0.0, mdast-util-from-markdown@^1.1.0, mdast-util-from-markdown@^1.3.0:
-=======
 mdast-util-from-markdown@^1.0.0, mdast-util-from-markdown@^1.1.0:
->>>>>>> Stashed changes
   version "1.3.1"
   resolved "https://registry.npmjs.org/mdast-util-from-markdown/-/mdast-util-from-markdown-1.3.1.tgz"
   integrity sha512-4xTO/M8c82qBcnQc1tgpNtubGUW/Y1tBQ1B0i5CtSoelOLKFYlElIr3bvgREYYO5iRqbMY1YuqZng0GVOI8Qww==
@@ -5372,14 +5368,6 @@ safe-regex@^2.1.1:
   dependencies:
     regexp-tree "~0.1.1"
 
-<<<<<<< Updated upstream
-"safer-buffer@>= 2.1.2 < 3.0.0":
-  version "2.1.2"
-  resolved "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz"
-  integrity sha512-YZo3K82SD7Riyi0E1EQPojLz7kpepnSQI9IyPbHHg1XXXevb5dJI7tpyN2ADxGcQbHG7vcyRHk0cbwqcQriUtg==
-
-=======
->>>>>>> Stashed changes
 sass@^1.61.0:
   version "1.62.1"
   resolved "https://registry.npmjs.org/sass/-/sass-1.62.1.tgz"
@@ -6127,14 +6115,6 @@ web-namespaces@^2.0.0:
   resolved "https://registry.npmjs.org/web-namespaces/-/web-namespaces-2.0.1.tgz"
   integrity sha512-bKr1DkiNa2krS7qxNtdrtHAmzuYGFQLiQ13TsorsdT6ULTkPLKuu5+GsFpDlg6JFjUTwX2DyhMPG2be8uPrqsQ==
 
-<<<<<<< Updated upstream
-web-worker@^1.2.0:
-  version "1.2.0"
-  resolved "https://registry.npmjs.org/web-worker/-/web-worker-1.2.0.tgz"
-  integrity sha512-PgF341avzqyx60neE9DD+XS26MMNMoUQRz9NOZwW32nPQrF6p77f1htcnjBSEV8BGMKZ16choqUG4hyI0Hx7mA==
-
-=======
->>>>>>> Stashed changes
 which-boxed-primitive@^1.0.2:
   version "1.0.2"
   resolved "https://registry.npmjs.org/which-boxed-primitive/-/which-boxed-primitive-1.0.2.tgz"