agent-starter-pack 0.3.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

agents/live_api/app/agent.py

@@ -1,78 +0,0 @@
-# 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.
-
-import os
-
-import google
-import vertexai
-from google import genai
-from google.genai.types import (
-    Content,
-    LiveConnectConfig,
-)
-from langchain_google_vertexai import VertexAIEmbeddings
-
-from app.templates import FORMAT_DOCS, SYSTEM_INSTRUCTION
-from app.vector_store import get_vector_store
-
-# Constants
-VERTEXAI = os.getenv("VERTEXAI", "true").lower() == "true"
-LOCATION = "us-central1"
-EMBEDDING_MODEL = "text-embedding-004"
-MODEL_ID = "gemini-2.0-flash-exp"
-URLS = [
-    "https://cloud.google.com/architecture/deploy-operate-generative-ai-applications"
-]
-
-# Initialize Google Cloud clients
-credentials, project_id = google.auth.default()
-vertexai.init(project=project_id, location=LOCATION)
-
-
-if VERTEXAI:
-    genai_client = genai.Client(project=project_id, location=LOCATION, vertexai=True)
-else:
-    # API key should be set using GOOGLE_API_KEY environment variable
-    genai_client = genai.Client(http_options={"api_version": "v1alpha"})
-
-# Initialize vector store and retriever
-embedding = VertexAIEmbeddings(model_name=EMBEDDING_MODEL)
-vector_store = get_vector_store(embedding=embedding, urls=URLS)
-retriever = vector_store.as_retriever()
-
-
-def retrieve_docs(query: str) -> dict[str, str]:
-    """
-    Retrieves pre-formatted documents about MLOps (Machine Learning Operations),
-      Gen AI lifecycle, and production deployment best practices.
-
-    Args:
-        query: Search query string related to MLOps, Gen AI, or production deployment.
-
-    Returns:
-        A set of relevant, pre-formatted documents.
-    """
-    docs = retriever.invoke(query)
-    formatted_docs = FORMAT_DOCS.format(docs=docs)
-    return {"output": formatted_docs}
-
-
-# Configure tools available to the agent and live connection
-tool_functions = {"retrieve_docs": retrieve_docs}
-
-live_connect_config = LiveConnectConfig(
-    response_modalities=["AUDIO"],
-    tools=[retrieve_docs],
-    system_instruction=Content(parts=[{"text": SYSTEM_INSTRUCTION}]),
-)

agents/live_api/app/server.py

@@ -1,196 +0,0 @@
-# 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.
-
-import asyncio
-import json
-import logging
-from collections.abc import Callable
-from typing import Any, Literal
-
-import backoff
-from fastapi import FastAPI, WebSocket
-from fastapi.middleware.cors import CORSMiddleware
-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 websockets.exceptions import ConnectionClosedError
-
-from app.agent import MODEL_ID, genai_client, live_connect_config, tool_functions
-
-app = FastAPI()
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-logging_client = google_cloud_logging.Client()
-logger = logging_client.logger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-
-class GeminiSession:
-    """Manages bidirectional communication between a client and the Gemini model."""
-
-    def __init__(
-        self, session: Any, websocket: WebSocket, tool_functions: dict[str, Callable]
-    ) -> None:
-        """Initialize the Gemini 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
-
-    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.
-        """
-        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"
-                    )
-                else:
-                    logging.warning(f"Received unexpected input from client: {data}")
-            except ConnectionClosedError as e:
-                logging.warning(f"Client {self.user_id} closed connection: {e}")
-                break
-            except Exception as e:
-                logging.error(f"Error receiving from client {self.user_id}: {e!s}")
-                break
-
-    def _get_func(self, action_label: str) -> Callable | None:
-        """Get the tool function for a given action label."""
-        return None if action_label == "" else 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.
-
-        Args:
-            session: The Gemini session
-            tool_call: Tool call request from Gemini
-        """
-        for fc in tool_call.function_calls:
-            logging.debug(f"Calling tool function: {fc.name} with args: {fc.args}")
-            response = self._get_func(fc.name)(**fc.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.
-
-        Continuously receives messages from Gemini, forwards them to the client,
-        and handles any tool calls. Handles connection errors gracefully.
-        """
-        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)
-                await self._handle_tool_call(self.session, tool_call)
-
-
-def get_connect_and_run_callable(websocket: WebSocket) -> Callable:
-    """Create a callable that handles Gemini connection with retry logic.
-
-    Args:
-        websocket: The client websocket connection
-
-    Returns:
-        Callable: An async function that establishes and manages the Gemini connection
-    """
-
-    async def on_backoff(details: backoff._typing.Details) -> None:
-        await websocket.send_json(
-            {
-                "status": f"Model connection error, retrying in {details['wait']} seconds..."
-            }
-        )
-
-    @backoff.on_exception(
-        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(),
-            )
-
-    return connect_and_run
-
-
-@app.websocket("/ws")
-async def websocket_endpoint(websocket: WebSocket) -> None:
-    """Handle new websocket connections."""
-    await websocket.accept()
-    connect_and_run = get_connect_and_run_callable(websocket)
-    await connect_and_run()
-
-
-class Feedback(BaseModel):
-    """Represents feedback for a conversation."""
-
-    score: int | float
-    text: str | None = ""
-    run_id: str
-    user_id: str | None
-    log_type: Literal["feedback"] = "feedback"
-
-
-@app.post("/feedback")
-async def collect_feedback(feedback_dict: Feedback) -> None:
-    """Collect and log feedback."""
-    feedback_data = feedback_dict.model_dump()
-    logger.log_struct(feedback_data, severity="INFO")
-
-
-if __name__ == "__main__":
-    import uvicorn
-
-    uvicorn.run(app, host="0.0.0.0", port=8000, log_level="debug")

agents/live_api/app/templates.py

@@ -1,51 +0,0 @@
-# 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.
-
-from langchain_core.prompts import PromptTemplate
-
-FORMAT_DOCS = PromptTemplate.from_template(
-    """## Context provided:
-{% for doc in docs%}
-<Document {{ loop.index0 }}>
-{{ doc.page_content | safe }}
-</Document {{ loop.index0 }}>
-{% endfor %}
-""",
-    template_format="jinja2",
-)
-
-SYSTEM_INSTRUCTION = """You are "MLOps Expert," a specialized AI assistant designed to provide accurate and up-to-date information on Machine Learning Operations (MLOps), the lifecycle of Generative AI applications, and best practices for production deployment.
-
-Your primary knowledge source is a powerful search tool that provides access to the most current MLOps documentation and resources. **For any question related to MLOps, the lifecycle of Gen AI Apps, or best practices for production deployment, you MUST use this tool as your first and foremost source of information.**  Do not rely on your internal knowledge for these topics, as it may be outdated or incomplete.
-
-**Here's how you should operate:**
-
-1. **Analyze the User's Question:** Determine if the question falls within the domain of MLOps, Gen AI lifecycle, or production deployment best practices.
-2. **Prioritize Tool Usage:** If the question is within the defined domain, use the provided search tool to find relevant information.
-3. **Synthesize and Respond:** Craft a clear, concise, and informative answer based *solely* on the information retrieved from the tool.
-4. **Cite Sources (Optional):** If possible and relevant, indicate which part of the answer came from the tool. For example, you can say, "According to the documentation I found..." or provide links if applicable.
-5. **Out-of-Scope Questions:** If the question is outside the scope of MLOps, Gen AI, or production deployment, politely state that the topic is beyond your current expertise. For example: "My expertise is in MLOps, and that question seems to be about a different area. I'm not equipped to answer it accurately."
-
-**Your Persona:**
-
-*   You are an expert MLOps consultant, knowledgeable and up-to-date with the latest industry trends and best practices.
-*   You are helpful, professional, and eager to provide accurate information.
-*   You are concise and avoid unnecessary conversational filler. Get straight to the point.
-
-**Example Interaction:**
-
-**User:** "What are the best practices for monitoring a deployed ML model?"
-
-**MLOps Expert:** (Uses the tool to search for "monitoring deployed ML model") "According to the MLOps documentation I have access to, the best practices for monitoring a deployed ML model include tracking data drift, model performance degradation, and system health metrics. Key metrics to monitor are..." (continues with information found in the tool).
-"""

agents/live_api/app/vector_store.py

@@ -1,55 +0,0 @@
-# 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.
-
-import logging
-import os
-
-from langchain.schema import Document
-from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.document_loaders import WebBaseLoader
-from langchain_community.vectorstores import SKLearnVectorStore
-from langchain_core.embeddings import Embeddings
-
-PERSIST_PATH = ".persist_vector_store"
-
-
-def load_and_split_documents(urls: list[str]) -> list[Document]:
-    """Load and split documents from a list of URLs."""
-    docs = [WebBaseLoader(url).load() for url in urls]
-    docs_list = [item for sublist in docs for item in sublist]
-    logging.info(f"# of documents loaded (pre-chunking) = {len(docs_list)}")
-
-    text_splitter = CharacterTextSplitter(chunk_size=2000, chunk_overlap=50)
-    doc_splits = text_splitter.split_documents(docs_list)
-    logging.info(f"# of documents after split = {len(doc_splits)}")
-
-    return doc_splits
-
-
-def get_vector_store(
-    embedding: Embeddings, urls: list[str], persist_path: str = PERSIST_PATH
-) -> SKLearnVectorStore:
-    """Get or create a vector store."""
-
-    if os.path.exists(persist_path):
-        vector_store = SKLearnVectorStore(
-            embedding=embedding, persist_path=persist_path
-        )
-    else:
-        doc_splits = load_and_split_documents(urls=urls)
-        vector_store = SKLearnVectorStore.from_documents(
-            documents=doc_splits, embedding=embedding, persist_path=persist_path
-        )
-        vector_store.persist()
-    return vector_store

agents/live_api/template/.templateconfig.yaml

@@ -1,29 +0,0 @@
-# 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.
-
-description: "A real-time multimodal RAG agent powered by Gemini Live API, supporting audio/video/text chat with vector DB-backed responses"
-settings:
-  requires_data_ingestion: false
-  frontend_type: "live_api_react"
-  deployment_targets: ["cloud_run"]
-  extra_dependencies: [
-    "backoff~=2.2.1",
-    "beautifulsoup4~=4.12.3",
-    "google-genai~=1.8.0",
-    "jinja2~=3.1.4",
-    "langchain~=0.3.13",
-    "langchain-community~=0.3.13",
-    "langchain-google-vertexai~=2.0.9",
-    "scikit-learn>=1.0.0,<2.0.0",
-  ]

agents/live_api/tests/integration/test_server_e2e.py

@@ -1,254 +0,0 @@
-# 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.
-
-import asyncio
-import json
-import logging
-import os
-import subprocess
-import sys
-import threading
-import time
-import uuid
-from collections.abc import Iterator
-from typing import Any
-
-import pytest
-import requests
-import websockets.client
-
-# Configure logging
-logging.basicConfig(level=logging.DEBUG)
-logger = logging.getLogger(__name__)
-
-BASE_URL = "ws://127.0.0.1:8000/"
-WS_URL = BASE_URL + "ws"
-
-FEEDBACK_URL = "http://127.0.0.1:8000/feedback"
-HEADERS = {"Content-Type": "application/json"}
-
-
-def log_output(pipe: Any, log_func: Any) -> None:
-    """Log the output from the given pipe."""
-    for line in iter(pipe.readline, ""):
-        log_func(line.strip())
-
-
-def start_server() -> subprocess.Popen[str]:
-    """Start the FastAPI server using subprocess and log its output."""
-    command = [
-        sys.executable,
-        "-m",
-        "uvicorn",
-        "app.server:app",
-        "--host",
-        "0.0.0.0",
-        "--port",
-        "8000",
-    ]
-    env = os.environ.copy()
-    env["INTEGRATION_TEST"] = "TRUE"
-    process = subprocess.Popen(
-        command,
-        stdout=subprocess.PIPE,
-        stderr=subprocess.PIPE,
-        text=True,
-        bufsize=1,
-        env=env,
-    )
-
-    # Start threads to log stdout and stderr in real-time
-    threading.Thread(
-        target=log_output, args=(process.stdout, logger.info), daemon=True
-    ).start()
-    threading.Thread(
-        target=log_output, args=(process.stderr, logger.error), daemon=True
-    ).start()
-
-    return process
-
-
-def wait_for_server(timeout: int = 60, interval: int = 1) -> bool:
-    """Wait for the server to be ready."""
-    start_time = time.time()
-    while time.time() - start_time < timeout:
-        try:
-            response = requests.get("http://127.0.0.1:8000/docs", timeout=10)
-            if response.status_code == 200:
-                logger.info("Server is ready")
-                return True
-        except Exception:
-            pass
-        time.sleep(interval)
-    logger.error(f"Server did not become ready within {timeout} seconds")
-    return False
-
-
-@pytest.fixture(scope="session")
-def server_fixture(request: Any) -> Iterator[subprocess.Popen[str]]:
-    """Pytest fixture to start and stop the server for testing."""
-    logger.info("Starting server process")
-    server_process = start_server()
-    if not wait_for_server():
-        pytest.fail("Server failed to start")
-    logger.info("Server process started")
-
-    def stop_server() -> None:
-        logger.info("Stopping server process")
-        server_process.terminate()
-        server_process.wait()
-        logger.info("Server process stopped")
-
-    request.addfinalizer(stop_server)
-    yield server_process
-
-
-@pytest.mark.asyncio
-async def test_websocket_connection(server_fixture: subprocess.Popen[str]) -> None:
-    """Test the websocket connection and message exchange."""
-
-    async def send_message(websocket: Any, message: dict[str, Any]) -> None:
-        """Helper function to send messages and log them."""
-        await websocket.send(json.dumps(message))
-
-    async def receive_message(websocket: Any, timeout: float = 5.0) -> dict[str, Any]:
-        """Helper function to receive messages with timeout."""
-        try:
-            response = await asyncio.wait_for(websocket.recv(), timeout=timeout)
-            if isinstance(response, bytes):
-                return json.loads(response.decode())
-            if isinstance(response, str):
-                return json.loads(response)
-            return response
-        except asyncio.TimeoutError as exc:
-            raise TimeoutError(
-                f"No response received within {timeout} seconds"
-            ) from exc
-
-    try:
-        await asyncio.sleep(2)
-
-        async with websockets.connect(
-            WS_URL, ping_timeout=10, close_timeout=10
-        ) as websocket:
-            try:
-                # Wait for initial ready message
-                initial_response = None
-                for _ in range(10):
-                    try:
-                        initial_response = await receive_message(websocket, timeout=5.0)
-                        if (
-                            initial_response is not None
-                            and initial_response.get("status")
-                            == "Backend is ready for conversation"
-                        ):
-                            break
-                    except TimeoutError:
-                        if _ == 9:
-                            raise
-                        continue
-
-                assert (
-                    initial_response is not None
-                    and initial_response.get("status")
-                    == "Backend is ready for conversation"
-                )
-
-                # Send messages
-                setup_msg = {"setup": {"run_id": "test-run", "user_id": "test-user"}}
-                await send_message(websocket, setup_msg)
-
-                dummy_audio = bytes([0] * 1024)
-                audio_msg = {
-                    "realtimeInput": {
-                        "mediaChunks": [
-                            {
-                                "mimeType": "audio/pcm;rate=16000",
-                                "data": dummy_audio.hex(),
-                            }
-                        ]
-                    }
-                }
-                await send_message(websocket, audio_msg)
-
-                text_msg = {
-                    "clientContent": {
-                        "turns": [
-                            {"role": "user", "parts": [{"text": "Hello, how are you?"}]}
-                        ],
-                        "turnComplete": True,
-                    }
-                }
-                await send_message(websocket, text_msg)
-
-                # Collect responses with timeout
-                responses = []
-                try:
-                    while True:
-                        try:
-                            response = await receive_message(websocket, timeout=10.0)
-                            responses.append(response)
-                            if (
-                                len(responses) >= 3
-                            ):  # Exit after receiving enough responses
-                                break
-                        except TimeoutError:
-                            break
-                except asyncio.TimeoutError:
-                    logger.info("Response collection timed out")
-
-                # Verify responses
-                assert len(responses) > 0, "No responses received from server"
-                assert any(
-                    isinstance(r, dict) and "serverContent" in r for r in responses
-                )
-                logger.info(
-                    f"Test completed successfully. Received {len(responses)} responses"
-                )
-
-            finally:
-                # Ensure websocket is closed properly
-                await websocket.close()
-
-    except Exception as e:
-        logger.error(f"Test failed with error: {e!s}")
-        raise
-
-    finally:
-        # Clean up any remaining tasks
-        tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-        for task in tasks:
-            task.cancel()
-        if tasks:
-            await asyncio.gather(*tasks, return_exceptions=True)
-
-
-def test_collect_feedback(server_fixture: subprocess.Popen[str]) -> None:
-    """
-    Test the feedback collection endpoint (/feedback) to ensure it properly
-    logs the received feedback.
-    """
-    # Create sample feedback data
-    feedback_data = {
-        "score": 4,
-        "text": "Great response!",
-        "run_id": str(uuid.uuid4()),
-        "user_id": "user1",
-        "log_type": "feedback",
-    }
-
-    response = requests.post(
-        FEEDBACK_URL, json=feedback_data, headers=HEADERS, timeout=10
-    )
-    assert response.status_code == 200