llm-app

1
import functools
2
import os
3
from abc import ABC
4

5
import diskcache
6
import pathway as pw
7

8

9
class _Cache:
10
    """A simple cache"""
11

12
    def __init__(self) -> None:
13
        if cache_dir := os.environ.get("PATHWAY_CACHE_DIR"):
14
            self.cache = diskcache.Cache(cache_dir)
15
        elif cache_dir := os.environ.get("PATHWAY_PERSISTENT_STORAGE"):
16
            self.cache = diskcache.Cache(cache_dir)
17
        else:
18
            self.cache = {}
19

20
    def __call__(self, fun):
21
        base_name = f"{fun.__module__}_{fun.__qualname__}"
22

23
        @functools.wraps(fun)
24
        def wrapper(*args, **kwargs):
25
            key = f"{base_name}({(args, kwargs)})"
26
            if key not in self.cache:
27
                self.cache[key] = fun(*args, **kwargs)
28
            return self.cache[key]
29

30
        return wrapper
31

32

33
class BaseModel(ABC):
34
    def __init__(self):
35
        self.cache = _Cache()
36

37
    def __call__(self, text: str, **kwargs):
38
        raise NotImplementedError()
39

40
    def apply(
41
        self,
42
        text: pw.ColumnExpression | str,
43
        **kwargs,
44
    ) -> pw.ColumnExpression:
45
        return pw.apply_async(self.cache(self.__call__), text=text, **kwargs)
46