lumivor 0.1.7__py3-none-any.whl

lumivor/agent/tests.py

@@ -0,0 +1,204 @@
+import pytest
+
+from lumivor.agent.views import (
+    ActionResult,
+    AgentBrain,
+    AgentHistory,
+    AgentHistoryList,
+    AgentOutput,
+)
+from lumivor.browser.views import BrowserState, BrowserStateHistory, TabInfo
+from lumivor.controller.registry.service import Registry
+from lumivor.controller.views import ClickElementAction, DoneAction, ExtractPageContentAction
+from lumivor.dom.views import DOMElementNode
+
+
+@pytest.fixture
+def sample_browser_state():
+    return BrowserState(
+        url='https://example.com',
+        title='Example Page',
+        tabs=[TabInfo(url='https://example.com',
+                      title='Example Page', page_id=1)],
+        screenshot='screenshot1.png',
+        element_tree=DOMElementNode(
+            tag_name='root',
+            is_visible=True,
+            parent=None,
+            xpath='',
+            attributes={},
+            children=[],
+        ),
+        selector_map={},
+    )
+
+
+@pytest.fixture
+def action_registry():
+    registry = Registry()
+
+    # Register the actions we need for testing
+    @registry.action(
+        description='Click an element', requires_browser=True, param_model=ClickElementAction
+    )
+    def click_element(params: ClickElementAction, browser=None):
+        pass
+
+    @registry.action(
+        description='Extract page content',
+        requires_browser=True,
+        param_model=ExtractPageContentAction,
+    )
+    def extract_page_content(params: ExtractPageContentAction, browser=None):
+        pass
+
+    @registry.action(description='Mark task as done', param_model=DoneAction)
+    def done(params: DoneAction):
+        pass
+
+    # Create the dynamic ActionModel with all registered actions
+    return registry.create_action_model()
+
+
+@pytest.fixture
+def sample_history(action_registry):
+    # Create actions with nested params structure
+    click_action = action_registry(click_element={'index': 1})
+
+    extract_action = action_registry(extract_page_content={'value': 'text'})
+
+    done_action = action_registry(done={'text': 'Task completed'})
+
+    histories = [
+        AgentHistory(
+            model_output=AgentOutput(
+                current_state=AgentBrain(
+                    evaluation_previous_goal='None', memory='Started task', next_goal='Click button'
+                ),
+                action=[click_action],
+            ),
+            result=[ActionResult(is_done=False)],
+            state=BrowserStateHistory(
+                url='https://example.com',
+                title='Page 1',
+                tabs=[TabInfo(url='https://example.com',
+                              title='Page 1', page_id=1)],
+                screenshot='screenshot1.png',
+                interacted_element=[],
+            ),
+        ),
+        AgentHistory(
+            model_output=AgentOutput(
+                current_state=AgentBrain(
+                    evaluation_previous_goal='Clicked button',
+                    memory='Button clicked',
+                    next_goal='Extract content',
+                ),
+                action=[extract_action],
+            ),
+            result=[
+                ActionResult(
+                    is_done=False,
+                    extracted_content='Extracted text',
+                    error='Failed to extract completely',
+                )
+            ],
+            state=BrowserStateHistory(
+                url='https://example.com/page2',
+                title='Page 2',
+                tabs=[TabInfo(url='https://example.com/page2',
+                              title='Page 2', page_id=2)],
+                screenshot='screenshot2.png',
+                interacted_element=[],
+            ),
+        ),
+        AgentHistory(
+            model_output=AgentOutput(
+                current_state=AgentBrain(
+                    evaluation_previous_goal='Extracted content',
+                    memory='Content extracted',
+                    next_goal='Finish task',
+                ),
+                action=[done_action],
+            ),
+            result=[ActionResult(
+                is_done=True, extracted_content='Task completed', error=None)],
+            state=BrowserStateHistory(
+                url='https://example.com/page2',
+                title='Page 2',
+                tabs=[TabInfo(url='https://example.com/page2',
+                              title='Page 2', page_id=2)],
+                screenshot='screenshot3.png',
+                interacted_element=[],
+            ),
+        ),
+    ]
+    return AgentHistoryList(history=histories)
+
+
+def test_last_model_output(sample_history: AgentHistoryList):
+    last_output = sample_history.last_action()
+    print(last_output)
+    assert last_output == {'done': {'text': 'Task completed'}}
+
+
+def test_get_errors(sample_history: AgentHistoryList):
+    errors = sample_history.errors()
+    assert len(errors) == 1
+    assert errors[0] == 'Failed to extract completely'
+
+
+def test_final_result(sample_history: AgentHistoryList):
+    assert sample_history.final_result() == 'Task completed'
+
+
+def test_is_done(sample_history: AgentHistoryList):
+    assert sample_history.is_done() == True
+
+
+def test_urls(sample_history: AgentHistoryList):
+    urls = sample_history.urls()
+    assert 'https://example.com' in urls
+    assert 'https://example.com/page2' in urls
+
+
+def test_all_screenshots(sample_history: AgentHistoryList):
+    screenshots = sample_history.screenshots()
+    assert len(screenshots) == 3
+    assert screenshots == ['screenshot1.png',
+                           'screenshot2.png', 'screenshot3.png']
+
+
+def test_all_model_outputs(sample_history: AgentHistoryList):
+    outputs = sample_history.model_actions()
+    assert len(outputs) == 3
+    assert outputs[0] == {'click_element': {
+        'index': 1, 'xpath': '//button[1]'}}
+    assert outputs[1] == {'extract_page_content': {'value': 'text'}}
+    assert outputs[2] == {'done': {'text': 'Task completed'}}
+
+
+def test_all_model_outputs_filtered(sample_history: AgentHistoryList):
+    filtered = sample_history.model_actions_filtered(include=['click_element'])
+    assert len(filtered) == 1
+    assert filtered[0]['click_element']['index'] == 1
+
+
+def test_empty_history():
+    empty_history = AgentHistoryList(history=[])
+    assert empty_history.last_action() is None
+    assert empty_history.final_result() is None
+    assert empty_history.is_done() == False
+    assert len(empty_history.urls()) == 0
+
+
+# Add a test to verify action creation
+def test_action_creation(action_registry):
+    click_action = action_registry(click_element={'index': 1})
+
+    assert click_action.model_dump(exclude_none=True) == {
+        'click_element': {'index': 1}}
+
+
+# run this with:
+# pytest lumivor/agent/tests.py

lumivor/agent/views.py

@@ -0,0 +1,272 @@
+from __future__ import annotations
+
+import json
+import traceback
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Optional, Type
+
+from openai import RateLimitError
+from pydantic import BaseModel, ConfigDict, Field, ValidationError, create_model
+
+from lumivor.browser.views import BrowserStateHistory
+from lumivor.controller.registry.views import ActionModel
+from lumivor.dom.history_tree_processor.service import (
+    DOMElementNode,
+    DOMHistoryElement,
+    HistoryTreeProcessor,
+)
+from lumivor.dom.views import SelectorMap
+
+
+@dataclass
+class AgentStepInfo:
+    step_number: int
+    max_steps: int
+
+
+class ActionResult(BaseModel):
+    """Result of executing an action"""
+
+    is_done: Optional[bool] = False
+    extracted_content: Optional[str] = None
+    error: Optional[str] = None
+    # whether to include in past messages as context or not
+    include_in_memory: bool = False
+
+
+class AgentBrain(BaseModel):
+    """Current state of the agent"""
+
+    evaluation_previous_goal: str
+    memory: str
+    next_goal: str
+
+
+class AgentOutput(BaseModel):
+    """Output model for agent
+
+    @dev note: this model is extended with custom actions in AgentService. You can also use some fields that are not in this model as provided by the linter, as long as they are registered in the DynamicActions model.
+    """
+
+    model_config = ConfigDict(arbitrary_types_allowed=True)
+
+    current_state: AgentBrain
+    action: list[ActionModel]
+
+    @staticmethod
+    def type_with_custom_actions(custom_actions: Type[ActionModel]) -> Type['AgentOutput']:
+        """Extend actions with custom actions"""
+        return create_model(
+            'AgentOutput',
+            __base__=AgentOutput,
+            # Properly annotated field with no default
+            action=(list[custom_actions], Field(...)),
+            __module__=AgentOutput.__module__,
+        )
+
+
+class AgentHistory(BaseModel):
+    """History item for agent actions"""
+
+    model_output: AgentOutput | None
+    result: list[ActionResult]
+    state: BrowserStateHistory
+
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True, protected_namespaces=())
+
+    @staticmethod
+    def get_interacted_element(
+            model_output: AgentOutput, selector_map: SelectorMap
+    ) -> list[DOMHistoryElement | None]:
+        elements = []
+        for action in model_output.action:
+            index = action.get_index()
+            if index and index in selector_map:
+                el: DOMElementNode = selector_map[index]
+                elements.append(
+                    HistoryTreeProcessor.convert_dom_element_to_history_element(el))
+            else:
+                elements.append(None)
+        return elements
+
+    def model_dump(self, **kwargs) -> Dict[str, Any]:
+        """Custom serialization handling circular references"""
+
+        # Handle action serialization
+        model_output_dump = None
+        if self.model_output:
+            action_dump = [
+                action.model_dump(exclude_none=True) for action in self.model_output.action
+            ]
+            model_output_dump = {
+                'current_state': self.model_output.current_state.model_dump(),
+                'action': action_dump,  # This preserves the actual action data
+            }
+
+        return {
+            'model_output': model_output_dump,
+            'result': [r.model_dump(exclude_none=True) for r in self.result],
+            'state': self.state.to_dict(),
+        }
+
+
+class AgentHistoryList(BaseModel):
+    """List of agent history items"""
+
+    history: list[AgentHistory]
+
+    def __str__(self) -> str:
+        """Representation of the AgentHistoryList object"""
+        return f'AgentHistoryList(all_results={self.action_results()}, all_model_outputs={self.model_actions()})'
+
+    def __repr__(self) -> str:
+        """Representation of the AgentHistoryList object"""
+        return self.__str__()
+
+    def save_to_file(self, filepath: str | Path) -> None:
+        """Save history to JSON file with proper serialization"""
+        try:
+            Path(filepath).parent.mkdir(parents=True, exist_ok=True)
+            data = self.model_dump()
+            with open(filepath, 'w', encoding='utf-8') as f:
+                json.dump(data, f, indent=2)
+        except Exception as e:
+            raise e
+
+    def model_dump(self, **kwargs) -> Dict[str, Any]:
+        """Custom serialization that properly uses AgentHistory's model_dump"""
+        return {
+            'history': [h.model_dump(**kwargs) for h in self.history],
+        }
+
+    @classmethod
+    def load_from_file(
+            cls, filepath: str | Path, output_model: Type[AgentOutput]
+    ) -> 'AgentHistoryList':
+        """Load history from JSON file"""
+        with open(filepath, 'r', encoding='utf-8') as f:
+            data = json.load(f)
+        # loop through history and validate output_model actions to enrich with custom actions
+        for h in data['history']:
+            if h['model_output']:
+                if isinstance(h['model_output'], dict):
+                    h['model_output'] = output_model.model_validate(
+                        h['model_output'])
+                else:
+                    h['model_output'] = None
+            if 'interacted_element' not in h['state']:
+                h['state']['interacted_element'] = None
+        history = cls.model_validate(data)
+        return history
+
+    def last_action(self) -> None | dict:
+        """Last action in history"""
+        if self.history and self.history[-1].model_output:
+            return self.history[-1].model_output.action[-1].model_dump(exclude_none=True)
+        return None
+
+    def errors(self) -> list[str]:
+        """Get all errors from history"""
+        errors = []
+        for h in self.history:
+            errors.extend([r.error for r in h.result if r.error])
+        return errors
+
+    def final_result(self) -> None | str:
+        """Final result from history"""
+        if self.history and self.history[-1].result[-1].extracted_content:
+            return self.history[-1].result[-1].extracted_content
+        return None
+
+    def is_done(self) -> bool:
+        """Check if the agent is done"""
+        if (
+                self.history
+                and len(self.history[-1].result) > 0
+                and self.history[-1].result[-1].is_done
+        ):
+            return self.history[-1].result[-1].is_done
+        return False
+
+    def has_errors(self) -> bool:
+        """Check if the agent has any errors"""
+        return len(self.errors()) > 0
+
+    def urls(self) -> list[str]:
+        """Get all unique URLs from history"""
+        return [h.state.url for h in self.history if h.state.url]
+
+    def screenshots(self) -> list[str]:
+        """Get all screenshots from history"""
+        return [h.state.screenshot for h in self.history if h.state.screenshot]
+
+    def action_names(self) -> list[str]:
+        """Get all action names from history"""
+        return [list(action.keys())[0] for action in self.model_actions()]
+
+    def model_thoughts(self) -> list[AgentBrain]:
+        """Get all thoughts from history"""
+        return [h.model_output.current_state for h in self.history if h.model_output]
+
+    def model_outputs(self) -> list[AgentOutput]:
+        """Get all model outputs from history"""
+        return [h.model_output for h in self.history if h.model_output]
+
+    # get all actions with params
+    def model_actions(self) -> list[dict]:
+        """Get all actions from history"""
+        outputs = []
+
+        for h in self.history:
+            if h.model_output:
+                for action in h.model_output.action:
+                    output = action.model_dump(exclude_none=True)
+                    outputs.append(output)
+        return outputs
+
+    def action_results(self) -> list[ActionResult]:
+        """Get all results from history"""
+        results = []
+        for h in self.history:
+            results.extend([r for r in h.result if r])
+        return results
+
+    def extracted_content(self) -> list[str]:
+        """Get all extracted content from history"""
+        content = []
+        for h in self.history:
+            content.extend(
+                [r.extracted_content for r in h.result if r.extracted_content])
+        return content
+
+    def model_actions_filtered(self, include: list[str] = []) -> list[dict]:
+        """Get all model actions from history as JSON"""
+        outputs = self.model_actions()
+        result = []
+        for o in outputs:
+            for i in include:
+                if i == list(o.keys())[0]:
+                    result.append(o)
+        return result
+
+
+class AgentError:
+    """Container for agent error handling"""
+
+    VALIDATION_ERROR = 'Invalid model output format. Please follow the correct schema.'
+    RATE_LIMIT_ERROR = 'Rate limit reached. Waiting before retry.'
+    NO_VALID_ACTION = 'No valid action found'
+
+    @staticmethod
+    def format_error(error: Exception, include_trace: bool = False) -> str:
+        """Format error message based on error type and optionally include trace"""
+        message = ''
+        if isinstance(error, ValidationError):
+            return f'{AgentError.VALIDATION_ERROR}\nDetails: {str(error)}'
+        if isinstance(error, RateLimitError):
+            return AgentError.RATE_LIMIT_ERROR
+        if include_trace:
+            return f'{str(error)}\nStacktrace:\n{traceback.format_exc()}'
+        return f'{str(error)}'

lumivor/browser/browser.py

@@ -0,0 +1,208 @@
+"""
+Playwright browser on steroids.
+"""
+
+import asyncio
+import logging
+from dataclasses import dataclass, field
+
+from playwright._impl._api_structures import ProxySettings
+from playwright.async_api import Browser as PlaywrightBrowser
+from playwright.async_api import (
+    Playwright,
+    async_playwright,
+)
+
+from lumivor.browser.context import BrowserContext, BrowserContextConfig
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class BrowserConfig:
+    """
+    Configuration for the Browser.
+
+    Default values:
+            headless: True
+                    Whether to run browser in headless mode
+
+            disable_security: False
+                    Disable browser security features
+
+            extra_chromium_args: []
+                    Extra arguments to pass to the browser
+
+            wss_url: None
+                    Connect to a browser instance via WebSocket
+
+            chrome_instance_path: None
+                    Path to a Chrome instance to use to connect to your normal browser
+                    e.g. '/Applications/Google\ Chrome.app/Contents/MacOS/Google\ Chrome'
+    """
+
+    headless: bool = False
+    disable_security: bool = True
+    extra_chromium_args: list[str] = field(default_factory=list)
+    chrome_instance_path: str | None = None
+    wss_url: str | None = None
+
+    proxy: ProxySettings | None = field(default=None)
+    new_context_config: BrowserContextConfig = field(
+        default_factory=BrowserContextConfig)
+
+
+# @singleton: TODO - think about id singleton makes sense here
+# @dev By default this is a singleton, but you can create multiple instances if you need to.
+class Browser:
+    """
+    Playwright browser on steroids.
+
+    This is persistant browser factory that can spawn multiple browser contexts.
+    It is recommended to use only one instance of Browser per your application (RAM usage will grow otherwise).
+    """
+
+    def __init__(
+            self,
+            config: BrowserConfig = BrowserConfig(),
+    ):
+        logger.debug('Initializing new browser')
+        self.config = config
+        self.playwright: Playwright | None = None
+        self.playwright_browser: PlaywrightBrowser | None = None
+
+    async def new_context(
+            self, config: BrowserContextConfig = BrowserContextConfig()
+    ) -> BrowserContext:
+        """Create a browser context"""
+        return BrowserContext(config=config, browser=self)
+
+    async def get_playwright_browser(self) -> PlaywrightBrowser:
+        """Get a browser context"""
+        if self.playwright_browser is None:
+            return await self._init()
+
+        return self.playwright_browser
+
+    async def _init(self):
+        """Initialize the browser session"""
+        playwright = await async_playwright().start()
+        browser = await self._setup_browser(playwright)
+
+        self.playwright = playwright
+        self.playwright_browser = browser
+
+        return self.playwright_browser
+
+    async def _setup_browser(self, playwright: Playwright) -> PlaywrightBrowser:
+        """Sets up and returns a Playwright Browser instance with anti-detection measures."""
+        if self.config.wss_url:
+            browser = await playwright.chromium.connect(self.config.wss_url)
+            return browser
+        elif self.config.chrome_instance_path:
+            import subprocess
+
+            import requests
+
+            try:
+                # Check if browser is already running
+                response = requests.get(
+                    'http://localhost:9222/json/version', timeout=2)
+                if response.status_code == 200:
+                    logger.info('Reusing existing Chrome instance')
+                    browser = await playwright.chromium.connect_over_cdp(
+                        endpoint_url='http://localhost:9222',
+                        timeout=20000,  # 20 second timeout for connection
+                    )
+                    return browser
+            except requests.ConnectionError:
+                logger.debug(
+                    'No existing Chrome instance found, starting a new one')
+
+            # Start a new Chrome instance
+            subprocess.Popen(
+                [
+                    self.config.chrome_instance_path,
+                    '--remote-debugging-port=9222',
+                ],
+                stdout=subprocess.DEVNULL,
+                stderr=subprocess.DEVNULL,
+            )
+
+            # Attempt to connect again after starting a new instance
+            try:
+                browser = await playwright.chromium.connect_over_cdp(
+                    endpoint_url='http://localhost:9222',
+                    timeout=20000,  # 20 second timeout for connection
+                )
+                return browser
+            except Exception as e:
+                logger.error(
+                    f'Failed to start a new Chrome instance.: {str(e)}')
+                raise RuntimeError(
+                    ' To start chrome in Debug mode, you need to close all existing Chrome instances and try again otherwise we can not connect to the instance.'
+                )
+
+        else:
+            try:
+                disable_security_args = []
+                if self.config.disable_security:
+                    disable_security_args = [
+                        '--disable-web-security',
+                        '--disable-site-isolation-trials',
+                        '--disable-features=IsolateOrigins,site-per-process',
+                    ]
+
+                browser = await playwright.chromium.launch(
+                    headless=self.config.headless,
+                    args=[
+                        '--no-sandbox',
+                        '--disable-blink-features=AutomationControlled',
+                        '--disable-infobars',
+                        '--disable-background-timer-throttling',
+                        '--disable-popup-blocking',
+                        '--disable-backgrounding-occluded-windows',
+                        '--disable-renderer-backgrounding',
+                        '--disable-window-activation',
+                        '--disable-focus-on-load',
+                        '--no-first-run',
+                        '--no-default-browser-check',
+                        '--no-startup-window',
+                        '--window-position=0,0',
+                        # '--window-size=1280,1000',
+                    ]
+                    + disable_security_args
+                    + self.config.extra_chromium_args,
+                    proxy=self.config.proxy,
+                )
+
+                return browser
+            except Exception as e:
+                logger.error(
+                    f'Failed to initialize Playwright browser: {str(e)}')
+                raise
+
+    async def close(self):
+        """Close the browser instance"""
+        try:
+            if self.playwright_browser:
+                await self.playwright_browser.close()
+            if self.playwright:
+                await self.playwright.stop()
+        except Exception as e:
+            logger.debug(f'Failed to close browser properly: {e}')
+        finally:
+            self.playwright_browser = None
+            self.playwright = None
+
+    def __del__(self):
+        """Async cleanup when object is destroyed"""
+        try:
+            if self.playwright_browser or self.playwright:
+                loop = asyncio.get_running_loop()
+                if loop.is_running():
+                    loop.create_task(self.close())
+                else:
+                    asyncio.run(self.close())
+        except Exception as e:
+            logger.debug(f'Failed to cleanup browser in destructor: {e}')