lumivor 0.1.7__py3-none-any.whl

lumivor/README.md

@@ -0,0 +1,51 @@
+# Codebase Structure
+
+> The code structure inspired by https://github.com/Netflix/dispatch.
+
+Very good structure on how to make a scalable codebase is also in [this repo](https://github.com/zhanymkanov/fastapi-best-practices).
+
+Just a brief document about how we should structure our backend codebase.
+
+## Code Structure
+
+```markdown
+src/
+/<service name>/
+models.py
+services.py
+prompts.py
+views.py
+utils.py
+routers.py
+
+    	/_<subservice name>/
+```
+
+### Service.py
+
+Always a single file, except if it becomes too long - more than ~500 lines, split it into \_subservices
+
+### Views.py
+
+Always split the views into two parts
+
+```python
+# All
+...
+
+# Requests
+...
+
+# Responses
+...
+```
+
+If too long → split into multiple files
+
+### Prompts.py
+
+Single file; if too long → split into multiple files (one prompt per file or so)
+
+### Routers.py
+
+Never split into more than one file

lumivor/__init__.py

@@ -0,0 +1,25 @@
+from lumivor.dom.service import DomService as DomService
+from lumivor.controller.service import Controller as Controller
+from lumivor.browser.browser import BrowserConfig as BrowserConfig
+from lumivor.browser.browser import Browser as Browser
+from lumivor.agent.views import AgentHistoryList as AgentHistoryList
+from lumivor.agent.views import ActionResult as ActionResult
+from lumivor.agent.views import ActionModel as ActionModel
+from lumivor.agent.service import Agent as Agent
+from lumivor.agent.prompts import SystemPrompt as SystemPrompt
+from lumivor.logging_config import setup_logging
+
+setup_logging()
+
+
+__all__ = [
+    'Agent',
+    'Browser',
+    'BrowserConfig',
+    'Controller',
+    'DomService',
+    'SystemPrompt',
+    'ActionResult',
+    'ActionModel',
+    'AgentHistoryList',
+]

lumivor/agent/message_manager/service.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import logging
+from datetime import datetime
+from typing import List, Optional, Type
+
+from langchain_anthropic import ChatAnthropic
+from langchain_core.language_models import BaseChatModel
+from langchain_core.messages import (
+    AIMessage,
+    BaseMessage,
+    HumanMessage,
+)
+from langchain_openai import ChatOpenAI
+
+from lumivor.agent.message_manager.views import MessageHistory, MessageMetadata
+from lumivor.agent.prompts import AgentMessagePrompt, SystemPrompt
+from lumivor.agent.views import ActionResult, AgentOutput, AgentStepInfo
+from lumivor.browser.views import BrowserState
+
+logger = logging.getLogger(__name__)
+
+
+class MessageManager:
+    def __init__(
+            self,
+            llm: BaseChatModel,
+            task: str,
+            action_descriptions: str,
+            system_prompt_class: Type[SystemPrompt],
+            max_input_tokens: int = 128000,
+            estimated_tokens_per_character: int = 3,
+            image_tokens: int = 800,
+            include_attributes: list[str] = [],
+            max_error_length: int = 400,
+            max_actions_per_step: int = 10,
+            tool_call_in_content: bool = True,
+    ):
+        self.llm = llm
+        self.system_prompt_class = system_prompt_class
+        self.max_input_tokens = max_input_tokens
+        self.history = MessageHistory()
+        self.task = task
+        self.action_descriptions = action_descriptions
+        self.ESTIMATED_TOKENS_PER_CHARACTER = estimated_tokens_per_character
+        self.IMG_TOKENS = image_tokens
+        self.include_attributes = include_attributes
+        self.max_error_length = max_error_length
+
+        system_message = self.system_prompt_class(
+            self.action_descriptions,
+            current_date=datetime.now(),
+            max_actions_per_step=max_actions_per_step,
+        ).get_system_message()
+
+        self._add_message_with_tokens(system_message)
+        self.system_prompt = system_message
+        self.tool_call_in_content = tool_call_in_content
+        tool_calls = [
+            {
+                'name': 'AgentOutput',
+                'args': {
+                    'current_state': {
+                        'evaluation_previous_goal': 'Unknown - No previous actions to evaluate.',
+                        'memory': '',
+                        'next_goal': 'Obtain task from user',
+                    },
+                    'action': [],
+                },
+                'id': '',
+                'type': 'tool_call',
+            }
+        ]
+        if self.tool_call_in_content:
+            # openai throws error if tool_calls are not responded -> move to content
+            example_tool_call = AIMessage(
+                content=f'{tool_calls}',
+                tool_calls=[],
+            )
+        else:
+            example_tool_call = AIMessage(
+                content=f'',
+                tool_calls=tool_calls,
+            )
+
+        self._add_message_with_tokens(example_tool_call)
+
+        task_message = HumanMessage(content=f'Your task is: {task}')
+        self._add_message_with_tokens(task_message)
+
+    def add_state_message(
+            self,
+            state: BrowserState,
+            result: Optional[List[ActionResult]] = None,
+            step_info: Optional[AgentStepInfo] = None,
+    ) -> None:
+        """Add browser state as human message"""
+
+        # if keep in memory, add to directly to history and add state without result
+        if result:
+            for r in result:
+                if r.include_in_memory:
+                    if r.extracted_content:
+                        msg = HumanMessage(
+                            content='Action result: ' + str(r.extracted_content))
+                        self._add_message_with_tokens(msg)
+                    if r.error:
+                        msg = HumanMessage(
+                            content='Action error: ' +
+                            str(r.error)[-self.max_error_length:]
+                        )
+                        self._add_message_with_tokens(msg)
+                    result = None  # if result in history, we dont want to add it again
+
+        # otherwise add state message and result to next message (which will not stay in memory)
+        state_message = AgentMessagePrompt(
+            state,
+            result,
+            include_attributes=self.include_attributes,
+            max_error_length=self.max_error_length,
+            step_info=step_info,
+        ).get_user_message()
+        self._add_message_with_tokens(state_message)
+
+    def _remove_last_state_message(self) -> None:
+        """Remove last state message from history"""
+        if len(self.history.messages) > 2 and isinstance(
+                self.history.messages[-1].message, HumanMessage
+        ):
+            self.history.remove_message()
+
+    def add_model_output(self, model_output: AgentOutput) -> None:
+        """Add model output as AI message"""
+        tool_calls = [
+            {
+                'name': 'AgentOutput',
+                'args': model_output.model_dump(mode='json', exclude_unset=True),
+                        'id': '',
+                        'type': 'tool_call',
+            }
+        ]
+        if self.tool_call_in_content:
+            msg = AIMessage(
+                content=f'{tool_calls}',
+                tool_calls=[],
+            )
+        else:
+            msg = AIMessage(
+                content='',
+                tool_calls=tool_calls,
+            )
+
+        self._add_message_with_tokens(msg)
+
+    def get_messages(self) -> List[BaseMessage]:
+        """Get current message list, potentially trimmed to max tokens"""
+        self.cut_messages()
+        return [m.message for m in self.history.messages]
+
+    def cut_messages(self):
+        """Get current message list, potentially trimmed to max tokens"""
+        diff = self.history.total_tokens - self.max_input_tokens
+        if diff <= 0:
+            return None
+
+        msg = self.history.messages[-1]
+
+        # if list with image remove image
+        if isinstance(msg.message.content, list):
+            text = ''
+            for item in msg.message.content:
+                if 'image_url' in item:
+                    msg.message.content.remove(item)
+                    diff -= self.IMG_TOKENS
+                    msg.metadata.input_tokens -= self.IMG_TOKENS
+                    self.history.total_tokens -= self.IMG_TOKENS
+                    logger.debug(
+                        f'Removed image with {self.IMG_TOKENS} tokens - total tokens now: {
+                            self.history.total_tokens}/{self.max_input_tokens}'
+                    )
+                elif 'text' in item and isinstance(item, dict):
+                    text += item['text']
+            msg.message.content = text
+            self.history.messages[-1] = msg
+
+        if diff <= 0:
+            return None
+
+        # if still over, remove text from state message proportionally to the number of tokens needed with buffer
+        # Calculate the proportion of content to remove
+        proportion_to_remove = diff / msg.metadata.input_tokens
+        if proportion_to_remove > 0.99:
+            raise ValueError(
+                f'Max token limit reached - history is too long - reduce the system prompt or task less tasks or remove old messages. '
+                f'proportion_to_remove: {proportion_to_remove}'
+            )
+        logger.debug(
+            f'Removing {proportion_to_remove * 100:.2f}% of the last message  {proportion_to_remove *
+                                                                               msg.metadata.input_tokens:.2f} / {msg.metadata.input_tokens:.2f} tokens)'
+        )
+
+        content = msg.message.content
+        characters_to_remove = int(len(content) * proportion_to_remove)
+        content = content[:-characters_to_remove]
+
+        # remove tokens and old long message
+        self.history.remove_message(index=-1)
+
+        # new message with updated content
+        msg = HumanMessage(content=content)
+        self._add_message_with_tokens(msg)
+
+        last_msg = self.history.messages[-1]
+
+        logger.debug(
+            f'Added message with {last_msg.metadata.input_tokens} tokens - total tokens now: {
+                self.history.total_tokens}/{self.max_input_tokens} - total messages: {len(self.history.messages)}'
+        )
+
+    def _add_message_with_tokens(self, message: BaseMessage) -> None:
+        """Add message with token count metadata"""
+        token_count = self._count_tokens(message)
+        metadata = MessageMetadata(input_tokens=token_count)
+        self.history.add_message(message, metadata)
+
+    def _count_tokens(self, message: BaseMessage) -> int:
+        """Count tokens in a message using the model's tokenizer"""
+        tokens = 0
+        if isinstance(message.content, list):
+            for item in message.content:
+                if 'image_url' in item:
+                    tokens += self.IMG_TOKENS
+                elif isinstance(item, dict) and 'text' in item:
+                    tokens += self._count_text_tokens(item['text'])
+        else:
+            tokens += self._count_text_tokens(message.content)
+        return tokens
+
+    def _count_text_tokens(self, text: str) -> int:
+        """Count tokens in a text string"""
+        if isinstance(self.llm, (ChatOpenAI, ChatAnthropic)):
+            try:
+                tokens = self.llm.get_num_tokens(text)
+            except Exception:
+                tokens = (
+                    len(text) // self.ESTIMATED_TOKENS_PER_CHARACTER
+                )  # Rough estimate if no tokenizer available
+        else:
+            tokens = (
+                len(text) // self.ESTIMATED_TOKENS_PER_CHARACTER
+            )  # Rough estimate if no tokenizer available
+        return tokens

lumivor/agent/message_manager/tests.py

@@ -0,0 +1,246 @@
+import pytest
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_openai import AzureChatOpenAI, ChatOpenAI
+
+from lumivor.agent.message_manager.service import MessageManager
+from lumivor.agent.prompts import SystemPrompt
+from lumivor.agent.views import ActionResult
+from lumivor.browser.views import BrowserState, TabInfo
+from lumivor.dom.views import DOMElementNode, DOMTextNode
+
+
+@pytest.fixture(
+    params=[
+        ChatOpenAI(model='gpt-4o-mini'),
+        AzureChatOpenAI(model='gpt-4o', api_version='2024-02-15-preview'),
+        ChatAnthropic(
+            model_name='claude-3-5-sonnet-20240620', timeout=100, temperature=0.0, stop=None
+        ),
+    ],
+    ids=['gpt-4o-mini', 'gpt-4o', 'claude-3-5-sonnet'],
+)
+def message_manager(request: pytest.FixtureRequest):
+    llm = request.param
+    task = 'Test task'
+    action_descriptions = 'Test actions'
+    return MessageManager(
+        llm=llm,
+        task=task,
+        action_descriptions=action_descriptions,
+        system_prompt_class=SystemPrompt,
+        max_input_tokens=1000,
+        estimated_tokens_per_character=3,
+        image_tokens=800,
+    )
+
+
+def test_initial_messages(message_manager: MessageManager):
+    """Test that message manager initializes with system and task messages"""
+    messages = message_manager.get_messages()
+    assert len(messages) == 2
+    assert isinstance(messages[0], SystemMessage)
+    assert isinstance(messages[1], HumanMessage)
+    assert 'Test task' in messages[1].content
+
+
+def test_add_state_message(message_manager: MessageManager):
+    """Test adding browser state message"""
+    state = BrowserState(
+        url='https://test.com',
+        title='Test Page',
+        element_tree=DOMElementNode(
+            tag_name='div',
+            attributes={},
+            children=[],
+            is_visible=True,
+            parent=None,
+            xpath='//div',
+        ),
+        selector_map={},
+        tabs=[TabInfo(page_id=1, url='https://test.com', title='Test Page')],
+    )
+    message_manager.add_state_message(state)
+
+    messages = message_manager.get_messages()
+    assert len(messages) == 3
+    assert isinstance(messages[2], HumanMessage)
+    assert 'https://test.com' in messages[2].content
+
+
+def test_add_state_with_memory_result(message_manager: MessageManager):
+    """Test adding state with result that should be included in memory"""
+    state = BrowserState(
+        url='https://test.com',
+        title='Test Page',
+        element_tree=DOMElementNode(
+            tag_name='div',
+            attributes={},
+            children=[],
+            is_visible=True,
+            parent=None,
+            xpath='//div',
+        ),
+        selector_map={},
+        tabs=[TabInfo(page_id=1, url='https://test.com', title='Test Page')],
+    )
+    result = ActionResult(
+        extracted_content='Important content', include_in_memory=True)
+
+    message_manager.add_state_message(state, result)
+    messages = message_manager.get_messages()
+
+    # Should have system, task, extracted content, and state messages
+    assert len(messages) == 4
+    assert 'Important content' in messages[2].content
+    assert isinstance(messages[2], HumanMessage)
+    assert isinstance(messages[3], HumanMessage)
+    assert 'Important content' not in messages[3].content
+
+
+def test_add_state_with_non_memory_result(message_manager: MessageManager):
+    """Test adding state with result that should not be included in memory"""
+    state = BrowserState(
+        url='https://test.com',
+        title='Test Page',
+        element_tree=DOMElementNode(
+            tag_name='div',
+            attributes={},
+            children=[],
+            is_visible=True,
+            parent=None,
+            xpath='//div',
+        ),
+        selector_map={},
+        tabs=[TabInfo(page_id=1, url='https://test.com', title='Test Page')],
+    )
+    result = ActionResult(
+        extracted_content='Temporary content', include_in_memory=False)
+
+    message_manager.add_state_message(state, result)
+    messages = message_manager.get_messages()
+
+    # Should have system, task, and combined state+result message
+    assert len(messages) == 3
+    assert 'Temporary content' in messages[2].content
+    assert isinstance(messages[2], HumanMessage)
+
+
+@pytest.mark.skip('not sure how to fix this')
+@pytest.mark.parametrize('max_tokens', [100000, 10000, 5000])
+def test_token_overflow_handling_with_real_flow(message_manager: MessageManager, max_tokens):
+    """Test handling of token overflow in a realistic message flow"""
+    # Set more realistic token limit
+    message_manager.max_input_tokens = max_tokens
+
+    # Create a long sequence of interactions
+    for i in range(200):  # Simulate 40 steps of interaction
+        # Create state with varying content length
+        state = BrowserState(
+            url=f'https://test{i}.com',
+            title=f'Test Page {i}',
+            element_tree=DOMElementNode(
+                tag_name='div',
+                attributes={},
+                children=[
+                    DOMTextNode(
+                        # Increasing content length
+                        text=f'Content {j} ' * (10 + i),
+                        is_visible=True,
+                        parent=None,
+                    )
+                    for j in range(5)  # Multiple DOM items
+                ],
+                is_visible=True,
+                parent=None,
+                xpath='//div',
+            ),
+            selector_map={j: f'//div[{j}]' for j in range(5)},
+            tabs=[
+                TabInfo(page_id=1, url=f'https://test{i}.com', title=f'Test Page {i}')],
+        )
+
+        # Alternate between different types of results
+        result = None
+        if i % 2 == 0:  # Every other iteration
+            result = ActionResult(
+                extracted_content=f'Important content from step {i}' * 5,
+                include_in_memory=i % 4 == 0,  # Include in memory every 4th message
+            )
+
+        # Add state message
+        message_manager.add_state_message(state, result)
+
+        try:
+            messages = message_manager.get_messages()
+        except ValueError as e:
+            if 'Max token limit reached - history is too long' in str(e):
+                return  # If error occurs, end the test
+            else:
+                raise e
+
+        assert message_manager.history.total_tokens <= message_manager.max_input_tokens + 100
+
+        last_msg = messages[-1]
+        assert isinstance(last_msg, HumanMessage)
+
+        if i % 4 == 0:
+            assert isinstance(
+                message_manager.history.messages[-2].message, HumanMessage)
+        if i % 2 == 0 and not i % 4 == 0:
+            if isinstance(last_msg.content, list):
+                assert 'Current url: https://test' in last_msg.content[0]['text']
+            else:
+                assert 'Current url: https://test' in last_msg.content
+
+        # Add model output every time
+        from lumivor.agent.views import AgentBrain, AgentOutput
+        from lumivor.controller.registry.views import ActionModel
+
+        output = AgentOutput(
+            current_state=AgentBrain(
+                evaluation_previous_goal=f'Success in step {i}',
+                memory=f'Memory from step {i}',
+                next_goal=f'Goal for step {i+1}',
+            ),
+            action=[ActionModel()],
+        )
+        message_manager._remove_last_state_message()
+        message_manager.add_model_output(output)
+
+        # Get messages and verify after each addition
+        messages = [m.message for m in message_manager.history.messages]
+
+        # Verify token limit is respected
+
+        # Verify essential messages are preserved
+        # System prompt always first
+        assert isinstance(messages[0], SystemMessage)
+        assert isinstance(messages[1], HumanMessage)  # Task always second
+        assert 'Test task' in messages[1].content
+
+        # Verify structure of latest messages
+        # Last message should be model output
+        assert isinstance(messages[-1], AIMessage)
+        # Should contain current step info
+        assert f'step {i}' in messages[-1].content
+
+        # Log token usage for debugging
+        token_usage = message_manager.history.total_tokens
+        token_limit = message_manager.max_input_tokens
+        # print(f'Step {i}: Using {token_usage}/{token_limit} tokens')
+
+        # go through all messages and verify that the token count and total tokens is correct
+        total_tokens = 0
+        real_tokens = []
+        stored_tokens = []
+        for msg in message_manager.history.messages:
+            total_tokens += msg.metadata.input_tokens
+            stored_tokens.append(msg.metadata.input_tokens)
+            real_tokens.append(message_manager._count_tokens(msg.message))
+        assert total_tokens == sum(real_tokens)
+        assert stored_tokens == real_tokens
+        assert message_manager.history.total_tokens == total_tokens
+
+
+# pytest -s lumivor/agent/message_manager/tests.py

lumivor/agent/message_manager/views.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage
+from pydantic import BaseModel, Field
+
+
+class MessageMetadata(BaseModel):
+	"""Metadata for a message including token counts"""
+
+	input_tokens: int = 0
+
+
+class ManagedMessage(BaseModel):
+	"""A message with its metadata"""
+
+	message: BaseMessage
+	metadata: MessageMetadata = Field(default_factory=MessageMetadata)
+
+
+class MessageHistory(BaseModel):
+	"""Container for message history with metadata"""
+
+	messages: List[ManagedMessage] = Field(default_factory=list)
+	total_tokens: int = 0
+
+	def add_message(self, message: BaseMessage, metadata: MessageMetadata) -> None:
+		"""Add a message with metadata"""
+		self.messages.append(ManagedMessage(message=message, metadata=metadata))
+		self.total_tokens += metadata.input_tokens
+
+	def remove_message(self, index: int = -1) -> None:
+		"""Remove last message from history"""
+		if self.messages:
+			msg = self.messages.pop(index)
+			self.total_tokens -= msg.metadata.input_tokens