"""Lisette usage and cost monitoring"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_usage.ipynb.

# %% auto #0
__all__ = ['Usage', 'search_count', 'LisetteUsageLogger']

# %% ../nbs/01_usage.ipynb #5b6856c6
from litellm.integrations.custom_logger import CustomLogger
from fastcore.utils import *
import time
try: from fastlite import *
except ImportError: raise ImportError("Please install `fastlite` to use sqlite based lisette usage logging.")

# %% ../nbs/01_usage.ipynb #aed71558
class Usage: id:int; timestamp:float; model:str; user_id:str; prompt_tokens:int; completion_tokens:int; total_tokens:int; cached_tokens:int; cache_creation_tokens:int; cache_read_tokens:int; web_search_requests:int; response_cost:int

# %% ../nbs/01_usage.ipynb #4323da97
def search_count(r):
    if cnt := nested_idx(r.usage, 'server_tool_use', 'web_search_requests'): return cnt # Anthropic
    if meta := getattr(r, 'vertex_ai_grounding_metadata', None): # Gemini
        if meta and (queries := meta[0].get('webSearchQueries')): return len(queries)
    if cnt := nested_idx(r.usage, 'prompt_tokens_details', 'web_search_requests'): return cnt # streaming with `include_usage`
    return 0

# %% ../nbs/01_usage.ipynb #0ad2e088
class LisetteUsageLogger(CustomLogger):
    def __init__(self, db_path): 
        self.db = Database(db_path)
        self.usage = self.db.create(Usage)
    
    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): self._log_usage(response_obj, kwargs, start_time, end_time)
    def log_success_event(self, kwargs, response_obj, start_time, end_time):             self._log_usage(response_obj, kwargs, start_time, end_time)
    def _log_usage(self, response_obj, kwargs, start_time, end_time):
        response_cost = kwargs.get('response_cost') or nested_idx(kwargs,'standard_logging_object','cost_breakdown','total_cost')
        usage = response_obj.usage
        ptd   = usage.prompt_tokens_details
        self.usage.insert(Usage(timestamp=time.time(),
                                model=response_obj.model,
                                user_id=self.user_id_fn(),
                                prompt_tokens=usage.prompt_tokens,
                                completion_tokens=usage.completion_tokens,
                                total_tokens=usage.total_tokens,
                                cached_tokens=ptd.cached_tokens if ptd else 0, # used by gemini (read tokens)
                                cache_creation_tokens=nested_idx(usage, 'cache_creation_input_tokens'),
                                cache_read_tokens=nested_idx(usage, 'cache_read_input_tokens'), # used by anthropic
                                web_search_requests=search_count(response_obj),
                                response_cost=response_cost))
                  
    def user_id_fn(self): raise NotImplementedError('Please implement `LisetteUsageLogger.user_id_fn` before initializing, e.g using fastcore.patch.')

# %% ../nbs/01_usage.ipynb #9ef8ac82
@patch
def total_cost(self:Usage, sc=0.01): return self.response_cost + sc * ifnone(self.web_search_requests, 0)
