Преглед изворни кода

chore: make prompt generator max tokens configurable (#6693)

Sangmin Ahn пре 9 месеци
родитељ
комит
ecb9c311b5

+ 1 - 0
api/.env.example

@@ -183,6 +183,7 @@ UPLOAD_IMAGE_FILE_SIZE_LIMIT=10
 
 # Model Configuration
 MULTIMODAL_SEND_IMAGE_FORMAT=base64
+PROMPT_GENERATION_MAX_TOKENS=512
 
 # Mail configuration, support: resend, smtp
 MAIL_TYPE=

+ 5 - 1
api/controllers/console/app/generator.py

@@ -1,3 +1,5 @@
+import os
+
 from flask_login import current_user
 from flask_restful import Resource, reqparse
 
@@ -28,13 +30,15 @@ class RuleGenerateApi(Resource):
         args = parser.parse_args()
 
         account = current_user
+        PROMPT_GENERATION_MAX_TOKENS = int(os.getenv('PROMPT_GENERATION_MAX_TOKENS', '512'))
 
         try:
             rules = LLMGenerator.generate_rule_config(
                 tenant_id=account.current_tenant_id,
                 instruction=args['instruction'],
                 model_config=args['model_config'],
-                no_variable=args['no_variable']
+                no_variable=args['no_variable'],
+                rule_config_max_tokens=PROMPT_GENERATION_MAX_TOKENS
             )
         except ProviderTokenNotInitError as ex:
             raise ProviderNotInitializeError(ex.description)

+ 2 - 2
api/core/llm_generator/llm_generator.py

@@ -118,7 +118,7 @@ class LLMGenerator:
         return questions
 
     @classmethod
-    def generate_rule_config(cls, tenant_id: str, instruction: str, model_config: dict, no_variable: bool) -> dict:
+    def generate_rule_config(cls, tenant_id: str, instruction: str, model_config: dict, no_variable: bool, rule_config_max_tokens: int = 512) -> dict:
         output_parser = RuleConfigGeneratorOutputParser()
 
         error = ""
@@ -130,7 +130,7 @@ class LLMGenerator:
             "error": ""
         }
         model_parameters = {
-            "max_tokens": 512,
+            "max_tokens": rule_config_max_tokens,
             "temperature": 0.01
         }