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

polaris_ai/__init__.py

@@ -0,0 +1,34 @@
+# 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 extensions for Python AI frameworks."""
+
+from ._version import __version__
+from .log import (
+    PolarisAILogger,
+    PolarisAILogOptions,
+    configure_ai_logging,
+    get_ai_logger,
+    use_ai_logger,
+)
+
+__all__ = [
+    "PolarisAILogOptions",
+    "PolarisAILogger",
+    "configure_ai_logging",
+    "get_ai_logger",
+    "use_ai_logger",
+    "__version__",
+]

polaris_ai/_async.py

@@ -0,0 +1,39 @@
+# 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.
+
+"""Finish worker operations before returning cancellation to callers."""
+
+import asyncio
+
+
+async def drain_task(task):
+    while not task.done():
+        try:
+            await asyncio.shield(task)
+        except asyncio.CancelledError:
+            continue
+        except Exception:
+            break
+    if not task.cancelled():
+        task.exception()
+
+
+async def await_completion(operation):
+    task = asyncio.create_task(operation)
+    try:
+        return await asyncio.shield(task)
+    except asyncio.CancelledError:
+        await drain_task(task)
+        raise

polaris_ai/_version.py

@@ -0,0 +1,18 @@
+# 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.
+
+"""The version of polaris-python-ai."""
+
+__version__ = "0.0.1.dev2"

polaris_ai/a2a/__init__.py

@@ -0,0 +1,31 @@
+# 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.
+
+"""Framework-independent A2A registration and card discovery."""
+
+from polaris_ai._version import __version__
+from .registry import A2ATransport, PolarisA2ARegistrar, RegisteredInstance
+from .resolver import AgentCardNotFoundError, PolarisA2AResolver
+from .server import PolarisA2AServerRegistration
+
+__all__ = [
+    "A2ATransport",
+    "PolarisA2ARegistrar",
+    "RegisteredInstance",
+    "AgentCardNotFoundError",
+    "PolarisA2AResolver",
+    "PolarisA2AServerRegistration",
+    "__version__",
+]

polaris_ai/a2a/_card.py

@@ -0,0 +1,44 @@
+# 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.
+
+"""Use the official A2A SDK's protobuf JSON and legacy card compatibility."""
+
+import json
+
+try:
+    from a2a.client.card_resolver import parse_agent_card
+    from a2a.types import AgentCard
+except ImportError as error:
+    raise ImportError(
+        "A2A integration requires pip install 'polaris-python-ai[a2a]'."
+    ) from error
+
+from google.protobuf.json_format import MessageToDict
+
+META_AGENT_NAME = "ai-agent-name"
+META_AGENT_CARD = "a2a.agent.card"
+META_TRANSPORT = "a2a.transport"
+META_PATH = "a2a.path"
+MAX_METADATA_VALUE_CHARS = 65535
+
+
+def encode_card(card: AgentCard) -> str:
+    if not isinstance(card, AgentCard) or not card.name.strip():
+        raise ValueError("A native a2a.types.AgentCard with a name is required")
+    return json.dumps(MessageToDict(card), ensure_ascii=False, separators=(",", ":"))
+
+
+def decode_card(payload: str) -> AgentCard:
+    return parse_agent_card(json.loads(payload))

polaris_ai/a2a/registry.py

@@ -0,0 +1,240 @@
+# 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.
+
+"""Register A2A transports using the public Polaris Provider API."""
+
+import asyncio
+import concurrent.futures
+from collections.abc import Sequence
+from dataclasses import dataclass
+from threading import Event, RLock
+
+from polaris.api.provider import ProviderAPI
+from polaris.pkg.model.server import InstanceDeRegisterRequest, InstanceRegisterRequest
+
+from polaris_ai._async import await_completion, drain_task
+from polaris_ai.log import get_ai_logger
+
+from ._card import (
+    MAX_METADATA_VALUE_CHARS,
+    META_AGENT_CARD,
+    META_AGENT_NAME,
+    META_PATH,
+    META_TRANSPORT,
+    AgentCard,
+    encode_card,
+)
+
+
+@dataclass(frozen=True)
+class A2ATransport:
+    """An already listening A2A endpoint; host is the explicit business instance address."""
+
+    host: str
+    port: int
+    transport: str = "JSONRPC"
+    path: str = "/"
+    support_tls: bool = False
+
+
+@dataclass(frozen=True)
+class RegisteredInstance:
+    service: str
+    host: str
+    port: int
+    instance_id: str
+
+
+class PolarisA2ARegistrar:
+    """Track registrations and clean them up without closing the supplied SDKContext.
+
+    ``service_token`` is used for register, heartbeat and deregistration and is
+    distinct from the connector's authentication token. It may be empty for
+    registries without service authentication. Heartbeat and instance-not-found
+    re-registration are managed by ProviderAPI.register_instance.
+    """
+
+    def __init__(self, sdk_context, namespace: str = "default", *, service_token: str = "",
+                 ttl: int = 5, logger=None):
+        if sdk_context is None or sdk_context.is_destroy():
+            raise ValueError("An active SDKContext is required")
+        if not isinstance(service_token, str):
+            raise ValueError("service_token must be a string")
+        if type(ttl) is not int or not 0 < ttl <= 0xFFFFFFFF:
+            raise ValueError("ttl must be a positive uint32 number of seconds")
+        self.namespace = namespace.strip() or "default"
+        self.ttl = ttl
+        self.service_token = service_token.strip()
+        self.logger = logger if logger is not None else get_ai_logger("a2a")
+        self._context = sdk_context
+        self._provider = ProviderAPI(sdk_context)
+        self._registered: list[RegisteredInstance] = []
+        self._lock = RLock()
+        self._closed = False
+
+    @property
+    def registered(self) -> tuple[RegisteredInstance, ...]:
+        with self._lock:
+            return tuple(self._registered)
+
+    def _check_open(self):
+        if self._closed or self._context.is_destroy():
+            raise RuntimeError("A2A registrar or SDKContext is closed")
+
+    def register(self, agent_card: AgentCard, transports: Sequence[A2ATransport]):
+        """Register one instance per transport; roll back successes on partial failure."""
+        return self._register(agent_card, tuple(transports))
+
+    def _register(self, agent_card, transports, cancelled: Event | None = None):
+        with self._lock:
+            self._check_open()
+            requests = self._requests(agent_card, transports)
+            self.logger.debug(
+                "Registering A2A agent: namespace=%s service=%s transports=%d",
+                self.namespace, agent_card.name, len(requests),
+            )
+            created = []
+            try:
+                for request in requests:
+                    if cancelled is not None and cancelled.is_set():
+                        raise concurrent.futures.CancelledError()
+                    response = self._provider.register_instance(request)
+                    instance = RegisteredInstance(
+                        request.service, request.host, request.port, response.instance_id,
+                    )
+                    self._registered.append(instance)
+                    created.append(instance)
+                    self.logger.debug(
+                        "Registered A2A transport: service=%s host=%s port=%s instanceId=%s",
+                        instance.service, instance.host, instance.port, instance.instance_id,
+                    )
+                    if cancelled is not None and cancelled.is_set():
+                        raise concurrent.futures.CancelledError()
+                self.logger.info(
+                    "Registered A2A agent %s/%s with %d transport(s)",
+                    self.namespace, agent_card.name, len(created),
+                )
+                return tuple(created)
+            except BaseException:
+                self.logger.warning_throttled(
+                    ("register-rollback", agent_card.name),
+                    "Rolling back %d A2A transport registration(s) for %s/%s",
+                    len(created), self.namespace, agent_card.name,
+                )
+                self._deregister(created)
+                raise
+
+    def _requests(self, card, transports):
+        payload = encode_card(card)
+        if not transports:
+            raise ValueError("transports must not be empty")
+        requests = []
+        for transport in transports:
+            if not isinstance(transport, A2ATransport):
+                raise ValueError("transports must contain A2ATransport objects")
+            if (not isinstance(transport.host, str) or not transport.host.strip()
+                    or type(transport.port) is not int or not 0 < transport.port < 65536
+                    or type(transport.support_tls) is not bool
+                    or not isinstance(transport.transport, str) or not transport.transport.strip()
+                    or not isinstance(transport.path, str) or not transport.path.startswith("/")):
+                raise ValueError("A2A transport requires a host, valid port, binding and absolute path")
+            metadata = {
+                META_AGENT_NAME: card.name,
+                META_AGENT_CARD: payload,
+                META_TRANSPORT: transport.transport,
+                META_PATH: transport.path,
+            }
+            for name, value in metadata.items():
+                # Java String.length counts UTF-16 code units, including supplementary characters.
+                if len(value.encode("utf-16-le")) // 2 > MAX_METADATA_VALUE_CHARS:
+                    raise ValueError(f"Metadata {name!r} exceeds the 65535 character limit")
+            request = InstanceRegisterRequest(
+                namespace=self.namespace, service=card.name, service_token=self.service_token,
+                host=transport.host, port=transport.port, version=card.version,
+                protocol="https" if transport.support_tls else "http",
+                metadata=metadata, ttl=self.ttl, auto_heartbeat=True,
+            )
+            request.validate()
+            requests.append(request)
+        return requests
+
+    def _deregister(self, instances):
+        for instance in instances:
+            request = InstanceDeRegisterRequest(
+                namespace=self.namespace, service=instance.service, service_token=self.service_token,
+                host=instance.host, port=instance.port, instance_id=instance.instance_id,
+            )
+            try:
+                self._provider.deregister_instance(request)
+            except Exception as error:
+                self.logger.warning_throttled(
+                    ("deregister", instance.service, instance.host, instance.port),
+                    "Failed to deregister A2A agent %s at %s:%s: %s",
+                    instance.service, instance.host, instance.port, error,
+                )
+            else:
+                self.logger.debug(
+                    "Deregistered A2A transport: service=%s host=%s port=%s instanceId=%s",
+                    instance.service, instance.host, instance.port, instance.instance_id,
+                )
+                self._registered = [item for item in self._registered if item is not instance]
+
+    async def async_register(self, agent_card: AgentCard, transports: Sequence[A2ATransport]):
+        """Run SDK sync calls off the caller loop; cancellation waits for rollback."""
+        cancelled = Event()
+        task = asyncio.create_task(asyncio.to_thread(
+            self._register, agent_card, tuple(transports), cancelled,
+        ))
+        try:
+            return await asyncio.shield(task)
+        except asyncio.CancelledError:
+            cancelled.set()
+            await drain_task(task)
+            if not task.cancelled() and task.exception() is None:
+                # Cancellation can arrive after the worker's last check, before delivery.
+                cleanup = asyncio.create_task(asyncio.to_thread(self._rollback, task.result()))
+                await drain_task(cleanup)
+            raise
+
+    def _rollback(self, instances):
+        with self._lock:
+            self._deregister(instances)
+
+    def close(self):
+        """Retry every tracked deregistration; failures remain in ``registered``."""
+        with self._lock:
+            self._closed = True
+            self.logger.debug(
+                "Closing A2A registrar: namespace=%s remaining=%d",
+                self.namespace, len(self._registered),
+            )
+            self._deregister(tuple(self._registered))
+
+    async def async_close(self):
+        await await_completion(asyncio.to_thread(self.close))
+
+    def __enter__(self):
+        self._check_open()
+        return self
+
+    def __exit__(self, *_):
+        self.close()
+
+    async def __aenter__(self):
+        self._check_open()
+        return self
+
+    async def __aexit__(self, *_):
+        await self.async_close()

polaris_ai/a2a/resolver.py

@@ -0,0 +1,198 @@
+# 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.
+
+"""Sticky AgentCard resolution using public Polaris discovery APIs."""
+
+import asyncio
+from dataclasses import dataclass
+from threading import RLock
+
+from google.protobuf.json_format import ParseError
+from polaris.api.consumer import ConsumerAPI
+from polaris.pkg.model.error import SDKError
+from polaris.pkg.model.service import GetHealthyInstancesRequest
+
+from polaris_ai._async import await_completion
+from polaris_ai.log import get_ai_logger
+
+from ._card import META_AGENT_CARD, AgentCard, decode_card
+
+
+class AgentCardNotFoundError(LookupError):
+    """Discovery succeeded but no healthy instance supplied a usable matching AgentCard."""
+
+
+@dataclass(frozen=True)
+class _CachedCard:
+    host: str
+    port: int
+    card: AgentCard
+
+
+class PolarisA2AResolver:
+    """Resolve native AgentCards, retaining the selected healthy host and port.
+
+    Every lookup refreshes discovery and decodes the selected instance's metadata.
+    A Polaris SDKError can reuse the last successful card; a successful empty
+    discovery raises AgentCardNotFoundError. Card URLs are preserved as registered.
+    """
+
+    def __init__(self, sdk_context, namespace: str = "default", *, timeout: int = 0,
+                 logger=None):
+        if sdk_context is None or sdk_context.is_destroy():
+            raise ValueError("An active SDKContext is required")
+        if type(timeout) is not int or timeout < 0:
+            raise ValueError("timeout must be a nonnegative integer")
+        self.namespace = namespace.strip() or "default"
+        self.timeout = timeout
+        self.logger = logger if logger is not None else get_ai_logger("a2a")
+        self._context = sdk_context
+        self._consumer = ConsumerAPI(sdk_context)
+        self._cache: dict[str, _CachedCard] = {}
+        self._lock = RLock()
+        self._closed = False
+
+    def _check_open(self):
+        if self._closed or self._context.is_destroy():
+            raise RuntimeError("A2A resolver or SDKContext is closed")
+
+    def resolve(self, agent_name: str) -> AgentCard:
+        """Return the latest card for the selected instance or select the first usable peer."""
+        if not isinstance(agent_name, str) or not agent_name.strip():
+            raise ValueError("agent_name is required")
+        agent_name = agent_name.strip()
+        with self._lock:
+            self._check_open()
+            cached = self._cache.get(agent_name)
+            self.logger.debug(
+                "Resolving A2A agent: namespace=%s service=%s sticky=%s",
+                self.namespace, agent_name, cached is not None,
+            )
+            try:
+                response = self._consumer.get_healthy_instances(GetHealthyInstancesRequest(
+                    namespace=self.namespace, service=agent_name, timeout=self.timeout,
+                    exclude_offline_service=True,
+                ))
+            except SDKError:
+                if cached is None:
+                    raise
+                self.logger.warning_throttled(
+                    ("refresh", agent_name),
+                    "Failed to refresh agent %s; reusing %s:%s",
+                    agent_name, cached.host, cached.port,
+                )
+                return cached.card
+            instances = [] if response is None else response.get_instances()
+            self.logger.debug(
+                "Discovered A2A candidates: namespace=%s service=%s count=%d",
+                self.namespace, agent_name, len(instances),
+            )
+            instances = [i for i in instances if i.is_healthy() and not i.is_isolated()]
+            if cached is not None:
+                sticky = [i for i in instances
+                          if (i.get_host(), i.get_port()) == (cached.host, cached.port)]
+                instances = sticky + [i for i in instances if i not in sticky]
+            for instance in instances:
+                payload = instance.get_metadata().get(META_AGENT_CARD, "")
+                try:
+                    card = decode_card(payload)
+                except (ValueError, TypeError, ParseError) as error:
+                    self.logger.warning_throttled(
+                        ("invalid-card", agent_name, instance.get_host(), instance.get_port()),
+                        "Skipping invalid AgentCard for %s at %s:%s: %s",
+                        agent_name, instance.get_host(), instance.get_port(), error,
+                    )
+                    continue
+                if card.name != agent_name:
+                    self.logger.warning_throttled(
+                        ("wrong-name", agent_name, instance.get_host(), instance.get_port()),
+                        "Skipping AgentCard %r while resolving %r", card.name, agent_name,
+                    )
+                    continue
+                if not any(i.url.strip() and i.protocol_binding.strip()
+                           for i in card.supported_interfaces):
+                    self.logger.warning_throttled(
+                        ("no-endpoint", agent_name, instance.get_host(), instance.get_port()),
+                        "Skipping endpoint-less AgentCard for %s", agent_name,
+                    )
+                    continue
+                if cached is not None and (
+                    cached.host != instance.get_host()
+                    or cached.port != instance.get_port()
+                ):
+                    self.logger.info(
+                        "Switched A2A agent %s/%s from %s:%s to %s:%s",
+                        self.namespace, agent_name, cached.host, cached.port,
+                        instance.get_host(), instance.get_port(),
+                    )
+                else:
+                    self.logger.debug(
+                        "Selected A2A instance: service=%s host=%s port=%s",
+                        agent_name, instance.get_host(), instance.get_port(),
+                    )
+                self._cache[agent_name] = _CachedCard(instance.get_host(), instance.get_port(), card)
+                return card
+            self.logger.debug(
+                "No usable A2A AgentCard: namespace=%s service=%s",
+                self.namespace, agent_name,
+            )
+            raise AgentCardNotFoundError(f"No usable AgentCard for {self.namespace}/{agent_name}")
+
+    async def async_resolve(self, agent_name: str) -> AgentCard:
+        return await await_completion(asyncio.to_thread(self.resolve, agent_name))
+
+    def invalidate(self, agent_name: str):
+        with self._lock:
+            self.logger.debug("Invalidating A2A card: service=%s", agent_name)
+            self._cache.pop(agent_name, None)
+
+    def invalidate_all(self):
+        with self._lock:
+            self.logger.debug("Invalidating all A2A cards: count=%d", len(self._cache))
+            self._cache.clear()
+
+    @property
+    def consumer(self):
+        return self._consumer
+
+    def selected_endpoint(self, agent_name: str):
+        """Return the sticky instance host/port for an already-resolved Agent."""
+        with self._lock:
+            cached = self._cache.get(agent_name)
+            return None if cached is None else (cached.host, cached.port)
+
+    def close(self):
+        """Clear local cards while leaving the borrowed context available."""
+        with self._lock:
+            self._closed = True
+            self.logger.debug("Closing A2A resolver: cached=%d", len(self._cache))
+            self._cache.clear()
+
+    async def async_close(self):
+        await await_completion(asyncio.to_thread(self.close))
+
+    def __enter__(self):
+        self._check_open()
+        return self
+
+    def __exit__(self, *_):
+        self.close()
+
+    async def __aenter__(self):
+        self._check_open()
+        return self
+
+    async def __aexit__(self, *_):
+        await self.async_close()