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

ima_sdk/knowledge_base.py

@@ -0,0 +1,591 @@
+"""知识库管理器。
+
+封装 IMA 知识库相关的所有 API 操作，包括搜索、浏览、上传文件、导入 URL 等。
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from typing import Dict, List, Optional, Tuple
+
+from .client import ImaClient
+from .cos_uploader import cos_upload
+from .file_checker import check_file
+from . import logger
+from .types import (
+    AddableKnowledgeBaseInfo,
+    CheckRepeatedNameResult,
+    CosCredential,
+    CreateMediaResult,
+    FileInfo,
+    ImportURLResult,
+    KBFolderInfo,
+    KnowledgeBaseInfo,
+    KnowledgeInfo,
+    KnowledgeListResult,
+    SearchedKnowledgeBaseInfo,
+    SearchedKnowledgeInfo,
+)
+
+# 知识库 API 基础路径
+_BASE = "openapi/wiki/v1"
+
+
+class KnowledgeBaseManager:
+    """知识库管理器。
+
+    用法::
+
+        from ima_sdk import ImaClient, KnowledgeBaseManager
+
+        client = ImaClient()
+        kb = KnowledgeBaseManager(client)
+
+        # 搜索知识库
+        results = kb.search_knowledge_base("我的资料")
+
+        # 上传文件
+        kb.upload_file("/path/to/report.pdf", knowledge_base_id="xxx")
+    """
+
+    def __init__(self, client: ImaClient):
+        self._client = client
+
+    # ── 获取知识库信息 ────────────────────────────────────────────────────────
+
+    def get_knowledge_base(self, ids: List[str]) -> Dict[str, KnowledgeBaseInfo]:
+        """获取知识库信息。
+
+        Args:
+            ids: 知识库 ID 列表（1-20 个，不重复）
+
+        Returns:
+            知识库 ID → KnowledgeBaseInfo 的映射
+        """
+        data = self._client.post(f"{_BASE}/get_knowledge_base", {"ids": ids})
+        infos = data.get("infos", {})
+        logger.log_result("获取知识库信息", infos)
+        return {
+            kb_id: KnowledgeBaseInfo.from_dict(info) for kb_id, info in infos.items()
+        }
+
+    # ── 搜索知识库列表 ────────────────────────────────────────────────────────
+
+    def search_knowledge_base(
+        self,
+        query: str,
+        *,
+        limit: int = 50,
+        fetch_all: bool = False,
+    ) -> List[SearchedKnowledgeBaseInfo]:
+        """按名称搜索知识库。
+
+        Args:
+            query: 搜索关键词
+            limit: 每页数量（1-50）
+            fetch_all: 是否自动翻页获取全部结果
+
+        Returns:
+            搜索到的知识库列表
+        """
+        body = {"query": query}
+        logger.log_step("搜索知识库", f"query={query!r} fetch_all={fetch_all}")
+        if fetch_all:
+            return self._client.paginate_items(
+                f"{_BASE}/search_knowledge_base",
+                body,
+                "info_list",
+                SearchedKnowledgeBaseInfo.from_dict,
+                limit=limit,
+            )
+        data = self._client.post(
+            f"{_BASE}/search_knowledge_base",
+            {**body, "cursor": "", "limit": limit},
+        )
+        return [
+            SearchedKnowledgeBaseInfo.from_dict(i)
+            for i in data.get("info_list", [])
+        ]
+
+    # ── 获取可添加的知识库列表 ────────────────────────────────────────────────
+
+    def get_addable_list(
+        self, *, limit: int = 50, fetch_all: bool = False
+    ) -> List[AddableKnowledgeBaseInfo]:
+        """获取当前用户有权限添加内容的知识库列表。
+
+        Args:
+            limit: 每页数量（1-50）
+            fetch_all: 是否自动翻页获取全部结果
+
+        Returns:
+            可添加的知识库列表
+        """
+        logger.log_step("获取可添加知识库列表", f"fetch_all={fetch_all}")
+        if fetch_all:
+            return self._client.paginate_items(
+                f"{_BASE}/get_addable_knowledge_base_list",
+                {},
+                "addable_knowledge_base_list",
+                AddableKnowledgeBaseInfo.from_dict,
+                limit=limit,
+            )
+        data = self._client.post(
+            f"{_BASE}/get_addable_knowledge_base_list",
+            {"cursor": "", "limit": limit},
+        )
+        return [
+            AddableKnowledgeBaseInfo.from_dict(i)
+            for i in data.get("addable_knowledge_base_list", [])
+        ]
+
+    # ── 浏览知识库内容 ────────────────────────────────────────────────────────
+
+    def get_knowledge_list(
+        self,
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+        limit: int = 50,
+    ) -> KnowledgeListResult:
+        """浏览知识库内容（单页）。
+
+        Args:
+            knowledge_base_id: 知识库 ID
+            folder_id: 文件夹 ID（省略则列出根目录）
+            limit: 数量限制（1-50）
+
+        Returns:
+            KnowledgeListResult 包含知识条目、文件夹和翻页信息
+        """
+        logger.log_step("浏览知识库内容", f"kb_id={knowledge_base_id} folder_id={folder_id}")
+        body: dict = {
+            "cursor": "",
+            "limit": limit,
+            "knowledge_base_id": knowledge_base_id,
+        }
+        if folder_id:
+            body["folder_id"] = folder_id
+        data = self._client.post(f"{_BASE}/get_knowledge_list", body)
+        return KnowledgeListResult.from_dict(data)
+
+    def get_all_knowledge(
+        self,
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+        limit: int = 50,
+    ) -> Tuple[List[KnowledgeInfo], List[KBFolderInfo]]:
+        """翻页获取知识库所有内容。
+
+        Args:
+            knowledge_base_id: 知识库 ID
+            folder_id: 文件夹 ID
+            limit: 每页数量
+
+        Returns:
+            (知识条目列表, 文件夹列表) 的元组
+        """
+        logger.log_step("获取知识库全部内容", f"kb_id={knowledge_base_id} folder_id={folder_id}")
+        all_knowledge: List[KnowledgeInfo] = []
+        all_folders: List[KBFolderInfo] = []
+        body: dict = {"knowledge_base_id": knowledge_base_id}
+        if folder_id:
+            body["folder_id"] = folder_id
+
+        for data in self._client.paginate(
+            f"{_BASE}/get_knowledge_list", body, "knowledge_list", limit=limit
+        ):
+            for item in data.get("knowledge_list", []):
+                all_knowledge.append(KnowledgeInfo.from_dict(item))
+            for item in data.get("folder_list", []):
+                all_folders.append(KBFolderInfo.from_dict(item))
+
+        return all_knowledge, all_folders
+
+    # ── 搜索知识库内容 ────────────────────────────────────────────────────────
+
+    def search_knowledge(
+        self,
+        query: str,
+        knowledge_base_id: str,
+        *,
+        fetch_all: bool = False,
+    ) -> List[SearchedKnowledgeInfo]:
+        """在知识库中搜索内容。
+
+        Args:
+            query: 搜索关键词
+            knowledge_base_id: 知识库 ID
+            fetch_all: 是否自动翻页获取全部结果
+
+        Returns:
+            搜索到的知识条目列表
+        """
+        logger.log_step("搜索知识库内容", f"query={query!r} kb_id={knowledge_base_id}")
+        body = {"query": query, "knowledge_base_id": knowledge_base_id}
+        if fetch_all:
+            return self._client.paginate_items(
+                f"{_BASE}/search_knowledge",
+                body,
+                "info_list",
+                SearchedKnowledgeInfo.from_dict,
+            )
+        data = self._client.post(
+            f"{_BASE}/search_knowledge",
+            {**body, "cursor": ""},
+        )
+        return [
+            SearchedKnowledgeInfo.from_dict(i) for i in data.get("info_list", [])
+        ]
+
+    # ── 目录路径解析 ────────────────────────────────────────────────────────────
+
+    _FOLDER_MEDIA_TYPE = 99
+    _FOLDER_PREFIX = "folder_"
+
+    def _list_subfolders(
+        self, knowledge_base_id: str, folder_id: Optional[str] = None
+    ) -> List[Tuple[str, str]]:
+        """列出指定目录下的子文件夹。
+
+        API 返回的文件夹以 media_type=99、media_id 以 "folder_" 前缀出现在
+        knowledge_list 中。
+
+        Returns:
+            [(folder_name, folder_id), ...]
+        """
+        result = self.get_knowledge_list(
+            knowledge_base_id, folder_id=folder_id, limit=50
+        )
+        folders = []
+        for item in result.knowledge_list:
+            if item.media_id.startswith(self._FOLDER_PREFIX):
+                folders.append((item.title, item.media_id))
+        return folders
+
+    def resolve_folder_path(
+        self, knowledge_base_id: str, path: str
+    ) -> Optional[str]:
+        """按路径名逐级查找 folder_id。
+
+        Args:
+            knowledge_base_id: 知识库 ID
+            path: 目录路径，如 "论文/机器学习"，"/" 或空字符串表示根目录
+
+        Returns:
+            最终的 folder_id，根目录返回 None
+
+        Raises:
+            ValueError: 路径中某一级文件夹不存在
+        """
+        path = path.strip().strip("/")
+        if not path:
+            return None
+
+        parts = path.split("/")
+        current_folder_id: Optional[str] = None
+
+        for part in parts:
+            folders = self._list_subfolders(knowledge_base_id, current_folder_id)
+            matched = [fid for name, fid in folders if name == part]
+            if not matched:
+                existing = [name for name, _ in folders]
+                loc = f"folder_id={current_folder_id}" if current_folder_id else "根目录"
+                raise ValueError(
+                    f"在 {loc} 下找不到文件夹 '{part}'，现有文件夹: {existing or '(空)'}"
+                )
+            current_folder_id = matched[0]
+
+        return current_folder_id
+
+    # ── 检查文件名重复 ────────────────────────────────────────────────────────
+
+    def check_repeated_names(
+        self,
+        params: List[Dict],
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+    ) -> List[CheckRepeatedNameResult]:
+        """检查文件名是否在知识库中重复。
+
+        Args:
+            params: 待检查列表，每项包含 name 和 media_type
+                    例: [{"name": "report.pdf", "media_type": 1}]
+            knowledge_base_id: 知识库 ID
+            folder_id: 文件夹 ID（省略则检查根目录）
+
+        Returns:
+            检查结果列表
+        """
+        body: dict = {"params": params, "knowledge_base_id": knowledge_base_id}
+        if folder_id:
+            body["folder_id"] = folder_id
+        data = self._client.post(f"{_BASE}/check_repeated_names", body)
+        return [
+            CheckRepeatedNameResult.from_dict(r) for r in data.get("results", [])
+        ]
+
+    # ── 导入 URL ──────────────────────────────────────────────────────────────
+
+    def import_urls(
+        self,
+        urls: List[str],
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+    ) -> Dict[str, ImportURLResult]:
+        """导入网页或微信公众号文章到知识库。
+
+        Args:
+            urls: URL 列表（1-10 个）
+            knowledge_base_id: 知识库 ID
+            folder_id: 文件夹 ID（省略则添加到根目录）
+
+        Returns:
+            URL → ImportURLResult 的映射
+        """
+        body: dict = {
+            "urls": urls,
+            "knowledge_base_id": knowledge_base_id,
+            "folder_id": folder_id or knowledge_base_id,
+        }
+        data = self._client.post(f"{_BASE}/import_urls", body)
+        results = data.get("results", {})
+        return {
+            url: ImportURLResult.from_dict(info) for url, info in results.items()
+        }
+
+    # ── 创建媒体 ──────────────────────────────────────────────────────────────
+
+    def create_media(
+        self,
+        file_name: str,
+        file_size: int,
+        content_type: str,
+        knowledge_base_id: str,
+        file_ext: str,
+    ) -> CreateMediaResult:
+        """创建媒体，获取 COS 上传凭证。
+
+        这是文件上传流程的第一步。
+
+        Args:
+            file_name: 文件名称（最长 1024 字符）
+            file_size: 文件大小（字节）
+            content_type: MIME 类型
+            knowledge_base_id: 知识库 ID
+            file_ext: 文件后缀名（无点号，如 "pdf"）
+
+        Returns:
+            CreateMediaResult 包含 media_id 和 COS 凭证
+        """
+        data = self._client.post(
+            f"{_BASE}/create_media",
+            {
+                "file_name": file_name,
+                "file_size": file_size,
+                "content_type": content_type,
+                "knowledge_base_id": knowledge_base_id,
+                "file_ext": file_ext,
+            },
+        )
+        return CreateMediaResult.from_dict(data)
+
+    # ── 添加知识 ──────────────────────────────────────────────────────────────
+
+    def add_knowledge(
+        self,
+        media_type: int,
+        title: str,
+        knowledge_base_id: str,
+        *,
+        media_id: Optional[str] = None,
+        folder_id: Optional[str] = None,
+        file_info: Optional[FileInfo] = None,
+        note_content_id: Optional[str] = None,
+        web_url: Optional[str] = None,
+        session_id: Optional[str] = None,
+    ) -> str:
+        """添加知识到知识库。
+
+        Args:
+            media_type: 媒体类型（见 MediaType 枚举）
+            title: 标题
+            knowledge_base_id: 知识库 ID
+            media_id: 文件上传时必填，create_media 返回的 ID
+            folder_id: 文件夹 ID（省略则添加到根目录）
+            file_info: 文件信息（文件上传时必填）
+            note_content_id: 笔记 doc_id（media_type=11 时使用）
+            web_url: 网页 URL（media_type=2 时使用）
+            session_id: AI 会话 ID（media_type=12 时使用）
+
+        Returns:
+            媒体 ID
+        """
+        body: dict = {
+            "media_type": media_type,
+            "title": title,
+            "knowledge_base_id": knowledge_base_id,
+        }
+        if media_id:
+            body["media_id"] = media_id
+        if folder_id:
+            body["folder_id"] = folder_id
+        if file_info:
+            body["file_info"] = file_info.to_dict()
+        if note_content_id:
+            body["note_info"] = {"content_id": note_content_id}
+        if web_url:
+            body["web_info"] = {"content_id": web_url}
+        if session_id:
+            body["session_info"] = {"content_id": session_id}
+
+        data = self._client.post(f"{_BASE}/add_knowledge", body)
+        return data.get("media_id", "")
+
+    # ── 上传预检（dry run）─────────────────────────────────────────────────────
+
+    def dry_run_upload(
+        self,
+        file_path: str,
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+        content_type: Optional[str] = None,
+    ) -> dict:
+        """只执行文件预检和重名检查，不实际上传。
+
+        Args:
+            file_path: 本地文件路径
+            knowledge_base_id: 目标知识库 ID
+            folder_id: 目标文件夹 ID
+            content_type: 可选的 MIME 类型
+
+        Returns:
+            包含预检结果和重名检查结果的字典
+        """
+        check = check_file(file_path, content_type)
+        result: dict = {
+            "file_check": {
+                "passed": check.passed,
+                "file_name": check.file_name,
+                "file_ext": check.file_ext,
+                "file_size": check.file_size,
+                "media_type": check.media_type,
+                "content_type": check.content_type,
+                "reason": check.reason,
+            },
+        }
+        if not check.passed:
+            return result
+
+        dup_results = self.check_repeated_names(
+            [{"name": check.file_name, "media_type": check.media_type}],
+            knowledge_base_id,
+            folder_id=folder_id,
+        )
+        result["duplicate_check"] = [
+            {"name": r.name, "is_repeated": r.is_repeated} for r in dup_results
+        ]
+        return result
+
+    # ── 完整文件上传流程 ──────────────────────────────────────────────────────
+
+    def upload_file(
+        self,
+        file_path: str,
+        knowledge_base_id: str,
+        *,
+        folder_id: Optional[str] = None,
+        content_type: Optional[str] = None,
+        skip_duplicate_check: bool = False,
+    ) -> str:
+        """完整的文件上传流程。
+
+        自动执行：预检 → 重名检查 → 创建媒体 → COS 上传 → 添加知识。
+
+        Args:
+            file_path: 本地文件路径
+            knowledge_base_id: 目标知识库 ID
+            folder_id: 目标文件夹 ID
+            content_type: 可选的 MIME 类型
+            skip_duplicate_check: 是否跳过重名检查
+
+        Returns:
+            上传后的 media_id
+
+        Raises:
+            ValueError: 文件预检失败
+            RuntimeError: COS 上传失败
+            ImaApiError: API 调用失败
+        """
+        # 1. 文件预检
+        logger.log_step("步骤 1/5: 文件预检", f"path={file_path}")
+        check = check_file(file_path, content_type)
+        if not check.passed:
+            raise ValueError(f"文件预检失败: {check.reason}")
+
+        logger.log(f"  预检通过: {check.file_name} ({check.file_ext}, "
+                   f"{check.file_size} bytes, media_type={check.media_type})")
+
+        # 2. 重名检查（可选）
+        if not skip_duplicate_check:
+            logger.log_step("步骤 2/5: 重名检查", f"name={check.file_name}")
+            dup_results = self.check_repeated_names(
+                [{"name": check.file_name, "media_type": check.media_type}],
+                knowledge_base_id,
+                folder_id=folder_id,
+            )
+            for r in dup_results:
+                if r.is_repeated:
+                    raise ValueError(f"文件名重复: {r.name}")
+            logger.log("  无重名文件")
+        else:
+            logger.log_step("步骤 2/5: 跳过重名检查")
+
+        # 3. 创建媒体，获取 COS 凭证
+        logger.log_step("步骤 3/5: 创建媒体", f"kb_id={knowledge_base_id}")
+        create_result = self.create_media(
+            file_name=check.file_name,
+            file_size=check.file_size,
+            content_type=check.content_type,
+            knowledge_base_id=knowledge_base_id,
+            file_ext=check.file_ext,
+        )
+
+        logger.log(f"  media_id={create_result.media_id}, "
+                   f"cos_key={create_result.cos_credential.cos_key}")
+
+        # 4. 上传到 COS
+        logger.log_step("步骤 4/5: 上传文件到 COS")
+        cos_upload(
+            file_path=check.file_path,
+            credential=create_result.cos_credential,
+            content_type=check.content_type,
+        )
+
+        logger.log("  COS 上传完成")
+
+        # 5. 添加知识到知识库
+        logger.log_step("步骤 5/5: 添加知识到知识库")
+        file_info = FileInfo(
+            cos_key=create_result.cos_credential.cos_key,
+            file_size=check.file_size,
+            file_name=check.file_name,
+            last_modify_time=int(os.path.getmtime(check.file_path)),
+        )
+
+        media_id = self.add_knowledge(
+            media_type=check.media_type,
+            title=check.file_name,
+            knowledge_base_id=knowledge_base_id,
+            media_id=create_result.media_id,
+            folder_id=folder_id,
+            file_info=file_info,
+        )
+
+        logger.log(f"  上传完成 media_id={media_id}")
+        return media_id

ima_sdk/logger.py

@@ -0,0 +1,89 @@
+"""日志工具模块。
+
+提供结构化的日志打印功能，仅输出到终端（stderr）。
+通过 enable() 函数或环境变量 IMA_SDK_VERBOSE=1 启用。
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import sys
+from typing import Any, Optional
+
+
+# 全局开关
+_enabled = os.environ.get("IMA_SDK_VERBOSE", "") == "1"
+
+
+def enable() -> None:
+    """启用日志输出。"""
+    global _enabled
+    _enabled = True
+
+
+def is_enabled() -> bool:
+    """返回日志是否启用。"""
+    return _enabled
+
+
+def _truncate(text: str, max_len: int = 500) -> str:
+    """截断过长文本。"""
+    if len(text) <= max_len:
+        return text
+    return text[:max_len] + f"... ({len(text)} chars)"
+
+
+def log(message: str) -> None:
+    """打印日志行到 stderr。"""
+    if _enabled:
+        print(f"[ima-sdk] {message}", file=sys.stderr)
+
+
+def log_step(step: str, detail: str = "") -> None:
+    """打印带步骤标记的日志。"""
+    if _enabled:
+        msg = f"[ima-sdk] >> {step}"
+        if detail:
+            msg += f"  {detail}"
+        print(msg, file=sys.stderr)
+
+
+def log_request(method: str, url: str, body: Optional[dict] = None) -> None:
+    """打印 HTTP 请求信息。"""
+    if not _enabled:
+        return
+    print(f"[ima-sdk] --> {method} {url}", file=sys.stderr)
+    if body is not None:
+        body_str = json.dumps(body, ensure_ascii=False)
+        print(f"[ima-sdk]     body: {_truncate(body_str)}", file=sys.stderr)
+
+
+def log_response(status: int, body_preview: str = "") -> None:
+    """打印 HTTP 响应信息。"""
+    if not _enabled:
+        return
+    msg = f"[ima-sdk] <-- {status}"
+    if body_preview:
+        msg += f"  {_truncate(body_preview, 300)}"
+    print(msg, file=sys.stderr)
+
+
+def log_result(label: str, value: Any) -> None:
+    """打印结果摘要。"""
+    if not _enabled:
+        return
+    if isinstance(value, list):
+        print(f"[ima-sdk]     {label}: [{len(value)} items]", file=sys.stderr)
+    elif isinstance(value, dict):
+        keys = list(value.keys())
+        print(f"[ima-sdk]     {label}: {{{', '.join(keys[:5])}}}", file=sys.stderr)
+    else:
+        text = str(value)
+        print(f"[ima-sdk]     {label}: {_truncate(text, 200)}", file=sys.stderr)
+
+
+def log_error(message: str) -> None:
+    """打印错误日志。"""
+    if _enabled:
+        print(f"[ima-sdk] !! {message}", file=sys.stderr)