polaris-python-ai 0.0.1.dev2__py3-none-any.whl

polaris_ai/agentscope/mcp.py

@@ -0,0 +1,598 @@
+# Tencent is pleased to support the open source community by making polaris-python-ai available.
+#
+# Copyright (C) 2026 Tencent. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+"""Polaris discovery and contract filtering for AgentScope MCP clients."""
+
+from __future__ import annotations
+
+import asyncio
+import re
+import time
+import types
+
+from pydantic import PrivateAttr
+
+try:
+    from agentscope.mcp import HttpMCPConfig, MCPClient
+    from agentscope.message import ToolResultState
+except ImportError as error:
+    raise ImportError(
+        "Polaris MCP support requires pip install 'polaris-python-ai[mcp]'."
+    ) from error
+
+from polaris.api.consumer import ConsumerAPI
+from polaris.pkg.model.contract import (
+    GetServiceContractRequest,
+    ServiceFeatureStatus,
+    ServiceFeatureType,
+)
+from polaris.pkg.model.service import GetHealthyInstancesRequest, GetInstancesRequest, InstanceType
+
+from polaris_ai._async import await_completion
+from polaris_ai.log import get_ai_logger
+from polaris_ai.metrics import PolarisCallReporter
+
+MCP_STREAMABLE_HTTP = "mcp-streamable-http"
+MCP_SSE = "mcp-sse"
+MCP_PROTOCOLS = (MCP_STREAMABLE_HTTP, MCP_SSE)
+MCP_SERVER_NAME = "mcp-server-name"
+MCP_PROTOCOL_VERSION = "mcp-protocol-version"
+MCP_ENDPOINT_PATH = "mcp-endpoint-path"
+MCP_SCHEME = "mcp-scheme"
+MCPPoolKey = tuple[str, str, str, str, tuple[str, ...]]
+
+
+def _client_name(service: str) -> str:
+    name = re.sub(r"[^a-zA-Z0-9_-]", "-", service.strip())
+    return name or "polaris-mcp"
+
+
+def _versions(value: str) -> set[str]:
+    return {item.strip() for item in value.split(",") if item.strip()}
+
+
+def mcp_endpoint(instance, scheme: str = "") -> tuple[str, int, str]:
+    """Build the AgentScope MCP URL, port and native protocol for an instance."""
+    protocol = instance.get_protocol()
+    if protocol not in MCP_PROTOCOLS:
+        raise ValueError(f"Unsupported MCP protocol {protocol!r}")
+    selected_scheme = scheme or instance.get_metadata().get(MCP_SCHEME, "http")
+    if selected_scheme not in ("http", "https"):
+        raise ValueError("MCP scheme must be http or https")
+    path = instance.get_metadata().get(MCP_ENDPOINT_PATH) or (
+        "/sse" if protocol == MCP_SSE else "/mcp"
+    )
+    if not path.startswith("/") or path.startswith("//") or "?" in path or "#" in path:
+        raise ValueError("MCP endpoint path must be an absolute path without query or fragment")
+    host = instance.get_host()
+    if ":" in host and not host.startswith("["):
+        host = f"[{host}]"
+    return f"{selected_scheme}://{host}:{instance.get_port()}{path}", instance.get_port(), protocol
+
+
+class PolarisMCPClient(MCPClient):
+    """AgentScope MCP client that reports successful and failed Tool calls."""
+
+    _polaris_reporter = PrivateAttr(default=None)
+    _polaris_service = PrivateAttr(default="")
+    _polaris_host = PrivateAttr(default="")
+    _polaris_port = PrivateAttr(default=0)
+
+    def __init__(self, *, reporter=None, service: str, host: str, port: int, **data):
+        super().__init__(**data)
+        self._polaris_reporter = reporter
+        self._polaris_service = service
+        self._polaris_host = host
+        self._polaris_port = port
+
+    async def get_tool(self, name: str):
+        tool = await super().get_tool(name)
+        reporter = self._polaris_reporter
+        if reporter is None:
+            return tool
+        original = tool.call
+
+        async def report(*, started_at, success, timeout=False, ret_code=None):
+            await reporter.report(
+                service=self._polaris_service,
+                host=self._polaris_host,
+                port=self._polaris_port,
+                method=name,
+                instance_type=InstanceType.MCP,
+                started_at=started_at,
+                success=success,
+                ret_code=ret_code,
+                timeout=timeout,
+            )
+
+        async def reported_call(_tool, **kwargs):
+            started_at = time.monotonic()
+            try:
+                result = await original(**kwargs)
+            except BaseException as error:
+                await report(
+                    started_at=started_at,
+                    success=False,
+                    timeout=isinstance(error, (asyncio.TimeoutError, TimeoutError)),
+                )
+                raise
+
+            if isinstance(result, types.AsyncGeneratorType):
+                async def reported_stream():
+                    final_state = None
+                    completed = False
+                    try:
+                        async for chunk in result:
+                            if getattr(chunk, "is_last", False):
+                                final_state = getattr(chunk, "state", None)
+                                completed = True
+                            yield chunk
+                    except BaseException as error:
+                        await report(
+                            started_at=started_at,
+                            success=False,
+                            timeout=isinstance(error, (asyncio.TimeoutError, TimeoutError)),
+                        )
+                        raise
+                    else:
+                        success = completed and final_state not in (
+                            ToolResultState.ERROR,
+                            ToolResultState.INTERRUPTED,
+                            ToolResultState.DENIED,
+                        )
+                        await report(
+                            started_at=started_at,
+                            success=success,
+                            ret_code=None if success else -1,
+                        )
+                    finally:
+                        await await_completion(result.aclose())
+
+                return reported_stream()
+
+            success = getattr(result, "state", None) not in (
+                ToolResultState.ERROR,
+                ToolResultState.INTERRUPTED,
+                ToolResultState.DENIED,
+            )
+            await report(
+                started_at=started_at,
+                success=success,
+                ret_code=None if success else -1,
+            )
+            return result
+
+        tool.call = types.MethodType(reported_call, tool)
+        return tool
+
+
+class PolarisMCPResolver:
+    """Resolve MCP instances and compose native AgentScope MCP clients."""
+
+    def __init__(
+        self,
+        sdk_context,
+        namespace: str = "default",
+        *,
+        client_factory=PolarisMCPClient,
+        caller_service: str = "",
+        labels: str = "",
+        protocol_version: str = "",
+        scheme: str = "",
+        timeout: float = 30.0,
+        filter_contract: bool = True,
+        headers: dict[str, str] | None = None,
+        logger=None,
+        metrics_logger=None,
+    ):
+        if sdk_context is None or sdk_context.is_destroy():
+            raise ValueError("An active SDKContext is required")
+        self.namespace = namespace.strip() or "default"
+        self.consumer = ConsumerAPI(sdk_context)
+        self.client_factory = client_factory
+        self.caller_service = caller_service
+        self.labels = labels
+        self.protocol_version = protocol_version
+        self.scheme = scheme
+        self.timeout = timeout
+        self.filter_contract = filter_contract
+        self.headers = dict(headers or {})
+        self.logger = logger if logger is not None else get_ai_logger("mcp")
+        self.metrics_logger = (
+            metrics_logger if metrics_logger is not None else get_ai_logger("metrics")
+        )
+
+    async def async_candidates(self, service: str, *, protocol_version: str | None = None):
+        if not isinstance(service, str) or not service.strip():
+            raise ValueError("MCP service is required")
+        self.logger.debug(
+            "Discovering MCP candidates: namespace=%s service=%s protocolVersion=%s",
+            self.namespace, service,
+            self.protocol_version if protocol_version is None else protocol_version,
+        )
+        response = await self.consumer.async_get_healthy_instances(
+            GetHealthyInstancesRequest(
+                namespace=self.namespace,
+                service=service.strip(),
+                exclude_offline_service=True,
+            )
+        )
+        result = []
+        requested = _versions(
+            self.protocol_version if protocol_version is None else protocol_version
+        )
+        for instance in [] if response is None else response.get_instances():
+            if not instance.is_healthy() or instance.is_isolated():
+                continue
+            if instance.get_protocol() not in MCP_PROTOCOLS:
+                continue
+            versions = _versions(instance.get_metadata().get(MCP_PROTOCOL_VERSION, ""))
+            if requested and versions and not requested.intersection(versions):
+                continue
+            result.append(instance)
+        self.logger.debug(
+            "Discovered MCP candidates: service=%s healthy=%d",
+            service, len(result),
+        )
+        return result
+
+    async def async_disabled_tools(self, service: str, *, protocol: str = "", version: str = ""):
+        try:
+            contract = await self.consumer.async_get_service_contract(
+                GetServiceContractRequest(
+                    namespace=self.namespace,
+                    service=service,
+                    protocol=protocol,
+                    version=version,
+                )
+            )
+        except Exception as error:
+            self.logger.warning_throttled(
+                ("contract", service),
+                "MCP contract unavailable for %s; skipping Tool filtering: %s",
+                service,
+                error,
+            )
+            return []
+        if contract is None:
+            return []
+        disabled = [
+            feature.name
+            for feature in contract.service_features
+            if feature.type == ServiceFeatureType.MCP_TOOL
+            and feature.status == ServiceFeatureStatus.DISABLED
+        ]
+        self.logger.debug(
+            "Loaded MCP contract filter: service=%s protocol=%s version=%s disabledTools=%d",
+            service, protocol, version, len(disabled),
+        )
+        return disabled
+
+    async def async_create_client(
+        self,
+        service: str,
+        *,
+        protocol_version: str | None = None,
+        scheme: str | None = None,
+        headers: dict[str, str] | None = None,
+        timeout: float | None = None,
+        filter_contract: bool | None = None,
+    ):
+        instances = await self.async_candidates(
+            service,
+            protocol_version=(
+                self.protocol_version if protocol_version is None else protocol_version
+            ),
+        )
+        if not instances:
+            raise LookupError(f"No healthy MCP instance for {self.namespace}/{service}")
+        self.logger.debug(
+            "Selecting MCP instance: service=%s host=%s port=%s",
+            service, instances[0].get_host(), instances[0].get_port(),
+        )
+        return await self.async_client_for_instance(
+            service,
+            instances[0],
+            scheme=self.scheme if scheme is None else scheme,
+            headers=self.headers if headers is None else headers,
+            timeout=self.timeout if timeout is None else timeout,
+            filter_contract=(
+                self.filter_contract if filter_contract is None else filter_contract
+            ),
+        )
+
+    async def async_client_for_instance(
+        self,
+        service: str,
+        instance,
+        *,
+        scheme: str = "",
+        headers: dict[str, str] | None = None,
+        timeout: float | None = 30.0,
+        filter_contract: bool = True,
+        disabled_tools: list[str] | None = None,
+    ):
+        url, port, protocol = mcp_endpoint(instance, scheme)
+        self.logger.debug(
+            "Creating MCP client: service=%s endpoint=%s protocol=%s stateful=%s",
+            service, url, protocol, protocol == MCP_SSE,
+        )
+        if disabled_tools is None:
+            if filter_contract:
+                disabled = await self.async_disabled_tools(
+                    service,
+                    protocol=protocol,
+                    version=instance.get_version(),
+                )
+            else:
+                disabled = []
+        else:
+            disabled = list(disabled_tools)
+        reporter = PolarisCallReporter(
+            self.consumer,
+            self.namespace,
+            caller_service=self.caller_service,
+            labels=self.labels,
+            logger=self.metrics_logger,
+        )
+        client = self.client_factory(
+            name=_client_name(service),
+            is_stateful=protocol == MCP_SSE,
+            mcp_config=HttpMCPConfig(url=url, headers=headers, timeout=timeout),
+            disable_tools=disabled or None,
+            reporter=reporter,
+            service=service,
+            host=instance.get_host(),
+            port=port,
+        )
+        try:
+            if client.is_stateful:
+                await client.connect()
+        except BaseException:
+            try:
+                await await_completion(client.close())
+            except BaseException:
+                pass
+            raise
+        return client
+
+
+class PolarisMCPClientPool:
+    """Maintain one MCP client per watched healthy Polaris instance."""
+
+    def __init__(self, resolver: PolarisMCPResolver, service: str, **client_options):
+        self.resolver = resolver
+        self.service = service
+        options = {
+            "scheme": resolver.scheme,
+            "headers": resolver.headers,
+            "timeout": resolver.timeout,
+            "filter_contract": resolver.filter_contract,
+            **client_options,
+        }
+        configured_protocol = options.pop("protocol_version", resolver.protocol_version)
+        self.protocol_version = (
+            resolver.protocol_version if configured_protocol is None else configured_protocol
+        )
+        self.client_options = options
+        self._clients: dict[MCPPoolKey, MCPClient] = {}
+        self._watch = None
+        self._refresh_event = asyncio.Event()
+        self._task = None
+        self._lock = asyncio.Lock()
+        self._closed = False
+        self._cursor = 0
+
+    @property
+    def clients(self):
+        return tuple(self._clients.values())
+
+    async def start(self):
+        """Load clients once and subscribe to service cache lifecycle events."""
+        if self._closed:
+            raise RuntimeError("MCP client pool is closed")
+        if self._task is not None:
+            return self
+        loop = asyncio.get_running_loop()
+
+        def changed(_event):
+            loop.call_soon_threadsafe(self._refresh_event.set)
+
+        self._watch = await self.resolver.consumer.async_watch_service(
+            GetInstancesRequest(
+                namespace=self.resolver.namespace,
+                service=self.service,
+                exclude_offline_service=True,
+            ),
+            changed,
+        )
+        try:
+            await self.refresh()
+            self._task = asyncio.create_task(self._run(), name=f"polaris-mcp-pool-{self.service}")
+            self.resolver.logger.info(
+                "Started MCP client pool: service=%s clients=%d",
+                self.service, len(self._clients),
+            )
+        except BaseException:
+            try:
+                await await_completion(self._cleanup_start_failure())
+            except asyncio.CancelledError:
+                pass
+            raise
+        return self
+
+    async def _cleanup_start_failure(self):
+        clients, self._clients = tuple(self._clients.values()), {}
+        watch, self._watch = self._watch, None
+        try:
+            await self._close_clients(clients)
+        finally:
+            if watch is not None:
+                await self.resolver.consumer.async_unwatch_service(watch)
+
+    async def _run(self):
+        while not self._closed:
+            await self._refresh_event.wait()
+            if self._closed:
+                return
+            self._refresh_event.clear()
+            try:
+                await self.refresh()
+            except Exception as error:
+                self.resolver.logger.warning_throttled(
+                    ("pool-refresh", self.service),
+                    "Failed to refresh MCP pool for %s: %s", self.service, error,
+                )
+
+    async def refresh(self):
+        """Reconcile current healthy instances with the native client pool."""
+        async with self._lock:
+            if self._closed:
+                return
+            self.resolver.logger.debug(
+                "Refreshing MCP client pool: service=%s currentClients=%d",
+                self.service, len(self._clients),
+            )
+            instances = await self.resolver.async_candidates(
+                self.service,
+                protocol_version=self.protocol_version,
+            )
+            desired: dict[MCPPoolKey, tuple[object, str, list[str]]] = {}
+            for instance in instances:
+                try:
+                    url, _port, protocol = mcp_endpoint(
+                        instance, self.client_options.get("scheme", "")
+                    )
+                except ValueError as error:
+                    self.resolver.logger.warning_throttled(
+                        ("invalid-instance", self.service, instance.get_id()),
+                        "Skipping invalid MCP instance %s: %s", instance.get_id(), error,
+                    )
+                    continue
+                disabled = (
+                    await self.resolver.async_disabled_tools(
+                        self.service,
+                        protocol=protocol,
+                        version=instance.get_version(),
+                    )
+                    if self.client_options["filter_contract"] else []
+                )
+                key: MCPPoolKey = (
+                    url,
+                    protocol,
+                    instance.get_version(),
+                    instance.get_metadata().get(MCP_PROTOCOL_VERSION, ""),
+                    tuple(sorted(set(disabled))),
+                )
+                desired[key] = (instance, url, disabled)
+
+            next_clients: dict[MCPPoolKey, MCPClient] = {}
+            created = []
+            try:
+                for key, (instance, _url, disabled) in desired.items():
+                    client = self._clients.get(key)
+                    if client is None:
+                        client = await self.resolver.async_client_for_instance(
+                            self.service,
+                            instance,
+                            disabled_tools=disabled,
+                            **self.client_options,
+                        )
+                        created.append(client)
+                    next_clients[key] = client
+            except BaseException:
+                try:
+                    await self._close_clients(created)
+                except asyncio.CancelledError:
+                    pass
+                raise
+
+            if self._closed:
+                await self._close_clients(created)
+                self.resolver.logger.debug(
+                    "Discarded MCP clients created during close: service=%s count=%d",
+                    self.service, len(created),
+                )
+                return
+
+            previous_keys = set(self._clients)
+            next_keys = set(next_clients)
+            previous, self._clients = self._clients, next_clients
+            await self._close_clients(
+                tuple(client for key, client in previous.items() if key not in next_clients)
+            )
+            self.resolver.logger.debug(
+                "MCP client pool reconciled: service=%s added=%s removed=%s",
+                self.service, sorted(next_keys - previous_keys), sorted(previous_keys - next_keys),
+            )
+            if next_keys != previous_keys:
+                self.resolver.logger.info(
+                    "MCP client pool changed: service=%s before=%d after=%d added=%d removed=%d",
+                    self.service, len(previous_keys), len(next_keys),
+                    len(next_keys - previous_keys), len(previous_keys - next_keys),
+                )
+
+    async def _close_clients(self, clients):
+        async def close_all():
+            for client in clients:
+                try:
+                    await client.close()
+                except Exception as error:
+                    self.resolver.logger.warning_throttled(
+                        ("close-client", self.service, getattr(client, "name", "")),
+                        "Failed to close MCP client for %s: %s", self.service, error,
+                    )
+
+        await await_completion(close_all())
+
+    async def _close_resources(self, task):
+        if task is not None:
+            await task
+        async with self._lock:
+            if self._watch is not None:
+                watch = self._watch
+                await self.resolver.consumer.async_unwatch_service(watch)
+                if self._watch is watch:
+                    self._watch = None
+            clients, self._clients = tuple(self._clients.values()), {}
+            await self._close_clients(clients)
+
+    def choose_client(self):
+        """Return the next connected client using round-robin selection."""
+        clients = list(self._clients.values())
+        if not clients:
+            raise LookupError(f"No MCP client is available for {self.service}")
+        self._cursor %= len(clients)
+        client = clients[self._cursor]
+        self._cursor = (self._cursor + 1) % len(clients)
+        return client
+
+    async def close(self):
+        """Stop watch updates and close all owned MCP clients."""
+        self._closed = True
+        self._refresh_event.set()
+        task, self._task = self._task, None
+        self.resolver.logger.debug(
+            "Closing MCP client pool: service=%s clients=%d",
+            self.service, len(self._clients),
+        )
+        await await_completion(self._close_resources(task))
+        self.resolver.logger.info("Closed MCP client pool: service=%s", self.service)
+
+    async def __aenter__(self):
+        return await self.start()
+
+    async def __aexit__(self, *_):
+        await self.close()

polaris_ai/agentscope/middleware.py

@@ -0,0 +1,103 @@
+# Tencent is pleased to support the open source community by making polaris-python-ai available.
+#
+# Copyright (C) 2026 Tencent. All rights reserved.
+#
+# Licensed under the BSD 3-Clause License (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# https://opensource.org/licenses/BSD-3-Clause
+#
+# Unless required by applicable law or agreed to in writing, software distributed
+# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
+# CONDITIONS OF ANY KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations under the License.
+
+"""AgentScope middleware helpers for Polaris-backed agents."""
+
+from weakref import WeakKeyDictionary
+
+from agentscope.middleware import MiddlewareBase
+from agentscope.skill import SkillLoaderBase
+
+from polaris_ai.log import get_ai_logger
+
+_SHARED_SNAPSHOT_LOADERS = WeakKeyDictionary()
+
+
+def get_shared_snapshot_loader(loader):
+    """Return the snapshot associated with a loader, if middleware created it."""
+    return _SHARED_SNAPSHOT_LOADERS.get(loader)
+
+
+class _SnapshotSkillLoader(SkillLoaderBase):
+    """Expose the middleware's last successful Skill snapshot to a Toolkit."""
+
+    def __init__(self):
+        self.last_skills = ()
+
+    async def list_skills(self):
+        return list(self.last_skills)
+
+
+class PolarisDynamicSkillMiddleware(MiddlewareBase):
+    """Refresh Polaris Skills before each reply while reusing one Agent.
+
+    The loader applies its refresh interval, so replies within that interval do
+    not query Polaris again. Refresh failures are logged and the reply proceeds
+    with the loader's last successful state.
+    """
+
+    def __init__(self, loader: SkillLoaderBase, *, logger=None):
+        if not isinstance(loader, SkillLoaderBase):
+            raise TypeError("loader must implement SkillLoaderBase")
+        self.loader = loader
+        self.logger = logger if logger is not None else get_ai_logger("agentscope")
+        self.last_skills = ()
+        snapshot = get_shared_snapshot_loader(loader)
+        if snapshot is None:
+            snapshot = _SnapshotSkillLoader()
+            _SHARED_SNAPSHOT_LOADERS[loader] = snapshot
+        self._snapshot_loader = snapshot
+
+    async def on_reply(self, agent, input_kwargs, next_handler):
+        self.logger.debug("Refreshing Polaris Skills before Agent reply")
+        try:
+            self.last_skills = tuple(await self.loader.list_skills())
+        except Exception as error:
+            self.logger.warning_throttled(
+                ("refresh", id(self.loader)),
+                "Failed to refresh Polaris Skills before reply: %s", error,
+            )
+        else:
+            self._snapshot_loader.last_skills = self.last_skills
+            self.logger.debug(
+                "Refreshed Polaris Skill snapshot: count=%d names=%s",
+                len(self.last_skills), [skill.name for skill in self.last_skills],
+            )
+        self._sync_toolkit(agent)
+        async for event in next_handler():
+            yield event
+
+    def _sync_toolkit(self, agent):
+        toolkit = getattr(agent, "toolkit", None)
+        if toolkit is None:
+            self.logger.debug("Agent has no Toolkit; skipping Polaris Skill snapshot binding")
+            return
+        for group in getattr(toolkit, "tool_groups", ()):
+            if group.name != "basic":
+                continue
+            for index, item in enumerate(group.skills_or_loaders):
+                if item is self._snapshot_loader:
+                    self.logger.debug(
+                        "Polaris Skill snapshot already bound to Toolkit: category=%s",
+                        group.name,
+                    )
+                    return
+                if item is self.loader:
+                    group.skills_or_loaders[index] = self._snapshot_loader
+                    self.logger.debug(
+                        "Bound Polaris Skill snapshot to Toolkit: category=%s index=%d",
+                        group.name, index,
+                    )
+                    return