agent-starter-pack 0.2.2__py3-none-any.whl → 0.3.0__py3-none-any.whl

agents/adk_base/notebooks/adk_app_testing.ipynb

@@ -0,0 +1,305 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# ADK Application Testing\n",
+    "\n",
+    "This notebook demonstrates how to test an ADK (Agent Development Kit) application.\n",
+    "It covers both local and remote testing, both with Agent Engine and Cloud Run.\n",
+    "\n",
+    "<img src=\"https://github.com/GoogleCloudPlatform/agent-starter-pack/blob/main/docs/images/adk_logo.png?raw=true\" width=\"400\">\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Import libraries"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import json\n",
+    "\n",
+    "import requests\n",
+    "import vertexai.agent_engines"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## If you are using Agent Engine\n",
+    "See more documentation at [Agent Engine Overview](https://cloud.google.com/vertex-ai/generative-ai/docs/agent-engine/overview)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Local Testing\n",
+    "\n",
+    "You can import directly the AgentEngineApp class within your environment. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from app.agent import root_agent\n",
+    "from app.agent_engine_app import AgentEngineApp\n",
+    "\n",
+    "agent_engine = AgentEngineApp(agent=root_agent)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for event in agent_engine.stream_query(message=\"hi!\", user_id=\"test\"):\n",
+    "    print(event)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Remote Testing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Replace with your Agent Engine ID\n",
+    "AGENT_ENGINE_ID = \"projects/PROJECT_ID/locations/us-central1/reasoningEngines/ENGINE_ID\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "remote_agent_engine = vertexai.agent_engines.get(AGENT_ENGINE_ID)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "for event in remote_agent_engine.stream_query(message=\"hi!\", user_id=\"test\"):\n",
+    "    print(event)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## If you are using Cloud Run"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Local Testing\n",
+    "\n",
+    "> You can run the application locally via the `make backend` command."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Create a session\n",
+    " Create a new session with user preferences and state information\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "user_id = \"test_user_123\"\n",
+    "session_id = \"test_session_456\"\n",
+    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
+    "\n",
+    "session_url = f\"http://127.0.0.1:8000/apps/app/users/{user_id}/sessions/{session_id}\"\n",
+    "headers = {\"Content-Type\": \"application/json\"}\n",
+    "\n",
+    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
+    "print(f\"Session creation status code: {session_response.status_code}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Send a message\n",
+    "Send a message to the backend service and receive a streaming response\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "message_data = {\n",
+    "    \"app_name\": \"app\",\n",
+    "    \"user_id\": user_id,\n",
+    "    \"session_id\": session_id,\n",
+    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
+    "    \"streaming\": True,\n",
+    "}\n",
+    "\n",
+    "message_url = \"http://127.0.0.1:8000/run_sse\"\n",
+    "message_response = requests.post(\n",
+    "    message_url, headers=headers, json=message_data, stream=True\n",
+    ")\n",
+    "\n",
+    "print(f\"Message send status code: {message_response.status_code}\")\n",
+    "\n",
+    "# Print streamed response\n",
+    "for line in message_response.iter_lines():\n",
+    "    if line:\n",
+    "        line_str = line.decode(\"utf-8\")\n",
+    "        if line_str.startswith(\"data: \"):\n",
+    "            event_json = line_str[6:]\n",
+    "            event = json.loads(event_json)\n",
+    "            print(f\"Received event: {event}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "#### Remote Testing\n",
+    "\n",
+    "For more information about authenticating HTTPS requests to Cloud Run services, see:\n",
+    "[Cloud Run Authentication Documentation](https://cloud.google.com/run/docs/triggering/https-request)\n",
+    "\n",
+    "Remote testing involves using a deployed service URL instead of localhost.\n",
+    "\n",
+    "Authentication is handled using GCP identity tokens instead of local credentials."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "ID_TOKEN = get_ipython().getoutput(\"gcloud auth print-identity-token -q\")[0]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SERVICE_URL = \"YOUR_SERVICE_URL_HERE\"  # Replace with your Cloud Run service URL"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "You'll need to first create a Session"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "user_id = \"test_user_123\"\n",
+    "session_id = \"test_session_456\"\n",
+    "session_data = {\"state\": {\"preferred_language\": \"English\", \"visit_count\": 1}}\n",
+    "\n",
+    "session_url = f\"{SERVICE_URL}/apps/app/users/{user_id}/sessions/{session_id}\"\n",
+    "headers = {\"Content-Type\": \"application/json\", \"Authorization\": f\"Bearer {ID_TOKEN}\"}\n",
+    "\n",
+    "session_response = requests.post(session_url, headers=headers, json=session_data)\n",
+    "print(f\"Session creation status code: {session_response.status_code}\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Then you will be able to send a message"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "message_data = {\n",
+    "    \"app_name\": \"app\",\n",
+    "    \"user_id\": user_id,\n",
+    "    \"session_id\": session_id,\n",
+    "    \"new_message\": {\"role\": \"user\", \"parts\": [{\"text\": \"Hello! Weather in New york?\"}]},\n",
+    "    \"streaming\": True,\n",
+    "}\n",
+    "\n",
+    "message_url = f\"{SERVICE_URL}/run_sse\"\n",
+    "message_response = requests.post(\n",
+    "    message_url, headers=headers, json=message_data, stream=True\n",
+    ")\n",
+    "\n",
+    "print(f\"Message send status code: {message_response.status_code}\")\n",
+    "\n",
+    "# Print streamed response\n",
+    "for line in message_response.iter_lines():\n",
+    "    if line:\n",
+    "        line_str = line.decode(\"utf-8\")\n",
+    "        if line_str.startswith(\"data: \"):\n",
+    "            event_json = line_str[6:]\n",
+    "            event = json.loads(event_json)\n",
+    "            print(f\"Received event: {event}\")"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.8"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

agents/adk_base/template/.templateconfig.yaml

@@ -0,0 +1,21 @@
+# 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: "An agent implementing a base ReAct agent using Google's Agent Development Kit"
+settings:
+  requires_data_ingestion: false
+  deployment_targets: ["agent_engine", "cloud_run"]
+  extra_dependencies: ["google-adk~=0.1.0"]
+  tags: ["adk"]
+  frontend_type: "streamlit_adk"

agents/adk_base/tests/integration/test_agent.py

@@ -0,0 +1,58 @@
+# 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.
+
+# mypy: disable-error-code="union-attr"
+from google.adk.agents.run_config import RunConfig, StreamingMode
+from google.adk.runners import Runner
+from google.adk.sessions import InMemorySessionService
+from google.genai import types
+
+from app.agent import root_agent
+
+
+def test_agent_stream() -> None:
+    """
+    Integration test for the agent stream functionality.
+    Tests that the agent returns valid streaming responses.
+    """
+
+    session_service = InMemorySessionService()
+
+    session = session_service.create_session(user_id="test_user", app_name="test")
+    runner = Runner(agent=root_agent, session_service=session_service, app_name="test")
+
+    message = types.Content(
+        role="user", parts=[types.Part.from_text(text="Why is the sky blue?")]
+    )
+
+    events = list(
+        runner.run(
+            new_message=message,
+            user_id="test_user",
+            session_id=session.id,
+            run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+        )
+    )
+    assert len(events) > 0, "Expected at least one message"
+
+    has_text_content = False
+    for event in events:
+        if (
+            event.content
+            and event.content.parts
+            and any(part.text for part in event.content.parts)
+        ):
+            has_text_content = True
+            break
+    assert has_text_content, "Expected at least one message with text content"

agents/agentic_rag/README.md

@@ -14,6 +14,7 @@ The agent implements the following architecture:
 
 ### Key Features
 
+- **Built on Agent Development Kit (ADK):** ADK is a flexible, modular framework for developing and deploying AI agents. It integrates with the Google ecosystem and Gemini models, supporting various LLMs and open-source AI tools, enabling both simple and complex agent architectures.
 - **Flexible Datastore Options:** Choose between Vertex AI Search or Vertex AI Vector Search for efficient data storage and retrieval based on your specific needs.
 - **Automated Data Ingestion Pipeline:** Automates the process of ingesting data from input sources.
 - **Custom Embeddings:** Generates embeddings using Vertex AI Embeddings and incorporates them into your data for enhanced semantic search.

agents/agentic_rag/app/agent.py

@@ -17,25 +17,22 @@ import os
 
 import google
 import vertexai
-from langchain_core.documents import Document
-from langchain_core.messages import BaseMessage
-from langchain_core.runnables import RunnableConfig
-from langchain_core.tools import tool
-from langchain_google_vertexai import ChatVertexAI, VertexAIEmbeddings
-from langgraph.graph import END, MessagesState, StateGraph
-from langgraph.prebuilt import ToolNode
+from google.adk.agents import Agent
+from langchain_google_vertexai import VertexAIEmbeddings
 
 from app.retrievers import get_compressor, get_retriever
-from app.templates import format_docs, inspect_conversation_template, rag_template
+from app.templates import format_docs
 
 EMBEDDING_MODEL = "text-embedding-005"
 LOCATION = "us-central1"
 LLM = "gemini-2.0-flash-001"
 
-# Initialize Google Cloud and Vertex AI
 credentials, project_id = google.auth.default()
-vertexai.init(project=project_id, location=LOCATION)
+os.environ.setdefault("GOOGLE_CLOUD_PROJECT", project_id)
+os.environ.setdefault("GOOGLE_CLOUD_LOCATION", LOCATION)
+os.environ.setdefault("GOOGLE_GENAI_USE_VERTEXAI", "True")
 
+vertexai.init(project=project_id, location=LOCATION)
 embedding = VertexAIEmbeddings(
     project=project_id, location=LOCATION, model_name=EMBEDDING_MODEL
 )
@@ -45,7 +42,7 @@ EMBEDDING_COLUMN = "embedding"
 TOP_K = 5
 
 data_store_region = os.getenv("DATA_STORE_REGION", "us")
-data_store_id = os.getenv("DATA_STORE_ID", "sample-datastore")
+data_store_id = os.getenv("DATA_STORE_ID", "{{cookiecutter.project_name}}-datastore")
 
 retriever = get_retriever(
     project_id=project_id,
@@ -56,9 +53,15 @@ retriever = get_retriever(
     max_documents=10,
 )
 {% elif cookiecutter.datastore_type == "vertex_ai_vector_search" %}
-vector_search_index = os.getenv("VECTOR_SEARCH_INDEX")
-vector_search_index_endpoint = os.getenv("VECTOR_SEARCH_INDEX_ENDPOINT")
-vector_search_bucket = os.getenv("VECTOR_SEARCH_BUCKET")
+vector_search_index = os.getenv(
+    "VECTOR_SEARCH_INDEX", "{{cookiecutter.project_name}}-vector-search"
+)
+vector_search_index_endpoint = os.getenv(
+    "VECTOR_SEARCH_INDEX_ENDPOINT", "{{cookiecutter.project_name}}-vector-search-endpoint"
+)
+vector_search_bucket = os.getenv(
+    "VECTOR_SEARCH_BUCKET", f"{project_id}-{{cookiecutter.project_name}}-vs"
+)
 
 retriever = get_retriever(
     project_id=project_id,
@@ -74,8 +77,7 @@ compressor = get_compressor(
 )
 
 
-@tool(response_format="content_and_artifact")
-def retrieve_docs(query: str) -> tuple[str, list[Document]]:
+def retrieve_docs(query: str) -> str:
     """
     Useful for retrieving relevant documents based on a query.
     Use this when you need additional information to answer a question.
@@ -84,78 +86,31 @@ def retrieve_docs(query: str) -> tuple[str, list[Document]]:
         query (str): The user's question or search query.
 
     Returns:
-        List[Document]: A list of the top-ranked Document objects, limited to TOP_K (5) results.
-    """
-    # Use the retriever to fetch relevant documents based on the query
-    retrieved_docs = retriever.invoke(query)
-    # Re-rank docs with Vertex AI Rank for better relevance
-    ranked_docs = compressor.compress_documents(documents=retrieved_docs, query=query)
-    # Format ranked documents into a consistent structure for LLM consumption
-    formatted_docs = format_docs.format(docs=ranked_docs)
-    return (formatted_docs, ranked_docs)
-
-
-@tool
-def should_continue() -> None:
-    """
-    Use this tool if you determine that you have enough context to respond to the questions of the user.
+        str: Formatted string containing relevant document content retrieved and ranked based on the query.
     """
-    return None
-
-
-tools = [retrieve_docs, should_continue]
-
-llm = ChatVertexAI(model=LLM, temperature=0, max_tokens=1024, streaming=True)
-
-# Set up conversation inspector
-inspect_conversation = inspect_conversation_template | llm.bind_tools(
-    tools, tool_choice="any"
+    try:
+        # Use the retriever to fetch relevant documents based on the query
+        retrieved_docs = retriever.invoke(query)
+        # Re-rank docs with Vertex AI Rank for better relevance
+        ranked_docs = compressor.compress_documents(
+            documents=retrieved_docs, query=query
+        )
+        # Format ranked documents into a consistent structure for LLM consumption
+        formatted_docs = format_docs.format(docs=ranked_docs)
+    except Exception as e:
+        return f"Calling retrieval tool with query:\n\n{query}\n\nraised the following error:\n\n{type(e)}: {e}"
+
+    return formatted_docs
+
+
+instruction = """You are an AI assistant for question-answering tasks.
+Answer to the best of your ability using the context provided.
+Leverage the Tools you are provided to answer questions.
+If you already know the answer to a question, you can respond directly without using the tools."""
+
+root_agent = Agent(
+    name="root_agent",
+    model="gemini-2.0-flash",
+    instruction=instruction,
+    tools=[retrieve_docs],
 )
-
-# Set up response chain
-response_chain = rag_template | llm
-
-
-def inspect_conversation_node(
-    state: MessagesState, config: RunnableConfig
-) -> dict[str, BaseMessage]:
-    """Inspects the conversation state and returns the next message using the conversation inspector."""
-    response = inspect_conversation.invoke(state, config)
-    return {"messages": response}
-
-
-def generate_node(
-    state: MessagesState, config: RunnableConfig
-) -> dict[str, BaseMessage]:
-    """Generates a response using the RAG template and returns it as a message."""
-    response = response_chain.invoke(state, config)
-    return {"messages": response}
-
-
-# Flow:
-# 1. Start with agent node that inspects conversation using inspect_conversation_node
-# 2. Agent node connects to tools node which can either:
-#    - Retrieve relevant docs using retrieve_docs tool
-#    - End tool usage with should_continue tool
-# 3. Tools node connects to generate node which produces final response
-# 4. Generate node connects to END to complete the workflow
-
-workflow = StateGraph(MessagesState)
-workflow.add_node("agent", inspect_conversation_node)
-workflow.add_node("generate", generate_node)
-workflow.set_entry_point("agent")
-
-workflow.add_node(
-    "tools",
-    ToolNode(
-        tools=tools,
-        # With False, tool errors won't be caught by LangGraph
-        handle_tool_errors=True,
-    ),
-)
-workflow.add_edge("agent", "tools")
-workflow.add_edge("tools", "generate")
-
-workflow.add_edge("generate", END)
-
-agent = workflow.compile()

agents/agentic_rag/app/templates.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 
 from langchain_core.prompts import (
-    ChatPromptTemplate,
-    MessagesPlaceholder,
     PromptTemplate,
 )
 
@@ -28,26 +26,3 @@ format_docs = PromptTemplate.from_template(
 """,
     template_format="jinja2",
 )
-
-inspect_conversation_template = ChatPromptTemplate.from_messages(
-    [
-        (
-            "system",
-            """You are an AI assistant tasked with analyzing the conversation"""
-            """ and determining the best course of action.""",
-        ),
-        MessagesPlaceholder(variable_name="messages"),
-    ]
-)
-
-rag_template = ChatPromptTemplate.from_messages(
-    [
-        (
-            "system",
-            """You are an AI assistant for question-answering tasks."""
-            """ Answer to the best of your ability using the context provided."""
-            """ Leverage the Tools you are provided to answer questions.""",
-        ),
-        MessagesPlaceholder(variable_name="messages"),
-    ]
-)