Fix: mypy errors.

aws-samples · Dec 12, 2024 · 8023c93 · 8023c93
1 parent 493b012
commit 8023c93
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 20 deletions.
diff --git 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
@@ -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,
             }