update

main.py

@@ -1,10 +1,12 @@
+# -*- coding: utf-8 -*-
 import os
 import json
 import openai
+import httpx
 from datetime import datetime
 
 
-def load_config(config_file="config.json"):
+def load_config(config_file="ai_config.json"):
     try:
         with open(config_file, 'r', encoding='utf-8') as f:
             return json.load(f)
@@ -17,27 +19,37 @@ def load_config(config_file="config.json"):
 
 
 def create_prompt_file():
-    if not os.path.exists("prompt.txt"):
-        with open("prompt.txt", 'w', encoding='utf-8') as f:
+    if not os.path.exists("alpha_prompt.txt"):
+        with open("alpha_prompt.txt", 'w', encoding='utf-8') as f:
             f.write("")
-        print("已创建prompt.txt，请填入提示词后重新运行")
+        print("已创建alpha_prompt.txt，请填入提示词后重新运行")
         return False
     return True
 
 
 def read_prompt():
-    with open("prompt.txt", 'r', encoding='utf-8') as f:
+    with open("alpha_prompt.txt", 'r', encoding='utf-8') as f:
         prompt = f.read().strip()
     if not prompt:
-        print("prompt.txt是空的，请填入提示词")
+        print("alpha_prompt.txt是空的，请填入提示词")
         return None
     return prompt
 
+def read_operator():
+    with open("wqb_operator.txt", 'r', encoding='utf-8') as f:
+        operator = f.read().strip()
+    if not operator:
+        print("operator.txt是空的，请填入操作符")
+        return None
+    return operator
+
 
 def create_result_folder():
-    if not os.path.exists("ai_result"):
-        os.makedirs("ai_result")
-    return "ai_result"
+    # 修改这里：创建 generated_alpha 文件夹
+    folder_name = "generated_alpha"
+    if not os.path.exists(folder_name):
+        os.makedirs(folder_name)
+    return folder_name
 
 
 def call_siliconflow(api_key, prompt, model, base_url):
@@ -90,8 +102,27 @@ def get_ai_config(config):
 
     return service_name, service_config
 
+def get_user_info():
+    # 获取用户信息
+    with open('ai_config.json', 'r') as f:
+        config = json.load(f)
+        siliconflow = config['siliconflow']
+        token = siliconflow['api_keys']
+
+    headers = {"Authorization": f"Bearer {token}"}
+
+    url = "https://api.siliconflow.cn/v1/user/info"
+
+    response = httpx.get(url, headers=headers)
+
+    data = response.json()['data']
+    balance = data['balance']
+    print(f"余额: {balance}")
+
 
 def main():
+    get_user_info()
+
     config = load_config()
     if not config:
         return
@@ -107,6 +138,11 @@ def main():
     if not config_result:
         return
 
+    operator = read_operator()
+
+    if operator:
+        prompt = prompt + "\n\n以下是我的账号有权限使用的操作符以及操作符的使用方法, 请严格按照操作符, 进行生成,组合因子\n\n" + operator
+
     service_name, service_config = config_result
 
     print(f"使用服务: {service_name}")