Den 204 use dendrite exceptions in sdk (#32)

* update ask-page exception handling and api call

* add new exceptions and rename old

* refactor return None instead of exception

* update download exception

* fix missing none check for page click and fill

* make screenshot optional for exception

* change default retry count to 2

* update default kwarg for page click and fill

* update default get_element retry count to 2 from 3

* rewrite Exception to DendriteException

Co-authored-by: Paul Sanders <[email protected]>

Author: coffecup25

dendrite_sdk/__init__.py

@@ -3,8 +3,7 @@
 from ._core.dendrite_element import DendriteElement
 from ._core.dendrite_page import DendritePage
 from ._core.models.response import DendriteElementsResponse
-from ._exceptions.dendrite_exception import DendriteException
-from ._exceptions.incorrect_outcome_exception import IncorrectOutcomeException
+
 
 logger.disable("dendrite_python_sdk")
 
@@ -13,6 +12,4 @@
     "DendriteElement",
     "DendritePage",
     "DendriteElementsResponse",
-    "DendriteException",
-    "IncorrectOutcomeException",
 ]

dendrite_sdk/_api/browser_api_client.py

@@ -2,6 +2,7 @@
 
 from loguru import logger
 from dendrite_sdk._core.models.authentication import AuthSession
+from dendrite_sdk._api.response.get_element_response import GetElementResponse
 from dendrite_sdk._api.dto.ask_page_dto import AskPageDTO
 from dendrite_sdk._api.dto.authenticate_dto import AuthenticateDTO
 from dendrite_sdk._api.dto.get_elements_dto import GetElementsDTO
@@ -33,11 +34,13 @@ async def upload_auth_session(self, dto: UploadAuthSessionDTO):
             "actions/upload-auth-session", data=dto.dict(), method="POST"
         )
 
-    async def get_interactions_selector(self, dto: GetElementsDTO) -> dict:
+    async def get_interactions_selector(
+        self, dto: GetElementsDTO
+    ) -> GetElementResponse:
         res = await self.send_request(
             "actions/get-interaction-selector", data=dto.dict(), method="POST"
         )
-        return res.json()
+        return GetElementResponse(**res.json())
 
     async def make_interaction(self, dto: MakeInteractionDTO) -> InteractionResponse:
         res = await self.send_request(
@@ -67,7 +70,9 @@ async def ask_page(self, dto: AskPageDTO) -> AskPageResponse:
         )
         res_dict = res.json()
         return AskPageResponse(
-            description=res_dict["description"], return_data=res_dict["return_data"]
+            status=res_dict["status"],
+            description=res_dict["description"],
+            return_data=res_dict["return_data"],
         )
 
     async def try_run_cached(

dendrite_sdk/_api/response/ask_page_response.py

@@ -1,10 +1,11 @@
-from typing import Generic, TypeVar
+from typing import Generic, Literal, TypeVar
 from pydantic import BaseModel
 
 
 T = TypeVar("T")
 
 
 class AskPageResponse(BaseModel, Generic[T]):
+    status: Literal["success", "error"]
     return_data: T
     description: str

dendrite_sdk/_api/response/get_element_response.py

@@ -0,0 +1,8 @@
+from typing import List, Optional
+
+from pydantic import BaseModel
+
+
+class GetElementResponse(BaseModel):
+    selectors: Optional[List[str]] = None
+    message: str = ""

dendrite_sdk/_common/constants.py

@@ -64,7 +64,7 @@
     "--font-render-hinting=none",
     "--disable-logging",
     "--enable-surface-synchronization",
-    #"--run-all-compositor-stages-before-draw",
+    # "--run-all-compositor-stages-before-draw",
     "--disable-threaded-animation",
     "--disable-threaded-scrolling",
     "--disable-checker-imaging",

dendrite_sdk/_common/event_sync.py

@@ -1,6 +1,8 @@
 import asyncio
 from typing import Generic, Optional, TypeVar
 
+from dendrite_sdk._exceptions.dendrite_exception import DendriteException
+
 
 T = TypeVar("T")
 
@@ -65,7 +67,7 @@ async def get_data(self, timeout: float = 30000) -> T:
         try:
             await asyncio.wait_for(self.event.wait(), timeout * 0.001)
             if not self.data:
-                raise Exception(f"No {self.data.__class__.__name__} was found.")
+                raise DendriteException(f"No {self.data.__class__.__name__} was found.")
             data = self.data
 
             self.data = None

dendrite_sdk/_core/_base_browser.py

@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 import sys
-from typing import Any, Generic, List, Optional, TypeVar, Union
+from typing import Any, List, Optional, Union
 from uuid import uuid4
 import os
 from loguru import logger
@@ -21,13 +21,17 @@
 
 from dendrite_sdk._core.dendrite_page import DendritePage
 from dendrite_sdk._common.constants import STEALTH_ARGS
-from dendrite_sdk._core.models.download_interface import DownloadInterface
 from dendrite_sdk._core.models.authentication import (
     AuthSession,
 )
 from dendrite_sdk._core.models.llm_config import LLMConfig
 from dendrite_sdk._api.browser_api_client import BrowserAPIClient
-from dendrite_sdk._exceptions.dendrite_exception import BrowserNotLaunchedError
+from dendrite_sdk._exceptions.dendrite_exception import (
+    BrowserNotLaunchedError,
+    DendriteException,
+    IncorrectOutcomeError,
+    MissingApiKeyError,
+)
 
 
 class BaseDendriteBrowser(ABC):
@@ -81,17 +85,23 @@ def __init__(
         if not dendrite_api_key or dendrite_api_key == "":
             dendrite_api_key = os.environ.get("DENDRITE_API_KEY", "")
             if not dendrite_api_key or dendrite_api_key == "":
-                raise Exception("Dendrite API key is required to use DendriteBrowser")
+                raise MissingApiKeyError(
+                    "Dendrite API key is required to use DendriteBrowser"
+                )
 
         if not anthropic_api_key or anthropic_api_key == "":
             anthropic_api_key = os.environ.get("ANTHROPIC_API_KEY", "")
             if anthropic_api_key == "":
-                raise Exception("Anthropic API key is required to use DendriteBrowser")
+                raise MissingApiKeyError(
+                    "Anthropic API key is required to use DendriteBrowser"
+                )
 
         if not openai_api_key or openai_api_key == "":
             openai_api_key = os.environ.get("OPENAI_API_KEY", "")
             if not openai_api_key or openai_api_key == "":
-                raise Exception("OpenAI API key is required to use DendriteBrowser")
+                raise MissingApiKeyError(
+                    "OpenAI API key is required to use DendriteBrowser"
+                )
 
         self._id = uuid4().hex
         self._auth_data: Optional[AuthSession] = None
@@ -188,8 +198,8 @@ async def goto(
             try:
                 prompt = f"We are checking if we have arrived on the expected type of page. If it is apparent that we have arrived on the wrong page, output an error. Here is the description: '{expected_page}'"
                 await active_page.ask(prompt, bool)
-            except Exception as e:
-                raise Exception(f"Incorrect navigation, reason: {e}")
+            except DendriteException as e:
+                raise IncorrectOutcomeError(f"Incorrect navigation, reason: {e}")
 
         return active_page
 
@@ -265,7 +275,7 @@ async def add_cookies(self, cookies):
             Exception: If the browser context is not initialized.
         """
         if not self.browser_context:
-            raise Exception("Browser context not initialized")
+            raise DendriteException("Browser context not initialized")
 
         await self.browser_context.add_cookies(cookies)
 
@@ -292,9 +302,11 @@ async def close(self):
                 )
                 await self._browser_api_client.upload_auth_session(dto)
             await self.browser_context.close()
-
-        if self._playwright:
-            await self._playwright.stop()
+        try:
+            if self._playwright:
+                await self._playwright.stop()
+        except AttributeError:
+            pass
 
     def _is_launched(self):
         """

dendrite_sdk/_core/_utils.py

@@ -79,35 +79,3 @@ def get_frame_context(page: Page, iframe_path: str) -> Union[FrameLocator, Page]
     for iframe_id in iframe_path_list:
         frame_context = frame_context.frame_locator(f"[tf623_id='{iframe_id}']")
     return frame_context
-
-
-async def get_all_elements_from_selector(
-    page: "DendritePage", selector: str
-) -> List[DendriteElement]:
-    dendrite_elements: List[DendriteElement] = []
-    soup = await page._get_soup()
-    elements = soup.select(selector)
-
-    for element in elements:
-        frame = page._get_context(element)
-        d_id = element.get("d-id", "")
-        locator = frame.locator(f"xpath=//*[@d-id='{d_id}']")
-
-        if not d_id:
-            continue
-
-        if isinstance(d_id, list):
-            d_id = d_id[0]
-
-        dendrite_elements.append(
-            DendriteElement(
-                d_id,
-                locator,
-                page.dendrite_browser,
-            )
-        )
-
-    if len(dendrite_elements) == 0:
-        raise Exception(f"No elements found for selector '{selector}'")
-
-    return dendrite_elements

dendrite_sdk/_core/dendrite_element.py

@@ -8,15 +8,14 @@
 from loguru import logger
 from playwright.async_api import Locator
 
+from dendrite_sdk._exceptions.dendrite_exception import IncorrectOutcomeError
+
 if TYPE_CHECKING:
     from dendrite_sdk._core._base_browser import BaseDendriteBrowser
 from dendrite_sdk._core.models.page_diff_information import PageDiffInformation
 from dendrite_sdk._core._type_spec import Interaction
 from dendrite_sdk._api.response.interaction_response import InteractionResponse
 from dendrite_sdk._api.dto.make_interaction_dto import MakeInteractionDTO
-from dendrite_sdk._exceptions.incorrect_outcome_exception import (
-    IncorrectOutcomeException,
-)
 
 
 def perform_action(interaction_type: Interaction):
@@ -81,7 +80,7 @@ async def wrapper(
             res = await self._browser_api_client.make_interaction(dto)
 
             if res.status == "failed":
-                raise IncorrectOutcomeException(
+                raise IncorrectOutcomeError(
                     message=res.message,
                     screenshot_base64=page_delta_information.page_after.screenshot_base64,
                 )

dendrite_sdk/_core/dendrite_page.py

@@ -37,7 +37,10 @@
 
 
 from dendrite_sdk._core._managers.screenshot_manager import ScreenshotManager
-from dendrite_sdk._exceptions.dendrite_exception import DendriteException
+from dendrite_sdk._exceptions.dendrite_exception import (
+    DendriteException,
+    PageConditionNotMet,
+)
 
 
 from dendrite_sdk._core._utils import (
@@ -247,7 +250,7 @@ async def _generate_dendrite_ids(self):
                 )
                 tries += 1
 
-        raise Exception("Failed to add d-ids to DOM.")
+        raise DendriteException("Failed to add d-ids to DOM.")
 
     async def scroll_through_entire_page(self) -> None:
         """
@@ -289,30 +292,31 @@ async def wait_for(
         )  # HACK: Wait for page to load slightly when running first time
         while num_attempts < max_retries:
             num_attempts += 1
+            start_time = time.time()
+
+            page_information = await self._get_page_information()
+            prompt = f"Prompt: '{prompt}'\n\nReturn a boolean that determines if the requested information or thing is available on the page."
             try:
-                start_time = time.time()
-                page_information = await self._get_page_information()
-                prompt = f"Prompt: '{prompt}'\n\nReturn a boolean that determines if the requested information or thing is available on the page."
                 res = await self.ask(prompt, bool)
-                elapsed_time = (
-                    time.time() - start_time
-                ) * 1000  # Convert to milliseconds
-
-                if res:
-                    return res
-
-                if elapsed_time >= timeout:
-                    # If the response took longer than the timeout, continue immediately
-                    continue
-                else:
-                    # Otherwise, wait for the remaining time
-                    await asyncio.sleep((timeout - elapsed_time) * 0.001)
-            except Exception as e:
-                logger.debug(f"Waited for page, but got this exception: {e}")
+            except DendriteException as e:
+                logger.debug(
+                    f"Attempt {num_attempts}/{max_retries} failed: {e.message}"
+                )
+
+            elapsed_time = (time.time() - start_time) * 1000  # Convert to milliseconds
+
+            if res:
+                return res
+
+            if elapsed_time >= timeout:
+                # If the response took longer than the timeout, continue immediately
                 continue
+            else:
+                # Otherwise, wait for the remaining time
+                await asyncio.sleep((timeout - elapsed_time) * 0.001)
 
         page_information = await self._get_page_information()
-        raise DendriteException(
+        raise PageConditionNotMet(
             message=f"Retried {max_retries} times but failed to wait for the requested condition.",
             screenshot_base64=page_information.screenshot_base64,
         )
@@ -326,7 +330,7 @@ async def click(
         timeout: int = 2000,
         force: bool = False,
         *args,
-        kwargs,
+        kwargs={},
     ) -> InteractionResponse:
         """
         Clicks an element on the page based on the provided prompt.
@@ -356,6 +360,13 @@ async def click(
             max_retries=max_retries,
             timeout=timeout,
         )
+
+        if not element:
+            raise DendriteException(
+                message=f"No element found with the prompt: {prompt}",
+                screenshot_base64="",
+            )
+
         return await element.click(
             expected_outcome=expected_outcome,
             timeout=timeout,
@@ -373,7 +384,7 @@ async def fill(
         max_retries: int = 3,
         timeout: int = 2000,
         *args,
-        kwargs,
+        kwargs={},
     ) -> InteractionResponse:
         """
         Fills an element on the page with the provided value based on the given prompt.
@@ -404,6 +415,12 @@ async def fill(
             timeout=timeout,
         )
 
+        if not element:
+            raise DendriteException(
+                message=f"No element found with the prompt: {prompt}",
+                screenshot_base64="",
+            )
+
         return await element.fill(
             value,
             expected_outcome=expected_outcome,
@@ -473,7 +490,9 @@ async def _expand_iframes(self, page_source: BeautifulSoup):
         """
         await expand_iframes(self.playwright_page, page_source)
 
-    async def _get_all_elements_from_selector(self, selector: str):
+    async def _get_all_elements_from_selector(
+        self, selector: str
+    ) -> List[DendriteElement]:
         dendrite_elements: List[DendriteElement] = []
         soup = await self._get_soup()
         elements = soup.select(selector)
@@ -497,9 +516,6 @@ async def _get_all_elements_from_selector(self, selector: str):
                 )
             )
 
-        if len(dendrite_elements) == 0:
-            raise Exception(f"No elements found for selector '{selector}'")
-
         return dendrite_elements
 
     async def _dump_html(self, path: str) -> None: