{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Google Cloud Vertex AI \n",
    "\n",
    "Note: This is separate from the Google PaLM integration. Google has chosen to offer an enterprise version of PaLM through GCP, and this supports the models made available through there. \n",
    "\n",
    "By default, Google Cloud [does not use](https://cloud.google.com/vertex-ai/docs/generative-ai/data-governance#foundation_model_development) customer data to train its foundation models as part of Google Cloud`s AI/ML Privacy Commitment. More details about how Google processes data can also be found in [Google's Customer Data Processing Addendum (CDPA)](https://cloud.google.com/terms/data-processing-addendum).\n",
    "\n",
    "To use `Google Cloud Vertex AI` PaLM you must have the `google-cloud-aiplatform` Python package installed and either:\n",
    "- Have credentials configured for your environment (gcloud, workload identity, etc...)\n",
    "- Store the path to a service account JSON file as the GOOGLE_APPLICATION_CREDENTIALS environment variable\n",
    "\n",
    "This codebase uses the `google.auth` library which first looks for the application credentials variable mentioned above, and then looks for system-level auth.\n",
    "\n",
    "For more information, see: \n",
    "- https://cloud.google.com/docs/authentication/application-default-credentials#GAC\n",
    "- https://googleapis.dev/python/google-auth/latest/reference/google.auth.html#module-google.auth\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "#!pip install langchain google-cloud-aiplatform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.chat_models import ChatVertexAI\n",
    "from langchain.prompts import ChatPromptTemplate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = ChatVertexAI()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = \"You are a helpful assistant who translate English to French\"\n",
    "human = \"Translate this sentence from English to French. I love programming.\"\n",
    "prompt = ChatPromptTemplate.from_messages([(\"system\", system), (\"human\", human)])\n",
    "messages = prompt.format_messages()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=\" J'aime la programmation.\", additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat(messages)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If we want to construct a simple chain that takes user specified parameters:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "system = (\n",
    "    \"You are a helpful assistant that translates {input_language} to {output_language}.\"\n",
    ")\n",
    "human = \"{text}\"\n",
    "prompt = ChatPromptTemplate.from_messages([(\"system\", system), (\"human\", human)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' 私はプログラミングが大好きです。', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chain = prompt | chat\n",
    "chain.invoke(\n",
    "    {\n",
    "        \"input_language\": \"English\",\n",
    "        \"output_language\": \"Japanese\",\n",
    "        \"text\": \"I love programming\",\n",
    "    }\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "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": [
    "## Code generation chat models\n",
    "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": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "chat = ChatVertexAI(\n",
    "    model_name=\"codechat-bison\", max_output_tokens=1000, temperature=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " ```python\n",
      "def is_prime(x): \n",
      "    if (x <= 1): \n",
      "        return False\n",
      "    for i in range(2, x): \n",
      "        if (x % i == 0): \n",
      "            return False\n",
      "    return True\n",
      "```\n"
     ]
    }
   ],
   "source": [
    "# For simple string in string out usage, we can use the `predict` method:\n",
    "print(chat.predict(\"Write a Python function to identify all prime numbers\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Asynchronous calls\n",
    "\n",
    "We can make asynchronous calls via the `agenerate` and `ainvoke` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "import asyncio\n",
    "\n",
    "# import nest_asyncio\n",
    "# nest_asyncio.apply()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LLMResult(generations=[[ChatGeneration(text=\" J'aime la programmation.\", generation_info=None, message=AIMessage(content=\" J'aime la programmation.\", additional_kwargs={}, example=False))]], llm_output={}, run=[RunInfo(run_id=UUID('223599ef-38f8-4c79-ac6d-a5013060eb9d'))])"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "chat = ChatVertexAI(\n",
    "    model_name=\"chat-bison\",\n",
    "    max_output_tokens=1000,\n",
    "    temperature=0.7,\n",
    "    top_p=0.95,\n",
    "    top_k=40,\n",
    ")\n",
    "\n",
    "asyncio.run(chat.agenerate([messages]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content=' अहं प्रोग्रामिंग प्रेमामि', additional_kwargs={}, example=False)"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "asyncio.run(\n",
    "    chain.ainvoke(\n",
    "        {\n",
    "            \"input_language\": \"English\",\n",
    "            \"output_language\": \"Sanskrit\",\n",
    "            \"text\": \"I love programming\",\n",
    "        }\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Streaming calls\n",
    "\n",
    "We can also stream outputs via the `stream` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      " 1. China (1,444,216,107)\n",
      "2. India (1,393,409,038)\n",
      "3. United States (332,403,650)\n",
      "4. Indonesia (273,523,615)\n",
      "5. Pakistan (220,892,340)\n",
      "6. Brazil (212,559,409)\n",
      "7. Nigeria (206,139,589)\n",
      "8. Bangladesh (164,689,383)\n",
      "9. Russia (145,934,462)\n",
      "10. Mexico (128,932,488)\n",
      "11. Japan (126,476,461)\n",
      "12. Ethiopia (115,063,982)\n",
      "13. Philippines (109,581,078)\n",
      "14. Egypt (102,334,404)\n",
      "15. Vietnam (97,338,589)"
     ]
    }
   ],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [(\"human\", \"List out the 15 most populous countries in the world\")]\n",
    ")\n",
    "messages = prompt.format_messages()\n",
    "for chunk in chat.stream(messages):\n",
    "    sys.stdout.write(chunk.content)\n",
    "    sys.stdout.flush()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "poetry-venv",
   "language": "python",
   "name": "poetry-venv"
  },
  "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.9.1"
  },
  "vscode": {
   "interpreter": {
    "hash": "cc99336516f23363341912c6723b01ace86f02e26b4290be1efc0677e2e2ec24"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
