Removed type requirement from types

rusiaaman · Dec 30, 2024 · 0ac067f · 0ac067f
1 parent afae1bd
commit 0ac067f
Show file tree

Hide file tree

Showing 4 changed files with 188 additions and 226 deletions.
diff --git a/src/wcgw/client/tools.py b/src/wcgw/client/tools.py
@@ -1127,7 +1127,7 @@ def take_help_of_ai_assistant(
 
 
 def which_tool(args: str) -> TOOLS:
-    adapter = TypeAdapter[TOOLS](TOOLS)
+    adapter = TypeAdapter[TOOLS](TOOLS, config={"extra": "forbid"})
     return adapter.validate_python(json.loads(args))
 
 
@@ -1180,7 +1180,7 @@ def get_tool_output(
 ) -> tuple[list[str | ImageData | DoneFlag], float]:
     global IS_IN_DOCKER, TOOL_CALLS
     if isinstance(args, dict):
-        adapter = TypeAdapter[TOOLS](TOOLS)
+        adapter = TypeAdapter[TOOLS](TOOLS, config={"extra": "forbid"})
         arg = adapter.validate_python(args)
     else:
         arg = args

diff --git a/src/wcgw/types_.py b/src/wcgw/types_.py
@@ -1,6 +1,14 @@
 from typing import Literal, Optional, Sequence
 
-from pydantic import BaseModel
+from pydantic import BaseModel as PydanticBaseModel
+
+
+class NoExtraArgs(PydanticBaseModel):
+    class Config:
+        extra = "forbid"
+
+
+BaseModel = NoExtraArgs
 
 
 class BashCommand(BaseModel):
@@ -14,7 +22,6 @@ class BashCommand(BaseModel):
 
 
 class BashInteraction(BaseModel):
-    type: Literal["BashInteraction"]
     send_text: Optional[str] = None
     send_specials: Optional[Sequence[Specials]] = None
     send_ascii: Optional[Sequence[int]] = None
@@ -23,7 +30,6 @@ class BashInteraction(BaseModel):
 
 class ReadImage(BaseModel):
     file_path: str
-    type: Literal["ReadImage"]
 
 
 class WriteIfEmpty(BaseModel):
@@ -33,7 +39,6 @@ class WriteIfEmpty(BaseModel):
 
 class ReadFiles(BaseModel):
     file_paths: list[str]
-    type: Literal["ReadFiles"]
 
 
 class FileEditFindReplace(BaseModel):
@@ -52,26 +57,22 @@ class FileEdit(BaseModel):
 
 
 class Initialize(BaseModel):
-    type: Literal["Initialize"]
     any_workspace_path: str
     initial_files_to_read: list[str]
 
 
 class GetScreenInfo(BaseModel):
-    type: Literal["GetScreenInfo"]
     docker_image_id: str
 
 
 class ScreenShot(BaseModel):
-    type: Literal["ScreenShot"]
     take_after_delay_seconds: int
 
 
 class MouseMove(BaseModel):
     x: int
     y: int
     do_left_click_on_move: bool
-    type: Literal["MouseMove"]
 
 
 class LeftClickDrag(BaseModel):