llamaindex with bedrock

llamaindex-bedrock.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "39894001-c5a8-4fd0-9070-848108dd8fbf",
   "metadata": {},
   "source": [
    "## LlamaIndex <> Bedrock Integrations\n",
    "\n",
    "This demo notebook shows how you can use LlamaIndex together with Amazon Bedrock. It provides following examples as mentioned in the [doc](https://docs.llamaindex.ai/en/latest/examples/llm/bedrock.html#):\n",
    "- Basic usage with `complete` and `stream_complete`\n",
    "- Basic usage with `chat` and `stream_chat`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f1a410d-8eac-473c-9f5f-be96f70435eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install boto3 -q"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4897bea-a3bf-4b34-9323-5fb3420ee2aa",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.llms import Bedrock, ChatMessage\n",
    "import boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f48d427d-24ff-42ad-9103-b9d2e65e879d",
   "metadata": {},
   "outputs": [],
   "source": [
    "profile_name = \"us-east-1\"  # Update it with your AWS Profile name\n",
    "model = \"amazon.titan-text-express-v1\"\n",
    "\n",
    "llm = Bedrock(model=model, profile_name=profile_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b179cdc-bb98-4ef2-9944-9cbd97e6ce64",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "CompletionResponse(text=\"the founder, CEO, and chief engineer of SpaceX; an angel investor, CEO, and product architect of Tesla, Inc.; Twitter's owner and CTO; the Boring Company's creator; the founder, CEO, and chief engineer of Neuralink and ******; and the co-founder of the X-Corp. and the PayPal Holdings, Inc. conglomerate. He was born in Pretoria, South Africa, and immigrated to the United States in 1989. Musk is known for his ambitious and controversial projects, including the development of electric vehicles, solar energy, and space exploration. He is also a prominent figure in the technology industry, with a net worth of over $200 billion.\", additional_kwargs={}, raw={'inputTextTokenCount': 5, 'results': [{'tokenCount': 149, 'outputText': \"the founder, CEO, and chief engineer of SpaceX; an angel investor, CEO, and product architect of Tesla, Inc.; Twitter's owner and CTO; the Boring Company's creator; the founder, CEO, and chief engineer of Neuralink and ******; and the co-founder of the X-Corp. and the PayPal Holdings, Inc. conglomerate. He was born in Pretoria, South Africa, and immigrated to the United States in 1989. Musk is known for his ambitious and controversial projects, including the development of electric vehicles, solar energy, and space exploration. He is also a prominent figure in the technology industry, with a net worth of over $200 billion.\", 'completionReason': 'FINISH'}]}, delta=None)"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "llm.complete(\"Elon Musk is \")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8c85f23-16a3-4cb2-8b45-614fbba1ec52",
   "metadata": {},
   "source": [
    "### Stream complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abe21336-1e8f-4be1-b443-1d08a719401a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "son of a preacher and grew up in a small town in Iowa. He was a rebellious teenager who often got into trouble. Despite his struggles, he found solace in music and began playing the guitar.\n",
      "\n",
      "At the age of 16, Paul moved to Boston to pursue his music career. He was rejected by many record labels, but he continued to play and write songs. He eventually signed a deal with Columbia Records and released his debut album, \"The Beatles,\" in 1963.\n",
      "\n",
      "The album was an instant hit and launched Paul McCartney into stardom. He became known for his songwriting and his unique voice, which blended pop and rock. He went on to release many successful albums, including \"Let It Be,\" \"Abbey Road,\" and \"Yesterday.\"\n",
      "\n",
      "Paul McCartney is one of the most successful musicians of all time. He has sold over 200 million records worldwide and has received numerous awards, including 18 Grammy Awards. He is also known for his philanthropic work, including his support for animal rights and the fight against poverty.\n",
      "\n",
      "Paul McCartney is a living legend who has left an indelible mark on the music industry. His music continues to inspire generations of musicians and fans, and he remains one of the most beloved and respected figures in popular culture."
     ]
    }
   ],
   "source": [
    "resp = llm.stream_complete(\"Paul Graham is \")\n",
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b66667be-bd5c-40e3-8757-79128831da27",
   "metadata": {},
   "source": [
    "## Chat with a list of messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e24244-7768-4ec6-9b0f-5cc37a05966d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ChatResponse(message=ChatMessage(role=<MessageRole.ASSISTANT: 'assistant'>, content='Once upon a time, there was a pirate named Captain Blackbeard. He was a notorious pirate who sailed the high seas with his crew, searching for treasure and adventure.\\n\\nOne day, Captain Blackbeard and his crew stumbled upon a Spanish galleon. The galleon was loaded with gold, silver, and precious jewels, and Captain Blackbeard knew that he had to have it. He ordered his crew to board the galleon and take everything they could.\\n\\nThe boarding was fierce, and the Spanish sailors fought back with all their might. But Captain Blackbeard was a skilled sailor, and he managed to outmaneuver his opponents. In the end, Captain Blackbeard and his crew took the galleon and sailed back to their ship, laden with treasure.\\n\\nCaptain Blackbeard was known for his ruthless and cunning nature, and he was feared by many sailors on the high seas. He was also known for his love of luxury and his lavish lifestyle. He would often dress in fancy clothes and wear expensive jewelry, and he would entertain his crew with grand parties and feasts.\\n\\nBut despite his fearsome reputation, Captain Blackbeard was a fair and just leader. He treated his crew with respect and kindness, and he always ensured that they were paid well for their hard work. He was also known for his generosity, and he would often share his treasure with the people of the towns that he visited.\\n\\nOne day, Captain Blackbeard and his crew were sailing off the coast of North Carolina when they were spotted by a British warship. The British ship was armed with cannons and had a large crew of sailors. Captain Blackbeard knew that he was in trouble, but he was not afraid. He ordered his crew to fire at the British ship, and the two ships fought for hours.\\n\\nIn the end, Captain Blackbeard and his crew were defeated. They were captured and taken to London, where they were tried and convicted of piracy. Captain Blackbeard was sentenced to death, but he was pardoned by the British government and allowed to return to his home in the Caribbean.\\n\\nCaptain Blackbeard lived the rest of his life in peace and luxury, and he is remembered as one of the most famous pirates of all time. He was a skilled sailor, a ruthless pirate, and a generous leader, and he will always be remembered as a legend of the high seas.', additional_kwargs={}), raw={'inputTextTokenCount': 31, 'results': [{'tokenCount': 505, 'outputText': 'Once upon a time, there was a pirate named Captain Blackbeard. He was a notorious pirate who sailed the high seas with his crew, searching for treasure and adventure.\\n\\nOne day, Captain Blackbeard and his crew stumbled upon a Spanish galleon. The galleon was loaded with gold, silver, and precious jewels, and Captain Blackbeard knew that he had to have it. He ordered his crew to board the galleon and take everything they could.\\n\\nThe boarding was fierce, and the Spanish sailors fought back with all their might. But Captain Blackbeard was a skilled sailor, and he managed to outmaneuver his opponents. In the end, Captain Blackbeard and his crew took the galleon and sailed back to their ship, laden with treasure.\\n\\nCaptain Blackbeard was known for his ruthless and cunning nature, and he was feared by many sailors on the high seas. He was also known for his love of luxury and his lavish lifestyle. He would often dress in fancy clothes and wear expensive jewelry, and he would entertain his crew with grand parties and feasts.\\n\\nBut despite his fearsome reputation, Captain Blackbeard was a fair and just leader. He treated his crew with respect and kindness, and he always ensured that they were paid well for their hard work. He was also known for his generosity, and he would often share his treasure with the people of the towns that he visited.\\n\\nOne day, Captain Blackbeard and his crew were sailing off the coast of North Carolina when they were spotted by a British warship. The British ship was armed with cannons and had a large crew of sailors. Captain Blackbeard knew that he was in trouble, but he was not afraid. He ordered his crew to fire at the British ship, and the two ships fought for hours.\\n\\nIn the end, Captain Blackbeard and his crew were defeated. They were captured and taken to London, where they were tried and convicted of piracy. Captain Blackbeard was sentenced to death, but he was pardoned by the British government and allowed to return to his home in the Caribbean.\\n\\nCaptain Blackbeard lived the rest of his life in peace and luxury, and he is remembered as one of the most famous pirates of all time. He was a skilled sailor, a ruthless pirate, and a generous leader, and he will always be remembered as a legend of the high seas.', 'completionReason': 'FINISH'}]}, delta=None, additional_kwargs={})"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "messages = [\n",
    "    ChatMessage(role=\"system\", content=\"You are a pirate with a colorful personality\"),\n",
    "    ChatMessage(role=\"user\", content=\"Tell me a story\"),\n",
    "]\n",
    "resp = llm.chat(messages)\n",
    "resp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e9d9402-1f49-4f9f-8380-645cf3443eb7",
   "metadata": {},
   "source": [
    "### Stream Chat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cb516df-718e-46ec-96d7-713a3ce4bfa8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Once upon a time, there was a pirate named Captain Jack Sparrow. He was a skilled sailor and a master of disguise, and he loved nothing more than stealing treasure from other ships.\n",
      "\n",
      "One day, Captain Jack Sparrow and his crew were sailing through the Caribbean Sea when they came across a mysterious ship. The ship was flying a black flag, and the crew was wearing all black. Captain Jack Sparrow was intrigued and decided to approach the ship.\n",
      "\n",
      "As he approached the ship, Captain Jack Sparrow saw that the crew was made up of pirates from all over the world. They were all dressed in their best pirate gear and were ready to fight. Captain Jack Sparrow knew that he had to be careful, but he was also excited to see what kind of treasure the pirates had on board.\n",
      "\n",
      "Captain Jack Sparrow and his crew boarded the mysterious ship and began searching for treasure. They found a chest full of gold coins, a map to a hidden treasure, and a magical sword. They were so excited that they couldn't believe their luck.\n",
      "\n",
      "But their luck didn't last long. The pirates on board the mysterious ship were not happy with Captain Jack Sparrow and his crew taking their treasure. They attacked the crew, and Captain Jack Sparrow was forced to fight for his life.\n",
      "\n",
      "Captain Jack Sparrow was a skilled fighter, and he was able to hold his own against the pirates. But they were too many, and Captain Jack Sparrow was eventually captured.\n",
      "\n",
      "The pirates took Captain Jack Sparrow to their captain, a ruthless pirate named Blackbeard. Blackbeard was known for his cruelty and his love of treasure. He was determined to make Captain Jack Sparrow suffer for taking his treasure.\n",
      "\n",
      "Blackbeard tied Captain Jack Sparrow to a mast and left him there to die. But Captain Jack Sparrow was not one to give up easily. He used his wit and his cunning to escape from his captors.\n",
      "\n",
      "Captain Jack Sparrow managed to escape from Blackbeard's ship and make his way back to his own ship. He was injured and tired, but he was still alive. He knew that he had to find a way to get revenge on Blackbeard and his crew.\n",
      "\n",
      "Captain Jack Sparrow and his crew set out on a mission to find Blackbeard and his treasure. They sailed through the Caribbean Sea, encountering all sorts of obstacles and challenges. They fought pirates, they navigated treacherous waters, and they even had to"
     ]
    }
   ],
   "source": [
    "resp = llm.stream_chat(messages)\n",
    "for r in resp:\n",
    "    print(r.delta, end=\"\")"
   ]
  }
 ],
 "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}