Huge refactoring

windowizer.py

@@ -0,0 +1,100 @@
+from dataclasses import dataclass
+from typing import Callable, Iterable, Any
+
+@dataclass
+class Window[T]:
+    """Window for AI processing."""
+
+    past: list[T]
+    """Part of data from previous window (can be empty)"""
+
+    present: list[T]
+    """Data for current window (can NOT be empty)"""
+
+    future: list[T]
+    """Part of data from the next window (can be empty)"""
+
+    context: dict[str, Any]
+    """Context that is preserved between LLM iterations"""
+
+
+class Windowizer[T]:
+    """This class builds windows from input data"""
+
+    @staticmethod
+    def _build_list_with_limits(items: Iterable[T], min_count: int, max_size: int, size_func: Callable[[T], int]) -> list[T]:
+        result: list[T] = []
+        total_size: int = 0
+        for i in items:
+            result.append(i)
+            total_size += size_func(i)
+            if len(result) < min_count:
+                continue
+            if total_size >= max_size:
+                break
+        return result
+
+    @staticmethod
+    def default_size_func(elem: T) -> int:
+        """Default size function just returns length of the `repr` result."""
+        return len(repr(elem))
+
+    def __init__(self,
+                 *,
+                 main_min_count: int = 2,
+                 main_max_size: int = 1024,
+                 side_min_count: int = 2,
+                 side_max_size: int = 256,
+                 size_func: Callable[[T], int] = default_size_func) -> None:
+        """Create Windowizer.
+        
+        Args:
+            - main_min_count - minimum count of items in `present`
+            - main_max_size - maximum total size of items in `present`
+            - side_min_count - minimum count of items in `past` and `future`
+            - side_max_size - maximum total size of items in `past` and `future`
+            - size_func - function that will be used to get item size
+        """
+        self._main_min_count = main_min_count
+        self._main_max_size = main_max_size
+        self._side_min_count = side_min_count
+        self._side_max_size = side_max_size
+        self._size_func = size_func
+
+    def windowize(self, items: list[T]) -> list[Window[T]]:
+        """Builds windows from items.
+
+        `min_count_*` has more priority than `max_size_*`. It would be possible
+        to build windows of 0 items otherwise.
+        """
+        result: list[Window[T]] = []
+        # create windows
+        window_start = 0
+        while window_start < len(items):
+            window = Window[T](past=[], present=[], future=[], context={})
+            # build `present`
+            window.present = self._build_list_with_limits(
+                items[window_start:],
+                min_count=self._main_min_count,
+                max_size=self._main_max_size,
+                size_func=self._size_func
+            )
+            # build `past`
+            window.past = self._build_list_with_limits(
+                reversed(items[:window_start]),
+                min_count=self._side_min_count,
+                max_size=self._side_max_size,
+                size_func=self._size_func
+            )
+            window.past.reverse()
+            # build `future`
+            window.future = self._build_list_with_limits(
+                items[window_start+len(window.present):],
+                min_count=self._side_min_count,
+                max_size=self._side_max_size,
+                size_func=self._size_func
+            )
+            window_start += len(window.present)
+            result.append(window)
+        # resulting windows
+        return result