1 files changed, 27 insertions, 11 deletions

diff --git a/autogpts/autogpt/autogpt/core/planning/prompt_strategies/next_ability.py b/autogpts/autogpt/autogpt/core/planning/prompt_strategies/next_ability.py
index 6efed7bee..9672f8e5b 100644
--- a/autogpts/autogpt/autogpt/core/planning/prompt_strategies/next_ability.py
+++ b/autogpts/autogpt/autogpt/core/planning/prompt_strategies/next_ability.py
@@ -6,7 +6,7 @@ from autogpt.core.prompting import PromptStrategy
 from autogpt.core.prompting.schema import ChatPrompt, LanguageModelClassification
 from autogpt.core.prompting.utils import json_loads, to_numbered_list
 from autogpt.core.resource.model_providers import (
-    AssistantChatMessageDict,
+    AssistantChatMessage,
     ChatMessage,
     CompletionModelFunction,
 )
@@ -41,25 +41,36 @@ class NextAbility(PromptStrategy):
         "{additional_info}\n\n"
         "Additionally, you should consider the following:\n"
         "{user_input}\n\n"
-        "Your task of {task_objective} is complete when the following acceptance criteria have been met:\n"
+        "Your task of {task_objective} is complete when the following acceptance"
+        " criteria have been met:\n"
         "{acceptance_criteria}\n\n"
         "Please choose one of the provided functions to accomplish this task. "
-        "Some tasks may require multiple functions to accomplish. If that is the case, choose the function that "
-        "you think is most appropriate for the current situation given your progress so far."
+        "Some tasks may require multiple functions to accomplish. If that is the case,"
+        " choose the function that you think is most appropriate for the current"
+        " situation given your progress so far."
     )
 
     DEFAULT_ADDITIONAL_ABILITY_ARGUMENTS = {
         "motivation": JSONSchema(
             type=JSONSchema.Type.STRING,
-            description="Your justification for choosing choosing this function instead of a different one.",
+            description=(
+                "Your justification for choosing choosing this function instead of a "
+                "different one."
+            ),
         ),
         "self_criticism": JSONSchema(
             type=JSONSchema.Type.STRING,
-            description="Thoughtful self-criticism that explains why this function may not be the best choice.",
+            description=(
+                "Thoughtful self-criticism that explains why this function may not be "
+                "the best choice."
+            ),
         ),
         "reasoning": JSONSchema(
             type=JSONSchema.Type.STRING,
-            description="Your reasoning for choosing this function taking into account the `motivation` and weighing the `self_criticism`.",
+            description=(
+                "Your reasoning for choosing this function taking into account the "
+                "`motivation` and weighing the `self_criticism`."
+            ),
         ),
     }
 
@@ -124,7 +135,9 @@ class NextAbility(PromptStrategy):
         template_kwargs["additional_info"] = to_numbered_list(
             [memory.summary() for memory in task.context.memories]
             + [info for info in task.context.supplementary_info],
-            no_items_response="There is no additional information available at this time.",
+            no_items_response=(
+                "There is no additional information available at this time."
+            ),
             **template_kwargs,
         )
         template_kwargs["user_input"] = to_numbered_list(
@@ -158,7 +171,7 @@ class NextAbility(PromptStrategy):
 
     def parse_response_content(
         self,
-        response_content: AssistantChatMessageDict,
+        response_content: AssistantChatMessage,
     ) -> dict:
         """Parse the actual text response from the objective model.
 
@@ -170,9 +183,12 @@ class NextAbility(PromptStrategy):
 
         """
         try:
-            function_name = response_content["tool_calls"][0]["function"]["name"]
+            if not response_content.tool_calls:
+                raise ValueError("LLM did not call any function")
+
+            function_name = response_content.tool_calls[0].function.name
             function_arguments = json_loads(
-                response_content["tool_calls"][0]["function"]["arguments"]
+                response_content.tool_calls[0].function.arguments
             )
             parsed_response = {
                 "motivation": function_arguments.pop("motivation"),