llama-index-llms-bedrock-converse 0.8.3__tar.gz → 0.9.1__tar.gz

{llama_index_llms_bedrock_converse-0.8.3 → llama_index_llms_bedrock_converse-0.9.1}/PKG-INFO

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

{llama_index_llms_bedrock_converse-0.8.3 → llama_index_llms_bedrock_converse-0.9.1}/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.3 → llama_index_llms_bedrock_converse-0.9.1}/llama_index/llms/bedrock_converse/utils.py

@@ -350,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.
@@ -374,6 +375,10 @@ 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
@@ -440,6 +445,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,
@@ -459,15 +466,21 @@ 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 guardrail_identifier and guardrail_version:
         converse_kwargs["guardrailConfig"] = {}
         converse_kwargs["guardrailConfig"]["guardrailIdentifier"] = guardrail_identifier
         converse_kwargs["guardrailConfig"]["guardrailVersion"] = guardrail_version
         if trace:
             converse_kwargs["guardrailConfig"]["trace"] = trace
+
     converse_kwargs = join_two_dicts(
         converse_kwargs,
         {
@@ -493,6 +506,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,
@@ -513,9 +528,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.3 → llama_index_llms_bedrock_converse-0.9.1}/pyproject.toml

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