llama-index-llms-bedrock-converse 0.8.2__tar.gz → 0.9.0__tar.gz

{llama_index_llms_bedrock_converse-0.8.2 → llama_index_llms_bedrock_converse-0.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llama-index-llms-bedrock-converse
-Version: 0.8.2
+Version: 0.9.0
 Summary: llama-index llms bedrock converse integration
 Author-email: Your Name <you@example.com>
 License-Expression: MIT

{llama_index_llms_bedrock_converse-0.8.2 → llama_index_llms_bedrock_converse-0.9.0}/llama_index/llms/bedrock_converse/base.py

@@ -138,6 +138,14 @@ class BedrockConverse(FunctionCallingLLM):
         default=60.0,
         description="The timeout for the Bedrock API request in seconds. It will be used for both connect and read timeouts.",
     )
+    system_prompt_caching: bool = Field(
+        default=False,
+        description="Whether to cache the system prompt. If you are using a system prompt, you should set this to True.",
+    )
+    tool_caching: bool = Field(
+        default=False,
+        description="Whether to cache the tools. If you are using tools, you should set this to True.",
+    )
     guardrail_identifier: Optional[str] = Field(
         description="The unique identifier of the guardrail that you want to use. If you don't provide a value, no guardrail is applied to the invocation."
     )
@@ -182,6 +190,8 @@ class BedrockConverse(FunctionCallingLLM):
         additional_kwargs: Optional[Dict[str, Any]] = None,
         callback_manager: Optional[CallbackManager] = None,
         system_prompt: Optional[str] = None,
+        system_prompt_caching: Optional[bool] = False,
+        tool_caching: Optional[bool] = False,
         messages_to_prompt: Optional[Callable[[Sequence[ChatMessage]], str]] = None,
         completion_to_prompt: Optional[Callable[[str], str]] = None,
         pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT,
@@ -212,6 +222,8 @@ class BedrockConverse(FunctionCallingLLM):
             model=model,
             callback_manager=callback_manager,
             system_prompt=system_prompt,
+            system_prompt_caching=system_prompt_caching,
+            tool_caching=tool_caching,
             messages_to_prompt=messages_to_prompt,
             completion_to_prompt=completion_to_prompt,
             pydantic_program_mode=pydantic_program_mode,
@@ -362,6 +374,8 @@ class BedrockConverse(FunctionCallingLLM):
             client=self._client,
             messages=converse_messages,
             system_prompt=system_prompt,
+            system_prompt_caching=self.system_prompt_caching,
+            tool_caching=self.tool_caching,
             max_retries=self.max_retries,
             stream=False,
             guardrail_identifier=self.guardrail_identifier,
@@ -408,6 +422,8 @@ class BedrockConverse(FunctionCallingLLM):
             client=self._client,
             messages=converse_messages,
             system_prompt=system_prompt,
+            system_prompt_caching=self.system_prompt_caching,
+            tool_caching=self.tool_caching,
             max_retries=self.max_retries,
             stream=True,
             guardrail_identifier=self.guardrail_identifier,
@@ -544,6 +560,8 @@ class BedrockConverse(FunctionCallingLLM):
             config=self._config,
             messages=converse_messages,
             system_prompt=system_prompt,
+            system_prompt_caching=self.system_prompt_caching,
+            tool_caching=self.tool_caching,
             max_retries=self.max_retries,
             stream=False,
             guardrail_identifier=self.guardrail_identifier,
@@ -592,6 +610,8 @@ class BedrockConverse(FunctionCallingLLM):
             config=self._config,
             messages=converse_messages,
             system_prompt=system_prompt,
+            system_prompt_caching=self.system_prompt_caching,
+            tool_caching=self.tool_caching,
             max_retries=self.max_retries,
             stream=True,
             guardrail_identifier=self.guardrail_identifier,
@@ -723,6 +743,7 @@ class BedrockConverse(FunctionCallingLLM):
         verbose: bool = False,
         allow_parallel_tool_calls: bool = False,
         tool_required: bool = False,
+        tool_caching: bool = False,
         tool_choice: Optional[dict] = None,
         **kwargs: Any,
     ) -> Dict[str, Any]:
@@ -737,7 +758,10 @@ class BedrockConverse(FunctionCallingLLM):
 
         # convert Llama Index tools to AWS Bedrock Converse tools
         tool_config = tools_to_converse_tools(
-            tools, tool_choice=tool_choice, tool_required=tool_required
+            tools,
+            tool_choice=tool_choice,
+            tool_required=tool_required,
+            tool_caching=tool_caching,
         )
 
         return {

{llama_index_llms_bedrock_converse-0.8.2 → llama_index_llms_bedrock_converse-0.9.0}/llama_index/llms/bedrock_converse/utils.py

@@ -29,6 +29,7 @@ HUMAN_PREFIX = "\n\nHuman:"
 ASSISTANT_PREFIX = "\n\nAssistant:"
 
 BEDROCK_MODELS = {
+    "amazon.nova-premier-v1:0": 1000000,
     "amazon.nova-pro-v1:0": 300000,
     "amazon.nova-lite-v1:0": 300000,
     "amazon.nova-micro-v1:0": 128000,
@@ -80,6 +81,7 @@ BEDROCK_MODELS = {
 }
 
 BEDROCK_FUNCTION_CALLING_MODELS = (
+    "amazon.nova-premier-v1:0",
     "amazon.nova-pro-v1:0",
     "amazon.nova-lite-v1:0",
     "amazon.nova-micro-v1:0",
@@ -108,6 +110,7 @@ BEDROCK_FUNCTION_CALLING_MODELS = (
 )
 
 BEDROCK_INFERENCE_PROFILE_SUPPORTED_MODELS = (
+    "amazon.nova-premier-v1:0",
     "amazon.nova-pro-v1:0",
     "amazon.nova-lite-v1:0",
     "amazon.nova-micro-v1:0",
@@ -347,6 +350,7 @@ def tools_to_converse_tools(
     tools: List["BaseTool"],
     tool_choice: Optional[dict] = None,
     tool_required: bool = False,
+    tool_caching: bool = False,
 ) -> Dict[str, Any]:
     """
     Converts a list of tools to AWS Bedrock Converse tools.
@@ -371,6 +375,8 @@ def tools_to_converse_tools(
             "inputSchema": {"json": tool.metadata.get_parameters_dict()},
         }
         converse_tools.append({"toolSpec": tool_dict})
+        if tool_caching:
+            converse_tools.append({"cachePoint": {"type": "default"}})
     return {
         "tools": converse_tools,
         # https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ToolChoice.html
@@ -437,6 +443,8 @@ def converse_with_retry(
     messages: Sequence[Dict[str, Any]],
     max_retries: int = 3,
     system_prompt: Optional[str] = None,
+    system_prompt_caching: bool = False,
+    tool_caching: bool = False,
     max_tokens: int = 1000,
     temperature: float = 0.1,
     stream: bool = False,
@@ -456,9 +464,16 @@ def converse_with_retry(
         },
     }
     if system_prompt:
-        converse_kwargs["system"] = [{"text": system_prompt}]
+        system_messages: list[dict[str, Any]] = [{"text": system_prompt}]
+        if system_prompt_caching:
+            system_messages.append({"cachePoint": {"type": "default"}})
+        converse_kwargs["system"] = system_messages
     if tool_config := kwargs.get("tools"):
         converse_kwargs["toolConfig"] = tool_config
+        if tool_caching and "tools" in converse_kwargs["toolConfig"]:
+            converse_kwargs["toolConfig"]["tools"].append(
+                {"cachePoint": {"type": "default"}}
+            )
     if guardrail_identifier and guardrail_version:
         converse_kwargs["guardrailConfig"] = {}
         converse_kwargs["guardrailConfig"]["guardrailIdentifier"] = guardrail_identifier
@@ -490,6 +505,8 @@ async def converse_with_retry_async(
     messages: Sequence[Dict[str, Any]],
     max_retries: int = 3,
     system_prompt: Optional[str] = None,
+    system_prompt_caching: bool = False,
+    tool_caching: bool = False,
     max_tokens: int = 1000,
     temperature: float = 0.1,
     stream: bool = False,
@@ -510,9 +527,16 @@ async def converse_with_retry_async(
         },
     }
     if system_prompt:
-        converse_kwargs["system"] = [{"text": system_prompt}]
+        system_messages: list[dict[str, Any]] = [{"text": system_prompt}]
+        if system_prompt_caching:
+            system_messages.append({"cachePoint": {"type": "default"}})
+        converse_kwargs["system"] = system_messages
     if tool_config := kwargs.get("tools"):
         converse_kwargs["toolConfig"] = tool_config
+        if tool_caching and "tools" in converse_kwargs["toolConfig"]:
+            converse_kwargs["toolConfig"]["tools"].append(
+                {"cachePoint": {"type": "default"}}
+            )
     if guardrail_identifier and guardrail_version:
         converse_kwargs["guardrailConfig"] = {}
         converse_kwargs["guardrailConfig"]["guardrailIdentifier"] = guardrail_identifier

{llama_index_llms_bedrock_converse-0.8.2 → llama_index_llms_bedrock_converse-0.9.0}/pyproject.toml

@@ -29,7 +29,7 @@ dev = [
 
 [project]
 name = "llama-index-llms-bedrock-converse"
-version = "0.8.2"
+version = "0.9.0"
 description = "llama-index llms bedrock converse integration"
 authors = [{name = "Your Name", email = "you@example.com"}]
 requires-python = ">=3.9,<4.0"