From 8023c93f1b6fc6b4a0902e233769ec034f3c0a91 Mon Sep 17 00:00:00 2001
From: Yukinobu Mine <yukimine@amazon.co.jp>
Date: Fri, 13 Dec 2024 01:20:23 +0900
Subject: [PATCH] Fix: mypy errors.

---
 backend/app/agents/tools/agent_tool.py | 26 +++++++++-----------------
 backend/app/bedrock.py                 |  9 ++++++---
 2 files changed, 15 insertions(+), 20 deletions(-)

diff --git a/backend/app/agents/tools/agent_tool.py b/backend/app/agents/tools/agent_tool.py
index cc32123e..d2187004 100644
--- a/backend/app/agents/tools/agent_tool.py
+++ b/backend/app/agents/tools/agent_tool.py
@@ -42,29 +42,21 @@ def run_result_to_tool_result_content_model(
         for related_document in run_result["related_documents"]
     ]
     if is_nova_model(model=model) and len(result_contents) > 1:
+        contents: list[str | dict[str, JsonValue]] = []
+        for result_content in result_contents:
+            if isinstance(result_content, JsonToolResultModel):
+                contents.append(result_content.json_)
+
+            elif isinstance(result_content, TextToolResultModel):
+                contents.append(result_content.text)
+
         return ToolResultContentModel(
             content_type="toolResult",
             body=ToolResultContentModelBody(
                 tool_use_id=run_result["tool_use_id"],
                 content=[
                     TextToolResultModel(
-                        text=json.dumps(
-                            [
-                                content
-                                for result_content in result_contents
-                                for content in (
-                                    [result_content.json_]
-                                    if isinstance(result_content, JsonToolResultModel)
-                                    else (
-                                        [result_content.text]
-                                        if isinstance(
-                                            result_content, TextToolResultModel
-                                        )
-                                        else []
-                                    )
-                                )
-                            ]
-                        ),
+                        text=json.dumps(contents),
                     ),
                 ],
                 status=run_result["status"],
diff --git a/backend/app/bedrock.py b/backend/app/bedrock.py
index 1255c8f0..a381971c 100644
--- a/backend/app/bedrock.py
+++ b/backend/app/bedrock.py
@@ -132,12 +132,15 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
     ]
 
     # Prepare model-specific parameters
+    inference_config: InferenceConfigurationTypeDef
+    additional_model_request_fields: dict[str, Any]
+    system_prompts: list[SystemContentBlockTypeDef]
     if is_nova_model(model):
         # Special handling for Nova models
         inference_config, additional_model_request_fields = _prepare_nova_model_params(
             model, generation_params
         )
-        system_prompts: list[SystemContentBlockTypeDef] = (
+        system_prompts = (
             [
                 {
                     "text": "\n\n".join(instructions),
@@ -149,7 +152,7 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
 
     else:
         # Standard handling for non-Nova models
-        inference_config: InferenceConfigurationTypeDef = {
+        inference_config = {
             "maxTokens": (
                 generation_params.max_tokens
                 if generation_params
@@ -178,7 +181,7 @@ def process_content(c: ContentModel, role: str) -> list[ContentBlockTypeDef]:
                 else DEFAULT_GENERATION_CONFIG["top_k"]
             )
         }
-        system_prompts: list[SystemContentBlockTypeDef] = [
+        system_prompts = [
             {
                 "text": instruction,
             }