agent-starter-pack 0.15.6__py3-none-any.whl → 0.16.0__py3-none-any.whl

src/deployment_targets/cloud_run/{{cookiecutter.agent_directory}}/server.py

@@ -11,26 +11,29 @@
 # 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.
-{% if cookiecutter.agent_name == "live_api" %}
+{% if cookiecutter.agent_name == "adk_live" %}
 import asyncio
 import json
 import logging
 from collections.abc import Callable
 from pathlib import Path
-from typing import Any, Literal
 
 import backoff
 from fastapi import FastAPI, HTTPException, WebSocket
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
+from google.adk.agents.live_request_queue import LiveRequest, LiveRequestQueue
+from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
+from google.adk.memory.in_memory_memory_service import InMemoryMemoryService
+from google.adk.runners import Runner
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
 from google.cloud import logging as google_cloud_logging
-from google.genai import types
-from google.genai.types import LiveServerToolCall
-from pydantic import BaseModel
+from vertexai.agent_engines import _utils
 from websockets.exceptions import ConnectionClosedError
 
-from .agent import MODEL_ID, genai_client, live_connect_config, tool_functions
+from .agent import root_agent
+from .utils.typing import Feedback
 
 app = FastAPI()
 app.add_middleware(
@@ -56,117 +59,160 @@ logger = logging_client.logger(__name__)
 logging.basicConfig(level=logging.INFO)
 
 
-class GeminiSession:
-    """Manages bidirectional communication between a client and the Gemini model."""
+# Initialize ADK services
+session_service = InMemorySessionService()
+artifact_service = InMemoryArtifactService()
+memory_service = InMemoryMemoryService()
 
-    def __init__(
-        self, session: Any, websocket: WebSocket, tool_functions: dict[str, Callable]
-    ) -> None:
-        """Initialize the Gemini session.
+# Initialize ADK runner
+runner = Runner(
+    agent=root_agent,
+    session_service=session_service,
+    artifact_service=artifact_service,
+    memory_service=memory_service,
+    app_name="live-app",
+)
+
+
+class AgentSession:
+    """Manages bidirectional communication between a client and the agent."""
+
+    def __init__(self, websocket: WebSocket) -> None:
+        """Initialize the agent session.
 
         Args:
-            session: The Gemini session
             websocket: The client websocket connection
-            user_id: Unique identifier for this client
-            tool_functions: Dictionary of available tool functions
         """
-        self.session = session
         self.websocket = websocket
-        self.run_id = "n/a"
-        self.user_id = "n/a"
-        self.tool_functions = tool_functions
-        self._tool_tasks: list[asyncio.Task] = []
+        self.input_queue: asyncio.Queue[dict] = asyncio.Queue()
+        self.user_id: str | None = None
+        self.session_id: str | None = None
 
     async def receive_from_client(self) -> None:
-        """Listen for and process messages from the client.
-
-        Continuously receives messages and forwards audio data to Gemini.
-        Handles connection errors gracefully.
-        """
+        """Listen for messages from the client and put them in the queue."""
         while True:
             try:
-                data = await self.websocket.receive_json()
-
-                if isinstance(data, dict) and (
-                    "realtimeInput" in data or "clientContent" in data
-                ):
-                    await self.session._ws.send(json.dumps(data))
-                elif "setup" in data:
-                    self.run_id = data["setup"]["run_id"]
-                    self.user_id = data["setup"]["user_id"]
-                    logger.log_struct(
-                        {**data["setup"], "type": "setup"}, severity="INFO"
-                    )
+                message = await self.websocket.receive()
+
+                if "text" in message:
+                    data = json.loads(message["text"])
+
+                    if isinstance(data, dict):
+                        # Skip setup messages - they're for backend logging only
+                        if "setup" in data:
+                            logger.log_struct(
+                                {**data["setup"], "type": "setup"}, severity="INFO"
+                            )
+                            logging.info(
+                                "Received setup message (not forwarding to agent)"
+                            )
+                            continue
+
+                        # Forward message to agent engine
+                        await self.input_queue.put(data)
+                    else:
+                        logging.warning(
+                            f"Received unexpected JSON structure from client: {data}"
+                        )
+
+                elif "bytes" in message:
+                    # Handle binary data
+                    await self.input_queue.put({"binary_data": message["bytes"]})
+
                 else:
-                    logging.warning(f"Received unexpected input from client: {data}")
+                    logging.warning(
+                        f"Received unexpected message type from client: {message}"
+                    )
+
             except ConnectionClosedError as e:
-                logging.warning(f"Client {self.user_id} closed connection: {e}")
+                logging.warning(f"Client closed connection: {e}")
+                break
+            except json.JSONDecodeError as e:
+                logging.error(f"Error parsing JSON from client: {e}")
                 break
             except Exception as e:
-                logging.error(f"Error receiving from client {self.user_id}: {e!s}")
+                logging.error(f"Error receiving from client: {e!s}")
                 break
 
-    def _get_func(self, action_label: str | None) -> Callable | None:
-        """Get the tool function for a given action label."""
-        if action_label is None or action_label == "":
-            return None
-        return self.tool_functions.get(action_label)
-
-    async def _handle_tool_call(
-        self, session: Any, tool_call: LiveServerToolCall
-    ) -> None:
-        """Process tool calls from Gemini and send back responses."""
-        if tool_call.function_calls is None:
-            logging.debug("No function calls in tool_call")
-            return
-
-        for fc in tool_call.function_calls:
-            logging.debug(f"Calling tool function: {fc.name} with args: {fc.args}")
-            func = self._get_func(fc.name)
-            if func is None:
-                logging.error(f"Function {fc.name} not found")
-                continue
-            args = fc.args if fc.args is not None else {}
-
-            # Handle both async and sync functions appropriately
-            if asyncio.iscoroutinefunction(func):
-                # Function is already async
-                response = await func(**args)
-            else:
-                # Run sync function in a thread pool to avoid blocking
-                response = await asyncio.to_thread(func, **args)
-
-            tool_response = types.LiveClientToolResponse(
-                function_responses=[
-                    types.FunctionResponse(name=fc.name, id=fc.id, response=response)
-                ]
-            )
-            logging.debug(f"Tool response: {tool_response}")
-            await session.send(input=tool_response)
-
-    async def receive_from_gemini(self) -> None:
-        """Listen for and process messages from Gemini without blocking."""
-        while result := await self.session._ws.recv(decode=False):
-            await self.websocket.send_bytes(result)
-            raw_message = json.loads(result)
-            if "toolCall" in raw_message:
-                message = types.LiveServerMessage.model_validate(raw_message)
-                tool_call = LiveServerToolCall.model_validate(message.tool_call)
-                # Create a separate task to handle the tool call without blocking
-                task = asyncio.create_task(
-                    self._handle_tool_call(self.session, tool_call)
+    async def run_agent(self) -> None:
+        """Run the agent with the input queue using bidi_stream_query protocol."""
+        try:
+            # Send setupComplete immediately
+            setup_complete_response: dict = {"setupComplete": {}}
+            await self.websocket.send_json(setup_complete_response)
+
+            # Wait for first request with user_id
+            first_request = await self.input_queue.get()
+            self.user_id = first_request.get("user_id")
+            if not self.user_id:
+                raise ValueError("The first request must have a user_id.")
+
+            self.session_id = first_request.get("session_id")
+            first_live_request = first_request.get("live_request")
+
+            # Create session if needed
+            if not self.session_id:
+                session = await session_service.create_session(
+                    app_name="live-app",
+                    user_id=self.user_id,
                 )
-                self._tool_tasks.append(task)
+                self.session_id = session.id
+
+            # Create LiveRequestQueue
+            live_request_queue = LiveRequestQueue()
+
+            # Add first live request if present
+            if first_live_request and isinstance(first_live_request, dict):
+                live_request_queue.send(LiveRequest.model_validate(first_live_request))
+
+            # Forward requests from input_queue to live_request_queue
+            async def _forward_requests() -> None:
+                while True:
+                    request = await self.input_queue.get()
+                    live_request = LiveRequest.model_validate(request)
+                    live_request_queue.send(live_request)
+
+            # Forward events from agent to websocket
+            async def _forward_events() -> None:
+                events_async = runner.run_live(
+                    user_id=self.user_id,
+                    session_id=self.session_id,
+                    live_request_queue=live_request_queue,
+                )
+                async for event in events_async:
+                    event_dict = _utils.dump_event_for_json(event)
+                    await self.websocket.send_json(event_dict)
+
+                    # Check for error responses
+                    if isinstance(event_dict, dict) and "error" in event_dict:
+                        logging.error(f"Agent error: {event_dict['error']}")
+                        break
+
+            # Run both tasks
+            requests_task = asyncio.create_task(_forward_requests())
+
+            try:
+                await _forward_events()
+            finally:
+                requests_task.cancel()
+                try:
+                    await requests_task
+                except asyncio.CancelledError:
+                    pass
+
+        except Exception as e:
+            logging.error(f"Error in agent: {e}")
+            await self.websocket.send_json({"error": str(e)})
 
 
 def get_connect_and_run_callable(websocket: WebSocket) -> Callable:
-    """Create a callable that handles Gemini connection with retry logic.
+    """Create a callable that handles agent connection with retry logic.
 
     Args:
         websocket: The client websocket connection
 
     Returns:
-        Callable: An async function that establishes and manages the Gemini connection
+        Callable: An async function that establishes and manages the agent connection
     """
 
     async def on_backoff(details: backoff._typing.Details) -> None:
@@ -180,18 +226,14 @@ def get_connect_and_run_callable(websocket: WebSocket) -> Callable:
         backoff.expo, ConnectionClosedError, max_tries=10, on_backoff=on_backoff
     )
     async def connect_and_run() -> None:
-        async with genai_client.aio.live.connect(
-            model=MODEL_ID, config=live_connect_config
-        ) as session:
-            await websocket.send_json({"status": "Backend is ready for conversation"})
-            gemini_session = GeminiSession(
-                session=session, websocket=websocket, tool_functions=tool_functions
-            )
-            logging.info("Starting bidirectional communication")
-            await asyncio.gather(
-                gemini_session.receive_from_client(),
-                gemini_session.receive_from_gemini(),
-            )
+        logging.info("Starting ADK agent")
+        session = AgentSession(websocket)
+
+        logging.info("Starting bidirectional communication with agent")
+        await asyncio.gather(
+            session.receive_from_client(),
+            session.run_agent(),
+        )
 
     return connect_and_run
 
@@ -204,15 +246,38 @@ async def websocket_endpoint(websocket: WebSocket) -> None:
     await connect_and_run()
 
 
-class Feedback(BaseModel):
-    """Represents feedback for a conversation."""
+@app.get("/")
+async def serve_frontend_root() -> FileResponse:
+    """Serve the frontend index.html at the root path."""
+    index_file = frontend_build_dir / "index.html"
+    if index_file.exists():
+        return FileResponse(str(index_file))
+    raise HTTPException(
+        status_code=404,
+        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
+    )
+
+
+@app.get("/{full_path:path}")
+async def serve_frontend_spa(full_path: str) -> FileResponse:
+    """Catch-all route to serve the frontend for SPA routing.
+
+    This ensures that client-side routes are handled by the React app.
+    Excludes API routes (ws, feedback) and static assets.
+    """
+    # Don't intercept API routes
+    if full_path.startswith(("ws", "feedback", "static", "api")):
+        raise HTTPException(status_code=404, detail="Not found")
 
-    score: int | float
-    text: str | None = ""
-    run_id: str
-    user_id: str | None
-    log_type: Literal["feedback"] = "feedback"
-{% elif "adk" in cookiecutter.tags %}
+    # Serve index.html for all other routes (SPA routing)
+    index_file = frontend_build_dir / "index.html"
+    if index_file.exists():
+        return FileResponse(str(index_file))
+    raise HTTPException(
+        status_code=404,
+        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
+    )
+{% elif cookiecutter.is_adk %}
 import os
 
 import google.auth
@@ -358,7 +423,7 @@ def stream_messages(
     set_tracing_properties(config)
     input_dict = input.model_dump()
 
-    for data in agent.stream(input_dict, config=config, stream_mode="messages"):  # type: ignore[arg-type]
+    for data in agent.stream(input_dict, config=config, stream_mode="messages"):
         yield dumps(data) + "\n"
 
 
@@ -397,40 +462,7 @@ def collect_feedback(feedback: Feedback) -> dict[str, str]:
     """
     logger.log_struct(feedback.model_dump(), severity="INFO")
     return {"status": "success"}
-{% if cookiecutter.agent_name == "live_api" %}
-
-@app.get("/")
-async def serve_frontend_root() -> FileResponse:
-    """Serve the frontend index.html at the root path."""
-    index_file = frontend_build_dir / "index.html"
-    if index_file.exists():
-        return FileResponse(str(index_file))
-    raise HTTPException(
-        status_code=404,
-        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
-    )
-
 
-@app.get("/{full_path:path}")
-async def serve_frontend_spa(full_path: str) -> FileResponse:
-    """Catch-all route to serve the frontend for SPA routing.
-
-    This ensures that client-side routes are handled by the React app.
-    Excludes API routes (ws, feedback) and static assets.
-    """
-    # Don't intercept API routes
-    if full_path.startswith(("ws", "feedback", "static", "api")):
-        raise HTTPException(status_code=404, detail="Not found")
-
-    # Serve index.html for all other routes (SPA routing)
-    index_file = frontend_build_dir / "index.html"
-    if index_file.exists():
-        return FileResponse(str(index_file))
-    raise HTTPException(
-        status_code=404,
-        detail="Frontend not built. Run 'npm run build' in the frontend directory.",
-    )
-{% endif %}
 
 # Main execution
 if __name__ == "__main__":