fix: make completion action optional

lavague-ai · Oct 1, 2024 · c158919 · c158919
1 parent f7aefc3
commit c158919
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/lavague-sdk/lavague/sdk/trajectory/base.py b/lavague-sdk/lavague/sdk/trajectory/base.py
@@ -17,8 +17,9 @@ def is_running(self):
 
     def next_action(self):
         ret = self._controller.next_step(self.run_id)
-        self.actions.append(ret.action)
         self.status = ret.run_status
+        if ret.action:
+            self.actions.append(ret.action)
         return ret.action
 
     def run_to_completion(self):

diff --git a/lavague-sdk/lavague/sdk/trajectory/model.py b/lavague-sdk/lavague/sdk/trajectory/model.py
@@ -1,5 +1,5 @@
 from enum import Enum
-from typing import List, Tuple
+from typing import List, Tuple, Optional
 from pydantic import BaseModel, SerializeAsAny
 from lavague.sdk.action import Action
 
@@ -31,4 +31,4 @@ def write_to_file(self, file_path: str):
 
 class StepCompletion(BaseModel):
     run_status: RunStatus
-    action: Action
+    action: Optional[Action]