fix type issues

autogpt/autogpt/agents/agent.py

@@ -32,7 +32,6 @@ from forge.components.user_interaction import UserInteractionComponent
 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 autogpt.app.log_cycle import (
 )
 
 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 @@ class Agent(BaseAgent[OneShotAgentActionProposal], Configurable[AgentSettings]):
         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)
 

autogpt/autogpt/agents/prompt_strategies/code_flow.py

@@ -172,6 +172,7 @@ class CodeFlowAgentPromptStrategy(PromptStrategy):
 
     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 @@ class CodeFlowAgentPromptStrategy(PromptStrategy):
         ).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:

forge/forge/command/command.py

@@ -1,9 +1,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import Callable, Concatenate, Generic, ParamSpec, TypeVar, cast
-
-from forge.agent.protocols import CommandProvider
+from typing import Callable, Generic, ParamSpec, TypeVar
 
 from .parameter import CommandParameter
 
@@ -24,14 +22,12 @@ class Command(Generic[P, CO]):
         self,
         names: list[str],
         description: str,
-        method: Callable[P, CO] | Callable[Concatenate[CommandProvider, P], CO],
+        method: Callable[P, CO],
         parameters: list[CommandParameter],
     ):
         self.names = names
         self.description = description
-        # Method technically has a `self` parameter, but we can ignore that
-        # since Python passes it internally.
-        self.method = cast(Callable[P, CO], method)
+        self.method = method
         self.parameters = parameters
 
         # Check if all parameters are provided
@@ -49,10 +45,10 @@ class Command(Generic[P, CO]):
         return inspect.iscoroutinefunction(self.method)
 
     @property
-    def return_type(self) -> type:
+    def return_type(self) -> str:
         _type = inspect.signature(self.method).return_annotation
         if _type == inspect.Signature.empty:
-            return None
+            return "None"
         return _type.__name__
 
     @property

forge/forge/command/decorator.py

@@ -1,7 +1,7 @@
 import inspect
 import logging
 import re
-from typing import Callable, Concatenate, Optional
+from typing import Callable, Concatenate, Optional, TypeVar, cast
 
 from forge.agent.protocols import CommandProvider
 from forge.models.json_schema import JSONSchema
@@ -10,14 +10,17 @@ from .command import CO, Command, CommandParameter, P
 
 logger = logging.getLogger(__name__)
 
+_CP = TypeVar("_CP", bound=CommandProvider)
+
 
 def command(
     names: Optional[list[str]] = None,
     description: Optional[str] = None,
     parameters: Optional[dict[str, JSONSchema]] = None,
-) -> Callable[
-    [Callable[P, CO] | Callable[Concatenate[CommandProvider, P], CO]], Command[P, CO]
-]:
+) -> (
+    Callable[[Callable[Concatenate[_CP, P], CO]], Command[P, CO]]
+    | Callable[[Callable[P, CO]], Command[P, CO]]
+):
     """
     Make a `Command` from a function or a method on a `CommandProvider`.
     All parameters are optional if the decorated function has a fully featured
@@ -35,7 +38,7 @@ def command(
     """
 
     def decorator(
-        func: Callable[P, CO] | Callable[Concatenate[CommandProvider, P], CO]
+        func: Callable[P, CO] | Callable[Concatenate[_CP, P], CO]
     ) -> Command[P, CO]:
         # If names is not provided, use the function name
         _names = names or [func.__name__]
@@ -71,17 +74,18 @@ def command(
         ]
 
         # Wrap func with Command
-        command = Command[P, CO](
+        command = Command(
             names=_names,
             description=_description,
-            method=func,
+            # Method technically has a `self` parameter, but we can ignore that
+            # since Python passes it internally.
+            method=cast(Callable[P, CO], func),
             parameters=typed_parameters,
         )
 
         return command
 
-    # https://github.com/microsoft/pyright/issues/7369#issuecomment-2155968296
-    return decorator  # pyright: ignore
+    return decorator
 
 
 def get_clean_description_from_docstring(docstring: str) -> str: