Backup-bdg/OpenHands

1import json2from dataclasses import dataclass3from typing import Iterable4 5from openhands.core.logger import openhands_logger as logger6from openhands.events.event import Event, EventSource7from openhands.events.event_filter import EventFilter8from openhands.events.event_store_abc import EventStoreABC9from openhands.events.serialization.event import event_from_dict10from openhands.storage.files import FileStore11from openhands.storage.locations import (12    get_conversation_dir,13    get_conversation_event_filename,14    get_conversation_events_dir,15)16from openhands.utils.shutdown_listener import should_continue17 18 19@dataclass(frozen=True)20class _CachePage:21    events: list[dict] | None22    start: int23    end: int24 25    def covers(self, global_index: int) -> bool:26        if global_index < self.start:27            return False28        if global_index >= self.end:29            return False30        return True31 32    def get_event(self, global_index: int) -> Event | None:33        # If there was not actually a cached page, return None34        if not self.events:35            return None36        local_index = global_index - self.start37        return event_from_dict(self.events[local_index])38 39 40_DUMMY_PAGE = _CachePage(None, 1, -1)41 42 43@dataclass44class EventStore(EventStoreABC):45    """46    A stored list of events backing a conversation47    """48 49    sid: str50    file_store: FileStore51    user_id: str | None52    cur_id: int = -1  # We fix this in post init if it is not specified53    cache_size: int = 2554 55    def __post_init__(self) -> None:56        if self.cur_id >= 0:57            return58        events = []59        try:60            events_dir = get_conversation_events_dir(self.sid, self.user_id)61            events = self.file_store.list(events_dir)62        except FileNotFoundError:63            logger.debug(f'No events found for session {self.sid} at {events_dir}')64 65        if not events:66            self.cur_id = 067            return68 69        # if we have events, we need to find the highest id to prepare for new events70        for event_str in events:71            id = self._get_id_from_filename(event_str)72            if id >= self.cur_id:73                self.cur_id = id + 174 75    def search_events(76        self,77        start_id: int = 0,78        end_id: int | None = None,79        reverse: bool = False,80        filter: EventFilter | None = None,81        limit: int | None = None,82    ) -> Iterable[Event]:83        """84        Retrieve events from the event stream, optionally filtering out events of a given type85        and events marked as hidden.86 87        Args:88            start_id: The ID of the first event to retrieve. Defaults to 0.89            end_id: The ID of the last event to retrieve. Defaults to the last event in the stream.90            reverse: Whether to retrieve events in reverse order. Defaults to False.91            filter: EventFilter to use92 93        Yields:94            Events from the stream that match the criteria.95        """96 97        if end_id is None:98            end_id = self.cur_id99        else:100            end_id += 1  # From inclusive to exclusive101 102        if reverse:103            step = -1104            start_id, end_id = end_id, start_id105            start_id -= 1106            end_id -= 1107        else:108            step = 1109 110        cache_page = _DUMMY_PAGE111        num_results = 0112        for index in range(start_id, end_id, step):113            if not should_continue():114                return115            if not cache_page.covers(index):116                cache_page = self._load_cache_page_for_index(index)117            event = cache_page.get_event(index)118            if event is None:119                try:120                    event = self.get_event(index)121                except FileNotFoundError:122                    event = None123            if event:124                if not filter or filter.include(event):125                    yield event126                    num_results += 1127                    if limit and limit <= num_results:128                        return129 130    def get_event(self, id: int) -> Event:131        filename = self._get_filename_for_id(id, self.user_id)132        content = self.file_store.read(filename)133        data = json.loads(content)134        return event_from_dict(data)135 136    def get_latest_event(self) -> Event:137        return self.get_event(self.cur_id - 1)138 139    def get_latest_event_id(self) -> int:140        return self.cur_id - 1141 142    def filtered_events_by_source(self, source: EventSource) -> Iterable[Event]:143        for event in self.get_events():144            if event.source == source:145                yield event146 147    def _get_filename_for_id(self, id: int, user_id: str | None) -> str:148        return get_conversation_event_filename(self.sid, id, user_id)149 150    def _get_filename_for_cache(self, start: int, end: int) -> str:151        return f'{get_conversation_dir(self.sid, self.user_id)}event_cache/{start}-{end}.json'152 153    def _load_cache_page(self, start: int, end: int) -> _CachePage:154        """Read a page from the cache. Reading individual events is slow when there are a lot of them, so we use pages."""155        cache_filename = self._get_filename_for_cache(start, end)156        try:157            content = self.file_store.read(cache_filename)158            events = json.loads(content)159        except FileNotFoundError:160            events = None161        page = _CachePage(events, start, end)162        return page163 164    def _load_cache_page_for_index(self, index: int) -> _CachePage:165        offset = index % self.cache_size166        index -= offset167        return self._load_cache_page(index, index + self.cache_size)168 169    @staticmethod170    def _get_id_from_filename(filename: str) -> int:171        try:172            return int(filename.split('/')[-1].split('.')[0])173        except ValueError:174            logger.warning(f'get id from filename ({filename}) failed.')175            return -1176