From e75ed8b4ae6608f85020c2f1a2662befc651dd56 Mon Sep 17 00:00:00 2001 From: jack Date: Fri, 24 Apr 2026 18:56:48 +0800 Subject: [PATCH] ++ --- fetch-alphas/my_alpha_list.csv | 2 + fetch-alphas/simple_alpha_fetcher.py | 264 -- fetch-alphas/simple_wqb_get_alphas.py | 385 +++ fetch-alphas/wqb-get-alphas.py | 4 +- .../hide_low_sharpe_alpha.py | 4 +- simple72/Tranformer/Transformer.py | 1174 +++++--- simple72/Tranformer/ace_lib.py | 98 +- simple72/Tranformer/validator.py | 2526 +++++++++-------- simple72/config.json | 3 +- simple72/config_kimi.json | 16 + simple72/config_minimax.json | 16 + simple72/config_ollama.json | 16 + simple72/fix_log_levels.py | 46 + simple72/main.py | 493 +++- simple72/templates/app.js | 137 +- simple72/templates/index.html | 33 +- 16 files changed, 3094 insertions(+), 2123 deletions(-) create mode 100644 fetch-alphas/my_alpha_list.csv delete mode 100644 fetch-alphas/simple_alpha_fetcher.py create mode 100644 fetch-alphas/simple_wqb_get_alphas.py create mode 100644 simple72/config_kimi.json create mode 100644 simple72/config_minimax.json create mode 100644 simple72/config_ollama.json create mode 100644 simple72/fix_log_levels.py diff --git a/fetch-alphas/my_alpha_list.csv b/fetch-alphas/my_alpha_list.csv new file mode 100644 index 0000000..5f5a1cf --- /dev/null +++ b/fetch-alphas/my_alpha_list.csv @@ -0,0 +1,2 @@ +id,type,author,settings,regular,dateCreated,dateSubmitted,dateModified,name,favorite,hidden,color,category,tags,classifications,grade,stage,status,is,os,train,test,prod,competitions,themes,pyramids,pyramidThemes,team,osmosisPoints,origin +rKrVV75J,REGULAR,YC93384,"{'instrumentType': 'EQUITY', 'region': 'USA', 'universe': 'TOP3000', 'delay': 1, 'decay': 0, 'neutralization': 'SLOW_AND_FAST', 'truncation': 0.01, 'pasteurization': 'ON', 'unitHandling': 'VERIFY', 'nanHandling': 'OFF', 'maxTrade': 'ON', 'maxPosition': 'OFF', 'language': 'FASTEXPR', 'visualization': False, 'startDate': '2014-01-01', 'endDate': '2023-12-31'}","{'code': 'subtract(ts_mean(subtract(implied_volatility_call_120, implied_volatility_put_90), 60), reverse(ts_mean(subtract(implied_volatility_call_120, implied_volatility_put_90), 60)))', 'description': None, 'operatorCount': 6}",2026-04-23T11:31:50-04:00,,2026-04-23T11:31:50-04:00,,False,False,,,[],"[{'id': 'DATA_USAGE:SINGLE_DATA_SET', 'name': 'Single Data Set Alpha'}]",,IS,UNSUBMITTED,"{'pnl': 9443011, 'bookSize': 20000000, 'longCount': 1680, 'shortCount': 1449, 'turnover': 0.1125, 'returns': 0.0946, 'drawdown': 0.0431, 'margin': 0.001683, 'sharpe': 1.9, 'fitness': 1.65, 'startDate': '2014-01-01', 'checks': [{'name': 'LOW_SHARPE', 'result': 'PASS', 'limit': 1.58, 'value': 1.9}, {'name': 'LOW_FITNESS', 'result': 'PASS', 'limit': 1.0, 'value': 1.65}, {'name': 'LOW_TURNOVER', 'result': 'PASS', 'limit': 0.01, 'value': 0.1125}, {'name': 'HIGH_TURNOVER', 'result': 'PASS', 'limit': 0.7, 'value': 0.1125}, {'name': 'CONCENTRATED_WEIGHT', 'result': 'PASS'}, {'name': 'LOW_SUB_UNIVERSE_SHARPE', 'result': 'FAIL', 'limit': 0.82, 'value': 0.64}, {'name': 'SELF_CORRELATION', 'result': 'PENDING'}, {'name': 'DATA_DIVERSITY', 'result': 'PENDING'}, {'name': 'PROD_CORRELATION', 'result': 'PENDING'}, {'name': 'REGULAR_SUBMISSION', 'result': 'PENDING'}, {'name': 'HT_TURNOVER', 'result': 'WARNING', 'limit': 0.2, 'value': 0.1125}, {'name': 'HT_HIGH_TURNOVER_RETURNS_RATIO', 'result': 'WARNING', 'limit': 0.75, 'value': 0.3076}, {'name': 'HT_PNL_REALIZATION_HORIZON', 'result': 'WARNING', 'limit': 20, 'value': 76}, {'name': 'LOW_2Y_SHARPE', 'result': 'PASS', 'value': 1.95, 'limit': 1.58}, {'result': 'PASS', 'name': 'MATCHES_PYRAMID', 'effective': 1, 'multiplier': 1.3, 'pyramids': [{'name': 'USA/D1/OPTION', 'multiplier': 1.3}]}, {'result': 'WARNING', 'name': 'MATCHES_THEMES', 'themes': [{'id': 'KymLz14', 'multiplier': 1.0, 'name': 'All regions/D1 Power Pool Apr`26 2'}, {'id': 'lBqMzOB', 'multiplier': 3.0, 'name': 'Investable HTVR Theme'}]}, {'name': 'OSMOSIS_ALLOCATION', 'result': 'WARNING'}]}",,,,,,,,,,,PLATFORM diff --git a/fetch-alphas/simple_alpha_fetcher.py b/fetch-alphas/simple_alpha_fetcher.py deleted file mode 100644 index 355b6ee..0000000 --- a/fetch-alphas/simple_alpha_fetcher.py +++ /dev/null @@ -1,264 +0,0 @@ -import httpx -from httpx import BasicAuth, Timeout -import pandas as pd -from tqdm import tqdm -from pathlib import Path -import logging -import time -import random - - -class SimpleAlphaFetcher: - def __init__(self, base_path=None): - """ - 初始化 Alpha 获取器 - - Args: - base_path: 输出文件保存路径 - """ - self.client = None - self.base_path = Path(base_path) if base_path else Path.cwd() - self.logger = self._setup_logger() - - # 登录 - self.login() - - def _setup_logger(self): - """设置日志记录器""" - logger = logging.getLogger(__name__) - logger.setLevel(logging.INFO) - - if not logger.handlers: - handler = logging.StreamHandler() - formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') - handler.setFormatter(formatter) - logger.addHandler(handler) - - return logger - - def login(self): - """登录 WorldQuant Brain API""" - try: - # 从 nacos 获取账号密码 - with httpx.Client(timeout=10.0) as temp_client: - nacos_resp = temp_client.get( - 'http://192.168.31.41:30848/nacos/v1/cs/configs?dataId=wq_account&group=quantify' - ) - - if nacos_resp.status_code != 200: - self.logger.error('获取账号密码失败') - return False - - config = nacos_resp.json() - username = config.get('user_name') - password = config.get('password') - - if not username or not password: - self.logger.error('账号密码不完整') - return False - - self.logger.info(f"正在登录账户: {username}") - - # 创建客户端并设置超时 - timeout = Timeout(connect=30.0, read=60.0, write=30.0, pool=30.0) - self.client = httpx.Client( - auth=BasicAuth(username, password), - timeout=timeout - ) - - # 发送登录请求 - response = self.client.post('https://api.worldquantbrain.com/authentication') - - if response.status_code == 201: - self.logger.info("登录成功!") - return True - else: - self.logger.error(f"登录失败: {response.status_code} - {response.text}") - self.client.close() - self.client = None - return False - - except Exception as e: - self.logger.error(f"登录异常: {e}") - return False - - def fetch_alphas( - self, - max_pages: int = 100, - limit: int = 100, - delay: int = 1, - region: str = "USA", - universe: str = "TOP3000", - hidden: str = "false", - output_file_name: str = "alpha_list.csv", - mode: str = "w", - max_retries: int = 3, - ): - """ - 获取 Alpha 列表 - - Args: - max_pages: 最大爬取页数 - limit: 每页数量 - delay: API 请求的延迟设置 - region: 市场区域 - universe: 股票池 - hidden: 是否搜索已隐藏的 Alpha - output_file_name: 输出的 CSV 文件名 - mode: 写入模式："w"(覆盖写入) 或 "a"(追加写入) - max_retries: 请求失败时的最大重试次数 - - Returns: - list: Alpha 列表 - """ - if not self.client: - self.logger.error("客户端未登录，无法执行获取") - return [] - - fetched_alphas = [] - offset = 0 - total_accessed = 0 - pages_fetched = 0 - - # 先获取总数 - count_url = ( - f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&hidden={hidden}" - f"&limit=1&settings.delay={delay}&settings.region={region}&status=UNSUBMITTED%1FIS_FAIL&settings.universe={universe}" - ) - - total_available = 0 - for attempt in range(max_retries): - try: - count_response = self.client.get(count_url) - total_available = count_response.json()["count"] - break - except Exception as e: - self.logger.warning(f"获取 Alpha 总数失败 (尝试 {attempt+1}/{max_retries}): {e}") - if attempt < max_retries - 1: - time.sleep(random.uniform(3, 5)) - else: - self.logger.error("获取 Alpha 总数最终失败") - return [] - - if total_available == 0: - self.logger.warning("未找到任何 Alpha") - return [] - - # 计算实际最大页数 - actual_max_pages = min(max_pages, (total_available + limit - 1) // limit) - self.logger.info(f"共找到 {total_available} 个 Alpha，计划爬取 {actual_max_pages} 页...") - - pbar = tqdm(total=actual_max_pages, desc="爬取 Alpha 页面", unit="页") - - while pages_fetched < actual_max_pages: - # 构建请求 URL - url = ( - f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&limit={limit}" - f"&offset={offset}&settings.delay={delay}&settings.region={region}&hidden={hidden}&status=UNSUBMITTED%1FIS_FAIL&settings.universe={universe}" - ) - - try: - # 使用重试机制 - response = None - for attempt in range(max_retries): - try: - response = self.client.get(url) - break - except Exception as e: - self.logger.warning(f"请求失败 (尝试 {attempt+1}/{max_retries}): {e}") - if attempt < max_retries - 1: - time.sleep(random.uniform(3, 5)) - else: - raise - - if response.status_code == 400: - self.logger.warning(f"遇到 API 限制 (offset={offset})，停止获取") - break - - response_data = response.json() - - if not isinstance(response_data, dict) or "results" not in response_data: - self.logger.error(f"API 返回了意外的数据: {response_data}") - break - - alphas = response_data["results"] - - if not alphas: - break - - fetched_alphas.extend(alphas) - total_accessed += len(alphas) - pages_fetched += 1 - - # 更新进度条 - pbar.update(1) - pbar.set_postfix({ - "Region": region, - "已获取": total_accessed, - "本页": len(alphas), - }) - - if len(alphas) < limit: - self.logger.info(f"本页数据不足 {limit} 条，已到达末尾") - break - - offset += limit - - # 添加延迟避免请求过快 - time.sleep(0.5) - - except Exception as e: - self.logger.error(f"请求失败: {e}") - break - - pbar.close() - - if not fetched_alphas: - self.logger.warning("未获取到任何 Alpha！") - return [] - - # 按 fitness 排序 - df = pd.DataFrame(fetched_alphas) - df["temp_fitness"] = df.apply( - lambda row: row["is"].get("fitness", 0) if isinstance(row.get("is"), dict) else 0, - axis=1 - ) - df_sorted = df.sort_values(by="temp_fitness", ascending=False) - df_sorted = df_sorted.drop("temp_fitness", axis=1) - - output_path = self.base_path / output_file_name - - if mode == "w": - df_sorted.to_csv(output_path, index=False) - elif mode == "a": - df_sorted.to_csv(output_path, mode="a", index=False, header=False) - - self.logger.info(f"Alpha 列表已保存！共 {len(fetched_alphas)} 条记录！\n文件路径: {output_path}") - - return fetched_alphas - - def close(self): - """关闭客户端""" - if self.client: - self.client.close() - - -# 使用示例 -if __name__ == "__main__": - # 创建获取器实例（会自动登录） - fetcher = SimpleAlphaFetcher() - - if fetcher.client: - # 获取 Alpha 列表 - results = fetcher.fetch_alphas( - max_pages=100, # 最多爬取 100 页 - limit=100, # 每页 100 条 - region="USA", # 美股市场 - universe="TOP3000", # TOP3000 股票池 - output_file_name="alpha_list.csv" # 输出文件名 - ) - - print(f"成功获取 {len(results) if results else 0} 个 Alpha") - - # 关闭客户端 - fetcher.close() diff --git a/fetch-alphas/simple_wqb_get_alphas.py b/fetch-alphas/simple_wqb_get_alphas.py new file mode 100644 index 0000000..668ab07 --- /dev/null +++ b/fetch-alphas/simple_wqb_get_alphas.py @@ -0,0 +1,385 @@ +import httpx +from httpx import BasicAuth, Timeout +import pandas as pd +from tqdm import tqdm +from pathlib import Path +import logging +from typing import List, Optional +import time + +class AlphaManager: + def __init__(self, credentials_file='account.txt', base_path=None): + """ + 初始化Alpha管理器 + + Args: + credentials_file: 账号文件路径（备用） + base_path: 输出文件保存路径 + """ + self.client = None + self.base_path = Path(base_path) if base_path else Path.cwd() + self.logger = self._setup_logger() + + # 登录 + self.login(credentials_file) + + def _setup_logger(self): + """设置日志记录器""" + logger = logging.getLogger(__name__) + logger.setLevel(logging.INFO) + + if not logger.handlers: + handler = logging.StreamHandler() + formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') + handler.setFormatter(formatter) + logger.addHandler(handler) + + return logger + + def login(self, credentials_file=None): + """登录WorldQuant Brain API""" + try: + # 从nacos获取账号密码 + with httpx.Client(timeout=10.0) as temp_client: + nacos_resp = temp_client.get( + 'http://192.168.31.41:30848/nacos/v1/cs/configs?dataId=wq_account&group=quantify' + ) + + if nacos_resp.status_code != 200: + self.logger.error('获取账号密码失败') + return False + + config = nacos_resp.json() + username = config.get('user_name') + password = config.get('password') + + if not username or not password: + self.logger.error('账号密码不完整') + return False + + self.logger.info(f"正在登录账户: {username}") + + # 创建客户端并设置超时（关键修复） + # 设置更长的超时时间：连接30秒，读取60秒 + timeout = Timeout(connect=30.0, read=60.0, write=30.0, pool=30.0) + self.client = httpx.Client( + auth=BasicAuth(username, password), + timeout=timeout + ) + + # 发送登录请求 + response = self.client.post('https://api.worldquantbrain.com/authentication') + + if response.status_code == 201: + self.logger.info("登录成功!") + return True + else: + self.logger.error(f"登录失败: {response.status_code} - {response.text}") + self.client.close() + self.client = None + return False + + except Exception as e: + self.logger.error(f"登录异常: {e}") + return False + + def update_alpha_color(self, alpha_id: str, color: str) -> bool: + """标记Alpha颜色""" + if not self.client: + self.logger.error("客户端未登录") + return False + + try: + update_data = {"color": color} + response = self.client.patch( + f"https://api.worldquantbrain.com/alphas/{alpha_id}", + json=update_data + ) + return response.status_code == 200 + except Exception as e: + self.logger.error(f"标记颜色失败: {e}") + return False + + def wechat_check_corr_message(self, message: str): + """微信通知（可选功能）""" + self.logger.info(f"通知消息: {message}") + + def _make_request_with_retry(self, url: str, max_retries: int = 3, retry_delay: float = 2.0): + """ + 带重试机制的请求 + + Args: + url: 请求URL + max_retries: 最大重试次数 + retry_delay: 重试延迟（秒） + + Returns: + httpx.Response 或 None + """ + for attempt in range(max_retries): + try: + response = self.client.get(url) + return response + except Exception as e: + self.logger.warning(f"请求失败 (尝试 {attempt+1}/{max_retries}): {e}") + if attempt < max_retries - 1: + time.sleep(retry_delay) + else: + self.logger.error(f"请求最终失败: {url}") + raise + return None + + def get_alphas( + self, + total_alphas: int = 5, + limit: int = 100, + delay: int = 1, + region: str = "USA", + universe: str = "TOP3000", + required_fields: Optional[List[str]] = None, + match_mode: str = "all", + min_sharpe: Optional[float] = None, + min_fitness: Optional[float] = None, + hidden: str = "false", + submittable: bool = False, + auto_color: bool = False, + color: str = "GREEN", + output_file_name: str = "alpha_search_list.csv", + mode: str = "w", + max_retries: int = 3, + ): + """ + 搜索Alpha并筛选 + + Args: + total_alphas: 最多获取多少个Alpha + limit: 每次API请求获取的数量 + delay: API请求的延迟设置 + region: 市场区域，如 "USA", "CHINA" + universe: 股票池，如 "TOP3000" + required_fields: 关键词列表，如 ['put', 'call'] + match_mode: 关键词匹配模式："all"(全匹配) 或 "any"(任一匹配) + min_sharpe: 最小夏普比率阈值 + min_fitness: 最小适应度阈值 + hidden: 是否搜索已隐藏的Alpha ("true"/"false") + submittable: 是否只筛选可提交（无FAIL检查项）的Alpha + auto_color: 是否自动给符合条件的Alpha标记颜色 + color: 标记的颜色 + output_file_name: 输出的CSV文件名 + mode: 写入模式："w"(覆盖写入) 或 "a"(追加写入) + max_retries: 请求失败时的最大重试次数 + + Returns: + list: 符合条件的Alpha列表 + """ + if not self.client: + self.logger.error("客户端未登录，无法执行搜索") + return [] + + + + # 验证颜色参数 + valid_colors = [None, "GREEN", "YELLOW", "RED", "BLUE", "PURPLE", "ORANGE"] + if color not in valid_colors: + raise ValueError(f"颜色必须是以下之一: {valid_colors}") + + fetched_alphas = [] + offset = 0 + total_accessed = 0 + colored_count = 0 + + # 先获取总数（带重试） + count_url = ( + f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&hidden={hidden}" + f"&limit=1&settings.delay={delay}&settings.region={region}&status=UNSUBMITTED%1FIS_FAIL&settings.universe={universe}" + ) + + total_available = 0 + for attempt in range(max_retries): + try: + count_response = self.client.get(count_url) + total_available = count_response.json()["count"] + break + except Exception as e: + self.logger.warning(f"获取Alpha总数失败 (尝试 {attempt+1}/{max_retries}): {e}") + if attempt < max_retries - 1: + time.sleep(2.0) + else: + self.logger.error("获取Alpha总数最终失败") + return [] + + if total_available == 0: + self.logger.warning("未找到任何Alpha") + return [] + + self.logger.info(f"共找到 {total_available} 个Alpha，开始筛选...") + + pbar = tqdm(total=min(total_available, 10000), desc="扫描Alpha", unit="条") + + while len(fetched_alphas) < total_alphas and offset < total_available: + # 构建请求URL + url = ( + f"https://api.worldquantbrain.com/users/self/alphas?stage=IS&limit={limit}" + f"&offset={offset}&settings.delay={delay}&settings.region={region}&hidden={hidden}&status=UNSUBMITTED%1FIS_PASS&settings.universe={universe}" + ) + + try: + # 使用重试机制 + response = None + for attempt in range(max_retries): + try: + response = self.client.get(url) + break + except Exception as e: + self.logger.warning(f"请求失败 (尝试 {attempt+1}/{max_retries}): {e}") + if attempt < max_retries - 1: + time.sleep(2.0) + else: + raise + + if response.status_code == 400: + self.logger.warning(f"遇到API限制 (offset={offset})，停止获取更多数据") + break + + response_data = response.json() + + if not isinstance(response_data, dict) or "results" not in response_data: + self.logger.error(f"API返回了意外的数据: {response_data}") + break + + alphas = response_data["results"] + + if not alphas: + break + + total_accessed += len(alphas) + + # 关键词筛选 + if required_fields: + if match_mode == "all": + filtered_alphas = [ + alpha for alpha in alphas + if all(field in alpha["regular"]["code"] for field in required_fields) + ] + elif match_mode == "any": + filtered_alphas = [ + alpha for alpha in alphas + if any(field in alpha["regular"]["code"] for field in required_fields) + ] + else: + raise ValueError("match_mode 必须是 'all' 或 'any'") + else: + filtered_alphas = alphas + + # 进一步筛选夏普比率和fitness + final_filtered = [] + for alpha in filtered_alphas: + sharpe = alpha.get("is", {}).get("sharpe", 0) + fitness = alpha.get("is", {}).get("fitness", 0) + + sharpe_ok = (min_sharpe is None) or (sharpe is not None and abs(sharpe) >= min_sharpe) + fitness_ok = (min_fitness is None) or (fitness is not None and abs(fitness) >= min_fitness) + + if sharpe_ok and fitness_ok: + if submittable: + checks = alpha.get("is", {}).get("checks", []) + fail_count = sum(1 for check in checks if check.get("result") == "FAIL") + + if fail_count == 0 and auto_color: + alpha_id = alpha.get("id") + if alpha_id: + success = self.update_alpha_color(alpha_id, color) + if success: + colored_count += 1 + + if fail_count <= 1: + final_filtered.append(alpha) + else: + final_filtered.append(alpha) + + fetched_alphas.extend(final_filtered) + + # 更新进度条 + pbar.update(len(alphas)) + pbar.set_postfix({ + "Region": region, + "已扫描": total_accessed, + "找到": len(fetched_alphas), + "本次": len(final_filtered), + "标记": colored_count, + }) + + if len(alphas) < limit: + break + + offset += limit + + # 添加延迟避免请求过快 + time.sleep(0.5) + + except Exception as e: + self.logger.error(f"请求失败: {e}") + break + + pbar.close() + + # 输出标记统计 + if auto_color and colored_count > 0: + self.logger.info(f"共标记了 {colored_count} 个Alpha颜色为{color}") + self.wechat_check_corr_message(f"共标记了 {colored_count} 个Alpha颜色为{color}") + + alpha_list = fetched_alphas[:total_alphas] + + if not alpha_list: + self.logger.warning("未找到任何符合条件的Alpha！请检查筛选条件是否过于严格。") + else: + df = pd.DataFrame(alpha_list) + df["temp_fitness"] = df.apply( + lambda row: row["is"].get("fitness", 0) if isinstance(row.get("is"), dict) else 0, + axis=1 + ) + df_sorted = df.sort_values(by="temp_fitness", ascending=False) + df_sorted = df_sorted.drop("temp_fitness", axis=1) + + output_path = self.base_path / output_file_name + + if mode == "w": + df_sorted.to_csv(output_path, index=False) + elif mode == "a": + df_sorted.to_csv(output_path, mode="a", index=False, header=False) + + self.logger.info(f"批量回测初步检测结果已经下载！共{len(alpha_list)}条记录！\n{output_path}文件名保存！") + + return alpha_list + + def close(self): + """关闭客户端""" + if self.client: + self.client.close() + + +# 使用示例 +if __name__ == "__main__": + # 创建管理器实例（会自动登录） + manager = AlphaManager() + + if manager.client: + # 搜索包含"put"关键词的可提交Alpha + results = manager.get_alphas( + total_alphas=20, + region="", + universe="", + required_fields=['put'], + match_mode="all", + submittable=True, + auto_color=True, + min_sharpe = 1.0, + min_fitness = 0.6, + color="GREEN", + output_file_name="my_alpha_list.csv" + ) + + print(f"找到 {len(results) if results else 0} 个符合条件的Alpha") + + # 关闭客户端 + manager.close() \ No newline at end of file diff --git a/fetch-alphas/wqb-get-alphas.py b/fetch-alphas/wqb-get-alphas.py index a0dadd9..443baf8 100644 --- a/fetch-alphas/wqb-get-alphas.py +++ b/fetch-alphas/wqb-get-alphas.py @@ -135,7 +135,7 @@ class AlphaManager: limit: int = 100, delay: int = 1, region: str = "USA", - universe: str = "TOP3000", + universe: str = "TOP500", required_fields: Optional[List[str]] = None, match_mode: str = "all", min_sharpe: Optional[float] = None, @@ -371,6 +371,8 @@ if __name__ == "__main__": match_mode="all", submittable=True, auto_color=True, + min_sharpe = None, + min_fitness = None, color="GREEN", output_file_name="my_alpha_list.csv" ) diff --git a/hide_low_sharpe_alpha/hide_low_sharpe_alpha.py b/hide_low_sharpe_alpha/hide_low_sharpe_alpha.py index 9f09a00..a324580 100644 --- a/hide_low_sharpe_alpha/hide_low_sharpe_alpha.py +++ b/hide_low_sharpe_alpha/hide_low_sharpe_alpha.py @@ -17,7 +17,7 @@ MAX_RETRIES = 3 # 最大重试次数 RETRY_DELAY_MIN = 3 # 重试最小等待时间（秒） RETRY_DELAY_MAX = 5 # 重试最大等待时间（秒） MAX_LIMIT = 100 # 每页最大数量 -MAX_PAGE = 5 # 最大页数 +MAX_PAGE = 2 # 最大页数 def login(): @@ -185,7 +185,7 @@ def batch_hide_alphas(client, alphas): else: fail_count += 1 - time.sleep(0.5) + time.sleep(1) print("=" * 50) print(f"批量隐藏完成!") diff --git a/simple72/Tranformer/Transformer.py b/simple72/Tranformer/Transformer.py index e7aa27e..ff9e7a0 100755 --- a/simple72/Tranformer/Transformer.py +++ b/simple72/Tranformer/Transformer.py @@ -1,3 +1,12 @@ +import logging +import io +import getpass +import itertools +import threading +import time +from urllib.parse import urljoin +from pathlib import Path +import pandas as pd import requests import json import sys @@ -5,7 +14,7 @@ import asyncio import os import openai import re -from typing import Optional, Union # Added this import +from typing import Optional, Union # Added this import try: from .validator_hooks import is_valid_template_expr, has_empty_datafield_candidates except Exception: @@ -17,6 +26,8 @@ except Exception: has_empty_datafield_candidates = None # --- Validation wrappers to integrate into the pipeline --- + + def _filter_valid_templates( proposed_templates: dict, operators_meta, @@ -58,18 +69,22 @@ def _should_skip_due_to_empty_candidates(populated_info: dict) -> bool: return has_empty_datafield_candidates(populated_info) except Exception: return False -import logging -import pandas as pd -import os -from pathlib import Path -from urllib.parse import urljoin -import time -import threading -import itertools -import getpass -import io -import validator as val -from ace_lib import get_instrument_type_region_delay + + +logger = logging.getLogger(__name__) +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(logging.Formatter( + '%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + +try: + from . import validator as val + from .ace_lib import get_instrument_type_region_delay +except ImportError: + import validator as val + from ace_lib import get_instrument_type_region_delay # Force stdout/stderr to use utf-8 on Windows to avoid UnicodeEncodeError if sys.platform.startswith('win'): try: @@ -87,20 +102,20 @@ password = None DATA_CATEGORIES = None - # 加载模板总结文件 -template_summary_path = os.path.join(os.path.dirname(__file__), "template_summary.md") +template_summary_path = os.path.join( + os.path.dirname(__file__), "template_summary.md") try: with open(template_summary_path, "r", encoding="utf-8") as f: template_summary = f.read() - print(f"✓ 已加载模板总结文件: {template_summary_path}", flush=True) + logger.info(f"✓ 已加载模板总结文件: {template_summary_path}") except FileNotFoundError: - print(f"⚠ 模板总结文件不存在: {template_summary_path}，使用内置模板", flush=True) + logger.warning(f"⚠ 模板总结文件不存在: {template_summary_path}，使用内置模板") template_summary = """# BRAIN论坛Alpha模板精华总结 请创建 template_summary.md 文件""" except Exception as e: - print(f"⚠ 加载模板总结文件失败: {e}，使用内置模板", flush=True) + logger.error(f"⚠ 加载模板总结文件失败: {e}，使用内置模板") template_summary = """# BRAIN论坛Alpha模板精华总结 请检查 template_summary.md 文件""" @@ -127,14 +142,15 @@ class SingleSession(requests.Session): def get_relogin_lock(self): return self._relogin_lock + def load_template_summary(file_path: Optional[str] = None) -> str: """ Loads the template summary from a file or returns the built-in template summary. - + Args: file_path: Optional path to a .txt or .md file containing the template summary. If None or file doesn't exist, returns the built-in template summary. - + Returns: str: The template summary content. """ @@ -144,15 +160,15 @@ def load_template_summary(file_path: Optional[str] = None) -> str: if file_path_obj.exists() and file_path_obj.is_file(): with open(file_path_obj, 'r', encoding='utf-8') as f: content = f.read() - print(f"✓ 成功从文件加载模板总结: {file_path}", flush=True) + logger.info(f"✓ 成功从文件加载模板总结: {file_path}") return content else: - print(f"⚠ 警告: 文件不存在: {file_path}，将使用内置模板总结", flush=True) + logger.warning(f"⚠ 警告: 文件不存在: {file_path}，将使用内置模板总结") except Exception as e: - print(f"⚠ 警告: 读取文件时出错: {e}，将使用内置模板总结", flush=True) - + logger.warning(f"⚠ 警告: 读取文件时出错: {e}，将使用内置模板总结") + # 返回内置的模板总结 - print("✓ 使用内置模板总结", flush=True) + logger.info("✓ 使用内置模板总结") return template_summary @@ -174,110 +190,113 @@ def get_credentials() -> tuple[str, str]: # please input your own BRAIN Credentials into the function return (username, password) + def get_token_from_auth_server() -> str: # 声明使用全局变量 global LLM_API_KEY # please input your own LLM Gateway token into the function, please note, we are using kimi-k2.5 model return LLM_API_KEY + def interactive_input() -> dict: """ 交互式输入函数，收集所有必要的配置信息。 - + Returns: dict: 包含所有配置信息的字典 """ - print("\n" + "="*60, flush=True) - print("欢迎使用 Alpha Transformer 交互式配置", flush=True) - print("此程序在于让您输入一个Alpha ID即可通过历史总结的Alpha模板,转化成更多的表达式", flush=True) - print("72变,助您腾云驾雾", flush=True) - print("如果你想修改模型，则可以使用新模型的url和api key", flush=True) - print("不同模型效果不同，默认的kimi可能会产生语法错误，请检查生成的模板文件进行甄别", flush=True) - print("强烈推荐你使用自己总结的模板文档，效果会更好", flush=True) - print("="*60 + "\n", flush=True) - + logger.info("\n" + "="*60) + logger.info("欢迎使用 Alpha Transformer 交互式配置") + logger.info("此程序在于让您输入一个Alpha ID即可通过历史总结的Alpha模板,转化成更多的表达式") + logger.info("72变,助您腾云驾雾") + logger.info("如果你想修改模型，则可以使用新模型的url和api key") + logger.info("不同模型效果不同，默认的kimi可能会产生语法错误，请检查生成的模板文件进行甄别") + logger.info("强烈推荐你使用自己总结的模板文档，效果会更好") + logger.info("="*60 + "\n") + config = {} - + # 1. 询问 LLM 模型名称 - print("【1/6】LLM 模型配置", flush=True) - print("如果你想修改模型，则可以使用新模型的名称", flush=True) + logger.info("【1/6】LLM 模型配置") + logger.info("如果你想修改模型，则可以使用新模型的名称") default_model = "kimi-k2.5" model_input = input(f"请输入 LLM 模型名称 (直接回车使用默认值: {default_model}): ").strip() config['LLM_model_name'] = model_input if model_input else default_model - print(f"✓ LLM 模型名称: {config['LLM_model_name']}\n", flush=True) - + logger.info(f"✓ LLM 模型名称: {config['LLM_model_name']}\n") + # 2. 询问 LLM API Key - print("【2/6】LLM API Key 配置", flush=True) + logger.info("【2/6】LLM API Key 配置") api_key = getpass.getpass("请输入 LLM API Key (输入时不会显示): ").strip() if not api_key: - print("⚠ 警告: API Key 为空，程序可能无法正常工作", flush=True) + logger.warning("⚠ 警告: API Key 为空，程序可能无法正常工作") config['LLM_API_KEY'] = api_key - print("✓ API Key 已设置\n", flush=True) - + logger.info("✓ API Key 已设置\n") + # 3. 询问 LLM Base URL - print("【3/6】LLM Base URL 配置", flush=True) - print("提示：不同模型有不同的URL", flush=True) + logger.info("【3/6】LLM Base URL 配置") + logger.info("提示：不同模型有不同的URL") default_url = "https://api.moonshot.cn/v1" url_input = input(f"请输入 LLM Base URL (直接回车使用默认值: {default_url}): ").strip() config['llm_base_url'] = url_input if url_input else default_url - print(f"✓ LLM Base URL: {config['llm_base_url']}\n", flush=True) - + logger.info(f"✓ LLM Base URL: {config['llm_base_url']}\n") + # 4. 询问 BRAIN 平台用户名 - print("【4/6】BRAIN 平台认证信息", flush=True) + logger.info("【4/6】BRAIN 平台认证信息") username_input = input("请输入 BRAIN 平台用户名/邮箱: ").strip() if not username_input: - print("⚠ 警告: 用户名为空，程序可能无法正常工作", flush=True) + logger.warning("⚠ 警告: 用户名为空，程序可能无法正常工作") config['username'] = username_input - print("✓ 用户名已设置\n", flush=True) - + logger.info("✓ 用户名已设置\n") + # 5. 询问 BRAIN 平台密码 password_input = getpass.getpass("请输入 BRAIN 平台密码 (输入时不会显示): ").strip() if not password_input: - print("⚠ 警告: 密码为空，程序可能无法正常工作", flush=True) + logger.warning("⚠ 警告: 密码为空，程序可能无法正常工作") config['password'] = password_input - print("✓ 密码已设置\n", flush=True) - + logger.info("✓ 密码已设置\n") + # 6. 询问模板总结文件路径 - print("【5/6】模板总结文件配置", flush=True) - print("强烈推荐你使用自己总结的模板文档，效果会更好", flush=True) - print("提示: 如果您有 template_summary 的 .txt 或 .md 文件，请输入完整路径", flush=True) - print(" 如果没有，直接回车将使用内置模板总结", flush=True) + logger.info("【5/6】模板总结文件配置") + logger.info("强烈推荐你使用自己总结的模板文档，效果会更好") + logger.info("提示: 如果您有 template_summary 的 .txt 或 .md 文件，请输入完整路径") + logger.info(" 如果没有，直接回车将使用内置模板总结") template_path = input("请输入模板总结文件路径 (直接回车使用内置模板): ").strip() config['template_summary_path'] = template_path if template_path else None if template_path: - print(f"✓ 将尝试从文件加载: {template_path}\n", flush=True) + logger.info(f"✓ 将尝试从文件加载: {template_path}\n") else: - print("✓ 将使用内置模板总结\n", flush=True) - + logger.info("✓ 将使用内置模板总结\n") + # 7. 询问 Alpha ID - print("【6/7】Alpha ID 配置", flush=True) + logger.info("【6/7】Alpha ID 配置") alpha_id = input("请输入要处理的 Alpha ID: ").strip() if not alpha_id: - print("❌ 错误: Alpha ID 不能为空", flush=True) + logger.error("❌ 错误: Alpha ID 不能为空") sys.exit(1) config['alpha_id'] = alpha_id - print(f"✓ Alpha ID: {alpha_id}\n", flush=True) - + logger.info(f"✓ Alpha ID: {alpha_id}\n") + # 8. 询问 Top N 参数（仅数据字段） - print("【7/7】候选数量配置 (Top N)", flush=True) - print("提示: 此参数控制为每个占位符生成的数据字段候选数量", flush=True) - + logger.info("【7/7】候选数量配置 (Top N)") + logger.info("提示: 此参数控制为每个占位符生成的数据字段候选数量") + # Datafield top_n default_datafield_topn = 50 - datafield_topn_input = input(f"请输入数据字段候选数量 (直接回车使用默认值: {default_datafield_topn}): ").strip() + datafield_topn_input = input( + f"请输入数据字段候选数量 (直接回车使用默认值: {default_datafield_topn}): ").strip() try: - config['top_n_datafield'] = int(datafield_topn_input) if datafield_topn_input else default_datafield_topn + config['top_n_datafield'] = int( + datafield_topn_input) if datafield_topn_input else default_datafield_topn except ValueError: - print(f"⚠ 警告: 输入无效，使用默认值: {default_datafield_topn}", flush=True) + logger.warning(f"⚠ 警告: 输入无效，使用默认值: {default_datafield_topn}") config['top_n_datafield'] = default_datafield_topn - print(f"✓ 数据字段候选数量: {config['top_n_datafield']}\n", flush=True) - - print("="*60, flush=True) - print("配置完成！开始处理...", flush=True) - print("="*60 + "\n", flush=True) - - return config + logger.info(f"✓ 数据字段候选数量: {config['top_n_datafield']}\n") + logger.info("="*60) + logger.info("配置完成！开始处理...") + logger.info("="*60 + "\n") + + return config def expand_dict_columns(data: pd.DataFrame) -> pd.DataFrame: @@ -290,15 +309,18 @@ def expand_dict_columns(data: pd.DataFrame) -> pd.DataFrame: Returns: pandas.DataFrame: A new DataFrame with expanded columns. """ - dict_columns = list(filter(lambda x: isinstance(data[x].iloc[0], dict), data.columns)) + dict_columns = list(filter(lambda x: isinstance( + data[x].iloc[0], dict), data.columns)) new_columns = pd.concat( - [data[col].apply(pd.Series).rename(columns=lambda x: f"{col}_{x}") for col in dict_columns], + [data[col].apply(pd.Series).rename( + columns=lambda x: f"{col}_{x}") for col in dict_columns], axis=1, ) data = pd.concat([data, new_columns], axis=1) return data + def start_session() -> SingleSession: """ Start a new session with the WorldQuant BRAIN platform. @@ -316,13 +338,13 @@ def start_session() -> SingleSession: s = SingleSession() s.auth = get_credentials() r = s.post(brain_api_url + "/authentication") - print(f"New session created (ID: {id(s)}) with authentication response: {r.status_code}, {r.json()} (新会话已创建)", flush=True) + logger.info( + f"New session created (ID: {id(s)}) with authentication response: {r.status_code}, {r.json()} (新会话已创建)") if r.status_code == requests.status_codes.codes.unauthorized: if r.headers["WWW-Authenticate"] == "persona": - print( - "Complete biometrics authentication and press any key to continue (请完成生物识别认证并按任意键继续): \n" - + urljoin(r.url, r.headers["Location"]) - + "\n" + logger.info( + "Complete biometrics authentication and press any key to continue (请完成生物识别认证并按任意键继续): \n" + urljoin( + r.url, r.headers["Location"]) + "\n" ) input() s.post(urljoin(r.url, r.headers["Location"])) @@ -334,10 +356,11 @@ def start_session() -> SingleSession: else: break else: - print("\nIncorrect email or password (邮箱或密码错误)\n", flush=True) + logger.info("\nIncorrect email or password (邮箱或密码错误)\n") return start_session() return s + def get_data_categories(s: SingleSession) -> list[dict]: """ Fetch and cache data categories from the BRAIN API. @@ -345,7 +368,7 @@ def get_data_categories(s: SingleSession) -> list[dict]: global DATA_CATEGORIES if DATA_CATEGORIES is not None: return DATA_CATEGORIES - + try: brain_api_url = "https://api.worldquantbrain.com" response = s.get(brain_api_url + "/data-categories") @@ -359,9 +382,10 @@ def get_data_categories(s: SingleSession) -> list[dict]: DATA_CATEGORIES = [] return DATA_CATEGORIES except Exception as e: - print(f"Error fetching data categories: {e}", flush=True) + logger.info(f"Error fetching data categories: {e}") return [] + def get_datafields( s: SingleSession, instrument_type: str = "EQUITY", @@ -394,20 +418,20 @@ def get_datafields( """ brain_api_url = "https://api.worldquantbrain.com" type_param = f"&type={data_type}" if data_type != "ALL" else "" - + url_template = ( brain_api_url + "/data-fields?" + f"&instrumentType={instrument_type}" + f"®ion={region}&delay={str(delay)}&universe={universe}{type_param}&limit=50" ) - + if dataset_id: url_template += f"&dataset.id={dataset_id}" - + if len(search) > 0: url_template += f"&search={search}" - + url_template += "&offset={x}" count = 0 @@ -415,18 +439,18 @@ def get_datafields( try: count = s.get(url_template.format(x=0)).json()["count"] except Exception as e: - print(f"Error getting count: {e}", flush=True) + logger.info(f"Error getting count: {e}") return pd.DataFrame() - + if count == 0: - print( + logger.info( f"No fields found (未找到字段): region={region}, delay={str(delay)}, universe={universe}, " f"type={data_type}, dataset.id={dataset_id}" ) return pd.DataFrame() else: if category: - count = 500 # Search deeper if filtering + count = 500 # Search deeper if filtering else: count = 100 @@ -440,7 +464,7 @@ def get_datafields( try: resp = s.get(url_template.format(x=x)) while resp.status_code == 429: - print("status_code 429, sleep 3 seconds", flush=True) + logger.info("status_code 429, sleep 3 seconds") time.sleep(3) resp = s.get(url_template.format(x=x)) if resp.status_code == 200 and "results" in resp.json(): @@ -455,16 +479,16 @@ def get_datafields( results = datafields.json().get("results", []) if not results: break - + if category: if isinstance(category, list): filtered_results = [ - item for item in results + item for item in results if isinstance(item.get('category'), dict) and item['category'].get('id') in category ] else: filtered_results = [ - item for item in results + item for item in results if isinstance(item.get('category'), dict) and item['category'].get('id') == category ] datafields_list.append(filtered_results) @@ -474,8 +498,9 @@ def get_datafields( else: datafields_list.append(results) - datafields_list_flat = [item for sublist in datafields_list for item in sublist] - + datafields_list_flat = [ + item for sublist in datafields_list for item in sublist] + if not datafields_list_flat: return pd.DataFrame() @@ -483,6 +508,7 @@ def get_datafields( datafields_df = expand_dict_columns(datafields_df) return datafields_df + def set_alpha_properties( s: SingleSession, alpha_id: str, @@ -523,7 +549,7 @@ def set_alpha_properties( params.setdefault("selection", {})["description"] = selection_desc if combo_desc != "None": # Assuming "None" is the default string value for combo_desc params.setdefault("combo", {})["description"] = combo_desc - + response = s.patch(brain_api_url + "/alphas/" + alpha_id, json=params) return response @@ -537,6 +563,7 @@ def extract_placeholders(template_expression: str) -> list[str]: # Only match placeholders of the form `` or `` with alphanumeric and underscores return re.findall(r'(<[A-Za-z0-9_]+/>)', template_expression) + def parse_alpha_code(alpha_code: str, all_operators: list[dict]) -> tuple[list[str], list[str]]: """ Parses the alpha code to extract operators and data fields. @@ -547,7 +574,7 @@ def parse_alpha_code(alpha_code: str, all_operators: list[dict]) -> tuple[list[s alpha_code = re.sub(r"#.*", "", alpha_code) operators_names = [op['name'] for op in all_operators] - + found_operators = [] found_datafields = [] @@ -559,15 +586,16 @@ def parse_alpha_code(alpha_code: str, all_operators: list[dict]) -> tuple[list[s for identifier in identifiers: if identifier in operators_names: found_operators.append(identifier) - elif not (identifier.isdigit() or identifier.lower() in ['true', 'false', 'null', 'nan', 'if', 'else', 'for', 'while', 'return', 'and', 'or', 'not', 'in', 'is', 'try', 'except', 'finally', 'with', 'as', 'def', 'class', 'import', 'from', 'yield', 'lambda', 'global', 'nonlocal', 'break', 'continue', 'pass', 'async', 'await', 'raise', 'assert', 'del', 'print', 'input', 'len', 'min', 'max', 'sum', 'abs', 'round', 'int', 'float', 'str', 'list', 'dict', 'set', 'tuple', 'range', 'map', 'filter', 'zip', 'open', 'file', 'type', 'id', 'dir', 'help', 'object', 'super', 'issubclass', 'isinstance', 'hasattr', 'getattr', 'setattr', 'delattr', '__import__', 'None', 'True', 'False']): + elif not (identifier.isdigit() or identifier.lower() in ['true', 'false', 'null', 'nan', 'if', 'else', 'for', 'while', 'return', 'and', 'or', 'not', 'in', 'is', 'try', 'except', 'finally', 'with', 'as', 'def', 'class', 'import', 'from', 'yield', 'lambda', 'global', 'nonlocal', 'break', 'continue', 'pass', 'async', 'await', 'raise', 'assert', 'del', 'logger.info', 'input', 'len', 'min', 'max', 'sum', 'abs', 'round', 'int', 'float', 'str', 'list', 'dict', 'set', 'tuple', 'range', 'map', 'filter', 'zip', 'open', 'file', 'type', 'id', 'dir', 'help', 'object', 'super', 'issubclass', 'isinstance', 'hasattr', 'getattr', 'setattr', 'delattr', '__import__', 'None', 'True', 'False']): found_datafields.append(identifier) - + # Remove duplicates found_operators = list(set(found_operators)) found_datafields = list(set(found_datafields)) return found_operators, found_datafields + async def generate_alpha_description(alpha_id: str, brain_session: SingleSession) -> str: """ Generates and potentially enriches the description of a given Alpha ID from the WorldQuant BRAIN API. @@ -585,26 +613,29 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession async def call_llm_new(prompt: str) -> dict: # 声明使用全局变量 global LLM_model_name, LLM_API_KEY, llm_base_url - print(f"\n[call_llm_new] 准备调用 LLM...", flush=True) - print(f"[call_llm_new] 模型: {LLM_model_name}", flush=True) - print(f"[call_llm_new] llm_base_url: {llm_base_url}", flush=True) - + logger.info(f"\n[call_llm_new] 准备调用 LLM...") + logger.info(f"[call_llm_new] 模型: {LLM_model_name}") + logger.info(f"[call_llm_new] llm_base_url: {llm_base_url}") + try: - print(f"[call_llm_new] 正在获取 LLM token...", flush=True) + logger.info(f"[call_llm_new] 正在获取 LLM token...") llm_api_key = get_token_from_auth_server() llm_base_url_value = llm_base_url # 使用全局变量 - print(f"[call_llm_new] 创建 LLM 客户端...", flush=True) - llm_client = openai.AsyncOpenAI(base_url=llm_base_url_value, api_key=llm_api_key) - print("[call_llm_new] LLM Gateway Authentication successful. (LLM网关认证成功)", flush=True) + logger.info(f"[call_llm_new] 创建 LLM 客户端...") + llm_client = openai.AsyncOpenAI( + base_url=llm_base_url_value, api_key=llm_api_key) + logger.info( + "[call_llm_new] LLM Gateway Authentication successful. (LLM网关认证成功)") except Exception as e: - print(f"[call_llm_new] ❌ LLM Gateway Authentication failed (LLM网关认证失败): {e}", flush=True) + logger.error( + f"[call_llm_new] ❌ LLM Gateway Authentication failed (LLM网关认证失败): {e}") sys.exit(1) - print("[call_llm_new] --- Calling LLM... (正在调用LLM...) ---", flush=True) - print(f"[call_llm_new] Prompt 长度: {len(prompt)} 字符", flush=True) + logger.info("[call_llm_new] --- Calling LLM... (正在调用LLM...) ---") + logger.info(f"[call_llm_new] Prompt 长度: {len(prompt)} 字符") try: # Await the async create call - print(f"[call_llm_new] 正在发送请求到 LLM...", flush=True) + logger.info(f"[call_llm_new] 正在发送请求到 LLM...") response = await llm_client.chat.completions.create( model=LLM_model_name, messages=[ @@ -613,7 +644,7 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession ], # response_format={"type": "json_object"}, ) - print(f"[call_llm_new] 收到 LLM 响应", flush=True) + logger.info(f"[call_llm_new] 收到 LLM 响应") # The async client may return a nested structure. Try to extract content robustly. content = None @@ -623,7 +654,8 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession choices = response.get('choices') if choices and isinstance(choices, list): msg = choices[0].get('message') or choices[0] - content = msg.get('content') if isinstance(msg, dict) else None + content = msg.get('content') if isinstance( + msg, dict) else None elif 'content' in response: content = response.get('content') else: @@ -647,46 +679,50 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession # Return wrapped string if not JSON return {"text": content} - print(f"[call_llm_new] ✓ 成功返回结果", flush=True) + logger.info(f"[call_llm_new] ✓ 成功返回结果") return {} except Exception as e: - print(f"[call_llm_new] ❌ Error calling LLM (调用LLM出错): {e}", flush=True) + logger.error(f"[call_llm_new] ❌ Error calling LLM (调用LLM出错): {e}") import traceback - print(f"[call_llm_new] 错误详情: {traceback.format_exc()}", flush=True) + logger.error(f"[call_llm_new] 错误详情: {traceback.format_exc()}") return {} - print(f"\n[Alpha Description] 开始获取 Alpha {alpha_id} 的详情...", flush=True) - + logger.info(f"\n[Alpha Description] 开始获取 Alpha {alpha_id} 的详情...") + try: brain_api_url = "https://api.worldquantbrain.com" alpha_url = f"{brain_api_url}/alphas/{alpha_id}" - print(f"[Alpha Description] 请求 URL: {alpha_url}", flush=True) + logger.info(f"[Alpha Description] 请求 URL: {alpha_url}") response = brain_session.get(alpha_url) - print(f"[Alpha Description] 响应状态码: {response.status_code}", flush=True) + logger.info(f"[Alpha Description] 响应状态码: {response.status_code}") response.raise_for_status() # Raise an exception for HTTP errors alpha_data = response.json() - print(f"[Alpha Description] 成功获取 Alpha 数据", flush=True) + logger.info(f"[Alpha Description] 成功获取 Alpha 数据") settings = alpha_data.get('settings', {}) - expression_dict = alpha_data.get('regular', alpha_data.get('combo', None)) - + expression_dict = alpha_data.get( + 'regular', alpha_data.get('combo', None)) + if not expression_dict or 'code' not in expression_dict: - print(f"Error: Alpha expression code not found for Alpha ID (未找到Alpha表达式代码): {alpha_id}", flush=True) + logger.info( + f"Error: Alpha expression code not found for Alpha ID (未找到Alpha表达式代码): {alpha_id}") return json.dumps({}) alpha_code = expression_dict['code'] - current_description = expression_dict.get('description', '') + current_description = expression_dict.get('description', ) # 1. Get all operators for parsing (no filter as per feedback) operators_data = get_brain_operators() all_operators = operators_data.get('operators', []) - + # 2. Parse the code to get operators and datafields - found_operators_names, found_datafields_names = parse_alpha_code(alpha_code, all_operators) + found_operators_names, found_datafields_names = parse_alpha_code( + alpha_code, all_operators) # 3. Get descriptions for operators - operator_descriptions = {op['name']: op.get('description', 'No description available.') for op in all_operators if op['name'] in found_operators_names} - + operator_descriptions = {op['name']: op.get( + 'description', 'No description available.') for op in all_operators if op['name'] in found_operators_names} + # 4. Get descriptions for datafields datafield_descriptions = {} if found_datafields_names: @@ -698,13 +734,15 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession for df_name in found_datafields_names: # get_datafields returns a DataFrame, so we need to process it - datafield_df = get_datafields(s=brain_session, instrument_type=instrument_type, region=region, delay=delay, universe=universe, search=df_name) + datafield_df = get_datafields(s=brain_session, instrument_type=instrument_type, + region=region, delay=delay, universe=universe, search=df_name) if not datafield_df.empty: # Assuming the first result is the most relevant - datafield_descriptions[df_name] = datafield_df.iloc[0].get('description', 'No description available.') + datafield_descriptions[df_name] = datafield_df.iloc[0].get( + 'description', 'No description available.') else: datafield_descriptions[df_name] = 'No description found.' - + # 5. Use LLM to judge if current description is good judgment_prompt = f""" Given the following alpha code, its current description, and descriptions of its operators and datafields: @@ -727,9 +765,10 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession Is the current description good enough? Respond with 'yes' or 'no' in a JSON object: {{"judgment": "yes/no"}} A "good" description should clearly explain the investment idea, rationale for data used, and rationale for operators used. """ - + judgment_response = await call_llm_new(judgment_prompt) - is_description_good = judgment_response.get("judgment", "no").lower() == "yes" + is_description_good = judgment_response.get( + "judgment", "no").lower() == "yes" new_description = current_description if not is_description_good: @@ -752,11 +791,13 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession Alpha Settings: {json.dumps(settings, indent=2)} """ - + generated_description_response = await call_llm_new(generation_prompt) # Assuming LLM returns a string directly or a JSON with a 'description' key - new_description = generated_description_response.get("description", generated_description_response) - if isinstance(new_description, dict): # Handle cases where LLM might return a dict directly + new_description = generated_description_response.get( + "description", generated_description_response) + # Handle cases where LLM might return a dict directly + if isinstance(new_description, dict): new_description = json.dumps(new_description, indent=2) # 7. Override this new description and patch the alpha @@ -765,8 +806,9 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession alpha_id=alpha_id, regular_desc=new_description ) - print(f"Alpha {alpha_id} description updated on platform. (Alpha描述已在平台更新)", flush=True) - + logger.info( + f"Alpha {alpha_id} description updated on platform. (Alpha描述已在平台更新)") + if 'regular' in alpha_data: alpha_data['regular']['description'] = new_description elif 'combo' in alpha_data: @@ -778,15 +820,17 @@ async def generate_alpha_description(alpha_id: str, brain_session: SingleSession }) except requests.exceptions.RequestException as e: - print(f"Error during API request (API请求出错): {e}", flush=True) + logger.info(f"Error during API request (API请求出错): {e}") return json.dumps({}) except json.JSONDecodeError: - print("Error: Could not decode JSON response from API. (无法解析API的JSON响应)", flush=True) + logger.info( + "Error: Could not decode JSON response from API. (无法解析API的JSON响应)") return json.dumps({}) except Exception as e: - print(f"An unexpected error occurred (发生意外错误): {e}", flush=True) + logger.info(f"An unexpected error occurred (发生意外错误): {e}") return json.dumps({}) + def get_brain_operators(scope_filters: Optional[list[str]] = None) -> dict: """ Retrieves the list of available operators from the WorldQuant BRAIN API, @@ -808,9 +852,10 @@ def get_brain_operators(scope_filters: Optional[list[str]] = None) -> dict: response.raise_for_status() # Raise an exception for HTTP errors operators_list = response.json() - + if not isinstance(operators_list, list): - print(f"Error: Expected a list of operators, but received type (预期运算符列表，但收到类型): {type(operators_list)}", flush=True) + logger.info( + f"Error: Expected a list of operators, but received type (预期运算符列表，但收到类型): {type(operators_list)}") return {} if scope_filters: @@ -829,15 +874,19 @@ def get_brain_operators(scope_filters: Optional[list[str]] = None) -> dict: } except requests.exceptions.RequestException as e: - print(f"Error during API request for operators (获取运算符时API请求出错): {e}", flush=True) + logger.info( + f"Error during API request for operators (获取运算符时API请求出错): {e}") return {} except json.JSONDecodeError: - print("Error: Could not decode JSON response from operators API. (无法解析运算符API的JSON响应)", flush=True) + logger.info( + "Error: Could not decode JSON response from operators API. (无法解析运算符API的JSON响应)") return {} except Exception as e: - print(f"An unexpected error occurred while getting operators (获取运算符时发生意外错误): {e}", flush=True) + logger.info( + f"An unexpected error occurred while getting operators (获取运算符时发生意外错误): {e}") return {} + async def call_llm(prompt: str, llm_client: openai.AsyncOpenAI, max_retries: int = 3) -> dict: """ Interface with a Large Language Model to process prompts and get a JSON response. @@ -846,17 +895,18 @@ async def call_llm(prompt: str, llm_client: openai.AsyncOpenAI, max_retries: int # 声明使用全局变量 global LLM_model_name if not llm_client: - print("LLM client not initialized. Please check authentication. (LLM客户端未初始化，请检查认证)", flush=True) + logger.info( + "LLM client not initialized. Please check authentication. (LLM客户端未初始化，请检查认证)") return {} - - print("\n[LLM Call] 准备调用 LLM API...", flush=True) - print(f"[LLM Call] 模型: {LLM_model_name}", flush=True) - print(f"[LLM Call] Prompt 长度: {len(prompt)} 字符", flush=True) - print("[LLM Call] 正在发送请求...", flush=True) - + + logger.info("\n[LLM Call] 准备调用 LLM API...") + logger.info(f"[LLM Call] 模型: {LLM_model_name}") + logger.info(f"[LLM Call] Prompt 长度: {len(prompt)} 字符") + logger.info("[LLM Call] 正在发送请求...") + for attempt in range(max_retries): try: - print(f"[LLM Call] 第 {attempt + 1} 次尝试...", flush=True) + logger.info(f"[LLM Call] 第 {attempt + 1} 次尝试...") response = await llm_client.chat.completions.create( model=LLM_model_name, # Or your preferred model messages=[ @@ -865,38 +915,68 @@ async def call_llm(prompt: str, llm_client: openai.AsyncOpenAI, max_retries: int ], # response_format={"type": "json_object"}, ) - print(f"[LLM Call] 收到响应，状态: OK", flush=True) + logger.info(f"[LLM Call] 收到响应，状态: OK") content = response.choices[0].message.content - print(f"[LLM Call] 响应内容长度: {len(content)} 字符", flush=True) + logger.info(f"[LLM Call] 响应内容长度: {len(content)} 字符") + + # with open(f'llm_resp{str(int(time.time()))}.md', 'w', encoding='utf-8') as f: + # f.write(content) + + # 增强的 JSON 清洗逻辑 + original_content = content + + # 首先移除 ... 标签及其内容 + if "" in content and "" in content: + think_start = content.find("") + think_end = content.find("") + len("") + content = content[:think_start] + content[think_end:] + logger.info(f"[LLM Call] 移除了 标签") + + content = content.strip() # Try to clean markdown code blocks if present if "```json" in content: content = content.split("```json")[1].split("```")[0].strip() - print(f"[LLM Call] 清理了 JSON markdown 标记", flush=True) + logger.info(f"[LLM Call] 清理了 JSON markdown 标记") elif "```" in content: content = content.split("```")[1].split("```")[0].strip() - print(f"[LLM Call] 清理了 markdown 标记", flush=True) - - print(f"[LLM Call] 解析 JSON...", flush=True) + logger.info(f"[LLM Call] 清理了 markdown 标记") + else: + # 如果 markdown 清理失败，尝试通过大括号提取 JSON + logger.info(f"[LLM Call] 未找到 markdown 标记，尝试大括号提取...") + start_idx = content.find('{') + end_idx = content.rfind('}') + if start_idx != -1 and end_idx != -1 and start_idx < end_idx: + content = content[start_idx:end_idx+1] + logger.info(f"[LLM Call] 通过大括号提取 JSON，长度: {len(content)} 字符") + + logger.info(f"[LLM Call] 解析 JSON...") result = json.loads(content) - print(f"[LLM Call] JSON 解析成功，返回 {len(result)} 个结果", flush=True) + logger.info(f"[LLM Call] JSON 解析成功，返回 {len(result)} 个结果") return result except json.JSONDecodeError as e: - print(f"⚠ JSON Decode Error (Attempt {attempt + 1}/{max_retries}): {e}", flush=True) + logger.info(f"start *" * 88) + logger.warning( + f"⚠ JSON Decode Error (Attempt {attempt + 1}/{max_retries}): {e}\n\n{content}\n\n") + logger.info(f"end *" * 88) if attempt == max_retries - 1: - print(f"❌ Failed to parse JSON after {max_retries} attempts. Raw content: {content[:100]}...", flush=True) + logger.error( + f"❌ Failed to parse JSON after {max_retries} attempts. Raw content: {content[:100]}...") except Exception as e: - print(f"⚠ LLM Call Error (Attempt {attempt + 1}/{max_retries}): {e}", flush=True) + logger.warning( + f"⚠ LLM Call Error (Attempt {attempt + 1}/{max_retries}): {e}") if attempt == max_retries - 1: - print(f"❌ Failed to call LLM after {max_retries} attempts.", flush=True) + logger.error( + f"❌ Failed to call LLM after {max_retries} attempts.") raise Exception(f"LLM 调用失败: {e}") - + # Wait before retrying (2 seconds for MiniMax 529) - print(f"⏳ 等待 2 秒后重试...", flush=True) + logger.info(f"⏳ 等待 2 秒后重试...") await asyncio.sleep(2) - + return {} + def has_valid_placeholders(template_str: str) -> bool: """检查模板字符串是否包含有效的占位符""" import re @@ -909,9 +989,9 @@ async def propose_alpha_templates_with_retry(alpha_details: dict, template_summa 使用重试机制生成 Alpha 模板，确保包含占位符 """ if not alpha_details.get('expression'): - print("Error: Alpha expression is missing. (错误：缺少Alpha表达式)", flush=True) + logger.info("Error: Alpha expression is missing. (错误：缺少Alpha表达式)") return {} - + data_type_instruction = "" if user_data_type == "MATRIX": data_type_instruction = "\n**Important Note on Data Type:**\nThe user has specified the data type as **MATRIX**. Please do NOT use any vector-type operators (e.g., `vec_avg`, `vec_sum`) in your proposed templates, as they will raise errors for MATRIX type data in BRAIN. Note: 'MATRIX' is just a system identifier and does not refer to mathematical matrices." @@ -951,58 +1031,59 @@ Based on the structure and potential economic rationale of the seed alpha, by th Now, generate the JSON object with your proposed templates. Remember: USE PLACEHOLDERS like ``, NOT actual field names! """ - print(f"\n[Step 1/5] 正在调用 LLM 生成 Alpha 模板...", flush=True) - print(f" - 模型: {LLM_model_name}", flush=True) - print(f" - 数据类型: {user_data_type}", flush=True) - print(f" - 最大重试次数: {max_retries}", flush=True) + logger.info(f"\n[Step 1/5] 正在调用 LLM 生成 Alpha 模板...") + logger.info(f" - 模型: {LLM_model_name}") + logger.info(f" - 数据类型: {user_data_type}") + logger.info(f" - 最大重试次数: {max_retries}") alpha_expr = alpha_details.get('expression', {}) if isinstance(alpha_expr, dict): alpha_expr = alpha_expr.get('code', 'N/A') - print(f" - 种子 Alpha: {str(alpha_expr)[:50]}...", flush=True) - + logger.info(f" - 种子 Alpha: {str(alpha_expr)[:50]}...") + # 重试机制 for attempt in range(1, max_retries + 1): try: - print(f"\n [尝试 {attempt}/{max_retries}] 调用 LLM...", flush=True) + logger.info(f"\n [尝试 {attempt}/{max_retries}] 调用 LLM...") proposed_templates = await call_llm(prompt, llm_client) - + # 验证是否包含占位符 valid_templates = {} invalid_templates = [] - + for template_expr, explanation in proposed_templates.items(): if has_valid_placeholders(template_expr): valid_templates[template_expr] = explanation else: invalid_templates.append(template_expr) - + if valid_templates: - print(f" ✓ 成功生成 {len(valid_templates)} 个有效模板（含占位符）", flush=True) + logger.info(f" ✓ 成功生成 {len(valid_templates)} 个有效模板（含占位符）") if invalid_templates: - print(f" ⚠ 丢弃 {len(invalid_templates)} 个无效模板（无占位符）", flush=True) + logger.warning( + f" ⚠ 丢弃 {len(invalid_templates)} 个无效模板（无占位符）") return valid_templates else: - print(f" ✗ 所有模板均无占位符，需要重试", flush=True) + logger.warning(f" ✗ 所有模板均无占位符，需要重试") if invalid_templates: - print(f" 无效模板示例: {invalid_templates[0][:80]}...", flush=True) - + logger.info(f" 无效模板示例: {invalid_templates[0][:80]}...") + if attempt < max_retries: - print(f" ↻ 等待重试...", flush=True) + logger.info(f" ↻ 等待重试...") await asyncio.sleep(1) # 短暂延迟避免请求过快 - + except Exception as e: - print(f" ✗ 调用 LLM 时发生错误: {e}", flush=True) + logger.error(f" ✗ 调用 LLM 时发生错误: {e}") # 529 是 MiniMax 的特色，继续重试 if "overloaded" in str(e) or "529" in str(e): - print(f" ⚠ MiniMax 529 错误，继续重试...", flush=True) + logger.warning(f" ⚠ MiniMax 529 错误，继续重试...") if attempt < max_retries: - print(f" ↻ 等待 2 秒后重试...", flush=True) + logger.info(f" ↻ 等待 2 秒后重试...") await asyncio.sleep(2) - + # 20次都失败了 - print(f"\n⚠⚠⚠ 警告: 经过 {max_retries} 次重试，仍未能生成包含占位符的模板！", flush=True) - print(f" 可能原因: LLM 未遵循指令，或模型不支持此格式。", flush=True) - print(f" 建议: 检查 LLM 模型是否正确，或手动修改 prompt。", flush=True) + logger.warning(f"\n⚠⚠⚠ 警告: 经过 {max_retries} 次重试，仍未能生成包含占位符的模板！") + logger.info(f" 可能原因: LLM 未遵循指令，或模型不支持此格式。") + logger.info(f" 建议: 检查 LLM 模型是否正确，或手动修改 prompt。") return {} @@ -1013,6 +1094,7 @@ async def propose_alpha_templates(alpha_details: dict, template_summary: str, ll """ return await propose_alpha_templates_with_retry(alpha_details, template_summary, llm_client, user_data_type, max_retries) + async def propose_datafield_keywords(template_expression: str, template_explanation: str, placeholder: str, llm_client: openai.AsyncOpenAI, user_category: Optional[Union[str, list]] = None) -> list[str]: """ Uses an LLM to propose search keywords for finding data fields. @@ -1049,9 +1131,10 @@ Explanation: "measures the time-series evolution of a fund's relative rank on a Now, generate the JSON array of search keywords for the given placeholder. """ - print(f"--- Calling LLM to get keywords for placeholder (正在调用LLM获取占位符关键词): {placeholder} ---", flush=True) + logger.info( + f"--- Calling LLM to get keywords for placeholder (正在调用LLM获取占位符关键词): {placeholder} ---") response = await call_llm(prompt, llm_client) - print(f"AI使用如下提示词获取搜索关键词推荐：{prompt}", flush=True) + logger.info(f"AI使用如下提示词获取搜索关键词推荐：{prompt}") # Accept either a direct list or a dict containing a 'keywords' key if isinstance(response, list) and all(isinstance(item, str) for item in response): return response @@ -1060,97 +1143,147 @@ Explanation: "measures the time-series evolution of a fund's relative rank on a for key in ('keywords', 'data', 'result', 'items'): if key in response and isinstance(response[key], list) and all(isinstance(i, str) for i in response[key]): return response[key] - print(f"Warning: LLM did not return a valid list of strings for keywords (警告：LLM未返回有效的关键词列表). Got: {response}", flush=True) + logger.warning( + f"Warning: LLM did not return a valid list of strings for keywords (警告：LLM未返回有效的关键词列表). Got: {response}") return [] + async def get_datafield_candidates(s: SingleSession, alpha_details: dict, template_expression: str, template_explanation: str, placeholder: str, llm_client: openai.AsyncOpenAI, top_n: int = 50, user_region: Optional[str] = None, user_universe: Optional[str] = None, user_delay: Optional[int] = None, user_category: Optional[Union[str, list]] = None, user_data_type: str = "MATRIX") -> list[dict]: """ Gets candidate data fields for a placeholder by using an LLM to generate search keywords - and then calling the BRAIN API's get_datafields to retrieve the top N results for each keyword. + and then searching in local cache file. """ keywords = await propose_datafield_keywords(template_expression, template_explanation, placeholder, llm_client, user_category=user_category) if not keywords: - print(f"Could not generate keywords for placeholder (无法生成占位符关键词): {placeholder}", flush=True) + logger.info( + f"Could not generate keywords for placeholder (无法生成占位符关键词): {placeholder}") return [] - print(f"LLM-proposed keywords for '{placeholder}' (LLM提议的关键词): {keywords}", flush=True) + logger.info( + f"LLM-proposed keywords for '{placeholder}' (LLM提议的关键词): {keywords}") - # Extract settings from alpha_details for the get_datafields call + # Extract settings from alpha_details for the cache file path settings = alpha_details.get('settings', {}) - print(f"Alpha settings for datafield search (用于数据字段搜索的Alpha设置):", flush=True) + logger.info(f"Alpha settings for datafield search (用于数据字段搜索的Alpha设置):") instrument_type = settings.get('instrumentType', 'EQUITY') - + if user_region: region = user_region elif 'region' in settings: region = settings['region'] else: - print(f"❌ Error: Could not determine 'region' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的地区，Alpha设置中缺失且用户未提供)", flush=True) + logger.error(f"❌ Error: Could not determine 'region' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的地区，Alpha设置中缺失且用户未提供)") return [] - print(f" 数据地区: {region}", flush=True) - + logger.info(f" 数据地区: {region}") + if user_universe: universe = user_universe elif 'universe' in settings: universe = settings['universe'] else: - print(f"❌ Error: Could not determine 'universe' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的范围，Alpha设置中缺失且用户未提供)", flush=True) + logger.error(f"❌ Error: Could not determine 'universe' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的范围，Alpha设置中缺失且用户未提供)") return [] - print(f" 数据范围: {universe}", flush=True) - + logger.info(f" 数据范围: {universe}") + if user_delay is not None: delay = user_delay elif 'delay' in settings: delay = settings['delay'] else: - print(f"❌ Error: Could not determine 'delay' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的Delay，Alpha设置中缺失且用户未提供)", flush=True) + logger.error(f"❌ Error: Could not determine 'delay' for datafield search. It is missing in Alpha settings and not provided by user. (错误：无法确定数据搜索的Delay，Alpha设置中缺失且用户未提供)") return [] - print(f" Delay: {delay} 类别", flush=True) - + logger.info(f" Delay: {delay} 类别") + if user_category: - print(f" Category Filter: {user_category}", flush=True) + logger.info(f" Category Filter: {user_category}") - # Use asyncio.gather to make parallel API calls for efficiency - tasks = [] + # 检查本地缓存文件 (从项目根目录的 dataset 文件夹读取) + dataset_dir = Path(__file__).parent.parent / "dataset" + + # 确定要读取的类别列表（前端总是传递类别列表） + categories_to_read = user_category if isinstance(user_category, list) else [user_category] if user_category else [] + + if not categories_to_read: + logger.error(f"❌ 未指定数据类别，请先选择类别或点击'不筛选(默认)'") + raise ValueError("未指定数据类别，请先选择类别") + + logger.info(f"[Cache] 准备读取 {len(categories_to_read)} 个类别的数据字段: {categories_to_read}") + + # 读取多个类别的缓存文件并合并 + all_dataframes = [] + missing_files = [] + used_cache_files = [] # 记录使用的缓存文件 + + for cat in categories_to_read: + cache_filename = f"datafields_cache_{region}_{universe}_D{delay}_{cat}.csv" + cache_path = dataset_dir / cache_filename + + if cache_path.exists(): + used_cache_files.append(cache_filename) + logger.info(f"[Cache] 读取类别 '{cat}' 的缓存: {cache_filename}") + df = pd.read_csv(cache_path) + all_dataframes.append(df) + logger.info(f"[Cache] 类别 '{cat}' 包含 {len(df)} 个字段") + else: + missing_files.append(cache_filename) + logger.warning(f"[Cache] 类别 '{cat}' 的缓存文件不存在: {cache_filename}") + + # 输出使用的缓存文件汇总 + logger.info(f"[Cache] ==========================================") + logger.info(f"[Cache] 使用的缓存文件 ({len(used_cache_files)} 个):") + for i, filename in enumerate(used_cache_files, 1): + logger.info(f"[Cache] {i}. {filename}") + logger.info(f"[Cache] ==========================================") + + if not all_dataframes: + logger.error(f"❌ 所有数据字段缓存文件都不存在") + logger.error(f"❌ 请先点击'下载数据字段缓存'按钮下载数据字段") + raise FileNotFoundError(f"数据字段缓存文件不存在: {', '.join(missing_files)}，请先下载数据字段缓存") + + # 合并所有数据框 + all_datafields_df = pd.concat(all_dataframes, ignore_index=True) + # 去重 + all_datafields_df.drop_duplicates(subset=['id'], inplace=True) + logger.info(f"[Cache] 合并后共有 {len(all_datafields_df)} 个唯一字段") + + # 在本地缓存中搜索匹配关键词的字段 + matched_results = [] for keyword in keywords: - tasks.append( - asyncio.to_thread(get_datafields, - s=s, - instrument_type=instrument_type, - region=region, - delay=delay, - universe=universe, - search=keyword, - category=user_category if user_category else "", - data_type=user_data_type - ) + # 在 id 和 description 中搜索关键词 + keyword_lower = keyword.lower() + mask = ( + all_datafields_df['id'].str.lower().str.contains(keyword_lower, na=False) | + all_datafields_df.get('description', '').str.lower().str.contains(keyword_lower, na=False) ) - - results = await asyncio.gather(*tasks) - - # Process results to get top N from each keyword search - top_results_per_keyword = [] - for res_df in results: - if not res_df.empty: - top_results_per_keyword.append(res_df.head(top_n)) + matched = all_datafields_df[mask] + if not matched.empty: + matched_results.append(matched.head(top_n)) + logger.info(f"[Cache] 关键词 '{keyword}' 匹配到 {len(matched)} 个字段") + else: + logger.info(f"[Cache] 关键词 '{keyword}' 未匹配到字段") candidate_datafields = [] - if top_results_per_keyword: - # Concatenate the top N results from all keywords - combined_df = pd.concat(top_results_per_keyword, ignore_index=True) - # Remove duplicates from the combined list + if matched_results: + # 合并所有匹配结果 + combined_df = pd.concat(matched_results, ignore_index=True) + # 去重 combined_df.drop_duplicates(subset=['id'], inplace=True) - # Format the final list of candidates + # 限制总数 + combined_df = combined_df.head(top_n * len(keywords)) + # 格式化结果 candidate_datafields = combined_df[['id', 'description']].to_dict(orient='records') + logger.info(f"[Cache] 最终返回 {len(candidate_datafields)} 个候选字段") return candidate_datafields + async def get_group_datafield_candidates(template_expression: str, template_explanation: str, placeholder: str, llm_client: openai.AsyncOpenAI, top_n: int = 3) -> list[dict]: """ Uses an LLM to select suitable group data fields from a predefined list. """ - predefined_group_fields = ["industry", "subindustry", "sector", "market", "exchange"] - + predefined_group_fields = ["industry", + "subindustry", "sector", "market", "exchange"] + prompt = f""" As a quantitative researcher, you need to select the most relevant group data fields for an alpha template placeholder. Based on the template's logic and the placeholder's name, select {top_n} group fields from the following list that are most suitable: {predefined_group_fields}. @@ -1172,13 +1305,17 @@ async def get_group_datafield_candidates(template_expression: str, template_expl Now, generate the JSON array of selected group data fields. """ - print(f"--- Calling LLM to select group datafields for placeholder (正在调用LLM选择分组数据字段): {placeholder} ---", flush=True) + logger.info( + f"--- Calling LLM to select group datafields for placeholder (正在调用LLM选择分组数据字段): {placeholder} ---") response = await call_llm(prompt, llm_client) - + if isinstance(response, list) and all(isinstance(item, str) for item in response): return [{"name": field} for field in response[:top_n]] - print(f"Warning: LLM did not return a valid list of strings for group datafields (警告：LLM未返回有效的分组数据字段列表). Got: {response}", flush=True) - return [{"name": field} for field in predefined_group_fields[:top_n]] # Fallback to default if LLM fails + logger.warning( + f"Warning: LLM did not return a valid list of strings for group datafields (警告：LLM未返回有效的分组数据字段列表). Got: {response}") + # Fallback to default if LLM fails + return [{"name": field} for field in predefined_group_fields[:top_n]] + async def get_operator_candidates(template_expression: str, template_explanation: str, placeholder: str, llm_client: openai.AsyncOpenAI, top_n: int = 3) -> list[dict]: """ @@ -1189,11 +1326,12 @@ async def get_operator_candidates(template_expression: str, template_explanation all_operators = operators_data.get('operators', []) if not all_operators: - print("No REGULAR scope operators found. (未找到REGULAR范围的运算符)", flush=True) + logger.info("No REGULAR scope operators found. (未找到REGULAR范围的运算符)") return [] # Create a summary of available operators for the LLM - operator_names_and_descriptions = "\n".join([f"- {op['name']}: {op.get('description', 'No description available.')}" for op in all_operators]) + operator_names_and_descriptions = "\n".join( + [f"- {op['name']}: {op.get('description', 'No description available.')}" for op in all_operators]) prompt = f""" As a quantitative finance expert, you need to select the most relevant operators for an alpha template placeholder. @@ -1219,7 +1357,8 @@ async def get_operator_candidates(template_expression: str, template_explanation Now, generate the JSON array of selected operators. """ - print(f"--- Calling LLM to select operator candidates for placeholder (正在调用LLM选择运算符候选): {placeholder} ---", flush=True) + logger.info( + f"--- Calling LLM to select operator candidates for placeholder (正在调用LLM选择运算符候选): {placeholder} ---") response = await call_llm(prompt, llm_client) if isinstance(response, list) and all(isinstance(item, str) for item in response): @@ -1228,13 +1367,16 @@ async def get_operator_candidates(template_expression: str, template_explanation for selected_name in response: for op in all_operators: if op['name'] == selected_name: - selected_ops_details.append({"name": op['name'], "description": op.get('description', '')}) + selected_ops_details.append( + {"name": op['name'], "description": op.get('description', )}) break return selected_ops_details[:top_n] - - print(f"Warning: LLM did not return a valid list of strings for operator candidates (警告：LLM未返回有效的运算符候选列表). Got: {response}", flush=True) + + logger.warning( + f"Warning: LLM did not return a valid list of strings for operator candidates (警告：LLM未返回有效的运算符候选列表). Got: {response}") # Fallback to a default set if LLM fails - return [{"name": op['name'], "description": op.get('description', '')} for op in all_operators[:top_n]] + return [{"name": op['name'], "description": op.get('description', )} for op in all_operators[:top_n]] + async def get_parameter_candidates(param_type: str, template_expression: str, template_explanation: str, placeholder: str, llm_client: openai.AsyncOpenAI) -> list[dict]: """ @@ -1270,13 +1412,15 @@ async def get_parameter_candidates(param_type: str, template_expression: str, te Now, generate the JSON array of numerical candidates. """ - print(f"--- Calling LLM to suggest candidates for {param_type} placeholder (正在调用LLM建议参数候选): {placeholder} ---", flush=True) + logger.info( + f"--- Calling LLM to suggest candidates for {param_type} placeholder (正在调用LLM建议参数候选): {placeholder} ---") response = await call_llm(prompt, llm_client) if isinstance(response, list) and all(isinstance(item, (int, float)) for item in response): return [{"value": val} for val in response] - print(f"Warning: LLM did not return a valid list of numbers for {param_type} candidates (警告：LLM未返回有效的数字候选列表). Got: {response}", flush=True) - + logger.warning( + f"Warning: LLM did not return a valid list of numbers for {param_type} candidates (警告：LLM未返回有效的数字候选列表). Got: {response}") + # Fallback to default if LLM fails if param_type == "integer_parameter": return [{"value": x} for x in [10, 20, 60, 120, 252]] @@ -1284,6 +1428,7 @@ async def get_parameter_candidates(param_type: str, template_expression: str, te return [{"value": x} for x in [0.01, 0.05, 0.1, 0.2, 0.5]] return [] + async def judge_placeholder_type(placeholder: str, template_expression: str, template_explanation: str, operator_summary: str, llm_client: openai.AsyncOpenAI) -> str: """ Uses an LLM to judge the type of placeholder (e.g., "data_field", "integer_parameter", "group_operator"). @@ -1323,67 +1468,75 @@ async def judge_placeholder_type(placeholder: str, template_expression: str, tem Now, classify the placeholder. """ - print(f"--- Calling LLM to judge type for placeholder (正在调用LLM判断占位符类型): {placeholder} ---", flush=True) - + logger.info( + f"--- Calling LLM to judge type for placeholder (正在调用LLM判断占位符类型): {placeholder} ---") + response = await call_llm(prompt, llm_client) return response.get("placeholder_type", "unknown") + async def populate_template(s: SingleSession, alpha_details: dict, template_expression: str, template_explanation: str, operator_summary: str, llm_client: openai.AsyncOpenAI, top_n_datafield: int = 50, user_region: Optional[str] = None, user_universe: Optional[str] = None, user_delay: Optional[int] = None, user_category: Optional[Union[str, list]] = None, user_data_type: str = "MATRIX") -> dict: """ Populates placeholders in an alpha template with candidate data fields, operators, or parameters. """ placeholders = extract_placeholders(template_expression) - + if not placeholders: - print("No placeholders found in the template. (模板中未找到占位符)", flush=True) + logger.info("No placeholders found in the template. (模板中未找到占位符)") return {} """ Populates placeholders in an alpha template with candidate data fields, operators, or parameters. """ placeholders = extract_placeholders(template_expression) - print(f"Found placeholders in template (在模板中找到占位符): {placeholders}", flush=True) - + logger.info(f"Found placeholders in template (在模板中找到占位符): {placeholders}") + populated_placeholders = {} for ph in placeholders: # Use LLM to judge placeholder type ph_type = await judge_placeholder_type(ph, template_expression, template_explanation, operator_summary, llm_client) - print(f"'{ph}' judged as type (判断类型为): {ph_type}", flush=True) + logger.info(f"'{ph}' judged as type (判断类型为): {ph_type}") if ph_type == "data_field": candidates = await get_datafield_candidates(s, alpha_details, template_expression, template_explanation, ph, llm_client, top_n=top_n_datafield, user_region=user_region, user_universe=user_universe, user_delay=user_delay, user_category=user_category, user_data_type=user_data_type) - populated_placeholders[ph] = {"type": "data_field", "candidates": candidates} + populated_placeholders[ph] = { + "type": "data_field", "candidates": candidates} elif ph_type == "group_data_field": candidates = await get_group_datafield_candidates(template_expression, template_explanation, ph, llm_client) - populated_placeholders[ph] = {"type": "group_data_field", "candidates": candidates} - elif ph_type in ["operator", "group_operator", "ts_operator","vector_operator"]: + populated_placeholders[ph] = { + "type": "group_data_field", "candidates": candidates} + elif ph_type in ["operator", "group_operator", "ts_operator", "vector_operator"]: candidates = await get_operator_candidates(template_expression, template_explanation, ph, llm_client) - populated_placeholders[ph] = {"type": ph_type, "candidates": candidates} + populated_placeholders[ph] = { + "type": ph_type, "candidates": candidates} elif ph_type in ["integer_parameter", "float_parameter"]: candidates = await get_parameter_candidates(ph_type, template_expression, template_explanation, ph, llm_client) - populated_placeholders[ph] = {"type": ph_type, "candidates": candidates} + populated_placeholders[ph] = { + "type": ph_type, "candidates": candidates} elif ph_type == "string_parameter": # Add logic for string_parameter if needed, for now it returns empty - populated_placeholders[ph] = {"type": "string_parameter", "candidates": []} + populated_placeholders[ph] = { + "type": "string_parameter", "candidates": []} else: - print(f"Could not determine type for placeholder (无法确定占位符类型): {ph} (LLM classified as {ph_type})", flush=True) + logger.info( + f"Could not determine type for placeholder (无法确定占位符类型): {ph} (LLM classified as {ph_type})") populated_placeholders[ph] = {"type": "unknown", "candidates": []} - + return populated_placeholders + def get_datafield_prefix(datafield_name: str) -> str: """Extracts the prefix from a datafield name (e.g., 'anl44_...' -> 'anl44').""" if '_' in datafield_name: return datafield_name.split('_')[0] return datafield_name - async def generate_new_alphas(alpha_description, brain_session, template_summary: Optional[str] = None, top_n_datafield: int = 50, user_region: Optional[str] = None, user_universe: Optional[str] = None, user_delay: Optional[int] = None, user_category: Optional[Union[str, list]] = None, user_data_type: str = "MATRIX", max_retries: int = 20): """ Main function to generate new alpha templates based on a seed alpha. - + Args: alpha_description: The alpha description JSON string. brain_session: The BRAIN session object. @@ -1393,7 +1546,7 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary """ # 声明使用全局变量 global LLM_model_name, LLM_API_KEY, llm_base_url - + # Load template summary if not provided if template_summary is None: template_summary = load_template_summary() @@ -1403,34 +1556,35 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary try: llm_api_key = get_token_from_auth_server() llm_base_url_value = llm_base_url # 使用全局变量 - llm_client = openai.AsyncOpenAI(base_url=llm_base_url_value, api_key=llm_api_key) - print("✓ LLM Gateway 认证成功", flush=True) + llm_client = openai.AsyncOpenAI( + base_url=llm_base_url_value, api_key=llm_api_key) + logger.info("✓ LLM Gateway 认证成功") except Exception as e: - print(f"❌ LLM Gateway 认证失败: {e}", flush=True) + logger.error(f"❌ LLM Gateway 认证失败: {e}") sys.exit(1) details = json.loads(alpha_description) - - if not details: - print(f"Failed to retrieve details for Alpha (获取Alpha详情失败)", flush=True) - sys.exit(1) - - print("Alpha Details Retrieved (已获取Alpha详情):", flush=True) - print(json.dumps(details, indent=4), flush=True) + if not details: + logger.error(f"Failed to retrieve details for Alpha (获取Alpha详情失败)") + sys.exit(1) + + logger.info("Alpha Details Retrieved (已获取Alpha详情):") + logger.info(json.dumps(details, indent=4)) # --- Step 4: Propose New Alpha Templates --- - print(f"\n{'='*60}", flush=True) - print("[Step 2/5] 正在生成 Alpha 模板提议...", flush=True) - print(f"{'='*60}", flush=True) + logger.info(f"\n{'='*60}") + logger.info("[Step 2/5] 正在生成 Alpha 模板提议...") + logger.info(f"{'='*60}") proposed_templates = await propose_alpha_templates(details, template_summary, llm_client, user_data_type=user_data_type, max_retries=max_retries) if not proposed_templates: - print("Failed to generate proposed alpha templates. (生成提议模板失败)", flush=True) + logger.error("Failed to generate proposed alpha templates. (生成提议模板失败)") sys.exit(1) - - print("\n--- Proposed Alpha Templates (JSON) (建议的Alpha模板,多样性会受到模型和模板总结文档的影响) ---", flush=True) - print(json.dumps(proposed_templates, indent=4), flush=True) + + logger.info( + "\n--- Proposed Alpha Templates (JSON) (建议的Alpha模板,多样性会受到模型和模板总结文档的影响) ---") + logger.info(json.dumps(proposed_templates, indent=4)) # --- Validation: Drop templates with suspicious literal identifiers --- try: @@ -1443,69 +1597,70 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary parse_alpha_code, ) except Exception as e: - print(f"⚠ 模板校验步骤出现异常，跳过校验: {e}", flush=True) + logger.warning(f"⚠ 模板校验步骤出现异常，跳过校验: {e}") if not proposed_templates: - print("❌ 所有模板在校验后被丢弃，无法继续。", flush=True) + logger.error("❌ 所有模板在校验后被丢弃，无法继续。") sys.exit(1) # --- Step 5: Process all proposed templates and gather candidates --- # --- Step 6: Prepare for Output --- - print(f"\n{'='*60}", flush=True) - print("[Step 3/5] 正在处理模板并收集候选数据字段...", flush=True) - print(f"{'='*60}", flush=True) - + logger.info(f"\n{'='*60}") + logger.info("[Step 3/5] 正在处理模板并收集候选数据字段...") + logger.info(f"{'='*60}") + # Ensure the output directory exists next to this script output_dir = Path(__file__).parent / "output" try: output_dir.mkdir(parents=True, exist_ok=True) - print(f"✓ 输出目录已准备: {output_dir}", flush=True) + logger.info(f"✓ 输出目录已准备: {output_dir}") except Exception as e: - print(f"Warning: could not create directory {output_dir}: {e}", flush=True) + logger.warning( + f"Warning: could not create directory {output_dir}: {e}") output_filepath = output_dir / f"Alpha_candidates.json" - + final_output = {} - + # --- Step 5: Process all proposed templates and gather candidates --- total_templates = len(proposed_templates) for idx, (template_expr, template_expl) in enumerate(proposed_templates.items(), 1): - print(f"\n[模板 {idx}/{total_templates}] 正在处理: '{template_expr[:60]}...'", flush=True) + logger.info( + f"\n[模板 {idx}/{total_templates}] 正在处理: '{template_expr[:60]}...'") try: populated_info = await populate_template(brain_session, details, template_expr, template_expl, operator_summary, llm_client, top_n_datafield=top_n_datafield, user_region=user_region, user_universe=user_universe, user_delay=user_delay, user_category=user_category, user_data_type=user_data_type) # Skip templates where any data_field placeholder has zero candidates if _should_skip_due_to_empty_candidates(populated_info): - print("⚠ 该模板存在数据字段候选为空的占位符，跳过此模板。", flush=True) + logger.warning("⚠ 该模板存在数据字段候选为空的占位符，跳过此模板。") continue - + final_output[template_expr] = { "template_explanation": template_expl, "seed_alpha_settings": details.get('settings', {}), "placeholder_candidates": populated_info } - + # --- Incremental Saving --- try: with output_filepath.open('w', encoding='utf-8') as f: json.dump(final_output, f, indent=4) - print(f"✓ Progress saved to {output_filepath.name}", flush=True) + logger.info(f"✓ Progress saved to {output_filepath.name}") except IOError as e: - print(f"⚠️ Warning: Failed to save progress: {e}", flush=True) - + logger.warning(f"⚠️ Warning: Failed to save progress: {e}") + except Exception as e: - print(f"❌ Error processing template '{template_expr}': {e}", flush=True) - print("Skipping this template and continuing...", flush=True) + logger.error(f"❌ Error processing template '{template_expr}': {e}") + logger.info("Skipping this template and continuing...") continue - print(f"\n{'='*60}", flush=True) - print("[Step 4/5] 正在生成 Alpha 表达式组合...", flush=True) - print(f"{'='*60}", flush=True) - print(f"✓ 已处理 {len(final_output)} 个有效模板", flush=True) - - print("\n--- Final Consolidated Output (最终合并输出) ---", flush=True) - print(json.dumps(final_output, indent=4), flush=True) + logger.info(f"\n{'='*60}") + logger.info("[Step 4/5] 正在生成 Alpha 表达式组合...") + logger.info(f"{'='*60}") + logger.info(f"✓ 已处理 {len(final_output)} 个有效模板") + logger.info("\n--- Final Consolidated Output (最终合并输出) ---") + logger.info(json.dumps(final_output, indent=4)) generated_expressions = set() @@ -1518,64 +1673,69 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary for placeholder, details in placeholder_candidates.items(): # Extract only the 'value' or 'name' from the candidates list if details["type"] == "data_field": - candidates_for_placeholders[placeholder] = [c["id"] for c in details["candidates"]] + candidates_for_placeholders[placeholder] = [ + c["id"] for c in details["candidates"]] elif details["type"] in ["integer_parameter", "float_parameter"]: - candidates_for_placeholders[placeholder] = [str(c["value"]) for c in details["candidates"]] + candidates_for_placeholders[placeholder] = [ + str(c["value"]) for c in details["candidates"]] elif details["type"] == "group_data_field": - candidates_for_placeholders[placeholder] = [c["name"] for c in details["candidates"]] + candidates_for_placeholders[placeholder] = [ + c["name"] for c in details["candidates"]] elif details["type"] == "operator": - candidates_for_placeholders[placeholder] = [c["name"] for c in details["candidates"]] + candidates_for_placeholders[placeholder] = [ + c["name"] for c in details["candidates"]] else: candidates_for_placeholders[placeholder] = [] - # --- Step 3: Implement logic to generate all alpha expression combinations from the candidates --- # Generate all possible combinations of placeholder values placeholder_names = list(candidates_for_placeholders.keys()) - all_combinations_values = list(itertools.product(*candidates_for_placeholders.values())) + all_combinations_values = list(itertools.product( + *candidates_for_placeholders.values())) for combination_values in all_combinations_values: - + # --- ATOM Mode --- datafield_values_in_combo = [] - placeholder_types = {ph: details["type"] for ph, details in placeholder_candidates.items()} - + placeholder_types = {ph: details["type"] + for ph, details in placeholder_candidates.items()} + for i, placeholder_name in enumerate(placeholder_names): if placeholder_types.get(placeholder_name) == 'data_field': datafield_values_in_combo.append(combination_values[i]) - + if len(datafield_values_in_combo) > 1: - first_prefix = get_datafield_prefix(datafield_values_in_combo[0]) + first_prefix = get_datafield_prefix( + datafield_values_in_combo[0]) if not all(get_datafield_prefix(df) == first_prefix for df in datafield_values_in_combo): continue # Skip this combination as prefixes do not match current_expression = template_expression for i, placeholder_name in enumerate(placeholder_names): - current_expression = current_expression.replace(placeholder_name, combination_values[i]) - + current_expression = current_expression.replace( + placeholder_name, combination_values[i]) + # Check for duplicates before adding if current_expression not in generated_expressions: generated_expressions.add(current_expression) # dump all unique generated expressions to a file, a list of strings in json file - print(f"\n{'='*60}", flush=True) - print("[Step 5/5] 正在验证生成的表达式...", flush=True) - print(f"{'='*60}", flush=True) - print(f"✓ 生成的唯一 Alpha 表达式总数: {len(generated_expressions)}", flush=True) - + logger.info(f"\n{'='*60}") + logger.info("[Step 5/5] 正在验证生成的表达式...") + logger.info(f"{'='*60}") + logger.info(f"✓ 生成的唯一 Alpha 表达式总数: {len(generated_expressions)}") + # output_filepath = output_dir / f"Alpha_generated_expressions.json" # try: # with output_filepath.open('w', encoding='utf-8') as f: # json.dump(list(generated_expressions), f, indent=4) - # print(f"\nGenerated expressions successfully written to {output_filepath} (生成的表达式已成功写入)") + # logger.info(f"\nGenerated expressions successfully written to {output_filepath} (生成的表达式已成功写入)") # except IOError as e: - # print(f"Error writing generated expressions to file {output_filepath} (写入生成的表达式出错): {e}") - - + # logger.info(f"Error writing generated expressions to file {output_filepath} (写入生成的表达式出错): {e}") validator = val.ExpressionValidator() - print("请注意，该文件仅用于验证表达式的格式正确性，\n不保证表达式在实际使用中的逻辑正确性或可执行性。\n", flush=True) - print("不在内置函数列表中的operator将无法检查，如有需要，请使用AI按需修改本源代码添加", flush=True) + logger.info("请注意，该文件仅用于验证表达式的格式正确性，\n不保证表达式在实际使用中的逻辑正确性或可执行性。\n") + logger.info("不在内置函数列表中的operator将无法检查，如有需要，请使用AI按需修改本源代码添加") expressions_data = list(generated_expressions) # 提取表达式列表 @@ -1585,23 +1745,24 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary elif isinstance(expressions_data, list): expressions = expressions_data else: - print("错误: JSON文件格式不正确，需要包含表达式列表", flush=True) + logger.error("错误: JSON文件格式不正确，需要包含表达式列表") return # 验证表达式 valid_expressions = [] invalid_expressions = [] - print(f"开始验证 {len(expressions)} 个表达式...", flush=True) + logger.info(f"开始验证 {len(expressions)} 个表达式...") for i, expr in enumerate(expressions, 1): if i % 10 == 0: - print(f"已验证 {i}/{len(expressions)} 个表达式", flush=True) - + logger.info(f"已验证 {i}/{len(expressions)} 个表达式") + result = validator.check_expression(expr) if result["valid"]: valid_expressions.append(expr) else: - invalid_expressions.append({"expression": expr, "errors": result["errors"]}) + invalid_expressions.append( + {"expression": expr, "errors": result["errors"]}) # 生成输出文件路径 name = "Alpha_generated_expressions" @@ -1609,32 +1770,32 @@ async def generate_new_alphas(alpha_description, brain_session, template_summary invalid_output_path = os.path.join(output_dir, f"{name}_error.json") # 保存结果到JSON文件 - print(f"\n验证完成！", flush=True) - print(f"有效表达式: {len(valid_expressions)}", flush=True) - print(f"无效表达式: {len(invalid_expressions)}", flush=True) + logger.info(f"\n验证完成！") + logger.info(f"有效表达式: {len(valid_expressions)}") + logger.info(f"无效表达式: {len(invalid_expressions)}") # 保存有效表达式 try: with open(valid_output_path, 'w', encoding='utf-8') as f: json.dump(valid_expressions, f, ensure_ascii=False, indent=2) - print(f"有效表达式已保存到: {valid_output_path}", flush=True) + logger.info(f"有效表达式已保存到: {valid_output_path}") except Exception as e: - print(f"错误: 保存有效表达式失败 - {e}", flush=True) + logger.error(f"错误: 保存有效表达式失败 - {e}") # 保存无效表达式 try: with open(invalid_output_path, 'w', encoding='utf-8') as f: json.dump(invalid_expressions, f, ensure_ascii=False, indent=2) - print(f"无效表达式已保存到: {invalid_output_path}，文件包含错误详情", flush=True) - print("查看该文件，你将获得修改模板的灵感，你可以定位到错误的模板并在APP里修改", flush=True) + logger.info(f"无效表达式已保存到: {invalid_output_path}，文件包含错误详情") + logger.info("查看该文件，你将获得修改模板的灵感，你可以定位到错误的模板并在APP里修改") except Exception as e: - print(f"错误: 保存无效表达式失败 - {e}", flush=True) + logger.error(f"错误: 保存无效表达式失败 - {e}") - print("请注意，该文件仅用于验证表达式的格式正确性，\n不保证表达式在实际使用中的逻辑正确性或可执行性。\n", flush=True) - print("不在内置函数列表中的operator将无法检查，如有需要，请使用AI按需修改validator源代码添加", flush=True) + logger.info("请注意，该文件仅用于验证表达式的格式正确性，\n不保证表达式在实际使用中的逻辑正确性或可执行性。\n") + logger.info("不在内置函数列表中的operator将无法检查，如有需要，请使用AI按需修改validator源代码添加") - print("不同模型效果不同，默认的kimi模型可能会产生Alpha语法错误，请检查生成的模板文件进行甄别", flush=True) - print("下一步，请下载已完成的模板，放入APP首页进行解析和语法检查，强烈建议生成表达式后手动尝试回测", flush=True) + logger.info("不同模型效果不同，默认的kimi模型可能会产生Alpha语法错误，请检查生成的模板文件进行甄别") + logger.info("下一步，请下载已完成的模板，放入APP首页进行解析和语法检查，强烈建议生成表达式后手动尝试回测") async def main(): @@ -1649,76 +1810,79 @@ async def main(): try: with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) - print(f"✓ 已从命令行参数加载配置: {config_path}", flush=True) + logger.info(f"✓ 已从命令行参数加载配置: {config_path}") # Ensure all required fields are present or set defaults if 'top_n_datafield' not in config: config['top_n_datafield'] = 50 if 'template_summary_path' not in config: config['template_summary_path'] = None except Exception as e: - print(f"❌ 加载配置文件失败: {e}", flush=True) + logger.error(f"❌ 加载配置文件失败: {e}") sys.exit(1) else: - print(f"❌ 配置文件不存在: {config_path}", flush=True) + logger.error(f"❌ 配置文件不存在: {config_path}") sys.exit(1) else: # --- Step 0: 交互式输入收集配置信息 --- - print("\n" + "="*60, flush=True) - print("交互式配置输入模式", flush=True) - print("="*60 + "\n", flush=True) - + logger.info("\n" + "="*60) + logger.info("交互式配置输入模式") + logger.info("="*60 + "\n") + config = interactive_input() - + # 设置全局变量 global LLM_model_name, LLM_API_KEY, llm_base_url, username, password - - print("\n[Config] 正在设置全局变量...", flush=True) - - required_config_fields = ['LLM_model_name', 'LLM_API_KEY', 'llm_base_url', 'username', 'password', 'alpha_id'] + + logger.info("\n[Config] 正在设置全局变量...") + + required_config_fields = ['LLM_model_name', 'LLM_API_KEY', + 'llm_base_url', 'username', 'password', 'alpha_id'] missing_fields = [f for f in required_config_fields if f not in config] - + if missing_fields: - print(f"❌ [Config] 配置缺少必填字段: {missing_fields}", flush=True) - print(f"❌ [Config] 当前配置内容: {list(config.keys())}", flush=True) + logger.error(f"❌ [Config] 配置缺少必填字段: {missing_fields}") + logger.error(f"❌ [Config] 当前配置内容: {list(config.keys())}") sys.exit(1) - + LLM_model_name = config['LLM_model_name'] LLM_API_KEY = config['LLM_API_KEY'] llm_base_url = config['llm_base_url'] username = config['username'] password = config['password'] - - print(f"✓ [Config] LLM_model_name: {LLM_model_name}", flush=True) - print(f"✓ [Config] llm_base_url: {llm_base_url}", flush=True) - print(f"✓ [Config] username: {username}", flush=True) - print(f"✓ [Config] alpha_id: {config['alpha_id']}", flush=True) - + + logger.info(f"✓ [Config] LLM_model_name: {LLM_model_name}") + logger.info(f"✓ [Config] llm_base_url: {llm_base_url}") + logger.info(f"✓ [Config] username: {username}") + logger.info(f"✓ [Config] alpha_id: {config['alpha_id']}") + # --- Step 1: 加载模板总结 --- - template_summary = load_template_summary(config.get('template_summary_path')) - + template_summary = load_template_summary( + config.get('template_summary_path')) + # --- Step 2: 启动 BRAIN 会话 --- - print("--- 正在启动 BRAIN 会话... ---", flush=True) + logger.info("--- 正在启动 BRAIN 会话... ---") s = start_session() - + # --- Step 3: 认证 LLM Gateway --- llm_client = None - print("--- 正在认证 LLM Gateway... ---", flush=True) + logger.info("--- 正在认证 LLM Gateway... ---") try: llm_api_key = get_token_from_auth_server() llm_base_url_value = llm_base_url - llm_client = openai.AsyncOpenAI(base_url=llm_base_url_value, api_key=llm_api_key) - print("✓ LLM Gateway 认证成功", flush=True) + llm_client = openai.AsyncOpenAI( + base_url=llm_base_url_value, api_key=llm_api_key) + logger.info("✓ LLM Gateway 认证成功") except Exception as e: - print(f"❌ LLM Gateway 认证失败: {e}", flush=True) + logger.error(f"❌ LLM Gateway 认证失败: {e}") sys.exit(1) # --- Step 4: 获取 Alpha 详情 --- alpha_id = config['alpha_id'] - print(f"\n--- 正在获取 Alpha ID: {alpha_id} 的详情... ---", flush=True) + logger.info(f"\n--- 正在获取 Alpha ID: {alpha_id} 的详情... ---") # --- Step 4.5: 交互式选择数据字段范围 --- if len(sys.argv) > 1: - user_datafield_config = { + user_datafield_config = { 'user_region': config.get('user_region'), 'user_universe': config.get('user_universe'), 'user_delay': config.get('user_delay'), @@ -1730,8 +1894,8 @@ async def main(): details_str = await generate_alpha_description(alpha_id, brain_session=s) await generate_new_alphas( - alpha_description=details_str, - brain_session=s, + alpha_description=details_str, + brain_session=s, template_summary=template_summary, top_n_datafield=config.get('top_n_datafield', 50), user_region=user_datafield_config.get('user_region'), @@ -1741,51 +1905,53 @@ async def main(): user_data_type=user_datafield_config.get('user_data_type', 'MATRIX'), max_retries=config.get('max_retries', 20) ) - + + def interactive_datafield_selection(s: SingleSession) -> dict: """ Interactively ask the user for datafield search configuration (Region, Universe, Delay). """ - print("\n" + "="*60, flush=True) - print("【附加配置】数据字段搜索范围配置", flush=True) - print("正在获取有效的 Region/Universe/Delay 组合...", flush=True) - + logger.info("\n" + "="*60) + logger.info("【附加配置】数据字段搜索范围配置") + logger.info("正在获取有效的 Region/Universe/Delay 组合...") + try: df = get_instrument_type_region_delay(s) except Exception as e: - print(f"⚠ 获取配置选项失败: {e}", flush=True) - print("将使用 Seed Alpha 的默认设置", flush=True) + logger.warning(f"⚠ 获取配置选项失败: {e}") + logger.info("将使用 Seed Alpha 的默认设置") return {} # Filter for EQUITY only as per current logic df_equity = df[df['InstrumentType'] == 'EQUITY'] - + if df_equity.empty: - print("未找到 EQUITY 类型的配置选项。", flush=True) + logger.info("未找到 EQUITY 类型的配置选项。") return {} - + # 1. Select Region regions = df_equity['Region'].unique().tolist() - print(f"\n可用地区 (Region): {regions}", flush=True) + logger.info(f"\n可用地区 (Region): {regions}") region_input = input(f"请输入地区 (直接回车使用 Seed Alpha 默认值): ").strip() - + selected_region = None if region_input: if region_input in regions: selected_region = region_input else: - print(f"⚠ 输入无效，将使用默认值", flush=True) - + logger.warning(f"⚠ 输入无效，将使用默认值") + # 2. Select Delay # If region is selected, filter delays for that region if selected_region: - delays = df_equity[df_equity['Region'] == selected_region]['Delay'].unique().tolist() + delays = df_equity[df_equity['Region'] == + selected_region]['Delay'].unique().tolist() else: delays = df_equity['Delay'].unique().tolist() - - print(f"\n可用延迟 (Delay): {delays}", flush=True) + + logger.info(f"\n可用延迟 (Delay): {delays}") delay_input = input(f"请输入延迟 (直接回车使用 Seed Alpha 默认值): ").strip() - + selected_delay = None if delay_input: try: @@ -1793,14 +1959,15 @@ def interactive_datafield_selection(s: SingleSession) -> dict: if d_val in delays: selected_delay = d_val else: - print(f"⚠ 输入不在列表中，将使用默认值", flush=True) + logger.warning(f"⚠ 输入不在列表中，将使用默认值") except ValueError: - print(f"⚠ 输入无效，将使用默认值", flush=True) + logger.warning(f"⚠ 输入无效，将使用默认值") # 3. Select Universe # If region and delay are selected, filter universes if selected_region and selected_delay is not None: - subset = df_equity[(df_equity['Region'] == selected_region) & (df_equity['Delay'] == selected_delay)] + subset = df_equity[(df_equity['Region'] == selected_region) & ( + df_equity['Delay'] == selected_delay)] if not subset.empty: universes = subset.iloc[0]['Universe'] else: @@ -1811,73 +1978,74 @@ def interactive_datafield_selection(s: SingleSession) -> dict: for u_list in df_equity['Universe']: universes.update(u_list) universes = list(universes) - - print(f"\n可用范围 (Universe): {universes}", flush=True) + + logger.info(f"\n可用范围 (Universe): {universes}") universe_input = input(f"请输入范围 (直接回车使用 Seed Alpha 默认值): ").strip() - + selected_universe = None if universe_input: if universe_input in universes: selected_universe = universe_input else: - print(f"⚠ 输入无效，将使用默认值", flush=True) - + logger.warning(f"⚠ 输入无效，将使用默认值") + # 4. Select Category - print("\n正在获取数据类别 (Data Categories)...", flush=True) + logger.info("\n正在获取数据类别 (Data Categories)...") categories = get_data_categories(s) - + selected_category = None if categories: - print("\n可用类别 (Categories):", flush=True) + logger.info("\n可用类别 (Categories):") for i, cat in enumerate(categories): - print(f"{i+1}. {cat['name']} (ID: {cat['id']})", flush=True) - + logger.info(f"{i+1}. {cat['name']} (ID: {cat['id']})") + cat_input = input(f"请输入类别编号或ID (多个用逗号分隔, 直接回车不筛选): ").strip() - + if cat_input: selected_categories = [] inputs = [x.strip() for x in cat_input.split(',')] - + for inp in inputs: # Check if input is an index if inp.isdigit(): idx = int(inp) - 1 if 0 <= idx < len(categories): selected_categories.append(categories[idx]['id']) - print(f"已选择类别: {categories[idx]['name']}", flush=True) + logger.info(f"已选择类别: {categories[idx]['name']}") else: # Check if input is an ID found = False for cat in categories: if cat['id'] == inp: selected_categories.append(cat['id']) - print(f"已选择类别: {cat['name']}", flush=True) + logger.info(f"已选择类别: {cat['name']}") found = True break if not found: - print(f"⚠ 输入无效: {inp}", flush=True) - + logger.warning(f"⚠ 输入无效: {inp}") + if selected_categories: selected_category = selected_categories else: - print(f"⚠ 未选择有效类别，将不筛选类别", flush=True) + logger.warning(f"⚠ 未选择有效类别，将不筛选类别") else: - print("⚠ 无法获取类别列表，跳过类别选择", flush=True) + logger.warning("⚠ 无法获取类别列表，跳过类别选择") # 5. Select Data Type - print("\n可用数据类型 (Data Type): [MATRIX, VECTOR]", flush=True) + logger.info("\n可用数据类型 (Data Type): [MATRIX, VECTOR]") data_type_input = input(f"请输入数据类型 (直接回车默认 MATRIX): ").strip().upper() - + selected_data_type = "MATRIX" if data_type_input == "VECTOR": - print("⚠ 警告: 请确保您输入的原型Alpha中正确地使用了vector operator,否则极容易造成数据类型错误", flush=True) + logger.warning( + "⚠ 警告: 请确保您输入的原型Alpha中正确地使用了vector operator,否则极容易造成数据类型错误") confirm = input("确认使用 VECTOR 吗? (y/n): ").strip().lower() if confirm == 'y': selected_data_type = "VECTOR" else: - print("已取消 VECTOR 选择，使用默认值 MATRIX", flush=True) + logger.info("已取消 VECTOR 选择，使用默认值 MATRIX") elif data_type_input and data_type_input != "MATRIX": - print(f"⚠ 输入无效，将使用默认值 MATRIX", flush=True) + logger.warning(f"⚠ 输入无效，将使用默认值 MATRIX") return { 'user_region': selected_region, @@ -1887,10 +2055,124 @@ def interactive_datafield_selection(s: SingleSession) -> dict: 'user_data_type': selected_data_type } + +async def run_transformer(config: dict) -> dict: + """ + 可导入的异步函数，用于执行 Alpha 生成流程。 + + Args: + config: 配置字典，包含以下必需字段： + - LLM_model_name: LLM 模型名称 + - LLM_API_KEY: LLM API 密钥 + - llm_base_url: LLM 服务基础 URL + - username: BRAIN 用户名 + - password: BRAIN 密码 + - alpha_id: Alpha ID + 以及可选字段： + - top_n_datafield: 数据字段数量限制（默认 50） + - template_summary_path: 模板总结文件路径 + - user_region: 用户指定的地区 + - user_universe: 用户指定的范围 + - user_delay: 用户指定的延迟 + - user_category: 用户指定的类别 + - user_data_type: 数据类型（默认 MATRIX） + - max_retries: 最大重试次数（默认 20） + + Returns: + dict: 执行结果，包含 success 状态和相关信息 + """ + # 设置全局变量（保持与原有代码兼容） + global LLM_model_name, LLM_API_KEY, llm_base_url, username, password + + logger.info("\n[Config] 正在设置全局变量...") + + required_config_fields = ['LLM_model_name', 'LLM_API_KEY', + 'llm_base_url', 'username', 'password', 'alpha_id'] + missing_fields = [f for f in required_config_fields if f not in config] + + if missing_fields: + logger.error(f"❌ [Config] 配置缺少必填字段: {missing_fields}") + logger.error(f"❌ [Config] 当前配置内容: {list(config.keys())}") + raise ValueError(f"Missing required config fields: {missing_fields}") + + LLM_model_name = config['LLM_model_name'] + LLM_API_KEY = config['LLM_API_KEY'] + llm_base_url = config['llm_base_url'] + username = config['username'] + password = config['password'] + + logger.info(f"✓ [Config] LLM_model_name: {LLM_model_name}") + logger.info(f"✓ [Config] llm_base_url: {llm_base_url}") + logger.info(f"✓ [Config] username: {username}") + logger.info(f"✓ [Config] alpha_id: {config['alpha_id']}") + + # --- Step 1: 加载模板总结 --- + template_summary = load_template_summary( + config.get('template_summary_path')) + + # --- Step 2: 启动 BRAIN 会话 --- + logger.info("--- 正在启动 BRAIN 会话... ---") + s = start_session() + + # --- Step 3: 认证 LLM Gateway --- + llm_client = None + logger.info("--- 正在认证 LLM Gateway... ---") + try: + llm_api_key = get_token_from_auth_server() + llm_base_url_value = llm_base_url + llm_client = openai.AsyncOpenAI( + base_url=llm_base_url_value, api_key=llm_api_key) + logger.info("✓ LLM Gateway 认证成功") + except Exception as e: + logger.error(f"❌ LLM Gateway 认证失败: {e}") + raise + + # --- Step 4: 获取 Alpha 详情 --- + alpha_id = config['alpha_id'] + logger.info(f"\n--- 正在获取 Alpha ID: {alpha_id} 的详情... ---") + + # --- Step 4.5: 使用配置中的数据字段范围 --- + user_datafield_config = { + 'user_region': config.get('user_region'), + 'user_universe': config.get('user_universe'), + 'user_delay': config.get('user_delay'), + 'user_category': config.get('user_category'), + 'user_data_type': config.get('user_data_type', 'MATRIX') + } + + details_str = await generate_alpha_description(alpha_id, brain_session=s) + await generate_new_alphas( + alpha_description=details_str, + brain_session=s, + template_summary=template_summary, + top_n_datafield=config.get('top_n_datafield', 50), + user_region=user_datafield_config.get('user_region'), + user_universe=user_datafield_config.get('user_universe'), + user_delay=user_datafield_config.get('user_delay'), + user_category=user_datafield_config.get('user_category'), + user_data_type=user_datafield_config.get('user_data_type', 'MATRIX'), + max_retries=config.get('max_retries', 20) + ) + + return {"success": True, "message": "Alpha generation completed"} + + +def run_transformer_sync(config: dict) -> dict: + """ + 同步版本的 run_transformer，方便直接调用。 + + Args: + config: 配置字典，详见 run_transformer 文档 + + Returns: + dict: 执行结果 + """ + return asyncio.run(run_transformer(config)) + + if __name__ == "__main__": # To allow asyncio to run in environments like Jupyter notebooks if sys.platform.startswith('win') and sys.version_info[:2] >= (3, 8): asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy()) - - asyncio.run(main()) + asyncio.run(main()) diff --git a/simple72/Tranformer/ace_lib.py b/simple72/Tranformer/ace_lib.py index bac64a9..6bdb679 100755 --- a/simple72/Tranformer/ace_lib.py +++ b/simple72/Tranformer/ace_lib.py @@ -1,4 +1,3 @@ -import getpass import json import logging import os @@ -6,23 +5,85 @@ import threading import time from functools import partial from multiprocessing.pool import ThreadPool -from pathlib import Path from typing import Literal, Optional, Union from urllib.parse import urljoin import pandas as pd import requests import tqdm -from helpful_functions import ( - expand_dict_columns, - save_pnl, - save_simulation_result, - save_yearly_stats, -) - DEV = False +def save_simulation_result(result: dict) -> None: + """ + Save the simulation result to a JSON file in the 'simulation_results' folder. + + Args: + result (dict): A dictionary containing the simulation result for an alpha. + """ + alpha_id = result["id"] + region = result["settings"]["region"] + folder_path = "simulation_results/" + file_path = os.path.join(folder_path, f"{alpha_id}_{region}") + + os.makedirs(folder_path, exist_ok=True) + + with open(file_path, "w", encoding="utf-8") as file: + json.dump(result, file) + + +def save_pnl(pnl_df: pd.DataFrame, alpha_id: str, region: str) -> None: + """ + Save the PnL data for an alpha to a CSV file in the 'alphas_pnl' folder. + + Args: + pnl_df (pandas.DataFrame): The DataFrame containing PnL data. + alpha_id (str): The ID of the alpha. + region (str): The region for which the PnL data was generated. + """ + folder_path = "alphas_pnl/" + file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv") + os.makedirs(folder_path, exist_ok=True) + + pnl_df.to_csv(file_path) + + +def save_yearly_stats(yearly_stats: pd.DataFrame, alpha_id: str, region: str): + """ + Save the yearly statistics for an alpha to a CSV file in the 'yearly_stats' folder. + + Args: + yearly_stats (pandas.DataFrame): The DataFrame containing yearly statistics. + alpha_id (str): The ID of the alpha. + region (str): The region for which the statistics were generated. + """ + folder_path = "yearly_stats/" + file_path = os.path.join(folder_path, f"{alpha_id}_{region}.csv") + os.makedirs(folder_path, exist_ok=True) + + yearly_stats.to_csv(file_path, index=False) + + +def expand_dict_columns(data: pd.DataFrame) -> pd.DataFrame: + """ + Expand dictionary columns in a DataFrame into separate columns. + + Args: + data (pandas.DataFrame): The input DataFrame with dictionary columns. + + Returns: + pandas.DataFrame: A new DataFrame with expanded columns. + """ + dict_columns = list(filter(lambda x: isinstance(data[x].iloc[0], dict), data.columns)) + new_columns = pd.concat( + [data[col].apply(pd.Series).rename(columns=lambda x: f"{col}_{x}") for col in dict_columns], + axis=1, + ) + + data = pd.concat([data, new_columns], axis=1) + return data + + class SingleSession(requests.Session): _instance = None _lock = threading.Lock() @@ -132,7 +193,7 @@ def start_session() -> SingleSession: logger.debug(f"New session created (ID: {id(s)}) with authentication response: {r.status_code}, {r.json()}") if r.status_code == requests.status_codes.codes.unauthorized: if r.headers["WWW-Authenticate"] == "persona": - print( + logger.info( "Complete biometrics authentication and press any key to continue: \n" + urljoin(r.url, r.headers["Location"]) + "\n" @@ -1192,6 +1253,7 @@ def get_datafields( dataset_id: str = "", data_type: str = "MATRIX", search: str = "", + category: str = "", ) -> pd.DataFrame: """ Retrieve available datafields based on specified parameters. @@ -1219,7 +1281,14 @@ def get_datafields( + f"®ion={region}&delay={str(delay)}&universe={universe}{type_param}&dataset.id={dataset_id}&limit=50" + "&offset={x}" ) - count = s.get(url_template.format(x=0)).json()["count"] + try: + response = s.get(url_template.format(x=0)) + response_json = response.json() + count = response_json.get("count", 0) + except Exception as e: + logger.warning(f"获取数据字段数量失败: {e}") + count = 0 + if count == 0: logger.warning( f"No fields found: region={region}, delay={str(delay)}, universe={universe}, " @@ -1238,13 +1307,15 @@ def get_datafields( ) count = 100 + logger.info(url_template.format(x=0)) + max_try = 5 datafields_list = [] for x in range(0, count, 50): for _ in range(max_try): datafields = s.get(url_template.format(x=x)) while datafields.status_code == 429: - print("status_code 429, sleep 3 seconds") + logger.info("status_code 429, sleep 3 seconds") time.sleep(3) datafields = s.get(url_template.format(x=x)) if "results" in datafields.json(): @@ -1510,5 +1581,4 @@ def main(): simulate_alpha_list(s, alpha_list) -if __name__ == "__main__": - main() + diff --git a/simple72/Tranformer/validator.py b/simple72/Tranformer/validator.py index 2fdaa4c..f1f1385 100755 --- a/simple72/Tranformer/validator.py +++ b/simple72/Tranformer/validator.py @@ -1,1261 +1,1265 @@ -#!/usr/bin/env python3 -# -*- coding: utf-8 -*- -""" -表达式验证器 - 使用抽象语法树验证字符串表达式格式是否正确 - -本模块实现了一个能够检测字符串表达式格式是否正确的系统，基于PLY(Python Lex-Yacc) -构建词法分析器和语法分析器，识别表达式中的操作符、函数和字段，并验证其格式正确性。 -""" - -import re -import sys -import json -import os -from typing import List, Dict, Any, Optional, Tuple - -# 尝试导入PLY库，如果不存在则提供安装提示 -try: - import ply.lex as lex - import ply.yacc as yacc -except ImportError: - print("错误: 需要安装PLY库。请运行 'pip install ply' 来安装。") - sys.exit(1) - -# 1. 定义支持的操作符和函数 -supported_functions = { - # Group 类别函数 - 'group_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - # group_mean(x, w, group) - 'group_mean': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, - 'group_median': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_rank': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_vector_proj': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, - 'group_normalize': {'min_args': 2, 'max_args': 5, 'arg_types': ['expression', 'category', 'expression', 'expression', 'expression']}, - 'group_extra': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, - 'group_backfill': {'min_args': 3, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['x', 'cat', 'days', 'std']}, - 'group_scale': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_count': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_zscore': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_std_dev': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_sum': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_neutralize': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, - 'group_multi_regression': {'min_args': 4, 'max_args': 9, 'arg_types': ['expression'] * 9}, - 'group_cartesian_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['category', 'category']}, - 'combo_a': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, - - # Transformational 类别函数 - 'right_tail': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'bucket': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 第二个参数可以是string类型的range参数 - 'tail': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression']}, - 'left_tail': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'trade_when': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, - 'generate_stats': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - - # Cross Sectional 类别函数 - 'winsorize': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'std']}, - 'rank': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'regression_proj': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'vector_neut': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'regression_neut': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'multi_regression': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # 支持多个自变量 - - # Time Series 类别函数 - 'ts_std_dev': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_mean': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_delay': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_corr': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, - 'ts_zscore': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_returns': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'mode'], 'keyword_only': True}, - 'ts_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_backfill(x, d) - 'ts_backfill': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'd']}, - 'days_from_last_change': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'last_diff_value': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_scale(x, d, constant=0) - 'ts_scale': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'constant'], 'keyword_only': True}, - # Platform: ts_entropy(x, d) - 'ts_entropy': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'd']}, - 'ts_step': {'min_args': 1, 'max_args': 1, 'arg_types': ['number']}, - 'ts_sum': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_co_kurtosis': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, - 'inst_tvr': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_decay_exp_window': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'factor'], 'keyword_only': True}, - 'ts_av_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_kurtosis': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_min_max_diff(x, d, f=0.5) - 'ts_min_max_diff': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'f'], 'keyword_only': True}, - 'ts_arg_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_min_max_cps(x, d, f=2) - 'ts_min_max_cps': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'f'], 'keyword_only': True}, - # Platform: ts_rank(x, d, constant=0) - 'ts_rank': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'constant'], 'keyword_only': True}, - 'ts_ir': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_theilsen': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, - # Platform: hump_decay(x, p=0) - 'hump_decay': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'p'], 'keyword_only': True}, - # Platform: ts_weighted_decay(x, k=0.5) - 'ts_weighted_decay': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'k'], 'keyword_only': True}, - # Platform: ts_quantile(x, d, driver="gaussian") - 'ts_quantile': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'string'], 'param_names': ['x', 'd', 'driver'], 'keyword_only': True}, - 'ts_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_count_nans': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_covariance': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, - 'ts_co_skewness': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, - 'ts_min_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_decay_linear(x, d, dense=false) - 'ts_decay_linear': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'boolean'], 'param_names': ['x', 'd', 'dense'], 'keyword_only': True}, - # Platform: jump_decay(x, d, sensitivity=0.5, force=0.1) - 'jump_decay': {'min_args': 2, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'd', 'sensitivity', 'force'], 'keyword_only': True}, - # Platform: ts_moment(x, d, k=0) - 'ts_moment': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'k'], 'keyword_only': True}, - 'ts_arg_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_regression': {'min_args': 3, 'max_args': 5, 'arg_types': ['expression', 'expression', 'number', 'number', 'number'], 'param_names': ['y', 'x', 'd', 'lag', 'rettype'], 'keyword_only': True}, - 'ts_skewness': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_max_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'kth_element': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'number', 'number']}, - 'hump': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'hump']}, - 'ts_median': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_delta': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - # Platform: ts_poly_regression(y, x, d, k=1) and k must be keyword if provided - 'ts_poly_regression': {'min_args': 3, 'max_args': 4, 'arg_types': ['expression', 'expression', 'number', 'number'], 'param_names': ['y', 'x', 'd', 'k'], 'keyword_only': True, 'keyword_only_from': 3}, - 'ts_target_tvr_decay': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, - 'ts_target_tvr_delta_limit': {'min_args': 2, 'max_args': 5, 'arg_types': ['expression', 'expression', 'number', 'number', 'number'], 'param_names': ['x', 'y', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, - 'ts_target_tvr_hump': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, - # Platform: ts_delta_limit(x, y, limit_volume=0.1) - 'ts_delta_limit': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'limit_volume'], 'keyword_only': True}, - - # Special 类别函数 - 'inst_pnl': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'self_corr': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'in': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 - 'universe_size': {'min_args': 0, 'max_args': 0, 'arg_types': []}, - - # Missing functions from operators.py - 'quantile': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'driver', 'sigma']}, # quantile(x, driver = gaussian, sigma = 1.0) - 'normalize': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'boolean', 'number']}, # normalize(x, useStd = false, limit = 0.0) - 'zscore': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # zscore(x) - - # Logical 类别函数 - 'or': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 - 'and': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 - 'not': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # 注意：这是关键字 - 'is_nan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'is_not_nan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'less': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'greater': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'is_finite': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'if_else': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, - 'not_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'less_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'greater_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - - # Vector 类别函数 - 'vec_kurtosis': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_min': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_count': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_sum': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_skewness': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_max': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_avg': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_range': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_choose': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'nth']}, - 'vec_powersum': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'constant']}, - 'vec_stddev': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_percentage': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'percentage']}, - 'vec_ir': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'vec_norm': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'ts_percentage': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'percentage']}, - 'signed_power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - 'ts_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, - - # Additional functions from test cases - 'rank_by_side': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'rate', 'scale']}, - 'log_diff': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'nan_mask': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, - 'ts_partial_corr': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'number']}, - 'ts_triple_corr': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'number']}, - 'clamp': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'lower', 'upper']}, - 'keep': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'condition', 'period']}, - 'replace': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'target', 'dest']}, - 'filter': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'h', 't']}, - 'one_side': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'string'], 'param_names': ['x', 'side']}, - 'scale_down': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'constant']}, - - # Arithmetic 类别函数 - # add(x, y, ..., filter=false) - # NOTE: add() is variadic (>=2 terms) with an optional boolean filter flag. - # We validate it with custom logic in validate_function(). - 'add': {'min_args': 2, 'max_args': 101, 'arg_types': ['expression'] * 101}, - 'multiply': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 99 + ['boolean'], 'param_names': ['x', 'y', 'filter']}, # multiply(x, y, ..., filter=false) - 'sign': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'subtract': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'boolean']}, # subtract(x, y, filter=false) - 'pasteurize': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'log': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'purify': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'arc_tan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'max': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # max(x, y, ...) - 'to_nan': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'boolean']}, # to_nan(x, value=0, reverse=false) - 'abs': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'sigmoid': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'divide': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # divide(x, y) - 'min': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # min(x, y, ...) - 'tanh': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'nan_out': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'lower', 'upper']}, # nan_out(x, lower=0, upper=0) - 'signed_power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # signed_power(x, y) - 'inverse': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'round': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'sqrt': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 's_log_1p': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'reverse': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # -x - 'power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # power(x, y) - 'densify': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - 'floor': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, - # Appended missing operators - 'arc_cos': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, - 'arc_sin': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, - 'ceiling': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, - 'exp': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, - 'fraction': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, - 'round_down': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'f']}, - 'is_not_finite': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'negate': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'ts_rank_gmean_amean_diff': {'min_args': 5, 'max_args': 5, 'arg_types': ['expression', 'expression', 'expression', 'expression', 'number'], 'param_names': ['input1', 'input2', 'input3', '...', 'd']}, - 'ts_vector_neut': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'd']}, - 'ts_vector_proj': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'd']}, - 'scale': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['x', 'scale', 'longscale', 'shortscale']}, - 'generalized_rank': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['open', 'm']}, - 'rank_gmean_amean_diff': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['input1', 'input2', 'input3', '...']}, - 'truncate': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'maxPercent']}, - 'vector_proj': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'y']}, - 'vec_filter': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['vec', 'value']}, - 'group_coalesce': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['original_group', 'group2', 'group3', '…']}, - 'group_percentage': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'category', 'expression'], 'param_names': ['x', 'group', 'percentage']}, - 'group_vector_neut': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'y', 'g']}, - 'convert': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'mode']}, - 'reduce_avg': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, - 'reduce_choose': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['input', 'nth', 'ignoreNan']}, - 'reduce_count': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, - 'reduce_ir': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_kurtosis': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_max': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_min': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_norm': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_percentage': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'percentage']}, - 'reduce_powersum': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['input', 'constant', 'precise']}, - 'reduce_range': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_skewness': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, - 'reduce_stddev': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, - 'reduce_sum': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, -} - -# 2. 定义group类型字段 -group_fields = { - 'sector', 'subindustry', 'industry', 'exchange', 'country', 'market' -} - -# 3. 有效类别集合 -valid_categories = group_fields - -# 4. 字段命名模式 - 只校验字段是不是数字字母下划线组成 -field_patterns = [ - re.compile(r'^[a-zA-Z0-9_]+$'), # 只允许数字、字母和下划线组成的字段名 -] - -# 4. 抽象语法树节点类型 -class ASTNode: - """抽象语法树节点基类""" - def __init__(self, node_type: str, children: Optional[List['ASTNode']] = None, - value: Optional[Any] = None, line: Optional[int] = None): - self.node_type = node_type # 'function', 'operator', 'field', 'number', 'expression' - self.children = children or [] - self.value = value - self.line = line - - def __str__(self) -> str: - return f"ASTNode({self.node_type}, {self.value}, line={self.line})" - - def __repr__(self) -> str: - return self.__str__() - -class ExpressionValidator: - """表达式验证器类""" - - def __init__(self): - """初始化词法分析器和语法分析器""" - # 构建词法分析器 - self.lexer = lex.lex(module=self, debug=False) - # 构建语法分析器 - self.parser = yacc.yacc(module=self, debug=False) - # 错误信息存储 - self.errors = [] - # Cache for unit inference (unit/scalar/category) - self._unit_cache: Dict[int, str] = {} - # Cache for derived category detection (bucket/group_cartesian_product outputs) - self._derived_category_cache: Dict[int, bool] = {} - - # 词法分析器规则 - tokens = ('FUNCTION', 'FIELD', 'NUMBER', 'LPAREN', 'RPAREN', - 'PLUS', 'MINUS', 'TIMES', 'DIVIDE', 'COMMA', 'CATEGORY', - 'EQUAL', 'ASSIGN', 'IDENTIFIER', 'STRING', 'GREATER', 'LESS', 'GREATEREQUAL', 'LESSEQUAL', 'NOTEQUAL', 'BOOLEAN') - - # 忽略空白字符 - t_ignore = ' \t\n' - - # 操作符 - 注意顺序很重要，长的操作符要放在前面 - t_PLUS = r'\+' - t_MINUS = r'-' - t_TIMES = r'\*' - t_DIVIDE = r'/' - t_LPAREN = r'\(' - t_RPAREN = r'\)' - t_COMMA = r',' - t_EQUAL = r'==' - t_NOTEQUAL = r'!=' - t_GREATEREQUAL = r'>=' - t_LESSEQUAL = r'<=' - t_GREATER = r'>' - t_LESS = r'<' - t_ASSIGN = r'=' - - # 数字（整数和浮点数） - def t_NUMBER(self, t): - r'\d+\.?\d*' - if '.' in t.value: - t.value = float(t.value) - else: - t.value = int(t.value) - return t - - # 字符串 - 需要放在所有其他标识符规则之前 - def t_STRING(self, t): - r"'[^']*'|\"[^\"]*\"" - # 去除引号 - t.value = t.value[1:-1] - return t - - # 函数和字段名 - def t_IDENTIFIER(self, t): - r'[a-zA-Z_][a-zA-Z0-9_]*' - # 检查是否为布尔值 - if t.value.lower() in {'true', 'false'}: - t.type = 'BOOLEAN' - t.value = t.value.lower() # 转换为小写以保持一致性 - else: - # 查看当前token后面的字符，判断是否为参数名（后面跟着'='） - lexpos = t.lexpos - next_chars = '' - if lexpos + len(t.value) < len(t.lexer.lexdata): - # 查看当前token后面的字符，跳过空格 - next_pos = lexpos + len(t.value) - while next_pos < len(t.lexer.lexdata) and t.lexer.lexdata[next_pos].isspace(): - next_pos += 1 - if next_pos < len(t.lexer.lexdata): - next_chars = t.lexer.lexdata[next_pos:next_pos+1] - - # 如果后面跟着'='，则为参数名 - if next_chars == '=': - t.type = 'IDENTIFIER' - # 如果后面跟着'('，则为函数名 - elif next_chars == '(': - t.type = 'FUNCTION' - t.value = t.value.lower() # 转换为小写以保持一致性 - # 检查是否为参数名（支持更多参数名） - elif t.value in {'std', 'k', 'lambda_min', 'lambda_max', 'target_tvr', 'range', 'buckets', 'lag', 'rettype', 'mode', 'nth', 'constant', 'percentage', 'driver', 'sigma', 'rate', 'scale', 'filter', 'lower', 'upper', 'target', 'dest', 'event', 'sensitivity', 'force', 'h', 't', 'period', 'stddev', 'factor', 'k', 'useStd', 'limit', 'gaussian', 'uniform', 'cauchy'}: - t.type = 'IDENTIFIER' - # 检查是否为函数名（不区分大小写） - elif t.value.lower() in supported_functions: - t.type = 'FUNCTION' - t.value = t.value.lower() # 转换为小写以保持一致性 - # 检查是否为有效类别 - elif t.value in valid_categories: - t.type = 'CATEGORY' - # 检查是否为字段名 - elif self._is_valid_field(t.value): - t.type = 'FIELD' - else: - # 其他标识符，保留为IDENTIFIER类型 - t.type = 'IDENTIFIER' - return t - - # 行号跟踪 - def t_newline(self, t): - r'\n+' - t.lexer.lineno += len(t.value) - - # 错误处理 - def t_error(self, t): - if t: - # 检查是否为非法字符 - if not re.match(r'[a-zA-Z0-9_\+\-\*/\(\)\,\s=<>!]', t.value[0]): - # 这是一个非法字符 - self.errors.append(f"非法字符 '{t.value[0]}' (行 {t.lexer.lineno})") - else: - # 这是一个非法标记 - self.errors.append(f"非法标记 '{t.value}' (行 {t.lexer.lineno})") - # 跳过这个字符，继续处理 - t.lexer.skip(1) - else: - self.errors.append("词法分析器到达文件末尾") - - # 语法分析器规则 - def p_expression(self, p): - """expression : comparison - | expression EQUAL comparison - | expression NOTEQUAL comparison - | expression GREATER comparison - | expression LESS comparison - | expression GREATEREQUAL comparison - | expression LESSEQUAL comparison""" - if len(p) == 2: - p[0] = p[1] - else: - p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) - - def p_comparison(self, p): - """comparison : term - | comparison PLUS term - | comparison MINUS term""" - if len(p) == 2: - p[0] = p[1] - else: - p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) - - def p_term(self, p): - """term : factor - | term TIMES factor - | term DIVIDE factor""" - if len(p) == 2: - p[0] = p[1] - else: - p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) - - def p_factor(self, p): - """factor : NUMBER - | STRING - | FIELD - | CATEGORY - | IDENTIFIER - | BOOLEAN - | MINUS factor - | LPAREN expression RPAREN - | function_call""" - if len(p) == 2: - # 数字、字符串、字段、类别或标识符 - if p.slice[1].type == 'NUMBER': - p[0] = ASTNode('number', value=p[1]) - elif p.slice[1].type == 'STRING': - p[0] = ASTNode('string', value=p[1]) - elif p.slice[1].type == 'FIELD': - p[0] = ASTNode('field', value=p[1]) - elif p.slice[1].type == 'CATEGORY': - p[0] = ASTNode('category', value=p[1]) - elif p.slice[1].type == 'BOOLEAN': - p[0] = ASTNode('boolean', value=p[1]) - elif p.slice[1].type == 'IDENTIFIER': - p[0] = ASTNode('identifier', value=p[1]) - else: - p[0] = p[1] - elif len(p) == 3: - # 一元负号 - p[0] = ASTNode('unop', [p[2]], {'op': p[1]}) - elif len(p) == 4: - # 括号表达式 - p[0] = p[2] - else: - # 函数调用 - p[0] = p[1] - - def p_function_call(self, p): - '''function_call : FUNCTION LPAREN args RPAREN''' - p[0] = ASTNode('function', p[3], p[1]) - - def p_args(self, p): - '''args : arg_list - | empty''' - if len(p) == 2 and p[1] is not None: - p[0] = p[1] - else: - p[0] = [] - - def p_arg_list(self, p): - '''arg_list : arg - | arg_list COMMA arg''' - if len(p) == 2: - p[0] = [p[1]] - else: - p[0] = p[1] + [p[3]] - - def p_arg(self, p): - '''arg : expression - | IDENTIFIER ASSIGN expression''' - if len(p) == 2: - p[0] = {'type': 'positional', 'value': p[1]} - else: - p[0] = {'type': 'named', 'name': p[1], 'value': p[3]} - - def p_empty(self, p): - '''empty :''' - p[0] = None - - # 语法错误处理 - def p_error(self, p): - if p: - self.errors.append(f"语法错误在位置 {p.lexpos}: 非法标记 '{p.value}'") - else: - self.errors.append("语法错误: 表达式不完整") - - def _is_valid_field(self, field_name: str) -> bool: - """检查字段名是否符合模式""" - for pattern in field_patterns: - if pattern.match(field_name): - return True - return False - - def validate_function(self, node: ASTNode, is_in_group_arg: bool = False) -> List[str]: - """验证函数调用的参数数量和类型""" - function_name = node.value - args = node.children - function_info = supported_functions.get(function_name) - - if not function_info: - return [f"未知函数: {function_name}"] - - # Custom validation for variadic functions with optional flags - if function_name == 'add': - return self._validate_add(args, is_in_group_arg) - - errors = [] - - # Keyword-only enforcement for optional parameters. - # If enabled, only the required leading arguments can be positional. - keyword_only_from = function_info.get('keyword_only_from') - if keyword_only_from is None and function_info.get('keyword_only'): - keyword_only_from = function_info.get('min_args', 0) - - # 检查参数数量 - if len(args) < function_info['min_args']: - errors.append(f"函数 {function_name} 需要至少 {function_info['min_args']} 个参数，但只提供了 {len(args)}") - elif len(args) > function_info['max_args']: - errors.append(f"函数 {function_name} 最多接受 {function_info['max_args']} 个参数，但提供了 {len(args)}") - - # 处理参数验证 - # 跟踪已使用的位置参数索引 - positional_index = 0 - - # 对于所有函数，支持命名参数 - for arg in args: - if isinstance(arg, dict): - if arg['type'] == 'named': - # 命名参数 - if 'param_names' in function_info and arg['name'] in function_info['param_names']: - # 查找参数在param_names中的索引 - param_index = function_info['param_names'].index(arg['name']) - if param_index < len(function_info['arg_types']): - expected_type = function_info['arg_types'][param_index] - arg_errors = self._validate_arg_type(arg['value'], expected_type, param_index, function_name, is_in_group_arg) - errors.extend(arg_errors) - # 对于winsorize函数，支持std和clip参数 - elif function_name == 'winsorize' and arg['name'] in ['std', 'clip']: - arg_errors = self._validate_arg_type(arg['value'], 'number', 0, function_name, is_in_group_arg) - errors.extend(arg_errors) - # 对于bucket函数，支持'range'和'buckets'参数 - elif function_name == 'bucket' and arg['name'] in ['range', 'buckets']: - # range和buckets参数应该是string类型 - arg_errors = self._validate_arg_type(arg['value'], 'string', 1, function_name, is_in_group_arg) - errors.extend(arg_errors) - else: - errors.append(f"函数 {function_name} 不存在参数 '{arg['name']}'") - elif arg['type'] == 'positional': - # 位置参数（字典形式） - if keyword_only_from is not None and positional_index >= keyword_only_from: - param_name = None - if 'param_names' in function_info and positional_index < len(function_info['param_names']): - param_name = function_info['param_names'][positional_index] - if param_name: - errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数 '{param_name}='") - else: - errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数") - else: - # 验证位置参数的类型 - if positional_index < len(function_info['arg_types']): - expected_type = function_info['arg_types'][positional_index] - arg_errors = self._validate_arg_type(arg['value'], expected_type, positional_index, function_name, is_in_group_arg) - errors.extend(arg_errors) - positional_index += 1 - else: - # 其他字典类型参数 - errors.append(f"参数 {positional_index+1} 格式错误") - positional_index += 1 - else: - # 位置参数（直接ASTNode形式） - if keyword_only_from is not None and positional_index >= keyword_only_from: - param_name = None - if 'param_names' in function_info and positional_index < len(function_info['param_names']): - param_name = function_info['param_names'][positional_index] - if param_name: - errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数 '{param_name}='") - else: - errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数") - else: - # 验证位置参数的类型 - if positional_index < len(function_info['arg_types']): - expected_type = function_info['arg_types'][positional_index] - arg_errors = self._validate_arg_type(arg, expected_type, positional_index, function_name, is_in_group_arg) - errors.extend(arg_errors) - positional_index += 1 - - return errors - - def _validate_arg_type(self, arg: ASTNode, expected_type: str, arg_index: int, function_name: str, is_in_group_arg: bool = False) -> List[str]: - """验证参数类型是否符合预期""" - errors = [] - - def _is_number_like(node: ASTNode) -> bool: - if node is None: - return False - if node.node_type == 'number': - return True - if node.node_type == 'unop' and isinstance(node.value, dict) and node.value.get('op') in {'-', '+'}: - if node.children and hasattr(node.children[0], 'node_type'): - return _is_number_like(node.children[0]) - return False - - # Unit compatibility check - # bucket()/group_cartesian_product() output a derived category (grouping key). - # It can only be consumed where a category/grouping key is expected. - if self._is_derived_category(arg) and expected_type != 'category': - errors.append( - f"Incompatible unit for input of \"{function_name}\" at index {arg_index}, expected \"Unit[]\", found \"Unit[Group:1]\"" - ) - return errors - - # 首先检查是否是group类型字段，如果是则只能用于Group类型函数 - # 但是如果当前函数是group_xxx或在group函数的参数链中，则允许使用 - if arg.node_type == 'category' and arg.value in group_fields: - if not (function_name.startswith('group_') or is_in_group_arg): - errors.append(f"Group类型字段 '{arg.value}' 只能用于Group类型函数的参数中") - - # 然后验证参数类型是否符合预期 - if expected_type == 'expression': - # 表达式可以是任何有效的AST节点 - pass - elif expected_type == 'number': - # 允许 -1 这类一元负号数字常量（解析为 unop(number)） - if not _is_number_like(arg): - errors.append(f"参数 {arg_index+1} 应该是一个数字，但得到 {arg.node_type}") - elif expected_type == 'boolean': - # 布尔值可以是 true/false 或数字（0/1） - if arg.node_type not in {'boolean', 'number'}: - errors.append(f"参数 {arg_index+1} 应该是一个布尔值（true/false 或 0/1），但得到 {arg.node_type}") - elif expected_type == 'field': - if arg.node_type != 'field' and arg.node_type != 'category': - # 允许field或category作为字段参数 - errors.append(f"参数 {arg_index+1} 应该是一个字段，但得到 {arg.node_type}") - elif arg.node_type == 'field' and not self._is_valid_field(arg.value): - errors.append(f"无效的字段名: {arg.value}") - elif expected_type == 'category': - if not function_name.startswith('group_'): - # 非group函数的category参数必须是category类型且在valid_categories中 - if arg.node_type != 'category': - errors.append(f"参数 {arg_index+1} 应该是一个类别，但得到 {arg.node_type}") - elif arg.value not in valid_categories: - errors.append(f"无效的类别: {arg.value}") - # group函数的category参数可以是任何类型（field、category等），不进行类型校验 - - return errors - - def _infer_unit(self, node: ASTNode) -> str: - """Infer the Unit kind of an AST node. - - Returns: - 'unit' - regular numeric time-series Unit[] - 'scalar' - literals (numbers/booleans/strings) - 'category' - category/grouping keys (industry/sector or derived via bucket/cartesian) - """ - if node is None: - return 'unit' - - cache_key = id(node) - cached = self._unit_cache.get(cache_key) - if cached is not None: - return cached - - unit = 'unit' - - if node.node_type in {'number', 'boolean', 'string'}: - unit = 'scalar' - elif node.node_type in {'field', 'identifier'}: - unit = 'unit' - elif node.node_type == 'category': - unit = 'category' - elif node.node_type in {'unop', 'binop'}: - child_units = [self._infer_unit(child) for child in node.children if hasattr(child, 'node_type')] - unit = 'category' if 'category' in child_units else 'unit' - elif node.node_type == 'function': - fname = node.value - if fname in {'bucket', 'group_cartesian_product'}: - unit = 'category' - else: - first_arg = None - for child in node.children: - if isinstance(child, dict): - if child.get('type') == 'positional': - first_arg = child.get('value') - break - else: - first_arg = child - break - if hasattr(first_arg, 'node_type'): - unit = self._infer_unit(first_arg) - else: - unit = 'unit' - - self._unit_cache[cache_key] = unit - return unit - - def _is_derived_category(self, node: ASTNode) -> bool: - """Return True if node is a derived category/grouping key (e.g., bucket/cartesian output).""" - if node is None: - return False - - cache_key = id(node) - cached = self._derived_category_cache.get(cache_key) - if cached is not None: - return cached - - derived = False - if node.node_type == 'function': - if node.value in {'bucket', 'group_cartesian_product'}: - derived = True - else: - function_info = supported_functions.get(node.value, {}) - arg_types = function_info.get('arg_types', []) - param_names = function_info.get('param_names', []) - - positional_index = 0 - for child in node.children: - if isinstance(child, dict): - if child.get('type') == 'named': - name = child.get('name') - value = child.get('value') - - expected_type = None - if name in param_names: - param_index = param_names.index(name) - if param_index < len(arg_types): - expected_type = arg_types[param_index] - - if expected_type == 'category': - continue - - if self._is_derived_category(value): - derived = True - break - elif child.get('type') == 'positional': - value = child.get('value') - expected_type = arg_types[positional_index] if positional_index < len(arg_types) else None - - if expected_type != 'category' and self._is_derived_category(value): - derived = True - break - positional_index += 1 - else: - expected_type = arg_types[positional_index] if positional_index < len(arg_types) else None - if expected_type != 'category' and self._is_derived_category(child): - derived = True - break - positional_index += 1 - elif node.node_type in {'unop', 'binop'}: - derived = any( - self._is_derived_category(child) - for child in node.children - if hasattr(child, 'node_type') - ) - - self._derived_category_cache[cache_key] = derived - return derived - - def _validate_add(self, args: List[Any], is_in_group_arg: bool = False) -> List[str]: - """Validate add(x, y, ..., filter=false). - - Rules: - - At least 2 positional expression terms. - - Optional filter flag can be provided as: - - named argument: filter= - - last positional argument: or 0/1 - """ - errors: List[str] = [] - - if len(args) < 2: - return [f"函数 add 需要至少 2 个参数，但只提供了 {len(args)}"] - - named_filter_nodes: List[ASTNode] = [] - positional_nodes: List[ASTNode] = [] - - for arg in args: - if isinstance(arg, dict) and arg.get('type') == 'named': - name = arg.get('name') - value = arg.get('value') - if name != 'filter': - errors.append(f"函数 add 不存在参数 '{name}'") - continue - if not hasattr(value, 'node_type'): - errors.append("函数 add 的参数 filter 格式错误") - continue - named_filter_nodes.append(value) - elif isinstance(arg, dict) and arg.get('type') == 'positional': - value = arg.get('value') - if hasattr(value, 'node_type'): - positional_nodes.append(value) - else: - errors.append("函数 add 的位置参数格式错误") - elif hasattr(arg, 'node_type'): - positional_nodes.append(arg) - else: - errors.append("函数 add 的参数格式错误") - - if len(named_filter_nodes) > 1: - errors.append("函数 add 的参数 'filter' 只能出现一次") - - positional_filter_node: Optional[ASTNode] = None - # Only infer a positional filter flag when: - # - no named filter is provided - # - there are at least 3 positional args (x, y, filter) - # - the last arg is boolean or numeric 0/1 - if not named_filter_nodes and len(positional_nodes) >= 3: - last = positional_nodes[-1] - if last.node_type == 'boolean' or (last.node_type == 'number' and last.value in {0, 1}): - positional_filter_node = positional_nodes.pop() - - if len(positional_nodes) < 2: - errors.append(f"函数 add 需要至少 2 个输入项（不含filter），但只提供了 {len(positional_nodes)}") - - for idx, node in enumerate(positional_nodes): - errors.extend(self._validate_arg_type(node, 'expression', idx, 'add', is_in_group_arg)) - - if positional_filter_node is not None and named_filter_nodes: - errors.append("函数 add 的 filter 不能同时用位置参数和命名参数传递") - if positional_filter_node is not None: - errors.extend(self._validate_arg_type(positional_filter_node, 'boolean', len(positional_nodes), 'add', is_in_group_arg)) - if named_filter_nodes: - errors.extend(self._validate_arg_type(named_filter_nodes[0], 'boolean', len(positional_nodes), 'add', is_in_group_arg)) - - return errors - - def validate_ast(self, ast: Optional[ASTNode], is_in_group_arg: bool = False) -> List[str]: - """递归验证抽象语法树""" - if not ast: - return ["无法解析表达式"] - - errors = [] - - # 根据节点类型进行验证 - if ast.node_type == 'function': - # 检查当前函数是否是group函数 - is_group_function = ast.value.startswith('group_') - # 确定当前是否在group函数的参数链中 - current_in_group_arg = is_in_group_arg or is_group_function - # 验证函数 - function_errors = self.validate_function(ast, current_in_group_arg) - errors.extend(function_errors) - - # 递归验证子节点时使用current_in_group_arg - for child in ast.children: - if isinstance(child, dict): - # 命名参数，验证其值 - if 'value' in child and hasattr(child['value'], 'node_type'): - child_errors = self.validate_ast(child['value'], current_in_group_arg) - errors.extend(child_errors) - elif hasattr(child, 'node_type'): - child_errors = self.validate_ast(child, current_in_group_arg) - errors.extend(child_errors) - elif ast.node_type in ['unop', 'binop']: - # 对操作符的子节点进行验证 - for child in ast.children: - if hasattr(child, 'node_type'): - child_errors = self.validate_ast(child, is_in_group_arg) - errors.extend(child_errors) - elif ast.node_type == 'field': - # 验证字段名 - if not self._is_valid_field(ast.value): - errors.append(f"无效的字段名: {ast.value}") - else: - # 递归验证子节点 - for child in ast.children: - if isinstance(child, dict): - # 命名参数，验证其值 - if 'value' in child and hasattr(child['value'], 'node_type'): - child_errors = self.validate_ast(child['value'], is_in_group_arg) - errors.extend(child_errors) - elif hasattr(child, 'node_type'): - child_errors = self.validate_ast(child, is_in_group_arg) - errors.extend(child_errors) - - return errors - - def _process_semicolon_expression(self, expression: str) -> Tuple[bool, str]: - """处理带有分号的表达式，将其转换为不带分号的简化形式 - - Args: - expression: 要处理的表达式字符串 - - Returns: - Tuple[bool, str]: (是否成功, 转换后的表达式或错误信息) - """ - def _top_level_equals_positions(stmt: str) -> List[int]: - """返回所有“顶层赋值”等号位置。 - - 仅统计括号外（()[]{})、引号外、且不属于比较操作符(==,!=,<=,>=)的 '='。 - 这样可以避免把关键字参数（如 rettype=0）误判为赋值语句。 - """ - positions: List[int] = [] - paren_depth = 0 - bracket_depth = 0 - brace_depth = 0 - in_single_quote = False - in_double_quote = False - escape = False - - for i, ch in enumerate(stmt): - if escape: - escape = False - continue - if ch == '\\': - escape = True - continue - - if in_single_quote: - if ch == "'": - in_single_quote = False - continue - if in_double_quote: - if ch == '"': - in_double_quote = False - continue - - if ch == "'": - in_single_quote = True - continue - if ch == '"': - in_double_quote = True - continue - - if ch == '(': - paren_depth += 1 - continue - if ch == ')': - paren_depth = max(0, paren_depth - 1) - continue - if ch == '[': - bracket_depth += 1 - continue - if ch == ']': - bracket_depth = max(0, bracket_depth - 1) - continue - if ch == '{': - brace_depth += 1 - continue - if ch == '}': - brace_depth = max(0, brace_depth - 1) - continue - - if paren_depth or bracket_depth or brace_depth: - continue - - if ch != '=': - continue - - prev_ch = stmt[i - 1] if i > 0 else '' - next_ch = stmt[i + 1] if i + 1 < len(stmt) else '' - if prev_ch in ['=', '!', '<', '>'] or next_ch == '=': - continue - - positions.append(i) - - return positions - - def _keyword_arg_names(stmt: str): - """提取函数调用中的命名参数名（如 rettype=0 中的 rettype）。 - - 只收集括号/中括号/大括号内部出现的 name= 形式，避免把脚本级赋值误当作命名参数。 - """ - names = set() - paren_depth = 0 - bracket_depth = 0 - brace_depth = 0 - in_single_quote = False - in_double_quote = False - escape = False - - i = 0 - while i < len(stmt): - ch = stmt[i] - - if escape: - escape = False - i += 1 - continue - if ch == '\\': - escape = True - i += 1 - continue - - if in_single_quote: - if ch == "'": - in_single_quote = False - i += 1 - continue - if in_double_quote: - if ch == '"': - in_double_quote = False - i += 1 - continue - - if ch == "'": - in_single_quote = True - i += 1 - continue - if ch == '"': - in_double_quote = True - i += 1 - continue - - if ch == '(': - paren_depth += 1 - i += 1 - continue - if ch == ')': - paren_depth = max(0, paren_depth - 1) - i += 1 - continue - if ch == '[': - bracket_depth += 1 - i += 1 - continue - if ch == ']': - bracket_depth = max(0, bracket_depth - 1) - i += 1 - continue - if ch == '{': - brace_depth += 1 - i += 1 - continue - if ch == '}': - brace_depth = max(0, brace_depth - 1) - i += 1 - continue - - inside_container = bool(paren_depth or bracket_depth or brace_depth) - - if inside_container and (ch.isalpha() or ch == '_'): - start = i - i += 1 - while i < len(stmt) and (stmt[i].isalnum() or stmt[i] == '_'): - i += 1 - name = stmt[start:i] - - j = i - while j < len(stmt) and stmt[j].isspace(): - j += 1 - - if j < len(stmt) and stmt[j] == '=': - next_ch = stmt[j + 1] if j + 1 < len(stmt) else '' - if next_ch != '=': - names.add(name.lower()) - continue - - i += 1 - - return names - - # 检查表达式是否以分号结尾 - if expression.strip().endswith(';'): - return False, "表达式不能以分号结尾" - - # 分割表达式为语句列表 - statements = [stmt.strip() for stmt in expression.split(';') if stmt.strip()] - if not statements: - return False, "表达式不能为空" - - # 存储变量赋值 - variables = {} - - # 处理每个赋值语句（除了最后一个） - for i, stmt in enumerate(statements[:-1]): - eq_positions = _top_level_equals_positions(stmt) - if not eq_positions: - return False, f"第{i+1}个语句必须是赋值语句（使用=符号）" - if len(eq_positions) > 1: - return False, f"第{i+1}个语句只能包含一个赋值符号（=）" - - real_equals_pos = eq_positions[0] - - # 分割变量名和值 - var_name = stmt[:real_equals_pos].strip() - var_value = stmt[real_equals_pos + 1:].strip() - - # 检查变量名是否有效 - if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', var_name): - return False, f"第{i+1}个语句的变量名'{var_name}'无效，只能包含字母、数字和下划线，且不能以数字开头" - - var_name_lower = var_name.lower() # 变量名不区分大小写 - - # 检查变量名是否在后续表达式中使用 - # 这里不需要，因为后面的表达式会检查 - - # 检查变量值中使用的变量是否已经定义 - # 简单检查：提取所有可能的变量名 - kw_names = _keyword_arg_names(var_value) - used_vars = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', var_value) - for used_var in used_vars: - used_var_lower = used_var.lower() - if used_var_lower in kw_names: - continue - if used_var_lower not in variables: - # 检查是否是函数名 - if used_var not in supported_functions: - # 对于单个字母或简单单词，不自动视为字段名，要求先定义 - if len(used_var) <= 2: - return False, f"第{i+1}个语句中使用的变量'{used_var}'未在之前定义" - # 对于较长的字段名，仍然允许作为字段名 - elif not self._is_valid_field(used_var): - return False, f"第{i+1}个语句中使用的变量'{used_var}'未在之前定义" - - # 将之前定义的变量替换到当前值中 - for existing_var, existing_val in variables.items(): - # 使用单词边界匹配，避免替换到其他单词的一部分 - var_value = re.sub(rf'\b{existing_var}\b', existing_val, var_value) - - # 存储变量 - variables[var_name_lower] = var_value - - # 处理最后一个语句（实际的表达式） - final_stmt = statements[-1] - - # 检查最后一个语句是否是赋值语句 - if _top_level_equals_positions(final_stmt): - return False, "最后一个语句不能是赋值语句" - - # 检查最后一个语句中使用的变量是否已经定义 - kw_names = _keyword_arg_names(final_stmt) - used_vars = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', final_stmt) - for used_var in used_vars: - used_var_lower = used_var.lower() - if used_var_lower in kw_names: - continue - if used_var_lower not in variables: - # 检查是否是函数名 - if used_var not in supported_functions: - # 对于单个字母或简单单词，不自动视为字段名，要求先定义 - if len(used_var) <= 2: - return False, f"最后一个语句中使用的变量'{used_var}'未在之前定义" - - # 允许直接使用字段名/类别名（如 close/industry） - if self._is_valid_field(used_var) or used_var_lower in valid_categories or used_var_lower in group_fields: - continue - - return False, f"最后一个语句中使用的变量'{used_var}'未在之前定义" - - # 将变量替换到最后一个表达式中 - final_expr = final_stmt - for var_name, var_value in variables.items(): - final_expr = re.sub(rf'\b{var_name}\b', var_value, final_expr) - - return True, final_expr - - def check_expression(self, expression: str) -> Dict[str, Any]: - """ - 检查表达式格式是否正确 - - Args: - expression: 要验证的表达式字符串 - - Returns: - 包含验证结果的字典 - """ - # 重置错误列表 - self.errors = [] - # Reset unit inference cache for this expression - self._unit_cache = {} - self._derived_category_cache = {} - - try: - expression = expression.strip() - if not expression: - return { - 'valid': False, - 'errors': ['表达式不能为空'], - 'tokens': [], - 'ast': None - } - - # 处理带有分号的表达式 - if ';' in expression: - success, result = self._process_semicolon_expression(expression) - if not success: - return { - 'valid': False, - 'errors': [result], - 'tokens': [], - 'ast': None - } - expression = result - - # 重置词法分析器的行号 - self.lexer.lineno = 1 - - # 词法分析（用于调试） - self.lexer.input(expression) - tokens = [] - # 调试：打印识别的标记 - # print(f"\n调试 - 表达式: {expression}") - # print("识别的标记:") - for token in self.lexer: - # print(f" - 类型: {token.type}, 值: '{token.value}', 位置: {token.lexpos}") - tokens.append(token) - - # 重新设置词法分析器的输入，以便语法分析器使用 - self.lexer.input(expression) - self.lexer.lineno = 1 - - # 语法分析 - ast = self.parser.parse(expression, lexer=self.lexer) - - # 验证AST - validation_errors = self.validate_ast(ast) - - # 合并所有错误 - all_errors = self.errors + validation_errors - - # 检查括号是否匹配 - bracket_count = 0 - for char in expression: - if char == '(': - bracket_count += 1 - elif char == ')': - bracket_count -= 1 - if bracket_count < 0: - all_errors.append("括号不匹配: 右括号过多") - break - if bracket_count > 0: - all_errors.append("括号不匹配: 左括号过多") - - return { - 'valid': len(all_errors) == 0, - 'errors': all_errors, - 'tokens': tokens, - 'ast': ast - } - except Exception as e: - return { - 'valid': False, - 'errors': [f"解析错误: {str(e)}"], - 'tokens': [], - 'ast': None - } - - - +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +表达式验证器 - 使用抽象语法树验证字符串表达式格式是否正确 + +本模块实现了一个能够检测字符串表达式格式是否正确的系统，基于PLY(Python Lex-Yacc) +构建词法分析器和语法分析器，识别表达式中的操作符、函数和字段，并验证其格式正确性。 +""" + +import re +import sys +import json +import os +from typing import List, Dict, Any, Optional, Tuple + +import logging + +logger = logging.getLogger(__name__) + +# 尝试导入PLY库，如果不存在则提供安装提示 +try: + import ply.lex as lex + import ply.yacc as yacc +except ImportError: + logger.error("错误: 需要安装PLY库。请运行 'pip install ply' 来安装。") + sys.exit(1) + +# 1. 定义支持的操作符和函数 +supported_functions = { + # Group 类别函数 + 'group_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + # group_mean(x, w, group) + 'group_mean': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, + 'group_median': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_rank': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_vector_proj': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, + 'group_normalize': {'min_args': 2, 'max_args': 5, 'arg_types': ['expression', 'category', 'expression', 'expression', 'expression']}, + 'group_extra': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'category']}, + 'group_backfill': {'min_args': 3, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['x', 'cat', 'days', 'std']}, + 'group_scale': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_count': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_zscore': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_std_dev': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_sum': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_neutralize': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'category']}, + 'group_multi_regression': {'min_args': 4, 'max_args': 9, 'arg_types': ['expression'] * 9}, + 'group_cartesian_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['category', 'category']}, + 'combo_a': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, + + # Transformational 类别函数 + 'right_tail': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'bucket': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 第二个参数可以是string类型的range参数 + 'tail': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression']}, + 'left_tail': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'trade_when': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, + 'generate_stats': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + + # Cross Sectional 类别函数 + 'winsorize': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'std']}, + 'rank': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'regression_proj': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'vector_neut': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'regression_neut': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'multi_regression': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # 支持多个自变量 + + # Time Series 类别函数 + 'ts_std_dev': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_mean': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_delay': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_corr': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, + 'ts_zscore': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_returns': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'mode'], 'keyword_only': True}, + 'ts_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_backfill(x, d) + 'ts_backfill': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'd']}, + 'days_from_last_change': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'last_diff_value': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_scale(x, d, constant=0) + 'ts_scale': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'constant'], 'keyword_only': True}, + # Platform: ts_entropy(x, d) + 'ts_entropy': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'd']}, + 'ts_step': {'min_args': 1, 'max_args': 1, 'arg_types': ['number']}, + 'ts_sum': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_co_kurtosis': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, + 'inst_tvr': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_decay_exp_window': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'factor'], 'keyword_only': True}, + 'ts_av_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_kurtosis': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_min_max_diff(x, d, f=0.5) + 'ts_min_max_diff': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'f'], 'keyword_only': True}, + 'ts_arg_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_max': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_min_max_cps(x, d, f=2) + 'ts_min_max_cps': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'f'], 'keyword_only': True}, + # Platform: ts_rank(x, d, constant=0) + 'ts_rank': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'constant'], 'keyword_only': True}, + 'ts_ir': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_theilsen': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, + # Platform: hump_decay(x, p=0) + 'hump_decay': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'p'], 'keyword_only': True}, + # Platform: ts_weighted_decay(x, k=0.5) + 'ts_weighted_decay': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'k'], 'keyword_only': True}, + # Platform: ts_quantile(x, d, driver="gaussian") + 'ts_quantile': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'string'], 'param_names': ['x', 'd', 'driver'], 'keyword_only': True}, + 'ts_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_count_nans': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_covariance': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, + 'ts_co_skewness': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number']}, + 'ts_min_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_decay_linear(x, d, dense=false) + 'ts_decay_linear': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'boolean'], 'param_names': ['x', 'd', 'dense'], 'keyword_only': True}, + # Platform: jump_decay(x, d, sensitivity=0.5, force=0.1) + 'jump_decay': {'min_args': 2, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'd', 'sensitivity', 'force'], 'keyword_only': True}, + # Platform: ts_moment(x, d, k=0) + 'ts_moment': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'k'], 'keyword_only': True}, + 'ts_arg_min': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_regression': {'min_args': 3, 'max_args': 5, 'arg_types': ['expression', 'expression', 'number', 'number', 'number'], 'param_names': ['y', 'x', 'd', 'lag', 'rettype'], 'keyword_only': True}, + 'ts_skewness': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_max_diff': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'kth_element': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'number', 'number']}, + 'hump': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'hump']}, + 'ts_median': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_delta': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + # Platform: ts_poly_regression(y, x, d, k=1) and k must be keyword if provided + 'ts_poly_regression': {'min_args': 3, 'max_args': 4, 'arg_types': ['expression', 'expression', 'number', 'number'], 'param_names': ['y', 'x', 'd', 'k'], 'keyword_only': True, 'keyword_only_from': 3}, + 'ts_target_tvr_decay': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, + 'ts_target_tvr_delta_limit': {'min_args': 2, 'max_args': 5, 'arg_types': ['expression', 'expression', 'number', 'number', 'number'], 'param_names': ['x', 'y', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, + 'ts_target_tvr_hump': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'number', 'number', 'number'], 'param_names': ['x', 'lambda_min', 'lambda_max', 'target_tvr'], 'keyword_only': True}, + # Platform: ts_delta_limit(x, y, limit_volume=0.1) + 'ts_delta_limit': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'limit_volume'], 'keyword_only': True}, + + # Special 类别函数 + 'inst_pnl': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'self_corr': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'in': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 + 'universe_size': {'min_args': 0, 'max_args': 0, 'arg_types': []}, + + # Missing functions from operators.py + 'quantile': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'driver', 'sigma']}, # quantile(x, driver = gaussian, sigma = 1.0) + 'normalize': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'boolean', 'number']}, # normalize(x, useStd = false, limit = 0.0) + 'zscore': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # zscore(x) + + # Logical 类别函数 + 'or': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 + 'and': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # 注意：这是关键字 + 'not': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # 注意：这是关键字 + 'is_nan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'is_not_nan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'less': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'greater': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'is_finite': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'if_else': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression']}, + 'not_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'less_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'greater_equal': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + + # Vector 类别函数 + 'vec_kurtosis': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_min': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_count': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_sum': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_skewness': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_max': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_avg': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_range': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_choose': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'nth']}, + 'vec_powersum': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'constant']}, + 'vec_stddev': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_percentage': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'percentage']}, + 'vec_ir': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'vec_norm': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'ts_percentage': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'd', 'percentage']}, + 'signed_power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + 'ts_product': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number']}, + + # Additional functions from test cases + 'rank_by_side': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'number', 'number'], 'param_names': ['x', 'rate', 'scale']}, + 'log_diff': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'nan_mask': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, + 'ts_partial_corr': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'number']}, + 'ts_triple_corr': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'number']}, + 'clamp': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'lower', 'upper']}, + 'keep': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'condition', 'period']}, + 'replace': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'target', 'dest']}, + 'filter': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'h', 't']}, + 'one_side': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'string'], 'param_names': ['x', 'side']}, + 'scale_down': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'number'], 'param_names': ['x', 'constant']}, + + # Arithmetic 类别函数 + # add(x, y, ..., filter=false) + # NOTE: add() is variadic (>=2 terms) with an optional boolean filter flag. + # We validate it with custom logic in validate_function(). + 'add': {'min_args': 2, 'max_args': 101, 'arg_types': ['expression'] * 101}, + 'multiply': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 99 + ['boolean'], 'param_names': ['x', 'y', 'filter']}, # multiply(x, y, ..., filter=false) + 'sign': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'subtract': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'boolean']}, # subtract(x, y, filter=false) + 'pasteurize': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'log': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'purify': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'arc_tan': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'max': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # max(x, y, ...) + 'to_nan': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'boolean']}, # to_nan(x, value=0, reverse=false) + 'abs': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'sigmoid': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'divide': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # divide(x, y) + 'min': {'min_args': 2, 'max_args': 100, 'arg_types': ['expression'] * 100}, # min(x, y, ...) + 'tanh': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'nan_out': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'lower', 'upper']}, # nan_out(x, lower=0, upper=0) + 'signed_power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # signed_power(x, y) + 'inverse': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'round': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'sqrt': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 's_log_1p': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'reverse': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, # -x + 'power': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression']}, # power(x, y) + 'densify': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + 'floor': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression']}, + # Appended missing operators + 'arc_cos': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, + 'arc_sin': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, + 'ceiling': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, + 'exp': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, + 'fraction': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['x']}, + 'round_down': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'f']}, + 'is_not_finite': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'negate': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'ts_rank_gmean_amean_diff': {'min_args': 5, 'max_args': 5, 'arg_types': ['expression', 'expression', 'expression', 'expression', 'number'], 'param_names': ['input1', 'input2', 'input3', '...', 'd']}, + 'ts_vector_neut': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'd']}, + 'ts_vector_proj': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'number'], 'param_names': ['x', 'y', 'd']}, + 'scale': {'min_args': 1, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['x', 'scale', 'longscale', 'shortscale']}, + 'generalized_rank': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['open', 'm']}, + 'rank_gmean_amean_diff': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['input1', 'input2', 'input3', '...']}, + 'truncate': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'maxPercent']}, + 'vector_proj': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'y']}, + 'vec_filter': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['vec', 'value']}, + 'group_coalesce': {'min_args': 4, 'max_args': 4, 'arg_types': ['expression', 'expression', 'expression', 'expression'], 'param_names': ['original_group', 'group2', 'group3', '…']}, + 'group_percentage': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'category', 'expression'], 'param_names': ['x', 'group', 'percentage']}, + 'group_vector_neut': {'min_args': 3, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['x', 'y', 'g']}, + 'convert': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['x', 'mode']}, + 'reduce_avg': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, + 'reduce_choose': {'min_args': 2, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['input', 'nth', 'ignoreNan']}, + 'reduce_count': {'min_args': 2, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, + 'reduce_ir': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_kurtosis': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_max': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_min': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_norm': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_percentage': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'percentage']}, + 'reduce_powersum': {'min_args': 1, 'max_args': 3, 'arg_types': ['expression', 'expression', 'expression'], 'param_names': ['input', 'constant', 'precise']}, + 'reduce_range': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_skewness': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, + 'reduce_stddev': {'min_args': 1, 'max_args': 2, 'arg_types': ['expression', 'expression'], 'param_names': ['input', 'threshold']}, + 'reduce_sum': {'min_args': 1, 'max_args': 1, 'arg_types': ['expression'], 'param_names': ['input']}, +} + +# 2. 定义group类型字段 +group_fields = { + 'sector', 'subindustry', 'industry', 'exchange', 'country', 'market' +} + +# 3. 有效类别集合 +valid_categories = group_fields + +# 4. 字段命名模式 - 只校验字段是不是数字字母下划线组成 +field_patterns = [ + re.compile(r'^[a-zA-Z0-9_]+$'), # 只允许数字、字母和下划线组成的字段名 +] + +# 4. 抽象语法树节点类型 +class ASTNode: + """抽象语法树节点基类""" + def __init__(self, node_type: str, children: Optional[List['ASTNode']] = None, + value: Optional[Any] = None, line: Optional[int] = None): + self.node_type = node_type # 'function', 'operator', 'field', 'number', 'expression' + self.children = children or [] + self.value = value + self.line = line + + def __str__(self) -> str: + return f"ASTNode({self.node_type}, {self.value}, line={self.line})" + + def __repr__(self) -> str: + return self.__str__() + +class ExpressionValidator: + """表达式验证器类""" + + def __init__(self): + """初始化词法分析器和语法分析器""" + # 构建词法分析器 + self.lexer = lex.lex(module=self, debug=False) + # 构建语法分析器 + self.parser = yacc.yacc(module=self, debug=False) + # 错误信息存储 + self.errors = [] + # Cache for unit inference (unit/scalar/category) + self._unit_cache: Dict[int, str] = {} + # Cache for derived category detection (bucket/group_cartesian_product outputs) + self._derived_category_cache: Dict[int, bool] = {} + + # 词法分析器规则 + tokens = ('FUNCTION', 'FIELD', 'NUMBER', 'LPAREN', 'RPAREN', + 'PLUS', 'MINUS', 'TIMES', 'DIVIDE', 'COMMA', 'CATEGORY', + 'EQUAL', 'ASSIGN', 'IDENTIFIER', 'STRING', 'GREATER', 'LESS', 'GREATEREQUAL', 'LESSEQUAL', 'NOTEQUAL', 'BOOLEAN') + + # 忽略空白字符 + t_ignore = ' \t\n' + + # 操作符 - 注意顺序很重要，长的操作符要放在前面 + t_PLUS = r'\+' + t_MINUS = r'-' + t_TIMES = r'\*' + t_DIVIDE = r'/' + t_LPAREN = r'\(' + t_RPAREN = r'\)' + t_COMMA = r',' + t_EQUAL = r'==' + t_NOTEQUAL = r'!=' + t_GREATEREQUAL = r'>=' + t_LESSEQUAL = r'<=' + t_GREATER = r'>' + t_LESS = r'<' + t_ASSIGN = r'=' + + # 数字（整数和浮点数） + def t_NUMBER(self, t): + r'\d+\.?\d*' + if '.' in t.value: + t.value = float(t.value) + else: + t.value = int(t.value) + return t + + # 字符串 - 需要放在所有其他标识符规则之前 + def t_STRING(self, t): + r"'[^']*'|\"[^\"]*\"" + # 去除引号 + t.value = t.value[1:-1] + return t + + # 函数和字段名 + def t_IDENTIFIER(self, t): + r'[a-zA-Z_][a-zA-Z0-9_]*' + # 检查是否为布尔值 + if t.value.lower() in {'true', 'false'}: + t.type = 'BOOLEAN' + t.value = t.value.lower() # 转换为小写以保持一致性 + else: + # 查看当前token后面的字符，判断是否为参数名（后面跟着'='） + lexpos = t.lexpos + next_chars = '' + if lexpos + len(t.value) < len(t.lexer.lexdata): + # 查看当前token后面的字符，跳过空格 + next_pos = lexpos + len(t.value) + while next_pos < len(t.lexer.lexdata) and t.lexer.lexdata[next_pos].isspace(): + next_pos += 1 + if next_pos < len(t.lexer.lexdata): + next_chars = t.lexer.lexdata[next_pos:next_pos+1] + + # 如果后面跟着'='，则为参数名 + if next_chars == '=': + t.type = 'IDENTIFIER' + # 如果后面跟着'('，则为函数名 + elif next_chars == '(': + t.type = 'FUNCTION' + t.value = t.value.lower() # 转换为小写以保持一致性 + # 检查是否为参数名（支持更多参数名） + elif t.value in {'std', 'k', 'lambda_min', 'lambda_max', 'target_tvr', 'range', 'buckets', 'lag', 'rettype', 'mode', 'nth', 'constant', 'percentage', 'driver', 'sigma', 'rate', 'scale', 'filter', 'lower', 'upper', 'target', 'dest', 'event', 'sensitivity', 'force', 'h', 't', 'period', 'stddev', 'factor', 'k', 'useStd', 'limit', 'gaussian', 'uniform', 'cauchy'}: + t.type = 'IDENTIFIER' + # 检查是否为函数名（不区分大小写） + elif t.value.lower() in supported_functions: + t.type = 'FUNCTION' + t.value = t.value.lower() # 转换为小写以保持一致性 + # 检查是否为有效类别 + elif t.value in valid_categories: + t.type = 'CATEGORY' + # 检查是否为字段名 + elif self._is_valid_field(t.value): + t.type = 'FIELD' + else: + # 其他标识符，保留为IDENTIFIER类型 + t.type = 'IDENTIFIER' + return t + + # 行号跟踪 + def t_newline(self, t): + r'\n+' + t.lexer.lineno += len(t.value) + + # 错误处理 + def t_error(self, t): + if t: + # 检查是否为非法字符 + if not re.match(r'[a-zA-Z0-9_\+\-\*/\(\)\,\s=<>!]', t.value[0]): + # 这是一个非法字符 + self.errors.append(f"非法字符 '{t.value[0]}' (行 {t.lexer.lineno})") + else: + # 这是一个非法标记 + self.errors.append(f"非法标记 '{t.value}' (行 {t.lexer.lineno})") + # 跳过这个字符，继续处理 + t.lexer.skip(1) + else: + self.errors.append("词法分析器到达文件末尾") + + # 语法分析器规则 + def p_expression(self, p): + """expression : comparison + | expression EQUAL comparison + | expression NOTEQUAL comparison + | expression GREATER comparison + | expression LESS comparison + | expression GREATEREQUAL comparison + | expression LESSEQUAL comparison""" + if len(p) == 2: + p[0] = p[1] + else: + p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) + + def p_comparison(self, p): + """comparison : term + | comparison PLUS term + | comparison MINUS term""" + if len(p) == 2: + p[0] = p[1] + else: + p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) + + def p_term(self, p): + """term : factor + | term TIMES factor + | term DIVIDE factor""" + if len(p) == 2: + p[0] = p[1] + else: + p[0] = ASTNode('binop', [p[1], p[3]], {'op': p[2]}) + + def p_factor(self, p): + """factor : NUMBER + | STRING + | FIELD + | CATEGORY + | IDENTIFIER + | BOOLEAN + | MINUS factor + | LPAREN expression RPAREN + | function_call""" + if len(p) == 2: + # 数字、字符串、字段、类别或标识符 + if p.slice[1].type == 'NUMBER': + p[0] = ASTNode('number', value=p[1]) + elif p.slice[1].type == 'STRING': + p[0] = ASTNode('string', value=p[1]) + elif p.slice[1].type == 'FIELD': + p[0] = ASTNode('field', value=p[1]) + elif p.slice[1].type == 'CATEGORY': + p[0] = ASTNode('category', value=p[1]) + elif p.slice[1].type == 'BOOLEAN': + p[0] = ASTNode('boolean', value=p[1]) + elif p.slice[1].type == 'IDENTIFIER': + p[0] = ASTNode('identifier', value=p[1]) + else: + p[0] = p[1] + elif len(p) == 3: + # 一元负号 + p[0] = ASTNode('unop', [p[2]], {'op': p[1]}) + elif len(p) == 4: + # 括号表达式 + p[0] = p[2] + else: + # 函数调用 + p[0] = p[1] + + def p_function_call(self, p): + '''function_call : FUNCTION LPAREN args RPAREN''' + p[0] = ASTNode('function', p[3], p[1]) + + def p_args(self, p): + '''args : arg_list + | empty''' + if len(p) == 2 and p[1] is not None: + p[0] = p[1] + else: + p[0] = [] + + def p_arg_list(self, p): + '''arg_list : arg + | arg_list COMMA arg''' + if len(p) == 2: + p[0] = [p[1]] + else: + p[0] = p[1] + [p[3]] + + def p_arg(self, p): + '''arg : expression + | IDENTIFIER ASSIGN expression''' + if len(p) == 2: + p[0] = {'type': 'positional', 'value': p[1]} + else: + p[0] = {'type': 'named', 'name': p[1], 'value': p[3]} + + def p_empty(self, p): + '''empty :''' + p[0] = None + + # 语法错误处理 + def p_error(self, p): + if p: + self.errors.append(f"语法错误在位置 {p.lexpos}: 非法标记 '{p.value}'") + else: + self.errors.append("语法错误: 表达式不完整") + + def _is_valid_field(self, field_name: str) -> bool: + """检查字段名是否符合模式""" + for pattern in field_patterns: + if pattern.match(field_name): + return True + return False + + def validate_function(self, node: ASTNode, is_in_group_arg: bool = False) -> List[str]: + """验证函数调用的参数数量和类型""" + function_name = node.value + args = node.children + function_info = supported_functions.get(function_name) + + if not function_info: + return [f"未知函数: {function_name}"] + + # Custom validation for variadic functions with optional flags + if function_name == 'add': + return self._validate_add(args, is_in_group_arg) + + errors = [] + + # Keyword-only enforcement for optional parameters. + # If enabled, only the required leading arguments can be positional. + keyword_only_from = function_info.get('keyword_only_from') + if keyword_only_from is None and function_info.get('keyword_only'): + keyword_only_from = function_info.get('min_args', 0) + + # 检查参数数量 + if len(args) < function_info['min_args']: + errors.append(f"函数 {function_name} 需要至少 {function_info['min_args']} 个参数，但只提供了 {len(args)}") + elif len(args) > function_info['max_args']: + errors.append(f"函数 {function_name} 最多接受 {function_info['max_args']} 个参数，但提供了 {len(args)}") + + # 处理参数验证 + # 跟踪已使用的位置参数索引 + positional_index = 0 + + # 对于所有函数，支持命名参数 + for arg in args: + if isinstance(arg, dict): + if arg['type'] == 'named': + # 命名参数 + if 'param_names' in function_info and arg['name'] in function_info['param_names']: + # 查找参数在param_names中的索引 + param_index = function_info['param_names'].index(arg['name']) + if param_index < len(function_info['arg_types']): + expected_type = function_info['arg_types'][param_index] + arg_errors = self._validate_arg_type(arg['value'], expected_type, param_index, function_name, is_in_group_arg) + errors.extend(arg_errors) + # 对于winsorize函数，支持std和clip参数 + elif function_name == 'winsorize' and arg['name'] in ['std', 'clip']: + arg_errors = self._validate_arg_type(arg['value'], 'number', 0, function_name, is_in_group_arg) + errors.extend(arg_errors) + # 对于bucket函数，支持'range'和'buckets'参数 + elif function_name == 'bucket' and arg['name'] in ['range', 'buckets']: + # range和buckets参数应该是string类型 + arg_errors = self._validate_arg_type(arg['value'], 'string', 1, function_name, is_in_group_arg) + errors.extend(arg_errors) + else: + errors.append(f"函数 {function_name} 不存在参数 '{arg['name']}'") + elif arg['type'] == 'positional': + # 位置参数（字典形式） + if keyword_only_from is not None and positional_index >= keyword_only_from: + param_name = None + if 'param_names' in function_info and positional_index < len(function_info['param_names']): + param_name = function_info['param_names'][positional_index] + if param_name: + errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数 '{param_name}='") + else: + errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数") + else: + # 验证位置参数的类型 + if positional_index < len(function_info['arg_types']): + expected_type = function_info['arg_types'][positional_index] + arg_errors = self._validate_arg_type(arg['value'], expected_type, positional_index, function_name, is_in_group_arg) + errors.extend(arg_errors) + positional_index += 1 + else: + # 其他字典类型参数 + errors.append(f"参数 {positional_index+1} 格式错误") + positional_index += 1 + else: + # 位置参数（直接ASTNode形式） + if keyword_only_from is not None and positional_index >= keyword_only_from: + param_name = None + if 'param_names' in function_info and positional_index < len(function_info['param_names']): + param_name = function_info['param_names'][positional_index] + if param_name: + errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数 '{param_name}='") + else: + errors.append(f"函数 {function_name} 的第{positional_index+1}个参数必须使用命名参数") + else: + # 验证位置参数的类型 + if positional_index < len(function_info['arg_types']): + expected_type = function_info['arg_types'][positional_index] + arg_errors = self._validate_arg_type(arg, expected_type, positional_index, function_name, is_in_group_arg) + errors.extend(arg_errors) + positional_index += 1 + + return errors + + def _validate_arg_type(self, arg: ASTNode, expected_type: str, arg_index: int, function_name: str, is_in_group_arg: bool = False) -> List[str]: + """验证参数类型是否符合预期""" + errors = [] + + def _is_number_like(node: ASTNode) -> bool: + if node is None: + return False + if node.node_type == 'number': + return True + if node.node_type == 'unop' and isinstance(node.value, dict) and node.value.get('op') in {'-', '+'}: + if node.children and hasattr(node.children[0], 'node_type'): + return _is_number_like(node.children[0]) + return False + + # Unit compatibility check + # bucket()/group_cartesian_product() output a derived category (grouping key). + # It can only be consumed where a category/grouping key is expected. + if self._is_derived_category(arg) and expected_type != 'category': + errors.append( + f"Incompatible unit for input of \"{function_name}\" at index {arg_index}, expected \"Unit[]\", found \"Unit[Group:1]\"" + ) + return errors + + # 首先检查是否是group类型字段，如果是则只能用于Group类型函数 + # 但是如果当前函数是group_xxx或在group函数的参数链中，则允许使用 + if arg.node_type == 'category' and arg.value in group_fields: + if not (function_name.startswith('group_') or is_in_group_arg): + errors.append(f"Group类型字段 '{arg.value}' 只能用于Group类型函数的参数中") + + # 然后验证参数类型是否符合预期 + if expected_type == 'expression': + # 表达式可以是任何有效的AST节点 + pass + elif expected_type == 'number': + # 允许 -1 这类一元负号数字常量（解析为 unop(number)） + if not _is_number_like(arg): + errors.append(f"参数 {arg_index+1} 应该是一个数字，但得到 {arg.node_type}") + elif expected_type == 'boolean': + # 布尔值可以是 true/false 或数字（0/1） + if arg.node_type not in {'boolean', 'number'}: + errors.append(f"参数 {arg_index+1} 应该是一个布尔值（true/false 或 0/1），但得到 {arg.node_type}") + elif expected_type == 'field': + if arg.node_type != 'field' and arg.node_type != 'category': + # 允许field或category作为字段参数 + errors.append(f"参数 {arg_index+1} 应该是一个字段，但得到 {arg.node_type}") + elif arg.node_type == 'field' and not self._is_valid_field(arg.value): + errors.append(f"无效的字段名: {arg.value}") + elif expected_type == 'category': + if not function_name.startswith('group_'): + # 非group函数的category参数必须是category类型且在valid_categories中 + if arg.node_type != 'category': + errors.append(f"参数 {arg_index+1} 应该是一个类别，但得到 {arg.node_type}") + elif arg.value not in valid_categories: + errors.append(f"无效的类别: {arg.value}") + # group函数的category参数可以是任何类型（field、category等），不进行类型校验 + + return errors + + def _infer_unit(self, node: ASTNode) -> str: + """Infer the Unit kind of an AST node. + + Returns: + 'unit' - regular numeric time-series Unit[] + 'scalar' - literals (numbers/booleans/strings) + 'category' - category/grouping keys (industry/sector or derived via bucket/cartesian) + """ + if node is None: + return 'unit' + + cache_key = id(node) + cached = self._unit_cache.get(cache_key) + if cached is not None: + return cached + + unit = 'unit' + + if node.node_type in {'number', 'boolean', 'string'}: + unit = 'scalar' + elif node.node_type in {'field', 'identifier'}: + unit = 'unit' + elif node.node_type == 'category': + unit = 'category' + elif node.node_type in {'unop', 'binop'}: + child_units = [self._infer_unit(child) for child in node.children if hasattr(child, 'node_type')] + unit = 'category' if 'category' in child_units else 'unit' + elif node.node_type == 'function': + fname = node.value + if fname in {'bucket', 'group_cartesian_product'}: + unit = 'category' + else: + first_arg = None + for child in node.children: + if isinstance(child, dict): + if child.get('type') == 'positional': + first_arg = child.get('value') + break + else: + first_arg = child + break + if hasattr(first_arg, 'node_type'): + unit = self._infer_unit(first_arg) + else: + unit = 'unit' + + self._unit_cache[cache_key] = unit + return unit + + def _is_derived_category(self, node: ASTNode) -> bool: + """Return True if node is a derived category/grouping key (e.g., bucket/cartesian output).""" + if node is None: + return False + + cache_key = id(node) + cached = self._derived_category_cache.get(cache_key) + if cached is not None: + return cached + + derived = False + if node.node_type == 'function': + if node.value in {'bucket', 'group_cartesian_product'}: + derived = True + else: + function_info = supported_functions.get(node.value, {}) + arg_types = function_info.get('arg_types', []) + param_names = function_info.get('param_names', []) + + positional_index = 0 + for child in node.children: + if isinstance(child, dict): + if child.get('type') == 'named': + name = child.get('name') + value = child.get('value') + + expected_type = None + if name in param_names: + param_index = param_names.index(name) + if param_index < len(arg_types): + expected_type = arg_types[param_index] + + if expected_type == 'category': + continue + + if self._is_derived_category(value): + derived = True + break + elif child.get('type') == 'positional': + value = child.get('value') + expected_type = arg_types[positional_index] if positional_index < len(arg_types) else None + + if expected_type != 'category' and self._is_derived_category(value): + derived = True + break + positional_index += 1 + else: + expected_type = arg_types[positional_index] if positional_index < len(arg_types) else None + if expected_type != 'category' and self._is_derived_category(child): + derived = True + break + positional_index += 1 + elif node.node_type in {'unop', 'binop'}: + derived = any( + self._is_derived_category(child) + for child in node.children + if hasattr(child, 'node_type') + ) + + self._derived_category_cache[cache_key] = derived + return derived + + def _validate_add(self, args: List[Any], is_in_group_arg: bool = False) -> List[str]: + """Validate add(x, y, ..., filter=false). + + Rules: + - At least 2 positional expression terms. + - Optional filter flag can be provided as: + - named argument: filter= + - last positional argument: or 0/1 + """ + errors: List[str] = [] + + if len(args) < 2: + return [f"函数 add 需要至少 2 个参数，但只提供了 {len(args)}"] + + named_filter_nodes: List[ASTNode] = [] + positional_nodes: List[ASTNode] = [] + + for arg in args: + if isinstance(arg, dict) and arg.get('type') == 'named': + name = arg.get('name') + value = arg.get('value') + if name != 'filter': + errors.append(f"函数 add 不存在参数 '{name}'") + continue + if not hasattr(value, 'node_type'): + errors.append("函数 add 的参数 filter 格式错误") + continue + named_filter_nodes.append(value) + elif isinstance(arg, dict) and arg.get('type') == 'positional': + value = arg.get('value') + if hasattr(value, 'node_type'): + positional_nodes.append(value) + else: + errors.append("函数 add 的位置参数格式错误") + elif hasattr(arg, 'node_type'): + positional_nodes.append(arg) + else: + errors.append("函数 add 的参数格式错误") + + if len(named_filter_nodes) > 1: + errors.append("函数 add 的参数 'filter' 只能出现一次") + + positional_filter_node: Optional[ASTNode] = None + # Only infer a positional filter flag when: + # - no named filter is provided + # - there are at least 3 positional args (x, y, filter) + # - the last arg is boolean or numeric 0/1 + if not named_filter_nodes and len(positional_nodes) >= 3: + last = positional_nodes[-1] + if last.node_type == 'boolean' or (last.node_type == 'number' and last.value in {0, 1}): + positional_filter_node = positional_nodes.pop() + + if len(positional_nodes) < 2: + errors.append(f"函数 add 需要至少 2 个输入项（不含filter），但只提供了 {len(positional_nodes)}") + + for idx, node in enumerate(positional_nodes): + errors.extend(self._validate_arg_type(node, 'expression', idx, 'add', is_in_group_arg)) + + if positional_filter_node is not None and named_filter_nodes: + errors.append("函数 add 的 filter 不能同时用位置参数和命名参数传递") + if positional_filter_node is not None: + errors.extend(self._validate_arg_type(positional_filter_node, 'boolean', len(positional_nodes), 'add', is_in_group_arg)) + if named_filter_nodes: + errors.extend(self._validate_arg_type(named_filter_nodes[0], 'boolean', len(positional_nodes), 'add', is_in_group_arg)) + + return errors + + def validate_ast(self, ast: Optional[ASTNode], is_in_group_arg: bool = False) -> List[str]: + """递归验证抽象语法树""" + if not ast: + return ["无法解析表达式"] + + errors = [] + + # 根据节点类型进行验证 + if ast.node_type == 'function': + # 检查当前函数是否是group函数 + is_group_function = ast.value.startswith('group_') + # 确定当前是否在group函数的参数链中 + current_in_group_arg = is_in_group_arg or is_group_function + # 验证函数 + function_errors = self.validate_function(ast, current_in_group_arg) + errors.extend(function_errors) + + # 递归验证子节点时使用current_in_group_arg + for child in ast.children: + if isinstance(child, dict): + # 命名参数，验证其值 + if 'value' in child and hasattr(child['value'], 'node_type'): + child_errors = self.validate_ast(child['value'], current_in_group_arg) + errors.extend(child_errors) + elif hasattr(child, 'node_type'): + child_errors = self.validate_ast(child, current_in_group_arg) + errors.extend(child_errors) + elif ast.node_type in ['unop', 'binop']: + # 对操作符的子节点进行验证 + for child in ast.children: + if hasattr(child, 'node_type'): + child_errors = self.validate_ast(child, is_in_group_arg) + errors.extend(child_errors) + elif ast.node_type == 'field': + # 验证字段名 + if not self._is_valid_field(ast.value): + errors.append(f"无效的字段名: {ast.value}") + else: + # 递归验证子节点 + for child in ast.children: + if isinstance(child, dict): + # 命名参数，验证其值 + if 'value' in child and hasattr(child['value'], 'node_type'): + child_errors = self.validate_ast(child['value'], is_in_group_arg) + errors.extend(child_errors) + elif hasattr(child, 'node_type'): + child_errors = self.validate_ast(child, is_in_group_arg) + errors.extend(child_errors) + + return errors + + def _process_semicolon_expression(self, expression: str) -> Tuple[bool, str]: + """处理带有分号的表达式，将其转换为不带分号的简化形式 + + Args: + expression: 要处理的表达式字符串 + + Returns: + Tuple[bool, str]: (是否成功, 转换后的表达式或错误信息) + """ + def _top_level_equals_positions(stmt: str) -> List[int]: + """返回所有“顶层赋值”等号位置。 + + 仅统计括号外（()[]{})、引号外、且不属于比较操作符(==,!=,<=,>=)的 '='。 + 这样可以避免把关键字参数（如 rettype=0）误判为赋值语句。 + """ + positions: List[int] = [] + paren_depth = 0 + bracket_depth = 0 + brace_depth = 0 + in_single_quote = False + in_double_quote = False + escape = False + + for i, ch in enumerate(stmt): + if escape: + escape = False + continue + if ch == '\\': + escape = True + continue + + if in_single_quote: + if ch == "'": + in_single_quote = False + continue + if in_double_quote: + if ch == '"': + in_double_quote = False + continue + + if ch == "'": + in_single_quote = True + continue + if ch == '"': + in_double_quote = True + continue + + if ch == '(': + paren_depth += 1 + continue + if ch == ')': + paren_depth = max(0, paren_depth - 1) + continue + if ch == '[': + bracket_depth += 1 + continue + if ch == ']': + bracket_depth = max(0, bracket_depth - 1) + continue + if ch == '{': + brace_depth += 1 + continue + if ch == '}': + brace_depth = max(0, brace_depth - 1) + continue + + if paren_depth or bracket_depth or brace_depth: + continue + + if ch != '=': + continue + + prev_ch = stmt[i - 1] if i > 0 else '' + next_ch = stmt[i + 1] if i + 1 < len(stmt) else '' + if prev_ch in ['=', '!', '<', '>'] or next_ch == '=': + continue + + positions.append(i) + + return positions + + def _keyword_arg_names(stmt: str): + """提取函数调用中的命名参数名（如 rettype=0 中的 rettype）。 + + 只收集括号/中括号/大括号内部出现的 name= 形式，避免把脚本级赋值误当作命名参数。 + """ + names = set() + paren_depth = 0 + bracket_depth = 0 + brace_depth = 0 + in_single_quote = False + in_double_quote = False + escape = False + + i = 0 + while i < len(stmt): + ch = stmt[i] + + if escape: + escape = False + i += 1 + continue + if ch == '\\': + escape = True + i += 1 + continue + + if in_single_quote: + if ch == "'": + in_single_quote = False + i += 1 + continue + if in_double_quote: + if ch == '"': + in_double_quote = False + i += 1 + continue + + if ch == "'": + in_single_quote = True + i += 1 + continue + if ch == '"': + in_double_quote = True + i += 1 + continue + + if ch == '(': + paren_depth += 1 + i += 1 + continue + if ch == ')': + paren_depth = max(0, paren_depth - 1) + i += 1 + continue + if ch == '[': + bracket_depth += 1 + i += 1 + continue + if ch == ']': + bracket_depth = max(0, bracket_depth - 1) + i += 1 + continue + if ch == '{': + brace_depth += 1 + i += 1 + continue + if ch == '}': + brace_depth = max(0, brace_depth - 1) + i += 1 + continue + + inside_container = bool(paren_depth or bracket_depth or brace_depth) + + if inside_container and (ch.isalpha() or ch == '_'): + start = i + i += 1 + while i < len(stmt) and (stmt[i].isalnum() or stmt[i] == '_'): + i += 1 + name = stmt[start:i] + + j = i + while j < len(stmt) and stmt[j].isspace(): + j += 1 + + if j < len(stmt) and stmt[j] == '=': + next_ch = stmt[j + 1] if j + 1 < len(stmt) else '' + if next_ch != '=': + names.add(name.lower()) + continue + + i += 1 + + return names + + # 检查表达式是否以分号结尾 + if expression.strip().endswith(';'): + return False, "表达式不能以分号结尾" + + # 分割表达式为语句列表 + statements = [stmt.strip() for stmt in expression.split(';') if stmt.strip()] + if not statements: + return False, "表达式不能为空" + + # 存储变量赋值 + variables = {} + + # 处理每个赋值语句（除了最后一个） + for i, stmt in enumerate(statements[:-1]): + eq_positions = _top_level_equals_positions(stmt) + if not eq_positions: + return False, f"第{i+1}个语句必须是赋值语句（使用=符号）" + if len(eq_positions) > 1: + return False, f"第{i+1}个语句只能包含一个赋值符号（=）" + + real_equals_pos = eq_positions[0] + + # 分割变量名和值 + var_name = stmt[:real_equals_pos].strip() + var_value = stmt[real_equals_pos + 1:].strip() + + # 检查变量名是否有效 + if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', var_name): + return False, f"第{i+1}个语句的变量名'{var_name}'无效，只能包含字母、数字和下划线，且不能以数字开头" + + var_name_lower = var_name.lower() # 变量名不区分大小写 + + # 检查变量名是否在后续表达式中使用 + # 这里不需要，因为后面的表达式会检查 + + # 检查变量值中使用的变量是否已经定义 + # 简单检查：提取所有可能的变量名 + kw_names = _keyword_arg_names(var_value) + used_vars = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', var_value) + for used_var in used_vars: + used_var_lower = used_var.lower() + if used_var_lower in kw_names: + continue + if used_var_lower not in variables: + # 检查是否是函数名 + if used_var not in supported_functions: + # 对于单个字母或简单单词，不自动视为字段名，要求先定义 + if len(used_var) <= 2: + return False, f"第{i+1}个语句中使用的变量'{used_var}'未在之前定义" + # 对于较长的字段名，仍然允许作为字段名 + elif not self._is_valid_field(used_var): + return False, f"第{i+1}个语句中使用的变量'{used_var}'未在之前定义" + + # 将之前定义的变量替换到当前值中 + for existing_var, existing_val in variables.items(): + # 使用单词边界匹配，避免替换到其他单词的一部分 + var_value = re.sub(rf'\b{existing_var}\b', existing_val, var_value) + + # 存储变量 + variables[var_name_lower] = var_value + + # 处理最后一个语句（实际的表达式） + final_stmt = statements[-1] + + # 检查最后一个语句是否是赋值语句 + if _top_level_equals_positions(final_stmt): + return False, "最后一个语句不能是赋值语句" + + # 检查最后一个语句中使用的变量是否已经定义 + kw_names = _keyword_arg_names(final_stmt) + used_vars = re.findall(r'\b[a-zA-Z_][a-zA-Z0-9_]*\b', final_stmt) + for used_var in used_vars: + used_var_lower = used_var.lower() + if used_var_lower in kw_names: + continue + if used_var_lower not in variables: + # 检查是否是函数名 + if used_var not in supported_functions: + # 对于单个字母或简单单词，不自动视为字段名，要求先定义 + if len(used_var) <= 2: + return False, f"最后一个语句中使用的变量'{used_var}'未在之前定义" + + # 允许直接使用字段名/类别名（如 close/industry） + if self._is_valid_field(used_var) or used_var_lower in valid_categories or used_var_lower in group_fields: + continue + + return False, f"最后一个语句中使用的变量'{used_var}'未在之前定义" + + # 将变量替换到最后一个表达式中 + final_expr = final_stmt + for var_name, var_value in variables.items(): + final_expr = re.sub(rf'\b{var_name}\b', var_value, final_expr) + + return True, final_expr + + def check_expression(self, expression: str) -> Dict[str, Any]: + """ + 检查表达式格式是否正确 + + Args: + expression: 要验证的表达式字符串 + + Returns: + 包含验证结果的字典 + """ + # 重置错误列表 + self.errors = [] + # Reset unit inference cache for this expression + self._unit_cache = {} + self._derived_category_cache = {} + + try: + expression = expression.strip() + if not expression: + return { + 'valid': False, + 'errors': ['表达式不能为空'], + 'tokens': [], + 'ast': None + } + + # 处理带有分号的表达式 + if ';' in expression: + success, result = self._process_semicolon_expression(expression) + if not success: + return { + 'valid': False, + 'errors': [result], + 'tokens': [], + 'ast': None + } + expression = result + + # 重置词法分析器的行号 + self.lexer.lineno = 1 + + # 词法分析（用于调试） + self.lexer.input(expression) + tokens = [] + # 调试：打印识别的标记 + # logger.info(f"\n调试 - 表达式: {expression}") + # print("识别的标记:") + for token in self.lexer: + # print(f" - 类型: {token.type}, 值: '{token.value}', 位置: {token.lexpos}") + tokens.append(token) + + # 重新设置词法分析器的输入，以便语法分析器使用 + self.lexer.input(expression) + self.lexer.lineno = 1 + + # 语法分析 + ast = self.parser.parse(expression, lexer=self.lexer) + + # 验证AST + validation_errors = self.validate_ast(ast) + + # 合并所有错误 + all_errors = self.errors + validation_errors + + # 检查括号是否匹配 + bracket_count = 0 + for char in expression: + if char == '(': + bracket_count += 1 + elif char == ')': + bracket_count -= 1 + if bracket_count < 0: + all_errors.append("括号不匹配: 右括号过多") + break + if bracket_count > 0: + all_errors.append("括号不匹配: 左括号过多") + + return { + 'valid': len(all_errors) == 0, + 'errors': all_errors, + 'tokens': tokens, + 'ast': ast + } + except Exception as e: + return { + 'valid': False, + 'errors': [f"解析错误: {str(e)}"], + 'tokens': [], + 'ast': None + } + + + diff --git a/simple72/config.json b/simple72/config.json index 5cbfe58..553c16a 100644 --- a/simple72/config.json +++ b/simple72/config.json @@ -10,6 +10,7 @@ }, "transformer": { "top_n_datafield": 30, - "data_type": "MATRIX" + "data_type": "MATRIX", + "alpha_id": "pw13nJlo" } } diff --git a/simple72/config_kimi.json b/simple72/config_kimi.json new file mode 100644 index 0000000..1ac968e --- /dev/null +++ b/simple72/config_kimi.json @@ -0,0 +1,16 @@ +{ + "brain": { + "username": "jack0210_@hotmail.com", + "password": "!QAZ2wsx+0913" + }, + "llm": { + "api_key": "sk-wwgvbK1sq0vaoCCZ3TahxtJKrs11MKKBMPZdgTQRqcwh55Kq", + "base_url": "https://api.moonshot.cn/v1", + "model": "kimi-k2.5" + }, + "transformer": { + "top_n_datafield": 30, + "data_type": "MATRIX", + "alpha_id": "QPJ1q0mQ" + } +} diff --git a/simple72/config_minimax.json b/simple72/config_minimax.json new file mode 100644 index 0000000..d82e37c --- /dev/null +++ b/simple72/config_minimax.json @@ -0,0 +1,16 @@ +{ + "brain": { + "username": "jack0210_@hotmail.com", + "password": "!QAZ2wsx+0913" + }, + "llm": { + "api_key": "sk-cp-l_as8mjqPhsOIny9IFKZ8jzA92z1c0eRwchldhEf4KzQjs9cjVknV2o7VNCcvYUXsXFq7uF4aSgp2RxxmUHLXwPGKgIvzedM70_XUIXiBB3gu_UmLDQLfh4", + "base_url": "https://api.minimaxi.com/v1", + "model": "MiniMax-M2.7" + }, + "transformer": { + "top_n_datafield": 30, + "data_type": "MATRIX", + "alpha_id": "QPJ1q0mQ" + } +} diff --git a/simple72/config_ollama.json b/simple72/config_ollama.json new file mode 100644 index 0000000..7a63343 --- /dev/null +++ b/simple72/config_ollama.json @@ -0,0 +1,16 @@ +{ + "brain": { + "username": "jack0210_@hotmail.com", + "password": "!QAZ2wsx+0913" + }, + "llm": { + "api_key": "1234", + "base_url": "https://127.0.0.1:11434/v1", + "model": "kimi-k2.5:cloud" + }, + "transformer": { + "top_n_datafield": 30, + "data_type": "MATRIX", + "alpha_id": "QPJ1q0mQ" + } +} diff --git a/simple72/fix_log_levels.py b/simple72/fix_log_levels.py new file mode 100644 index 0000000..4587c30 --- /dev/null +++ b/simple72/fix_log_levels.py @@ -0,0 +1,46 @@ +#!/usr/bin/env python3 +""" +批量修复日志级别脚本 +将 logger.info 中应该为 warning 或 error 的改为正确的级别 +""" + +import re + +def fix_log_levels(file_path): + with open(file_path, 'r', encoding='utf-8') as f: + content = f.read() + + # 定义需要改为 warning 的模式 + warning_patterns = [ + (r'logger\.info\("(⚠[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.warning("\1")'), + (r'logger\.info\(f"(⚠[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.warning(f"\1")'), + (r'logger\.info\("([^"]*Warning[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.warning("\1")'), + (r'logger\.info\(f"([^"]*Warning[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.warning(f"\1")'), + ] + + # 定义需要改为 error 的模式 + error_patterns = [ + (r'logger\.info\("(❌[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error("\1")'), + (r'logger\.info\(f"(❌[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error(f"\1")'), + (r'logger\.info\("([^"]*Failed[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error("\1")'), + (r'logger\.info\(f"([^"]*Failed[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error(f"\1")'), + (r'logger\.info\("(错误:[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error("\1")'), + (r'logger\.info\(f"(错误:[^"]*)"\s*,\s*flush\s*=\s*True\)', r'logger.error(f"\1")'), + ] + + # 应用 warning 模式 + for pattern, replacement in warning_patterns: + content = re.sub(pattern, replacement, content, flags=re.IGNORECASE) + + # 应用 error 模式 + for pattern, replacement in error_patterns: + content = re.sub(pattern, replacement, content, flags=re.IGNORECASE) + + with open(file_path, 'w', encoding='utf-8') as f: + f.write(content) + + print(f"✓ 已修复: {file_path}") + +if __name__ == "__main__": + fix_log_levels("/Users/jack/source/mySpace/mycode/my_project/py/alpha_odoo/alpha_tools/simple72/Tranformer/Transformer.py") + print("✓ 所有日志级别修复完成！") diff --git a/simple72/main.py b/simple72/main.py index 2b46836..10e4e9e 100644 --- a/simple72/main.py +++ b/simple72/main.py @@ -1,16 +1,27 @@ -# FastAPI 应用主入口 +import shutil from fastapi import FastAPI, Request -from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse, FileResponse -from fastapi.staticfiles import StaticFiles +from fastapi.responses import HTMLResponse, JSONResponse, FileResponse from fastapi.templating import Jinja2Templates import sys import os import json -import subprocess import uuid -import asyncio import zipfile from datetime import datetime +import logging +import openai +from Tranformer.ace_lib import SingleSession, get_instrument_type_region_delay, get_datafields, get_datasets +from Tranformer.Transformer import run_transformer +import asyncio +import pandas as pd + +# 配置日志 +logger = logging.getLogger("alpha_transformer") +if not logger.handlers: + handler = logging.StreamHandler(sys.stdout) + handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')) + logger.addHandler(handler) + logger.setLevel(logging.INFO) # 创建 FastAPI 应用实例 app = FastAPI(title="Alpha Transformer", version="1.0.0") @@ -37,6 +48,20 @@ async def get_app_js(): # 配置 Jinja2 模板引擎，用于渲染 HTML 页面 templates = Jinja2Templates(directory="templates") +# 程序启动前清空 Tranformer/output 目录 +if os.path.exists("Tranformer/output"): + for filename in os.listdir("Tranformer/output"): + file_path = os.path.join("Tranformer/output", filename) + try: + if os.path.isfile(file_path) or os.path.islink(file_path): + os.unlink(file_path) + elif os.path.isdir(file_path): + shutil.rmtree(file_path) + except Exception as e: + logger.error(f"⚠ 清空目录 {file_path} 时出错: {e}") + + + # 存储正在运行的任务信息（task_id -> 任务状态） transformer_tasks = {} @@ -53,13 +78,13 @@ def load_config(): try: with open(config_path, 'r', encoding='utf-8') as f: config = json.load(f) - print(f"✓ 已加载配置文件: {config_path}") + logger.info(f"✓ 已加载配置文件: {config_path}") return config except Exception as e: - print(f"⚠ 加载配置文件失败: {e}") + logger.warning(f"⚠ 加载配置文件失败: {e}") return {} else: - print(f"⚠ 配置文件不存在: {config_path}") + logger.warning(f"⚠ 配置文件不存在: {config_path}") return {} # 启动时加载配置 @@ -120,17 +145,17 @@ async def generate_alpha(request: Request): 生成 Alpha 变种的 API 端点 接收前端表单数据，启动 Transformer 脚本执行 Alpha 生成任务 """ - print("=" * 50) - print("收到生成变种请求") + logger.info("=" * 50) + logger.info("收到生成变种请求") try: # 解析请求数据 data = await request.json() - print(f"请求数据: alpha_id={data.get('alpha_id')}, llm_model={data.get('llm_model')}") + logger.info(f"请求数据: alpha_id={data.get('alpha_id')}, llm_model={data.get('llm_model')}") # 生成唯一任务 ID task_id = str(uuid.uuid4()) - print(f"生成任务 ID: {task_id}") + logger.info(f"生成任务 ID: {task_id}") # 定义必须提交的字段 required_fields = [ @@ -145,7 +170,7 @@ async def generate_alpha(request: Request): # 检查必填字段是否完整 for field in required_fields: if not data.get(field): - print(f"缺少必填字段: {field}") + logger.warning(f"缺少必填字段: {field}") return JSONResponse( status_code=400, content={"success": False, "error": f"Missing required field: {field}"} @@ -154,7 +179,7 @@ async def generate_alpha(request: Request): # 获取脚本所在目录和 Transformer 子目录 script_dir = os.path.dirname(os.path.abspath(__file__)) transformer_dir = os.path.join(script_dir, 'Tranformer') - print(f"Transformer 目录: {transformer_dir}") + logger.info(f"Transformer 目录: {transformer_dir}") # 构建传递给 Transformer 脚本的配置 config = { @@ -172,38 +197,147 @@ async def generate_alpha(request: Request): "user_data_type": data.get('user_data_type', 'MATRIX'), "max_retries": int(data.get('max_retries', 20)) } - print(f"配置已构建: LLM_model={config['LLM_model_name']}, alpha_id={config['alpha_id']}, max_retries={config['max_retries']}") + logger.info(f"配置已构建: LLM_model={config['LLM_model_name']}, alpha_id={config['alpha_id']}, max_retries={config['max_retries']}") - # 将配置写入临时 JSON 文件，供 Transformer 脚本读取 - config_path = os.path.join(transformer_dir, f'config_{task_id}.json') - with open(config_path, 'w', encoding='utf-8') as f: - json.dump(config, f, indent=4) - print(f"配置文件已写入: {config_path}") + # 检查数据字段缓存是否存在，不存在则自动下载 + dataset_dir = os.path.join(script_dir, 'dataset') + os.makedirs(dataset_dir, exist_ok=True) - try: - # 启动 Transformer.py 子进程执行 Alpha 生成 - print(f"启动 Transformer 脚本...") - process = subprocess.run( - [sys.executable, '-u', os.path.join(transformer_dir, 'Transformer.py'), config_path], - cwd=transformer_dir, - capture_output=True, - text=True, - timeout=600, - env={**os.environ, "PYTHONIOENCODING": "utf-8"} + user_category = data.get('user_category', []) + categories_to_check = user_category if isinstance(user_category, list) else [user_category] if user_category else [] + + # 创建会话来获取数据集信息 + session = SingleSession() + session.auth = (data.get('brain_username'), data.get('brain_password')) + brain_api_url = "https://api.worldquantbrain.com" + auth_response = session.post(brain_api_url + "/authentication") + + if auth_response.status_code != 201: + return JSONResponse( + status_code=401, + content={"success": False, "error": "Authentication failed"} ) - print(f"Transformer 脚本执行完成，返回码: {process.returncode}") - - # 打印 Transformer 的输出到终端 - if process.stdout: - print(f"\n{'='*50}") - print("Transformer 输出:") - print(f"{'='*50}") - print(process.stdout) - if process.stderr: - print(f"\n{'='*50}") - print("Transformer 错误:") - print(f"{'='*50}") - print(process.stderr) + + # 获取所有可用的数据集 + datasets_df = get_datasets( + s=session, + instrument_type="EQUITY", + region=data.get('user_region', 'USA'), + delay=int(data.get('user_delay', 1)), + universe=data.get('user_universe', 'TOP3000'), + theme="ALL" + ) + + if datasets_df.empty: + return JSONResponse( + status_code=404, + content={"success": False, "error": "No datasets found"} + ) + + # 使用 category_name 列（category 列是字典类型） + category_column = 'category_name' if 'category_name' in datasets_df.columns else 'category' + + # 如果没有指定类别，使用所有类别 + if not categories_to_check: + if category_column == 'category': + # category 是字典，提取 name + categories_to_check = datasets_df[category_column].apply(lambda x: x['name'] if isinstance(x, dict) and 'name' in x else x).unique().tolist() + else: + categories_to_check = datasets_df[category_column].unique().tolist() + logger.info(f"未指定类别，将使用所有 {len(categories_to_check)} 个类别: {categories_to_check}") + + # 构建类别到 dataset_id 的映射 + category_to_datasets = {} + for cat in categories_to_check: + if category_column == 'category': + # category 是字典，需要匹配 name 字段 + cat_datasets = datasets_df[datasets_df[category_column].apply(lambda x: x.get('name') if isinstance(x, dict) else x) == cat] + else: + cat_datasets = datasets_df[datasets_df[category_column] == cat] + if not cat_datasets.empty: + category_to_datasets[cat] = cat_datasets['id'].tolist() + + # 检查每个类别的缓存文件 + missing_categories = [] + for cat in categories_to_check: + cache_filename = f"datafields_cache_{data.get('user_region', 'USA')}_{data.get('user_universe', 'TOP3000')}_D{data.get('user_delay', 1)}_{cat}.csv" + cache_path = os.path.join(dataset_dir, cache_filename) + if not os.path.exists(cache_path): + missing_categories.append(cat) + + # 如果有缺失的缓存，检查是否都能对应到数据集 + if missing_categories: + # 检查是否有类别在数据集中不存在 + invalid_categories = [cat for cat in missing_categories if cat not in category_to_datasets] + if invalid_categories: + error_msg = f"以下类别在数据集中不存在，无法下载: {', '.join(invalid_categories)}。请检查类别名称或先点击'下载数据字段缓存'按钮。" + logger.error(error_msg) + return JSONResponse( + status_code=400, + content={"success": False, "error": error_msg} + ) + + logger.info(f"检测到缺失的数据字段缓存: {missing_categories}，自动下载中...") + + downloaded_count = 0 + failed_categories = [] + + for cat in missing_categories: + dataset_ids = category_to_datasets[cat] + all_datafields = [] + + try: + # 下载该类别下的所有数据集 + for dataset_id in dataset_ids: + logger.info(f"下载数据集 {dataset_id} (类别: {cat})...") + datafields_df = get_datafields( + s=session, + instrument_type="EQUITY", + region=data.get('user_region', 'USA'), + delay=int(data.get('user_delay', 1)), + universe=data.get('user_universe', 'TOP3000'), + dataset_id=dataset_id, + data_type=data.get('user_data_type', 'MATRIX') + ) + + if not datafields_df.empty: + all_datafields.append(datafields_df) + logger.info(f"数据集 {dataset_id} 包含 {len(datafields_df)} 个字段") + + if all_datafields: + # 合并该类别下的所有数据字段 + combined_df = pd.concat(all_datafields, ignore_index=True) + combined_df.drop_duplicates(subset=['id'], inplace=True) + + cache_filename = f"datafields_cache_{data.get('user_region', 'USA')}_{data.get('user_universe', 'TOP3000')}_D{data.get('user_delay', 1)}_{cat}.csv" + cache_path = os.path.join(dataset_dir, cache_filename) + combined_df.to_csv(cache_path, index=False) + downloaded_count += 1 + logger.info(f"自动下载完成: {cache_filename} ({len(combined_df)} 个字段)") + else: + failed_categories.append(cat) + logger.error(f"类别 '{cat}' 没有获取到任何数据字段") + + except Exception as e: + failed_categories.append(cat) + logger.error(f"下载类别 '{cat}' 失败: {e}") + + # 如果有下载失败的类别，返回错误 + if failed_categories: + error_msg = f"以下类别下载失败: {', '.join(failed_categories)}。请检查网络连接或手动下载。" + logger.error(error_msg) + return JSONResponse( + status_code=500, + content={"success": False, "error": error_msg} + ) + + logger.info(f"自动下载完成，共下载 {downloaded_count}/{len(missing_categories)} 个类别") + + # 直接调用 Transformer 模块（替代 subprocess） + logger.info(f"启动 Transformer 执行...") + try: + transformer_result = await run_transformer(config) + logger.info(f"Transformer 执行完成: {transformer_result}") # 定义输出文件路径 output_file = os.path.join(transformer_dir, 'output', 'Alpha_generated_expressions_success.json') @@ -214,56 +348,47 @@ async def generate_alpha(request: Request): result = { "success": True, "alpha_id": data.get('alpha_id'), - "stdout": process.stdout, - "stderr": process.stderr, - "return_code": process.returncode + "stdout": "", + "stderr": "", + "return_code": 0 } # 读取成功生成的表达式 if os.path.exists(output_file): - print(f"读取成功表达式文件: {output_file}") + logger.info(f"读取成功表达式文件: {output_file}") with open(output_file, 'r', encoding='utf-8') as f: result['expressions_success'] = json.load(f) else: - print(f"成功表达式文件不存在: {output_file}") + logger.info(f"成功表达式文件不存在: {output_file}") result['expressions_success'] = [] # 读取候选表达式 if os.path.exists(candidates_file): - print(f"读取候选表达式文件: {candidates_file}") + logger.info(f"读取候选表达式文件: {candidates_file}") with open(candidates_file, 'r', encoding='utf-8') as f: result['candidates'] = json.load(f) else: - print(f"候选表达式文件不存在: {candidates_file}") + logger.info(f"候选表达式文件不存在: {candidates_file}") result['candidates'] = [] # 读取生成失败的表达式 if os.path.exists(error_file): - print(f"读取错误表达式文件: {error_file}") + logger.info(f"读取错误表达式文件: {error_file}") with open(error_file, 'r', encoding='utf-8') as f: result['expressions_error'] = json.load(f) else: - print(f"错误表达式文件不存在: {error_file}") + logger.info(f"错误表达式文件不存在: {error_file}") result['expressions_error'] = [] - print(f"成功: {len(result['expressions_success'])} 个, 候选: {len(result['candidates'])} 个, 错误: {len(result['expressions_error'])} 个") - print("=" * 50) + logger.info(f"成功: {len(result['expressions_success'])} 个, 候选: {len(result['candidates'])} 个, 错误: {len(result['expressions_error'])} 个") + logger.info("=" * 50) return JSONResponse(content=result) - finally: - # 清理临时配置文件 - if os.path.exists(config_path): - os.remove(config_path) - print(f"已清理临时配置文件: {config_path}") - - except subprocess.TimeoutExpired: - print("任务执行超时 (600秒)") - return JSONResponse( - status_code=408, - content={"success": False, "error": "Task timeout (600s)"} - ) + except Exception as transformer_error: + logger.error(f"Transformer 执行失败: {transformer_error}") + raise except Exception as e: - print(f"执行异常: {str(e)}") + logger.warning(f"执行异常: {str(e)}") return JSONResponse( status_code=500, content={"success": False, "error": str(e)} @@ -293,10 +418,6 @@ async def login_and_fetch_options(request: Request): if transformer_dir not in sys.path: sys.path.append(transformer_dir) - # 导入必要的模块 - from Tranformer.ace_lib import SingleSession, get_instrument_type_region_delay - import pandas as pd - # 创建新的会话实例 session = SingleSession() session.auth = (username, password) @@ -356,7 +477,7 @@ async def login_and_fetch_options(request: Request): }) except Exception as e: - print(f"登录获取选项失败: {str(e)}") + logger.warning(f"登录获取选项失败: {str(e)}") return JSONResponse( status_code=500, content={'success': False, 'error': str(e)} @@ -374,6 +495,7 @@ async def test_llm_connection(request: Request): """ 测试 LLM 连接 接收 LLM 配置，尝试连接并返回测试结果 + 对 529 状态码会自动重试，默认重试 3 次 """ try: data = await request.json() @@ -381,6 +503,7 @@ async def test_llm_connection(request: Request): api_key = data.get('llm_api_key') base_url = data.get('llm_base_url') model = data.get('llm_model') + max_retries = data.get('max_retries', 3) if not api_key or not base_url or not model: return JSONResponse( @@ -388,51 +511,215 @@ async def test_llm_connection(request: Request): content={"success": False, "error": "Missing required LLM configuration"} ) - print(f"测试 LLM 连接: {base_url}, 模型: {model}") - - # 导入 openai - import openai + logger.info(f"测试 LLM 连接: {base_url}, 模型: {model}, 最大重试次数: {max_retries}") # 创建客户端 client = openai.AsyncOpenAI( api_key=api_key, base_url=base_url ) - - # 尝试发送一个简单的请求 - response = await client.chat.completions.create( - model=model, - messages=[ - {"role": "user", "content": "Hello, this is a connection test. Reply with 'OK' only."} - ], - max_tokens=10, - timeout=30 - ) - - # 检查响应 - if response and response.choices and len(response.choices) > 0: - content = response.choices[0].message.content - print(f"✓ LLM 连接测试成功: {content[:50]}...") - return JSONResponse(content={ - "success": True, - "message": "LLM 连接成功", - "response": content[:100] - }) - else: - return JSONResponse( - status_code=500, - content={"success": False, "error": "Empty response from LLM"} - ) - + + # 重试机制 + last_error = None + for attempt in range(1, max_retries + 1): + try: + response = await client.chat.completions.create( + model=model, + messages=[{"role": "user", "content": "Hello, this is a connection test. Reply with 'OK' only."}], + max_tokens=10, + timeout=30 + ) + + # 检查响应 + if response and response.choices and len(response.choices) > 0: + content = response.choices[0].message.content + logger.info(f"✓ LLM 连接测试成功: {content[:50]}...") + return JSONResponse(content={ + "success": True, + "message": "LLM 连接成功", + "response": content[:100] + }) + else: + return JSONResponse( + status_code=500, + content={"success": False, "error": "Empty response from LLM"} + ) + + except Exception as e: + error_msg = str(e) + last_error = e + + # 检查是否是 529 错误（MiniMax 过载错误） + if "529" in error_msg or "overloaded" in error_msg.lower(): + if attempt < max_retries: + logger.warning(f"⚠ LLM 返回 529 错误，第 {attempt}/{max_retries} 次尝试失败，2 秒后重试...") + await asyncio.sleep(2) + continue + else: + logger.error(f"✗ LLM 连接测试失败，已重试 {max_retries} 次，仍然返回 529 错误") + return JSONResponse( + status_code=529, + content={"success": False, "error": f"LLM service overloaded (529) after {max_retries} retries"} + ) + else: + # 其他错误直接抛出 + raise + except Exception as e: error_msg = str(e) - print(f"✗ LLM 连接测试失败: {error_msg}") + logger.warning(f"✗ LLM 连接测试失败: {error_msg}") return JSONResponse( status_code=500, content={"success": False, "error": error_msg} ) +@app.post("/api/download-datafields") +async def download_datafields(request: Request): + """ + 下载数据字段缓存 + 全量获取指定 region/delay/universe/category 的数据字段并保存到 CSV + """ + try: + data = await request.json() + username = data.get('username') + password = data.get('password') + region = data.get('region') + delay = data.get('delay') + universe = data.get('universe') + data_type = data.get('data_type', 'MATRIX') + category = data.get('category', []) # 类别列表 + + if not all([username, password, region, delay, universe]): + return JSONResponse( + status_code=400, + content={"success": False, "error": "Missing required parameters"} + ) + + # 创建会话 + session = SingleSession() + session.auth = (username, password) + + brain_api_url = "https://api.worldquantbrain.com" + response = session.post(brain_api_url + "/authentication") + + if response.status_code != 201: + return JSONResponse( + status_code=401, + content={"success": False, "error": "Authentication failed"} + ) + + # 获取所有数据集信息 + datasets_df = get_datasets( + s=session, + instrument_type="EQUITY", + region=region, + delay=delay, + universe=universe, + theme="ALL" + ) + + if datasets_df.empty: + return JSONResponse( + status_code=404, + content={"success": False, "error": "No datasets found"} + ) + + # 处理类别参数（前端总是传递类别列表） + categories_to_download = category if isinstance(category, list) else [category] if category else [] + + # 如果没有指定类别，使用所有类别 + if not categories_to_download: + if 'category_id' in datasets_df.columns: + categories_to_download = datasets_df['category_id'].unique().tolist() + elif 'category_name' in datasets_df.columns: + categories_to_download = datasets_df['category_name'].unique().tolist() + else: + # category 是字典，提取 id + categories_to_download = datasets_df['category'].apply(lambda x: x['id'] if isinstance(x, dict) and 'id' in x else x).unique().tolist() + logger.info(f"未指定类别，将使用所有 {len(categories_to_download)} 个类别") + + # 构建类别到 dataset_id 的映射（使用 category_id 匹配） + category_to_datasets = {} + for cat in categories_to_download: + if 'category_id' in datasets_df.columns: + cat_datasets = datasets_df[datasets_df['category_id'] == cat] + elif 'category_name' in datasets_df.columns: + cat_datasets = datasets_df[datasets_df['category_name'] == cat] + else: + cat_datasets = datasets_df[datasets_df['category'].apply(lambda x: x.get('id') if isinstance(x, dict) else x) == cat] + if not cat_datasets.empty: + category_to_datasets[cat] = cat_datasets['id'].tolist() + + # 保存到 CSV (dataset 文件夹) + dataset_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'dataset') + os.makedirs(dataset_dir, exist_ok=True) + + logger.info(f"开始下载数据字段: region={region}, delay={delay}, universe={universe}, type={data_type}, categories={categories_to_download}") + + total_count = 0 + saved_files = [] + + # 每个类别单独下载并保存 + for cat in categories_to_download: + if cat not in category_to_datasets: + logger.warning(f"类别 '{cat}' 在数据集中不存在，跳过") + continue + + dataset_ids = category_to_datasets[cat] + all_datafields = [] + + # 下载该类别下的所有数据集 + for dataset_id in dataset_ids: + logger.info(f"下载数据集 {dataset_id} (类别: {cat})...") + datafields_df = get_datafields( + s=session, + instrument_type="EQUITY", + region=region, + delay=delay, + universe=universe, + dataset_id=dataset_id, + data_type=data_type + ) + + if not datafields_df.empty: + all_datafields.append(datafields_df) + logger.info(f"数据集 {dataset_id} 包含 {len(datafields_df)} 个字段") + + if all_datafields: + # 合并该类别下的所有数据字段 + combined_df = pd.concat(all_datafields, ignore_index=True) + combined_df.drop_duplicates(subset=['id'], inplace=True) + + cache_filename = f"datafields_cache_{region}_{universe}_D{delay}_{cat}.csv" + cache_path = os.path.join(dataset_dir, cache_filename) + combined_df.to_csv(cache_path, index=False) + total_count += len(combined_df) + saved_files.append(cache_filename) + logger.info(f"类别 '{cat}' 数据字段已保存: {cache_path}, 共 {len(combined_df)} 个字段") + else: + logger.warning(f"类别 '{cat}' 未找到数据字段") + + if total_count == 0: + return JSONResponse( + status_code=404, + content={"success": False, "error": "No datafields found for any category"} + ) + + return JSONResponse(content={ + "success": True, + "count": total_count, + "cache_files": saved_files + }) + + except Exception as e: + logger.error(f"下载数据字段失败: {e}") + return JSONResponse( + status_code=500, + content={"success": False, "error": str(e)} + ) + + @app.get("/api/download/{alpha_id}") async def download_results(alpha_id: str): """ diff --git a/simple72/templates/app.js b/simple72/templates/app.js index 22f11b7..3a93c01 100644 --- a/simple72/templates/app.js +++ b/simple72/templates/app.js @@ -34,6 +34,9 @@ document.addEventListener('DOMContentLoaded', async () => { if (config.transformer.data_type) { document.getElementById('dataType').value = config.transformer.data_type; } + if (config.transformer.alpha_id) { + document.getElementById('alphaId').value = config.transformer.alpha_id; + } } } } catch (error) { @@ -46,6 +49,7 @@ const submitBtn = document.getElementById('submitBtn'); const downloadBtn = document.getElementById('downloadBtn'); const loginAndFetchBtn = document.getElementById('loginAndFetchBtn'); const testLLMBtn = document.getElementById('testLLMBtn'); +const downloadDatafieldsBtn = document.getElementById('downloadDatafieldsBtn'); const regionSelect = document.getElementById('region'); const delaySelect = document.getElementById('delay'); const universeSelect = document.getElementById('universe'); @@ -178,17 +182,20 @@ function populateCategoryButtons(categories) { function toggleCategory(btn) { const allBtn = document.getElementById('cat-all'); const isAllBtn = (btn === allBtn); - + if (isAllBtn) { + // 点击"不筛选"按钮：选中所有具体类别 allBtn.style.backgroundColor = '#000080'; allBtn.style.color = 'white'; - + const otherBtns = categoryButtons.querySelectorAll('button:not(#cat-all)'); otherBtns.forEach(b => { - b.style.backgroundColor = '#c0c0c0'; - b.style.color = 'black'; + // 选中所有具体类别按钮 + b.style.backgroundColor = '#000080'; + b.style.color = 'white'; }); } else { + // 点击具体类别按钮 if (btn.style.backgroundColor === 'rgb(0, 0, 128)') { btn.style.backgroundColor = '#c0c0c0'; btn.style.color = 'black'; @@ -196,7 +203,8 @@ function toggleCategory(btn) { btn.style.backgroundColor = '#000080'; btn.style.color = 'white'; } - + + // 检查是否有选中的具体类别 const anySelected = categoryButtons.querySelectorAll('button:not(#cat-all)'); let hasSelected = false; anySelected.forEach(b => { @@ -204,34 +212,72 @@ function toggleCategory(btn) { hasSelected = true; } }); - - if (hasSelected) { - allBtn.style.backgroundColor = '#c0c0c0'; - allBtn.style.color = 'black'; - } else { + + // 如果没有选中任何具体类别，恢复"不筛选"按钮 + if (!hasSelected) { allBtn.style.backgroundColor = '#000080'; allBtn.style.color = 'white'; + // 自动选中所有类别 + anySelected.forEach(b => { + b.style.backgroundColor = '#000080'; + b.style.color = 'white'; + }); + } else { + allBtn.style.backgroundColor = '#c0c0c0'; + allBtn.style.color = 'black'; } } } +// 获取选中的类别列表 +function getSelectedCategories() { + const allBtn = document.getElementById('cat-all'); + const otherBtns = categoryButtons.querySelectorAll('button:not(#cat-all)'); + + // 如果"不筛选"按钮被选中，返回所有具体类别的值 + if (allBtn.style.backgroundColor === 'rgb(0, 0, 128)') { + const allCategories = []; + otherBtns.forEach(b => { + allCategories.push(b.dataset.value); + }); + return allCategories; + } + + // 否则返回选中的具体类别 + const selected = []; + otherBtns.forEach(b => { + if (b.style.backgroundColor === 'rgb(0, 0, 128)') { + selected.push(b.dataset.value); + } + }); + return selected; +} + regionSelect.addEventListener('change', () => { const selectedRegion = regionSelect.value; - - delaySelect.innerHTML = ''; + universeSelect.innerHTML = ''; - delaySelect.disabled = true; universeSelect.disabled = true; - + if (selectedRegion && optionsData[selectedRegion]) { const delays = Object.keys(optionsData[selectedRegion]); + // 保留 delaySelect 的当前值（如果有效） + const currentDelay = delaySelect.value; + delaySelect.innerHTML = ''; delays.forEach(delay => { const option = document.createElement('option'); option.value = delay; option.textContent = delay; + if (delay === currentDelay) { + option.selected = true; + } delaySelect.appendChild(option); }); delaySelect.disabled = false; + // 触发 delay change 事件来更新 universe + if (currentDelay && optionsData[selectedRegion][currentDelay]) { + delaySelect.dispatchEvent(new Event('change')); + } } }); @@ -309,12 +355,13 @@ form.addEventListener('submit', async (e) => { if (result.success) { const successCount = result.expressions_success ? result.expressions_success.length : 0; - const candidateCount = result.candidates ? result.candidates.length : 0; + const candidatesData = result.candidates || {}; + const candidateTemplatesCount = Object.keys(candidatesData).length; const errorCount = result.expressions_error ? result.expressions_error.length : 0; // 显示下载按钮 downloadBtn.style.display = 'block'; - downloadBtn.textContent = '下载结果 (' + successCount + '成功, ' + candidateCount + '候选)'; + downloadBtn.textContent = '下载结果 (' + successCount + '成功, ' + candidateTemplatesCount + '模板)'; downloadBtn.onclick = function() { const alphaId = document.getElementById('alphaId').value.trim(); window.location.href = '/api/download/' + alphaId; @@ -332,3 +379,61 @@ form.addEventListener('submit', async (e) => { submitBtn.textContent = '生成变种'; } }); + +// 下载数据字段缓存按钮 +downloadDatafieldsBtn.addEventListener('click', async () => { + const region = regionSelect.value; + const delay = delaySelect.value; + const universe = universeSelect.value; + + if (!region || !delay || !universe) { + alert('请先选择地区、Delay和股票池'); + return; + } + + const brainUsername = document.getElementById('brainUsername').value.trim(); + const brainPassword = document.getElementById('brainPassword').value.trim(); + + if (!brainUsername || !brainPassword) { + alert('请先输入BRAIN用户名和密码'); + return; + } + + // 获取选中的类别（使用新函数） + const selectedCategories = getSelectedCategories(); + console.log('选中的类别:', selectedCategories); + + downloadDatafieldsBtn.disabled = true; + downloadDatafieldsBtn.textContent = '正在下载...'; + + try { + const requestBody = { + username: brainUsername, + password: brainPassword, + region: region, + delay: parseInt(delay), + universe: universe, + data_type: dataTypeSelect.value, + category: selectedCategories // 总是传递类别列表 + }; + + const response = await fetch('/api/download-datafields', { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + }, + body: JSON.stringify(requestBody) + }); + + const result = await response.json(); + + if (!result.success) { + alert('下载失败: ' + (result.error || '未知错误')); + } + } catch (error) { + alert('下载出错: ' + error.message); + } finally { + downloadDatafieldsBtn.disabled = false; + downloadDatafieldsBtn.textContent = '下载数据字段缓存'; + } +}); diff --git a/simple72/templates/index.html b/simple72/templates/index.html index ea99c49..fa6c378 100644 --- a/simple72/templates/index.html +++ b/simple72/templates/index.html @@ -14,15 +14,6 @@
-
-

使用说明

-
    -
  • 输入您的种子Alpha ID，系统将基于BRAIN平台生成多个变种
    • -
  • 配置您的LLM（支持OpenAI兼容接口）和BRAIN凭证
    • -
  • 点击生成按钮，系统将自动完成所有处理并返回结果
    • -
-
-

BRAIN 凭证

@@ -92,7 +83,7 @@
+ value="10" min="1" max="100">
@@ -117,9 +108,10 @@
- - + +
@@ -139,7 +131,11 @@
- + + + @@ -148,7 +144,14 @@ 下载结果 (ZIP) - +
+

使用说明

+
    +
  • 输入您的种子Alpha ID，系统将基于BRAIN平台生成多个变种
    • +
  • 配置您的LLM（支持OpenAI兼容接口）和BRAIN凭证
    • +
  • 点击生成按钮，系统将自动完成所有处理并返回结果
    • +
+