From e99e9b6181f091a9625ef9b922dac15dd5f0a885 Mon Sep 17 00:00:00 2001
From: Reinier van der Leer <pwuts@agpt.co>
Date: Sat, 7 Oct 2023 19:10:18 -0700
Subject: [PATCH] AutoGPT/plugins: Support full parameter defs for plugin
 commands

---
 .../autogpt/agents/utils/prompt_scratchpad.py | 20 ++++++++++++-------
 .../autogpt/autogpt/core/utils/json_schema.py |  2 ++
 2 files changed, 15 insertions(+), 7 deletions(-)

diff --git a/autogpts/autogpt/autogpt/agents/utils/prompt_scratchpad.py b/autogpts/autogpt/autogpt/agents/utils/prompt_scratchpad.py
index f42f5eb9..07b32425 100644
--- a/autogpts/autogpt/autogpt/agents/utils/prompt_scratchpad.py
+++ b/autogpts/autogpt/autogpt/agents/utils/prompt_scratchpad.py
@@ -33,7 +33,7 @@ class PromptScratchpad(BaseModel):
         self,
         name: str,
         description: str,
-        params: dict[str, str],
+        params: dict[str, str | dict],
         function: Callable,
     ) -> None:
         """
@@ -50,15 +50,20 @@ class PromptScratchpad(BaseModel):
             function (callable, optional): A callable function to be called when
                 the command is executed. Defaults to None.
         """
-        for p, t in params.items():
+        for p, s in params.items():
             invalid = False
-            if t not in JSONSchema.Type._value2member_map_:
+            if type(s) == str and s not in JSONSchema.Type._value2member_map_:
                 invalid = True
                 logger.warning(
                     f"Cannot add command '{name}':"
-                    f" parameter '{p}' has invalid type '{t}'."
+                    f" parameter '{p}' has invalid type '{s}'."
                     f" Valid types are: {JSONSchema.Type._value2member_map_.keys()}"
                 )
+            elif isinstance(s, dict):
+                try:
+                    JSONSchema.from_dict(s)
+                except KeyError:
+                    invalid = True
             if invalid:
                 return
 
@@ -66,9 +71,10 @@ class PromptScratchpad(BaseModel):
             name=name,
             description=description,
             parameters={
-                # TODO: require plugins to specify parameters as a JSON schema
-                name: JSONSchema(type=JSONSchema.Type._value2member_map_[type])
-                for name, type in params.items()
+                name: JSONSchema(type=JSONSchema.Type._value2member_map_[spec])
+                if type(spec) == str
+                else JSONSchema.from_dict(spec)
+                for name, spec in params.items()
             },
             method=function,
         )
diff --git a/autogpts/autogpt/autogpt/core/utils/json_schema.py b/autogpts/autogpt/autogpt/core/utils/json_schema.py
index 9529f912..743d70d5 100644
--- a/autogpts/autogpt/autogpt/core/utils/json_schema.py
+++ b/autogpts/autogpt/autogpt/core/utils/json_schema.py
@@ -148,6 +148,8 @@ class JSONSchema(BaseModel):
             if not self.properties:
                 return "Record<string, any>"
             return self.to_typescript_object_interface()
+        elif self.enum:
+            return " | ".join(repr(v) for v in self.enum)
         else:
             raise NotImplementedError(
                 f"JSONSchema.typescript_type does not support Type.{self.type.name} yet"