{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use BigQuery DataFrames to run Anthropic LLM at scale\n",
    "\n",
    "<table align=\"left\">\n",
    "\n",
    "  <td>\n",
    "    <a href=\"https://colab.research.google.com/github/googleapis/python-bigquery-dataframes/blob/main/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb\">\n",
    "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/colab-logo.png\" alt=\"Colab logo\"> Run in Colab\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb\">\n",
    "      <img src=\"https://raw.githubusercontent.com/googleapis/python-bigquery-dataframes/refs/heads/main/third_party/logo/github-logo.png\" width=\"32\" alt=\"GitHub logo\">\n",
    "      View on GitHub\n",
    "    </a>\n",
    "  </td>\n",
    "  <td>\n",
    "    <a href=\"https://console.cloud.google.com/bigquery/import?url=https://github.com/googleapis/python-bigquery-dataframes/blob/main/notebooks/remote_functions/remote_function_vertex_claude_model.ipynb\">\n",
    "      <img src=\"https://encrypted-tbn0.gstatic.com/images?q=tbn:ANd9GcTW1gvOovVlbZAIZylUtf5Iu8-693qS1w5NJw&s\" alt=\"BQ logo\" width=\"35\">\n",
    "      Open in BQ Studio\n",
    "    </a>\n",
    "  </td>\n",
    "</table>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "\n",
    "# Python 3.13 is not yet a supported runtime for remote functions.\n",
    "# See: https://cloud.google.com/functions/docs/runtime-support#python for the supported runtimes.\n",
    "if sys.version_info >= (3, 13, 0):\n",
    "    sys.exit(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overview\n",
    "\n",
    "Anthropic Claude models are available as APIs on Vertex AI ([docs](https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude)).\n",
    "\n",
    "To run the Claude models at large scale data we can utilze the BigQuery\n",
    "DataFrames remote functions ([docs](https://cloud.google.com/bigquery/docs/use-bigquery-dataframes#remote-functions)).\n",
    "BigQuery DataFrames provides a simple pythonic interface `remote_function` to\n",
    "deploy the user code as a BigQuery remote function and then invoke it at scale\n",
    "by utilizing the parallel distributed computing architecture of BigQuery and\n",
    "Google Cloud Function.\n",
    "\n",
    "In this notebook we showcase one such example. For the demonstration purpose we\n",
    "use a small amount of data, but the example generalizes for large data. Check out\n",
    "various IO APIs provided by BigQuery DataFrames [here](https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.pandas#bigframes_pandas_read_gbq)\n",
    "to see how you could create a DataFrame from your Big Data sitting in a BigQuery\n",
    "table or GCS bucket."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up a claude model in Vertex\n",
    "\n",
    "https://cloud.google.com/vertex-ai/generative-ai/docs/partner-models/use-claude#before_you_begin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Install Anthropic with Vertex if needed\n",
    "\n",
    "Uncomment the following cell and run the cell to install anthropic python\n",
    "package with vertex extension if you don't already have it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip install anthropic[vertex] --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define project and location for GCP integration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "PROJECT = \"bigframes-dev\" # replace with your project\n",
    "LOCATION = \"us-east5\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Initialize BigQuery DataFrames dataframe\n",
    "\n",
    "BigQuery DataFrames is a set of open source Python libraries that let you take\n",
    "advantage of BigQuery data processing by using familiar Python APIs.\n",
    "See for more details https://cloud.google.com/bigquery/docs/bigquery-dataframes-introduction."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import BigQuery DataFrames pandas module and initialize it with your project\n",
    "# and location\n",
    "\n",
    "import bigframes.pandas as bpd\n",
    "bpd.options.bigquery.project = PROJECT\n",
    "bpd.options.bigquery.location = LOCATION"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's use a DataFrame with small amount of inline data for demo purpose.\n",
    "You could create a DataFrame from your own data. See APIs like `read_gbq`,\n",
    "`read_csv`, `read_json` etc. at https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.pandas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Query job c4c27713-51c8-4293-8454-5c904df79318 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:c4c27713-51c8-4293-8454-5c904df79318&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Query job 0b1b71d8-8546-45f2-b403-707161fe4002 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:0b1b71d8-8546-45f2-b403-707161fe4002&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>questions</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What is the capital of France?</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Explain the concept of photosynthesis in simpl...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Write a haiku about artificial intelligence.</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 1 columns</p>\n",
       "</div>[3 rows x 1 columns in total]"
      ],
      "text/plain": [
       "                                           questions\n",
       "0                     What is the capital of France?\n",
       "1  Explain the concept of photosynthesis in simpl...\n",
       "2       Write a haiku about artificial intelligence.\n",
       "\n",
       "[3 rows x 1 columns]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = bpd.DataFrame({\"questions\": [\n",
    "   \"What is the capital of France?\",\n",
    "   \"Explain the concept of photosynthesis in simple terms.\",\n",
    "   \"Write a haiku about artificial intelligence.\"\n",
    " ]})\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Use BigQuery DataFrames `remote_function`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a remote function from a custom python function that takes a prompt\n",
    "and returns the output of the claude LLM running in Vertex. We will be using\n",
    "`max_batching_rows=1` to control parallelization. This ensures that a single\n",
    "prompt is processed per batch in the underlying cloud function so that the batch\n",
    "processing does not time out. An ideal value for `max_batching_rows` depends on\n",
    "the complexity of the prompts in the real use case and should be discovered\n",
    "through offline experimentation. Check out the API for other ways to control\n",
    "parallelization https://cloud.google.com/python/docs/reference/bigframes/latest/bigframes.pandas#bigframes_pandas_remote_function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Query job 488a116f-44b2-4ff7-9f95-bd36473dab0f is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:488a116f-44b2-4ff7-9f95-bd36473dab0f&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "@bpd.remote_function(packages=[\"anthropic[vertex]\", \"google-auth[requests]\"],\n",
    "                     max_batching_rows=1, \n",
    "                     bigquery_connection=\"bigframes-dev.us-east5.bigframes-rf-conn\", # replace with your connection\n",
    "                     cloud_function_service_account=\"default\",\n",
    ")\n",
    "def anthropic_transformer(message: str) -> str:\n",
    "  from anthropic import AnthropicVertex\n",
    "  client = AnthropicVertex(region=LOCATION, project_id=PROJECT)\n",
    "\n",
    "  message = client.messages.create(\n",
    "              max_tokens=1024,\n",
    "              messages=[\n",
    "                  {\n",
    "                      \"role\": \"user\",\n",
    "                      \"content\": message,\n",
    "                  }\n",
    "              ],\n",
    "              model=\"claude-3-5-sonnet@20240620\",\n",
    "          )\n",
    "  content_text = message.content[0].text if message.content else \"\"\n",
    "  return content_text"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'bigframes-dev._b52b272a35b88e236e1f96fbe3f560c83a8fee85.bigframes_session265649_de1176dd4c57f40ba959503af3981682'"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Print the BigQuery remote function created\n",
    "anthropic_transformer.bigframes_remote_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'projects/bigframes-dev/locations/us-east5/functions/bigframes-session265649-de1176dd4c57f40ba959503af3981682'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Print the cloud function created\n",
    "anthropic_transformer.bigframes_cloud_function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "Query job 58b230a8-6536-4bac-ab02-dcf574692dd6 is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:58b230a8-6536-4bac-ab02-dcf574692dd6&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Query job 46d6a1e9-426a-4615-8eb5-98d34d08ec07 is DONE. 1.3 kB processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:46d6a1e9-426a-4615-8eb5-98d34d08ec07&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Query job cf8fcbaa-b233-47cd-b4e3-60876b24879f is DONE. 0 Bytes processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:cf8fcbaa-b233-47cd-b4e3-60876b24879f&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Query job 2620a544-d1df-4b30-bec4-4221e79ddf0f is DONE. 1.4 kB processed. <a target=\"_blank\" href=\"https://console.cloud.google.com/bigquery?project=bigframes-dev&j=bq:us-east5:2620a544-d1df-4b30-bec4-4221e79ddf0f&page=queryresults\">Open Job</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>questions</th>\n",
       "      <th>answers</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>What is the capital of France?</td>\n",
       "      <td>The capital of France is Paris.</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>Explain the concept of photosynthesis in simpl...</td>\n",
       "      <td>Photosynthesis is the process plants use to ma...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Write a haiku about artificial intelligence.</td>\n",
       "      <td>Here's a haiku about artificial intelligence:\n",
       "...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>3 rows × 2 columns</p>\n",
       "</div>[3 rows x 2 columns in total]"
      ],
      "text/plain": [
       "                                           questions  \\\n",
       "0                     What is the capital of France?   \n",
       "1  Explain the concept of photosynthesis in simpl...   \n",
       "2       Write a haiku about artificial intelligence.   \n",
       "\n",
       "                                             answers  \n",
       "0                    The capital of France is Paris.  \n",
       "1  Photosynthesis is the process plants use to ma...  \n",
       "2  Here's a haiku about artificial intelligence:\n",
       "...  \n",
       "\n",
       "[3 rows x 2 columns]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Apply the remote function on the user data\n",
    "df[\"answers\"] = df[\"questions\"].apply(anthropic_transformer)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Clean Up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bpd.close_session()"
   ]
  }
 ],
 "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
