"""Lisette Core"""

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

# %% auto 0
__all__ = ['effort', 'patch_litellm', 'mk_msg', 'mk_msgs', 'stream_with_complete', 'lite_mk_func', 'cite_footnote',
           'cite_footnotes', 'Chat', 'astream_with_complete', 'AsyncChat', 'aformat_stream', 'adisplay_stream']

# %% ../nbs/00_core.ipynb
import asyncio, base64, json, litellm, mimetypes
from typing import Optional
from html import escape
from IPython.display import Markdown
from litellm import (acompletion, completion, stream_chunk_builder, Message,
                     ModelResponse, ModelResponseStream, get_model_info)
from litellm.utils import function_to_dict, StreamingChoices, Delta
from toolslm.funccall import mk_ns, call_func, call_func_async, get_schema
from fastcore.utils import *
from fastcore import imghdr

# %% ../nbs/00_core.ipynb
def patch_litellm():
    "Patch litellm.ModelResponseBase such that `id` and `created` are fixed."
    from litellm.types.utils import ModelResponseBase
    @patch
    def __init__(self: ModelResponseBase, id=None, created=None, *args, **kwargs): 
        self._orig___init__(id='chatcmpl-xxx', created=1000000000, *args, **kwargs)

    @patch
    def __setattr__(self: ModelResponseBase, name, value):
        if name == 'id': value = 'chatcmpl-xxx'
        elif name == 'created': value = 1000000000
        self._orig___setattr__(name, value)

# %% ../nbs/00_core.ipynb
@patch
def _repr_markdown_(self: litellm.ModelResponse):
    message = self.choices[0].message
    content = ''
    if mc:=message.content: content += mc[0]['text'] if isinstance(mc,list) else mc
    if message.tool_calls:
        tool_calls = [f"\n\n🔧 {nested_idx(tc,'function','name')}({nested_idx(tc,'function','arguments')})\n" for tc in message.tool_calls]
        content += "\n".join(tool_calls)
    if not content: content = str(message)
    details = [
        f"id: `{self.id}`",
        f"model: `{self.model}`",
        f"finish_reason: `{self.choices[0].finish_reason}`"
    ]
    if hasattr(self, 'usage') and self.usage: details.append(f"usage: `{self.usage}`")
    det_str = '\n- '.join(details)
    
    return f"""{content}

<details>

- {det_str}

</details>"""

# %% ../nbs/00_core.ipynb
def _mk_img(data:bytes)->tuple:
    "Convert image bytes to a base64 encoded image"
    img = base64.b64encode(data).decode("utf-8")
    mtype = mimetypes.types_map["."+imghdr.what(None, h=data)]
    return img, mtype

# %% ../nbs/00_core.ipynb
def _is_img(data): 
    return isinstance(data, bytes) and bool(imghdr.what(None, data))

def _add_cache_control(msg,          # LiteLLM formatted msg
                       cache=False,  # Enable Anthropic caching
                       ttl=None):    # Cache TTL: '5m' (default) or '1h'
    "cache `msg` with default time-to-live (ttl) of 5minutes ('5m'), but can be set to '1h'."
    if not cache: return msg
    if isinstance(msg["content"], str): 
        msg["content"] = [{"type": "text", "text": msg["content"]}]
    cache_control = {"type": "ephemeral"}
    if ttl is not None: cache_control["ttl"] = ttl
    if isinstance(msg["content"], list) and msg["content"]:
        msg["content"][-1]["cache_control"] = cache_control
    return msg

def _remove_cache_ckpts(msg):
    "remove unnecessary cache checkpoints."
    if isinstance(msg["content"], list) and msg["content"]:
        msg["content"][-1].pop('cache_control', None)
    return msg

def _mk_content(o):
    if isinstance(o, str): return {'type':'text','text':o.strip() or '.'}
    if _is_img(o): 
        img, mtype = _mk_img(o)
        return {"type": "image_url", "image_url": f"data:{mtype};base64,{img}"}
    return o

def mk_msg(content,      # Content: str, bytes (image), list of mixed content, or dict w 'role' and 'content' fields
           role="user",  # Message role if content isn't already a dict/Message
           cache=False,  # Enable Anthropic caching
           ttl=None):    # Cache TTL: '5m' (default) or '1h'
    "Create a LiteLLM compatible message."
    if isinstance(content, dict) or isinstance(content, Message): return content
    if isinstance(content, list) and len(content) == 1 and isinstance(content[0], str): c = content[0]
    elif isinstance(content, list): c = [_mk_content(o) for o in content]
    else: c = content
    return _add_cache_control({"role": role, "content": c}, cache=cache, ttl=ttl)

# %% ../nbs/00_core.ipynb
def mk_msgs(msgs,                       # List of messages (each: str, bytes, list, or dict w 'role' and 'content' fields)
            cache=False,                # Enable Anthropic caching
            ttl=None,                   # Cache TTL: '5m' (default) or '1h'
            cache_last_ckpt_only=True   # Only cache the last message
           ):
    "Create a list of LiteLLM compatible messages."
    if not msgs: return []
    if not isinstance(msgs, list): msgs = [msgs]
    res,role = [],'user'
    for m in msgs:
        res.append(msg:=mk_msg(m, role=role,cache=cache))
        role = 'assistant' if msg['role'] in ('user','function', 'tool') else 'user'
    if cache_last_ckpt_only: res = [_remove_cache_ckpts(m) for m in res]
    if res and cache: res[-1] = _add_cache_control(res[-1], cache=cache, ttl=ttl)
    return res

# %% ../nbs/00_core.ipynb
def stream_with_complete(gen, postproc=noop):
    "Extend streaming response chunks with the complete response"
    chunks = []
    for chunk in gen:
        chunks.append(chunk)
        yield chunk
    postproc(chunks)
    return stream_chunk_builder(chunks)

# %% ../nbs/00_core.ipynb
def lite_mk_func(f):
    if isinstance(f, dict): return f
    return {'type':'function', 'function':get_schema(f, pname='parameters')}

# %% ../nbs/00_core.ipynb
def _lite_call_func(tc,ns,raise_on_err=True):
    try: fargs = json.loads(tc.function.arguments)
    except Exception as e: raise ValueError(f"Failed to parse function arguments: {tc.function.arguments}") from e
    res = call_func(tc.function.name, fargs,ns=ns)
    return {"tool_call_id": tc.id, "role": "tool", "name": tc.function.name, "content": str(res)}

# %% ../nbs/00_core.ipynb
def _has_search(m):
    i = get_model_info(m)
    return bool(i['search_context_cost_per_query'] or i['supports_web_search'])

# %% ../nbs/00_core.ipynb
def cite_footnote(msg):
    if not (delta:=nested_idx(msg, 'choices', 0, 'delta')): return
    if citation:= nested_idx(delta, 'provider_specific_fields', 'citation'):
        title = citation['title'].replace('"', '\\"')
        delta.content = f'[*]({citation["url"]} "{title}") '
        
def cite_footnotes(stream_list):
    "Add markdown footnote citations to stream deltas"
    for msg in stream_list: cite_footnote(msg)

# %% ../nbs/00_core.ipynb
effort = AttrDict({o[0]:o for o in ('low','medium','high')})

# %% ../nbs/00_core.ipynb
def _mk_prefill(pf): return ModelResponseStream([StreamingChoices(delta=Delta(content=pf,role='assistant'))])

# %% ../nbs/00_core.ipynb
class Chat:
    def __init__(
        self,
        model:str,                # LiteLLM compatible model name 
        sp='',                    # System prompt
        temp=0,                   # Temperature
        search=False,             # Search (l,m,h), if model supports it
        tools:list=None,          # Add tools
        hist:list=None,           # Chat history
        ns:Optional[dict]=None,   # Custom namespace for tool calling 
        cache=False               # Anthropic prompt caching
    ):
        "LiteLLM chat client."
        self.model = model
        hist,tools = mk_msgs(hist),listify(tools)
        if ns is None and tools: ns = mk_ns(tools)
        elif ns is None: ns = globals()
        self.tool_schemas = [lite_mk_func(t) for t in tools] if tools else None
        store_attr()
    
    def _prep_msg(self, msg=None, prefill=None):
        "Prepare the messages list for the API call"
        sp = [{"role": "system", "content": self.sp}] if self.sp else []
        if msg: self.hist = mk_msgs(self.hist+[msg], cache=self.cache)
        pf = [{"role":"assistant","content":prefill}] if prefill else []
        return sp + self.hist + pf

    def _call(self, msg=None, prefill=None, temp=None, think=None, search=None, stream=False, max_steps=2, step=1, final_prompt=None, tool_choice=None, **kwargs):
        "Internal method that always yields responses"
        if step>max_steps: return
        if not get_model_info(self.model)["supports_assistant_prefill"]: prefill=None
        if _has_search(self.model) and (s:=ifnone(search,self.search)): kwargs['web_search_options'] = {"search_context_size": effort[s]}
        else: _=kwargs.pop('web_search_options',None)
        res = completion(model=self.model, messages=self._prep_msg(msg, prefill), stream=stream, 
                         tools=self.tool_schemas, reasoning_effort = effort.get(think), tool_choice=tool_choice,
                         # temperature is not supported when reasoning
                         temperature=None if think else ifnone(temp,self.temp),
                         **kwargs)
        if stream:
            if prefill: yield _mk_prefill(prefill)
            res = yield from stream_with_complete(res,postproc=cite_footnotes)
        m = res.choices[0].message
        if prefill: m.content = prefill + m.content
        self.hist.append(m)
        yield res

        if tcs := m.tool_calls:
            tool_results=[_lite_call_func(tc, ns=self.ns) for tc in tcs]
            self.hist+=tool_results
            for r in tool_results: yield r
            if step>=max_steps-1: prompt,tool_choice,search = final_prompt,'none',False
            else: prompt = None
            yield from self._call(
                prompt, prefill, temp, think, search, stream, max_steps, step+1,
                final_prompt, tool_choice, **kwargs)
    
    def __call__(self,
                 msg=None,          # Message str, or list of multiple message parts
                 prefill=None,      # Prefill AI response if model supports it
                 temp=None,         # Override temp set on chat initialization
                 think=None,        # Thinking (l,m,h)
                 search=None,       # Override search set on chat initialization (l,m,h)
                 stream=False,      # Stream results
                 max_steps=2, # Maximum number of tool calls
                 final_prompt=None, # Final prompt when tool calls have ran out 
                 return_all=False,  # Returns all intermediate ModelResponses if not streaming and has tool calls
                 **kwargs):
        "Main call method - handles streaming vs non-streaming"
        result_gen = self._call(msg, prefill, temp, think, search, stream, max_steps, 1, final_prompt, **kwargs)     
        if stream: return result_gen              # streaming
        elif return_all: return list(result_gen)  # toolloop behavior
        else: return last(result_gen)             # normal chat behavior

# %% ../nbs/00_core.ipynb
async def _alite_call_func(tc, ns, raise_on_err=True):
    try: fargs = json.loads(tc.function.arguments)
    except Exception as e: raise ValueError(f"Failed to parse function arguments: {tc.function.arguments}") from e
    res = await call_func_async(tc.function.name, fargs, ns=ns)
    return {"tool_call_id": tc.id, "role": "tool", "name": tc.function.name, "content": str(res)}

# %% ../nbs/00_core.ipynb
@asave_iter
async def astream_with_complete(self, agen, postproc=noop):
    chunks = []
    async for chunk in agen:
        chunks.append(chunk)
        postproc(chunk)
        yield chunk
    self.value = stream_chunk_builder(chunks)

# %% ../nbs/00_core.ipynb
class AsyncChat(Chat):
    async def _call(self, msg=None, prefill=None, temp=None, think=None, search=None, stream=False, max_steps=2, step=1, final_prompt=None, tool_choice=None, **kwargs):
        if step>max_steps+1: return
        if not get_model_info(self.model)["supports_assistant_prefill"]: prefill=None
        if _has_search(self.model) and (s:=ifnone(search,self.search)): kwargs['web_search_options'] = {"search_context_size": effort[s]}
        else: _=kwargs.pop('web_search_options',None)
        res = await acompletion(model=self.model, messages=self._prep_msg(msg, prefill), stream=stream,
                         tools=self.tool_schemas, reasoning_effort=effort.get(think), tool_choice=tool_choice,
                         # temperature is not supported when reasoning
                         temperature=None if think else ifnone(temp,self.temp), 
                         **kwargs)
        if stream:
            if prefill: yield _mk_prefill(prefill)
            res = astream_with_complete(res,postproc=cite_footnote)
            async for chunk in res: yield chunk
            res = res.value
        m=res.choices[0].message
        if prefill: m.content = prefill + m.content
        yield res
        self.hist.append(m)

        if tcs := m.tool_calls:
            tool_results = []
            for tc in tcs:
                result = await _alite_call_func(tc, ns=self.ns)
                tool_results.append(result)
                yield result
            self.hist+=tool_results
            if step>=max_steps-1: prompt,tool_choice,search = final_prompt,'none',False
            else: prompt = None
            async for result in self._call(
                prompt, prefill, temp, think, search, stream, max_steps, step+1,
                final_prompt, tool_choice=tool_choice, **kwargs):
                    yield result
    
    async def __call__(self,
                       msg=None,          # Message str, or list of multiple message parts
                       prefill=None,      # Prefill AI response if model supports it
                       temp=None,         # Override temp set on chat initialization
                       think=None,        # Thinking (l,m,h)
                       search=None,       # Override search set on chat initialization (l,m,h)
                       stream=False,      # Stream results
                       max_steps=2, # Maximum number of tool calls
                       final_prompt=None, # Final prompt when tool calls have ran out 
                       return_all=False,  # Returns all intermediate ModelResponses if not streaming and has tool calls
                       **kwargs):
        result_gen = self._call(msg, prefill, temp, think, search, stream, max_steps, 1, final_prompt, **kwargs)
        if stream or return_all: return result_gen
        async for res in result_gen: pass
        return res # normal chat behavior only return last msg

# %% ../nbs/00_core.ipynb
def _clean_str(text):
    "Clean content to prevent breaking surrounding markdown formatting."
    return escape(str(text)).replace('`', '').replace('\n', ' ').replace('|', ' ')

# %% ../nbs/00_core.ipynb
def _trunc_str(s, mx=2000, replace="…"):
    "Truncate `s` to `mx` chars max, adding `replace` if truncated"
    s = str(s).strip()
    return s[:mx]+replace if len(s)>mx else s

# %% ../nbs/00_core.ipynb
async def aformat_stream(rs):
    "Format the response stream for markdown display."
    think = False
    async for o in rs:
        if isinstance(o, ModelResponseStream):
            d = o.choices[0].delta
            if nested_idx(d, 'reasoning_content'): 
                think = True
                yield '🧠'
            elif think:
                think = False
                yield '\n\n'
            if c := d.content: yield c
        elif isinstance(o, ModelResponse) and (c := getattr(o.choices[0].message, 'tool_calls', None)):
            fn = first(c).function
            yield f"\n<details class='tool-usage-details'>\n\n `{fn.name}({_trunc_str(fn.arguments)})`\n"
        elif isinstance(o, dict) and 'tool_call_id' in o: 
            yield f"  - `{_trunc_str(_clean_str(o.get('content')))}`\n\n</details>\n\n"

# %% ../nbs/00_core.ipynb
async def adisplay_stream(rs):
    "Use IPython.display to markdown display the response stream."
    md = ''
    async for o in aformat_stream(rs): 
        md+=o
        display(Markdown(md),clear=True)
