agentengine-sdk-python 0.6.2__py3-none-any.whl → 0.6.9__py3-none-any.whl

ksadk/knowledge_base/client.py

@@ -0,0 +1,297 @@
+"""KnowledgeBaseClient - 金山云知识库检索客户端
+
+封装 AICP RetrieveKnowledge API，通过 kingsoftcloud SDK 调用。
+支持从环境变量自动配置，由各框架 Runner 自动集成。
+
+环境变量:
+    KSADK_KB_DATASET_ID: 知识库 ID (必填，存在即启用)
+    KSADK_KB_ACCESS_KEY: AK (可选，默认取 KSYUN_ACCESS_KEY)
+    KSADK_KB_SECRET_KEY: SK (可选，默认取 KSYUN_SECRET_KEY)
+    KSADK_KB_REGION: 区域 (默认 cn-beijing-6)
+    KSADK_KB_ENDPOINT: API 端点 (默认 aicp.api.ksyun.com)
+    KSADK_KB_TOP_K: 返回结果数 (默认 5)
+    KSADK_KB_SEARCH_METHOD: 检索方法 (默认 intelligence_search)
+    KSADK_KB_SCORE_THRESHOLD: 分数阈值 (可选)
+    KSADK_KB_RERANKING_ENABLE: 是否启用重排序 (默认 false)
+
+使用示例:
+    # 自动从环境变量配置 (推荐)
+    client = KnowledgeBaseClient.from_env()
+    results = client.search("如何使用知识库？")
+
+    # 手动创建
+    client = KnowledgeBaseClient(
+        dataset_id="your-kb-id",
+        access_key="ak",
+        secret_key="sk",
+    )
+"""
+
+import json
+import logging
+import os
+from typing import Any, List, Optional
+
+from pydantic import BaseModel, ConfigDict
+
+from ksadk.common.aicp_env import resolve_aicp_connection
+
+logger = logging.getLogger(__name__)
+
+
+class KnowledgeBaseResult(BaseModel):
+    """单条知识库检索结果"""
+
+    content: str = ""
+    score: float = 0.0
+    segment_id: str = ""
+    document_id: str = ""
+    document_name: str = ""
+    position: int = 0
+    answer: str = ""
+    keywords: List[str] = []
+
+
+class KnowledgeBaseClient(BaseModel):
+    """金山云知识库检索客户端
+
+    Attributes:
+        dataset_id: 知识库 ID (DatasetId)
+        access_key: 访问密钥 ID (AK)
+        secret_key: 访问密钥 (SK)
+        region: API 区域
+        endpoint: API 端点
+        top_k: 返回结果数
+        search_method: 检索方法
+        score_threshold: 分数阈值
+        score_threshold_enabled: 是否启用阈值过滤
+        reranking_enable: 是否启用重排序
+    """
+
+    dataset_id: str
+    access_key: str = ""
+    secret_key: str = ""
+    region: str = "cn-beijing-6"
+    endpoint: str = "aicp.api.ksyun.com"
+    scheme: str = "https"
+    top_k: int = 5
+    search_method: str = "intelligence_search"
+    score_threshold: float = 0.0
+    score_threshold_enabled: bool = False
+    reranking_enable: bool = False
+
+    _aicp_client: Any = None
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    def model_post_init(self, __context: Any) -> None:
+        if not self.dataset_id:
+            raise ValueError("dataset_id is required for KnowledgeBaseClient")
+
+        if not self.access_key or not self.secret_key:
+            logger.warning(
+                "KnowledgeBaseClient: AK/SK not provided, "
+                "API calls will fail. Set KSADK_KB_ACCESS_KEY/KSADK_KB_SECRET_KEY "
+                "or KSYUN_ACCESS_KEY/KSYUN_SECRET_KEY."
+            )
+
+    def _get_client(self):
+        """懒加载 AICP 客户端"""
+        if self._aicp_client is not None:
+            return self._aicp_client
+
+        try:
+            from ksyun.common import credential
+            from ksyun.common.profile.client_profile import ClientProfile
+            from ksyun.common.profile.http_profile import HttpProfile
+        except ImportError:
+            raise ImportError(
+                "kingsoftcloud-sdk-python is required for knowledge base. "
+                "Install it with: pip install kingsoftcloud-sdk-python"
+            )
+
+        # 尝试导入 aicp client (多版本 fallback)
+        aicp_module = None
+        for version in ["v20251114", "v20251212", "v20240612"]:
+            try:
+                aicp_module = __import__(
+                    f"ksyun.client.aicp.{version}.client",
+                    fromlist=["AicpClient"],
+                )
+                logger.debug(f"Using aicp client version: {version}")
+                break
+            except ImportError:
+                continue
+
+        if aicp_module is None:
+            raise ImportError(
+                "Cannot import ksyun.client.aicp client. "
+                "Ensure kingsoftcloud-sdk-python is installed and up to date."
+            )
+
+        cred = credential.Credential(self.access_key, self.secret_key)
+
+        http_profile = HttpProfile()
+        http_profile.endpoint = self.endpoint
+        http_profile.reqMethod = "POST"
+        http_profile.reqTimeout = 60
+        http_profile.scheme = self.scheme
+
+        client_profile = ClientProfile()
+        client_profile.httpProfile = http_profile
+
+        self._aicp_client = aicp_module.AicpClient(
+            cred, self.region, profile=client_profile
+        )
+
+        # 强制覆写 API 版本为 RetrieveKnowledge 所需的 2025-11-14
+        # SDK 的 _apiVersion 由导入的模块版本决定，可能不匹配
+        self._aicp_client._apiVersion = "2025-11-14"
+
+        logger.info(
+            f"KnowledgeBaseClient initialized: "
+            f"dataset_id={self.dataset_id}, region={self.region}, "
+            f"endpoint={self.endpoint}"
+        )
+        return self._aicp_client
+
+    def _build_params(self, query: str, top_k: Optional[int] = None) -> dict:
+        """构建 RetrieveKnowledge 请求参数 (JSON 嵌套格式)"""
+        effective_top_k = top_k if top_k is not None else self.top_k
+
+        params = {
+            "DatasetId": self.dataset_id,
+            "Query": query,
+            "RetrievalModel": {
+                "SearchMethod": self.search_method,
+                "TopK": effective_top_k,
+                "RerankingEnable": self.reranking_enable,
+            },
+        }
+
+        if self.score_threshold_enabled:
+            params["RetrievalModel"]["ScoreThresholdEnabled"] = True
+            params["RetrievalModel"]["ScoreThreshold"] = self.score_threshold
+
+        return params
+
+    def _parse_response(self, response: str) -> List[KnowledgeBaseResult]:
+        """解析 RetrieveKnowledge 响应"""
+        try:
+            data = json.loads(response) if isinstance(response, str) else response
+        except (json.JSONDecodeError, TypeError):
+            logger.error(f"Failed to parse response: {str(response)[:200]}")
+            return []
+
+        records = data.get("Records", [])
+        results = []
+
+        for record in records:
+            segment = record.get("Segment", {})
+            document = segment.get("Document", {})
+
+            results.append(
+                KnowledgeBaseResult(
+                    content=segment.get("Content", ""),
+                    score=record.get("Score", 0.0),
+                    segment_id=segment.get("Id", ""),
+                    document_id=segment.get("DocumentId", ""),
+                    document_name=document.get("Name", ""),
+                    position=segment.get("Position", 0),
+                    answer=segment.get("Answer", ""),
+                    keywords=segment.get("Keywords", []),
+                )
+            )
+
+        return results
+
+    def search(
+        self, query: str, top_k: Optional[int] = None
+    ) -> List[KnowledgeBaseResult]:
+        """检索知识库
+
+        Args:
+            query: 检索关键词
+            top_k: 返回结果数 (覆盖默认值)
+
+        Returns:
+            匹配的文档片段列表
+        """
+        client = self._get_client()
+        params = self._build_params(query, top_k)
+
+        logger.info(
+            f"Searching knowledge base: dataset_id={self.dataset_id}, "
+            f"query='{query[:50]}'"
+        )
+
+        try:
+            response = client.call(
+                "RetrieveKnowledge", params, options={"IsPostJson": True}
+            )
+            results = self._parse_response(response)
+            logger.info(
+                f"Knowledge base returned {len(results)} results "
+                f"for query='{query[:50]}'"
+            )
+            return results
+        except Exception as e:
+            logger.error(f"Knowledge base search failed: {e}")
+            raise
+
+    @classmethod
+    def from_env(cls) -> "KnowledgeBaseClient":
+        """从环境变量创建 KnowledgeBaseClient
+
+        Returns:
+            KnowledgeBaseClient 实例
+
+        Raises:
+            ValueError: 如果 KSADK_KB_DATASET_ID 未设置
+        """
+        dataset_id = os.environ.get("KSADK_KB_DATASET_ID", "")
+        if not dataset_id:
+            raise ValueError(
+                "KSADK_KB_DATASET_ID environment variable is required "
+                "to enable knowledge base."
+            )
+
+        access_key = (
+            os.environ.get("KSADK_KB_ACCESS_KEY")
+            or os.environ.get("KSYUN_ACCESS_KEY")
+            or os.environ.get("KSYUN_SECRET_ID", "")
+        )
+        secret_key = (
+            os.environ.get("KSADK_KB_SECRET_KEY")
+            or os.environ.get("KSYUN_SECRET_KEY")
+            or os.environ.get("KSYUN_SECRET_KEY", "")
+        )
+
+        score_threshold_str = os.environ.get("KSADK_KB_SCORE_THRESHOLD", "")
+        score_threshold = float(score_threshold_str) if score_threshold_str else 0.0
+        score_threshold_enabled = bool(score_threshold_str)
+
+        reranking_str = os.environ.get("KSADK_KB_RERANKING_ENABLE", "false")
+        reranking_enable = reranking_str.lower() in ("true", "1", "yes")
+        connection = resolve_aicp_connection("KSADK_KB")
+
+        return cls(
+            dataset_id=dataset_id,
+            access_key=access_key,
+            secret_key=secret_key,
+            region=connection["region"],
+            endpoint=connection["endpoint"],
+            scheme=connection["scheme"],
+            top_k=int(os.environ.get("KSADK_KB_TOP_K", "5")),
+            search_method=os.environ.get(
+                "KSADK_KB_SEARCH_METHOD", "intelligence_search"
+            ),
+            score_threshold=score_threshold,
+            score_threshold_enabled=score_threshold_enabled,
+            reranking_enable=reranking_enable,
+        )
+
+    @staticmethod
+    def is_configured() -> bool:
+        """检查环境变量是否已配置知识库"""
+        return bool(os.environ.get("KSADK_KB_DATASET_ID"))

ksadk/knowledge_base/langchain_tool.py

@@ -0,0 +1,55 @@
+"""LangChain / LangGraph 知识库工具包装
+
+将通用 search_knowledge 函数包装为 LangChain Tool。
+用户在 LangGraph/LangChain 的 agent 中一行 import 即可使用。
+
+使用示例:
+    from ksadk.knowledge_base.langchain_tool import search_knowledge_base
+    tools = [search_knowledge_base, ...]
+"""
+
+import logging
+
+from ksadk.knowledge_base.tool import search_knowledge
+
+logger = logging.getLogger(__name__)
+
+
+def create_langchain_tool():
+    """创建 LangChain Tool
+
+    Returns:
+        LangChain @tool 装饰的函数
+    """
+    try:
+        from langchain_core.tools import tool
+
+        @tool
+        def search_knowledge_base(query: str) -> str:
+            """搜索知识库获取相关信息。
+
+            当需要查找专业知识、文档内容或特定领域信息时使用此工具。
+            会自动从已配置的金山云知识库中检索最相关的内容。
+
+            Args:
+                query: 检索关键词或问题
+            """
+            return search_knowledge(query)
+
+        return search_knowledge_base
+
+    except ImportError:
+        logger.warning(
+            "langchain-core not installed, returning raw function. "
+            "Install with: pip install langchain-core"
+        )
+
+        def search_knowledge_base(query: str) -> str:
+            """搜索知识库获取相关信息"""
+            return search_knowledge(query)
+
+        return search_knowledge_base
+
+
+# 预创建实例，方便直接 import 使用
+search_knowledge_base = create_langchain_tool()

ksadk/knowledge_base/service.py

@@ -0,0 +1,66 @@
+"""Platform-level knowledge base service helpers."""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from ksadk.knowledge_base.client import KnowledgeBaseClient, KnowledgeBaseResult
+
+logger = logging.getLogger(__name__)
+
+
+def format_knowledge_results(results: list[KnowledgeBaseResult]) -> str:
+    if not results:
+        return "未找到相关知识库内容。"
+
+    formatted_parts: list[str] = []
+    for i, result in enumerate(results, 1):
+        part = f"[{i}] "
+        if result.document_name:
+            part += f"(来源: {result.document_name}) "
+        part += result.content
+        if result.answer:
+            part += f"\n    答案: {result.answer}"
+        formatted_parts.append(part)
+
+    return "\n\n".join(formatted_parts)
+
+
+class KnowledgeBaseService:
+    def __init__(self, client: KnowledgeBaseClient | None = None):
+        self._client = client
+
+    @classmethod
+    def from_env(cls) -> "KnowledgeBaseService":
+        return cls(KnowledgeBaseClient.from_env())
+
+    @staticmethod
+    def is_configured() -> bool:
+        return KnowledgeBaseClient.is_configured()
+
+    def _get_client(self) -> KnowledgeBaseClient:
+        if self._client is None:
+            self._client = KnowledgeBaseClient.from_env()
+        return self._client
+
+    def search(self, query: str, top_k: Optional[int] = None) -> list[KnowledgeBaseResult]:
+        return self._get_client().search(query, top_k)
+
+    def search_text(self, query: str, top_k: Optional[int] = None) -> str:
+        try:
+            return format_knowledge_results(self.search(query, top_k))
+        except Exception as exc:
+            logger.error("search_knowledge failed: %s", exc)
+            return f"知识库检索失败: {exc}"
+
+    def build_context(self, query: str, top_k: Optional[int] = None) -> dict[str, str] | None:
+        normalized = str(query or "").strip()
+        if not normalized:
+            return None
+        if not self.is_configured():
+            return None
+        return {
+            "query": normalized,
+            "formatted_text": self.search_text(normalized, top_k),
+        }

ksadk/knowledge_base/tool.py

@@ -0,0 +1,50 @@
+"""通用知识库检索函数
+
+提供跨框架的 search_knowledge 函数，由各框架工具包装调用。
+客户端实例通过单例模式懒加载，自动从环境变量配置。
+
+使用示例:
+    from ksadk.knowledge_base.tool import search_knowledge
+    result = search_knowledge("如何配置知识库？")
+    print(result)
+"""
+
+import logging
+from typing import Optional
+
+from ksadk.knowledge_base.service import KnowledgeBaseService
+
+logger = logging.getLogger(__name__)
+
+# 单例客户端
+_service: Optional[KnowledgeBaseService] = None
+
+
+def _get_or_create_service() -> KnowledgeBaseService:
+    """获取或创建知识库 service (单例)"""
+    global _service
+    if _service is None:
+        _service = KnowledgeBaseService.from_env()
+        logger.info("KnowledgeBaseService singleton created from environment")
+    return _service
+
+
+def search_knowledge(query: str, top_k: Optional[int] = None) -> str:
+    """检索知识库
+
+    从环境变量自动配置的知识库中检索相关内容。
+    返回格式化的文本结果，供 LLM 参考。
+
+    Args:
+        query: 检索关键词/问题
+        top_k: 返回结果数 (可选，覆盖环境变量配置)
+
+    Returns:
+        格式化的检索结果文本
+    """
+    try:
+        service = _get_or_create_service()
+        return service.search_text(query, top_k)
+    except Exception as e:
+        logger.error(f"search_knowledge failed: {e}")
+        return f"知识库检索失败: {e}"

ksadk/markdown.py

@@ -0,0 +1,169 @@
+from __future__ import annotations
+
+from typing import Any
+
+
+_FENCE_MARKER_CHARS = ("`", "~")
+_LIST_PREFIXES = ("- ", "* ", "+ ")
+
+
+def repair_markdown(text: Any, *, enabled: bool = False) -> str:
+    """尽量修复 LLM 输出中的常见 Markdown 形态问题。
+
+    这是业务侧按需调用的保守修复工具。KsADK runtime 默认保留模型原文，
+    不会自动调用这个函数。
+    """
+
+    if text is None:
+        return ""
+    value = str(text)
+    if not enabled:
+        return value
+    if value == "":
+        return ""
+
+    normalized = value.replace("\r\n", "\n").replace("\r", "\n")
+    lines = _strip_terminal_empty_lines(normalized.split("\n"))
+    if not lines:
+        return ""
+
+    lines = _normalize_block_spacing(lines)
+    lines = _close_unclosed_fence(lines)
+
+    return "\n".join(lines) + "\n"
+
+
+def _strip_terminal_empty_lines(lines: list[str]) -> list[str]:
+    start = 0
+    end = len(lines)
+    while start < end and lines[start] == "":
+        start += 1
+    while end > start and lines[end - 1] == "":
+        end -= 1
+    return lines[start:end]
+
+
+def _normalize_block_spacing(lines: list[str]) -> list[str]:
+    result: list[str] = []
+    in_fence = False
+    fence_marker = ""
+    previous_block = ""
+
+    for index, line in enumerate(lines):
+        stripped = line.lstrip()
+        block = _line_block_type(line, in_fence)
+
+        if block == "fence":
+            marker = _fence_marker(stripped) or ""
+            if not in_fence:
+                _append_blank_before_block(result, previous_block)
+                in_fence = True
+                fence_marker = marker
+            elif _is_closing_fence(marker, fence_marker):
+                in_fence = False
+                fence_marker = ""
+        elif not in_fence:
+            if block in {"table", "list"}:
+                _append_blank_before_block(result, previous_block)
+            elif block == "paragraph":
+                _append_blank_after_block(result, previous_block)
+
+        result.append(line)
+
+        if line == "":
+            previous_block = ""
+        elif not in_fence or block == "fence":
+            previous_block = block
+
+        next_line = lines[index + 1] if index + 1 < len(lines) else None
+        if next_line is None or in_fence:
+            continue
+        next_block = _line_block_type(next_line, in_fence)
+        if block in {"table", "list"} and next_block == "paragraph":
+            _append_blank_once(result)
+            previous_block = ""
+
+    return result
+
+
+def _append_blank_before_block(result: list[str], previous_block: str) -> None:
+    if not result or result[-1] == "":
+        return
+    if previous_block in {"", "table", "list", "fence"}:
+        return
+    result.append("")
+
+
+def _append_blank_after_block(result: list[str], previous_block: str) -> None:
+    if result and result[-1] != "" and previous_block in {"list", "table", "fence"}:
+        result.append("")
+
+
+def _append_blank_once(result: list[str]) -> None:
+    if result and result[-1] != "":
+        result.append("")
+
+
+def _close_unclosed_fence(lines: list[str]) -> list[str]:
+    open_marker = ""
+    for line in lines:
+        marker = _fence_marker(line.lstrip())
+        if marker is None:
+            continue
+        if not open_marker:
+            open_marker = marker
+        elif _is_closing_fence(marker, open_marker):
+            open_marker = ""
+    if open_marker:
+        return [*lines, open_marker]
+    return lines
+
+
+def _line_block_type(line: str, in_fence: bool) -> str:
+    if line == "":
+        return ""
+    stripped = line.lstrip()
+    if _fence_marker(stripped):
+        return "fence"
+    if in_fence:
+        return "code"
+    if _is_table_line(stripped):
+        return "table"
+    if _is_list_line(stripped):
+        return "list"
+    return "paragraph"
+
+
+def _fence_marker(stripped_line: str) -> str | None:
+    for char in _FENCE_MARKER_CHARS:
+        if not stripped_line.startswith(char * 3):
+            continue
+        length = 0
+        for value in stripped_line:
+            if value != char:
+                break
+            length += 1
+        return char * length
+    return None
+
+
+def _is_closing_fence(candidate: str, open_marker: str) -> bool:
+    return (
+        bool(candidate)
+        and bool(open_marker)
+        and candidate[0] == open_marker[0]
+        and len(candidate) >= len(open_marker)
+    )
+
+
+def _is_table_line(stripped_line: str) -> bool:
+    return stripped_line.startswith("|") and stripped_line.endswith("|") and stripped_line.count("|") >= 2
+
+
+def _is_list_line(stripped_line: str) -> bool:
+    if stripped_line.startswith(_LIST_PREFIXES):
+        return True
+    dot_index = stripped_line.find(". ")
+    if dot_index <= 0:
+        return False
+    return stripped_line[:dot_index].isdigit()