Add prompt caching (#891)

* Reflect the cost of cache read/write to the price calculation.

* Implement prompt caching.
- Backend
  - Add a field `use_prompt_caching` to models and schemas of custom bot.
    - `BotModel`
    - `BotInput`
    - `BotOutput`
    - `BotModifyInput`
    - `BotModifyOutput`
  - Add a column `UsePromptCaching` to DynamoDB table `BotTableV3`.
  - Use prompt caching if enabled and the model supports it.
- Frontend
  - Add a field `usePromptCaching` to schemas of custom bot.
    - `BotDetails`
    - `RegisterBotRequest`
    - `RegisterBotResponse`
    - `UpdateBotRequest`
    - `UpdateBotResponse`
  - Add 'Prompt Caching' section to `BotKbEditPage`.

* [Debug]Print token count and price when received `STREAMING_END`.

* Reformat modified python codes.
- `backend/app/config.py`
- `backend/app/bedrock.py`
- `backend/app/repositories/custom_bot.py`
- `backend/app/usecases/chat.py`

* Use `ExpandableDrawerGroup` for prompt caching settings.

* Revert "[Debug]Print token count and price when received `STREAMING_END`."

This reverts commit ba5e584.

* Refactor data structure of agent settings.

* Rename `usePromptCaching` to `promptCachingEnabled`.
- `use_prompt_caching` -> `prompt_caching_enabled`
- Change `BotModel.prompt_caching_enabled` to a non-nullable type.

* Change `BotModifyInput.prompt_caching_enabled` to a non-nullable type.

* Add `token_count` and `price` to the payload of `STREAMING_END` notification.

Author: Yukinobu-Mine

backend/app/bedrock.py

@@ -2,7 +2,7 @@
 
 import logging
 import os
-from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple, TypeGuard
+from typing import TYPE_CHECKING, Any, Dict, Optional, Literal, Tuple, TypeGuard
 
 from app.config import (
     BEDROCK_PRICING,
@@ -30,6 +30,7 @@
         InferenceConfigurationTypeDef,
         MessageTypeDef,
         SystemContentBlockTypeDef,
+        ToolTypeDef,
     )
 
 
@@ -81,6 +82,31 @@ def is_tooluse_supported(model: type_model_name) -> bool:
     ]
 
 
+def is_prompt_caching_supported(
+    model: type_model_name, target: Literal["system", "message", "tool"]
+) -> bool:
+    if target == "tool":
+        return model in [
+            "claude-v4-opus",
+            "claude-v4-sonnet",
+            "claude-v3.7-sonnet",
+            "claude-v3.5-sonnet-v2",
+            "claude-v3.5-haiku",
+        ]
+
+    else:
+        return model in [
+            "claude-v4-opus",
+            "claude-v4-sonnet",
+            "claude-v3.7-sonnet",
+            "claude-v3.5-sonnet-v2",
+            "claude-v3.5-haiku",
+            "amazon-nova-pro",
+            "amazon-nova-lite",
+            "amazon-nova-micro",
+        ]
+
+
 def _prepare_deepseek_model_params(
     model: type_model_name, generation_params: Optional[GenerationParamsModel] = None
 ) -> Tuple[InferenceConfigurationTypeDef, None]:
@@ -263,6 +289,7 @@ def compose_args_for_converse_api(
     tools: dict[str, AgentTool] | None = None,
     stream: bool = True,
     enable_reasoning: bool = False,
+    prompt_caching_enabled: bool = False,
 ) -> ConverseStreamRequestTypeDef:
     def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
         # Drop unsigned reasoning blocks only for DeepSeek R1
@@ -303,6 +330,16 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
         for message in messages
         if _is_conversation_role(message.role)
     ]
+    tool_specs: list[ToolTypeDef] | None = (
+        [
+            {
+                "toolSpec": tool.to_converse_spec(),
+            }
+            for tool in tools.values()
+        ]
+        if tools
+        else None
+    )
 
     # Prepare model-specific parameters
     inference_config: InferenceConfigurationTypeDef
@@ -457,6 +494,41 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
             if len(instruction) > 0
         ]
 
+    if prompt_caching_enabled and not (
+        tool_specs and not is_prompt_caching_supported(model, target="tool")
+    ):
+        if is_prompt_caching_supported(model, "system") and len(system_prompts) > 0:
+            system_prompts.append(
+                {
+                    "cachePoint": {
+                        "type": "default",
+                    },
+                }
+            )
+
+        if is_prompt_caching_supported(model, target="message"):
+            for order, message in enumerate(
+                filter(lambda m: m["role"] == "user", reversed(arg_messages))
+            ):
+                if order >= 2:
+                    break
+
+                message["content"] = [
+                    *(message["content"]),
+                    {
+                        "cachePoint": {"type": "default"},
+                    },
+                ]
+
+        if is_prompt_caching_supported(model, target="tool") and tool_specs:
+            tool_specs.append(
+                {
+                    "cachePoint": {
+                        "type": "default",
+                    },
+                }
+            )
+
     # Construct the base arguments
     args: ConverseStreamRequestTypeDef = {
         "inferenceConfig": inference_config,
@@ -480,14 +552,9 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
             args["guardrailConfig"]["streamProcessingMode"] = "async"
 
     # NOTE: Some models doesn't support tool use. https://docs.aws.amazon.com/bedrock/latest/userguide/conversation-inference-supported-models-features.html
-    if tools:
+    if tool_specs:
         args["toolConfig"] = {
-            "tools": [
-                {
-                    "toolSpec": tool.to_converse_spec(),
-                }
-                for tool in tools.values()
-            ],
+            "tools": tool_specs,
         }
 
     return args
@@ -519,6 +586,8 @@ def calculate_price(
     model: type_model_name,
     input_tokens: int,
     output_tokens: int,
+    cache_read_input_tokens: int,
+    cache_write_input_tokens: int,
     region: str = BEDROCK_REGION,
 ) -> float:
     input_price = (
@@ -531,8 +600,29 @@ def calculate_price(
         .get(model, {})
         .get("output", BEDROCK_PRICING["default"][model]["output"])
     )
+    cache_read_input_price = (
+        BEDROCK_PRICING.get(region, {})
+        .get(model, {})
+        .get(
+            "cache_read_input",
+            BEDROCK_PRICING["default"][model].get("cache_read_input", input_price),
+        )
+    )
+    cache_write_input_price = (
+        BEDROCK_PRICING.get(region, {})
+        .get(model, {})
+        .get(
+            "cache_write_input",
+            BEDROCK_PRICING["default"][model].get("cache_write_input", input_price),
+        )
+    )
 
-    return input_price * input_tokens / 1000.0 + output_price * output_tokens / 1000.0
+    return (
+        input_price * input_tokens / 1000.0
+        + output_price * output_tokens / 1000.0
+        + cache_read_input_price * cache_read_input_tokens / 1000.0
+        + cache_write_input_price * cache_write_input_tokens / 1000.0
+    )
 
 
 def get_model_id(

backend/app/config.py

@@ -63,19 +63,64 @@ class EmbeddingConfig(TypedDict):
 # See: https://aws.amazon.com/bedrock/pricing/
 BEDROCK_PRICING = {
     "us-east-1": {
-        "claude-v4-opus": {"input": 0.015, "output": 0.075},
-        "claude-v4-sonnet": {"input": 0.003, "output": 0.015},
+        "claude-v4-opus": {
+            "input": 0.015,
+            "output": 0.075,
+            "cache_write_input": 0.01875,
+            "cache_read_input": 0.0015,
+        },
+        "claude-v4-sonnet": {
+            "input": 0.003,
+            "output": 0.015,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
         "claude-v3-haiku": {"input": 0.00025, "output": 0.00125},
-        "claude-v3.5-haiku": {"input": 0.001, "output": 0.005},
-        "claude-v3.5-sonnet": {"input": 0.00300, "output": 0.01500},
-        "claude-v3.5-sonnet-v2": {"input": 0.00300, "output": 0.01500},
-        "claude-v3.7-sonnet": {"input": 0.00300, "output": 0.01500},
+        "claude-v3.5-haiku": {
+            "input": 0.001,
+            "output": 0.005,
+            "cache_write_input": 0.001,
+            "cache_read_input": 0.00008,
+        },
+        "claude-v3.5-sonnet": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
+        "claude-v3.5-sonnet-v2": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
+        "claude-v3.7-sonnet": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
         "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002},
         "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007},
         "mistral-large": {"input": 0.004, "output": 0.012},
-        "amazon-nova-pro": {"input": 0.0008, "output": 0.0032},
-        "amazon-nova-lite": {"input": 0.00006, "output": 0.00024},
-        "amazon-nova-micro": {"input": 0.000035, "output": 0.00014},
+        "amazon-nova-pro": {
+            "input": 0.0008,
+            "output": 0.0032,
+            "cache_write_input": 0.0008,
+            "cache_read_input": 0.0002,
+        },
+        "amazon-nova-lite": {
+            "input": 0.00006,
+            "output": 0.00024,
+            "cache_write_input": 0.00006,
+            "cache_read_input": 0.000015,
+        },
+        "amazon-nova-micro": {
+            "input": 0.000035,
+            "output": 0.00014,
+            "cache_write_input": 0.000035,
+            "cache_read_input": 0.00000875,
+        },
         "deepseek-r1": {"input": 0.00135, "output": 0.0054},
         # Meta Llama 3 models (US region)
         "llama3-3-70b-instruct": {"input": 0.00072, "output": 0.00072},
@@ -85,17 +130,47 @@ class EmbeddingConfig(TypedDict):
         "llama3-2-90b-instruct": {"input": 0.00072, "output": 0.00072},
     },
     "us-west-2": {
-        "claude-v4-opus": {"input": 0.015, "output": 0.075},
-        "claude-v4-sonnet": {"input": 0.003, "output": 0.015},
-        "claude-v3.7-sonnet": {"input": 0.00300, "output": 0.01500},
+        "claude-v4-opus": {
+            "input": 0.015,
+            "output": 0.075,
+            "cache_write_input": 0.01875,
+            "cache_read_input": 0.0015,
+        },
+        "claude-v4-sonnet": {
+            "input": 0.003,
+            "output": 0.015,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
+        "claude-v3.7-sonnet": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
         "claude-v3-opus": {"input": 0.01500, "output": 0.07500},
         "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002},
         "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007},
         "mistral-large": {"input": 0.004, "output": 0.012},
         "mistral-large-2": {"input": 0.002, "output": 0.06},
-        "amazon-nova-pro": {"input": 0.0008, "output": 0.0032},
-        "amazon-nova-lite": {"input": 0.00006, "output": 0.00024},
-        "amazon-nova-micro": {"input": 0.000035, "output": 0.00014},
+        "amazon-nova-pro": {
+            "input": 0.0008,
+            "output": 0.0032,
+            "cache_write_input": 0.0008,
+            "cache_read_input": 0.0002,
+        },
+        "amazon-nova-lite": {
+            "input": 0.00006,
+            "output": 0.00024,
+            "cache_write_input": 0.00006,
+            "cache_read_input": 0.000015,
+        },
+        "amazon-nova-micro": {
+            "input": 0.000035,
+            "output": 0.00014,
+            "cache_write_input": 0.000035,
+            "cache_read_input": 0.00000875,
+        },
         "deepseek-r1": {"input": 0.00135, "output": 0.0054},
         # Meta Llama 3 models (US region)
         "llama3-3-70b-instruct": {"input": 0.00072, "output": 0.00072},
@@ -106,21 +181,66 @@ class EmbeddingConfig(TypedDict):
     },
     "ap-northeast-1": {},
     "default": {
-        "claude-v4-opus": {"input": 0.015, "output": 0.075},
-        "claude-v4-sonnet": {"input": 0.003, "output": 0.015},
+        "claude-v4-opus": {
+            "input": 0.015,
+            "output": 0.075,
+            "cache_write_input": 0.01875,
+            "cache_read_input": 0.0015,
+        },
+        "claude-v4-sonnet": {
+            "input": 0.003,
+            "output": 0.015,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
         "claude-v3-haiku": {"input": 0.00025, "output": 0.00125},
-        "claude-v3.5-haiku": {"input": 0.001, "output": 0.005},
-        "claude-v3.5-sonnet": {"input": 0.00300, "output": 0.01500},
-        "claude-v3.5-sonnet-v2": {"input": 0.00300, "output": 0.01500},
-        "claude-v3.7-sonnet": {"input": 0.00300, "output": 0.01500},
+        "claude-v3.5-haiku": {
+            "input": 0.001,
+            "output": 0.005,
+            "cache_write_input": 0.001,
+            "cache_read_input": 0.00008,
+        },
+        "claude-v3.5-sonnet": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
+        "claude-v3.5-sonnet-v2": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
+        "claude-v3.7-sonnet": {
+            "input": 0.00300,
+            "output": 0.01500,
+            "cache_write_input": 0.00375,
+            "cache_read_input": 0.0003,
+        },
         "claude-v3-opus": {"input": 0.01500, "output": 0.07500},
         "mistral-7b-instruct": {"input": 0.00015, "output": 0.0002},
         "mixtral-8x7b-instruct": {"input": 0.00045, "output": 0.0007},
         "mistral-large": {"input": 0.004, "output": 0.012},
         "mistral-large-2": {"input": 0.002, "output": 0.06},
-        "amazon-nova-pro": {"input": 0.0008, "output": 0.0032},
-        "amazon-nova-lite": {"input": 0.00006, "output": 0.00024},
-        "amazon-nova-micro": {"input": 0.000035, "output": 0.00014},
+        "amazon-nova-pro": {
+            "input": 0.0008,
+            "output": 0.0032,
+            "cache_write_input": 0.0008,
+            "cache_read_input": 0.0002,
+        },
+        "amazon-nova-lite": {
+            "input": 0.00006,
+            "output": 0.00024,
+            "cache_write_input": 0.00006,
+            "cache_read_input": 0.000015,
+        },
+        "amazon-nova-micro": {
+            "input": 0.000035,
+            "output": 0.00014,
+            "cache_write_input": 0.000035,
+            "cache_read_input": 0.00000875,
+        },
         "deepseek-r1": {"input": 0.00135, "output": 0.0054},
         # Meta Llama 3 models (US region)
         "llama3-3-70b-instruct": {"input": 0.00072, "output": 0.00072},