"""Record your LLM calls and make your notebooks fast again."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['enable_reclm']

# %% ../nbs/00_core.ipynb 3
import hashlib, httpx, json, random, time
from importlib import util as ilib_util, import_module
from unittest.mock import patch as mpatch
from fastcore.utils import *

# %% ../nbs/00_core.ipynb 16
class _RecLMClient(httpx.Client):
    "Record and reuse your LLM calls."
    pass

# %% ../nbs/00_core.ipynb 22
@patch
def _req_data(self:_RecLMClient, req): return {'method':req.method, 'url':str(req.url), 'content':req.content.decode()}

@patch
def _hash(self:_RecLMClient, req): return hashlib.sha256(f"{json.dumps(self._req_data(req), sort_keys=True)}".encode()).hexdigest()

# %% ../nbs/00_core.ipynb 33
class _CacheLock:
    "Lock for our cache file."
    def __init__(self, fp):
        self.lock_fp=str(fp)+'.lock'
        self.max_retries=50
    
    def __enter__(self):
        retries=0
        while retries<self.max_retries:
            try:
                with open(self.lock_fp,'x'): return self
            except FileExistsError:
                time.sleep(0.05+random.random()*0.05)
                retries += 1
        raise TimeoutError(f"Could not acquire lock after {retries} attempts")
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        os.unlink(self.lock_fp)

# %% ../nbs/00_core.ipynb 42
def _cache_dir(): 
    try: return Config.find('settings.ini').config_path
    except AttributeError: return Path.cwd().resolve()  

# %% ../nbs/00_core.ipynb 49
@patch
def send(self:_RecLMClient, req, **kwargs):
    "Fetch `req` from the cache. If it doesn't exist, call the LLM and cache the response."
    is_stream = kwargs.get('stream')
    _hash = self._hash(req)
    with _CacheLock(self.cpth): cache = json.loads(self.cpth.read_text() or '{}')
    if resp:= cache.get(_hash): return httpx.Response(request=req, status_code=resp['status_code'], **{'content' if is_stream else 'json': resp['response']})
    resp = super(_RecLMClient,self).send(req, **kwargs)
    content = b''.join(list(resp.iter_bytes())).decode() if is_stream else resp.json()
    cache[_hash] = {'request':self._req_data(req), 'status_code':resp.status_code, 'response':content}
    with _CacheLock(self.cpth): self.cpth.write_text(json.dumps(cache))
    return httpx.Response(status_code=resp.status_code, content=content, request=req)

# %% ../nbs/00_core.ipynb 53
@patch
def __init__(self:_RecLMClient, *args, cache_dir, **kwargs):
    super(_RecLMClient,self).__init__(*args, **kwargs)
    self.cpth = Path(cache_dir) / ('updated_reclm.json' if os.getenv('CLEAN_RECLM_CACHE') else 'reclm.json')
    with _CacheLock(self.cpth): self.cpth.touch(exist_ok=True)

# %% ../nbs/00_core.ipynb 59
class _RecLMAsyncClient(httpx.AsyncClient):
    "Record and reuse your async LLM calls."
    def __init__(self:_RecLMClient, *args, cache_dir, **kwargs):
        super(_RecLMAsyncClient,self).__init__(*args, **kwargs)
        self.cpth = Path(cache_dir) / ('updated_reclm.json' if os.getenv('CLEAN_RECLM_CACHE') else 'reclm.json')
        with _CacheLock(self.cpth): self.cpth.touch(exist_ok=True)

    def _req_data(self, req): return {'method':req.method, 'url':str(req.url), 'content':req.content.decode()}
    def _hash(self, req): return hashlib.sha256(f"{json.dumps(self._req_data(req), sort_keys=True)}".encode()).hexdigest()

    async def send(self:_RecLMClient, req, **kwargs):
        "Fetch `req` from the cache. If it doesn't exist, call the LLM and cache the response."
        is_stream = kwargs.get('stream')
        _hash = self._hash(req)
        with _CacheLock(self.cpth): cache = json.loads(self.cpth.read_text() or '{}')
        if resp:= cache.get(_hash): return httpx.Response(request=req, status_code=resp['status_code'], **{'content' if is_stream else 'json': resp['response']})
        resp = await super(_RecLMAsyncClient,self).send(req, **kwargs)
        content = b''.join([c async for c in resp.aiter_bytes()]).decode() if is_stream else resp.json()
        cache[_hash] = {'request':self._req_data(req), 'status_code':resp.status_code, 'response':content}
        with _CacheLock(self.cpth): self.cpth.write_text(json.dumps(cache))
        return httpx.Response(status_code=resp.status_code, content=content, request=req)

# %% ../nbs/00_core.ipynb 61
_ispatched = False

def _mk_patch(sdk,cls,**kws):
    "Patch `sdk.cls.__init__` so that the `http_client` is a reclm client."
    init = getattr(import_module(sdk), cls).__init__
    reclm_client = _RecLMAsyncClient if cls.startswith('Async') else _RecLMClient
    return mpatch(f'{sdk}.{cls}.__init__', lambda *a,**k: init(*a, http_client=reclm_client(**kws), **k))

def enable_reclm(cache_dir=None):
    "Set the OpenAI and Anthropic `http_client` to a reclm client."
    if cache_dir and not Path(cache_dir).is_dir(): raise ValueError('`cache_dir` must be a directory.')
    kws = dict(cache_dir=cache_dir or _cache_dir())
    def _mk_patches(sdk, clss): return [_mk_patch(sdk,cls,**kws) for cls in clss if ilib_util.find_spec(sdk)]
    global _ispatched
    if _ispatched: return  # do not double-patch
    _ispatched = True
    ant_patches = _mk_patches('anthropic', ['Anthropic','AsyncAnthropic'])
    oai_patches = _mk_patches('openai', ['OpenAI','AsyncOpenAI'])
    for p in ant_patches+oai_patches: p.start()
    return
