Den 444 Improve cache speed (#85)

Author: coffecup25

dendrite/__init__.py

@@ -18,7 +18,7 @@
 
 fmt = "<green>{time: HH:mm:ss.SSS}</green> | <level>{level: <8}</level>- <level>{message}</level>"
 
-logger.add(sys.stderr, level="WARNING", format=fmt)
+logger.add(sys.stderr, level="DEBUG", format=fmt)
 
 
 __all__ = [

dendrite/async_api/_core/_utils.py

@@ -1,10 +1,12 @@
-from typing import Union, List, TYPE_CHECKING
-from playwright.async_api import FrameLocator, ElementHandle, Error
+from typing import Optional, Union, List, TYPE_CHECKING
+from playwright.async_api import FrameLocator, ElementHandle, Error, Frame
 from bs4 import BeautifulSoup
 from loguru import logger
 
+from dendrite.async_api._api.response.get_element_response import GetElementResponse
 from dendrite.async_api._core._type_spec import PlaywrightPage
 from dendrite.async_api._core.dendrite_element import AsyncElement
+from dendrite.async_api._core.models.response import AsyncElementsResponse
 
 if TYPE_CHECKING:
     from dendrite.async_api._core.dendrite_page import AsyncPage
@@ -18,51 +20,40 @@
 async def expand_iframes(
     page: PlaywrightPage,
     page_soup: BeautifulSoup,
-    iframe_path: str = "",
-    frame: Union[ElementHandle, None] = None,
 ):
-
-    if frame is None:
-        iframes = await page.query_selector_all("iframe")
-    else:
-        content_frame = await frame.content_frame()
-        if not content_frame:
-            return
-        iframes = await content_frame.query_selector_all("iframe")
-    for iframe in iframes:
-        # TODO: kolla om iframe inte har doc eller body, skippa då
-        iframe_id = await iframe.get_attribute("d-id")
+    async def get_iframe_path(frame: Frame):
+        path_parts = []
+        current_frame = frame
+        while current_frame.parent_frame is not None:
+            iframe_element = await current_frame.frame_element()
+            iframe_id = await iframe_element.get_attribute("d-id")
+            if iframe_id is None:
+                # If any iframe_id in the path is None, we cannot build the path
+                return None
+            path_parts.insert(0, iframe_id)
+            current_frame = current_frame.parent_frame
+        return "|".join(path_parts)
+
+    for frame in page.frames:
+        if frame.parent_frame is None:
+            continue  # Skip the main frame
+        iframe_element = await frame.frame_element()
+        iframe_id = await iframe_element.get_attribute("d-id")
         if iframe_id is None:
             continue
-
-        new_iframe_path = ""
-        if iframe_path:
-            new_iframe_path = f"{iframe_path}|"
-        new_iframe_path = f"{new_iframe_path}{iframe_id}"
-
+        iframe_path = await get_iframe_path(frame)
+        if iframe_path is None:
+            continue
         try:
-            content_frame = await iframe.content_frame()
-
-            if content_frame is None:
-                continue
-
-            await content_frame.evaluate(
-                GENERATE_DENDRITE_IDS_IFRAME_SCRIPT, {"frame_path": new_iframe_path}
+            await frame.evaluate(
+                GENERATE_DENDRITE_IDS_IFRAME_SCRIPT, {"frame_path": iframe_path}
             )
-
-            frame_content = await content_frame.content()
-
-            frame_tree = BeautifulSoup(frame_content, "html.parser")
+            frame_content = await frame.content()
+            frame_tree = BeautifulSoup(frame_content, "lxml")
             mild_strip_in_place(frame_tree)
             merge_iframe_to_page(iframe_id, page_soup, frame_tree)
-            await expand_iframes(
-                page,
-                page_soup,
-                new_iframe_path,
-                iframe,
-            )
         except Error as e:
-            logger.debug(f"Error getting content frame for iframe {iframe_id}: {e}")
+            logger.debug(f"Error processing frame {iframe_id}: {e}")
             continue
 
 
@@ -79,11 +70,54 @@ def merge_iframe_to_page(
     iframe_element.replace_with(iframe)
 
 
-def get_frame_context(
-    page: PlaywrightPage, iframe_path: str
-) -> Union[FrameLocator, PlaywrightPage]:
-    iframe_path_list = iframe_path.split("|")
-    frame_context = 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_soup(
+    selector: str, soup: BeautifulSoup, page: "AsyncPage"
+) -> List[AsyncElement]:
+    dendrite_elements: List[AsyncElement] = []
+
+    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(
+            AsyncElement(d_id, locator, page.dendrite_browser, page._browser_api_client)
+        )
+
+    return dendrite_elements
+
+
+async def get_elements_from_selectors_soup(
+    page: "AsyncPage",
+    soup: BeautifulSoup,
+    res: GetElementResponse,
+    only_one: bool,
+) -> Union[Optional[AsyncElement], List[AsyncElement], AsyncElementsResponse]:
+    if isinstance(res.selectors, dict):
+        result = {}
+        for key, selectors in res.selectors.items():
+            for selector in selectors:
+                dendrite_elements = await _get_all_elements_from_selector_soup(
+                    selector, soup, page
+                )
+                if len(dendrite_elements) > 0:
+                    result[key] = dendrite_elements[0]
+                    break
+        return AsyncElementsResponse(result)
+    elif isinstance(res.selectors, list):
+        for selector in reversed(res.selectors):
+            dendrite_elements = await _get_all_elements_from_selector_soup(
+                selector, soup, page
+            )
+
+            if len(dendrite_elements) > 0:
+                return dendrite_elements[0] if only_one else dendrite_elements
+
+    return None

dendrite/async_api/_core/dendrite_page.py

@@ -357,8 +357,18 @@ async def _get_soup(self) -> BeautifulSoup:
         page_source = await self.playwright_page.content()
         soup = BeautifulSoup(page_source, "lxml")
         await self._expand_iframes(soup)
+        self._previous_soup = soup
         return soup
 
+    async def _get_previous_soup(self) -> BeautifulSoup:
+        """
+        Retrieves the page source generated by the latest _get_soup() call as a Beautiful soup object. If it hasn't been called yet, it will call it.
+        """
+
+        if self._previous_soup is None:
+            return await self._get_soup()
+        return self._previous_soup
+
     async def _expand_iframes(self, page_source: BeautifulSoup):
         """
         Expands iframes in the given page source to make their content accessible.

dendrite/async_api/_core/mixin/get_element.py

@@ -7,6 +7,7 @@
 from dendrite.async_api._api.dto.get_elements_dto import GetElementsDTO
 from dendrite.async_api._api.response.get_element_response import GetElementResponse
 from dendrite.async_api._api.dto.get_elements_dto import CheckSelectorCacheDTO
+from dendrite.async_api._core._utils import get_elements_from_selectors_soup
 from dendrite.async_api._core.dendrite_element import AsyncElement
 from dendrite.async_api._core.models.response import AsyncElementsResponse
 from dendrite.async_api._core.protocol.page_protocol import DendritePageProtocol
@@ -202,7 +203,10 @@ async def _get_element(
         start_time = time.time()
 
         # First, let's check if there is a cached selector
-        cache_available = await test_if_cache_available(self, prompt_or_elements)
+        page = await self._get_page()
+        cache_available = await test_if_cache_available(
+            self, prompt_or_elements, page.url
+        )
 
         # If we have cached elements, attempt to use them with an exponentation backoff
         if cache_available and use_cache == True:
@@ -244,16 +248,14 @@ async def _get_element(
 
 
 async def test_if_cache_available(
-    obj: DendritePageProtocol,
-    prompt_or_elements: Union[str, Dict[str, str]],
+    obj: DendritePageProtocol, prompt_or_elements: Union[str, Dict[str, str]], url: str
 ) -> bool:
-    page = await obj._get_page()
-    page_information = await page.get_page_information(include_screenshot=False)
     dto = CheckSelectorCacheDTO(
-        url=page_information.url,
+        url=url,
         prompt=prompt_or_elements,
     )
     cache_available = await obj._get_browser_api_client().check_selector_cache(dto)
+
     return cache_available.exists
 
 
@@ -297,7 +299,9 @@ async def attempt_with_backoff(
             return None
 
         if res.status == "success":
-            response = await get_elements_from_selectors(obj, res, only_one)
+            response = await get_elements_from_selectors_soup(
+                page, await page._get_previous_soup(), res, only_one
+            )
             if response:
                 return response
 

dendrite/sync_api/_core/_utils.py

@@ -1,51 +1,52 @@
-from typing import Union, List, TYPE_CHECKING
-from playwright.sync_api import FrameLocator, ElementHandle, Error
+from typing import Optional, Union, List, TYPE_CHECKING
+from playwright.sync_api import FrameLocator, ElementHandle, Error, Frame
 from bs4 import BeautifulSoup
 from loguru import logger
+from dendrite.sync_api._api.response.get_element_response import GetElementResponse
 from dendrite.sync_api._core._type_spec import PlaywrightPage
 from dendrite.sync_api._core.dendrite_element import Element
+from dendrite.sync_api._core.models.response import ElementsResponse
 
 if TYPE_CHECKING:
     from dendrite.sync_api._core.dendrite_page import Page
 from dendrite.sync_api._core._js import GENERATE_DENDRITE_IDS_IFRAME_SCRIPT
 from dendrite.sync_api._dom.util.mild_strip import mild_strip_in_place
 
 
-def expand_iframes(
-    page: PlaywrightPage,
-    page_soup: BeautifulSoup,
-    iframe_path: str = "",
-    frame: Union[ElementHandle, None] = None,
-):
-    if frame is None:
-        iframes = page.query_selector_all("iframe")
-    else:
-        content_frame = frame.content_frame()
-        if not content_frame:
-            return
-        iframes = content_frame.query_selector_all("iframe")
-    for iframe in iframes:
-        iframe_id = iframe.get_attribute("d-id")
+def expand_iframes(page: PlaywrightPage, page_soup: BeautifulSoup):
+
+    def get_iframe_path(frame: Frame):
+        path_parts = []
+        current_frame = frame
+        while current_frame.parent_frame is not None:
+            iframe_element = current_frame.frame_element()
+            iframe_id = iframe_element.get_attribute("d-id")
+            if iframe_id is None:
+                return None
+            path_parts.insert(0, iframe_id)
+            current_frame = current_frame.parent_frame
+        return "|".join(path_parts)
+
+    for frame in page.frames:
+        if frame.parent_frame is None:
+            continue
+        iframe_element = frame.frame_element()
+        iframe_id = iframe_element.get_attribute("d-id")
         if iframe_id is None:
             continue
-        new_iframe_path = ""
-        if iframe_path:
-            new_iframe_path = f"{iframe_path}|"
-        new_iframe_path = f"{new_iframe_path}{iframe_id}"
+        iframe_path = get_iframe_path(frame)
+        if iframe_path is None:
+            continue
         try:
-            content_frame = iframe.content_frame()
-            if content_frame is None:
-                continue
-            content_frame.evaluate(
-                GENERATE_DENDRITE_IDS_IFRAME_SCRIPT, {"frame_path": new_iframe_path}
+            frame.evaluate(
+                GENERATE_DENDRITE_IDS_IFRAME_SCRIPT, {"frame_path": iframe_path}
             )
-            frame_content = content_frame.content()
-            frame_tree = BeautifulSoup(frame_content, "html.parser")
+            frame_content = frame.content()
+            frame_tree = BeautifulSoup(frame_content, "lxml")
             mild_strip_in_place(frame_tree)
             merge_iframe_to_page(iframe_id, page_soup, frame_tree)
-            expand_iframes(page, page_soup, new_iframe_path, iframe)
         except Error as e:
-            logger.debug(f"Error getting content frame for iframe {iframe_id}: {e}")
+            logger.debug(f"Error processing frame {iframe_id}: {e}")
             continue
 
 
@@ -57,11 +58,44 @@ def merge_iframe_to_page(iframe_id: str, page: BeautifulSoup, iframe: BeautifulS
     iframe_element.replace_with(iframe)
 
 
-def get_frame_context(
-    page: PlaywrightPage, iframe_path: str
-) -> Union[FrameLocator, PlaywrightPage]:
-    iframe_path_list = iframe_path.split("|")
-    frame_context = page
-    for iframe_id in iframe_path_list:
-        frame_context = frame_context.frame_locator(f"[tf623_id='{iframe_id}']")
-    return frame_context
+def _get_all_elements_from_selector_soup(
+    selector: str, soup: BeautifulSoup, page: "Page"
+) -> List[Element]:
+    dendrite_elements: List[Element] = []
+    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(
+            Element(d_id, locator, page.dendrite_browser, page._browser_api_client)
+        )
+    return dendrite_elements
+
+
+def get_elements_from_selectors_soup(
+    page: "Page", soup: BeautifulSoup, res: GetElementResponse, only_one: bool
+) -> Union[Optional[Element], List[Element], ElementsResponse]:
+    if isinstance(res.selectors, dict):
+        result = {}
+        for key, selectors in res.selectors.items():
+            for selector in selectors:
+                dendrite_elements = _get_all_elements_from_selector_soup(
+                    selector, soup, page
+                )
+                if len(dendrite_elements) > 0:
+                    result[key] = dendrite_elements[0]
+                    break
+        return ElementsResponse(result)
+    elif isinstance(res.selectors, list):
+        for selector in reversed(res.selectors):
+            dendrite_elements = _get_all_elements_from_selector_soup(
+                selector, soup, page
+            )
+            if len(dendrite_elements) > 0:
+                return dendrite_elements[0] if only_one else dendrite_elements
+    return None

dendrite/sync_api/_core/dendrite_page.py

@@ -299,8 +299,17 @@ def _get_soup(self) -> BeautifulSoup:
         page_source = self.playwright_page.content()
         soup = BeautifulSoup(page_source, "lxml")
         self._expand_iframes(soup)
+        self._previous_soup = soup
         return soup
 
+    def _get_previous_soup(self) -> BeautifulSoup:
+        """
+        Retrieves the page source generated by the latest _get_soup() call as a Beautiful soup object. If it hasn't been called yet, it will call it.
+        """
+        if self._previous_soup is None:
+            return self._get_soup()
+        return self._previous_soup
+
     def _expand_iframes(self, page_source: BeautifulSoup):
         """
         Expands iframes in the given page source to make their content accessible.