diff --git a/docs/examples/embeddings/bedrock.ipynb b/docs/examples/embeddings/bedrock.ipynb
index 7c80f79709078681b560da8fe3131f722e308db6..81880921d4d39d9bb16dadb73993513932299711 100644
--- a/docs/examples/embeddings/bedrock.ipynb
+++ b/docs/examples/embeddings/bedrock.ipynb
@@ -32,12 +32,12 @@
"metadata": {},
"outputs": [],
"source": [
- "embed_model = BedrockEmbedding.from_credentials(\n",
+ "embed_model = BedrockEmbedding(\n",
" aws_access_key_id=os.getenv(\"AWS_ACCESS_KEY_ID\"),\n",
" aws_secret_access_key=os.getenv(\"AWS_SECRET_ACCESS_KEY\"),\n",
" aws_session_token=os.getenv(\"AWS_SESSION_TOKEN\"),\n",
- " aws_region=\"<aws-region>\",\n",
- " aws_profile=\"<aws-profile>\",\n",
+ " region_name=\"<aws-region>\",\n",
+ " profile_name=\"<aws-profile>\",\n",
")"
]
},
@@ -47,7 +47,8 @@
"metadata": {},
"outputs": [],
"source": [
- "embedding = embed_model.get_text_embedding(\"hello world\")"
+ "embedding = embed_model.get_text_embedding(\"hello world\")\n",
+ "print(embedding)"
]
},
{
@@ -63,7 +64,28 @@
"cell_type": "code",
"execution_count": null,
"metadata": {},
- "outputs": [],
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "{\n",
+ " \"amazon\": [\n",
+ " \"amazon.titan-embed-text-v1\",\n",
+ " \"amazon.titan-embed-g1-text-02\",\n",
+ " \"cohere.embed-english-v3\",\n",
+ " \"cohere.embed-multilingual-v3\"\n",
+ " ],\n",
+ " \"cohere\": [\n",
+ " \"amazon.titan-embed-text-v1\",\n",
+ " \"amazon.titan-embed-g1-text-02\",\n",
+ " \"cohere.embed-english-v3\",\n",
+ " \"cohere.embed-multilingual-v3\"\n",
+ " ]\n",
+ "}\n"
+ ]
+ }
+ ],
"source": [
"from llama_index.embeddings import BedrockEmbedding\n",
"import json\n",
@@ -88,9 +110,7 @@
"source": [
"from llama_index.embeddings import BedrockEmbedding\n",
"\n",
- "model = BedrockEmbedding().from_credentials(\n",
- " model_name=\"amazon.titan-embed-g1-text-02\"\n",
- ")\n",
+ "model = BedrockEmbedding(model_name=\"amazon.titan-embed-g1-text-02\")\n",
"embeddings = model.get_text_embedding(\"hello world\")\n",
"print(embeddings)"
]
@@ -104,21 +124,41 @@
"### cohere.embed-english-v3"
]
},
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Embed text for search"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "model = BedrockEmbedding(model_name=\"cohere.embed-english-v3\")\n",
+ "coherePayload = [\"This is a test document\", \"This is another test document\"]\n",
+ "embeddings = model.get_text_embedding_batch(coherePayload)\n",
+ "print(embeddings)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Embed query for question answering"
+ ]
+ },
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
- "model = BedrockEmbedding().from_credentials(\n",
- " model_name=\"cohere.embed-english-v3\"\n",
- ")\n",
- "coherePayload = {\n",
- " \"texts\": [\"This is a test document\", \"This is another test document\"],\n",
- " \"input_type\": \"search_document\",\n",
- " \"truncate\": \"NONE\",\n",
- "}\n",
- "embeddings = model.get_text_embedding(coherePayload)\n",
+ "model = BedrockEmbedding(model_name=\"cohere.embed-english-v3\")\n",
+ "coherePayload = \"What is gravity?\"\n",
+ "embeddings = model._get_query_embedding(coherePayload)\n",
"print(embeddings)"
]
},
@@ -135,18 +175,12 @@
"metadata": {},
"outputs": [],
"source": [
- "model = BedrockEmbedding().from_credentials(\n",
- " model_name=\"cohere.embed-multilingual-v3\"\n",
- ")\n",
- "coherePayload = {\n",
- " \"texts\": [\n",
- " \"This is a test document\",\n",
- " \"తెలà±à°—ౠఅనేది à°¦à±à°°à°¾à°µà°¿à°¡ à°à°¾à°·à°² à°•à±à°Ÿà±à°‚బానికి చెందిన à°à°¾à°·.\",\n",
- " ],\n",
- " \"input_type\": \"search_document\",\n",
- " \"truncate\": \"NONE\",\n",
- "}\n",
- "embeddings = model.get_text_embedding(coherePayload)\n",
+ "model = BedrockEmbedding(model_name=\"cohere.embed-multilingual-v3\")\n",
+ "coherePayload = [\n",
+ " \"This is a test document\",\n",
+ " \"తెలà±à°—ౠఅనేది à°¦à±à°°à°¾à°µà°¿à°¡ à°à°¾à°·à°² à°•à±à°Ÿà±à°‚బానికి చెందిన à°à°¾à°·.\",\n",
+ "]\n",
+ "embeddings = model.get_text_embedding_batch(coherePayload)\n",
"print(embeddings)"
]
}
diff --git a/docs/examples/llm/bedrock.ipynb b/docs/examples/llm/bedrock.ipynb
index 90f96aff3f362ec89bcd374fdecd6d84124e4448..0d3c4e57118e121bb2fdc546cd821d02bc4200e6 100644
--- a/docs/examples/llm/bedrock.ipynb
+++ b/docs/examples/llm/bedrock.ipynb
@@ -49,7 +49,7 @@
"metadata": {},
"outputs": [],
"source": [
- "!pip install llama-index"
+ "%pip install llama-index"
]
},
{
@@ -61,9 +61,8 @@
"source": [
"from llama_index.llms import Bedrock\n",
"\n",
- "profile_name = \"Your aws profile name\"\n",
"resp = Bedrock(\n",
- " model=\"amazon.titan-text-express-v1\", profile_name=profile_name\n",
+ " model=\"amazon.titan-text-express-v1\",\n",
").complete(\"Paul Graham is \")"
]
},
@@ -77,8 +76,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "\n",
- "Paul Graham is a computer scientist and entrepreneur, best known for co-founding the Silicon Valley startup incubator Y Combinator. He is also a prominent writer and speaker on technology and business topics, and his essays have been collected in a book titled \"Hackers & Painters.\"\n"
+ "son of John and Mary Graham.\n"
]
}
],
@@ -111,7 +109,7 @@
"]\n",
"\n",
"resp = Bedrock(\n",
- " model=\"amazon.titan-text-express-v1\", profile_name=profile_name\n",
+ " model=\"amazon.titan-text-express-v1\",\n",
").chat(messages)"
]
},
@@ -125,23 +123,27 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "assistant: Alright, matey! Here's a story for you:\n",
+ "assistant:  Ahoy, matey! Here's a tale for you:\n",
"\n",
- "Once upon a time, there was a pirate named Captain Jack Sparrow who sailed the seas in search of his next adventure. He was a notorious rogue with a reputation for being unpredictable and a bit of a scallywag.\n",
+ "Once upon a time, in a faraway land, there was a pirate named Captain Jack Sparrow. He was the most notorious pirate on the seven seas, with a reputation for his wit, charm, and daring exploits.\n",
"\n",
- "One day, Captain Jack encountered a group of treasure-hunting rivals who were also after the same treasure. The rivals tried to steal the treasure from Captain Jack, but he outsmarted them and managed to keep the treasure for himself.\n",
+ "Captain Jack Sparrow was born on the island of Tortuga, where he grew up learning the ways of the sea from his father, Captain Teague Sparrow. As a young boy, Jack was fascinated by the stories of pirates and their adventures, and he dreamed of one day setting sail on his own ship.\n",
"\n",
- "However, Captain Jack soon discovered that the treasure he had stolen was cursed. Every time he tried to use it, it would cause him some sort of trouble or inconvenience. For example, whenever he tried to spend it, it would turn into a pile of sand or a bunch of sea turtles.\n",
+ "One day, when Jack was just a teenager, his father was captured by the Spanish Navy. Jack was determined to rescue his father, and he set out on a quest to find the Spanish galleon that had taken him.\n",
"\n",
- "Despite the curse, Captain Jack was determined to find a way to break it. He set out on a journey to find a wise old seer who could help him lift the curse. Along the way, he encountered all sorts of strange and magical creatures, including a talking parrot and a sea witch.\n",
+ "After a long and treacherous journey, Jack finally found the Spanish galleon. He boarded the ship and fought his way to his father's cell, where he rescued him and escaped the ship with his father.\n",
"\n",
- "Finally, Captain Jack found the seer and explained his predicament. The seer told him that the only way to break the curse was to return the treasure to its rightful owner.\n",
+ "From that day on, Jack Sparrow became a pirate, sailing the high seas and seeking his fortune. He encountered many strange and exotic creatures, including mermaids, sea monsters, and even the legendary Kraken.\n",
"\n",
- "Captain Jack was hesitant at first, but he knew that it was the right thing to do. He set out on a new adventure to find the rightful owner of the treasure, and along the way, he discovered that sometimes the greatest treasures are not the ones that can be measured in gold or silver, but the ones that come with a sense of purpose and meaning.\n",
+ "One of Captain Jack Sparrow's most famous adventures was his encounter with the Black Pearl, a legendary pirate ship that was said to be cursed. Captain Jack Sparrow stole the Black Pearl from a group of Spanish pirates and set out on a quest to find the legendary Fountain of Youth.\n",
"\n",
- "And so, Captain Jack returned the treasure to its rightful owner, and the curse was lifted. He sailed off into the sunset, a hero who had learned that the true treasure of life is not found in material possessions, but in the experiences and connections we make with others.\n",
+ "Along the way, Jack encountered many challenges, including a group of pirates led by Captain Barbossa, who was determined to reclaim the Black Pearl. Captain Jack Sparrow and his crew fought bravely against Barbossa's pirates, and in the end, they were able to reclaim the Black Pearl and escape with their lives.\n",
"\n",
- "Yarr! Hope you enjoyed that tale, matey!\n"
+ "Captain Jack Sparrow was a skilled navigator and a master of disguise. He was able to fool many of his enemies, including the Spanish Navy, by using his wit and charm to outsmart them.\n",
+ "\n",
+ "In the end, Captain Jack Sparrow was a legendary pirate, and his adventures will be remembered for generations to come. He was a true hero, and his legacy will continue to inspire people to seek their fortune and adventure on the high seas.\n",
+ "\n",
+ "Would you like to hear more tales of adventure and bravery?\n"
]
}
],
@@ -174,7 +176,7 @@
"source": [
"from llama_index.llms import Bedrock\n",
"\n",
- "llm = Bedrock(model=\"amazon.titan-text-express-v1\", profile_name=profile_name)\n",
+ "llm = Bedrock(model=\"amazon.titan-text-express-v1\")\n",
"resp = llm.stream_complete(\"Paul Graham is \")"
]
},
@@ -188,14 +190,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "\n",
- "Paul Graham is a computer programmer, entrepreneur, investor, and writer, best known for co-founding the internet firm Y Combinator. He is also the author of several books, including \"The Innovator's Dilemma\" and \"On the Internet.\"\n",
- "\n",
- "Graham has been a strong supporter of the startup community and the concept of \"disruption\" in the technology sector. He has written extensively about the challenges faced by early-stage companies and the importance of creating new and innovative products.\n",
- "\n",
- "Graham is also known for his contrarian views on a variety of topics, including education, government, and the future of the internet. He has been an outspoken critic of the way higher education is administered in the United States and has advocated for a more experimental and entrepreneurial approach to learning.\n",
- "\n",
- "Overall, Paul Graham is a highly influential figure in the technology industry, known for his thoughtful and thought-provoking writing and his support for innovative startups and entrepreneurs."
+ "son of John and Mary Graham."
]
}
],
@@ -221,7 +216,9 @@
"source": [
"from llama_index.llms import Bedrock\n",
"\n",
- "llm = Bedrock(model=\"amazon.titan-text-express-v1\", profile_name=profile_name)\n",
+ "llm = Bedrock(\n",
+ " model=\"amazon.titan-text-express-v1\",\n",
+ ")\n",
"messages = [\n",
" ChatMessage(\n",
" role=\"system\", content=\"You are a pirate with a colorful personality\"\n",
@@ -241,15 +238,33 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "Once upon a time, there was a pirate with a colorful personality who sailed the high seas in search of adventure. She was known for her boldness, her wit, and her love of all things flashy and fancy. But beneath her swashbuckling exterior, there was a heart full of gold, and a desire to do good in the world.\n",
+ " Ahoy, matey! Here's a tale for you:\n",
"\n",
- "One day, while on her usual voyages, the pirate came across a small island in distress. The villagers were suffering from a terrible drought, and their crops were failing. The pirate knew that she had to help them, and so she set out to find a way to bring water to the island.\n",
+ "Once upon a time, there was a pirate named Captain Jack Sparrow. He was a notorious pirate on the high seas, known for his wit, his charm, and his love of adventure.\n",
"\n",
- "After much searching, the pirate discovered a hidden spring deep in the heart of the island. She worked tirelessly to build a system of pipes and aqueducts that would carry the spring water to the villages, and finally, after many long months of hard work, the drought was over, and the people were saved.\n",
+ "Captain Jack Sparrow was born on the Caribbean island of Tortuga. His father, Captain Teague Sparrow, was a notorious pirate himself, and Jack grew up learning the ways of the sea from his father.\n",
"\n",
- "The pirate was hailed as a hero, and the villagers threw a grand celebration in her honor. But she knew that her work was not yet done. She continued to sail the seas, seeking out other ways to help those in need, and to spread joy and happiness wherever she went.\n",
+ "At a young age, Jack showed a natural talent for sailing and navigation. He could navigate his ship through the toughest storms and find his way through the most treacherous waters.\n",
"\n",
- "And so, the pirate with the colorful personality lived out her days in a blaze of glory, inspiring others with her courage, her kindness, and her unquenchable sense of adventure."
+ "One day, Jack's father was captured by the Spanish Navy. The Spanish were determined to capture all the pirates on the Caribbean, and they were particularly interested in Captain Teague Sparrow.\n",
+ "\n",
+ "Jack was determined to rescue his father. He gathered a crew of his most trusted friends, and they set out on a daring mission to rescue Captain Teague Sparrow from the Spanish prison ship.\n",
+ "\n",
+ "The mission was perilous, and the crew faced many challenges. They had to navigate through treacherous waters, avoid the Spanish Navy, and deal with a host of other pirates and sea creatures.\n",
+ "\n",
+ "But Jack was a skilled captain, and he led his crew through the challenges with courage and determination. They finally reached the Spanish prison ship, and Jack and his crew attacked the guards.\n",
+ "\n",
+ "They managed to rescue Captain Teague Sparrow, but they were soon surrounded by the Spanish Navy. Jack and his crew fought bravely, but they were outnumbered and outgunned.\n",
+ "\n",
+ "Just when it seemed like all was lost, Jack had an idea. He used his wit and charm to distract the Spanish Navy, while his crew sneaked aboard their ship.\n",
+ "\n",
+ "They managed to steal the Spanish ship's treasure, and they set sail for Tortuga, their home island.\n",
+ "\n",
+ "Captain Jack Sparrow and his crew were hailed as heroes on Tortuga. They were celebrated for their bravery and their cunning, and they were known as the most feared pirates on the high seas.\n",
+ "\n",
+ "But Jack was not content with just being a hero. He was determined to find the legendary treasure of the Spanish Main. He knew that it was hidden somewhere on the ocean, and he was determined to find it.\n",
+ "\n",
+ "Jack and his crew set out on a new adventure, and they encountered many challenges along the way. They faced storms,"
]
}
],
@@ -275,7 +290,9 @@
"source": [
"from llama_index.llms import Bedrock\n",
"\n",
- "llm = Bedrock(model=\"amazon.titan-text-express-v1\", profile_name=profile_name)"
+ "llm = Bedrock(\n",
+ " model=\"amazon.titan-text-express-v1\", temperature=0.5, max_tokens=2048\n",
+ ")"
]
},
{
@@ -298,8 +315,54 @@
"name": "stdout",
"output_type": "stream",
"text": [
+ "an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
+ "Paul McCartney is known for his unique songwriting style, which combines elements of pop, rock, classical, and folk music. He is also known for his versatile guitar playing, which has inspired many musicians throughout the world.\n",
+ "\n",
+ "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
"\n",
- "Paul Graham is a computer scientist, entrepreneur, investor, and writer. He co-founded Viaweb, the first commercial web browser, and was a founder of Y Combinator, a startup accelerator. He is the author of several books, including \"The Art of Computer Programming\" and \"On Lisp.\" He is known for his essays on technology and business, and his perspective on the tech industry.\n"
+ "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
+ "\n",
+ "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
+ "\n",
+ "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
+ "\n",
+ "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
+ "\n",
+ "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
+ "\n",
+ "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n",
+ "\n",
+ "Paul McCartney is an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
+ "\n",
+ "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
+ "\n",
+ "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
+ "\n",
+ "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
+ "\n",
+ "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
+ "\n",
+ "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
+ "\n",
+ "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
+ "\n",
+ "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n",
+ "\n",
+ "Paul McCartney is an English singer, songwriter, and guitarist who is widely regarded as one of the most influential and important musicians of the 20th century. He is best known for his work with the Beatles, which he joined in 1962, and for his solo career, which began in 1965.\n",
+ "\n",
+ "One of Paul McCartney's most famous songs is \"Yesterday,\" which he wrote and performed with the Beatles in 1965. The song has become one of the most covered songs in history and has been praised for its beautiful melody and clever lyrics.\n",
+ "\n",
+ "Paul McCartney is also known for his activism and humanitarian work. He has been a champion for environmental causes, peace, and social justice throughout his career, and has used his platform to raise awareness about these issues. He has been awarded numerous honors, including the Nobel Peace Prize in 2019, for his contributions to music and society.\n",
+ "\n",
+ "Paul McCartney was born on September 18, 1942, in Liverpool, England. He grew up in a working-class family and showed an early interest in music. He began playing the guitar at the age of 12 and formed his first band, the Quarrymen, in 1957. The Quarrymen later became the Beatles, and Paul McCartney became the band's lead singer and guitarist.\n",
+ "\n",
+ "The Beatles' first hit song, \"Love Me Do,\" was released in 1962, and the band quickly became one of the most popular bands in the world. The Beatles' music was characterized by its catchy melodies, innovative production techniques, and clever lyrics, and the band's popularity continued to grow throughout the 1960s.\n",
+ "\n",
+ "Paul McCartney's solo career began in 1965 with the release of his debut album, \"McCartney.\" The album was a huge success, and Paul McCartney became known as a solo artist in his own right. He continued to release successful albums throughout the 1970s, including \"Band on the Run,\" \"Let It Be,\" and \"Abbey Road.\"\n",
+ "\n",
+ "In addition to his music career, Paul McCartney has been involved in many other ventures. He has written books, directed films, and even designed a line of clothing. He is also a member of the band Wings, which he formed in 1971 with his wife Linda McCartney.\n",
+ "\n",
+ "Paul McCartney is widely regarded as one of the greatest musicians of all time. His unique songwriting style, versatile guitar playing, and activism have made him a beloved figure in the music industry and a respected figure in society. He continues to perform and record new music, and his legacy as a musician and humanitarian will continue to inspire generations to come.\n"
]
}
],
@@ -329,7 +392,8 @@
" aws_access_key_id=\"AWS Access Key ID to use\",\n",
" aws_secret_access_key=\"AWS Secret Access Key to use\",\n",
" aws_session_token=\"AWS Session Token to use\",\n",
- " aws_region_name=\"AWS Region to use, eg. us-east-1\",\n",
+ " region_name=\"AWS Region to use, eg. us-east-1\",\n",
+ " profile_name=\"AWS Profile to use (if not using aws_access_key_id and aws_secret_access_key)\",\n",
")\n",
"\n",
"resp = llm.complete(\"Paul Graham is \")"
@@ -345,8 +409,7 @@
"name": "stdout",
"output_type": "stream",
"text": [
- "\n",
- "Paul Graham is an American computer scientist, entrepreneur, investor, and author, best known for co-founding Viaweb, the first commercial web browser. He was a co-founder of Netscape Communications and the creator of the Mozilla Foundation. He was also a Y Combinator partner and a prominent early-stage investor in companies such as Airbnb, Dropbox, Facebook, and Twitter.\n"
+ "son of John and Mary Graham.\n"
]
}
],
diff --git a/llama_index/embeddings/bedrock.py b/llama_index/embeddings/bedrock.py
index d9f3c7c0026d976ef131b03ee4ef4e6431033984..707764e632d911b7420c0d026b26ac668bcc4853 100644
--- a/llama_index/embeddings/bedrock.py
+++ b/llama_index/embeddings/bedrock.py
@@ -1,13 +1,12 @@
import json
-import os
-import warnings
from enum import Enum
-from typing import Any, Dict, List, Literal, Optional
+from typing import Any, Callable, Dict, List, Literal, Optional, Sequence
-from llama_index.bridge.pydantic import PrivateAttr
+from llama_index.bridge.pydantic import Field, PrivateAttr
from llama_index.callbacks.base import CallbackManager
-from llama_index.constants import DEFAULT_EMBED_BATCH_SIZE
from llama_index.core.embeddings.base import BaseEmbedding, Embedding
+from llama_index.core.llms.types import ChatMessage
+from llama_index.types import BaseOutputParser, PydanticProgramMode
class PROVIDERS(str, Enum):
@@ -33,169 +32,140 @@ PROVIDER_SPECIFIC_IDENTIFIERS = {
class BedrockEmbedding(BaseEmbedding):
+ model: str = Field(description="The modelId of the Bedrock model to use.")
+ profile_name: Optional[str] = Field(
+ description="The name of aws profile to use. If not given, then the default profile is used.",
+ exclude=True,
+ )
+ aws_access_key_id: Optional[str] = Field(
+ description="AWS Access Key ID to use", exclude=True
+ )
+ aws_secret_access_key: Optional[str] = Field(
+ description="AWS Secret Access Key to use", exclude=True
+ )
+ aws_session_token: Optional[str] = Field(
+ description="AWS Session Token to use", exclude=True
+ )
+ region_name: Optional[str] = Field(
+ description="AWS region name to use. Uses region configured in AWS CLI if not passed",
+ exclude=True,
+ )
+ botocore_session: Optional[Any] = Field(
+ description="Use this Botocore session instead of creating a new default one.",
+ exclude=True,
+ )
+ botocore_config: Optional[Any] = Field(
+ description="Custom configuration object to use instead of the default generated one.",
+ exclude=True,
+ )
+ max_retries: int = Field(
+ default=10, description="The maximum number of API retries.", gt=0
+ )
+ timeout: float = Field(
+ default=60.0,
+ description="The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts.",
+ )
+ additional_kwargs: Dict[str, Any] = Field(
+ default_factory=dict, description="Additional kwargs for the bedrock client."
+ )
_client: Any = PrivateAttr()
- _verbose: bool = PrivateAttr()
def __init__(
self,
- model_name: str = Models.TITAN_EMBEDDING,
- client: Any = None,
- embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
- callback_manager: Optional[CallbackManager] = None,
- verbose: bool = False,
- ):
- self._client = client
- self._verbose = verbose
-
- super().__init__(
- model_name=model_name,
- client=client,
- embed_batch_size=embed_batch_size,
- callback_manager=callback_manager,
- )
-
- @staticmethod
- def list_supported_models() -> Dict[str, List[str]]:
- list_models = {}
- for provider in PROVIDERS:
- list_models[provider.value] = [m.value for m in Models]
- return list_models
-
- @classmethod
- def class_name(self) -> str:
- return "BedrockEmbedding"
-
- def set_credentials(
- self,
- aws_region: Optional[str] = None,
- aws_access_key_id: Optional[str] = None,
- aws_secret_access_key: Optional[str] = None,
- aws_session_token: Optional[str] = None,
- aws_profile: Optional[str] = None,
- ) -> None:
- aws_region = aws_region or os.getenv("AWS_REGION")
- aws_access_key_id = aws_access_key_id or os.getenv("AWS_ACCESS_KEY_ID")
- aws_secret_access_key = aws_secret_access_key or os.getenv(
- "AWS_SECRET_ACCESS_KEY"
- )
- aws_session_token = aws_session_token or os.getenv("AWS_SESSION_TOKEN")
-
- if aws_region is None:
- warnings.warn(
- "AWS_REGION not found. Set environment variable AWS_REGION or set aws_region"
- )
-
- if aws_access_key_id is None:
- warnings.warn(
- "AWS_ACCESS_KEY_ID not found. Set environment variable AWS_ACCESS_KEY_ID or set aws_access_key_id"
- )
- assert aws_access_key_id is not None
-
- if aws_secret_access_key is None:
- warnings.warn(
- "AWS_SECRET_ACCESS_KEY not found. Set environment variable AWS_SECRET_ACCESS_KEY or set aws_secret_access_key"
- )
- assert aws_secret_access_key is not None
-
- if aws_session_token is None:
- warnings.warn(
- "AWS_SESSION_TOKEN not found. Set environment variable AWS_SESSION_TOKEN or set aws_session_token"
- )
- assert aws_session_token is not None
-
- session_kwargs = {
- "profile_name": aws_profile,
- "region_name": aws_region,
- "aws_access_key_id": aws_access_key_id,
- "aws_secret_access_key": aws_secret_access_key,
- "aws_session_token": aws_session_token,
- }
-
- try:
- import boto3
-
- session = boto3.Session(**session_kwargs)
- except ImportError:
- raise ImportError(
- "boto3 package not found, install with" "'pip install boto3'"
- )
-
- if "bedrock-runtime" in session.get_available_services():
- self._client = session.client("bedrock-runtime")
- else:
- self._client = session.client("bedrock")
-
- @classmethod
- def from_credentials(
- cls,
- model_name: str = Models.TITAN_EMBEDDING,
- aws_region: Optional[str] = None,
+ model: str = Models.TITAN_EMBEDDING,
+ profile_name: Optional[str] = None,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
aws_session_token: Optional[str] = None,
- aws_profile: Optional[str] = None,
- embed_batch_size: int = DEFAULT_EMBED_BATCH_SIZE,
+ region_name: Optional[str] = None,
+ client: Optional[Any] = None,
+ botocore_session: Optional[Any] = None,
+ botocore_config: Optional[Any] = None,
+ additional_kwargs: Optional[Dict[str, Any]] = None,
+ max_retries: int = 10,
+ timeout: float = 60.0,
callback_manager: Optional[CallbackManager] = None,
- verbose: bool = False,
- ) -> "BedrockEmbedding":
- """
- Instantiate using AWS credentials.
-
- Args:
- model_name (str) : Name of the model
- aws_access_key_id (str): AWS access key ID
- aws_secret_access_key (str): AWS secret access key
- aws_session_token (str): AWS session token
- aws_region (str): AWS region where the service is located
- aws_profile (str): AWS profile, when None, default profile is chosen automatically
-
- Example:
- .. code-block:: python
-
- from llama_index.embeddings import BedrockEmbedding
-
- # Define the model name
- model_name = "your_model_name"
-
- embeddings = BedrockEmbedding.from_credentials(
- model_name,
- aws_access_key_id,
- aws_secret_access_key,
- aws_session_token,
- aws_region,
- aws_profile,
- )
+ # base class
+ system_prompt: Optional[str] = None,
+ messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
+ completion_to_prompt: Optional[Callable[[str], str]] = None,
+ pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
+ output_parser: Optional[BaseOutputParser] = None,
+ **kwargs: Any,
+ ):
+ additional_kwargs = additional_kwargs or {}
- """
session_kwargs = {
- "profile_name": aws_profile,
- "region_name": aws_region,
+ "profile_name": profile_name,
+ "region_name": region_name,
"aws_access_key_id": aws_access_key_id,
"aws_secret_access_key": aws_secret_access_key,
"aws_session_token": aws_session_token,
+ "botocore_session": botocore_session,
}
-
+ config = None
try:
import boto3
-
+ from botocore.config import Config
+
+ config = (
+ Config(
+ retries={"max_attempts": max_retries, "mode": "standard"},
+ connect_timeout=timeout,
+ read_timeout=timeout,
+ )
+ if botocore_config is None
+ else botocore_config
+ )
session = boto3.Session(**session_kwargs)
except ImportError:
raise ImportError(
"boto3 package not found, install with" "'pip install boto3'"
)
- if "bedrock-runtime" in session.get_available_services():
- client = session.client("bedrock-runtime")
+ # Prior to general availability, custom boto3 wheel files were
+ # distributed that used the bedrock service to invokeModel.
+ # This check prevents any services still using those wheel files
+ # from breaking
+ if client is not None:
+ self._client = client
+ elif "bedrock-runtime" in session.get_available_services():
+ self._client = session.client("bedrock-runtime", config=config)
else:
- client = session.client("bedrock")
- return cls(
- client=client,
- model_name=model_name,
- embed_batch_size=embed_batch_size,
+ self._client = session.client("bedrock", config=config)
+
+ super().__init__(
+ model=model,
+ max_retries=max_retries,
+ timeout=timeout,
+ botocore_config=config,
+ profile_name=profile_name,
+ aws_access_key_id=aws_access_key_id,
+ aws_secret_access_key=aws_secret_access_key,
+ aws_session_token=aws_session_token,
+ region_name=region_name,
+ botocore_session=botocore_session,
+ additional_kwargs=additional_kwargs,
callback_manager=callback_manager,
- verbose=verbose,
+ system_prompt=system_prompt,
+ messages_to_prompt=messages_to_prompt,
+ completion_to_prompt=completion_to_prompt,
+ pydantic_program_mode=pydantic_program_mode,
+ output_parser=output_parser,
+ **kwargs,
)
+ @staticmethod
+ def list_supported_models() -> Dict[str, List[str]]:
+ list_models = {}
+ for provider in PROVIDERS:
+ list_models[provider.value] = [m.value for m in Models]
+ return list_models
+
+ @classmethod
+ def class_name(self) -> str:
+ return "BedrockEmbedding"
+
def _get_embedding(self, payload: str, type: Literal["text", "query"]) -> Embedding:
if self._client is None:
self.set_credentials()
@@ -203,12 +173,12 @@ class BedrockEmbedding(BaseEmbedding):
if self._client is None:
raise ValueError("Client not set")
- provider = self.model_name.split(".")[0]
+ provider = self.model.split(".")[0]
request_body = self._get_request_body(provider, payload, type)
response = self._client.invoke_model(
body=request_body,
- modelId=self.model_name,
+ modelId=self.model,
accept="application/json",
contentType="application/json",
)
@@ -244,8 +214,6 @@ class BedrockEmbedding(BaseEmbedding):
}
"""
- if self._verbose:
- print("provider: ", provider, PROVIDERS.AMAZON)
if provider == PROVIDERS.AMAZON:
request_body = json.dumps({"inputText": payload})
elif provider == PROVIDERS.COHERE:
diff --git a/llama_index/llms/bedrock.py b/llama_index/llms/bedrock.py
index 32a0bc4bf18b076e1aba178fdb208aa1c353f7ea..4db7f7ad7b00c99dbfa072847cdb3823fd3b5072 100644
--- a/llama_index/llms/bedrock.py
+++ b/llama_index/llms/bedrock.py
@@ -3,6 +3,9 @@ from typing import Any, Callable, Dict, Optional, Sequence
from llama_index.bridge.pydantic import Field, PrivateAttr
from llama_index.callbacks import CallbackManager
+from llama_index.constants import (
+ DEFAULT_TEMPERATURE,
+)
from llama_index.core.llms.types import (
ChatMessage,
ChatResponse,
@@ -41,16 +44,33 @@ class Bedrock(LLM):
profile_name: Optional[str] = Field(
description="The name of aws profile to use. If not given, then the default profile is used."
)
- aws_access_key_id: Optional[str] = Field(description="AWS Access Key ID to use")
+ aws_access_key_id: Optional[str] = Field(
+ description="AWS Access Key ID to use", exclude=True
+ )
aws_secret_access_key: Optional[str] = Field(
- description="AWS Secret Access Key to use"
+ description="AWS Secret Access Key to use", exclude=True
+ )
+ aws_session_token: Optional[str] = Field(
+ description="AWS Session Token to use", exclude=True
+ )
+ region_name: Optional[str] = Field(
+ description="AWS region name to use. Uses region configured in AWS CLI if not passed",
+ exclude=True,
)
- aws_session_token: Optional[str] = Field(description="AWS Session Token to use")
- aws_region_name: Optional[str] = Field(
- description="AWS region name to use. Uses region configured in AWS CLI if not passed"
+ botocore_session: Optional[Any] = Field(
+ description="Use this Botocore session instead of creating a new default one.",
+ exclude=True,
+ )
+ botocore_config: Optional[Any] = Field(
+ description="Custom configuration object to use instead of the default generated one.",
+ exclude=True,
)
max_retries: int = Field(
- default=10, description="The maximum number of API retries."
+ default=10, description="The maximum number of API retries.", gt=0
+ )
+ timeout: float = Field(
+ default=60.0,
+ description="The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts.",
)
additional_kwargs: Dict[str, Any] = Field(
default_factory=dict,
@@ -64,16 +84,19 @@ class Bedrock(LLM):
def __init__(
self,
model: str,
- temperature: Optional[float] = 0.5,
+ temperature: Optional[float] = DEFAULT_TEMPERATURE,
max_tokens: Optional[int] = 512,
context_size: Optional[int] = None,
profile_name: Optional[str] = None,
aws_access_key_id: Optional[str] = None,
aws_secret_access_key: Optional[str] = None,
aws_session_token: Optional[str] = None,
- aws_region_name: Optional[str] = None,
- timeout: Optional[float] = None,
+ region_name: Optional[str] = None,
+ botocore_session: Optional[Any] = None,
+ client: Optional[Any] = None,
+ timeout: Optional[float] = 60.0,
max_retries: Optional[int] = 10,
+ botocore_config: Optional[Any] = None,
additional_kwargs: Optional[Dict[str, Any]] = None,
callback_manager: Optional[CallbackManager] = None,
system_prompt: Optional[str] = None,
@@ -81,6 +104,7 @@ class Bedrock(LLM):
completion_to_prompt: Optional[Callable[[str], str]] = None,
pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
output_parser: Optional[BaseOutputParser] = None,
+ **kwargs: Any,
) -> None:
if context_size is None and model not in BEDROCK_FOUNDATION_LLMS:
raise ValueError(
@@ -88,39 +112,45 @@ class Bedrock(LLM):
"model provided refers to a non-foundation model."
" Please specify the context_size"
)
+
+ session_kwargs = {
+ "profile_name": profile_name,
+ "region_name": region_name,
+ "aws_access_key_id": aws_access_key_id,
+ "aws_secret_access_key": aws_secret_access_key,
+ "aws_session_token": aws_session_token,
+ "botocore_session": botocore_session,
+ }
+ config = None
try:
import boto3
- import botocore
+ from botocore.config import Config
- except Exception as e:
- raise ImportError(
- "You must install the `boto3` package to use Bedrock."
- "Please `pip install boto3`"
- ) from e
- try:
- if not profile_name and aws_access_key_id:
- session = boto3.Session(
- aws_access_key_id=aws_access_key_id,
- aws_secret_access_key=aws_secret_access_key,
- aws_session_token=aws_session_token,
- region_name=aws_region_name,
+ config = (
+ Config(
+ retries={"max_attempts": max_retries, "mode": "standard"},
+ connect_timeout=timeout,
+ read_timeout=timeout,
)
- else:
- session = boto3.Session(profile_name=profile_name)
- # Prior to general availability, custom boto3 wheel files were
- # distributed that used the bedrock service to invokeModel.
- # This check prevents any services still using those wheel files
- # from breaking
- if "bedrock-runtime" in session.get_available_services():
- self._client = session.client("bedrock-runtime")
- else:
- self._client = session.client("bedrock")
-
- except botocore.exceptions.NoRegionError as e:
- raise ValueError(
- "If default region is not set in AWS CLI, you must provide"
- " the region_name argument to llama_index.llms.Bedrock"
+ if botocore_config is None
+ else botocore_config
)
+ session = boto3.Session(**session_kwargs)
+ except ImportError:
+ raise ImportError(
+ "boto3 package not found, install with" "'pip install boto3'"
+ )
+
+ # Prior to general availability, custom boto3 wheel files were
+ # distributed that used the bedrock service to invokeModel.
+ # This check prevents any services still using those wheel files
+ # from breaking
+ if client is not None:
+ self._client = client
+ elif "bedrock-runtime" in session.get_available_services():
+ self._client = session.client("bedrock-runtime", config=config)
+ else:
+ self._client = session.client("bedrock", config=config)
additional_kwargs = additional_kwargs or {}
callback_manager = callback_manager or CallbackManager([])
@@ -138,6 +168,7 @@ class Bedrock(LLM):
profile_name=profile_name,
timeout=timeout,
max_retries=max_retries,
+ botocore_config=config,
additional_kwargs=additional_kwargs,
callback_manager=callback_manager,
system_prompt=system_prompt,
diff --git a/tests/embeddings/test_bedrock.py b/tests/embeddings/test_bedrock.py
index aaa7ff741e51732fa4114ae6240ff3d38762a7db..d14df0be8cd87c3484c45ea91b2572d192bfa0fc 100644
--- a/tests/embeddings/test_bedrock.py
+++ b/tests/embeddings/test_bedrock.py
@@ -34,7 +34,7 @@ class TestBedrockEmbedding(TestCase):
)
bedrock_embedding = BedrockEmbedding(
- model_name=Models.TITAN_EMBEDDING,
+ model=Models.TITAN_EMBEDDING,
client=self.bedrock_client,
)
@@ -63,7 +63,7 @@ class TestBedrockEmbedding(TestCase):
)
bedrock_embedding = BedrockEmbedding(
- model_name=Models.COHERE_EMBED_ENGLISH_V3,
+ model=Models.COHERE_EMBED_ENGLISH_V3,
client=self.bedrock_client,
)
diff --git a/tests/llms/test_bedrock.py b/tests/llms/test_bedrock.py
index d26661d40f88b023e6b5f80d6ea936fc3f43d6ed..eece5960c1267870026c8b932430a287f0177eec 100644
--- a/tests/llms/test_bedrock.py
+++ b/tests/llms/test_bedrock.py
@@ -58,7 +58,7 @@ class MockStreamCompletionWithRetry:
) -> dict:
assert json.loads(request_body) == {
"inputText": self.expected_prompt,
- "textGenerationConfig": {"maxTokenCount": 512, "temperature": 0.5},
+ "textGenerationConfig": {"maxTokenCount": 512, "temperature": 0.1},
}
return {
"ResponseMetadata": {
@@ -84,27 +84,27 @@ class MockStreamCompletionWithRetry:
[
(
"amazon.titan-text-express-v1",
- '{"inputText": "test prompt", "textGenerationConfig": {"temperature": 0.5, "maxTokenCount": 512}}',
+ '{"inputText": "test prompt", "textGenerationConfig": {"temperature": 0.1, "maxTokenCount": 512}}',
'{"inputTextTokenCount": 3, "results": [{"tokenCount": 14, "outputText": "\\n\\nThis is indeed a test", "completionReason": "FINISH"}]}',
- '{"inputText": "user: test prompt\\nassistant: ", "textGenerationConfig": {"temperature": 0.5, "maxTokenCount": 512}}',
+ '{"inputText": "user: test prompt\\nassistant: ", "textGenerationConfig": {"temperature": 0.1, "maxTokenCount": 512}}',
),
(
"ai21.j2-grande-instruct",
- '{"prompt": "test prompt", "temperature": 0.5, "maxTokens": 512}',
+ '{"prompt": "test prompt", "temperature": 0.1, "maxTokens": 512}',
'{"completions": [{"data": {"text": "\\n\\nThis is indeed a test"}}]}',
- '{"prompt": "user: test prompt\\nassistant: ", "temperature": 0.5, "maxTokens": 512}',
+ '{"prompt": "user: test prompt\\nassistant: ", "temperature": 0.1, "maxTokens": 512}',
),
(
"cohere.command-text-v14",
- '{"prompt": "test prompt", "temperature": 0.5, "max_tokens": 512}',
+ '{"prompt": "test prompt", "temperature": 0.1, "max_tokens": 512}',
'{"generations": [{"text": "\\n\\nThis is indeed a test"}]}',
- '{"prompt": "user: test prompt\\nassistant: ", "temperature": 0.5, "max_tokens": 512}',
+ '{"prompt": "user: test prompt\\nassistant: ", "temperature": 0.1, "max_tokens": 512}',
),
(
"anthropic.claude-instant-v1",
- '{"prompt": "\\n\\nHuman: test prompt\\n\\nAssistant: ", "temperature": 0.5, "max_tokens_to_sample": 512}',
+ '{"prompt": "\\n\\nHuman: test prompt\\n\\nAssistant: ", "temperature": 0.1, "max_tokens_to_sample": 512}',
'{"completion": "\\n\\nThis is indeed a test"}',
- '{"prompt": "\\n\\nHuman: test prompt\\n\\nAssistant: ", "temperature": 0.5, "max_tokens_to_sample": 512}',
+ '{"prompt": "\\n\\nHuman: test prompt\\n\\nAssistant: ", "temperature": 0.1, "max_tokens_to_sample": 512}',
),
(
"meta.llama2-13b-chat-v1",
@@ -112,13 +112,13 @@ class MockStreamCompletionWithRetry:
"honest assistant. Always answer as helpfully as possible and follow "
"ALL given instructions. Do not speculate or make up information. Do "
"not reference any given instructions or context. \\n<</SYS>>\\n\\n "
- 'test prompt [/INST]", "temperature": 0.5, "max_gen_len": 512}',
+ 'test prompt [/INST]", "temperature": 0.1, "max_gen_len": 512}',
'{"generation": "\\n\\nThis is indeed a test"}',
'{"prompt": "<s> [INST] <<SYS>>\\n You are a helpful, respectful and '
"honest assistant. Always answer as helpfully as possible and follow "
"ALL given instructions. Do not speculate or make up information. Do "
"not reference any given instructions or context. \\n<</SYS>>\\n\\n "
- 'test prompt [/INST]", "temperature": 0.5, "max_gen_len": 512}',
+ 'test prompt [/INST]", "temperature": 0.1, "max_gen_len": 512}',
),
],
)
@@ -128,7 +128,7 @@ def test_model_basic(
llm = Bedrock(
model=model,
profile_name=None,
- aws_region_name="us-east-1",
+ region_name="us-east-1",
aws_access_key_id="test",
)
@@ -168,7 +168,7 @@ def test_model_streaming(monkeypatch: MonkeyPatch) -> None:
llm = Bedrock(
model="amazon.titan-text-express-v1",
profile_name=None,
- aws_region_name="us-east-1",
+ region_name="us-east-1",
aws_access_key_id="test",
)
test_prompt = "test prompt"