99 lines
3.5 KiB
Python
99 lines
3.5 KiB
Python
from pydantic import BaseModel
|
|
from typing import Callable, Iterable, Any
|
|
|
|
class Window[T](BaseModel):
|
|
"""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 = 10240,
|
|
side_min_count: int = 2,
|
|
side_max_size: int = 2560,
|
|
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 |