xparse-client 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

xparse_client/pipeline/destinations.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import json
+import logging
+import uuid
+import boto3
+
+from abc import ABC, abstractmethod
+from datetime import datetime
+from pathlib import Path
+from typing import List, Dict, Any
+
+from botocore.config import Config
+from pymilvus import MilvusClient
+
+
+logger = logging.getLogger(__name__)
+
+
+class Destination(ABC):
+    """数据目的地抽象基类"""
+
+    @abstractmethod
+    def write(self, data: List[Dict[str, Any]], metadata: Dict[str, Any]) -> bool:
+        """写入数据"""
+        raise NotImplementedError
+
+
+class MilvusDestination(Destination):
+    """Milvus/Zilliz 向量数据库目的地"""
+
+    def __init__(self, db_path: str, collection_name: str, dimension: int, api_key: str = None, token: str = None):
+        from pymilvus import DataType
+
+        self.db_path = db_path
+        self.collection_name = collection_name
+        self.dimension = dimension
+
+        client_kwargs = {'uri': db_path}
+        if api_key:
+            client_kwargs['token'] = api_key
+        elif token:
+            client_kwargs['token'] = token
+
+        self.client = MilvusClient(**client_kwargs)
+
+        if not self.client.has_collection(collection_name):
+            schema = self.client.create_schema(
+                auto_id=False,
+                enable_dynamic_field=True
+            )
+
+            schema.add_field(field_name="element_id", datatype=DataType.VARCHAR, max_length=128, is_primary=True)
+            schema.add_field(field_name="embeddings", datatype=DataType.FLOAT_VECTOR, dim=dimension)
+            schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=65535)
+            schema.add_field(field_name="record_id", datatype=DataType.VARCHAR, max_length=128)
+            schema.add_field(field_name="metadata", datatype=DataType.JSON)
+
+            index_params = self.client.prepare_index_params()
+            index_params.add_index(
+                field_name="embeddings",
+                index_type="AUTOINDEX",
+                metric_type="COSINE"
+            )
+
+            self.client.create_collection(
+                collection_name=collection_name,
+                schema=schema,
+                index_params=index_params
+            )
+            print(f"✓ Milvus/Zilliz 集合创建: {collection_name} (自定义 Schema)")
+        else:
+            print(f"✓ Milvus/Zilliz 集合存在: {collection_name}")
+
+        logger.info(f"Milvus/Zilliz 连接成功: {db_path}")
+
+    def write(self, data: List[Dict[str, Any]], metadata: Dict[str, Any]) -> bool:
+        try:
+            insert_data = []
+            for item in data:
+                metadata = item.get('metadata', {})
+                if 'embeddings' in item and item['embeddings']:
+                    element_id = item.get('element_id') or item.get('id') or str(uuid.uuid4())
+                    insert_data.append({
+                        'embeddings': item['embeddings'],
+                        'text': item.get('text', ''),
+                        'element_id': element_id,
+                        'record_id': metadata.get('record_id', ''),
+                        'metadata': metadata,
+                        'created_at': datetime.now().isoformat()
+                    })
+
+            if not insert_data:
+                print(f"  ! 警告: 没有有效的向量数据")
+                return False
+
+            self.client.insert(
+                collection_name=self.collection_name,
+                data=insert_data
+            )
+            print(f"  ✓ 写入 Milvus: {len(insert_data)} 条")
+            logger.info(f"写入 Milvus 成功: {len(insert_data)} 条")
+            return True
+        except Exception as e:
+            print(f"  ✗ 写入 Milvus 失败: {str(e)}")
+            logger.error(f"写入 Milvus 失败: {str(e)}")
+            return False
+
+
+class LocalDestination(Destination):
+    """本地文件系统目的地"""
+
+    def __init__(self, output_dir: str):
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"✓ 输出目录: {self.output_dir}")
+        logger.info(f"输出目录: {self.output_dir}")
+
+    def write(self, data: List[Dict[str, Any]], metadata: Dict[str, Any]) -> bool:
+        try:
+            file_name = metadata.get('file_name', 'output')
+            base_name = Path(file_name).stem
+            stage = metadata.get('stage')  # 用于区分中间结果的阶段
+            
+            # 如果是中间结果，在文件名中添加阶段标识
+            if stage:
+                output_file = self.output_dir / f"{base_name}_{stage}.json"
+            else:
+                output_file = self.output_dir / f"{base_name}.json"
+
+            with open(output_file, 'w', encoding='utf-8') as f:
+                json.dump(data, f, ensure_ascii=False, indent=2)
+
+            print(f"  ✓ 写入本地: {output_file}")
+            logger.info(f"写入本地成功: {output_file}")
+            return True
+        except Exception as e:
+            print(f"  ✗ 写入本地失败: {str(e)}")
+            logger.error(f"写入本地失败: {str(e)}")
+            return False
+
+
+class S3Destination(Destination):
+    """S3/MinIO 数据目的地"""
+
+    def __init__(self, endpoint: str, access_key: str, secret_key: str,
+                 bucket: str, prefix: str = '', region: str = 'us-east-1'):
+        self.endpoint = endpoint
+        self.bucket = bucket
+        self.prefix = prefix.strip('/') if prefix else ''
+
+        if self.endpoint == 'https://textin-minio-api.ai.intsig.net':
+            config = Config(signature_version='s3v4')
+        elif self.endpoint.endswith('aliyuncs.com'):
+            config = Config(signature_version='s3', s3={'addressing_style': 'virtual'})
+        else:
+            config = Config(signature_version='s3v4', s3={'addressing_style': 'virtual'})
+
+        self.client = boto3.client(
+            's3',
+            endpoint_url=endpoint,
+            aws_access_key_id=access_key,
+            aws_secret_access_key=secret_key,
+            region_name=region,
+            config=config
+        )
+
+        try:
+            self.client.head_bucket(Bucket=bucket)
+            test_key = f"{self.prefix}/empty.tmp" if self.prefix else f"empty.tmp"
+            self.client.put_object(
+                Bucket=bucket,
+                Key=test_key,
+                Body=b''
+            )
+            try:
+                self.client.delete_object(Bucket=bucket, Key=test_key)
+            except Exception:
+                pass
+
+            print(f"✓ S3 连接成功且可写: {endpoint}/{bucket}")
+            logger.info(f"S3 连接成功且可写: {endpoint}/{bucket}")
+        except Exception as e:
+            print(f"✗ S3 连接或写入测试失败: {str(e)}")
+            logger.error(f"S3 连接或写入测试失败: {str(e)}")
+            raise
+
+    def write(self, data: List[Dict[str, Any]], metadata: Dict[str, Any]) -> bool:
+        try:
+            file_name = metadata.get('file_name', 'output')
+            base_name = Path(file_name).stem
+            object_key = f"{self.prefix}/{base_name}.json" if self.prefix else f"{base_name}.json"
+
+            json_data = json.dumps(data, ensure_ascii=False, indent=2)
+            json_bytes = json_data.encode('utf-8')
+
+            self.client.put_object(
+                Bucket=self.bucket,
+                Key=object_key,
+                Body=json_bytes,
+                ContentType='application/json'
+            )
+
+            print(f"  ✓ 写入 S3: {self.endpoint}/{self.bucket}/{object_key}")
+            logger.info(f"写入 S3 成功: {self.endpoint}/{self.bucket}/{object_key}")
+            return True
+        except Exception as e:
+            print(f"  ✗ 写入 S3 失败: {str(e)}")
+            logger.error(f"写入 S3 失败: {str(e)}")
+            return False
+
+
+__all__ = [
+    'Destination',
+    'MilvusDestination',
+    'LocalDestination',
+    'S3Destination',
+]
+

xparse_client/pipeline/pipeline.py

@@ -0,0 +1,440 @@
+#!/usr/bin/env python
+# -*- encoding: utf-8 -*-
+
+import json
+import logging
+import time
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Dict, Any, Optional, Tuple, List, Union
+
+import requests
+
+from .config import ParseConfig, ChunkConfig, EmbedConfig, Stage, PipelineStats, PipelineConfig
+from .sources import Source, S3Source, LocalSource, FtpSource, SmbSource
+from .destinations import Destination, MilvusDestination, LocalDestination, S3Destination
+
+
+logger = logging.getLogger(__name__)
+
+
+class Pipeline:
+    """数据处理 Pipeline"""
+
+    def __init__(
+        self,
+        source: Source,
+        destination: Destination,
+        api_base_url: str = 'http://localhost:8000/api/xparse',
+        api_headers: Optional[Dict[str, str]] = None,
+        stages: Optional[List[Stage]] = None,
+        pipeline_config: Optional[PipelineConfig] = None,
+        intermediate_results_destination: Optional[Destination] = None
+    ):
+        self.source = source
+        self.destination = destination
+        self.api_base_url = api_base_url.rstrip('/')
+        self.api_headers = api_headers or {}
+        self.pipeline_config = pipeline_config or PipelineConfig()
+
+        # 处理 intermediate_results_destination 参数
+        # 如果直接传入了 intermediate_results_destination，优先使用它并自动启用中间结果保存
+        if intermediate_results_destination is not None:
+            self.pipeline_config.include_intermediate_results = True
+            self.pipeline_config.intermediate_results_destination = intermediate_results_destination
+        # 如果 pipeline_config 中已设置，使用 pipeline_config 中的值
+        elif self.pipeline_config.include_intermediate_results:
+            if not self.pipeline_config.intermediate_results_destination:
+                raise ValueError("当 include_intermediate_results 为 True 时，必须设置 intermediate_results_destination")
+
+        # 处理 stages 配置
+        if stages is None:
+            raise ValueError("必须提供 stages 参数")
+        
+        self.stages = stages
+
+        # 验证 stages
+        if not self.stages or self.stages[0].type != 'parse':
+            raise ValueError("stages 必须包含且第一个必须是 'parse' 类型")
+        
+        # 验证 embed config（如果存在）
+        for stage in self.stages:
+            if stage.type == 'embed' and isinstance(stage.config, EmbedConfig):
+                stage.config.validate()
+        
+        # 验证 intermediate_results_destination
+        if self.pipeline_config.include_intermediate_results:
+            # 验证是否为支持的 Destination 类型
+            from .destinations import Destination
+            if not isinstance(self.pipeline_config.intermediate_results_destination, Destination):
+                raise ValueError(f"intermediate_results_destination 必须是 Destination 类型")
+            self.intermediate_results_destination = self.pipeline_config.intermediate_results_destination
+
+        print("=" * 60)
+        print("Pipeline 初始化完成")
+        print(f"  Stages: {[s.type for s in self.stages]}")
+        for stage in self.stages:
+            print(f"    - {stage.type}: {stage.config}")
+        if self.pipeline_config.include_intermediate_results:
+            print(f"  Pipeline Config: 中间结果保存已启用")
+        print("=" * 60)
+
+    def _call_pipeline_api(self, file_bytes: bytes, file_name: str, data_source: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        url = f"{self.api_base_url}/pipeline"
+        max_retries = 3
+
+        for try_count in range(max_retries):
+            try:
+                files = {'file': (file_name or 'file', file_bytes)}
+                form_data = {}
+
+                # 将 stages 转换为 API 格式
+                stages_data = [stage.to_dict() for stage in self.stages]
+                form_data['stages'] = json.dumps(stages_data)
+                form_data['data_source'] = json.dumps(data_source, ensure_ascii=False)
+                
+                # 如果启用了中间结果保存，在请求中添加参数
+                if self.pipeline_config:
+                    form_data['config'] = json.dumps(self.pipeline_config.to_dict(), ensure_ascii=False)
+
+                response = requests.post(
+                    url,
+                    files=files,
+                    data=form_data,
+                    headers=self.api_headers,
+                    timeout=120
+                )
+
+                if response.status_code == 200:
+                    result = response.json()
+                    print(f"  ✓ Pipeline 接口返回 x_request_id: {result.get('x_request_id')}")
+                    if result.get('code') == 200 and 'data' in result:
+                        return result.get('data')
+                    return None
+                else:
+                    print(f"  ! API 错误 {response.status_code}, 重试 {try_count + 1}/{max_retries}")
+                    logger.warning(f"API 错误 {response.status_code}: pipeline")
+
+            except Exception as e:
+                print(f"  ! 请求异常: {str(e)}, 重试 {try_count + 1}/{max_retries}")
+                logger.error(f"API 请求异常 pipeline: {str(e)}")
+
+            if try_count < max_retries - 1:
+                time.sleep(2)
+
+        return None
+
+    def process_with_pipeline(self, file_bytes: bytes, file_name: str, data_source: Dict[str, Any]) -> Optional[Tuple[List[Dict[str, Any]], PipelineStats]]:
+        print(f"  → 调用 Pipeline 接口: {file_name}")
+        result = self._call_pipeline_api(file_bytes, file_name, data_source)
+
+        if result and 'elements' in result and 'stats' in result:
+            elements = result['elements']
+            stats_data = result['stats']
+
+            stats = PipelineStats(
+                original_elements=stats_data.get('original_elements', 0),
+                chunked_elements=stats_data.get('chunked_elements', 0),
+                embedded_elements=stats_data.get('embedded_elements', 0),
+                stages=self.stages  # 使用实际执行的 stages
+            )
+
+            # 如果启用了中间结果保存，处理中间结果
+            if self.pipeline_config.include_intermediate_results and 'intermediate_results' in result:
+                self._save_intermediate_results(result['intermediate_results'], file_name, data_source)
+
+            print(f"  ✓ Pipeline 完成:")
+            print(f"    - 原始元素: {stats.original_elements}")
+            print(f"    - 分块后: {stats.chunked_elements}")
+            print(f"    - 向量化: {stats.embedded_elements}")
+            logger.info(f"Pipeline 完成: {file_name}, {stats.embedded_elements} 个向量")
+
+            return elements, stats
+        else:
+            print(f"  ✗ Pipeline 失败")
+            logger.error(f"Pipeline 失败: {file_name}")
+            return None
+
+    def _save_intermediate_results(self, intermediate_results: List[Dict[str, Any]], file_name: str, data_source: Dict[str, Any]) -> None:
+        """保存中间结果
+        
+        Args:
+            intermediate_results: 中间结果数组，每个元素包含 stage 和 elements 字段
+            file_name: 文件名
+            data_source: 数据源信息
+        """
+        try:
+            # intermediate_results 是一个数组，每个元素是 {stage: str, elements: List}
+            for result_item in intermediate_results:
+                if 'stage' not in result_item or 'elements' not in result_item:
+                    logger.warning(f"中间结果项缺少 stage 或 elements 字段: {result_item}")
+                    continue
+                
+                stage = result_item['stage']
+                elements = result_item['elements']
+                
+                metadata = {
+                    'file_name': file_name,
+                    'stage': stage,
+                    'total_elements': len(elements),
+                    'processed_at': datetime.now().isoformat(),
+                    'data_source': data_source
+                }
+                
+                self.pipeline_config.intermediate_results_destination.write(elements, metadata)
+                print(f"  ✓ 保存 {stage.upper()} 中间结果: {len(elements)} 个元素")
+                logger.info(f"保存 {stage.upper()} 中间结果成功: {file_name}")
+
+        except Exception as e:
+            print(f"  ✗ 保存中间结果失败: {str(e)}")
+            logger.error(f"保存中间结果失败: {file_name}, {str(e)}")
+
+    def process_file(self, file_path: str) -> bool:
+        print(f"\n{'=' * 60}")
+        print(f"处理文件: {file_path}")
+        logger.info(f"开始处理文件: {file_path}")
+
+        try:
+            print(f"  → 读取文件...")
+            file_bytes, data_source = self.source.read_file(file_path)
+            data_source = data_source or {}
+            data_source['date_processed'] = datetime.now(timezone.utc).timestamp()
+            print(f"  ✓ 文件读取完成: {len(file_bytes)} bytes")
+
+            result = self.process_with_pipeline(file_bytes, file_path, data_source)
+            if not result:
+                return False
+
+            embedded_data, stats = result
+
+            print(f"  → 写入目的地...")
+            metadata = {
+                'file_name': file_path,
+                'total_elements': len(embedded_data),
+                'processed_at': datetime.now().isoformat(),
+                'data_source': data_source,
+                'stats': {
+                    'original_elements': stats.original_elements,
+                    'chunked_elements': stats.chunked_elements,
+                    'embedded_elements': stats.embedded_elements
+                }
+            }
+
+            success = self.destination.write(embedded_data, metadata)
+
+            if success:
+                print(f"\n✓✓✓ 文件处理成功: {file_path}")
+                logger.info(f"文件处理成功: {file_path}")
+            else:
+                print(f"\n✗✗✗ 文件处理失败: {file_path}")
+                logger.error(f"文件处理失败: {file_path}")
+
+            return success
+
+        except Exception as e:
+            print(f"\n✗✗✗ 处理异常: {str(e)}")
+            logger.error(f"处理文件异常 {file_path}: {str(e)}")
+            return False
+
+    def run(self):
+        start_time = time.time()
+
+        print("\n" + "=" * 60)
+        print("开始执行 Pipeline")
+        print("=" * 60)
+        logger.info("=" * 60)
+        logger.info("开始执行 Pipeline")
+
+        print("\n→ 列出文件...")
+        files = self.source.list_files()
+
+        if not files:
+            print("\n✗ 没有找到文件")
+            logger.info("没有找到文件")
+            return
+
+        total = len(files)
+        success_count = 0
+        fail_count = 0
+
+        for idx, file_path in enumerate(files, 1):
+            print(f"\n进度: [{idx}/{total}]")
+            logger.info(f"进度: [{idx}/{total}] - {file_path}")
+
+            try:
+                if self.process_file(file_path):
+                    success_count += 1
+                else:
+                    fail_count += 1
+            except Exception as e:
+                print(f"\n✗✗✗ 文件处理异常: {str(e)}")
+                logger.error(f"文件处理异常 {file_path}: {str(e)}")
+                fail_count += 1
+
+            if idx < total:
+                time.sleep(1)
+
+        elapsed = time.time() - start_time
+        print("\n" + "=" * 60)
+        print("Pipeline 执行完成！")
+        print("=" * 60)
+        print(f"总文件数: {total}")
+        print(f"成功: {success_count}")
+        print(f"失败: {fail_count}")
+        print(f"总耗时: {elapsed:.2f} 秒")
+        print("=" * 60)
+
+        logger.info("=" * 60)
+        logger.info(f"Pipeline 完成 - 总数:{total}, 成功:{success_count}, 失败:{fail_count}, 耗时:{elapsed:.2f}秒")
+        logger.info("=" * 60)
+
+
+def create_pipeline_from_config(config: Dict[str, Any]) -> Pipeline:
+    source_config = config['source']
+    if source_config['type'] == 's3':
+        source = S3Source(
+            endpoint=source_config['endpoint'],
+            access_key=source_config['access_key'],
+            secret_key=source_config['secret_key'],
+            bucket=source_config['bucket'],
+            prefix=source_config.get('prefix', ''),
+            region=source_config.get('region', 'us-east-1'),
+            pattern=source_config.get('pattern', '*')
+        )
+    elif source_config['type'] == 'local':
+        source = LocalSource(
+            directory=source_config['directory'],
+            pattern=source_config.get('pattern', '*')
+        )
+    elif source_config['type'] == 'ftp':
+        source = FtpSource(
+            host=source_config['host'],
+            port=source_config['port'],
+            username=source_config['username'],
+            password=source_config['password'],
+            pattern=source_config.get('pattern', '*')
+        )
+    elif source_config['type'] == 'smb':
+        source = SmbSource(
+            host=source_config['host'],
+            share_name=source_config['share_name'],
+            username=source_config['username'],
+            password=source_config['password'],
+            domain=source_config.get('domain', ''),
+            port=source_config.get('port', 445),
+            path=source_config.get('path', ''),
+            pattern=source_config.get('pattern', '*')
+        )
+    else:
+        raise ValueError(f"未知的 source 类型: {source_config['type']}")
+
+    dest_config = config['destination']
+    if dest_config['type'] in ['milvus', 'zilliz']:
+        destination = MilvusDestination(
+            db_path=dest_config['db_path'],
+            collection_name=dest_config['collection_name'],
+            dimension=dest_config['dimension'],
+            api_key=dest_config.get('api_key'),
+            token=dest_config.get('token')
+        )
+    elif dest_config['type'] == 'local':
+        destination = LocalDestination(
+            output_dir=dest_config['output_dir']
+        )
+    elif dest_config['type'] == 's3':
+        destination = S3Destination(
+            endpoint=dest_config['endpoint'],
+            access_key=dest_config['access_key'],
+            secret_key=dest_config['secret_key'],
+            bucket=dest_config['bucket'],
+            prefix=dest_config.get('prefix', ''),
+            region=dest_config.get('region', 'us-east-1')
+        )
+    else:
+        raise ValueError(f"未知的 destination 类型: {dest_config['type']}")
+
+    # 处理 stages 配置
+    if 'stages' not in config or not config['stages']:
+        raise ValueError("配置中必须包含 'stages' 字段")
+    
+    stages = []
+    for stage_cfg in config['stages']:
+        stage_type = stage_cfg.get('type')
+        stage_config_dict = stage_cfg.get('config', {})
+        
+        if stage_type == 'parse':
+            parse_cfg_copy = dict(stage_config_dict)
+            provider = parse_cfg_copy.pop('provider', 'textin')
+            stage_config = ParseConfig(provider=provider, **parse_cfg_copy)
+        elif stage_type == 'chunk':
+            stage_config = ChunkConfig(
+                strategy=stage_config_dict.get('strategy', 'basic'),
+                include_orig_elements=stage_config_dict.get('include_orig_elements', False),
+                new_after_n_chars=stage_config_dict.get('new_after_n_chars', 512),
+                max_characters=stage_config_dict.get('max_characters', 1024),
+                overlap=stage_config_dict.get('overlap', 0),
+                overlap_all=stage_config_dict.get('overlap_all', False)
+            )
+        elif stage_type == 'embed':
+            stage_config = EmbedConfig(
+                provider=stage_config_dict.get('provider', 'qwen'),
+                model_name=stage_config_dict.get('model_name', 'text-embedding-v3')
+            )
+        else:
+            raise ValueError(f"未知的 stage 类型: {stage_type}")
+        
+        stages.append(Stage(type=stage_type, config=stage_config))
+
+    # 创建 Pipeline 配置
+    pipeline_config = None
+    if 'pipeline_config' in config and config['pipeline_config']:
+        pipeline_cfg = config['pipeline_config']
+        include_intermediate_results = pipeline_cfg.get('include_intermediate_results', False)
+        intermediate_results_destination = None
+        
+        if include_intermediate_results:
+            if 'intermediate_results_destination' in pipeline_cfg:
+                dest_cfg = pipeline_cfg['intermediate_results_destination']
+                dest_type = dest_cfg.get('type')
+                
+                if dest_type == 'local':
+                    intermediate_results_destination = LocalDestination(
+                        output_dir=dest_cfg['output_dir']
+                    )
+                elif dest_type == 's3':
+                    intermediate_results_destination = S3Destination(
+                        endpoint=dest_cfg['endpoint'],
+                        access_key=dest_cfg['access_key'],
+                        secret_key=dest_cfg['secret_key'],
+                        bucket=dest_cfg['bucket'],
+                        prefix=dest_cfg.get('prefix', ''),
+                        region=dest_cfg.get('region', 'us-east-1')
+                    )
+                else:
+                    raise ValueError(f"不支持的 intermediate_results_destination 类型: '{dest_type}'，支持的类型: 'local', 's3'")
+            else:
+                raise ValueError("当 include_intermediate_results 为 True 时，必须设置 intermediate_results_destination")
+        
+        pipeline_config = PipelineConfig(
+            include_intermediate_results=include_intermediate_results,
+            intermediate_results_destination=intermediate_results_destination
+        )
+
+    # 创建 Pipeline
+    pipeline = Pipeline(
+        source=source,
+        destination=destination,
+        api_base_url=config.get('api_base_url', 'http://localhost:8000/api/xparse'),
+        api_headers=config.get('api_headers', {}),
+        stages=stages,
+        pipeline_config=pipeline_config
+    )
+
+    return pipeline
+
+
+__all__ = [
+    'Pipeline',
+    'create_pipeline_from_config',
+]
+