kg-agent/src/step2_graph_building/graph/production.py

"""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