++

main.py

@@ -8,6 +8,7 @@ import csv
 from datetime import datetime
 import jieba
 import time
+import sqlite3
 
 sys.path.append(os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator'))
 PROJECT_PATH = os.path.join(os.path.abspath(__file__).split('AlphaGenerator')[0] + 'AlphaGenerator')
@@ -15,11 +16,20 @@ 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 = 1
+USE_AI = 0
 
 TEMPERATURE = 0.1
 
-RANDOM_DATA_SETS_COUNT = 100
+RANDOM_DATA_SETS_COUNT = 0
+
+# 数据集筛选之后大于此数值, 则随机抽取x条数据
+MAX_DATA_COUNT = 800
+RANDOM_DATA_COUNT = 400
+
+# 数据库搜索字段
+REGION = 'USA'
+
+UNIVERSE = 'TOP3000'
 
 SILICONFLOW_API_KEY = "sk-pvdiisdowmuwkrpnxsrlhxaovicqibmlljwrwwvbbdjaitdl"
 SILICONFLOW_BASE_URL = "https://api.siliconflow.cn/v1"
@@ -27,7 +37,7 @@ MODELS = [
     'Pro/deepseek-ai/DeepSeek-V3.1-Terminus',
     'deepseek-ai/DeepSeek-V3.2-Exp',
     'Qwen/Qwen3-VL-235B-A22B-Instruct',
-    'MiniMaxAI/MiniMax-M2',
+    # 'MiniMaxAI/MiniMax-M2',
     # 'zai-org/GLM-4.6',
     # 'inclusionAI/Ring-flash-2.0',
     # 'zai-org/GLM-4.6',
@@ -110,6 +120,45 @@ def csvFileLoader(file_path, keys_text):
     return list(data_dict.values())
 
 
+def sqliteLoader(file_path, keys_text):
+    if not os.path.exists(file_path):
+        print(f"SQLite数据库文件不存在: {file_path}")
+        exit(1)
+    
+    data_dict = {}  # 使用字典来存储，以id为键
+    
+    try:
+        conn = sqlite3.connect(file_path)
+        cursor = conn.cursor()
+        
+        # 首先筛选符合 region 和 universe 条件的数据
+        cursor.execute("SELECT id, name, description, region, universe FROM data_sets WHERE region=? AND universe=?", 
+                      (REGION, UNIVERSE))
+        rows = cursor.fetchall()
+        
+        for row in rows:
+            row_id, name, description, region, universe = 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}"
+                        }
+        
+        conn.close()
+        
+        # 将字典的值转换为列表
+        return list(data_dict.values())
+        
+    except sqlite3.Error as e:
+        print(f"SQLite数据库错误: {e}")
+        exit(1)
+
+
 def extend_data_sets(file_path, original_data_sets):
     result = original_data_sets.copy()
 
@@ -342,7 +391,7 @@ def prepare_prompt(data_sets):
     prompt += "\n========================= 操作符结束 =======================================\n\n"
 
     prompt += "========================= 数据字段开始 =======================================\n"
-    prompt += "注意: data_set_name: 后面的是数据字段(可以使用), description: 此字段后面的是数据字段对应的描述或使用说明(不能使用), description_cn字段后面的内容是中文使用说明(不能使用)\n\n"
+    prompt += "注意: data_set_name: 后面的是数据字段(可以使用), description: 此字段后面的是数据字段对应的描述或使用说明(不能使用)\n\n"
     for data_set in data_sets:
         prompt += str(data_set) + '\n'
 
@@ -362,14 +411,22 @@ def main():
     # 将金融逻辑, 分割成标签
     keys_text = keysTextLoader()
 
-    # 分割好的标签, 搜索对应的数据集, 返回匹配到的结果
-    data_sets_path = os.path.join(PREPARE_PROMPT, "all_data_combined.csv")
-    result_data_sets = csvFileLoader(data_sets_path, keys_text)
+    # # 分割好的标签, 搜索对应的数据集, 返回匹配到的结果
+    # data_sets_path = os.path.join(PREPARE_PROMPT, "all_data_combined.csv")
+    # result_data_sets = csvFileLoader(data_sets_path, keys_text)
+
+    # if not result_data_sets:
+    #     print(f'搜索数据集为空, 程序退出')
+    #     exit(1)
+
+    data_sets_path = os.path.join(PREPARE_PROMPT, "data_sets.db")
+    result_data_sets = sqliteLoader(data_sets_path, keys_text)
 
     if not result_data_sets:
         print(f'搜索数据集为空, 程序退出')
         exit(1)
 
+
     # 扩展数据集
     mistakes_notebook_path = os.path.join(PREPARE_PROMPT, "all_data_combined.csv")
     if RANDOM_DATA_SETS_COUNT:
@@ -380,8 +437,9 @@ def main():
         print('=' * 100)
 
     print(f'从数据集中提取了 {len(result_data_sets)} 条数据')
-    if len(result_data_sets) > 500:
-        data_sets = random.sample(result_data_sets, 10)
+    if len(result_data_sets) > MAX_DATA_COUNT:
+        print(f'筛选数据集数量大于 {MAX_DATA_COUNT}, 随机选择其中的 {RANDOM_DATA_COUNT} 条')
+        data_sets = random.sample(result_data_sets, RANDOM_DATA_COUNT)
     else:
         data_sets = result_data_sets
 
@@ -399,4 +457,4 @@ def main():
 
 
 if __name__ == "__main__":
-    main()
+    main()