"""Template-driven production graph, independent of the legacy fixed schemas.""" from __future__ import annotations import hashlib import json import logging import math from collections import Counter, defaultdict from contextlib import closing from dataclasses import asdict, dataclass from datetime import date, datetime from pathlib import Path import numpy as np from openpyxl import load_workbook from step1_data_aggregation.template_mapping import load_templates class GraphValidationError(ValueError): def __init__(self, issues): self.issues = issues super().__init__(f"生产数据校验失败，共 {len(issues)} 项；见构图报告") @dataclass(frozen=True) class Rule: source: str target: str source_field: str target_field: str method: str edge: str location: str @property def id(self) -> str: return digest([self.source, self.target, self.source_field, self.target_field, self.method, self.edge]) def digest(value) -> str: return hashlib.sha256(json.dumps(value, ensure_ascii=False, sort_keys=True, allow_nan=False).encode()).hexdigest() def missing(value) -> bool: return value is None or isinstance(value, str) and not value.strip() def scalar(value): if isinstance(value, (datetime, date)): return value.isoformat() if isinstance(value, float) and not math.isfinite(value): return None return value def property_values(value): """Merged fields are alternatives, never concatenated matching text.""" return [v for v in (value if isinstance(value, list) else [value]) if not missing(v)] def merge_value(left, right): values = {json.dumps(v, ensure_ascii=False, sort_keys=True, allow_nan=False): v for v in property_values(left) + property_values(right)} ordered = [values[k] for k in sorted(values)] return ordered[0] if len(ordered) == 1 else ordered or None def neo4j_properties(props): """Neo4j arrays must be homogeneous; retain mixed-type originals as JSON.""" result = dict(props) mixed = {} for field, value in props.items(): if isinstance(value, list) and len({type(v) for v in value}) > 1: mixed[field] = value result[field] = [json.dumps(v, ensure_ascii=False) for v in value] if mixed: result['_kg_mixed_values_json'] = json.dumps(mixed, ensure_ascii=False) return result def load_rules(path: Path, fields: dict[str, list[str]]) -> list[Rule]: required = ['起点', '终点', '连接字段', '判断方法', '边名'] rules, seen = [], set() with closing(load_workbook(path, read_only=True, data_only=True)) as wb: for sheet in wb: rows = iter(sheet.values) headers = [str(v).strip() if v is not None else '' for v in next(rows, ())] if not any(headers): continue if len(headers) != 5 or set(headers) != set(required): raise ValueError(f'{path}:{sheet.title}: 必须包含且仅包含五个关系列 {required}') for rownum, values in enumerate(rows, 2): if all(missing(v) for v in values): continue row = dict(zip(headers, values)) loc = f'{path}:{sheet.title}!{rownum}' if any(missing(row.get(k)) for k in required): raise ValueError(f'{loc}: 关系定义有空字段') src, dst, link, method, edge = [str(row[k]).strip() for k in required] if src not in fields or dst not in fields: raise ValueError(f'{loc}: 未知节点类型 {src}/{dst}') # Resolve against real field names; hyphens inside a field are valid. pairs = [(a, b) for a in fields[src] for b in fields[dst] if f'{a}-{b}' == link] if len(pairs) != 1: raise ValueError(f'{loc}: 连接字段不存在或有歧义: {link}') if method not in {'相等', '包含', '评分'}: raise ValueError(f'{loc}: 未知判断方法 {method}') rule = Rule(src, dst, *pairs[0], method, edge, loc) if rule.id in seen: raise ValueError(f'{loc}: 重复关系规则') seen.add(rule.id) rules.append(rule) if not rules: raise ValueError(f'{path}: 没有关系规则') return rules def read_production(template_dir: Path, production_dir: Path, *, keys=None, validation_warnings=None): """Merge valid keys; skip missing keys and optionally collect warning details.""" templates = load_templates(template_dir) fields = {t.name: [f.template_field for f in t.fields] for t in templates} if set(keys or {}) - fields.keys(): raise ValueError('主键配置包含未知模板') nodes, stats, issues = {}, {}, [] for template in templates: name = template.name path = production_dir / f'{name}.xlsx' if any(f.startswith('_kg_') for f in fields[name]): raise ValueError(f'{name}: _kg_ 为保留属性前缀') primary = (keys or {}).get(name, []) if not isinstance(primary, list) or not primary or len(primary) != len(set(primary)) or any(f not in fields[name] for f in primary): raise ValueError(f'{name}: 无效主键配置') unique = {} source_count = 0 skipped_rows = 0 with closing(load_workbook(path, read_only=True, data_only=True)) as wb: if wb.sheetnames != ['数据']: raise ValueError(f'{path}: 需要且仅需要“数据”工作表') rows = iter(wb['数据'].values) headers = list(next(rows, ())) if headers != fields[name]: raise ValueError(f'{path}: 表头与当前模板不一致') for rownum, values in enumerate(rows, 2): if all(missing(v) for v in values): continue source_count += 1 props = {k: scalar(v) for k, v in zip(headers, values)} missing_fields = [k for k in primary if missing(props[k])] if missing_fields: issues.append({'type': 'missing_key', 'template': name, 'file': str(path), 'sheet': '数据', 'row': rownum, 'fields': missing_fields, 'action': 'skipped'}) skipped_rows += 1 continue identity = [props[k] for k in primary] node_id = digest([name, identity]) if node_id in unique: current = unique[node_id]['properties'] for field in headers: current[field] = merge_value(current[field], props[field]) unique[node_id]['rows'].append(rownum) else: unique[node_id] = {'id': node_id, 'properties': {k: merge_value(None, v) for k, v in props.items()}, 'rows': [rownum], 'file': str(path.resolve()), 'sheet': '数据'} nodes[name] = list(unique.values()) stats[name] = {'source_rows': source_count, 'node_count': len(unique), 'identity_fields': primary, 'identity_mode': 'key', 'skipped_rows': skipped_rows, 'merged_rows': source_count - skipped_rows - len(unique), 'merged_key_groups': sum(len(n['rows']) > 1 for n in unique.values()), 'multivalue_fields': dict(Counter( k for n in unique.values() for k, v in n['properties'].items() if isinstance(v, list))), 'merge_policy': 'distinct_nonempty_values', 'relation_value_policy': 'any_pair_matches'} if validation_warnings is not None: validation_warnings.extend(issues) logger = logging.getLogger(__name__) for issue in issues: logger.warning('警告：%s [数据] 第 %s 行缺少主键字段 %s，已跳过该记录', issue['file'], issue['row'], '、'.join(issue['fields'])) if issues: logger.warning('共跳过 %s 条缺少主键的记录，其余有效记录继续构图', len(issues)) return templates, fields, nodes, stats class LocalEncoder: def __init__(self): self.model = None self.cache = {} def encode(self, texts): uncached = list(dict.fromkeys(t for t in texts if t not in self.cache)) if uncached: if self.model is None: from sentence_transformers import SentenceTransformer from step2_graph_building.config import get_embedding_model_dir path = get_embedding_model_dir() if not path.is_dir(): raise ValueError(f'本地嵌入模型不存在: {path}') self.model = SentenceTransformer(str(path), device='cpu', local_files_only=True) vectors = self.model.encode(uncached, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False) self.cache.update(zip(uncached, vectors)) return np.asarray([self.cache[t] for t in texts]) def match_rule(rule, nodes, threshold, encoder): left, right = defaultdict(list), defaultdict(list) for group, name, field in ((left, rule.source, rule.source_field), (right, rule.target, rule.target_field)): for node in nodes[name]: for value in property_values(node['properties'][field]): if rule.method in {'包含', '评分'} and not isinstance(value, str): raise ValueError(f'{rule.location}: {field} 必须为文本以执行 {rule.method}') # Exact comparison preserves types and whitespace. key = (type(value).__name__, value) if rule.method == '相等' else value group[key].append(node['id']) def expand(a, b, score=None): for src in left[a]: for dst in right[b]: yield {'source': src, 'target': dst, 'score': score} if rule.method == '相等': for key in left.keys() & right.keys(): yield from expand(key, key) elif rule.method == '包含': for a in left: for b in right: if a in b or b in a: yield from expand(a, b) elif left and right: atexts, btexts = list(left), list(right) av, bv = encoder.encode(atexts), encoder.encode(btexts) # Normalize here as well so injected encoders obey cosine semantics. av = np.asarray(av, dtype=np.float64) bv = np.asarray(bv, dtype=np.float64) for vectors in (av, bv): norms = np.linalg.norm(vectors, axis=1, keepdims=True) if not np.isfinite(vectors).all() or (norms == 0).any(): raise ValueError('嵌入模型返回无效向量') vectors /= norms for start in range(0, len(av), 256): for offset in range(0, len(bv), 256): scores = av[start:start+256] @ bv[offset:offset+256].T for i, j in np.argwhere(scores > threshold): yield from expand(atexts[start+i], btexts[offset+j], float(scores[i, j])) def prepare_graph(template_dir: Path, production_dir: Path, relation_path: Path, *, threshold: float, keys=None, encoder=None): if not math.isfinite(threshold) or not -1 <= threshold <= 1: raise ValueError('相似度阈值必须是 [-1, 1] 内的有限数') validation_warnings = [] templates, fields, nodes, stats = read_production( template_dir, production_dir, keys=keys, validation_warnings=validation_warnings) rules = load_rules(relation_path, fields) encoder = encoder or LocalEncoder() edges, relation_stats = {}, [] for rule in rules: by_pair = {} for match in match_rule(rule, nodes, threshold, encoder): pair = (match['source'], match['target']) previous = by_pair.get(pair) if previous is None or (match['score'] is not None and match['score'] > previous['score']): by_pair[pair] = match matches = [by_pair[pair] for pair in sorted(by_pair)] edges[rule.id] = matches source_degree, target_degree = Counter(), Counter() for edge in matches: source_degree[edge['source']] += 1 target_degree[edge['target']] += 1 relation_stats.append({**asdict(rule), 'id': rule.id, 'count': len(matches), 'unmatched_source_nodes': len(nodes[rule.source]) - len(source_degree), 'cardinality': ('N' if max(target_degree.values(), default=0) > 1 else '1') + ':' + ('N' if max(source_degree.values(), default=0) > 1 else '1') if matches else 'unknown'}) return {'templates': templates, 'fields': fields, 'nodes': nodes, 'rules': rules, 'edges': edges, 'node_stats': stats, 'relation_stats': relation_stats, 'threshold': threshold, 'warnings': validation_warnings, 'skipped_rows': len(validation_warnings)} def identifier(value: str) -> str: return '`' + value.replace('`', '``') + '`' def write_graph(graph, driver, *, build_id: str, batch_size=500): """Stage a separately tagged build; never erase previous/user-owned graphs.""" if batch_size < 1: raise ValueError('batch_size must be positive') driver.execute_query('CREATE CONSTRAINT step2_identity IF NOT EXISTS ' 'FOR (n:_Step2Record) REQUIRE (n._kg_build, n._kg_id) IS UNIQUE') for name, nodes in graph['nodes'].items(): for start in range(0, len(nodes), batch_size): batch = [{'id': n['id'], 'props': {**neo4j_properties(n['properties']), '_kg_file': n['file'], '_kg_sheet': n['sheet'], '_kg_rows': n['rows']}} for n in nodes[start:start+batch_size]] driver.execute_query( f'UNWIND $rows AS row MERGE (n:_Step2Record:{identifier(name)} ' '{_kg_build: $build, _kg_id: row.id}) SET n += row.props', rows=batch, build=build_id) for rule in graph['rules']: edges = graph['edges'][rule.id] for start in range(0, len(edges), batch_size): driver.execute_query( 'UNWIND $rows AS row MATCH (a:_Step2Record {_kg_build: $build, _kg_id: row.source}), ' '(b:_Step2Record {_kg_build: $build, _kg_id: row.target}) ' f'MERGE (a)-[r:{identifier(rule.edge)} {{_kg_rule: $rule}}]->(b) ' 'SET r._kg_method=$method, r._kg_score=row.score, r._kg_threshold=$threshold', rows=edges[start:start+batch_size], build=build_id, rule=rule.id, method=rule.method, threshold=graph['threshold'] if rule.method == '评分' else None) result = driver.execute_query( 'MATCH (n:_Step2Record {_kg_build:$build}) OPTIONAL MATCH (n)-[r]->(m:_Step2Record {_kg_build:$build}) ' 'RETURN count(DISTINCT n) AS nodes, count(r) AS edges', build=build_id).records expected = {'nodes': sum(map(len, graph['nodes'].values())), 'edges': sum(map(len, graph['edges'].values()))} if result != [expected]: raise RuntimeError(f'Neo4j 数量校验失败: expected={expected}, actual={result}') return expected