# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_client.ipynb.

# %% auto 0
__all__ = ['valid_func', 'Client', 'Chat']

# %% ../nbs/03_client.ipynb 4
from .utils import *
from .core import *

import os

from google import genai
from google.genai import types

from fastcore.all import *



# %% ../nbs/03_client.ipynb 7
@patch
def __call__(self: genai.models.Models, 
             inps=None, # The inputs to be passed to the model
             sp:str='', # Optional system prompt
             temp:float=0.6, # Temperature
             maxtok:int|None=None, # Maximum number of tokens for the output
             stream:bool=False, # Stream response?
             stop:str|list[str]|None=None, # Stop sequence[s]
             tools=None, # A list of functions or tools to be passed to the model
             use_afc=True, # Use Google's automatic function calling? If False, functions will be converted to tools
             # `AUTO` lets the model decide whether tools to use, 
             # `ANY` forces the model to call a function `NONE` avoids any function calling
             tool_mode='AUTO',  
             **kwargs):
    """Call to a Gemini LLM"""
    kwargs["model"] = kwargs.get("model", getattr(self, "model", None))
    model = kwargs["model"]
    prepped_tools = None
    if tools:
        self._tools = tools
        prepped_tools = prep_tools(tools, toolify_everything=not use_afc)  
    config = self._genconf(sp=sp, temp=temp, maxtok=maxtok, stop=stop, tools=prepped_tools,
                           tool_mode=tool_mode, **kwargs)
    
    contents = mk_contents(inps, cli=kwargs.get('client', None))    
    gen_f = self.generate_content_stream if stream else self.generate_content
    r = gen_f(model=model, contents=contents, config=config if config else None)
    return self._stream(r) if stream else self._r(r)


@patch
@delegates(genai.models.Models.__call__)
def __call__(self: genai.Client, inps=None, **kwargs):
    return self.models(inps, client=self, **kwargs)

# %% ../nbs/03_client.ipynb 11
def Client(model:str, # The model to be used by default (can be overridden when generating)
           sp:str='', # System prompt
           temp:float=0.6, # Default temperature
           text_only:bool=False, # Suppress multimodality even if the model allows for it
          ): 
    """An extension of `google.genai.Client` with a series of quality of life improvements"""
    c = genai.Client(api_key=os.environ['GEMINI_API_KEY'])
    c.models.post_cbs = [c.models._call_tools]
    c.models.model, c.models.sp, c.models.temp, c.models.text_only = model, sp, temp, text_only
    return c

# %% ../nbs/03_client.ipynb 16
@patch
def _repr_markdown_(self:genai.models.Models):
    if not hasattr(self,'result'): return 'No results yet'
    msg = self.result._repr_markdown_()
    return f"""{msg}


|        | Input | Output | Cached |
|--------|------:|-------:|-------:|
| Tokens | {self.use.inp:,} | {self.use.out:,}  | {self.use.cached:,}  |
| **Totals** | **Tokens: {self.use.total:,}** | **${self.cost:.6f}** |  |
"""

@patch
def _repr_markdown_(self:genai.Client): return self.models._repr_markdown_()

# %% ../nbs/03_client.ipynb 19
@patch
def structured(self: genai.models.Models, inps, tool, model=None, **kwargs):
    _ = self(inps,  tools=[tool], use_afc=False, tool_mode="ANY", temp=0., stream=False, **kwargs)
    return [nested_idx(ct, "function_response", "response", "result") for ct in nested_idx(self, "result_content", -1, "parts") or []]

@patch
def structured(self: genai.Client, inps, tool, model=None):
    return self.models.structured(inps, tool, model)

# %% ../nbs/03_client.ipynb 30
valid_func = genai.chats._validate_response

@patch(as_prop=True)
def c(self: genai.chats.Chat): return self._modules

@patch(as_prop=True)
def h(self: genai.chats.Chat): return self._curated_history

@patch(as_prop=True)
def full_h(self: genai.chats.Chat): return self._comprehensive_history

@patch
def _rec_res(self: genai.chats.Chat, resp):
    if not getattr(self, "user_query", False): return
    resp_c = nested_idx(resp, "candidates", 0, "content")
    self.record_history(
        user_input=self.user_query,
        model_output=[resp_c] if resp_c else [],
        automatic_function_calling_history=resp.automatic_function_calling_history,
        is_valid=valid_func(resp)
    )

def Chat(model:str, # The model to be used 
           sp:str='', # System prompt
           temp:float=0.6, # Default temperature
           text_only:bool=False, # Suppress multimodality even if the model allows for it
           cli:genai.Client|None=None, # Optional Client to be passed (to keep track of usage)
          ): 
    """An extension of `google.genai.chats.Chat` with a series of quality of life improvements"""        
    c = Client(model, sp, temp, text_only) if cli is None else cli
    c.model, c.sp, c.temp, c.text_only = model, sp, temp, text_only
    chat = c.chats.create(model=c.model)
    chat.c.post_cbs.insert(0, chat._rec_res)
    return chat

# %% ../nbs/03_client.ipynb 32
@patch
@delegates(genai.Client.__call__, keep=True)
def __call__(self: genai.chats.Chat, inps=None, **kwargs):
    self.user_query = mk_content(inps) if inps else self.c.result_content[-1]
    return self.c(self.h + [self.user_query], **kwargs)

# %% ../nbs/03_client.ipynb 40
@patch(as_prop=True)
def use(self: genai.chats.Chat): return self.c.use

@patch(as_prop=True)
def cost(self: genai.chats.Chat): return self.c.cost

@patch
def _repr_markdown_(self: genai.chats.Chat):
    if not hasattr(self.c, 'result'): return 'No results yet'
    last_msg = self.c.result._repr_markdown_().split('<details>')[0]

    def content_repr(ct):
        r = ct.role
        cts = {'text': '', 'images': []}
        for part in ct.parts:
            if part.text is not None: cts['text'] += part.text
            if part.inline_data is not None:
                cts['images'].append(types.Image(image_bytes=part.inline_data.data, mime_type=part.inline_data.mime_type))
        return f"_**{r}**_: {cts['text']}" + f"{' '.join(['IMAGE_' + str(i) for i, _ in enumerate(cts['images'])])}"
    
    history = "\n\n".join([content_repr(ct) for ct in self.h])
    det = self.c._repr_markdown_().split('\n\n')[-1]
    return f"""{last_msg}

<details>
<summary>History</summary>

{history}
</details>

{det}"""
