feat: Exploring global caching for LM calls

[KCaverly]

I took a look into how we can integrate global joblib Memory caching into the BaseLM class.

It takes a little bit of misdirection, but the below BaseLM class will provide for global caching for all sub-classes that implement BaseLM.

_cache_memory = Memory(cachedir, verbose=0)


class BaseLM(BaseModel, ABC):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._cached = _cache_memory.cache(self._call)

    def __call__(self, prompt: str, **kwargs) -> list[str]:
        """Generates `n` predictions for the signature output."""
        return self._cached(prompt, **kwargs)

    @abstractmethod
    def _call(
        self,
        prompt: str,
        **kwargs,
    ) -> list[str]:
        """Generates `n` predictions for the signature output."""
        ...

    @abstractmethod
    def count_tokens(self, prompt: str) -> int:
        """Counts the number of tokens for a specific prompt."""
        ...

The only change to downstream LM implementations, is now the abstract method is called _call as opposed to __call__. Additionally, with this scenario the cache is based on both LM attributes and function arguments. This is essential in the LiteLLM case, as the model name is essential for appropriate hashing, but each subsequent modification in the class itself, in which attributes are added/removed, will use a new cache. To combat this, the _cache_memory.cache does have an ignore parameter, which we can use to isolate exactly what is necessary for cache hashing.

This also is grabbing the same cache_utils, we are using in the current version of dspy, and I have a few questions on that generally:

1. Would we be okay with just global joblib.Memory for all caching needs?

   • I noticed the current version is using nested caching for what appears to be Notebook first, then global caching. I am unsure exactly what benefit this is providing for.

2. Small note, but can we hide this cachedir behind a hidden folder?

   • Currently, the cachedir is stored at 'Path.home() / cachedir_joblib' can we change this to 'Path.home() / .cachedir_job' with this refactor, so all new LiteLLM caching calls are not cluttering up the visible files in the home directory.

Let me know what you think @CyrusOfEden, @okhat, I know cache misses are an important point for some folks.

[CyrusNuevoDia]

Honestly project root / .joblib_cache would be my fav here, but I'm not sure if there's a way to do it consistently, so home dir / .joblib_cache would be nice too. I'm in favour of the .prefix. @okhat ?

[KCaverly]

Added cache specific settings, with a caching default set to True.

[okhat]

What do you think of having the LMs implement forward instead of __call__ for this kind of reason? @CyrusOfEden

[okhat]

(Also honestly I think having a more mature caching system is growing on me. It should be possible to configure and pass around, I sort of like how LangChain does caching but it should be ON by default in some way)