feat: support optional query content (#1097)

Co-authored-by: Garfield Dai <dai.hai@foxmail.com>

api/commands.py

@@ -6,6 +6,7 @@ import string
 import time
 
 import click
+from tqdm import tqdm
 from flask import current_app
 from langchain.embeddings import OpenAIEmbeddings
 from werkzeug.exceptions import NotFound
@@ -21,9 +22,9 @@ from libs.password import password_pattern, valid_password, hash_password
 from libs.helper import email as email_validate
 from extensions.ext_database import db
 from libs.rsa import generate_key_pair
-from models.account import InvitationCode, Tenant
+from models.account import InvitationCode, Tenant, TenantAccountJoin
 from models.dataset import Dataset, DatasetQuery, Document
-from models.model import Account
+from models.model import Account, AppModelConfig, App
 import secrets
 import base64
 
@@ -439,6 +440,107 @@ def update_qdrant_indexes():
 
     click.echo(click.style('Congratulations! Update {} dataset indexes.'.format(create_count), fg='green'))
 
+@click.command('update_app_model_configs', help='Migrate data to support paragraph variable.')
+@click.option("--batch-size", default=500, help="Number of records to migrate in each batch.")
+def update_app_model_configs(batch_size):
+    pre_prompt_template = '{{default_input}}'
+    user_input_form_template = {
+        "en-US": [
+            {
+                "paragraph": {
+                    "label": "Query",
+                    "variable": "default_input",
+                    "required": False,
+                    "default": ""
+                }
+            }
+        ],
+        "zh-Hans": [
+            {
+                "paragraph": {
+                    "label": "查询内容",
+                    "variable": "default_input",
+                    "required": False,
+                    "default": ""
+                }
+            }
+        ]
+    }
+
+    click.secho("Start migrate old data that the text generator can support paragraph variable.", fg='green')
+
+    total_records = db.session.query(AppModelConfig) \
+        .join(App, App.app_model_config_id == AppModelConfig.id) \
+        .filter(App.mode == 'completion') \
+        .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') \
+                .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:
+                    # click.secho(f"Migrating data {data.id}, pre_prompt: {data.pre_prompt}, user_input_form: {data.user_input_form}", fg='green')
+
+                    if data.pre_prompt is None:
+                        data.pre_prompt = pre_prompt_template
+                    else:
+                        if pre_prompt_template in data.pre_prompt:
+                            continue
+                        data.pre_prompt += pre_prompt_template
+
+                    app_data = db.session.query(App) \
+                        .filter(App.id == data.app_id) \
+                        .one()
+                    
+                    account_data = db.session.query(Account) \
+                        .join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id) \
+                        .filter(TenantAccountJoin.role == 'owner') \
+                        .filter(TenantAccountJoin.tenant_id == app_data.tenant_id) \
+                        .one_or_none()
+
+                    if not account_data:
+                        continue
+
+                    if data.user_input_form is None or data.user_input_form == 'null':
+                        data.user_input_form = json.dumps(user_input_form_template[account_data.interface_language])
+                    else:
+                        raw_json_data = json.loads(data.user_input_form)
+                        raw_json_data.append(user_input_form_template[account_data.interface_language][0])
+                        data.user_input_form = json.dumps(raw_json_data)
+
+                    # click.secho(f"Updated data {data.id}, pre_prompt: {data.pre_prompt}, user_input_form: {data.user_input_form}", fg='green')
+
+                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)
     app.cli.add_command(reset_email)
@@ -448,4 +550,5 @@ def register_commands(app):
     app.cli.add_command(sync_anthropic_hosted_providers)
     app.cli.add_command(clean_unused_dataset_indexes)
     app.cli.add_command(create_qdrant_indexes)
-    app.cli.add_command(update_qdrant_indexes)
+    app.cli.add_command(update_qdrant_indexes)
+    app.cli.add_command(update_app_model_configs)

api/constants/model_template.py

@@ -38,7 +38,18 @@ model_templates = {
                     "presence_penalty": 0,
                     "frequency_penalty": 0
                 }
-            })
+            }),
+            'user_input_form': json.dumps([
+                {
+                    "paragraph": {
+                        "label": "Query",
+                        "variable": "query",
+                        "required": True,
+                        "default": ""
+                    }
+                }
+            ]),
+            'pre_prompt': '{{query}}'
         }
     },
 

api/controllers/console/app/completion.py

@@ -39,7 +39,7 @@ class CompletionMessageApi(Resource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('model_config', type=dict, required=True, location='json')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         args = parser.parse_args()

api/controllers/console/explore/completion.py

@@ -31,7 +31,7 @@ class CompletionApi(InstalledAppResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         args = parser.parse_args()
 

api/controllers/service_api/app/completion.py

@@ -27,7 +27,7 @@ class CompletionApi(AppApiResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         parser.add_argument('user', type=str, location='json')
         args = parser.parse_args()

api/controllers/web/completion.py

@@ -29,7 +29,7 @@ class CompletionApi(WebApiResource):
 
         parser = reqparse.RequestParser()
         parser.add_argument('inputs', type=dict, required=True, location='json')
-        parser.add_argument('query', type=str, location='json')
+        parser.add_argument('query', type=str, location='json', default='')
         parser.add_argument('response_mode', type=str, choices=['blocking', 'streaming'], location='json')
         args = parser.parse_args()
 

api/core/model_providers/models/llm/base.py

@@ -342,7 +342,7 @@ class BaseLLM(BaseProviderModel):
             if order == 'context_prompt':
                 prompt += context_prompt_content
             elif order == 'pre_prompt':
-                prompt += (pre_prompt_content + '\n\n') if pre_prompt_content else ''
+                prompt += pre_prompt_content
 
         query_prompt = prompt_rules['query_prompt'] if 'query_prompt' in prompt_rules else '{{query}}'
 

api/core/prompt/generate_prompts/baichuan_chat.json

@@ -8,6 +8,6 @@
     "pre_prompt",
     "histories_prompt"
   ],
-  "query_prompt": "用户：{{query}}",
+  "query_prompt": "\n\n用户：{{query}}",
   "stops": ["用户："]
 }

api/core/prompt/generate_prompts/common_chat.json

@@ -8,6 +8,6 @@
     "pre_prompt",
     "histories_prompt"
   ],
-  "query_prompt": "Human: {{query}}\n\nAssistant: ",
+  "query_prompt": "\n\nHuman: {{query}}\n\nAssistant: ",
   "stops": ["\nHuman:", "</histories>"]
 }

api/services/app_model_config_service.py

@@ -216,8 +216,8 @@ class AppModelConfigService:
         variables = []
         for item in config["user_input_form"]:
             key = list(item.keys())[0]
-            if key not in ["text-input", "select"]:
-                raise ValueError("Keys in user_input_form list can only be 'text-input' or 'select'")
+            if key not in ["text-input", "select", "paragraph"]:
+                raise ValueError("Keys in user_input_form list can only be 'text-input', 'paragraph'  or 'select'")
 
             form_item = item[key]
             if 'label' not in form_item:

api/services/completion_service.py

@@ -34,7 +34,7 @@ class CompletionService:
         inputs = args['inputs']
         query = args['query']
 
-        if not query:
+        if app_model.mode != 'completion' and not query:
             raise ValueError('query is required')
 
         query = query.replace('\x00', '')
@@ -347,8 +347,8 @@ class CompletionService:
                 if value not in options:
                     raise ValueError(f"{variable} in input form must be one of the following: {options}")
             else:
-                if 'max_length' in variable:
-                    max_length = variable['max_length']
+                if 'max_length' in input_config:
+                    max_length = input_config['max_length']
                     if len(value) > max_length:
                         raise ValueError(f'{variable} in input form must be less than {max_length} characters')