agentscope-runtime 0.2.0b1__py3-none-any.whl → 1.0.0__py3-none-any.whl

agentscope_runtime/tools/searches/modelstudio_search_lite.py

@@ -0,0 +1,310 @@
+# -*- coding: utf-8 -*-
+# pylint:disable=unused-argument, (unused-argument
+
+import json
+import os
+from typing import Any, Dict, List, Optional, Tuple
+
+import aiohttp
+from mcp.server.fastmcp import Context
+from pydantic import BaseModel, Field
+
+from .modelstudio_search import (
+    SEARCH_PAGE,
+    SEARCH_STRATEGY_SETTING,
+    SEARCH_TIMEOUT,
+    FieldValidator,
+)
+from ..utils.mcp_util import get_mcp_dash_request_id
+from ...engine.tracing import trace
+from ..base import Tool
+
+SEARCH_URL = os.getenv(
+    "SEARCH_URL",
+    "https://dashscope.aliyuncs.com/api/v1/indices/plugin/mcp_search",
+)
+SEARCH_STRATEGY = os.getenv("SEARCH_STRATEGY", "turbo")
+DASHSCOPE_API_KEY = os.getenv("DASHSCOPE_API_KEY", None)
+SEARCH_RULES = {
+    "url": {
+        "DROPOUT_ENTIRE_IF_MISSING": "DROPOUT_ENTIRE_IF_MISSING",
+        "FILTER_ITEMS_FROM_LIST": [
+            "uclink://www.uc.cn",
+            "qklink://www.quark.cn",
+            "https://page.sm.cn",
+        ],
+    },
+    "title": "DROPOUT_ENTIRE_IF_MISSING",
+}
+
+
+class SearchLiteInput(BaseModel):
+    """
+    Search Input.
+    """
+
+    query: str = Field(..., description="user query in the format of string")
+    count: int = Field(default=5, description="number of search results")
+    ctx: Optional[Context] = Field(
+        default=None,
+        description="HTTP request context containing headers",
+    )
+
+
+class SearchLiteOutput(BaseModel):
+    """
+    Search Output.
+    """
+
+    status: int = Field(
+        ...,
+        description="status code for success or failed, 0 is success",
+    )
+    pages: List[Dict[str, Any]] = Field(
+        ...,
+        description="Search results with url, title, etc,.",
+    )
+    tools: List[Dict[str, Any]] = Field(
+        ...,
+        description="The tool calling result during search",
+    )
+    request_id: Optional[str] = Field(
+        default=None,
+        title="Request ID",
+        description="请求ID",
+    )
+
+
+class ModelstudioSearchLite(Tool[SearchLiteInput, SearchLiteOutput]):
+    """
+    Search tool that calling dashscope for llm search result.
+    """
+
+    description = "搜索可用于查询百科知识、时事新闻、天气等信息"
+    name = "modelstudio_web_search"
+
+    @trace(trace_type="SEARCH", trace_name="modelstudio_search_lite")
+    async def _arun(
+        self,
+        args: SearchLiteInput,
+        **kwargs: Any,
+    ) -> SearchLiteOutput:
+        """
+        Run the search component with SearchLiteInput and return
+        SearchLiteOutput.
+        Args:
+            args: in SearchLiteInput format that generated by from llm
+            ctx: Optional MCP context containing request information
+            **kwargs: should include other arguments if needed,
+            'trace_event' in kwargs can be used for tracing_utils
+
+        Returns:SearchLiteOutput
+
+        """
+
+        if DASHSCOPE_API_KEY is None:
+            raise ValueError(
+                "DASHSCOPE_API_KEY is required for search component, "
+                "please get your dashscope api-key on Modelstudio platform",
+            )
+        trace_event = kwargs.pop("trace_event", None)
+        request_id = get_mcp_dash_request_id(args.ctx)
+
+        # call search engine to get search result
+        payload = ModelstudioSearchLite.generate_search_payload(
+            search_input=args,
+            request_id=request_id,
+            rid=request_id,
+        )
+        header = {
+            "Content-Type": "application/json",
+            "Accept-Encoding": "utf-8",
+            "Authorization": "Bearer " + DASHSCOPE_API_KEY,
+            "x-acs-req-uuid": request_id,
+        }
+        kwargs["context"] = {
+            "payload": payload,
+            "search_strategy": SEARCH_STRATEGY,
+            "timeout": SEARCH_TIMEOUT,
+        }
+
+        trace_event.on_log(
+            "",
+            **{
+                "step_suffix": "test_header_inject",
+                "payload": {
+                    "request_id": request_id,
+                    "dash_request_id": request_id,
+                    "search_header": header,
+                    "search_payload": payload,
+                    "search_url": SEARCH_URL,
+                },
+            },
+        )
+        try:
+            (
+                search_result,
+                extra_tool_info,
+                status,
+                result_message,
+            ) = await ModelstudioSearchLite.dashscope_search_kernel(
+                url=SEARCH_URL,
+                payload=payload,
+                headers=header,
+                timeout=SEARCH_TIMEOUT,
+            )
+            if trace_event:
+                trace_event.on_log(
+                    "",
+                    **{
+                        "step_suffix": "results",
+                        "payload": {
+                            "request_id": request_id,
+                            "search_status": status,
+                            "search_query_result": result_message,
+                            "search_result": search_result,
+                            "extra_tool_info": extra_tool_info,
+                        },
+                    },
+                )
+
+        except Exception as e:
+            print(f"Error: {e}")
+            return SearchLiteOutput(
+                status=1,
+                pages=[],
+                tools=[],
+                request_id=request_id,
+            )
+
+        # post process search results
+        (
+            search_items,
+            tool_items,
+        ) = ModelstudioSearchLite.post_process_search_detail(
+            search_results=search_result,
+            extra_tool_info=extra_tool_info,
+        )
+
+        return SearchLiteOutput(
+            status=0,
+            pages=search_items,
+            tools=tool_items,
+            request_id=request_id,
+        )
+
+    @staticmethod
+    def generate_search_payload(
+        search_input: SearchLiteInput,
+        request_id: str,
+        **kwargs: Any,
+    ) -> str:
+        search_strategy = SEARCH_STRATEGY
+
+        payload = {
+            "scene": SEARCH_STRATEGY_SETTING[search_strategy]["scene"],
+            "uq": search_input.query,
+            "rid": request_id,
+            "fields": [],
+            "page": int(SEARCH_PAGE),
+            "rows": search_input.count,
+            "headers": {
+                "__d_head_qto": SEARCH_STRATEGY_SETTING[search_strategy][
+                    "timeout"
+                ],
+            },
+        }
+        result = json.dumps(payload)
+        return result
+
+    @staticmethod
+    async def dashscope_search_kernel(
+        url: str,
+        payload: str,
+        headers: Dict,
+        timeout: int,
+        **kwargs: Any,
+    ) -> Tuple[List, List, int, str]:
+        extra_tool_info = []
+        results_list = []
+        results = {}
+        try:
+            # 使用异步HTTP客户端替代同步requests
+            timeout_config = aiohttp.ClientTimeout(total=timeout)
+            async with aiohttp.ClientSession(
+                timeout=timeout_config,
+            ) as session:
+                async with session.post(
+                    url,
+                    headers=headers,
+                    data=payload,
+                ) as response:
+                    results = await response.json()
+
+            result_message = results["message"]
+            if results["status"] == 0:
+                extra_tool_info = results["data"]["extras"].get(
+                    "toolResult",
+                    [],
+                )
+                results_list = results["data"]["docs"]
+        except Exception as e:
+            print(f"Error: {e}")
+            raise RuntimeError(
+                f"Unable to reach search source with error: {e} and result: "
+                f"{results}",
+            ) from e
+
+        return results_list, extra_tool_info, results["status"], result_message
+
+    @staticmethod
+    def post_process_search_detail(
+        search_results: List,
+        extra_tool_info: List,
+        **kwargs: Any,
+    ) -> Tuple[List[Dict[str, Any]], List[Dict[str, Any]]]:
+        field_validator = FieldValidator(SEARCH_RULES)
+        search_items = []
+        tool_items = []
+
+        try:
+            for doc in search_results:
+                tmp_search_result = {
+                    "url": doc.get("url", "") or "",
+                    "title": doc.get("title", "") or "",
+                    "icon": doc.get("hostlogo", "") or "",
+                    "site_name": doc.get("hostname", "") or "",
+                    "image": doc.get("image", "") or "",
+                }
+                filtered_search_result = field_validator.validate(
+                    tmp_search_result,
+                )
+                if filtered_search_result:
+                    url = doc.get("url")
+                    url = url.replace(" ", "%20").strip() or "expired_url"
+                    url = url.replace(
+                        "chatm6.sm.cn",
+                        "quark.sm.cn",
+                    )  # noqa E501
+                    search_items.append(
+                        {
+                            "snippet": doc.get("snippet", ""),
+                            "title": doc.get("title", ""),
+                            "url": url,
+                            "hostname": doc.get("hostname", ""),
+                            "hostlogo": doc.get("hostlogo", ""),
+                        },
+                    )
+            for item in extra_tool_info:
+                if item.get("tool"):
+                    tool_items.append(
+                        {
+                            "tool": item.get("tool"),
+                            "result": item.get("result", ""),
+                        },
+                    )
+        except Exception as e:
+            print(f"Error: {e}")
+            return [], []
+
+        return search_items, tool_items

agentscope_runtime/tools/utils/api_key_util.py

@@ -0,0 +1,45 @@
+# -*- coding: utf-8 -*-
+import os
+from enum import Enum
+from typing import Any, Optional
+
+
+class ApiNames(Enum):
+    """Enumeration of API key names."""
+
+    dashscope_api_key = "DASHSCOPE_API_KEY"
+
+
+def get_api_key(
+    api_enum: ApiNames,
+    key: Optional[str] = None,
+    **kwargs: Any,
+) -> str:
+    """Get API key from both input, kwargs, and env sources with priority
+    order.
+
+    Args:
+        api_enum (ApiNames): Enum of API name to retrieve.
+        key (Optional[str]): Default key value. Defaults to None.
+        **kwargs (Any): Additional keyword arguments that might contain the
+                API key.
+
+    Returns:
+        str: The API key value.
+
+    Raises:
+        AssertionError: If no API key is found from any source.
+    """
+    api_key = ""
+    if key is not None:
+        if kwargs.get(api_enum.name, "") != "":
+            if key != kwargs.get(api_enum.name):
+                # use runtime key instead of init key
+                api_key = kwargs.get(api_enum.name)
+        else:
+            api_key = key
+    else:
+        api_key = kwargs.get(api_enum.name, os.environ.get(api_enum.value, ""))
+
+    assert api_key != "", f"{api_enum.name} must be acquired"
+    return api_key

agentscope_runtime/tools/utils/crypto_utils.py

@@ -0,0 +1,99 @@
+# -*- coding: utf-8 -*-
+"""
+Crypto utilities for key format conversion.
+"""
+
+try:
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.backends import default_backend
+
+    CRYPTOGRAPHY_AVAILABLE = True
+except ImportError:
+    CRYPTOGRAPHY_AVAILABLE = False
+    serialization = None
+    default_backend = None
+
+
+def ensure_pkcs1_format(private_key_string: str) -> str:
+    """
+    Receive a private key string of unknown format and ensure the output is
+    in PKCS#1 format.
+
+    Supported input formats:
+    - PKCS#1 PEM format
+    - PKCS#8 PEM format
+
+    Args:
+        private_key_string: A string containing the private key (PEM format
+        or pure Base64)
+
+    Returns:
+        A string containing the private key in PKCS#1 PEM format.
+
+    Raises:
+        ImportError: Raised when the cryptography library is not installed.
+        ValueError: Raised when the private key string is invalid or empty.
+    """
+    # Check if cryptography library is available
+    if not CRYPTOGRAPHY_AVAILABLE:
+        raise ImportError(
+            "Please install the cryptography library: pip "
+            "install cryptography",
+        )
+
+    if not private_key_string or not private_key_string.strip():
+        raise ValueError("The private key string cannot be empty.")
+
+    key_string = private_key_string.strip()
+
+    try:
+        # Try to load directly as PEM format
+        if "-----BEGIN" in key_string:
+            private_key = serialization.load_pem_private_key(
+                data=key_string.encode("utf-8"),
+                password=None,
+                backend=default_backend(),
+            )
+        else:
+            # If it's pure Base64, first try PKCS#8 format
+            clean_base64 = "".join(key_string.split())
+
+            try:
+                # Attempt PKCS#8 format PEM wrapping
+                pkcs8_pem = (
+                    f"-----BEGIN {'PRIVATE'} {'KEY'}-----\n"
+                    f"{clean_base64}\n-----END {'PRIVATE'} {'KEY'}-----"
+                )
+                private_key = serialization.load_pem_private_key(
+                    data=pkcs8_pem.encode("utf-8"),
+                    password=None,
+                    backend=default_backend(),
+                )
+            except Exception:
+                # If PKCS#8 fails, try PKCS#1 format PEM wrapping
+                pkcs1_pem = (
+                    f"-----BEGIN RSA {'PRIVATE'} {'KEY'}-----\n"
+                    f"{clean_base64}\n-----END RSA {'PRIVATE'} {'KEY'}-----"
+                )
+                private_key = serialization.load_pem_private_key(
+                    data=pkcs1_pem.encode("utf-8"),
+                    password=None,
+                    backend=default_backend(),
+                )
+
+        # Force output in PKCS#1 format (Traditional OpenSSL format)
+        pkcs1_pem = private_key.private_bytes(
+            encoding=serialization.Encoding.PEM,
+            format=serialization.PrivateFormat.TraditionalOpenSSL,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+
+        return pkcs1_pem.decode("utf-8")
+
+    except Exception as e:
+        # If loading or conversion fails (e.g., the key is corrupted),
+        # raise an exception
+        raise ValueError(
+            f"Failed to parse or convert the provided private key. Please "
+            f"check if it is valid: {e}",
+        ) from e

agentscope_runtime/tools/utils/mcp_util.py

@@ -0,0 +1,35 @@
+# -*- coding: utf-8 -*-
+# pylint:disable=unused-argument
+
+import uuid
+from typing import Any
+
+from mcp.server.fastmcp import Context
+
+
+def get_mcp_dash_request_id(
+    ctx: Context,
+    **kwargs: Any,
+) -> str:
+    """
+    Return the MCP dash request ID.
+    Args:
+        ctx: MCP CONTEXT
+        **kwargs (Any): Additional keyword arguments.
+
+    Returns:
+        dash request ID.
+    """
+    dashscope_request_id = None
+    if ctx and ctx.request_context:
+        http_request = ctx.request_context.request
+        if http_request:
+            headers = dict(http_request.headers.items())
+            if headers:
+                dashscope_request_id = headers.get("dashscope_request_id")
+
+    if dashscope_request_id and dashscope_request_id.strip():
+        dashscope_request_id = dashscope_request_id.strip()
+    else:
+        dashscope_request_id = str(uuid.uuid4())
+    return dashscope_request_id

agentscope_runtime/version.py

@@ -1,2 +1,2 @@
 # -*- coding: utf-8 -*-
-__version__ = "v0.2.0b1"
+__version__ = "v1.0.0"