ai21: init package (#17592)

Co-authored-by: Asaf Gardin <[email protected]> Co-authored-by: etang <[email protected]> Co-authored-by: asafgardin <[email protected]>
langchain-ai · Feb 15, 2024 · 6cc6faa · 6cc6faa
1 parent 20a56fe
commit 6cc6faa
Show file tree

Hide file tree

Showing 32 changed files with 2,685 additions and 75 deletions.
diff --git a/.github/workflows/_integration_test.yml b/.github/workflows/_integration_test.yml
@@ -52,6 +52,7 @@ jobs:
       - name: Run integration tests
         shell: bash
         env:
+          AI21_API_KEY: ${{ secrets.AI21_API_KEY }}
           GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
           MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}

diff --git a/.github/workflows/_release.yml b/.github/workflows/_release.yml
@@ -166,6 +166,7 @@ jobs:
       - name: Run integration tests
         if: ${{ startsWith(inputs.working-directory, 'libs/partners/') }}
         env:
+          AI21_API_KEY: ${{ secrets.AI21_API_KEY }}
           GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
           ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
           MISTRAL_API_KEY: ${{ secrets.MISTRAL_API_KEY }}

diff --git a/docs/docs/integrations/chat/ai21.ipynb b/docs/docs/integrations/chat/ai21.ipynb
@@ -0,0 +1,131 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "e49f1e0d",
+   "metadata": {},
+   "source": [
+    "# ChatAI21\n",
+    "\n",
+    "This notebook covers how to get started with AI21 chat models.\n",
+    "\n",
+    "## Installation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "4c3bef91",
+   "metadata": {
+    "ExecuteTime": {
+     "end_time": "2024-02-15T06:50:44.929635Z",
+     "start_time": "2024-02-15T06:50:41.209704Z"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "!pip install -qU langchain-ai21"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "2b4f3e15",
+   "metadata": {},
+   "source": [
+    "## Environment Setup\n",
+    "\n",
+    "We'll need to get a [AI21 API key](https://docs.ai21.com/) and set the `AI21_API_KEY` environment variable:\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "62e0dbc3",
+   "metadata": {
+    "tags": []
+   },
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "from getpass import getpass\n",
+    "\n",
+    "os.environ[\"AI21_API_KEY\"] = getpass()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4828829d3da430ce",
+   "metadata": {
+    "collapsed": false
+   },
+   "source": [
+    "## Usage"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "39353473fce5dd2e",
+   "metadata": {
+    "collapsed": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "AIMessage(content='Bonjour, comment vas-tu?')"
+      ]
+     },
+     "execution_count": 1,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from langchain_ai21 import ChatAI21\n",
+    "from langchain_core.prompts import ChatPromptTemplate\n",
+    "\n",
+    "chat = ChatAI21(model=\"j2-ultra\")\n",
+    "\n",
+    "prompt = ChatPromptTemplate.from_messages(\n",
+    "    [\n",
+    "        (\"system\", \"You are a helpful assistant that translates English to French.\"),\n",
+    "        (\"human\", \"Translate this sentence from English to French. {english_text}.\"),\n",
+    "    ]\n",
+    ")\n",
+    "\n",
+    "chain = prompt | chat\n",
+    "chain.invoke({\"english_text\": \"Hello, how are you?\"})"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c159a79f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.11.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/docs/docs/integrations/llms/ai21.ipynb b/docs/docs/integrations/llms/ai21.ipynb
@@ -5,141 +5,108 @@
    "id": "9597802c",
    "metadata": {},
    "source": [
-    "# AI21\n",
+    "# AI21LLM\n",
     "\n",
-    "[AI21 Studio](https://docs.ai21.com/) provides API access to `Jurassic-2` large language models.\n",
+    "This example goes over how to use LangChain to interact with `AI21` models.\n",
     "\n",
-    "This example goes over how to use LangChain to interact with [AI21 models](https://docs.ai21.com/docs/jurassic-2-models)."
+    "## Installation"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
-   "id": "02be122d-04e8-4ec6-84d1-f1d8961d6828",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "\u001b[33mWARNING: There was an error checking the latest version of pip.\u001b[0m\u001b[33m\n",
-      "\u001b[0mNote: you may need to restart the kernel to use updated packages.\n"
-     ]
-    }
-   ],
-   "source": [
-    "# install the package:\n",
-    "%pip install --upgrade --quiet  ai21"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 1,
-   "id": "4229227e-6ca2-41ad-a3c3-5f29e3559091",
-   "metadata": {
-    "tags": []
-   },
+   "execution_count": null,
+   "id": "59c710c4",
+   "metadata": {},
    "outputs": [],
    "source": [
-    "# get AI21_API_KEY. Use https://studio.ai21.com/account/account\n",
-    "\n",
-    "from getpass import getpass\n",
-    "\n",
-    "AI21_API_KEY = getpass()"
+    "!pip install -qU langchain-ai21"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 7,
-   "id": "6fb585dd",
+   "cell_type": "markdown",
+   "id": "560a2f9254963fd7",
    "metadata": {
-    "tags": []
+    "collapsed": false
    },
-   "outputs": [],
    "source": [
-    "from langchain_community.llms import AI21\n",
-    "from langchain_core.prompts import PromptTemplate"
+    "## Environment Setup\n",
+    "\n",
+    "We'll need to get a [AI21 API key](https://docs.ai21.com/) and set the `AI21_API_KEY` environment variable:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 4,
    "id": "035dea0f",
    "metadata": {
     "tags": []
    },
    "outputs": [],
    "source": [
-    "template = \"\"\"Question: {question}\n",
-    "\n",
-    "Answer: Let's think step by step.\"\"\"\n",
+    "import os\n",
+    "from getpass import getpass\n",
     "\n",
-    "prompt = PromptTemplate.from_template(template)"
+    "os.environ[\"AI21_API_KEY\"] = getpass()"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 9,
-   "id": "3f3458d9",
-   "metadata": {
-    "tags": []
-   },
-   "outputs": [],
-   "source": [
-    "llm = AI21(ai21_api_key=AI21_API_KEY)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 10,
-   "id": "a641dbd9",
+   "cell_type": "markdown",
+   "id": "1891df96eb076e1a",
    "metadata": {
-    "tags": []
+    "collapsed": false
    },
-   "outputs": [],
    "source": [
-    "llm_chain = prompt | llm"
+    "## Usage"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
-   "id": "9f0b1960",
+   "execution_count": 6,
+   "id": "98f70927a87e4745",
    "metadata": {
-    "tags": []
+    "collapsed": false
    },
    "outputs": [
     {
      "data": {
       "text/plain": [
-       "'\\nThe Super Bowl in the year Justin Beiber was born was in the year 1991.\\nThe Super Bowl in 1991 was won by the Washington Redskins.\\nFinal answer: Washington Redskins'"
+       "'\\nLangChain is a decentralized blockchain network that leverages AI and machine learning to provide language translation services.'"
       ]
      },
-     "execution_count": 13,
+     "execution_count": 6,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "question = \"What NFL team won the Super Bowl in the year Justin Beiber was born?\"\n",
+    "from langchain_ai21 import AI21LLM\n",
+    "from langchain_core.prompts import PromptTemplate\n",
+    "\n",
+    "template = \"\"\"Question: {question}\n",
     "\n",
-    "llm_chain.invoke({\"question\": question})"
+    "Answer: Let's think step by step.\"\"\"\n",
+    "\n",
+    "prompt = PromptTemplate.from_template(template)\n",
+    "\n",
+    "model = AI21LLM(model=\"j2-ultra\")\n",
+    "\n",
+    "chain = prompt | model\n",
+    "\n",
+    "chain.invoke({\"question\": \"What is LangChain?\"})"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "22bce013",
+   "id": "a52f765c",
    "metadata": {},
    "outputs": [],
    "source": []
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.11.1 64-bit",
    "language": "python",
    "name": "python3"
   },
@@ -153,7 +120,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.13"
+   "version": "3.11.4"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "e971737741ff4ec9aff7dc6155a1060a59a8a6d52c757dbbe66bf8ee389494b1"
+   }
   }
  },
  "nbformat": 4,