diff --git a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb index 0a41447a53..b2966c404c 100644 --- a/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb +++ b/notebooks/generative_ai/bq_dataframes_llm_code_generation.ipynb @@ -370,83 +370,6 @@ "This section walks through a few steps required in order to use the model in your notebook." ] }, - { - "cell_type": "markdown", - "metadata": { - "id": "rS4VO1TGiO4G" - }, - "source": [ - "## Create a BigQuery Cloud resource connection\n", - "\n", - "You need to create a [Cloud resource connection](https://cloud.google.com/bigquery/docs/create-cloud-resource-connection) to enable BigQuery DataFrames to interact with Vertex AI services." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "KFPjDM4LVh96" - }, - "outputs": [], - "source": [ - "CONN_NAME = \"bqdf-llm\"\n", - "\n", - "client = bq_connection.ConnectionServiceClient()\n", - "new_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}\"\n", - "exists_conn_parent = f\"projects/{PROJECT_ID}/locations/{REGION}/connections/{CONN_NAME}\"\n", - "cloud_resource_properties = bq_connection.CloudResourceProperties({})\n", - "\n", - "try:\n", - " request = client.get_connection(\n", - " request=bq_connection.GetConnectionRequest(name=exists_conn_parent)\n", - " )\n", - " CONN_SERVICE_ACCOUNT = f\"serviceAccount:{request.cloud_resource.service_account_id}\"\n", - "except Exception:\n", - " connection = bq_connection.types.Connection(\n", - " {\"friendly_name\": CONN_NAME, \"cloud_resource\": cloud_resource_properties}\n", - " )\n", - " request = bq_connection.CreateConnectionRequest(\n", - " {\n", - " \"parent\": new_conn_parent,\n", - " \"connection_id\": CONN_NAME,\n", - " \"connection\": connection,\n", - " }\n", - " )\n", - " response = client.create_connection(request)\n", - " CONN_SERVICE_ACCOUNT = (\n", - " f\"serviceAccount:{response.cloud_resource.service_account_id}\"\n", - " )\n", - "print(CONN_SERVICE_ACCOUNT)" - ] - }, - { - "cell_type": "markdown", - "metadata": { - "id": "W6l6Ol2biU9h" - }, - "source": [ - "## Set permissions for the service account\n", - "\n", - "The resource connection service account requires certain project-level permissions:\n", - " - `roles/aiplatform.user` and `roles/bigquery.connectionUser`: These roles are required for the connection to create a model definition using the LLM model in Vertex AI ([documentation](https://cloud.google.com/bigquery/docs/generate-text#give_the_service_account_access)).\n", - " - `roles/run.invoker`: This role is required for the connection to have read-only access to Cloud Run services that back custom/remote functions ([documentation](https://cloud.google.com/bigquery/docs/remote-functions#grant_permission_on_function)).\n", - "\n", - "Set these permissions by running the following `gcloud` commands:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": { - "id": "d8wja24SVq6s" - }, - "outputs": [], - "source": [ - "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/bigquery.connectionUser'\n", - "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/aiplatform.user'\n", - "!gcloud projects add-iam-policy-binding {PROJECT_ID} --condition=None --no-user-output-enabled --member={CONN_SERVICE_ACCOUNT} --role='roles/run.invoker'" - ] - }, { "cell_type": "markdown", "metadata": { @@ -468,9 +391,7 @@ "source": [ "from bigframes.ml.llm import PaLM2TextGenerator\n", "\n", - "session = bf.get_global_session()\n", - "connection = f\"{PROJECT_ID}.{REGION}.{CONN_NAME}\"\n", - "model = PaLM2TextGenerator(session=session, connection_name=connection)" + "model = PaLM2TextGenerator()" ] }, { @@ -651,7 +572,7 @@ }, "outputs": [], "source": [ - "@bf.remote_function([str], str, bigquery_connection=CONN_NAME)\n", + "@bf.remote_function([str], str)\n", "def extract_code(text: str):\n", " try:\n", " res = text[text.find('\\n')+1:text.find('```', 3)]\n",