From 253231c3531b3488542c86b33eb0ba1fb10e197c Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <pwuts@agpt.co>
Date: Sat, 8 Jun 2024 20:40:50 +0200
Subject: [PATCH] fix type issues

---
 autogpt/autogpt/agents/agent.py                       | 10 +++++++---
 autogpt/autogpt/agents/prompt_strategies/code_flow.py | 11 ++++++-----
 forge/forge/command/command.py                        |  2 +-
 forge/forge/command/decorator.py                      |  2 +-
 4 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/autogpt/autogpt/agents/agent.py b/autogpt/autogpt/agents/agent.py
index 0082bc0adfc..43dab9e8a68 100644
--- a/autogpt/autogpt/agents/agent.py
+++ b/autogpt/autogpt/agents/agent.py
@@ -32,7 +32,6 @@
 from forge.components.watchdog import WatchdogComponent
 from forge.components.web import WebSearchComponent, WebSeleniumComponent
 from forge.file_storage.base import FileStorage
-from forge.llm.prompting import PromptStrategy
 from forge.llm.prompting.schema import ChatPrompt
 from forge.llm.prompting.utils import dump_prompt
 from forge.llm.providers import (
@@ -64,7 +63,10 @@
 )
 
 from .prompt_strategies.code_flow import CodeFlowAgentPromptStrategy
-from .prompt_strategies.one_shot import OneShotAgentActionProposal
+from .prompt_strategies.one_shot import (
+    OneShotAgentActionProposal,
+    OneShotAgentPromptStrategy,
+)
 
 if TYPE_CHECKING:
     from forge.config.config import Config
@@ -101,7 +103,9 @@ def __init__(
         llm_provider: MultiProvider,
         file_storage: FileStorage,
         legacy_config: Config,
-        prompt_strategy_class: type[PromptStrategy] = CodeFlowAgentPromptStrategy,
+        prompt_strategy_class: type[
+            OneShotAgentPromptStrategy | CodeFlowAgentPromptStrategy
+        ] = CodeFlowAgentPromptStrategy,
     ):
         super().__init__(settings)
 
diff --git a/autogpt/autogpt/agents/prompt_strategies/code_flow.py b/autogpt/autogpt/agents/prompt_strategies/code_flow.py
index eb458e85b69..2e5ccce216d 100644
--- a/autogpt/autogpt/agents/prompt_strategies/code_flow.py
+++ b/autogpt/autogpt/agents/prompt_strategies/code_flow.py
@@ -172,6 +172,7 @@ def build_system_prompt(
 
     def response_format_instruction(self) -> tuple[str, str]:
         response_schema = self.response_schema.copy(deep=True)
+        assert response_schema.properties
 
         # Unindent for performance
         response_format = re.sub(
@@ -295,15 +296,15 @@ async def parse_response_content(
         ).validate_code(parsed_response.python_code)
 
         # TODO: prevent combining finish with other functions
-        if re.search(r"finish\((.*?)\)", code_validation.functionCode):
-            finish_reason = re.search(
-                r"finish\((reason=)?(.*?)\)", code_validation.functionCode
-            ).group(2)
+        if _finish_call := re.search(
+            r"finish\((reason=)?(.*?)\)", code_validation.functionCode
+        ):
+            finish_reason = _finish_call.group(2)[1:-1]  # remove quotes
             result = OneShotAgentActionProposal(
                 thoughts=parsed_response.thoughts,
                 use_tool=AssistantFunctionCall(
                     name="finish",
-                    arguments={"reason": finish_reason[1:-1]},
+                    arguments={"reason": finish_reason},
                 ),
             )
         else:
diff --git a/forge/forge/command/command.py b/forge/forge/command/command.py
index bbb8b5a0ffc..a0548dc650a 100644
--- a/forge/forge/command/command.py
+++ b/forge/forge/command/command.py
@@ -24,7 +24,7 @@ def __init__(
         self,
         names: list[str],
         description: str,
-        method: Callable[P, CO] | Callable[Concatenate[CommandProvider, P], CO],
+        method: Callable[Concatenate[CommandProvider, P], CO] | Callable[P, CO],
         parameters: list[CommandParameter],
     ):
         self.names = names
diff --git a/forge/forge/command/decorator.py b/forge/forge/command/decorator.py
index 0a665330d91..753cd7afd83 100644
--- a/forge/forge/command/decorator.py
+++ b/forge/forge/command/decorator.py
@@ -16,7 +16,7 @@ def command(
     description: Optional[str] = None,
     parameters: Optional[dict[str, JSONSchema]] = None,
 ) -> Callable[
-    [Callable[P, CO] | Callable[Concatenate[CommandProvider, P], CO]], Command[P, CO]
+    [Callable[Concatenate[CommandProvider, P], CO] | Callable[P, CO]], Command[P, CO]
 ]:
     """
     Make a `Command` from a function or a method on a `CommandProvider`.