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

# %% auto 0
__all__ = ['def_model', 'def_sysp', 'payload', 'set_cell', 'get_cells', 'sigils', 'mk_dialog', 'display_stream', 'stream_resp',
           'AiMagic', 'create_magic', 'load_ipython_extension', 'create_ipython_config']

# %% ../00_core.ipynb
import html,litellm
from lisette import Chat,mk_msgs,StreamFormatter
from bs4 import BeautifulSoup
from collections.abc import Iterable
from IPython.display import Markdown,Javascript,clear_output
from IPython.paths import get_ipython_dir
from fastcore.utils import *
from fastcore.xml import *
from datetime import date
from execnb.nbio import *

# %% ../00_core.ipynb
def_model = 'anthropic/claude-haiku-4-5'
litellm.suppress_debug_info = True

# %% ../00_core.ipynb
def payload(ip, single=True, **kw):
    "kwargs wrapper for `IPython.payload_manager.write_payload`"
    ip.payload_manager.write_payload(kw, single=single)

# %% ../00_core.ipynb
def set_cell(ip, text, code=True, replace=False, execute=False, offset=1):
    "Create or replace an nb cell containing `text` offset from the active cell"
    if not code: execute=True
    ctype = 'code' if code else 'markdown'
    payload(ip, source='set_next_input', replace=replace, execute=execute, text=text, ctype=ctype, offset=offset)

# %% ../00_core.ipynb
from fastcore.xml import Output,Error,Message,Messages,CellSource,Context,Dialog,Variable,Variables,Outputs

# %% ../00_core.ipynb
def _get_output(o):
    "Get the output from the provided cell"
    ot = o['output_type']
    if ot in ('stream','execute_result','display_data'): return Output(o['text' if ot=='stream' else 'data'], type=ot)
    elif o['output_type']=='error': return Error(o['tb'], evalue=o['evalue'])
    raise Exception(o)

# %% ../00_core.ipynb
def _cellxml(c):
    "Cell `c` converted to XML"
    elems = [CellSource(c['source'])]
    outs = c.get('outputs', [])
    ol = [_get_output(o) for o in outs]
    if ol: elems.append(Outputs(*ol))
    return Message(*elems, type=c['cell_type'])

# %% ../00_core.ipynb
def _mk_var(var, name):
    "Create context variable."
    return Variable(repr(var), name=name, type=type(var).__name__) 

# %% ../00_core.ipynb
def _mk_var_elems(vs, ns):
    "Create context variables."
    elems = []
    for v in vs:
        if v not in ns: continue
        elems += [_mk_var(ns[v], v)]
    return elems

# %% ../00_core.ipynb
def _mk_ctx(cells):
    "Context for Claude nb `cells`"
    if not cells: return ''
    r = []
    elems = [_cellxml(c) for c in cells]
    dtxt = "*NB:* These messages have been executed since the last prompt."
    r.append(Messages(dtxt, *elems))
    res = to_xml(Context(*r), 0, True, do_escape=False)
    return res+'\n' if res else ''

# %% ../00_core.ipynb
def _keep_cell(src):
    li = first(src.splitlines())
    if not li: return False
    if re.match(r'%ai\s+reset\s*$', li): return True
    if re.match(r'%ai\s+skip\s*$', li): return False
    return not (li.startswith('%ai') and '\n' not in src)

# %% ../00_core.ipynb
def _get_cells(cells):
    matches = [i for i,o in enumerate(cells) if o.get('source','').startswith('%ai reset')]
    if matches: cells = cells[matches[-1]+1:]
    return [o for o in cells if _keep_cell(o.get('source',''))]

def get_cells(nbm, offset=1):
    return _get_cells(nbm['cells'][:nbm['idx'] + offset])

# %% ../00_core.ipynb
_pp = '''{context_goes_here}<instructions>
{instructions_go_here}
</instructions>

<task>
{prompt_goes_here}
</task>'''

_prose_pp = 'Respond to the request in the `task` below.'

# %% ../00_core.ipynb
def _mk_prompt(aic, rep, cells):
    "Prompt for AI cell `aic` and optional reply `rep`, with nb cell context `cells`, and optionally $`variable`s expanded"
    magic,prompt = (aic['source'] + '\n').split('\n', 1)
    prompt = prompt.strip()
    cmd,*rest = magic.split()
    inst = _prose_pp
    ctx = _mk_ctx(cells)
    fullp = _pp.format(prompt_goes_here=prompt, context_goes_here=ctx, instructions_go_here=inst)
    res = [fullp]
    if rep:
        src = rep['source'].strip()
        if rep['cell_type']=='code': src = f'```\n{src}\n```'
        res.append(src)
    return res

# %% ../00_core.ipynb
def _mk_var_ctx(vs, ns):
    "Context for Claude using variables `vs`, with variables from namespace `ns`"
    if not vs: return ''
    elems = _mk_var_elems(vs, ns)
    names = ', '.join(f"`{o}`" for o in vs)
    varstxt = "This 'variables' section shows the *CURRENT* values of variables that the user has chosen to share. The values are *NOT* previous values from earlier -- they are the latest, current values. The following variable names have been shared: " + names
    r = Variables(varstxt, *elems)
    return to_xml(r, 0, True, do_escape=False)

# %% ../00_core.ipynb
def sigils(s, sigil, expand):
    return set([re.sub(fr'\{sigil}`(\w+)`', r'\1', o) for o in re.findall(fr'\{sigil}`\w+`', s)] if expand else [])

def mk_dialog(cells, ns=None, expand=True):
    "Split `cells` into groups based on ai magics, and create Claude dialog messages"
    res = []
    n = len(cells) # we declare `n` because `cells` is updated within the loop
    for _ in range(n):
        m = first_match(cells, lambda o: o.get('source','').startswith('%%ai'), 0)
        aic = cells[m]
        rep = cells[m+1] if m<len(cells)-1 else None
        res += _mk_prompt(aic, rep, cells[1:m])
        cells = cells[m+1 if rep else m:]
        if len(cells)<2: break
    text = '\n\n'.join(str(o) for o in res)
#     print(text)
    if ns is None: ns = get_ipython().user_ns
    vars  = sigils(text, '$', expand)
    tools = [ns[o] for o in sigils(text, '&', True) if o in ns]
    vres = _mk_var_ctx(vars, ns)
    return res,vres,tools

# %% ../00_core.ipynb
from fastcore.xml import Message,Dialog

# %% ../00_core.ipynb
_roles = ('user','assistant')

def _prettify_dialog(dialog): return html.unescape(BeautifulSoup(dialog, 'html.parser').prettify())

def _get_dlg_content(c):
    if isinstance(c, list): return c
    # content is a dict when caching is enabled
    if isinstance(c, dict): return [m.get('text','') for m in c['content']]
    return c

def _get_dialog(dialog, vres):
    msgs = [Message(_get_dlg_content(o),role=_roles[i%2]) for i,o in enumerate(dialog)]
    if vres and msgs:
        msgs[-1] = vres+msgs[-1]
    return _prettify_dialog(to_xml(Dialog(*msgs), do_escape=False))

# %% ../00_core.ipynb
def display_stream(rs):
    "Use IPython.display to markdown display the response stream."
    fmt = StreamFormatter()
    md = ''
    dh = display(Markdown(''), display_id=True)
    try:
        for o in fmt.format_stream(rs): 
            md+=o
            dh.update(Markdown(md))
    except KeyboardInterrupt: pass
    return md

# %% ../00_core.ipynb
def stream_resp(model, dlg, vres, tools, shell, sp=None, trace=False, search='l', think='l', **kw):
    *msgs,prompt = dlg
    if vres: prompt = to_xml(vres, do_escape=False)+prompt
    *msgs,prompt = mk_msgs(msgs+[prompt])
    cts = ''
    if tools and 'gemini' in model: search=None
    c = Chat(model=model, hist=msgs, sp=sp, tools=listify(tools) or None, ns=shell.user_ns, **kw)
    r = c(prompt, stream=True, search=search, think=think)
    cts = display_stream(r)
    if cts.strip(): clear_output()
    if trace: print(f'{sp}\n\n----\n\n{_get_dialog(dlg)}')
    set_cell(shell or get_ipython(), cts.strip(), code=False)

# %% ../00_core.ipynb
if '__file__' not in globals(): __file__ = 'aimagic/core.py'
_parent = Path(__file__).parent
def_sysp = (_parent/'sysp.txt').read_text()

# %% ../00_core.ipynb
class AiMagic:
    def __init__(self, llm, shell=None, sysp=def_sysp, search='l', think='l', **kw):
        "Backend functionality for `create_magic`"
        self.kw = kw
        shell = shell or get_ipython()
        store_attr()

    def _handle(self, line):
        if not line: return
        cmd,*rest = map(str.strip, line.split(' ', 1))
        if cmd=='0': return 0
        f = getattr(self, '_'+cmd, None)
        if not f: return
        return f(*rest)

    def __call__(self, line='', cell='', temp=0.5):
        ns = self.shell.user_ns
        lres = self._handle(line)
        if lres is not None: return lres if lres else None
        cells = get_cells(ns['nbmeta'])
        if not cell: cell=''
        dlg,vres,tools = mk_dialog(cells, ns=ns, expand=True)
        if line.startswith('dialog'):
            return set_cell(self.shell, f"```html\n{_get_dialog(dlg)}\n```", code=False, execute=False)
        if not cell: return
        sp = self.sysp.format(today = date.today().strftime('%Y-%m-%d'))
        stream_resp(self.llm, dlg, vres, tools, sp=self.sysp, shell=self.shell, trace=line=='-t',
                    search=self.search, think=self.think, **self.kw)

# %% ../00_core.ipynb
@patch
def _model(self:AiMagic, rest=None):
    if rest: self.llm = rest
    print(f"Model: {self.llm}")
    return 0

@patch
def _think(self:AiMagic, rest=None):
    if rest: self.think = rest[0]
    print(f"Thinking: {self.think}")
    return 0

@patch
def _reset(self:AiMagic): return 0

@patch
def _obj(self:AiMagic): return self

# %% ../00_core.ipynb
def create_magic(llm=None, sysp=def_sysp, nm='ai', shell=None, search='l', think='l', **kw):
    "Create magic named `nm` using `llm` and sys prompt `sysp`"
    if not llm: llm = def_model
    nm = str(nm)
    if not shell: shell = get_ipython()
    r = AiMagic(llm, shell, sysp=sysp, search=search, think=think, **kw)

    js = '''require(['base/js/namespace', 'notebook/js/codecell'], function(j, c) {
    c.CodeCell.options_default.highlight_modes['magic_markdown'] = {'reg':['^%%ai']} ;
    j.notebook.get_cells().forEach(function(cell){
        if (cell.cell_type === 'code'){ cell.auto_highlight(); }
    });
});'''
    display(Javascript(js))
    
    def f(line, cell=None): return r(line, cell)
    shell.register_magic_function(partial(f), 'line_cell', nm)

# %% ../00_core.ipynb
def load_ipython_extension(ipython): create_magic(shell=ipython)

# %% ../00_core.ipynb
def create_ipython_config():
    ipython_dir = Path(get_ipython_dir())
    cf = ipython_dir/'profile_default'/'ipython_config.py'
    cf.parent.mkdir(parents=True, exist_ok=True)
    if cf.exists() and 'aimagic' in cf.read_text(): return print('aimagic already installed!')
    with cf.open(mode='a') as f: f.write("\nc.InteractiveShellApp.extensions.append('aimagic.core')\n\n")
    print(f"Jupyter config updated at {cf}")
