xparse-client 0.2.20__py3-none-any.whl → 0.3.0b1__py3-none-any.whl

xparse_client/connectors/destinations/qdrant.py

@@ -0,0 +1,238 @@
+"""Qdrant 向量数据库目的地（懒加载 qdrant-client）"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from typing import Any
+
+from ...exceptions import DestinationError
+from .base import Destination, VectorDestinationMixin
+
+logger = logging.getLogger(__name__)
+
+
+def _get_qdrant():
+    """懒加载 qdrant-client"""
+    try:
+        from qdrant_client import QdrantClient
+        from qdrant_client.models import (
+            Distance,
+            PayloadSchemaType,
+            PointStruct,
+            VectorParams,
+        )
+
+        return QdrantClient, Distance, PayloadSchemaType, PointStruct, VectorParams
+    except ImportError as e:
+        raise ImportError(
+            "使用 QdrantDestination 需要安装 qdrant-client: pip install xparse-client[qdrant]"
+        ) from e
+
+
+class QdrantDestination(Destination, VectorDestinationMixin):
+    """Qdrant 向量数据库目的地
+
+    支持 Qdrant 本地部署和 Qdrant Cloud。
+
+    Attributes:
+        url: Qdrant 服务地址
+        collection_name: Collection 名称
+        dimension: 向量维度
+
+    Example:
+        >>> # 本地 Qdrant
+        >>> dest = QdrantDestination(
+        ...     url="http://localhost:6333",
+        ...     collection_name="documents",
+        ...     dimension=1024,
+        ... )
+        >>>
+        >>> # Qdrant Cloud
+        >>> dest = QdrantDestination(
+        ...     url="https://xxx.qdrant.io",
+        ...     collection_name="documents",
+        ...     dimension=1024,
+        ...     api_key="your-api-key",
+        ... )
+    """
+
+    def __init__(
+        self,
+        url: str,
+        collection_name: str,
+        dimension: int,
+        api_key: str | None = None,
+        prefer_grpc: bool = False,
+    ) -> None:
+        """初始化 Qdrant 目的地
+
+        Args:
+            url: Qdrant 服务地址
+            collection_name: Collection 名称
+            dimension: 向量维度
+            api_key: API Key（用于 Qdrant Cloud）
+            prefer_grpc: 是否优先使用 gRPC，默认 False
+
+        Raises:
+            DestinationError: 连接失败
+        """
+        (
+            QdrantClient,
+            Distance,
+            PayloadSchemaType,
+            PointStruct,
+            VectorParams,
+        ) = _get_qdrant()
+
+        self.url = url
+        self.collection_name = collection_name
+        self.dimension = dimension
+        self._PointStruct = PointStruct
+
+        client_kwargs = {"url": url}
+        if api_key:
+            client_kwargs["api_key"] = api_key
+        if prefer_grpc:
+            client_kwargs["prefer_grpc"] = True
+
+        try:
+            self.client = QdrantClient(**client_kwargs)
+
+            # 检查或创建 Collection
+            collections = self.client.get_collections()
+            collection_exists = any(
+                col.name == collection_name for col in collections.collections
+            )
+
+            if not collection_exists:
+                self.client.create_collection(
+                    collection_name=collection_name,
+                    vectors_config=VectorParams(size=dimension, distance=Distance.COSINE),
+                )
+                # 为 record_id 创建索引
+                try:
+                    self.client.create_payload_index(
+                        collection_name=collection_name,
+                        field_name="record_id",
+                        field_schema=PayloadSchemaType.KEYWORD,
+                    )
+                except Exception as e:
+                    logger.warning(f"创建 record_id 索引失败: {e}")
+
+                logger.info(f"Qdrant Collection 创建: {collection_name}")
+            else:
+                logger.info(f"Qdrant Collection 已存在: {collection_name}")
+                # 确保索引存在
+                try:
+                    self.client.create_payload_index(
+                        collection_name=collection_name,
+                        field_name="record_id",
+                        field_schema=PayloadSchemaType.KEYWORD,
+                    )
+                except Exception:
+                    pass  # 索引可能已存在
+
+        except ImportError:
+            raise
+        except Exception as e:
+            raise DestinationError(
+                f"Qdrant 连接失败: {e}",
+                connector_type="qdrant",
+                operation="connect",
+                details={"url": url, "collection_name": collection_name},
+            ) from e
+
+    def write(self, data: list[dict[str, Any]], metadata: dict[str, Any]) -> bool:
+        """写入向量数据到 Qdrant
+
+        Args:
+            data: 包含 embeddings 的元素列表
+            metadata: 元数据，必须包含 record_id
+
+        Returns:
+            是否写入成功
+
+        Raises:
+            DestinationError: 写入失败
+        """
+        record_id = metadata.get("record_id")
+        if not record_id:
+            logger.warning("没有 record_id，跳过写入")
+            return False
+
+        try:
+            # 删除旧记录
+            try:
+                scroll_result = self.client.scroll(
+                    collection_name=self.collection_name,
+                    scroll_filter={
+                        "must": [{"key": "record_id", "match": {"value": record_id}}]
+                    },
+                    limit=10000,
+                )
+
+                if scroll_result[0]:
+                    point_ids = [point.id for point in scroll_result[0]]
+                    self.client.delete(
+                        collection_name=self.collection_name,
+                        points_selector=point_ids,
+                    )
+                    logger.info(
+                        f"删除 Qdrant 旧记录: record_id={record_id}, 数量={len(point_ids)}"
+                    )
+            except Exception as e:
+                logger.warning(f"删除旧记录失败: {e}")
+
+            # 准备插入数据
+            fixed_fields = {"embeddings", "text", "element_id", "record_id", "metadata"}
+            points = []
+
+            for item in data:
+                if "embeddings" not in item or not item["embeddings"]:
+                    continue
+
+                element_id = (
+                    item.get("element_id") or item.get("id") or str(uuid.uuid4())
+                )
+
+                # Qdrant 的 point id 需要是 UUID 或整数
+                try:
+                    point_id = str(uuid.UUID(element_id))
+                except (ValueError, TypeError):
+                    point_id = str(uuid.uuid5(uuid.NAMESPACE_URL, str(element_id)))
+
+                # 准备 payload
+                payload = self.prepare_payload(item, metadata, record_id, fixed_fields)
+
+                point = self._PointStruct(
+                    id=point_id,
+                    vector=item["embeddings"],
+                    payload=payload,
+                )
+                points.append(point)
+
+            if not points:
+                logger.warning("没有有效的向量数据")
+                return False
+
+            # 批量插入
+            self.client.upsert(
+                collection_name=self.collection_name,
+                points=points,
+            )
+            logger.info(f"写入 Qdrant: {len(points)} 条")
+            return True
+
+        except Exception as e:
+            raise DestinationError(
+                f"写入 Qdrant 失败: {e}",
+                connector_type="qdrant",
+                operation="write",
+            ) from e
+
+    def __repr__(self) -> str:
+        return f"<QdrantDestination url={self.url} collection={self.collection_name}>"
+
+
+__all__ = ["QdrantDestination"]

xparse_client/connectors/destinations/s3.py

@@ -0,0 +1,163 @@
+"""S3/MinIO 目的地（懒加载 boto3）"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import Any
+
+from ...exceptions import DestinationError
+from .base import Destination
+
+logger = logging.getLogger(__name__)
+
+
+def _get_boto3():
+    """懒加载 boto3"""
+    try:
+        import boto3
+        from botocore.config import Config
+
+        return boto3, Config
+    except ImportError as e:
+        raise ImportError(
+            "使用 S3Destination 需要安装 boto3: pip install xparse-client[s3]"
+        ) from e
+
+
+class S3Destination(Destination):
+    """S3/MinIO 目的地
+
+    将处理结果写入 S3 或兼容存储。
+
+    Attributes:
+        endpoint: S3 端点
+        bucket: 存储桶
+        prefix: 对象前缀
+
+    Example:
+        >>> dest = S3Destination(
+        ...     endpoint="https://s3.amazonaws.com",
+        ...     access_key="...",
+        ...     secret_key="...",
+        ...     bucket="my-bucket",
+        ...     prefix="output/",
+        ... )
+        >>> dest.write(elements, {"filename": "doc.pdf"})
+    """
+
+    def __init__(
+        self,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+        bucket: str,
+        prefix: str = "",
+        region: str = "us-east-1",
+    ) -> None:
+        """初始化 S3 目的地
+
+        Args:
+            endpoint: S3 端点 URL
+            access_key: 访问密钥
+            secret_key: 秘密密钥
+            bucket: 存储桶名称
+            prefix: 对象前缀
+            region: 区域
+
+        Raises:
+            DestinationError: 连接失败
+        """
+        boto3, Config = _get_boto3()
+
+        self.endpoint = endpoint
+        self.bucket = bucket
+        self.prefix = prefix.strip("/") if prefix else ""
+
+        # 配置签名版本
+        if endpoint == "https://textin-minio-api.ai.intsig.net":
+            config = Config(signature_version="s3v4")
+        elif endpoint.endswith("aliyuncs.com"):
+            config = Config(signature_version="s3", s3={"addressing_style": "virtual"})
+        elif endpoint.endswith("myhuaweicloud.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 "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
+
+            logger.info(f"S3 目的地连接成功: {endpoint}/{bucket}")
+
+        except Exception as e:
+            raise DestinationError(
+                f"S3 连接或写入测试失败: {e}",
+                connector_type="s3",
+                operation="connect",
+                details={"endpoint": endpoint, "bucket": bucket},
+            ) from e
+
+    def write(self, data: list[dict[str, Any]], metadata: dict[str, Any]) -> bool:
+        """写入数据到 S3
+
+        Args:
+            data: 要写入的数据列表
+            metadata: 元数据，包含 filename
+
+        Returns:
+            是否写入成功
+
+        Raises:
+            DestinationError: 写入失败
+        """
+        try:
+            filename = metadata.get("filename", "output")
+            base_name = Path(filename).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",
+            )
+
+            logger.info(f"写入 S3: {self.endpoint}/{self.bucket}/{object_key}")
+            return True
+
+        except Exception as e:
+            raise DestinationError(
+                f"写入 S3 失败: {e}",
+                connector_type="s3",
+                operation="write",
+            ) from e
+
+    def __repr__(self) -> str:
+        return f"<S3Destination endpoint={self.endpoint} bucket={self.bucket}>"
+
+
+__all__ = ["S3Destination"]

xparse_client/connectors/sources/__init__.py

@@ -0,0 +1,45 @@
+"""数据源连接器
+
+提供从各种数据源读取文件的能力。
+
+可用的数据源：
+- LocalSource: 本地文件系统
+- S3Source: S3/MinIO 对象存储（需要 pip install xparse-client[s3]）
+- FtpSource: FTP 服务器
+- SmbSource: SMB/CIFS 共享（需要 pip install xparse-client[smb]）
+
+Example:
+    >>> from xparse_client.connectors.sources import LocalSource
+    >>> source = LocalSource(directory="./docs", pattern=["*.pdf"])
+    >>> files = source.list_files()
+    >>> content, metadata = source.read_file(files[0])
+"""
+
+from .base import Source
+from .local import LocalSource
+
+# 懒加载其他 Source，避免强依赖
+__all__ = [
+    "Source",
+    "LocalSource",
+    "S3Source",
+    "FtpSource",
+    "SmbSource",
+]
+
+
+def __getattr__(name: str):
+    """懒加载 Source 类"""
+    if name == "S3Source":
+        from .s3 import S3Source
+
+        return S3Source
+    elif name == "FtpSource":
+        from .ftp import FtpSource
+
+        return FtpSource
+    elif name == "SmbSource":
+        from .smb import SmbSource
+
+        return SmbSource
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

xparse_client/connectors/sources/base.py

@@ -0,0 +1,74 @@
+"""数据源抽象基类"""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+class Source(ABC):
+    """数据源抽象基类
+
+    所有数据源必须实现 list_files() 和 read_file() 方法。
+
+    Example:
+        >>> class MySource(Source):
+        ...     def list_files(self) -> List[str]:
+        ...         return ["file1.pdf", "file2.pdf"]
+        ...
+        ...     def read_file(self, file_path: str) -> Tuple[bytes, Dict[str, Any]]:
+        ...         return b"content", {"url": file_path}
+    """
+
+    @abstractmethod
+    def list_files(self) -> list[str]:
+        """列出所有匹配的文件
+
+        Returns:
+            文件路径列表
+
+        Raises:
+            SourceError: 列出文件失败
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def read_file(self, file_path: str) -> tuple[bytes, dict[str, Any]]:
+        """读取文件内容
+
+        Args:
+            file_path: 文件路径（相对路径）
+
+        Returns:
+            (file_bytes, data_source) 元组
+            - file_bytes: 文件二进制内容
+            - data_source: 数据来源元信息，包含 url, version, date_created 等
+
+        Raises:
+            SourceError: 读取文件失败
+        """
+        raise NotImplementedError
+
+    def close(self) -> None:  # noqa: B027
+        """关闭连接
+
+        子类可以重写此方法来释放资源（如网络连接）。
+        """
+        pass
+
+    def __enter__(self) -> Source:
+        """上下文管理器入口"""
+        return self
+
+    def __exit__(self, *args: Any) -> None:
+        """上下文管理器退出，自动关闭连接"""
+        self.close()
+
+    def __repr__(self) -> str:
+        return f"<{self.__class__.__name__}>"
+
+
+__all__ = ["Source"]