local-coze 0.0.1__py3-none-any.whl

local_coze/s3/client.py

@@ -0,0 +1,580 @@
+"""
+S3 兼容存储模块
+提供对象存储的上传、下载、删除等功能
+"""
+
+import os
+import re
+import logging
+from pathlib import Path
+from typing import Optional, Any, Dict, List, Iterable
+from uuid import uuid4
+
+import boto3
+from botocore.exceptions import ClientError
+from boto3.s3.transfer import TransferConfig
+
+from .models import ListFilesResult
+
+logger = logging.getLogger(__name__)
+
+# 允许的文件名字符集（面向用户输入的约束）
+FILE_NAME_ALLOWED_RE = re.compile(r"^[A-Za-z0-9._\-/]+$")
+
+
+def _load_env() -> None:
+    """加载环境变量"""
+    try:
+        from dotenv import load_dotenv
+        load_dotenv()
+    except ImportError:
+        pass
+
+    try:
+        from coze_workload_identity import Client
+        client = Client()
+        env_vars = client.get_project_env_vars()
+        client.close()
+        for env_var in env_vars:
+            if env_var.key not in os.environ:
+                os.environ[env_var.key] = env_var.value
+    except Exception as e:
+        logger.debug(f"coze_workload_identity not available: {e}")
+
+
+class S3SyncStorage:
+    """S3 兼容存储实现"""
+
+    def __init__(
+        self,
+        *,
+        endpoint_url: Optional[str] = None,
+        access_key: str = "",
+        secret_key: str = "",
+        bucket_name: Optional[str] = None,
+        region: str = "cn-beijing",
+    ):
+        """
+        初始化 S3 存储客户端
+
+        Args:
+            endpoint_url: S3 兼容端点，默认从环境变量 COZE_BUCKET_ENDPOINT_URL 获取
+            access_key: 访问密钥，默认为空
+            secret_key: 秘密密钥，默认为空
+            bucket_name: 桶名称，默认从环境变量 COZE_BUCKET_NAME 获取
+            region: 区域，默认 cn-beijing
+        """
+        _load_env()
+        self.endpoint_url = os.environ.get("COZE_BUCKET_ENDPOINT_URL") or endpoint_url or ""
+        self.access_key = access_key
+        self.secret_key = secret_key
+        self.bucket_name = bucket_name or os.environ.get("COZE_BUCKET_NAME") or ""
+        self.region = region
+        self._client = None
+
+    def _get_client(self):
+        if self._client is None:
+            endpoint = self.endpoint_url
+            if not endpoint:
+                try:
+                    from coze_workload_identity import Client as CozeEnvClient
+                    coze_env_client = CozeEnvClient()
+                    env_vars = coze_env_client.get_project_env_vars()
+                    coze_env_client.close()
+                    for env_var in env_vars:
+                        if env_var.key == "COZE_BUCKET_ENDPOINT_URL":
+                            endpoint = env_var.value.replace("'", "'\\''")
+                            self.endpoint_url = endpoint
+                            break
+                except Exception as e:
+                    logger.error(f"Error loading COZE_BUCKET_ENDPOINT_URL: {e}")
+
+            if not endpoint:
+                logger.error("未配置存储端点：请设置 endpoint_url")
+                raise ValueError("未配置存储端点：请设置 endpoint_url")
+
+            client = boto3.client(
+                "s3",
+                endpoint_url=endpoint,
+                aws_access_key_id=self.access_key,
+                aws_secret_access_key=self.secret_key,
+                region_name=self.region,
+            )
+
+            # 注册 before-call 钩子，发送前注入 x-storage-token 头
+            def _inject_header(**kwargs):
+                try:
+                    from coze_workload_identity import Client as CozeClient
+                    coze_client = CozeClient()
+                    try:
+                        token = coze_client.get_access_token()
+                    except Exception as e:
+                        logger.error("Error loading COZE_WORKLOAD_IDENTITY_TOKEN: %s", e)
+                        token = None
+                        raise e
+                    finally:
+                        coze_client.close()
+                    params = kwargs.get("params", {})
+                    headers = params.setdefault("headers", {})
+                    headers["x-storage-token"] = token
+                except Exception as e:
+                    logger.error("Error loading COZE_WORKLOAD_IDENTITY_TOKEN: %s", e)
+
+            client.meta.events.register("before-call.s3", _inject_header)
+            self._client = client
+        return self._client
+
+    def _generate_object_key(self, *, original_name: str) -> str:
+        suffix = Path(original_name).suffix.lower()
+        stem = Path(original_name).stem
+        uniq = uuid4().hex[:8]
+        return f"{stem}_{uniq}{suffix}"
+
+    def _extract_logid(self, e: Exception) -> Optional[str]:
+        """从 ClientError 中提取 x-tt-logid"""
+        if isinstance(e, ClientError):
+            headers = (e.response or {}).get("ResponseMetadata", {}).get("HTTPHeaders", {})
+            return headers.get("x-tt-logid")
+        return None
+
+    def _error_msg(self, msg: str, e: Exception) -> str:
+        """构建带 logid 的错误信息"""
+        logid = self._extract_logid(e)
+        if logid:
+            return f"{msg}: {e} (x-tt-logid: {logid})"
+        return f"{msg}: {e}"
+
+    def _resolve_bucket(self, bucket: Optional[str]) -> str:
+        """统一解析 bucket 来源，确保得到有效桶名"""
+        target_bucket = bucket or os.environ.get("COZE_BUCKET_NAME") or self.bucket_name
+        if not target_bucket:
+            raise ValueError("未配置 bucket：请传入 bucket 或设置 COZE_BUCKET_NAME，或在实例化时提供 bucket_name")
+        return target_bucket
+
+    def _validate_file_name(self, name: str) -> None:
+        """校验 S3 对象命名：长度≤1024；允许 [A-Za-z0-9._-/]；不以 / 起止且不含 //"""
+        msg = (
+            "file name invalid: 文件名需满足以下 S3 对象命名规范："
+            "1) 长度 1–1024 字节；"
+            "2) 仅允许字母、数字、点(.)、下划线(_)、短横(-)、目录分隔符(/)；"
+            "3) 不允许空格或以下特殊字符：? # & % { } ^ [ ] ` \\ < > ~ | \" ' + = : ;；"
+            "4) 不以 / 开头或结尾，且不包含连续的 //；"
+            "示例：report_2025-12-11.pdf、images/photo-01.png。"
+        )
+
+        if not name or not name.strip():
+            raise ValueError(msg + "（原因：文件名为空）")
+
+        if len(name.encode("utf-8")) > 1024:
+            raise ValueError(msg + "（原因：长度超过 1024 字节）")
+
+        if name.startswith("/") or name.endswith("/"):
+            raise ValueError(msg + "（原因：以 / 开头或结尾）")
+        if "//" in name:
+            raise ValueError(msg + "（原因：包含连续的 //）")
+
+        if not FILE_NAME_ALLOWED_RE.match(name):
+            bad = re.findall(r"[^A-Za-z0-9._\-/]", name)
+            example = bad[0] if bad else "非法字符"
+            raise ValueError(msg + f"（原因：包含非法字符，例如：{example}）")
+
+    def upload_file(
+        self,
+        *,
+        file_content: bytes,
+        file_name: str,
+        content_type: str = "application/octet-stream",
+        bucket: Optional[str] = None,
+    ) -> str:
+        """
+        上传文件
+
+        Args:
+            file_content: 文件内容
+            file_name: 原始文件名
+            content_type: MIME 类型
+            bucket: 目标桶（可选）
+
+        Returns:
+            str: 对象 key
+        """
+        self._validate_file_name(file_name)
+        try:
+            client = self._get_client()
+            object_key = self._generate_object_key(original_name=file_name)
+            target_bucket = self._resolve_bucket(bucket)
+            client.put_object(Bucket=target_bucket, Key=object_key, Body=file_content, ContentType=content_type)
+            return object_key
+        except Exception as e:
+            logger.error(self._error_msg("Error uploading file to S3", e))
+            raise e
+
+    def delete_file(self, *, file_key: str, bucket: Optional[str] = None) -> bool:
+        """
+        删除文件
+
+        Args:
+            file_key: 对象 key
+            bucket: 目标桶（可选）
+
+        Returns:
+            bool: 是否成功
+        """
+        try:
+            client = self._get_client()
+            target_bucket = self._resolve_bucket(bucket)
+            client.delete_object(Bucket=target_bucket, Key=file_key)
+            return True
+        except Exception as e:
+            logger.error(self._error_msg("Error deleting file from S3", e))
+            raise e
+
+    def file_exists(self, *, file_key: str, bucket: Optional[str] = None) -> bool:
+        """
+        检查文件是否存在
+
+        Args:
+            file_key: 对象 key
+            bucket: 目标桶（可选）
+
+        Returns:
+            bool: 是否存在
+        """
+        try:
+            client = self._get_client()
+            target_bucket = self._resolve_bucket(bucket)
+            client.head_object(Bucket=target_bucket, Key=file_key)
+            return True
+        except ClientError as e:
+            code = (e.response or {}).get("Error", {}).get("Code", "")
+            if code in {"404", "NoSuchKey", "NotFound"}:
+                return False
+            logger.error(self._error_msg("Error checking file existence in S3", e))
+            return False
+        except Exception as e:
+            logger.error(self._error_msg("Error checking file existence in S3", e))
+            return False
+
+    def read_file(self, *, file_key: str, bucket: Optional[str] = None) -> bytes:
+        """
+        读取文件
+
+        Args:
+            file_key: 对象 key
+            bucket: 目标桶（可选）
+
+        Returns:
+            bytes: 文件内容
+        """
+        try:
+            client = self._get_client()
+            target_bucket = self._resolve_bucket(bucket)
+            resp = client.get_object(Bucket=target_bucket, Key=file_key)
+            body = resp.get("Body")
+            if body is None:
+                raise RuntimeError("S3 get_object returned no Body")
+            try:
+                return body.read()
+            finally:
+                try:
+                    body.close()
+                except Exception as ce:
+                    logger.debug("Failed to close S3 response body: %s", ce)
+        except Exception as e:
+            logger.error(self._error_msg("Error reading file from S3", e))
+            raise e
+
+    def list_files(
+        self,
+        *,
+        prefix: Optional[str] = None,
+        bucket: Optional[str] = None,
+        max_keys: int = 1000,
+        continuation_token: Optional[str] = None,
+    ) -> ListFilesResult:
+        """
+        列出对象
+
+        Args:
+            prefix: 前缀过滤
+            bucket: 目标桶（可选）
+            max_keys: 最大返回数量（1-1000）
+            continuation_token: 分页 token
+
+        Returns:
+            ListFilesResult: 包含 keys, is_truncated, next_continuation_token
+        """
+        try:
+            client = self._get_client()
+            target_bucket = self._resolve_bucket(bucket)
+            if max_keys <= 0 or max_keys > 1000:
+                raise ValueError("max_keys 必须在 1 到 1000 之间")
+
+            kwargs: Dict[str, Any] = {
+                "Bucket": target_bucket,
+                "MaxKeys": max_keys,
+                "Prefix": prefix,
+                "ContinuationToken": continuation_token,
+            }
+            kwargs = {k: v for k, v in kwargs.items() if v is not None}
+
+            resp = client.list_objects_v2(**kwargs)
+            contents = resp.get("Contents", []) or []
+            keys: List[str] = [item.get("Key") for item in contents if isinstance(item, dict) and item.get("Key")]
+            return {
+                "keys": keys,
+                "is_truncated": bool(resp.get("IsTruncated")),
+                "next_continuation_token": resp.get("NextContinuationToken"),
+            }
+        except ClientError as e:
+            code = (e.response or {}).get("Error", {}).get("Code", "")
+            logger.error(self._error_msg(f"Error listing files in S3 (code={code})", e))
+            raise e
+        except Exception as e:
+            logger.error(self._error_msg("Error listing files in S3", e))
+            raise e
+
+    def generate_presigned_url(
+        self,
+        *,
+        key: str,
+        bucket: Optional[str] = None,
+        expire_time: int = 86400,
+    ) -> str:
+        """
+        生成签名 URL
+
+        Args:
+            key: 对象 key
+            bucket: 目标桶（可选）
+            expire_time: 过期时间（秒），默认 86400
+
+        Returns:
+            str: 签名 URL
+        """
+        import json
+        import urllib.request as urllib_request
+
+        try:
+            from coze_workload_identity import Client as CozeClient
+            coze_client = CozeClient()
+            try:
+                token = coze_client.get_access_token()
+            finally:
+                try:
+                    coze_client.close()
+                except Exception:
+                    pass
+        except Exception as e:
+            logger.error(f"Error loading x-storage-token: {e}")
+            raise RuntimeError(f"获取 x-storage-token 失败: {e}")
+
+        try:
+            sign_base = os.environ.get("COZE_BUCKET_ENDPOINT_URL") or self.endpoint_url
+            if not sign_base:
+                raise ValueError("未配置签名端点：请设置 COZE_BUCKET_ENDPOINT_URL 或传入 endpoint_url")
+            sign_url_endpoint = sign_base.rstrip("/") + "/sign-url"
+
+            headers = {
+                "Content-Type": "application/json",
+                "x-storage-token": token,
+            }
+
+            target_bucket = self._resolve_bucket(bucket)
+            payload = {"bucket_name": target_bucket, "path": key, "expire_time": expire_time}
+            data = json.dumps(payload).encode("utf-8")
+            request = urllib_request.Request(sign_url_endpoint, data=data, headers=headers, method="POST")
+        except Exception as e:
+            logger.error(f"Error creating request for sign-url: {e}")
+            raise RuntimeError(f"创建 sign-url 请求失败: {e}")
+
+        try:
+            with urllib_request.urlopen(request) as resp:
+                resp_bytes = resp.read()
+                content_type = resp.headers.get("Content-Type", "")
+                text = resp_bytes.decode("utf-8", errors="replace")
+                if "application/json" in content_type or text.strip().startswith("{"):
+                    try:
+                        obj = json.loads(text)
+                    except Exception:
+                        return text
+                    data = obj.get("data")
+                    if isinstance(data, dict) and "url" in data:
+                        return data["url"]
+                    url_value = obj.get("url") or obj.get("signed_url") or obj.get("presigned_url")
+                    if url_value:
+                        return url_value
+                    raise ValueError("签名服务返回缺少 data.url/url 字段")
+                return text
+        except Exception as e:
+            raise RuntimeError(f"生成签名URL失败: {e}")
+
+    def stream_upload_file(
+        self,
+        *,
+        fileobj,
+        file_name: str,
+        content_type: str = "application/octet-stream",
+        bucket: Optional[str] = None,
+        multipart_chunksize: int = 5 * 1024 * 1024,
+        multipart_threshold: int = 5 * 1024 * 1024,
+        max_concurrency: int = 1,
+        use_threads: bool = False,
+    ) -> str:
+        """
+        流式上传文件对象
+
+        Args:
+            fileobj: 文件对象（需有 read() 方法）
+            file_name: 原始文件名
+            content_type: MIME 类型
+            bucket: 目标桶（可选）
+            multipart_chunksize: 分片大小（默认 5MB）
+            multipart_threshold: 触发分片上传的阈值（默认 5MB）
+            max_concurrency: 并发数（默认 1）
+            use_threads: 是否启用多线程（默认 False）
+
+        Returns:
+            str: 对象 key
+        """
+        try:
+            client = self._get_client()
+            target_bucket = self._resolve_bucket(bucket)
+            key = self._generate_object_key(original_name=file_name)
+
+            extra_args = {"ContentType": content_type} if content_type else {}
+            config = TransferConfig(
+                multipart_chunksize=multipart_chunksize,
+                multipart_threshold=multipart_threshold,
+                max_concurrency=max_concurrency,
+                use_threads=use_threads,
+            )
+            client.upload_fileobj(Fileobj=fileobj, Bucket=target_bucket, Key=key, ExtraArgs=extra_args, Config=config)
+            return key
+        except Exception as e:
+            logger.error(self._error_msg("Error streaming upload (fileobj) to S3", e))
+            raise e
+
+    def upload_from_url(
+        self,
+        *,
+        url: str,
+        bucket: Optional[str] = None,
+        timeout: int = 30,
+    ) -> str:
+        """
+        从 URL 下载并上传到 S3
+
+        Args:
+            url: 源文件 URL
+            bucket: 目标桶（可选）
+            timeout: HTTP 超时（秒），默认 30
+
+        Returns:
+            str: 对象 key
+        """
+        import urllib.request as urllib_request
+        from urllib.parse import urlparse, unquote
+
+        try:
+            request = urllib_request.Request(url)
+            with urllib_request.urlopen(request, timeout=timeout) as resp:
+                parsed = urlparse(url)
+                file_name = Path(unquote(parsed.path)).name or "file"
+                content_type = resp.headers.get("Content-Type", "application/octet-stream")
+                return self.stream_upload_file(
+                    fileobj=resp,
+                    file_name=file_name,
+                    content_type=content_type,
+                    bucket=bucket,
+                )
+        except Exception as e:
+            logger.error(self._error_msg("Error uploading from URL to S3", e))
+            raise e
+
+    def trunk_upload_file(
+        self,
+        *,
+        chunk_iter: Iterable[bytes],
+        file_name: str,
+        content_type: str = "application/octet-stream",
+        bucket: Optional[str] = None,
+        part_size: int = 5 * 1024 * 1024,
+    ) -> str:
+        """
+        分块流式上传
+
+        Args:
+            chunk_iter: 字节块迭代器
+            file_name: 原始文件名
+            content_type: MIME 类型
+            bucket: 目标桶（可选）
+            part_size: 分片大小（默认 5MB）
+
+        Returns:
+            str: 对象 key
+        """
+        client = self._get_client()
+        target_bucket = self._resolve_bucket(bucket)
+        key = self._generate_object_key(original_name=file_name)
+
+        # 初始化分片上传
+        try:
+            init_resp = client.create_multipart_upload(Bucket=target_bucket, Key=key, ContentType=content_type)
+            upload_id = init_resp["UploadId"]
+        except Exception as e:
+            logger.error(self._error_msg("create_multipart_upload failed", e))
+            raise e
+
+        parts = []
+        part_number = 1
+        buffer = bytearray()
+        try:
+            for chunk in chunk_iter:
+                if not chunk:
+                    continue
+                buffer.extend(chunk)
+                while len(buffer) >= part_size:
+                    data = bytes(buffer[:part_size])
+                    buffer = buffer[part_size:]
+                    resp = client.upload_part(
+                        Bucket=target_bucket,
+                        Key=key,
+                        UploadId=upload_id,
+                        PartNumber=part_number,
+                        Body=data,
+                    )
+                    parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
+                    part_number += 1
+
+            # 上传最后不足 part_size 的余量
+            if len(buffer) > 0:
+                resp = client.upload_part(
+                    Bucket=target_bucket,
+                    Key=key,
+                    UploadId=upload_id,
+                    PartNumber=part_number,
+                    Body=bytes(buffer),
+                )
+                parts.append({"PartNumber": part_number, "ETag": resp["ETag"]})
+
+            # 完成分片
+            client.complete_multipart_upload(
+                Bucket=target_bucket,
+                Key=key,
+                UploadId=upload_id,
+                MultipartUpload={"Parts": parts},
+            )
+            return key
+        except Exception as e:
+            logger.error(self._error_msg("multipart upload failed", e))
+            try:
+                client.abort_multipart_upload(Bucket=target_bucket, Key=key, UploadId=upload_id)
+            except Exception as ae:
+                logger.error(self._error_msg("abort_multipart_upload failed", ae))
+            raise e
+
+
+__all__ = [
+    "S3SyncStorage",
+]

local_coze/s3/models.py

@@ -0,0 +1,18 @@
+"""
+S3 模块数据模型
+"""
+
+from typing import Optional, List, TypedDict
+
+
+class ListFilesResult(TypedDict):
+    """list_files 的返回结构类型"""
+
+    keys: List[str]
+    is_truncated: bool
+    next_continuation_token: Optional[str]
+
+
+__all__ = [
+    "ListFilesResult",
+]

local_coze/search/__init__.py

@@ -0,0 +1,19 @@
+from .client import SearchClient
+from .models import (
+    SearchRequest,
+    SearchResponse,
+    SearchFilter,
+    WebItem,
+    ImageItem,
+    ImageInfo
+)
+
+__all__ = [
+    "SearchClient",
+    "SearchRequest",
+    "SearchResponse",
+    "SearchFilter",
+    "WebItem",
+    "ImageItem",
+    "ImageInfo"
+]