diff --git a/examples/vertex_ai_conversation/check_documents_in_datastore.ipynb b/examples/vertex_ai_conversation/check_documents_in_datastore.ipynb index 35cf46e4..d0ae03c9 100644 --- a/examples/vertex_ai_conversation/check_documents_in_datastore.ipynb +++ b/examples/vertex_ai_conversation/check_documents_in_datastore.ipynb @@ -9,17 +9,17 @@ "\n", "
\n",
- " \n",
+ " \n",
" Run in Colab\n", " \n", " | \n",
" \n",
- " \n",
+ " \n",
" View on GitHub\n", " \n", " | \n",
" \n",
- " \n",
+ " \n",
" Open in Vertex AI Workbench\n", " \n", " | \n",
@@ -29,23 +29,17 @@
]
},
{
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "vscode": {
- "languageId": "plaintext"
- }
- },
- "outputs": [],
+ "cell_type": "markdown",
+ "metadata": {},
"source": [
- "## Instructions:"
+ "## Instructions"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
- "## Setup:"
+ "## Setup"
]
},
{
@@ -91,11 +85,19 @@
"from google.api_core import operations_v1, grpc_helpers\n",
"from google.longrunning import operations_pb2\n",
"from typing import List, Optional\n",
+ "from google.api_core.client_options import ClientOptions\n",
+ "\n",
"\n",
"def list_documents(\n",
" project_id: str, location: str, datastore_id: str, rate_limit: int = 1):\n",
" \"\"\"Gets a list of docs in a datastore.\"\"\"\n",
- " client = discoveryengine_v1.DocumentServiceClient()\n",
+ " # Create a client\n",
+ " client_options = (\n",
+ " ClientOptions(api_endpoint=f\"{location}-discoveryengine.googleapis.com\")\n",
+ " if location != \"global\"\n",
+ " else None)\n",
+ " \n",
+ " client = discoveryengine_v1.DocumentServiceClient(client_options=client_options)\n",
"\n",
" request = discoveryengine_v1.ListDocumentsRequest(\n",
" parent=f'projects/{project_id}/locations/{location}/collections/default_collection/dataStores/{datastore_id}/branches/0',\n",