fix python engine

lavague-ai · Aug 26, 2024 · fe3eeb0 · fe3eeb0
1 parent 73cbca4
commit fe3eeb0
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/lavague-core/lavague/core/python_engine.py b/lavague-core/lavague/core/python_engine.py
@@ -39,6 +39,7 @@ class PythonEngine(BaseEngine):
     ocr_llm: BaseLLM
     batch_size: int
     confidence_threshold: float
+    fallback_theshold: float
     temp_screenshots_path: str
     n_search_attempts: int
 
@@ -54,6 +55,7 @@ def __init__(
         display: bool = False,
         batch_size: int = 5,
         confidence_threshold: float = 0.85,
+        fallback_threshold: float = 0.65,
         temp_screenshots_path="./tmp_screenshots",
         n_search_attemps=10,
     ):
@@ -71,6 +73,7 @@ def __init__(
         self.confidence_threshold = confidence_threshold
         self.temp_screenshots_path = temp_screenshots_path
         self.n_search_attempts = n_search_attemps
+        self.fallback_theshold = fallback_threshold
 
     @classmethod
     def from_context(cls, context: Context, driver: BaseDriver):
@@ -152,7 +155,7 @@ def perform_fallback(self, prompt, instruction) -> str:
                 image_documents=screenshots, prompt=prompt
             ).text.strip()
             output_dict = self.extract_json(output)
-            context_score = output_dict.get("score")
+            context_score = output_dict.get("score", 0)
             output = output_dict.get("ret")
             memory += output
 
@@ -208,7 +211,7 @@ def execute_instruction(self, instruction: str) -> ActionResult:
 
         try:
             if (
-                output_dict.get("score", 0) < self.confidence_threshold
+                output_dict.get("score", 0) < self.fallback_theshold
             ):  # use fallback method
                 output = self.perform_fallback(prompt=prompt, instruction=instruction)