Be able to use Codey models on Vertex AI (#6354)

Added the functionality to leverage 3 new Codey models from Vertex AI:
- code-bison - Code generation using the existing LLM integration
- code-gecko - Code completion using the existing LLM integration
- codechat-bison - Code chat using the existing chat_model integration

---------

Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>

docs/extras/modules/model_io/models/chat/integrations/google_vertex_ai_palm.ipynb

@@ -141,6 +141,73 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:09:25.423568Z",
+     "iopub.status.busy": "2023-06-17T21:09:25.423213Z",
+     "iopub.status.idle": "2023-06-17T21:09:25.429641Z",
+     "shell.execute_reply": "2023-06-17T21:09:25.429060Z",
+     "shell.execute_reply.started": "2023-06-17T21:09:25.423546Z"
+    },
+    "tags": []
+   },
+   "source": [
+    "You can now leverage the Codey API for code chat within Vertex AI. The model name is:\n",
+    "- codechat-bison: for code assistance"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:30:43.974841Z",
+     "iopub.status.busy": "2023-06-17T21:30:43.974431Z",
+     "iopub.status.idle": "2023-06-17T21:30:44.248119Z",
+     "shell.execute_reply": "2023-06-17T21:30:44.247362Z",
+     "shell.execute_reply.started": "2023-06-17T21:30:43.974820Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "chat = ChatVertexAI(model_name=\"codechat-bison\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:30:45.146093Z",
+     "iopub.status.busy": "2023-06-17T21:30:45.145752Z",
+     "iopub.status.idle": "2023-06-17T21:30:47.449126Z",
+     "shell.execute_reply": "2023-06-17T21:30:47.448609Z",
+     "shell.execute_reply.started": "2023-06-17T21:30:45.146069Z"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content='The following Python function can be used to identify all prime numbers up to a given integer:\\n\\n```\\ndef is_prime(n):\\n  \"\"\"\\n  Determines whether the given integer is prime.\\n\\n  Args:\\n    n: The integer to be tested for primality.\\n\\n  Returns:\\n    True if n is prime, False otherwise.\\n  \"\"\"\\n\\n  # Check if n is divisible by 2.\\n  if n % 2 == 0:\\n    return False\\n\\n  # Check if n is divisible by any integer from 3 to the square root', additional_kwargs={}, example=False)"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "messages = [\n",
+    "    HumanMessage(content=\"How do I create a python function to identify all prime numbers?\")\n",
+    "]\n",
+    "chat(messages)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,

docs/extras/modules/model_io/models/llms/integrations/google_vertex_ai_palm.ipynb

@@ -101,11 +101,80 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
+   "cell_type": "markdown",
    "metadata": {},
+   "source": [
+    "You can now leverage the Codey API for code generation within Vertex AI. The model names are:\n",
+    "- code-bison: for code suggestion\n",
+    "- code-gecko: for code completion"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:16:53.149438Z",
+     "iopub.status.busy": "2023-06-17T21:16:53.149065Z",
+     "iopub.status.idle": "2023-06-17T21:16:53.421824Z",
+     "shell.execute_reply": "2023-06-17T21:16:53.421136Z",
+     "shell.execute_reply.started": "2023-06-17T21:16:53.149415Z"
+    },
+    "tags": []
+   },
    "outputs": [],
-   "source": []
+   "source": [
+    "llm = VertexAI(model_name=\"code-bison\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:17:11.179077Z",
+     "iopub.status.busy": "2023-06-17T21:17:11.178686Z",
+     "iopub.status.idle": "2023-06-17T21:17:11.182499Z",
+     "shell.execute_reply": "2023-06-17T21:17:11.181895Z",
+     "shell.execute_reply.started": "2023-06-17T21:17:11.179052Z"
+    },
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "llm_chain = LLMChain(prompt=prompt, llm=llm)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "execution": {
+     "iopub.execute_input": "2023-06-17T21:18:47.024785Z",
+     "iopub.status.busy": "2023-06-17T21:18:47.024230Z",
+     "iopub.status.idle": "2023-06-17T21:18:49.352249Z",
+     "shell.execute_reply": "2023-06-17T21:18:49.351695Z",
+     "shell.execute_reply.started": "2023-06-17T21:18:47.024762Z"
+    },
+    "tags": []
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "'```python\\ndef is_prime(n):\\n  \"\"\"\\n  Determines if a number is prime.\\n\\n  Args:\\n    n: The number to be tested.\\n\\n  Returns:\\n    True if the number is prime, False otherwise.\\n  \"\"\"\\n\\n  # Check if the number is 1.\\n  if n == 1:\\n    return False\\n\\n  # Check if the number is 2.\\n  if n == 2:\\n    return True\\n\\n'"
+      ]
+     },
+     "execution_count": 15,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "question = \"Write a python function that identifies if the number is a prime number?\"\n",
+    "\n",
+    "llm_chain.run(question)"
+   ]
   }
  ],
  "metadata": {