++

main.py

@@ -16,10 +16,13 @@ PROJECT_PATH = os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0]
 PREPARE_PROMPT = os.path.join(PROJECT_PATH, 'prepare_prompt')
 KEYS_TEXT = os.path.join(PREPARE_PROMPT, 'keys_text.txt')
 
-USE_AI = 0
+# 是否使用AI生成
+USE_AI = 1
 
-TEMPERATURE = 0.1
+# 模型温度
+TEMPERATURE = 0.2
 
+# 随机补充数据集(数据集不够的情况下打开)
 RANDOM_DATA_SETS_COUNT = 0
 
 # 数据集筛选之后大于此数值, 则随机抽取x条数据
@@ -31,14 +34,21 @@ REGION = 'USA'
 
 UNIVERSE = 'TOP3000'
 
-LOAD_OPERATOR = 0
+# 读取操作符
+LOAD_OPERATOR = 1
+
+# 数据类别名
+CATEGORY_NAME_LIST = ['Analyst', 'Fundamental', 'Price Volume']
+
+# 基础提示词文件
+alpha_prompt_path = os.path.join(PREPARE_PROMPT, "alpha_prompt.txt")
 
 SILICONFLOW_API_KEY = "sk-pvdiisdowmuwkrpnxsrlhxaovicqibmlljwrwwvbbdjaitdl"
 SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
 MODELS = [
     'Pro/deepseek-ai/DeepSeek-V3.1-Terminus',
-    'deepseek-ai/DeepSeek-V3.2-Exp',
-    'Qwen/Qwen3-VL-235B-A22B-Instruct',
+    # 'deepseek-ai/DeepSeek-V3.2-Exp',
+    # 'Qwen/Qwen3-VL-235B-A22B-Instruct',
     # 'MiniMaxAI/MiniMax-M2',
     # 'zai-org/GLM-4.6',
     # 'inclusionAI/Ring-flash-2.0',
@@ -134,22 +144,32 @@ def sqliteLoader(file_path, keys_text):
         cursor = conn.cursor()
         
         # 首先筛选符合 region 和 universe 条件的数据
-        cursor.execute("SELECT id, name, description, region, universe FROM data_sets WHERE region=? AND universe=?", 
+        cursor.execute("SELECT id, name, description, region, universe, category_name FROM data_sets WHERE region=? AND universe=?", 
                       (REGION, UNIVERSE))
         rows = cursor.fetchall()
         
         for row in rows:
-            row_id, name, description, region, universe = row
+            row_id, name, description, region, universe, category_name = row
             # 检查关键词是否在 name 中
             for key in keys_text:
-                if key in name:
-                    item_id = str(row_id)
-                    if item_id not in data_dict:
-                        data_dict[item_id] = {
-                            'id': int(row_id),
-                            'data_set_name': f"可以使用:{name}",
-                            'description': f"不可使用,仅供参考:{description}"
-                        }
+                if CATEGORY_NAME_LIST:
+                    if key in name and category_name in CATEGORY_NAME_LIST:
+                        item_id = str(row_id)
+                        if item_id not in data_dict:
+                            data_dict[item_id] = {
+                                'id': int(row_id),
+                                'data_set_name': f"可以使用:{name}",
+                                'description': f"不可使用,仅供参考:{description}"
+                            }
+                else:
+                    if key in name:
+                        item_id = str(row_id)
+                        if item_id not in data_dict:
+                            data_dict[item_id] = {
+                                'id': int(row_id),
+                                'data_set_name': f"可以使用:{name}",
+                                'description': f"不可使用,仅供参考:{description}"
+                            }
         
         conn.close()
         
@@ -379,7 +399,6 @@ def prepare_prompt(data_sets):
     prompt = ''
 
     # 读取基础提示词
-    alpha_prompt_path = os.path.join(PREPARE_PROMPT, "alpha_prompt.txt")
     prompt += promptLoader(alpha_prompt_path)
 
     # 读取操作符