Copilot commented on code in PR #352: URL: https://github.com/apache/hugegraph-ai/pull/352#discussion_r3330081739
########## text2gremlin/AST_Text2Gremlin/base/Schema.py: ########## @@ -0,0 +1,238 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +""" +图数据库Schema管理模块。 + +负责解析Schema定义和CSV数据文件，为查询生成器提供图结构信息和真实数据实例。 +""" + +import json +import os +import random +from typing import Any + +import pandas as pd + + +class Schema: + def __init__(self, schema_file: str, data_dir: str): + self.data_dir = data_dir + self.vertices: dict[str, dict[str, Any]] = {} + self.edges: dict[str, dict[str, Any]] = {} + self.vertex_data: dict[str, pd.DataFrame] = {} + self.edge_data: dict[str, pd.DataFrame] = {} + + with open(schema_file, encoding="utf-8") as f: + schema_data = json.load(f) + + # 解析 schema 定义 + for item in schema_data.get("schema", []): + label = item["label"] + if item["type"] == "VERTEX": + self.vertices[label] = { + "primary": item.get("primary", None), + "properties": { + prop["name"]: {"type": prop["type"], "optional": prop.get("optional", False)} + for prop in item.get("properties", []) + }, + } + elif item["type"] == "EDGE": + self.edges[label] = { + "source": None, + "destination": None, + "properties": { + prop["name"]: {"type": prop["type"], "optional": prop.get("optional", False)} + for prop in item.get("properties", []) + }, + } + + # 2. 解析 files 定义，获取路径、header行数和边的端点 + self.vertex_files: dict[str, dict] = {} + self.edge_files: dict[str, dict] = {} + for file_info in schema_data.get("files", []): + label = file_info["label"] + path = os.path.join(self.data_dir, file_info["path"]) + header_rows = file_info.get("header", 1) # 获取header行数，默认为1 + + file_details = {"path": path, "header_rows": header_rows} + + is_edge = "SRC_ID" in file_info and "DST_ID" in file_info + if is_edge: + self.edge_files[label] = file_details + if label in self.edges: + self.edges[label]["source"] = file_info["SRC_ID"] + self.edges[label]["destination"] = file_info["DST_ID"] + else: + self.vertex_files[label] = file_details + + def _parse_custom_csv(self, file_path: str, header_line_index: int) -> pd.DataFrame: + """解析自定义多行表头的 CSV 文件。""" + try: + with open(file_path, encoding="utf-8") as f: + lines = f.readlines() + + # 从第二行解析列名 + header_line = lines[header_line_index - 1] + column_defs = header_line.strip().split(",") + column_names = [d.split(":")[0] for d in column_defs] + + # 处理重复的列名（为重复的列添加后缀） + seen = {} + unique_names = [] + for name in column_names: + if name in seen: + seen[name] += 1 + unique_names.append(f"{name}_{seen[name]}") + else: + seen[name] = 0 + unique_names.append(name) + column_names = unique_names + + # 从指定header行之后开始读取数据 + data_lines = lines[header_line_index:] + + if not data_lines: + return pd.DataFrame(columns=column_names) + + # 使用pandas从内存中的字符串列表读取数据 + from io import StringIO + + csv_data = StringIO("".join(data_lines)) + df = pd.read_csv(csv_data, header=None, names=column_names) + return df + + except (FileNotFoundError, IndexError) as e: + print(f"警告: 读取或解析文件失败: {file_path}, 错误: {e}") + return pd.DataFrame() + + def _load_vertex_data(self, label: str): + if label not in self.vertex_data and label in self.vertex_files: + file_info = self.vertex_files[label] + self.vertex_data[label] = self._parse_custom_csv(file_info["path"], file_info["header_rows"]) + + def _load_edge_data(self, label: str): + if label not in self.edge_data and label in self.edge_files: + file_info = self.edge_files[label] + self.edge_data[label] = self._parse_custom_csv(file_info["path"], file_info["header_rows"]) + + # --- Schema 查询方法 (保持不变) --- + def get_vertex_labels(self) -> list[str]: + return list(self.vertices.keys()) + + def get_edge_labels(self) -> list[str]: + return list(self.edges.keys()) + + def get_properties_with_type(self, label: str) -> list[dict[str, str]]: + props_dict = self.vertices.get(label, {}).get("properties", {}) or self.edges.get(label, {}).get( + "properties", {} + ) + return [{"name": name, "type": meta["type"]} for name, meta in props_dict.items()] + + def get_valid_steps(self, current_label: str, element_type: str = "vertex") -> list[dict]: + if element_type == "vertex": + if current_label not in self.vertices: + return [] + valid_steps = [] + outgoing = [l for l, e in self.edges.items() if e["source"] == current_label] + if outgoing: + valid_steps.append({"step": "out", "params": outgoing}) + incoming = [l for l, e in self.edges.items() if e["destination"] == current_label] + if incoming: + valid_steps.append({"step": "in", "params": incoming}) + props = self.get_properties_with_type(current_label) + if props: + valid_steps.append({"step": "properties", "params": [p["name"] for p in props]}) + valid_steps.append({"step": "has", "params": props}) + return valid_steps + return [] + + def get_step_result_label(self, start_label: str, step: dict) -> tuple[str, str]: + step_name, step_param = step.get("step"), step.get("param") + if step_name == "out": + if step_param not in self.edges: + raise KeyError(f"边标签 '{step_param}' 不存在于 schema 中") + return self.edges[step_param]["destination"], "vertex" + if step_name == "in": + if step_param not in self.edges: + raise KeyError(f"边标签 '{step_param}' 不存在于 schema 中") + return self.edges[step_param]["source"], "vertex" + if step_name in ["properties", "has", "values"]: + return start_label, "vertex" + return None, None + + def get_vertex_creation_info(self, label: str) -> dict: + if label not in self.vertices: + return {} + schema_info = self.vertices[label] + required = [name for name, meta in schema_info["properties"].items() if not meta["optional"]] + return {"primary": schema_info.get("primary"), "required": required} + + def get_edge_creation_info(self, label: str) -> tuple[str, str]: + if label in self.edges: + return (self.edges[label]["source"], self.edges[label]["destination"]) + return (None, None) + + def get_updatable_properties(self, label: str) -> list[dict[str, str]]: + if label not in self.vertices: + return [] + schema_info = self.vertices[label] + primary_key = schema_info.get("primary") + return [ + {"name": name, "type": meta["type"]} + for name, meta in schema_info["properties"].items() + if name != primary_key + ] + + def get_instance(self, label: str) -> dict: + """获取单个实例（保持向后兼容）""" + instances = self.get_instances(label, count=1) + return instances[0] if instances else {} + + def get_instances(self, label: str, count: int | None = None) -> list[dict]: + """获取多个实例 + + Args: + label: 标签名 + count: 要获取的实例数量，如果为None则随机选择2-5个 + + Returns: + 实例列表 + """ + + is_edge = label in self.edges + data_cache = self.edge_data if is_edge else self.vertex_data + load_func = self._load_edge_data if is_edge else self._load_vertex_data + + if label not in data_cache: + load_func(label) + + df = data_cache.get(label) + if df is None or df.empty: + return [] + + # 如果没有指定数量，随机选择2-5个 + if count is None: + count = random.randint(2, 5) Review Comment: The default sample size range `2-5` is a magic number embedded in a generic getter. Either lift it into a named module-level constant (e.g., `DEFAULT_SAMPLE_MIN`/`DEFAULT_SAMPLE_MAX`) or read it from the same combination-control config that governs other sampling sizes, so behavior is configurable and discoverable. ########## text2gremlin/AST_Text2Gremlin/base/GremlinBase.py: ########## @@ -0,0 +1,345 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +""" +Gremlin翻译引擎模块。 + +提供Gremlin术语到中文的智能翻译，负责生成自然流畅的中文查询描述。 +""" + +import os +import random + +from .gremlin.GremlinParser import GremlinParser + + +class GremlinBase: + def __init__(self, config): + """ + Gremlin 基础类，作为项目的“工具箱”和“字典”。 + """ + self.config = config + + # 从 GremlinParser 加载rule_names + self.rule_names = GremlinParser.ruleNames + + self.token_dict = {} + self.template = [] # 索引对应 token_dict 的值，每个元素是一个包含多种中文翻译模板的子列表。 + self._initialize_translation_templates() + + # 复用 schema_dict 加载同义词等 + self.schema_dict = {} + self._load_schema_translations() + + def get_rule_name(self, rule_index: int) -> str: + """根据索引获取 ANTLR 规则名。""" + if 0 <= rule_index < len(self.rule_names): + return self.rule_names[rule_index] + return "UnknownRule" + + def _load_schema_translations(self): + """加载schema翻译字典""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + + def collect_config_paths(method_name: str) -> list[str]: + if not hasattr(self.config, method_name): + return [] + try: + configured_paths = getattr(self.config, method_name)() + except Exception as e: + print(f"[INFO] Config path {method_name} not available: {e}") + return [] + if isinstance(configured_paths, list): + return configured_paths + if isinstance(configured_paths, str): + return [configured_paths] + return [] + + def existing_or_default(configured_paths: list[str], default_filename: str) -> list[str]: + existing_paths = [path for path in configured_paths if os.path.exists(path)] + if existing_paths: + return existing_paths + default_path = os.path.join(current_dir, "template", default_filename) + if os.path.exists(default_path): + return [default_path] + return [] + + schema_paths = existing_or_default(collect_config_paths("get_schema_dict_path"), "schema_dict.txt") + syn_paths = existing_or_default(collect_config_paths("get_syn_dict_path"), "syn_dict.txt") + existing_paths = schema_paths + syn_paths + + if existing_paths: + self.load_dict_from_file(existing_paths) + else: + print("[WARNING] No dictionary files found") + + def _initialize_translation_templates(self): + """ + 初始化 Gremlin 步骤的翻译模板。 + """ + # 定义模板数据 + templates_data = { + # --- 起始步骤 --- + "v": ["查询图中的所有顶点", "获取所有节点"], + "V": ["查询图中的所有顶点", "获取所有节点"], # 大写版本 + "e": ["查询图中的所有边", "获取所有关系"], + "E": ["查询图中的所有边", "获取所有关系"], # 大写版本 + "addv": ["添加一个标签为 '{}' 的新顶点"], + "addV": ["添加一个标签为 '{}' 的新顶点"], # 大写版本 + "adde": ["添加一条从一个顶点到另一个顶点的 '{}' 边"], + "addE": ["添加一条从一个顶点到另一个顶点的 '{}' 边"], # 大写版本 + # --- 导航步骤 --- + "out": ["从当前位置出发，沿着 '{}' 方向的出边前进", "找到 '{}' 类型的邻居"], + "in": ["从当前位置出发，沿着 '{}' 方向的入边前进", "找到拥有 '{}' 类型关系的来源"], + "both": ["沿着 '{}' 方向的双向边进行遍历"], + "outE": ["获取出边"], + "inE": ["获取入边"], + "bothE": ["获取双向边"], + "outV": ["从当前边，找到它的出射顶点", "获取边的头节点"], + "inV": ["从当前边，找到它的入射顶点", "获取边的尾节点"], + "otherV": ["获取边的另一端顶点"], + "bothV": ["从当前边，找到它的两个端点"], + # --- 过滤步骤 --- + "hasLabel": ["并筛选出标签为 '{}' 的元素"], + "has": ["并筛选出属性 '{}' 为 '{}' 的元素", "查找其中 '{}' 是 '{}' 的数据"], + "hasId": ["筛选ID为 '{}' 的元素"], + "hasKey": ["筛选包含键 '{}' 的元素"], + "hasValue": ["筛选包含值 '{}' 的元素"], + "where": ["并根据 '{}' 的条件进行过滤"], + "is": ["判断值是否为 '{}'"], + "not": ["取反"], + # --- 数值参数步骤 --- + "limit": ["并限制最多返回 {} 个结果", "取前 {} 条数据"], + "skip": ["跳过前 {} 个结果"], + "tail": ["取最后 {} 个结果"], + "sample": ["随机采样 {} 个结果"], + "range": ["取范围内的结果"], + # --- 转换步骤 --- + "dedup": ["并对结果进行去重"], + "order": ["然后对结果进行排序"], + "simplePath": ["过滤出简单路径"], + "cyclicPath": ["过滤出循环路径"], + # --- 属性访问步骤 --- + "values": ["然后获取它们的属性值", "提取属性值"], # 无参数版本 + "values_with_key": ["然后获取它们的 '{}' 属性值", "提取 '{}' 字段的值"], # 有参数版本 + "properties": ["获取属性对象"], + "properties_with_key": ["获取 '{}' 属性对象"], # 有参数版本 + "valueMap": ["然后以键值对的形式返回它们的属性"], + "valueMap_with_key": ["返回 '{}' 的键值对"], # 有参数版本 + "elementMap": ["获取元素映射"], + "keys": ["获取键"], + "key": ["获取键"], + "value": ["获取属性的值"], + # --- 聚合步骤 --- + "count": ["最后统计结果的总数"], + "group": ["然后对结果进行分组"], + "groupCount": ["分组并统计数量"], + "sum": ["求和"], + "mean": ["求平均值"], + "min": ["求最小值"], + "max": ["求最大值"], + # --- 终端步骤 --- + "tolist": ["转为列表"], + "toList": ["转为列表"], # 大写版本 + "toset": ["转为集合"], + "toSet": ["转为集合"], # 大写版本 + "next": ["获取下一个"], + "hasnext": ["判断是否有下一个"], + "hasNext": ["判断是否有下一个"], # 大写版本 + "trynext": ["尝试获取下一个"], + "tryNext": ["尝试获取下一个"], # 大写版本 + # --- 简单步骤 --- + "fold": ["折叠为列表"], + "unfold": ["展开列表"], + "iterate": ["迭代执行"], + "explain": ["解释查询计划"], + "profile": ["性能分析"], + # --- 其他步骤 --- + "label": ["然后获取它们的标签"], + "id": ["然后获取它们的ID"], + "path": ["然后返回完整的遍历路径"], + "project": ["然后将结果投影为指定字段"], + "by": ["按指定属性进行分组或投影"], + "property": ["并将其 '{}' 属性的值更新为 '{}'"], + "drop": ["最后将这些元素从图中删除", "移除这些数据"], + "as": ["标记为 '{}'", "标记"], + "select": ["选择标记的元素"], + "repeat": ["重复遍历"], + "until": ["直到满足条件"], + "times": ["重复指定次数"], + "emit": ["发射中间结果"], + "choose": ["条件分支"], + "coalesce": ["合并多个遍历"], + "optional": ["可选遍历"], + "union": ["联合多个遍历"], + "match": ["模式匹配"], + "flatMap": ["扁平映射"], + "map": ["映射转换"], + "tree": ["构建树结构"], + # --- 新增步骤 --- + "loops": ["获取当前循环的次数"], + "coin": ["以 {} 的概率保留结果"], + "filter": ["应用过滤条件"], + "and": ["应用逻辑与过滤"], + "or": ["应用逻辑或过滤"], + "aggregate": ["将它们聚合到名为 '{}' 的侧边变量中"], + "store": ["将它们存储到名为 '{}' 的侧边变量中"], + "sideEffect": ["执行附加操作"], + "cap": ["然后取出 '{}' 中存储的内容"], + "sack": ["获取携带的值"], + "barrier": ["等待所有结果到齐"], + "constant": ["映射为常量值"], + "identity": ["保持元素不变"], + "local": ["在本地作用域内执行"], + # --- 剩余15个步骤 --- + # 图算法 + "pageRank": ["计算PageRank值"], + "peerPressure": ["应用同伴压力算法"], + "connectedComponent": ["计算连通分量"], + "shortestPath": ["计算最短路径"], + # 工具步骤 + "math": ["应用数学表达式"], + "subgraph": ["提取子图"], + "timeLimit": ["设置时间限制"], + "inject": ["注入值到遍历"], + "call": ["调用服务"], + "io": ["执行IO操作"], + "mergeE": ["合并边"], + "mergeV": ["合并顶点"], + "with": ["配置选项"], + # 边修改 + "from": ["指定边的起始顶点"], + "to": ["指定边的目标顶点"], + # --- 谓词 --- + # 数值谓词 + "eq": ["等于"], + "neq": ["不等于"], + "gt": ["大于"], + "gte": ["大于等于"], + "lt": ["小于"], + "lte": ["小于等于"], + "between": ["在...之间"], + "inside": ["在范围内"], + "outside": ["在范围外"], + "within": ["在集合中"], + "without": ["不在集合中"], + # 文本谓词 + "startingWith": ["以...开始"], + "endingWith": ["以...结束"], + "containing": ["包含"], + "notStartingWith": ["不以...开始"], + "notEndingWith": ["不以...结束"], + "notContaining": ["不包含"], + "regex": ["匹配正则表达式"], + "notRegex": ["不匹配正则表达式"], + } + + # 填充 self.token_dict 和 self.template + # 保持原始大小写，不进行转换 + for index, (key, value) in enumerate(templates_data.items()): + self.token_dict[key] = index + self.template.append(value) Review Comment: `templates_data` contains both lowercase and PascalCase duplicate keys (e.g., `v`/`V`, `addv`/`addV`, `tolist`/`toList`) mapping to the same translation lists. This duplication is error-prone — updating one variant but not the other will cause inconsistent translations. Consider normalizing tokens at lookup time (e.g., maintain a single canonical map and add an alias map, or look up by `.lower()`/canonical name in `get_token_desc`). ########## text2gremlin/AST_Text2Gremlin/base/generator.py: ########## @@ -0,0 +1,439 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +""" +Gremlin语料库生成器主入口脚本。 + +从Gremlin查询模板生成大量多样化的查询-描述对，用于Text-to-Gremlin任务的训练数据。 +""" + +import json +import os +from datetime import datetime + +from antlr4 import CommonTokenStream, InputStream +from antlr4.error.ErrorListener import ErrorListener + +from .Config import Config +from .gremlin.GremlinLexer import GremlinLexer +from .gremlin.GremlinParser import GremlinParser +from .GremlinBase import GremlinBase +from .GremlinTransVisitor import GremlinTransVisitor +from .Schema import Schema +from .TraversalGenerator import TraversalGenerator + + +class SyntaxErrorListener(ErrorListener): + """私有错误监听器类，捕获语法错误。""" + + def __init__(self): + super().__init__() + self.has_error = False + self.error_message = "" + + def syntaxError(self, recognizer, offendingSymbol, line, column, msg, e): + """当语法错误发生时，此方法被调用。""" + self.has_error = True + self.error_message = f"Syntax Error at line {line}, column {column}: {msg}" + + +def check_gremlin_syntax(query_string: str) -> tuple[bool, str]: + """ + 检查给定的Gremlin查询语句的语法。 + + Args: + query_string: The Gremlin query to check. + + Returns: + A tuple containing: + - bool: True if syntax is correct, False otherwise. + - str: An error message if syntax is incorrect, or "Syntax OK" if correct. + """ + try: + input_stream = InputStream(query_string) + lexer = GremlinLexer(input_stream) + token_stream = CommonTokenStream(lexer) + parser = GremlinParser(token_stream) + + # 移除默认的控制台错误监听器 + lexer.removeErrorListeners() + parser.removeErrorListeners() + + # 添加自定义的监听器 + error_listener = SyntaxErrorListener() + lexer.addErrorListener(error_listener) + parser.addErrorListener(error_listener) + + # 尝试解析查询 + parser.queryList() + + if error_listener.has_error: + return (False, error_listener.error_message) + else: + return (True, "Syntax OK") + + except Exception as e: + return (False, f"Parser Exception: {e!s}") + + +def generate_corpus_from_template( + template_string: str, config: Config, schema: Schema, gremlin_base: GremlinBase, global_corpus_dict: dict +) -> tuple[int, dict]: + """ + 执行单个 Gremlin 模板字符串的完整 pipeline。 + + Args: + template_string: 用作模板的 Gremlin query。 + config: 加载的 Config 对象。 + schema: 加载的 Schema 对象。 + gremlin_base: 加载的 GremlinBase 对象。 + global_corpus_dict: 用于存储唯一 query-description 对的全局字典。 + + Returns: + tuple: (添加到全局语料库的新的唯一对的数量, 处理统计信息) + """ + # 初始化统计信息 + stats = { + "success": False, + "error_stage": "", + "error_message": "", + "generated_count": 0, + "new_pairs_count": 0, + "duplicate_count": 0, + "syntax_error_count": 0, + } + + try: + # ANTLR 解析为 AST,并提取模版 + visitor = GremlinTransVisitor() + recipe = visitor.parse_and_visit(template_string) + + if not recipe: + stats["error_stage"] = "recipe_extraction" + stats["error_message"] = "Recipe extraction failed" + return 0, stats + + if not hasattr(recipe, "steps") or not recipe.steps: + stats["error_stage"] = "recipe_validation" + stats["error_message"] = "Recipe has no steps" + return 0, stats + + # 泛化 + generator = TraversalGenerator(schema, recipe, gremlin_base) + corpus = generator.generate() + + if not corpus: + stats["error_stage"] = "generation" + stats["error_message"] = "Generator returned empty corpus" + return 0, stats + + stats["generated_count"] = len(corpus) + + # 语法检查 & 全局去重 + new_pairs_count = 0 + duplicate_count = 0 + syntax_error_count = 0 + + for query, description in corpus: + try: + # 先判重，避免对重复项做语法检查 + if query in global_corpus_dict: + duplicate_count += 1 + continue + + # 再进行语法检查 + is_valid, _error_msg = check_gremlin_syntax(query) + + if not is_valid: + syntax_error_count += 1 + continue + + # 新的查询且语法正确，添加到全局字典 + global_corpus_dict[query] = description + new_pairs_count += 1 + + except Exception: + syntax_error_count += 1 + continue + + # 更新统计信息 + stats["new_pairs_count"] = new_pairs_count + stats["duplicate_count"] = duplicate_count + stats["syntax_error_count"] = syntax_error_count + stats["success"] = True + + # 添加生成数量的警告信息 + if stats["generated_count"] > 5000: + stats["warning"] = f"由于本条模版的Recipe复杂,生成了大量查询({stats['generated_count']}条)" Review Comment: The threshold `5000` is duplicated between `generate_corpus_from_template` (line 181) and the per-template progress print in `generate_gremlin_corpus` (line 289). Extract a module-level constant (e.g., `LARGE_GENERATION_THRESHOLD = 5000`) to keep the two checks in sync. ########## text2gremlin/AST_Text2Gremlin/base/CombinationController.py: ########## @@ -0,0 +1,367 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +""" +组合爆炸控制器 + +提供统一的配置驱动的控制策略，适用于所有Gremlin步骤和谓词的泛化生成。 +""" + +import random + + +class CombinationController: + """组合爆炸控制器 - 基于配置文件的统一控制策略""" + + def __init__(self, config: dict): + """ + 初始化控制器 + + Args: + config: 从combination_control_config.json加载的配置字典 + """ + self.config = config + + # 验证必要配置项并加载 + try: + # 链长度分类阈值 + self.chain_thresholds = config["chain_thresholds"] + + # 随机增强控制 + self.random_enhancement = config["random_enhancement"] + + # 数据填充策略 + self.value_fill = config["value_fill_strategy"] + + # 属性泛化策略 + self.property_gen = config["property_generalization"] + except KeyError as e: + raise ValueError(f"缺少必要配置项: {e}") from None + + # 总数限制（可选） + self.max_total = config.get("max_total_combinations", {}) + + # 验证关键类别的存在性 + # chain_thresholds 只需要 short, medium, long（ultra 通过 else 分支隐式定义） + for category in ("short", "medium", "long"): + if category not in self.chain_thresholds: + raise ValueError(f"chain_thresholds 缺少 '{category}' 配置") Review Comment: The validation duplicates the literal category tuple in two places (here and lines 66). Extract `CHAIN_CATEGORIES = ('short', 'medium', 'long', 'ultra')` as a class/module constant and reuse it (and in `get_chain_category`) to keep them in sync. ########## text2gremlin/AST_Text2Gremlin/base/GremlinBase.py: ########## @@ -0,0 +1,345 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + + +""" +Gremlin翻译引擎模块。 + +提供Gremlin术语到中文的智能翻译，负责生成自然流畅的中文查询描述。 +""" + +import os +import random + +from .gremlin.GremlinParser import GremlinParser + + +class GremlinBase: + def __init__(self, config): + """ + Gremlin 基础类，作为项目的“工具箱”和“字典”。 + """ + self.config = config + + # 从 GremlinParser 加载rule_names + self.rule_names = GremlinParser.ruleNames + + self.token_dict = {} + self.template = [] # 索引对应 token_dict 的值，每个元素是一个包含多种中文翻译模板的子列表。 + self._initialize_translation_templates() + + # 复用 schema_dict 加载同义词等 + self.schema_dict = {} + self._load_schema_translations() + + def get_rule_name(self, rule_index: int) -> str: + """根据索引获取 ANTLR 规则名。""" + if 0 <= rule_index < len(self.rule_names): + return self.rule_names[rule_index] + return "UnknownRule" + + def _load_schema_translations(self): + """加载schema翻译字典""" + current_dir = os.path.dirname(os.path.abspath(__file__)) + + def collect_config_paths(method_name: str) -> list[str]: + if not hasattr(self.config, method_name): + return [] + try: + configured_paths = getattr(self.config, method_name)() + except Exception as e: + print(f"[INFO] Config path {method_name} not available: {e}") + return [] + if isinstance(configured_paths, list): + return configured_paths + if isinstance(configured_paths, str): + return [configured_paths] + return [] + + def existing_or_default(configured_paths: list[str], default_filename: str) -> list[str]: + existing_paths = [path for path in configured_paths if os.path.exists(path)] + if existing_paths: + return existing_paths + default_path = os.path.join(current_dir, "template", default_filename) + if os.path.exists(default_path): + return [default_path] + return [] + + schema_paths = existing_or_default(collect_config_paths("get_schema_dict_path"), "schema_dict.txt") + syn_paths = existing_or_default(collect_config_paths("get_syn_dict_path"), "syn_dict.txt") + existing_paths = schema_paths + syn_paths + + if existing_paths: + self.load_dict_from_file(existing_paths) + else: + print("[WARNING] No dictionary files found") + + def _initialize_translation_templates(self): + """ + 初始化 Gremlin 步骤的翻译模板。 + """ + # 定义模板数据 + templates_data = { + # --- 起始步骤 --- + "v": ["查询图中的所有顶点", "获取所有节点"], + "V": ["查询图中的所有顶点", "获取所有节点"], # 大写版本 + "e": ["查询图中的所有边", "获取所有关系"], + "E": ["查询图中的所有边", "获取所有关系"], # 大写版本 + "addv": ["添加一个标签为 '{}' 的新顶点"], + "addV": ["添加一个标签为 '{}' 的新顶点"], # 大写版本 + "adde": ["添加一条从一个顶点到另一个顶点的 '{}' 边"], + "addE": ["添加一条从一个顶点到另一个顶点的 '{}' 边"], # 大写版本 + # --- 导航步骤 --- + "out": ["从当前位置出发，沿着 '{}' 方向的出边前进", "找到 '{}' 类型的邻居"], + "in": ["从当前位置出发，沿着 '{}' 方向的入边前进", "找到拥有 '{}' 类型关系的来源"], + "both": ["沿着 '{}' 方向的双向边进行遍历"], + "outE": ["获取出边"], + "inE": ["获取入边"], + "bothE": ["获取双向边"], + "outV": ["从当前边，找到它的出射顶点", "获取边的头节点"], + "inV": ["从当前边，找到它的入射顶点", "获取边的尾节点"], + "otherV": ["获取边的另一端顶点"], + "bothV": ["从当前边，找到它的两个端点"], + # --- 过滤步骤 --- + "hasLabel": ["并筛选出标签为 '{}' 的元素"], + "has": ["并筛选出属性 '{}' 为 '{}' 的元素", "查找其中 '{}' 是 '{}' 的数据"], + "hasId": ["筛选ID为 '{}' 的元素"], + "hasKey": ["筛选包含键 '{}' 的元素"], + "hasValue": ["筛选包含值 '{}' 的元素"], + "where": ["并根据 '{}' 的条件进行过滤"], + "is": ["判断值是否为 '{}'"], + "not": ["取反"], + # --- 数值参数步骤 --- + "limit": ["并限制最多返回 {} 个结果", "取前 {} 条数据"], + "skip": ["跳过前 {} 个结果"], + "tail": ["取最后 {} 个结果"], + "sample": ["随机采样 {} 个结果"], + "range": ["取范围内的结果"], + # --- 转换步骤 --- + "dedup": ["并对结果进行去重"], + "order": ["然后对结果进行排序"], + "simplePath": ["过滤出简单路径"], + "cyclicPath": ["过滤出循环路径"], + # --- 属性访问步骤 --- + "values": ["然后获取它们的属性值", "提取属性值"], # 无参数版本 + "values_with_key": ["然后获取它们的 '{}' 属性值", "提取 '{}' 字段的值"], # 有参数版本 + "properties": ["获取属性对象"], + "properties_with_key": ["获取 '{}' 属性对象"], # 有参数版本 + "valueMap": ["然后以键值对的形式返回它们的属性"], + "valueMap_with_key": ["返回 '{}' 的键值对"], # 有参数版本 + "elementMap": ["获取元素映射"], + "keys": ["获取键"], + "key": ["获取键"], + "value": ["获取属性的值"], + # --- 聚合步骤 --- + "count": ["最后统计结果的总数"], + "group": ["然后对结果进行分组"], + "groupCount": ["分组并统计数量"], + "sum": ["求和"], + "mean": ["求平均值"], + "min": ["求最小值"], + "max": ["求最大值"], + # --- 终端步骤 --- + "tolist": ["转为列表"], + "toList": ["转为列表"], # 大写版本 + "toset": ["转为集合"], + "toSet": ["转为集合"], # 大写版本 + "next": ["获取下一个"], + "hasnext": ["判断是否有下一个"], + "hasNext": ["判断是否有下一个"], # 大写版本 + "trynext": ["尝试获取下一个"], + "tryNext": ["尝试获取下一个"], # 大写版本 + # --- 简单步骤 --- + "fold": ["折叠为列表"], + "unfold": ["展开列表"], + "iterate": ["迭代执行"], + "explain": ["解释查询计划"], + "profile": ["性能分析"], + # --- 其他步骤 --- + "label": ["然后获取它们的标签"], + "id": ["然后获取它们的ID"], + "path": ["然后返回完整的遍历路径"], + "project": ["然后将结果投影为指定字段"], + "by": ["按指定属性进行分组或投影"], + "property": ["并将其 '{}' 属性的值更新为 '{}'"], + "drop": ["最后将这些元素从图中删除", "移除这些数据"], + "as": ["标记为 '{}'", "标记"], + "select": ["选择标记的元素"], + "repeat": ["重复遍历"], + "until": ["直到满足条件"], + "times": ["重复指定次数"], + "emit": ["发射中间结果"], + "choose": ["条件分支"], + "coalesce": ["合并多个遍历"], + "optional": ["可选遍历"], + "union": ["联合多个遍历"], + "match": ["模式匹配"], + "flatMap": ["扁平映射"], + "map": ["映射转换"], + "tree": ["构建树结构"], + # --- 新增步骤 --- + "loops": ["获取当前循环的次数"], + "coin": ["以 {} 的概率保留结果"], + "filter": ["应用过滤条件"], + "and": ["应用逻辑与过滤"], + "or": ["应用逻辑或过滤"], + "aggregate": ["将它们聚合到名为 '{}' 的侧边变量中"], + "store": ["将它们存储到名为 '{}' 的侧边变量中"], + "sideEffect": ["执行附加操作"], + "cap": ["然后取出 '{}' 中存储的内容"], + "sack": ["获取携带的值"], + "barrier": ["等待所有结果到齐"], + "constant": ["映射为常量值"], + "identity": ["保持元素不变"], + "local": ["在本地作用域内执行"], + # --- 剩余15个步骤 --- + # 图算法 + "pageRank": ["计算PageRank值"], + "peerPressure": ["应用同伴压力算法"], + "connectedComponent": ["计算连通分量"], + "shortestPath": ["计算最短路径"], + # 工具步骤 + "math": ["应用数学表达式"], + "subgraph": ["提取子图"], + "timeLimit": ["设置时间限制"], + "inject": ["注入值到遍历"], + "call": ["调用服务"], + "io": ["执行IO操作"], + "mergeE": ["合并边"], + "mergeV": ["合并顶点"], + "with": ["配置选项"], + # 边修改 + "from": ["指定边的起始顶点"], + "to": ["指定边的目标顶点"], + # --- 谓词 --- + # 数值谓词 + "eq": ["等于"], + "neq": ["不等于"], + "gt": ["大于"], + "gte": ["大于等于"], + "lt": ["小于"], + "lte": ["小于等于"], + "between": ["在...之间"], + "inside": ["在范围内"], + "outside": ["在范围外"], + "within": ["在集合中"], + "without": ["不在集合中"], + # 文本谓词 + "startingWith": ["以...开始"], + "endingWith": ["以...结束"], + "containing": ["包含"], + "notStartingWith": ["不以...开始"], + "notEndingWith": ["不以...结束"], + "notContaining": ["不包含"], + "regex": ["匹配正则表达式"], + "notRegex": ["不匹配正则表达式"], + } + + # 填充 self.token_dict 和 self.template + # 保持原始大小写，不进行转换 + for index, (key, value) in enumerate(templates_data.items()): + self.token_dict[key] = index + self.template.append(value) + + def get_token_desc(self, token_key: str, *args) -> str: + """ + 根据 token 和参数获取一个随机的、格式化后的中文描述。 + """ + key = token_key # 保持原始大小写 + if key in self.token_dict: + index = self.token_dict[key] + # 随机选择一个模板 + selected_template = random.choice(self.template[index]) + try: + # 翻译参数中的schema术语 + translated_args = [] + for arg in args: + if isinstance(arg, str): + # 尝试翻译schema术语 + translated_arg = self.get_schema_desc(arg) + translated_args.append(translated_arg) + else: + translated_args.append(arg) + + # 使用翻译后的参数格式化模板 + return selected_template.format(*translated_args) + except (IndexError, KeyError): + # 如果参数数量不匹配，返回原始模板 + return selected_template + return "" # 如果 token 不存在，返回空字符串 + + # 复用的通用方法 + def merge_desc(self, desc_list: list) -> str: + """合并多个描述片段，移除空字符串并用合适的连接词连接。""" + # 过滤掉空字符串 + filtered_list = [s for s in desc_list if s and s.strip()] + return ",".join(filtered_list) + + def load_dict_from_file(self, file_paths: list): + """从文件加载字典，例如同义词词典。""" + for file_path in file_paths: + if not os.path.exists(file_path): + print(f"[WARNING] Dictionary file not found: {file_path}") + continue + with open(file_path, encoding="utf-8") as file: + for line in file: + elements = line.strip().split() + if elements: + key = elements[0] + values = elements[1:] + self.schema_dict[key] = values + + def get_schema_desc(self, key: str) -> str: + """从加载的字典中获取一个随机的同义词或描述。""" + try: + # 确保键存在 + if self.schema_dict.get(key): + return random.choice(self.schema_dict[key]) + return key # 如果没有同义词，返回原词 + except KeyError: + return key + + +if __name__ == "__main__": + # 临时创建config 对象，用于测试 + class MockConfig: + def get_schema_dict_path(self): + return "./template/schema_dict.txt" + + def get_syn_dict_path(self): + return "./template/syn_dict.txt" Review Comment: `MockConfig` returns plain strings, but `_load_schema_translations` calls `collect_config_paths`, which only treats list or str returns — that's fine — yet `Config.get_schema_dict_path` (the production code) returns a list per `config_example.json`. The inconsistency means the `__main__` smoke test exercises a different code path than production. Consider returning a list here to keep the smoke test aligned with real usage. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
