deepdata/my_ops/text_quality_evaluator/process.py

# -*- coding: utf-8 -*-
"""
Text Quality Evaluator Process Logic
基于 LLM 的长文本多维度质量评估算子
"""
import re
import time
import json
import logging
import os
import sys
from typing import Dict, Any, List
from concurrent.futures import ThreadPoolExecutor, as_completed

try:
    from deepdata.core.base_op import LLM
    from deepdata.common.utils.text_splitter import TextSplitter
except ImportError:
    try:
        executor_path = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..', 'runtime', 'python-executor'))
        if executor_path not in sys.path:
            sys.path.insert(0, executor_path)
    
        from deepdata.core.base_op import LLM
        from deepdata.common.utils.text_splitter import TextSplitter
    except Exception as e:
        print(f"{e}")
        sys.exit(-1)
    
# 内置默认配置
DEFAULT_DIMENSIONS = '[{"name": "连贯性", "desc": "逻辑通顺度"}, {"name": "信息密度", "desc": "有效信息占比"}]'
DEFAULT_PROMPT = "请作为资深数据专家，对以下文本进行多维度质量打分（0.0~1.0）：\n\n【待评估文本】\n{context}\n\n【评估维度】\n{dimensions}\n\n请严格按照格式返回对应数量的逗号分隔的浮点数，不要输出其他内容。"

class TextQualityEvaluator(LLM):
    """
    文本质量评估算子
    """

    def __init__(self, *args, **kwargs):
        # A. 解析 Headers (必须转为 Dict，否则 requests 报错)
        llm_headers = kwargs.get("LLMHeaders")
        if isinstance(llm_headers, str):
            try:
                llm_headers = json.loads(llm_headers)
            except:
                llm_headers = {"Content-Type": "application/json"}
        kwargs["LLMHeaders"] = llm_headers

        # B. 解析 Body (必须转为 Dict)
        llm_body = kwargs.get("LLMBody")
        if isinstance(llm_body, str):
            try:
                llm_body = json.loads(llm_body)
            except:
                llm_body = {}
        kwargs["LLMBody"] = llm_body
        
        # C. 转换 accessType (将 "true"/"false" 转为 布尔值)
        access_type = kwargs.get("accessType")
        if isinstance(access_type, str):
            kwargs["accessType"] = access_type.lower() == "true"
            
        # D. 替换模型名称
        model_name = kwargs.get("LLMModel", "qwen-plus")
        if isinstance(llm_body, dict):
            llm_body["model"] = model_name
            kwargs["LLMBody"] = llm_body

        # ==========================================
        # 3. 调用父类初始化
        # ==========================================
        super(TextQualityEvaluator, self).__init__(*args, **kwargs)
        
        self.task_id = kwargs.get("taskId", "default_id")
        self.max_workers = int(kwargs.get("max_workers", 5))
        self.chunk_size = int(kwargs.get("chunk_size", 2000))
        self.chunk_overlap = int(kwargs.get("chunk_overlap", 100))
        
        dim_conf = kwargs.get("dimensions")
        if not dim_conf: dim_conf = DEFAULT_DIMENSIONS
        if isinstance(dim_conf, str):
            try: self.dimensions = json.loads(dim_conf)
            except: self.dimensions = []
        else: self.dimensions = dim_conf
            
        self.prompt_template = kwargs.get("prompt_template") or DEFAULT_PROMPT
        self.pattern = r'\d+\.\d+'

    def _split_text(self, text: str) -> List[str]:
        """文本分块逻辑"""
        if not text: return []
        
        try:
            splitter = TextSplitter(1024 * 1024 * 10, self.chunk_size, self.chunk_overlap)
            return splitter.split_text(text)
        except Exception: 
            pass
        
        # 降级方案
        chunks = []
        start = 0
        while start < len(text):
            chunks.append(text[start : start + self.chunk_size])
            start += self.chunk_size - self.chunk_overlap
        return chunks

    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        """算子执行入口"""
        start_time = time.time()
        
        self.read_file_first(sample)
         
        text = sample.get(getattr(self, 'text_key', 'text'), "")
        if not text: return sample

        # 1. 分块
        chunks = self._split_text(text)
        print(f"📄 文本分块完成：共 {len(chunks)} 个块")

        # 2. 准备 Prompt
        dim_str = "\n".join([f"- {d.get('name', 'Unknown')}: {d.get('desc', '')}" for d in self.dimensions])
        prompts = [self.prompt_template.replace("{context}", c).replace("{dimensions}", dim_str) for c in chunks]

        # 3. 并发调用
        weighted_scores = [0.0] * len(self.dimensions)
        total_chars = 0
        
        with ThreadPoolExecutor(max_workers=self.max_workers) as executor:
            future_to_idx = {executor.submit(self._call_llm_with_retry, p): i for i, p in enumerate(prompts)}
            for future in as_completed(future_to_idx):
                idx = future_to_idx[future]
                try:
                    scores = future.result()
                    if scores:
                        for i, s in enumerate(scores):
                            weighted_scores[i] += s * len(chunks[idx])
                        total_chars += len(chunks[idx])
                except Exception as e:
                    print(f"❌ 块 {idx} 失败: {e}")

        # 4. 聚合
        final_scores = {}
        if total_chars > 0:
            for i, dim in enumerate(self.dimensions):
                final_scores[dim.get('name', f'Dim_{i}')] = round(weighted_scores[i] / total_chars, 4)

        print(f"✅ 评估完成 | 耗时: {time.time()-start_time:.2f}s | 得分: {final_scores}")

        output_record = {
            "evaluation_scores": final_scores,
            "dimensions_count": len(self.dimensions),
            # 添加关联元数据
            "source_file_id": sample.get(self.fileid_key),      # 原始文件 ID
            "source_file_name": sample.get(self.filename_key),  # 原始文件名
            "source_dataset_id": sample.get("dataset_id")                     
        }

        sample[self.text_key] = json.dumps(output_record, ensure_ascii=False)
        sample["target_type"] = "jsonl"
        base_name = os.path.splitext(sample.get(self.filename_key, ""))[0]
        sample[self.filename_key] = f"{base_name}_eval"
        
        '''
        # 独立存储文件
        try:
            self.save_sample([output_record], sample)
            sample_key = getattr(self, 'text_key', 'text')
            sample[sample_key] = ""
        except Exception as e:
            print(f"Failed to save sample: {e}")
        '''
        return sample

    def _call_llm_with_retry(self, prompt: str, retries: int = 2) -> List[float]:
        retry_count = 0
        while retry_count <= retries:
            try:
                response = self.llm(prompt)
                return self._parse_scores(response)
            except Exception as e:
                print(f"⚠️ LLM 请求失败 ({retry_count+1}): {e}")
                retry_count += 1
                time.sleep(1)
        return None

    def _parse_scores(self, text: str) -> List[float]:
        if not text: return None
        matches = re.findall(self.pattern, text)
        return [float(m) for m in matches if 0.0 <= float(m) <= 1.0] if matches else None

# ==========================================================
# 🧪 本地测试入口 (实际调用 Ollama)
# ==========================================================
if __name__ == "__main__":
    print("🚀 启动 TextQualityEvaluator 测试...")

    # 配置参数 (请确保 Ollama 正在运行且已下载模型)
    mock_kwargs = {
        # --- Ollama 连接配置 ---
        # 注意：如果你在 Docker 容器外运行，使用 localhost；如果在容器内运行，可能需要 host.docker.internal
        "LLMUrl": "http://ollama.ailab.ccoop.cc:18585/v1/chat/completions", 
        
        "accessType": "true",       # 必须为 "true" 才能绕过 DeepData 内部证书校验
        "isHttps": "false",
        "LLMModel": "qwen2.5:7b-instruct-q4_0",     # <--- 请修改为你 Ollama 中实际存在的模型
        
        # 请求体模板
        "LLMBody": '{"model": "{model}", "messages": [{"role": "user", "content": "{prompt}"}]}',
        
        # 业务参数
        "dimensions": '[{"name": "连贯性", "desc": "逻辑通顺度"}, {"name": "信息密度", "desc": "有效信息占比"}]',
        "max_workers": 1
    }

    # 模拟输入样本
    sample_text = "DeepData 是一个强大的数据处理平台，支持多种文档格式的解析与处理。" * 20
    mock_sample = {"text": sample_text}

    try:
        # 实例化算子 (这步会触发 LLM 初始化)
        evaluator = TextQualityEvaluator(**mock_kwargs)
        
        # 执行
        evaluator.execute(mock_sample)
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()