langchain/docs/docs/integrations/llms/ai21.ipynb

{
 "cells": [
  {
   "cell_type": "raw",
   "id": "602a52a4",
   "metadata": {},
   "source": [
    "---\n",
    "sidebar_label: AI21 Labs\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9597802c",
   "metadata": {},
   "source": [
    "# AI21LLM\n",
    "\n",
    "This example goes over how to use LangChain to interact with `AI21` Jurassic models. To use the Jamba model, use the [ChatAI21 object](https://python.langchain.com/docs/integrations/chat/ai21/) instead.\n",
    "\n",
    "[See a full list of AI21 models and tools on LangChain.](https://pypi.org/project/langchain-ai21/)\n",
    "\n",
    "## Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "59c710c4",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-05T20:58:42.397591Z",
     "start_time": "2024-03-05T20:58:40.944729Z"
    }
   },
   "outputs": [],
   "source": [
    "!pip install -qU langchain-ai21"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "560a2f9254963fd7",
   "metadata": {
    "collapsed": false
   },
   "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:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "035dea0f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-05T20:58:44.465443Z",
     "start_time": "2024-03-05T20:58:42.399724Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "if \"AI21_API_KEY\" not in os.environ:\n",
    "    os.environ[\"AI21_API_KEY\"] = getpass()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1891df96eb076e1a",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "## Usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "98f70927a87e4745",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-05T20:58:45.859265Z",
     "start_time": "2024-03-05T20:58:44.466637Z"
    },
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'\\nLangChain is a (database)\\nLangChain is a database for storing and processing documents'"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain_ai21 import AI21LLM\n",
    "from langchain_core.prompts import PromptTemplate\n",
    "\n",
    "template = \"\"\"Question: {question}\n",
    "\n",
    "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": "markdown",
   "id": "9965c10269159ed1",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "# AI21 Contextual Answer\n",
    "\n",
    "You can use AI21's contextual answers model to receives text or document, serving as a context,\n",
    "and a question and returns an answer based entirely on this context.\n",
    "\n",
    "This means that if the answer to your question is not in the document,\n",
    "the model will indicate it (instead of providing a false answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "411adf42eab80829",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-05T20:59:00.943426Z",
     "start_time": "2024-03-05T20:59:00.263497Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from langchain_ai21 import AI21ContextualAnswers\n",
    "\n",
    "tsm = AI21ContextualAnswers()\n",
    "\n",
    "response = tsm.invoke(input={\"context\": \"Your context\", \"question\": \"Your question\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af59ffdbf4964875",
   "metadata": {
    "collapsed": false
   },
   "source": [
    "You can also use it with chains and output parsers and vector DBs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "bc63830f921b4ac9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-03-05T20:59:07.719225Z",
     "start_time": "2024-03-05T20:59:07.102950Z"
    },
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from langchain_ai21 import AI21ContextualAnswers\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "\n",
    "tsm = AI21ContextualAnswers()\n",
    "chain = tsm | StrOutputParser()\n",
    "\n",
    "response = chain.invoke(\n",
    "    {\"context\": \"Your context\", \"question\": \"Your question\"},\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.11.1 64-bit",
   "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"
  },
  "vscode": {
   "interpreter": {
    "hash": "e971737741ff4ec9aff7dc6155a1060a59a8a6d52c757dbbe66bf8ee389494b1"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}