google-cloud-agentplatform 1.165.1.dev0__py3-none-any.whl

agentplatform/_genai/live.py

@@ -0,0 +1,64 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+#
+
+"""[Preview] Live API client."""
+
+import importlib
+import logging
+
+from typing import Optional, TYPE_CHECKING
+from types import ModuleType
+
+from google.genai import _api_module
+from google.genai import _common
+from google.genai._api_client import BaseApiClient
+
+logger = logging.getLogger("google_genai.live")
+
+if TYPE_CHECKING:
+    from agentplatform._genai import (
+        live_agent_engines as live_agent_engines_module,
+    )
+
+
+class AsyncLive(_api_module.BaseModule):
+    """[Preview] AsyncLive."""
+
+    def __init__(self, api_client: BaseApiClient):
+        super().__init__(api_client)
+        self._agent_engines: Optional[ModuleType] = None
+
+    @property
+    @_common.experimental_warning(
+        "The Vertex SDK GenAI agent engines module is experimental, "
+        "and may change in future versions."
+    )
+    def agent_engines(self) -> "live_agent_engines_module.AsyncLiveAgentEngines":
+        if self._agent_engines is None:
+            try:
+                # We need to lazy load the live_agent_engines module to handle
+                # the possibility of ImportError when dependencies are not
+                # installed.
+                self._agent_engines = importlib.import_module(
+                    ".live_agent_engines",
+                    __package__,
+                )
+            except ImportError as e:
+                raise ImportError(
+                    "The 'agent_engines' module requires 'additional packages'. "
+                    "Please install them using pip install "
+                    "google-cloud-aiplatform[agent_engines]"
+                ) from e
+        return self._agent_engines.AsyncLiveAgentEngines(self._api_client)  # type: ignore[no-any-return]

agentplatform/_genai/live_agent_engines.py

@@ -0,0 +1,179 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# 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.
+#
+
+"""Live AgentEngine API client."""
+
+import contextlib
+import json
+from typing import Any, AsyncIterator, Dict, Optional
+import google.auth
+
+from google.genai import _api_module
+from .types import QueryAgentEngineConfig, QueryAgentEngineConfigOrDict
+
+
+try:
+    from websockets.asyncio.client import ClientConnection
+    from websockets.asyncio.client import connect as ws_connect
+except ModuleNotFoundError:
+    # This try/except is for TAP, mypy complains about it which is why we have the type: ignore
+    from websockets.client import ClientConnection  # type: ignore
+    from websockets.client import connect as ws_connect  # type: ignore
+
+
+class AsyncLiveAgentEngineSession:
+    """AsyncLiveAgentEngineSession."""
+
+    def __init__(self, websocket: ClientConnection):
+        self._ws = websocket
+
+    async def send(self, query_input: Dict[str, Any]) -> None:
+        """Send a query input to the Agent.
+
+        Args:
+          query_input: A JSON serializable Python Dict to be send to the Agent.
+        """
+
+        try:
+            json_request = json.dumps({"bidi_stream_input": query_input})
+        except Exception as exc:
+            raise ValueError(
+                "Failed to encode query input to JSON in live_agent_engines: "
+                f"{str(query_input)}"
+            ) from exc
+        await self._ws.send(json_request)
+
+    async def receive(self) -> Any:
+        """Receive one response from the Agent.
+
+        Returns:
+          A response from the Agent.
+
+        Raises:
+          websockets.exceptions.ConnectionClosed: If the connection is closed.
+        """
+
+        response = await self._ws.recv()
+        try:
+            return json.loads(response)
+        except json.decoder.JSONDecodeError as exc:
+            raise ValueError(
+                "Failed to parse response to JSON in live_agent_engines: "
+                f"{str(response)}"
+            ) from exc
+
+    async def close(self) -> None:
+        """Close the connection."""
+        await self._ws.close()
+
+
+class AsyncLiveAgentEngines(_api_module.BaseModule):
+    """AsyncLiveAgentEngines.
+
+    Example usage:
+
+    .. code-block:: python
+
+      from pathlib import Path
+
+      from google import genai
+      from google.genai import types
+
+      class MyAgentEngine(client):
+        def bidi_stream_query(self, input_queue: asyncio.Queue):
+          while True:
+            input = await input_queue.get()
+            yield {"output": f"Agent received {input}!"}
+
+      client = agentplatform.Client(project="my-project", location="us-central1")
+      agent_engine = client.agent_engines.create(agent)
+
+      async with client.aio.live.agent_engines.connect(
+          agent_engine=agent_engine.api_resource.name,
+          setup={"class_method": "bidi_stream_query"},
+      ) as session:
+        await session.send(input={"input": "Hello world"})
+
+        response = await session.receive()
+        # {"output": "Agent received Hello world!"}
+        ...
+    """
+
+    @contextlib.asynccontextmanager
+    async def connect(
+        self,
+        *,
+        agent_engine: str,
+        config: Optional[QueryAgentEngineConfigOrDict] = None,
+    ) -> AsyncIterator[AsyncLiveAgentEngineSession]:
+        """Connect to the agent deployed to Agent Engine in a live (bidirectional streaming) session.
+
+        Args:
+          agent_engine: The resource name of the Agent Engine to use for the
+            live session.
+          config: The optional configuration for starting the live Agent Engine
+            session. Custom class_method and an optional initial input could be
+            provided. If no class_method is provided, the default class_method
+            "bidi_stream_query" will be used by the Agent Engine.
+
+        Yields:
+          An AsyncLiveAgentEngineSession object.
+        """
+        if isinstance(config, dict):
+            config = QueryAgentEngineConfig(**config)
+
+        agent_engine_resource_name = agent_engine
+        if not agent_engine_resource_name.startswith("projects/"):
+            agent_engine_resource_name = f"projects/{self._api_client.project}/locations/{self._api_client.location}/reasoningEngines/{agent_engine}"
+        request_dict = {"setup": {"name": agent_engine_resource_name}}
+        if config is not None and config.class_method:
+            request_dict["setup"]["class_method"] = config.class_method
+        if config is not None and config.input:
+            request_dict["setup"]["input"] = config.input  # type: ignore[assignment]
+
+        request = json.dumps(request_dict)
+
+        if not self._api_client._credentials:
+            # Get bearer token through Application Default Credentials.
+            creds, _ = google.auth.default(
+                scopes=["https://www.googleapis.com/auth/cloud-platform"]
+            )
+        else:
+            creds = self._api_client._credentials
+        # creds.valid is False, and creds.token is None
+        # Need to refresh credentials to populate those
+        if not (creds.token and creds.valid):
+            auth_req = google.auth.transport.requests.Request()
+            creds.refresh(auth_req)  # type: ignore[no-untyped-call]
+        bearer_token = creds.token
+
+        original_headers = self._api_client._http_options.headers
+        headers = original_headers.copy() if original_headers is not None else {}
+        headers["Authorization"] = f"Bearer {bearer_token}"
+
+        base_url = self._api_client._websocket_base_url()
+        if isinstance(base_url, bytes):
+            base_url = base_url.decode("utf-8")
+        uri = (
+            f"{base_url}/ws/google.cloud.aiplatform."
+            f"{self._api_client._http_options.api_version}"
+            ".ReasoningEngineExecutionService/BidiQueryReasoningEngine"
+        )
+
+        async with ws_connect(
+            uri, additional_headers=headers, **self._api_client._websocket_ssl_ctx
+        ) as ws:
+            await ws.send(request)
+            yield AsyncLiveAgentEngineSession(websocket=ws)