Den 281 extract as markdown (#58)

* began adding markdown functionality

* Changed all legacy naming to "extract". Added _last_frame_navigated_timestamp functionality to page to detect if page is still loading. Added markdownify

* Loading added to extract. Also added small sleep to screenshot manager.

* Cleaned up screenshot manager and added a screenshot mixin.

* update poetry.lock

---------

Co-authored-by: Arian Hanifi <[email protected]>

Author: charlesmaddock

dendrite_sdk/async_api/_api/browser_api_client.py

@@ -7,14 +7,14 @@
 from dendrite_sdk.async_api._api.dto.authenticate_dto import AuthenticateDTO
 from dendrite_sdk.async_api._api.dto.get_elements_dto import GetElementsDTO
 from dendrite_sdk.async_api._api.dto.make_interaction_dto import MakeInteractionDTO
-from dendrite_sdk.async_api._api.dto.scrape_page_dto import ScrapePageDTO
+from dendrite_sdk.async_api._api.dto.extract_dto import ExtractDTO
 from dendrite_sdk.async_api._api.dto.try_run_script_dto import TryRunScriptDTO
 from dendrite_sdk.async_api._api.dto.upload_auth_session_dto import UploadAuthSessionDTO
 from dendrite_sdk.async_api._api.response.ask_page_response import AskPageResponse
 from dendrite_sdk.async_api._api.response.interaction_response import (
     InteractionResponse,
 )
-from dendrite_sdk.async_api._api.response.scrape_page_response import ScrapePageResponse
+from dendrite_sdk.async_api._api.response.extract_response import ExtractResponse
 from dendrite_sdk.async_api._api._http_client import HTTPClient
 from dendrite_sdk._common._exceptions.dendrite_exception import (
     InvalidAuthSessionError,
@@ -55,12 +55,12 @@ async def make_interaction(self, dto: MakeInteractionDTO) -> InteractionResponse
             status=res_dict["status"], message=res_dict["message"]
         )
 
-    async def scrape_page(self, dto: ScrapePageDTO) -> ScrapePageResponse:
+    async def extract(self, dto: ExtractDTO) -> ExtractResponse:
         res = await self.send_request(
-            "actions/scrape-page", data=dto.dict(), method="POST"
+            "actions/extract-page", data=dto.dict(), method="POST"
         )
         res_dict = res.json()
-        return ScrapePageResponse(
+        return ExtractResponse(
             status=res_dict["status"],
             message=res_dict["message"],
             return_data=res_dict["return_data"],
@@ -79,9 +79,7 @@ async def ask_page(self, dto: AskPageDTO) -> AskPageResponse:
             return_data=res_dict["return_data"],
         )
 
-    async def try_run_cached(
-        self, dto: TryRunScriptDTO
-    ) -> Optional[ScrapePageResponse]:
+    async def try_run_cached(self, dto: TryRunScriptDTO) -> Optional[ExtractResponse]:
         res = await self.send_request(
             "actions/try-run-cached", data=dto.dict(), method="POST"
         )
@@ -92,7 +90,7 @@ async def try_run_cached(
         if loaded_value is None:
             return None
 
-        return ScrapePageResponse(
+        return ExtractResponse(
             status=res_dict["status"],
             message=res_dict["message"],
             return_data=loaded_value,

…dk/async_api/_api/dto/scrape_page_dto.py …te_sdk/async_api/_api/dto/extract_dto.pydendrite_sdk/async_api/_api/dto/scrape_page_dto.py renamed to dendrite_sdk/async_api/_api/dto/extract_dto.py dendrite_sdk/async_api/_api/dto/scrape_page_dto.py renamed to dendrite_sdk/async_api/_api/dto/extract_dto.py

@@ -5,7 +5,7 @@
 from dendrite_sdk.async_api._core.models.page_information import PageInformation
 
 
-class ScrapePageDTO(BaseModel):
+class ExtractDTO(BaseModel):
     page_information: PageInformation
     api_config: APIConfig
     prompt: str

…pi/_api/response/scrape_page_response.py …nc_api/_api/response/extract_response.pydendrite_sdk/async_api/_api/response/scrape_page_response.py renamed to dendrite_sdk/async_api/_api/response/extract_response.py dendrite_sdk/async_api/_api/response/scrape_page_response.py renamed to dendrite_sdk/async_api/_api/response/extract_response.py

@@ -7,7 +7,7 @@
 T = TypeVar("T")
 
 
-class ScrapePageResponse(BaseModel, Generic[T]):
+class ExtractResponse(BaseModel, Generic[T]):
     return_data: T
     message: str
     created_script: Optional[str] = None

dendrite_sdk/async_api/_core/_managers/screenshot_manager.py

@@ -3,25 +3,27 @@
 import os
 from typing import Tuple
 from uuid import uuid4
-from playwright.async_api import Page
 
 from dendrite_sdk.async_api._core._type_spec import PlaywrightPage
 
 
 class ScreenshotManager:
-    def __init__(self) -> None:
+    def __init__(self, page: PlaywrightPage) -> None:
         self.screenshot_before: str = ""
         self.screenshot_after: str = ""
+        self.page = page
 
-    async def take_full_page_screenshot(self, page: Page) -> str:
-        image_data = await page.screenshot(type="jpeg", full_page=True, timeout=30000)
+    async def take_full_page_screenshot(self) -> str:
+        image_data = await self.page.screenshot(
+            type="jpeg", full_page=True, timeout=30000
+        )
         if image_data is None:
             return ""
 
         return base64.b64encode(image_data).decode("utf-8")
 
-    async def take_viewport_screenshot(self, page: Page) -> str:
-        image_data = await page.screenshot(type="jpeg", timeout=30000)
+    async def take_viewport_screenshot(self) -> str:
+        image_data = await self.page.screenshot(type="jpeg", timeout=30000)
 
         if image_data is None:
             return ""
@@ -30,18 +32,6 @@ async def take_viewport_screenshot(self, page: Page) -> str:
 
         return reduced_base64
 
-    async def start_recording_diff(self, page: Page):
-        self.screenshot_before = await self.take_viewport_screenshot(page)
-
-    async def get_diff_images(
-        self,
-        page: PlaywrightPage,
-        wait_time=1,
-    ) -> Tuple[str, str]:
-        await asyncio.sleep(wait_time)
-        self.screenshot_after = await self.take_viewport_screenshot(page)
-        return self.screenshot_before, self.screenshot_after
-
     def store_screenshot(self, name, image_data):
         if not name:
             name = str(uuid4())

dendrite_sdk/async_api/_core/dendrite_browser.py

@@ -31,7 +31,9 @@
 from dendrite_sdk.async_api._core.mixin.fill_fields import FillFieldsMixin
 from dendrite_sdk.async_api._core.mixin.get_element import GetElementMixin
 from dendrite_sdk.async_api._core.mixin.keyboard import KeyboardMixin
+from dendrite_sdk.async_api._core.mixin.screenshot import ScreenshotMixin
 from dendrite_sdk.async_api._core.mixin.wait_for import WaitForMixin
+from dendrite_sdk.async_api._core.mixin.markdown import MarkdownMixin
 from dendrite_sdk.async_api._core.models.authentication import (
     AuthSession,
 )
@@ -47,8 +49,10 @@
 
 
 class AsyncDendrite(
-    ExtractionMixin,
+    ScreenshotMixin,
     WaitForMixin,
+    MarkdownMixin,
+    ExtractionMixin,
     AskMixin,
     FillFieldsMixin,
     ClickMixin,

dendrite_sdk/async_api/_core/dendrite_page.py

@@ -76,9 +76,16 @@ def __init__(
         browser_api_client: "BrowserAPIClient",
     ):
         self.playwright_page = page
-        self.screenshot_manager = ScreenshotManager()
+        self.screenshot_manager = ScreenshotManager(page)
         self.dendrite_browser = dendrite_browser
         self._browser_api_client = browser_api_client
+        self._last_frame_navigated_timestamp = time.time()
+
+        self.playwright_page.on("framenavigated", self._on_frame_navigated)
+
+    def _on_frame_navigated(self, frame):
+        if frame is self.playwright_page.main_frame:
+            self._last_frame_navigated_timestamp = time.time()
 
     @property
     def url(self):
@@ -246,15 +253,14 @@ async def get_page_information(self) -> PageInformation:
             PageInformation: An object containing the page's URL, raw HTML, and a screenshot in base64 format.
         """
 
-        base64 = await self.screenshot_manager.take_full_page_screenshot(
-            self.playwright_page
-        )
+        base64 = await self.screenshot_manager.take_full_page_screenshot()
         soup = await self._get_soup()
 
         return PageInformation(
             url=self.playwright_page.url,
             raw_html=str(soup),
             screenshot_base64=base64,
+            time_since_frame_navigated=self.get_time_since_last_frame_navigated(),
         )
 
     async def _generate_dendrite_ids(self):
@@ -393,3 +399,12 @@ async def _dump_html(self, path: str) -> None:
 
         with open(path, "w") as f:
             f.write(await self.playwright_page.content())
+
+    def get_time_since_last_frame_navigated(self) -> float:
+        """
+        Get the time elapsed since the last URL change.
+
+        Returns:
+            float: The number of seconds elapsed since the last URL change.
+        """
+        return time.time() - self._last_frame_navigated_timestamp

dendrite_sdk/async_api/_core/mixin/extract.py

@@ -1,5 +1,7 @@
+import asyncio
+import time
 from typing import Any, Optional, Type, overload
-from dendrite_sdk.async_api._api.dto.scrape_page_dto import ScrapePageDTO
+from dendrite_sdk.async_api._api.dto.extract_dto import ExtractDTO
 from dendrite_sdk.async_api._core._type_spec import (
     JsonSchema,
     PydanticModel,
@@ -25,6 +27,7 @@ async def extract(
         prompt: str,
         type_spec: Type[bool],
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> bool: ...
 
     @overload
@@ -33,6 +36,7 @@ async def extract(
         prompt: str,
         type_spec: Type[int],
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> int: ...
 
     @overload
@@ -41,6 +45,7 @@ async def extract(
         prompt: str,
         type_spec: Type[float],
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> float: ...
 
     @overload
@@ -49,6 +54,7 @@ async def extract(
         prompt: str,
         type_spec: Type[str],
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> str: ...
 
     @overload
@@ -57,6 +63,7 @@ async def extract(
         prompt: Optional[str],
         type_spec: Type[PydanticModel],
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> PydanticModel: ...
 
     @overload
@@ -65,6 +72,7 @@ async def extract(
         prompt: Optional[str],
         type_spec: JsonSchema,
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> JsonSchema: ...
 
     @overload
@@ -73,13 +81,15 @@ async def extract(
         prompt: str,
         type_spec: None = None,
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> Any: ...
 
     async def extract(
         self,
         prompt: Optional[str],
         type_spec: Optional[TypeSpec] = None,
         use_cache: bool = True,
+        timeout: int = 180,
     ) -> TypeSpec:
         """
         Extract data from a web page based on a prompt and optional type specification.
@@ -88,9 +98,13 @@ async def extract(
             prompt (Optional[str]): The prompt to guide the extraction.
             type_spec (Optional[TypeSpec], optional): The type specification for the extracted data.
             use_cache (bool, optional): Whether to use cached results. Defaults to True.
+            timeout (int, optional): The maximum time to wait for extraction in seconds. Defaults to 180 seconds, which is 3 minutes.
 
         Returns:
-            ScrapePageResponse: The extracted data wrapped in a ScrapePageResponse object.
+            ExtractResponse: The extracted data wrapped in a ExtractResponse object.
+
+        Raises:
+            TimeoutError: If the extraction process exceeds the specified timeout.
         """
         json_schema = None
         if type_spec:
@@ -99,22 +113,40 @@ async def extract(
         if prompt is None:
             prompt = ""
 
+        init_start_time = time.time()
         page = await self._get_page()
+
         page_information = await page.get_page_information()
-        scrape_dto = ScrapePageDTO(
+        extract_dto = ExtractDTO(
             page_information=page_information,
             api_config=self._get_dendrite_browser().api_config,
             prompt=prompt,
             return_data_json_schema=json_schema,
             use_screenshot=True,
             use_cache=use_cache,
         )
-        res = await self._get_browser_api_client().scrape_page(scrape_dto)
+
+        delay = 1
+        while True:
+            elapsed_time = time.time() - init_start_time
+            if elapsed_time > timeout:
+                raise TimeoutError(
+                    f"Extraction process exceeded the timeout of {timeout} seconds"
+                )
+
+            start_time = time.time()
+            res = await self._get_browser_api_client().extract(extract_dto)
+            request_time = time.time() - start_time
+
+            if res.status != "loading":
+                break
+
+            remaining_delay = max(0, delay - request_time)
+            await asyncio.sleep(remaining_delay)
+            delay = min(delay * 2, 20)  # Exponential back-off, capped at 60 seconds
 
         converted_res = res.return_data
         if type_spec is not None:
             converted_res = convert_to_type_spec(type_spec, res.return_data)
 
-        res = converted_res
-
-        return res
+        return converted_res

dendrite_sdk/async_api/_core/mixin/markdown.py

@@ -0,0 +1,17 @@
+from typing import Optional
+from dendrite_sdk.async_api._core.mixin.extract import ExtractionMixin
+from dendrite_sdk.async_api._core.protocol.page_protocol import DendritePageProtocol
+
+from markdownify import markdownify as md
+
+
+class MarkdownMixin(ExtractionMixin, DendritePageProtocol):
+    async def markdown(self, prompt: Optional[str] = None):
+        page = await self._get_page()
+        page_information = await page.get_page_information()
+        if prompt:
+            extract_prompt = f"Extract and return the html for this requested section of the website:\n\n{prompt}"
+            res = await self.extract(extract_prompt, str)
+            return md(res, heading_style="ATX")
+        else:
+            return md(page_information.raw_html)

dendrite_sdk/async_api/_core/mixin/screenshot.py

@@ -0,0 +1,20 @@
+from dendrite_sdk.async_api._core.protocol.page_protocol import DendritePageProtocol
+
+
+class ScreenshotMixin(DendritePageProtocol):
+
+    async def screenshot(self, full_page: bool = False) -> str:
+        """
+        Take a screenshot of the current page.
+
+        Args:
+            full_page (bool, optional): If True, captures the full page. If False, captures only the viewport. Defaults to False.
+
+        Returns:
+            str: A base64 encoded string of the screenshot in JPEG format.
+        """
+        page = await self._get_page()
+        if full_page:
+            return await page.screenshot_manager.take_full_page_screenshot()
+        else:
+            return await page.screenshot_manager.take_viewport_screenshot()

dendrite_sdk/async_api/_core/models/page_information.py

@@ -12,3 +12,4 @@ class PageInformation(BaseModel):
     url: str
     raw_html: str
     screenshot_base64: str
+    time_since_frame_navigated: float