ima-python-sdk 0.1.0__py3-none-any.whl

ima_sdk/client.py

@@ -0,0 +1,268 @@
+"""IMA OpenAPI 核心客户端。
+
+提供 HTTP 请求封装、认证管理和统一错误处理。
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import shlex
+import ssl
+import sys
+import time
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Any, Callable, Dict, Generator, List, Optional
+
+from . import logger
+
+
+class ImaApiError(Exception):
+    """IMA API 调用异常。"""
+
+    def __init__(self, retcode: int, errmsg: str, data: Any = None):
+        self.retcode = retcode
+        self.errmsg = errmsg
+        self.data = data
+        super().__init__(f"[{retcode}] {errmsg}")
+
+
+class ImaClient:
+    """IMA OpenAPI 核心客户端。
+
+    凭证来源优先级：
+    1. 构造参数
+    2. 环境变量 IMA_OPENAPI_CLIENTID / IMA_OPENAPI_APIKEY
+    3. 配置文件 ~/.config/ima/client_id / ~/.config/ima/api_key
+
+    用法::
+
+        client = ImaClient()
+        data = client.post("openapi/wiki/v1/get_knowledge_base", {"ids": ["xxx"]})
+    """
+
+    BASE_URL = "https://ima.qq.com"
+
+    def __init__(
+        self,
+        client_id: Optional[str] = None,
+        api_key: Optional[str] = None,
+        base_url: Optional[str] = None,
+        *,
+        curl_mode: bool = False,
+    ):
+        self.base_url = (base_url or self.BASE_URL).rstrip("/")
+        self.curl_mode = curl_mode
+        self._client_id = client_id or self._resolve_credential("client_id")
+        self._api_key = api_key or self._resolve_credential("api_key")
+
+        if not self._client_id or not self._api_key:
+            raise ValueError(
+                "缺少 IMA 凭证。请通过以下方式之一提供：\n"
+                "  1. ImaClient(client_id=..., api_key=...)\n"
+                "  2. 环境变量 IMA_OPENAPI_CLIENTID / IMA_OPENAPI_APIKEY\n"
+                "  3. 配置文件 ~/.config/ima/client_id / ~/.config/ima/api_key"
+            )
+
+    @staticmethod
+    def _resolve_credential(name: str) -> Optional[str]:
+        """按优先级解析凭证：环境变量 > 配置文件。"""
+        # 环境变量
+        env_key = f"IMA_OPENAPI_{name.upper()}"
+        # client_id -> IMA_OPENAPI_CLIENTID, api_key -> IMA_OPENAPI_APIKEY
+        env_key = env_key.replace("CLIENT_ID", "CLIENTID").replace("API_KEY", "APIKEY")
+        value = os.environ.get(env_key)
+        if value:
+            return value.strip()
+
+        # 配置文件
+        config_path = Path.home() / ".config" / "ima" / name
+        if config_path.is_file():
+            return config_path.read_text(encoding="utf-8").strip()
+
+        return None
+
+    def _build_headers(self) -> Dict[str, str]:
+        return {
+            "ima-openapi-clientid": self._client_id,
+            "ima-openapi-apikey": self._api_key,
+            "Content-Type": "application/json",
+        }
+
+    def _print_curl(self, method: str, url: str, headers: Dict[str, str], payload: bytes) -> None:
+        """输出等价的 curl 命令到 stderr。"""
+        parts = ["curl", "-X", method, shlex.quote(url)]
+        for k, v in headers.items():
+            parts.append(f"-H {shlex.quote(f'{k}: {v}')}")
+        if payload:
+            parts.append(f"-d {shlex.quote(payload.decode('utf-8'))}")
+        print(" \\\n  ".join(parts), file=sys.stderr)
+
+    # 连接级别重试配置（应对 SSL 握手瞬断）
+    _CONNECT_RETRIES = 3
+    _CONNECT_RETRY_DELAY = 1  # 秒
+
+    def post(self, path: str, body: Optional[dict] = None) -> dict:
+        """发送 POST 请求并返回响应 JSON。
+
+        API 直接返回业务数据 JSON，通过 HTTP 状态码判断成败。
+        对 SSL 握手瞬断自动重试。
+
+        Args:
+            path: API 路径，例如 ``openapi/wiki/v1/get_knowledge_base``
+            body: 请求体字典
+
+        Returns:
+            API 返回的 JSON 字典
+
+        Raises:
+            ImaApiError: HTTP 错误或响应体包含错误信息
+            urllib.error.URLError: 网络异常
+        """
+        url = f"{self.base_url}/{path.lstrip('/')}"
+        req_body = body or {}
+        payload = json.dumps(req_body, ensure_ascii=False).encode("utf-8")
+        logger.log_request("POST", url, req_body)
+        if self.curl_mode:
+            self._print_curl("POST", url, self._build_headers(), payload)
+
+        raw: str = ""
+        status: int = 0
+        for attempt in range(self._CONNECT_RETRIES):
+            req = urllib.request.Request(
+                url,
+                data=payload,
+                headers=self._build_headers(),
+                method="POST",
+            )
+            try:
+                with urllib.request.urlopen(req) as resp:
+                    raw = resp.read().decode("utf-8")
+                    status = resp.status
+                break
+            except urllib.error.HTTPError as e:
+                raw = e.read().decode("utf-8", errors="replace")
+                logger.log_error(f"HTTP {e.code}: {raw}")
+                raise ImaApiError(
+                    retcode=e.code,
+                    errmsg=raw,
+                ) from e
+            except urllib.error.URLError as e:
+                if isinstance(e.reason, ssl.SSLEOFError) and attempt < self._CONNECT_RETRIES - 1:
+                    logger.log(f"  SSL 连接中断，{self._CONNECT_RETRY_DELAY}s 后重试 ({attempt + 1}/{self._CONNECT_RETRIES})")
+                    time.sleep(self._CONNECT_RETRY_DELAY)
+                    continue
+                raise
+
+        result = json.loads(raw)
+
+        # 兼容三种响应格式：
+        # 1. 有 retcode 包装: {"retcode": 0, "errmsg": "ok", "data": {...}}
+        # 2. 有 code 包装:    {"code": 0, "msg": "success", "data": {...}}
+        # 3. 直接返回数据:    {"addable_knowledge_base_list": [...], "is_end": true}
+        if "retcode" in result:
+            retcode = result["retcode"]
+            if retcode != 0:
+                logger.log_error(f"API 错误 [{retcode}] {result.get('errmsg', '')}: {raw}")
+                raise ImaApiError(
+                    retcode=retcode,
+                    errmsg=result.get("errmsg", "未知错误"),
+                    data=result.get("data"),
+                )
+            logger.log_response(status, f"retcode=0, data keys={list(result.get('data', {}).keys())}")
+            return result.get("data", {})
+
+        if "code" in result:
+            code = result["code"]
+            if code != 0:
+                logger.log_error(f"API 错误 [{code}] {result.get('msg', '')}: {raw}")
+                raise ImaApiError(
+                    retcode=code,
+                    errmsg=result.get("msg", "未知错误"),
+                    data=result.get("data"),
+                )
+            logger.log_response(status, f"code=0, data keys={list(result.get('data', {}).keys())}")
+            return result.get("data", {})
+
+        logger.log_response(status, f"keys={list(result.keys())}")
+        return result
+
+    def post_raw(self, path: str, body: Optional[dict] = None) -> dict:
+        """发送 POST 请求并返回完整响应（不捕获 HTTP 错误）。"""
+        return self.post(path, body)
+
+    def paginate(
+        self,
+        path: str,
+        body: dict,
+        items_key: str,
+        *,
+        cursor_field: str = "cursor",
+        next_cursor_field: str = "next_cursor",
+        is_end_field: str = "is_end",
+        limit: int = 50,
+        limit_field: str = "limit",
+        initial_cursor: str = "",
+    ) -> Generator[dict, None, None]:
+        """游标翻页迭代器，逐页 yield data 字典。
+
+        自动处理游标翻页逻辑，每次 yield 一页的 data 字典。
+
+        Args:
+            path: API 路径
+            body: 请求体基础参数（不包含 cursor/limit）
+            items_key: data 中包含列表数据的字段名
+            cursor_field: 请求体中游标参数名
+            next_cursor_field: 返回中下一页游标字段名
+            is_end_field: 返回中是否结束字段名
+            limit: 每页数量
+            limit_field: 请求体中 limit 参数名
+            initial_cursor: 初始游标值
+        """
+        cursor = initial_cursor
+        page = 0
+        while True:
+            page += 1
+            req_body = {**body, cursor_field: cursor, limit_field: limit}
+            logger.log_step(f"翻页请求 (第 {page} 页)", f"cursor={cursor!r}")
+            data = self.post(path, req_body)
+            items_count = len(data.get(items_key, []))
+            logger.log(f"  本页获取 {items_count} 条 {items_key}")
+            yield data
+
+            if data.get(is_end_field, True):
+                logger.log(f"  已到最后一页 (共 {page} 页)")
+                break
+            next_cursor = data.get(next_cursor_field, "")
+            if not next_cursor:
+                logger.log(f"  无下一页游标，结束翻页 (共 {page} 页)")
+                break
+            cursor = next_cursor
+
+    def paginate_items(
+        self,
+        path: str,
+        body: dict,
+        items_key: str,
+        item_factory: Callable[[dict], Any],
+        **kwargs: Any,
+    ) -> List[Any]:
+        """翻页并收集所有条目到列表。
+
+        Args:
+            path: API 路径
+            body: 请求体基础参数
+            items_key: data 中包含列表数据的字段名
+            item_factory: 将单个条目字典转换为对象的工厂函数
+            **kwargs: 传递给 paginate 的额外参数
+
+        Returns:
+            所有页面合并后的条目列表
+        """
+        results: List[Any] = []
+        for data in self.paginate(path, body, items_key, **kwargs):
+            for item in data.get(items_key, []):
+                results.append(item_factory(item))
+        return results

ima_sdk/cos_uploader.py

@@ -0,0 +1,158 @@
+"""腾讯云 COS 文件上传模块。
+
+纯 Python 实现 COS PUT Object 操作，使用 HMAC-SHA1 签名。
+参考: https://cloud.tencent.com/document/product/436/7778
+"""
+
+from __future__ import annotations
+
+import hashlib
+import hmac
+import math
+import time
+import urllib.request
+from typing import Optional
+
+from .types import CosCredential
+from . import logger
+
+
+def _hmac_sha1(key: str, data: str) -> str:
+    """HMAC-SHA1 签名，返回十六进制字符串。"""
+    return hmac.new(
+        key.encode("utf-8"), data.encode("utf-8"), hashlib.sha1
+    ).hexdigest()
+
+
+def _sha1(data: str) -> str:
+    """SHA1 哈希，返回十六进制字符串。"""
+    return hashlib.sha1(data.encode("utf-8")).hexdigest()
+
+
+def _build_authorization(
+    secret_id: str,
+    secret_key: str,
+    method: str,
+    pathname: str,
+    headers: dict[str, str],
+    start_time: int,
+    expired_time: int,
+) -> str:
+    """构建 COS Authorization header。
+
+    算法步骤：
+    1. SignKey = HMAC-SHA1(SecretKey, KeyTime)
+    2. HttpString = method\\npathname\\nparams\\nheaders\\n
+    3. StringToSign = sha1\\nKeyTime\\nSHA1(HttpString)\\n
+    4. Signature = HMAC-SHA1(SignKey, StringToSign)
+    """
+    key_time = f"{start_time};{expired_time}"
+
+    # 1. SignKey
+    sign_key = _hmac_sha1(secret_key, key_time)
+
+    # 2. HttpString — 对 PUT 请求，无 query params
+    header_keys = sorted(headers.keys())
+    http_headers = "&".join(
+        f"{k.lower()}={urllib.request.quote(headers[k], safe='')}"
+        for k in header_keys
+    )
+    http_string = f"{method.lower()}\n{pathname}\n\n{http_headers}\n"
+
+    # 3. StringToSign
+    string_to_sign = f"sha1\n{key_time}\n{_sha1(http_string)}\n"
+
+    # 4. Signature
+    signature = _hmac_sha1(sign_key, string_to_sign)
+
+    # 5. 组装 Authorization
+    header_list = ";".join(k.lower() for k in header_keys)
+    return (
+        f"q-sign-algorithm=sha1"
+        f"&q-ak={secret_id}"
+        f"&q-sign-time={key_time}"
+        f"&q-key-time={key_time}"
+        f"&q-header-list={header_list}"
+        f"&q-url-param-list="
+        f"&q-signature={signature}"
+    )
+
+
+def cos_upload(
+    file_path: str,
+    credential: CosCredential,
+    content_type: str = "application/octet-stream",
+    start_time: Optional[int] = None,
+    expired_time: Optional[int] = None,
+) -> None:
+    """上传文件到腾讯云 COS。
+
+    Args:
+        file_path: 本地文件路径
+        credential: COS 上传凭证（由 create_media 返回）
+        content_type: 文件 MIME 类型
+        start_time: 签名开始时间（秒级时间戳），默认当前时间
+        expired_time: 签名过期时间，默认当前时间 + 1小时
+
+    Raises:
+        RuntimeError: 上传失败
+        FileNotFoundError: 文件不存在
+    """
+    with open(file_path, "rb") as f:
+        file_content = f.read()
+
+    logger.log_step("COS 上传", f"file={file_path} size={len(file_content)} bytes")
+
+    now = math.floor(time.time())
+    if start_time is None:
+        start_time = credential.start_time or now
+    if expired_time is None:
+        expired_time = credential.expired_time or (now + 3600)
+
+    hostname = f"{credential.bucket_name}.cos.{credential.region}.myqcloud.com"
+    pathname = f"/{credential.cos_key}"
+
+    # 参与签名的 headers
+    sign_headers = {
+        "content-length": str(len(file_content)),
+        "host": hostname,
+    }
+
+    authorization = _build_authorization(
+        secret_id=credential.secret_id,
+        secret_key=credential.secret_key,
+        method="PUT",
+        pathname=pathname,
+        headers=sign_headers,
+        start_time=start_time,
+        expired_time=expired_time,
+    )
+
+    url = f"https://{hostname}{pathname}"
+    logger.log(f"  PUT {url}")
+    logger.log(f"  content_type={content_type}")
+    req = urllib.request.Request(
+        url,
+        data=file_content,
+        headers={
+            "Content-Type": content_type,
+            "Content-Length": str(len(file_content)),
+            "Authorization": authorization,
+            "x-cos-security-token": credential.token,
+        },
+        method="PUT",
+    )
+
+    try:
+        with urllib.request.urlopen(req) as resp:
+            if resp.status < 200 or resp.status >= 300:
+                body = resp.read().decode("utf-8", errors="replace")
+                logger.log_error(f"COS 上传失败 (HTTP {resp.status}): {body}")
+                raise RuntimeError(
+                    f"COS 上传失败 (HTTP {resp.status}): {body}"
+                )
+            logger.log(f"  上传成功 (HTTP {resp.status})")
+    except urllib.error.HTTPError as e:
+        body = e.read().decode("utf-8", errors="replace")
+        logger.log_error(f"COS 上传失败 (HTTP {e.code}): {body}")
+        raise RuntimeError(f"COS 上传失败 (HTTP {e.code}): {body}") from e

ima_sdk/file_checker.py

@@ -0,0 +1,226 @@
+"""文件预检模块。
+
+在上传文件到知识库前，校验文件类型和大小是否符合要求。
+完整移植自 preflight-check.cjs 的逻辑。
+"""
+
+from __future__ import annotations
+
+import os
+from typing import Optional
+
+from .types import FileCheckResult, MediaType
+from . import logger
+
+# ── 扩展名 → media_type + content_type 映射 ──────────────────────────────────
+
+EXT_MAP: dict[str, dict] = {
+    "pdf": {"media_type": MediaType.PDF, "content_type": "application/pdf"},
+    "doc": {"media_type": MediaType.WORD, "content_type": "application/msword"},
+    "docx": {
+        "media_type": MediaType.WORD,
+        "content_type": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+    },
+    "ppt": {
+        "media_type": MediaType.PPT,
+        "content_type": "application/vnd.ms-powerpoint",
+    },
+    "pptx": {
+        "media_type": MediaType.PPT,
+        "content_type": "application/vnd.openxmlformats-officedocument.presentationml.presentation",
+    },
+    "xls": {
+        "media_type": MediaType.EXCEL,
+        "content_type": "application/vnd.ms-excel",
+    },
+    "xlsx": {
+        "media_type": MediaType.EXCEL,
+        "content_type": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+    },
+    "csv": {"media_type": MediaType.EXCEL, "content_type": "text/csv"},
+    "md": {"media_type": MediaType.MARKDOWN, "content_type": "text/markdown"},
+    "markdown": {"media_type": MediaType.MARKDOWN, "content_type": "text/markdown"},
+    "png": {"media_type": MediaType.IMAGE, "content_type": "image/png"},
+    "jpg": {"media_type": MediaType.IMAGE, "content_type": "image/jpeg"},
+    "jpeg": {"media_type": MediaType.IMAGE, "content_type": "image/jpeg"},
+    "webp": {"media_type": MediaType.IMAGE, "content_type": "image/webp"},
+    "txt": {"media_type": MediaType.TXT, "content_type": "text/plain"},
+    "xmind": {"media_type": MediaType.XMIND, "content_type": "application/x-xmind"},
+    "mp3": {"media_type": MediaType.AUDIO, "content_type": "audio/mpeg"},
+    "m4a": {"media_type": MediaType.AUDIO, "content_type": "audio/x-m4a"},
+    "wav": {"media_type": MediaType.AUDIO, "content_type": "audio/wav"},
+    "aac": {"media_type": MediaType.AUDIO, "content_type": "audio/aac"},
+}
+
+# ── content_type → media_type 反向映射 ────────────────────────────────────────
+
+CONTENT_TYPE_MAP: dict[str, int] = {}
+for _info in EXT_MAP.values():
+    ct = _info["content_type"]
+    if ct not in CONTENT_TYPE_MAP:
+        CONTENT_TYPE_MAP[ct] = _info["media_type"]
+
+# 额外 content_type 别名
+CONTENT_TYPE_MAP.update(
+    {
+        "text/x-markdown": MediaType.MARKDOWN,
+        "application/md": MediaType.MARKDOWN,
+        "application/markdown": MediaType.MARKDOWN,
+        "application/vnd.xmind.workbook": MediaType.XMIND,
+        "application/zip": MediaType.XMIND,  # xmind 可能为 zip
+    }
+)
+
+# ── 文件大小限制（字节）────────────────────────────────────────────────────────
+
+MB = 1024 * 1024
+
+SIZE_LIMITS: dict[int, int] = {
+    MediaType.EXCEL: 10 * MB,
+    MediaType.MARKDOWN: 10 * MB,
+    MediaType.TXT: 10 * MB,
+    MediaType.XMIND: 10 * MB,
+    MediaType.IMAGE: 30 * MB,
+}
+DEFAULT_SIZE_LIMIT = 200 * MB
+
+# ── 不支持的视频类型 ─────────────────────────────────────────────────────────
+
+UNSUPPORTED_VIDEO_EXTS = {
+    "mp4", "avi", "mov", "mkv", "wmv", "flv", "webm", "m4v", "rmvb", "rm", "3gp"
+}
+
+UNSUPPORTED_VIDEO_CONTENT_TYPES = {
+    "video/mp4",
+    "video/x-msvideo",
+    "video/quicktime",
+    "video/x-matroska",
+    "video/x-ms-wmv",
+    "video/x-flv",
+    "video/webm",
+}
+
+
+def _format_size(size_bytes: int) -> str:
+    """格式化文件大小。"""
+    if size_bytes < MB:
+        return f"{size_bytes / 1024:.1f} KB"
+    return f"{size_bytes / MB:.1f} MB"
+
+
+def check_file(
+    file_path: str, content_type: Optional[str] = None
+) -> FileCheckResult:
+    """对上传文件进行预检。
+
+    检查文件是否存在、是否为支持的类型、是否超过大小限制。
+
+    解析优先级：
+    1. 若 content_type 可识别 → 使用 content_type
+    2. 若 content_type 不可识别 → 回退到扩展名
+    3. 若无 content_type → 使用扩展名
+    4. 两者都无法识别 → 检查失败
+
+    Args:
+        file_path: 文件绝对路径
+        content_type: 可选的 MIME 类型
+
+    Returns:
+        FileCheckResult 对象，passed=True 表示通过预检
+    """
+    abs_path = os.path.abspath(file_path)
+    file_name = os.path.basename(abs_path)
+    logger.log_step("文件预检", f"{file_name}")
+
+    # 提取扩展名
+    _, ext_with_dot = os.path.splitext(file_name)
+    ext = ext_with_dot.lstrip(".").lower() if ext_with_dot else ""
+
+    base_result = FileCheckResult(
+        file_path=abs_path, file_name=file_name, file_ext=ext
+    )
+
+    # 1. 检查文件存在
+    if not os.path.isfile(abs_path):
+        base_result.reason = f"文件不存在: {abs_path}"
+        logger.log_error(base_result.reason)
+        return base_result
+
+    # 2. 检查不支持的视频类型
+    if ext in UNSUPPORTED_VIDEO_EXTS:
+        base_result.reason = (
+            f"不支持视频文件 (.{ext})，视频仅支持在 IMA 桌面端添加。"
+        )
+        logger.log_error(base_result.reason)
+        return base_result
+
+    if content_type and content_type in UNSUPPORTED_VIDEO_CONTENT_TYPES:
+        base_result.reason = (
+            f"不支持视频文件 ({content_type})，视频仅支持在 IMA 桌面端添加。"
+        )
+        logger.log_error(base_result.reason)
+        return base_result
+
+    # 3. 解析 media_type 和 content_type
+    media_type: Optional[int] = None
+    resolved_ct: Optional[str] = None
+
+    ct_media_type = CONTENT_TYPE_MAP.get(content_type) if content_type else None
+    ext_mapping = EXT_MAP.get(ext) if ext else None
+
+    if ct_media_type is not None:
+        # content_type 可识别
+        media_type = ct_media_type
+        resolved_ct = content_type
+    elif content_type:
+        # content_type 提供了但不可识别，回退到扩展名
+        if ext_mapping:
+            media_type = ext_mapping["media_type"]
+            resolved_ct = ext_mapping["content_type"]
+        else:
+            ext_hint = f"和文件扩展名 .{ext}" if ext else ""
+            base_result.reason = (
+                f"无法识别的 content_type: {content_type}{ext_hint}，"
+                "该文件类型不受支持。"
+            )
+            return base_result
+    else:
+        # 无 content_type，使用扩展名
+        if ext_mapping:
+            media_type = ext_mapping["media_type"]
+            resolved_ct = ext_mapping["content_type"]
+        elif ext:
+            base_result.reason = f"无法识别的文件扩展名 .{ext}，该文件类型不受支持。"
+            return base_result
+        else:
+            base_result.reason = (
+                "文件没有扩展名且未提供 content_type，无法判断文件类型。"
+            )
+            return base_result
+
+    # 4. 检查文件大小
+    file_size = os.path.getsize(abs_path)
+    size_limit = SIZE_LIMITS.get(media_type, DEFAULT_SIZE_LIMIT)
+
+    if file_size > size_limit:
+        base_result.file_size = file_size
+        base_result.media_type = media_type
+        base_result.content_type = resolved_ct
+        base_result.reason = (
+            f"文件大小 {_format_size(file_size)} "
+            f"超过此类型的限制 {_format_size(size_limit)}。"
+        )
+        return base_result
+
+    # 5. 全部通过
+    logger.log(f"  预检通过: ext={ext} media_type={media_type} "
+               f"content_type={resolved_ct} size={_format_size(file_size)}")
+    return FileCheckResult(
+        passed=True,
+        file_path=abs_path,
+        file_name=file_name,
+        file_ext=ext,
+        file_size=file_size,
+        media_type=media_type,
+        content_type=resolved_ct,
+    )