"""Access wordslab-notebooks Jupyterlab extension version, current notebook path, json content and cell id, and create or update cells."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_notebook.ipynb.

# %% auto 0
__all__ = ['prompt_template', 'FUNC_RE', 'VAR_RE', 'find_var', 'WordslabNotebook', 'ChatTurn']

# %% ../nbs/02_notebook.ipynb 2
from ast import literal_eval
import asyncio
from datetime import datetime
from functools import partial
from html import escape
from inspect import currentframe, getattr_static, getdoc, isfunction, ismethod, signature
import re
from textwrap import dedent
from types import ModuleType, FunctionType, MethodType, BuiltinFunctionType
import time
import typing

from IPython.core.getipython import get_ipython
from IPython.core.interactiveshell import InteractiveShell
from IPython.core.oinspect import Inspector
from IPython.display import display, HTML, Markdown, clear_output
from comm import create_comm
import nbformat

from fastcore.utils import L, patch
from fastcore.xml import to_xml, Note, Prompt, Code, Source, Outputs, User, Assistant, Out, Var
from toolslm.funccall import get_schema

from ollama import chat, ChatResponse

from .env import WordslabEnv

# %% ../nbs/02_notebook.ipynb 8
@patch
def _get_info(self: Inspector, obj, oname='', formatter=None, info=None, detail_level=0, omit_sections=()):
    """Custom Python formatter for ?? output:
    - suppress the useless code indentation
    - wrap the code in a markdown python code block
    - display the output as markdown
    """
    orig = self._orig__get_info(obj, oname=oname, formatter=formatter, info=info,
                                detail_level=detail_level, omit_sections=omit_sections)
    if detail_level == 0:
        return orig
    info_dict = self.info(obj, oname=oname, info=info, detail_level=detail_level)
    out = []
    if c := info_dict.get('source'): 
        out.append(f"\n```python\n{dedent(c)}\n```")
    if c := info_dict.get('file'): 
        out.append(f"**File:** `{c}`")
    return {'text/markdown': '\n\n'.join(out), 'text/html': '', 'text/plain': orig['text/plain']}

# %% ../nbs/02_notebook.ipynb 10
def _safe_str(obj, max_str_len=200):
    "Safely get the string representation of an object, truncating if it exceeds max_len."
    try:
        s = str(obj)
        return s[:max_str_len] + ("…" if len(s) > max_str_len else "")
    except Exception as e: 
        return f"<str error: {str(e)}>"


@patch
def user_items(self: InteractiveShell, max_str_len=200, xtra_ignore=()):
    """Get an overview of the variables & functions defined by the user so far in the notebook.
    The value addded by this function is to filter out all internal ipython and wordslab variables.
    Returns a tuple of dictionaries (user_variables, user_functions):
    - the keys are the variables or function names
    - the value is a truncated string representation of the variable value or the function signature
    The `max_str_len` parameter is used to truncate the string representation of the variables.
    The `xtra_ignore` parameter is used to hide additional names from the result. 
    """
    ns, nsh = self.user_ns, self.user_ns_hidden
    ignore = set()  # Add here the wordslab specific vars and funcs we want to hide
    ignore.add(xtra_ignore)
    rm_types = (
        type, FunctionType, ModuleType, MethodType, BuiltinFunctionType,
        getattr(typing, '_SpecialGenericAlias', ()),
        getattr(typing, '_GenericAlias', ()),
        getattr(typing, '_SpecialForm', ())
    )
    user_items = {k: v for k, v in ns.items()
                  if k not in ignore and k not in nsh}
    user_vars = {k: _safe_str(v, max_str_len=max_str_len)
                 for k, v in user_items.items() if not k.startswith('_') and not isinstance(v, rm_types)}
    user_fns = {k: str(signature(v)) for k, v in user_items.items()
                if isinstance(v, FunctionType) and v.__module__ == '__main__' and not k.startswith('__')}
    return user_vars, user_fns

# %% ../nbs/02_notebook.ipynb 12
@patch
def get_variables_values(self: InteractiveShell, var_names: list, literal=True):
    """Get a safe and serializable representation of variables values from the user namespace.
    This method preserves real Python values when they are safe literals, otherwise it falls back to strings.
    You can call it in two modes:
    - literal = True : Preserve actual Python values when safe, best for internal tools
    - literal = False : Force everything to strings, best for logging / UI display / debug output
    """
    ns = self.user_ns

    def _maybe_eval(o):
        try:
            literal_eval(repr(o))
            return o
        except:
            return str(o)
    return {v: _maybe_eval(ns[v]) if literal else str(ns[v]) for v in var_names if v in ns}

# %% ../nbs/02_notebook.ipynb 14
def _get_schema(ns: dict, t):
    "Check if tool `t` has errors."
    if t not in ns:
        return f"`{t}` not found. Did you run the cell where it is defined?"
    try:
        return {'type': 'function', 'function': get_schema(ns[t], pname='parameters')}
    except Exception as e:
        return f"`{t}`: {e}."

@patch
def get_tools_schemas_and_functions(self: InteractiveShell, func_names: list):
    """Get a json schema and a function object for the functions defined in the user namespace which can be used as tools."""
    ns = self.user_ns
    return {f: (_get_schema(ns, f), ns[f]) for f in func_names if f in ns}

# %% ../nbs/02_notebook.ipynb 19
def _find_frame_dict(var: str):
    "Find the dict (globals or locals) containing var"
    frame = currentframe().f_back.f_back
    while frame:
        if var in frame.f_globals:
            return frame.f_globals
        frame = frame.f_back
    raise ValueError(f"Could not find {var} in any scope")

def find_var(var: str):
    "Search for var in all frames of the call stack"
    return _find_frame_dict(var)[var]

# %% ../nbs/02_notebook.ipynb 20
class WordslabNotebook:
    """Jupyterlab notebook introspection and metaprogramming."""

    def __init__(self):
        # Check if we are running inside a Jupyter notebook
        if get_ipython() is None:
            raise RuntimeError("This class can only be used in the context of a Jupyter notebook")

        # Check if the wordslab-notebooks-lib Jupyterlab frontend extension is installed
        try:
            find_var("__wordslab_extension_version")
            self.jupyterlab_extension_installed = True
        except:
            self.jupyterlab_extension_installed = False

        # Initialize a communication channel with the frontend extension
        if self.jupyterlab_extension_installed:
            self._comm = WordslabNotebook.JupyterlabExtensionComm()

        # Default notebook assistant config
        self.chat_model = WordslabEnv().default_model_code
        self.chat_thinking = True

    def _ensure_jupyterlab_extension(self):
        if not self.jupyterlab_extension_installed:
            raise RuntimeError(
                "The JupyterLab extension for wordslab-notebooks is not activated: "
                "please execute `pip install wordslab-notebooks-lib` in JupyterLab virtual environment "
                "and refresh your web browser."
            )

    # --------------------------------
    # Frontend -> Kernel communication
    # --------------------------------

    # The frontend extension injects the following variables before each cell execution

    @property
    def jupyterlab_extension_version(self):
        """wordslab-notebooks-lib version number injected by the Jupyterlab frontend extension"""
        self._ensure_jupyterlab_extension()
        return find_var("__wordslab_extension_version")

    @property
    def path(self):
        """Relative path of the notebook .ipynb file in the notebook workspace"""
        self._ensure_jupyterlab_extension()
        return find_var("__notebook_path")

    @property
    def content(self):
        """Full content of the notebook returned as a NotebookNode object from the nbformat library"""
        self._ensure_jupyterlab_extension()
        return nbformat.from_dict(find_var("__notebook_content"))

    @property
    def cell_id(self):
        """Unique ID of the current notebook cell, useful to locate the current cell in the full notebook content"""
        self._ensure_jupyterlab_extension()
        return find_var("__cell_id")

    # --------------------------------
    # Kernel -> Frontend communication
    # --------------------------------

    # The kernel sends commands to the frontend extension through this comms channel

    class JupyterlabExtensionComm:
        def __init__(self, target_name='wordslab_notebooks', timeout=2.0):
            self.target_name = target_name
            self._init_comm()
            self.timeout = timeout
            self.result = None

        def _init_comm(self):
            self.comm = create_comm(target_name=self.target_name)
            self.comm.on_msg(self._on_msg)

        def _on_msg(self, msg):
            self.result = msg['content']['data']

        async def send(self, data):
            self.result = None
            self.comm.send(data)

            start = time.time()
            while self.result is None and (time.time() - start) < self.timeout:
                await asyncio.sleep(0.01)

            if self.result is None:
                self._init_comm()
                raise TimeoutError('No response from Jupyterlab frontend. If you just refreshed the browser, this timeout is expected: please retry running this cell.')

            return self.result

# %% ../nbs/02_notebook.ipynb 29
@patch
async def add_cell(
    self: WordslabNotebook,
    content: str,  # Content of the cell (i.e the prompt, code, or note cell text)
    placement: str = 'add_after',  # Can be 'add_after', 'add_before', 'at_start', 'at_end'
    cell_id: str = None,  # id of the cell that placement is relative to (if None, uses current cell)
    cell_type: str = 'note',  # Cell type, can be 'code', 'note', or 'prompt'
    notebook_path: str = ''  # Notebook to update, defaults to current notebook
):
    """Add a cell to the current notebook or any other opened notebook (`notebook_path`),
    at the start/end of the notebook or before/after any cell (`placement`and `cell_id`),
    with a `cell_type` (note|prompt|code) and `content` (text).
    Returns the new cell id."""
    self._ensure_jupyterlab_extension()
    if notebook_path and notebook_path != self.path:
        if placement not in ('at_start', 'at_end') and not cell_id:
            raise ValueError("`cell_id` or `placement='at_end'`/`placement='at_start'` must be provided when target notebook is different")
    if placement not in ('at_start', 'at_end') and not cell_id:
        cell_id = self.cell_id
    result = await self._comm.send({'action': 'create_cell', 'cell_type': cell_type, 'content': content, 'placement': placement, 'cell_id': cell_id, 'notebook_path': notebook_path})
    if 'success' in result and result['success']:
        return result['cell_id']
    elif 'error' in result:
        raise RuntimeError(result['error'])

# %% ../nbs/02_notebook.ipynb 49
@patch
async def update_cell(
    self: WordslabNotebook,
    cell_id: str = None,  # id of the cell to update (if None, uses current cell)
    content: str = None,  # Content of the cell (i.e the prompt, code, or note cell text)
    notebook_path: str = ''  # Notebook to update, defaults to current notebook
):
    """Update the cell identified by `cell_id`,
    in the current notebook or any other opened notebook (`notebook_path`),
    with a new `content`.
    Returns the updated cell id."""
    self._ensure_jupyterlab_extension()
    if not cell_id:
        raise ValueError("`cell_id` parameter is mandatory")
    result = await self._comm.send({'action': 'update_cell', 'cell_id': cell_id, 'content': content, 'notebook_path': notebook_path})
    if 'success' in result and result['success']:
        return result['cell_id']
    elif 'error' in result:
        raise RuntimeError(result['error'])

# %% ../nbs/02_notebook.ipynb 57
@patch
async def delete_cell(
    self: WordslabNotebook,
    cell_id: str = None,  # id of cell to delete
    notebook_path: str = ''  # Notebook to update, defaults to current notebook
):
    """"Update the cell identified by `cell_id`,
    in the current notebook or any other opened notebook (`notebook_path`).
    Returns the deleted cell id."""
    self._ensure_jupyterlab_extension()
    if not cell_id:
        raise ValueError("`cell_id` parameter is mandatory")
    result = await self._comm.send({'action': 'delete_cell', 'cell_id': cell_id, 'notebook_path': notebook_path})
    if 'success' in result and result['success']:
        return result['cell_id']
    elif 'error' in result:
        raise RuntimeError(result['error'])

# %% ../nbs/02_notebook.ipynb 60
@patch
async def run_cell(
    self: WordslabNotebook,
    cell_id: str = None,  # id of cell to execute
):
    """"Adds the cell identified by `cell_id` to the run queue, only in the current notebook (jupyterlab 'run-cell' command limitation).
    Returns the cell id.
    DOES NOT return the result of the execution: the target cell will only be run after the current cell execution finishes.
    Use the `read_cell` method later with the same `cell_id` to get the result of the execution."""
    self._ensure_jupyterlab_extension()
    if not cell_id:
        raise ValueError("`cell_id` parameter is mandatory")
    result = await self._comm.send({'action': 'run_cell', 'cell_id': cell_id})
    if 'success' in result and result['success']:
        return result['cell_id']
    elif 'error' in result:
        raise RuntimeError(result['error'])

# %% ../nbs/02_notebook.ipynb 64
@patch
def read_cell(
    self: WordslabNotebook,
    cell_id: str = None,  # id of cell to delete
):
    """"Read the text content of the cell identified by `cell_id`, only in the current notebook.
    Returns the text content of the cell as a single multiline string."""
    self._ensure_jupyterlab_extension()
    if not cell_id:
        raise ValueError("`cell_id` parameter is mandatory")
    cell = next((c for c in self.content.cells if c.id == cell_id), None)
    if not cell:
        raise ValueError(f"Cell not found: {cell_id}")
    return cell.source

# %% ../nbs/02_notebook.ipynb 68
@patch
def show_variables_and_functions(self: WordslabNotebook):
    """Display the variables and functions defined by the user so far in the notebook."""
    variables, functions = get_ipython().user_items()
    output = "<h4>Variables</h4><table><tr><th>Name</th><th>Value</th></tr>"
    for var, value in variables.items():
        output += f"<tr><td>{var}</td><td>{escape(value)}</td></tr>"
    output += "</table>"
    output += "<h4>Functions</h4><table><tr><th>Name</th><th>Signature</th></tr>"
    for func, sig in functions.items():
        output += f"<tr><td>{func}</td><td>{escape(sig)}</td></tr>"
    output += "</table>"
    display(HTML(output))

# %% ../nbs/02_notebook.ipynb 70
def _safe_getattr(obj, name):
    try:
        return getattr(obj, name)
    except:
        return None

def _safe_attr_doc(obj, name, max_str_len=200):
    # Try to get the static attribute from the class
    try:
        class_attr = getattr_static(obj.__class__, name)
    except AttributeError:
        class_attr = None
    # 1. Property
    if isinstance(class_attr, property):
        attr_doc = getdoc(class_attr)
    # 2. classmethod / staticmethod
    elif isinstance(class_attr, (classmethod, staticmethod)):
        attr_doc = getdoc(class_attr.__func__)
    # 3. Function (instance method)
    elif isfunction(class_attr):
        attr_doc = getdoc(class_attr)
    # 4. Descriptor with __doc__ (other descriptors)
    elif hasattr(class_attr, "__doc__") and class_attr.__doc__:
        attr_doc = getdoc(class_attr)
    # 5. Fallback to instance attribute
    elif hasattr(obj, name):
        instance_attr = getattr(obj, name)
        if not type(instance_attr).__module__ == "builtins" and hasattr(instance_attr, "__doc__") and instance_attr.__doc__:
            attr_doc = getdoc(instance_attr)
        else:
            attr_doc = ""
    # Max length
    if attr_doc is None:
        attr_doc = ""
    elif len(attr_doc) > max_str_len:
        attr_doc = attr_doc[:max_str_len] + "…"
    return attr_doc

@patch
def show_object_members(self: WordslabNotebook, obj):
    """Display the attributes and methods of a given python object"""
    obj_class = obj.__class__
    output = (f"<h3>Object of type: {obj_class.__name__}</h3>") 
    output += f"<pre>{getdoc(obj)}</pre>"

    obj_members = [(name, _safe_getattr(obj, name)) for name in dir(obj) if not name.startswith("_")]

    output += "<h4>Attributes</h4>"
    output += "<table><tr><th>Name</th><th>Type</th><th>Value</th><th>Doc</th></tr>"
    for name, value in obj_members:
        if name.startswith("_"):
            continue
        # Skip callables (handled as methods below)
        if callable(value):
            continue
        attr_type = type(value).__name__
        attr_value = _safe_str(value)
        attr_doc = _safe_attr_doc(obj, name)
        output += f"<tr><td>{name}</td><td>{escape(attr_type)}</td><td>{escape(attr_value)}</td><td>{escape(attr_doc)}</td></tr>"
    output += "</table>"

    output += "<h4>Methods</h4>"
    output += "<table><tr><th>Name</th><th>Signatue</th><th>Type</th><th>Doc</th></tr>"
    for name, value in obj_members:
        if name.startswith("_"):
            continue
        # Skip attributes (handled above)
        if not callable(value):
            continue
        # Determine method type
        method_type = "instance method"
        if isfunction(value):
            # Defined on the class
            class_attr = getattr(obj_class, name, None)
            if isinstance(class_attr, classmethod):
                method_type = "class method"
            elif isinstance(class_attr, staticmethod):
                method_type = "static method"
        elif ismethod(value):
            method_type = "instance method"
        # Signature
        try:
            sig = str(signature(value))
        except (ValueError, TypeError):
            sig = "(...)"
        # Doc
        method_doc = _safe_attr_doc(obj, name)
        output += f"<tr><td>{name}</td><td>{escape(sig)}</td><td>{method_type}</td><td>{escape(method_doc)}</td></tr>"
    output += "</table>"

    display(HTML(output))

# %% ../nbs/02_notebook.ipynb 73
@patch
def get_variables_values(self: WordslabNotebook, var_names: list):
    """Get a safe and serializable representation of variables values."""
    return get_ipython().get_variables_values(var_names=var_names)

@patch
def get_tools_schemas_and_functions(self: WordslabNotebook, func_names: list):
    """Get a json schema of functions which can be used as tools."""
    return get_ipython().get_tools_schemas_and_functions(func_names=func_names)

# %% ../nbs/02_notebook.ipynb 83
def _mime_bundle_to_text(data):
    "Get text from MIME bundle, preferring markdown over plain"
    if 'text/markdown' in data:
        return ('markdown', ''.join(list(data['text/markdown'])))
    if 'text/html' in data:
        return ('html', ''.join(list(data['text/html'])))
    if 'text/plain' in data: 
        return ('text', ''.join(list(data['text/plain'])))

def _cell_output_to_xml(o):
    "Convert single notebook output to XML format"
    # execute_result — the return value of the last expression, or calls to display()
    if hasattr(o, 'data'):
        mime, txt = _mime_bundle_to_text(o.data)
        if txt:
            return Out(txt, type=mime)
    # stream — stdout/stderr text (e.g., from print())
    if hasattr(o, 'text'):
        txt = o.text if isinstance(o.text, str) else ''.join(o.text)
        return Out(txt, type=o.get('name', 'stdout'))
    # Error - exceptions
    if hasattr(o, 'ename'):
        return Out(f"{o.ename}: {o.evalue}", type='error')

# %% ../nbs/02_notebook.ipynb 85
def _cell_to_xml(cell):
    "Convert notebook cell to concise XML format"
    src = ''.join(getattr(cell, 'source', ''))
    if cell.cell_type == 'markdown':
        return Note(src)
    elif cell.cell_type == 'code':
        out_items = L(getattr(cell,'outputs',[])).map(_cell_output_to_xml).filter()
        is_prompt = "wordslab_cell_type" in cell.metadata and cell.metadata["wordslab_cell_type"] == "prompt"
        if is_prompt:
            parts = [User(src)]
            if out_items:
                parts.append(Assistant(*out_items))
            return Prompt(*parts)
        else:
            parts = [Source(src)]
            if out_items:
                parts.append(Outputs(*out_items))
            return Code(*parts)
    else:
        return None

# %% ../nbs/02_notebook.ipynb 87
def _cells_to_notebook_xml(cells):
    cells_xml = [_cell_to_xml(c) for c in cells if c.cell_type in ('code', 'markdown')]
    return "\n".join(L(cells_xml).map(partial(to_xml, do_escape=False)))

# %% ../nbs/02_notebook.ipynb 89
@patch
def get_context_for_llm(self: WordslabNotebook):
    # Get current notebook state
    all_cells = self.content.cells
    current_cell_id = self.cell_id

    # Exclude all cells after the previous one
    previous_cells = all_cells[:next((i for i, c in enumerate(all_cells) if c.id == current_cell_id), len(all_cells))]

    # Then exclude all cells hidden from AI 
    context_cells = [cell for cell in previous_cells if "wordslab_hide_from_ai" not in cell.metadata]

    # Convert the remaining markdown and code cells to XML
    notebook_context_xml = _cells_to_notebook_xml(context_cells)

    # Return a XML string version of the notebook
    return to_xml(notebook_context_xml)

# %% ../nbs/02_notebook.ipynb 93
prompt_template = """
<system_instructions>

<role_and_behavior>
You are an AI assistant designed to help the user learn and solve problems interactively.
You work with the user step-by-step rather than just giving complete answers. You are especially good at:
- Breaking down complex topics into manageable pieces
- Helping the user work through coding problems in Python
- Encouraging the user to try things himself, with guidance when he needs it
- Adapting to the user level and interests
You are designed to be collaborative - you ask questions, check the user understanding, and let him explore ideas rather than just lecturing. 
You can help with teaching, coding, problem-solving, research, and creative projects. You assume that the user is very smart.
What sets you apart is your teaching approach - you focus on helping the user develop his skills rather than just giving him answers. 
You provide information in small chunks, check in frequently to see if things make sense, and encourage the user to try things himself.
</role_and_behavior>

<execution_context_info>
Your design is heavily inspired by the Solveit platform developed by Answer.ai, but adapted for a local and open source Jupyterlab environment.
You run in an interactive Juyter notebook environment where the user can take notes, write code, and chat with you.
This notebook mixes three types of cells:
- "note" cells (green border), which contain markdown text written by the user
- "prompt" cells (red border), which contain a user instruction and an assistant answer, both in markdown text
- "code" cells (blue border), which contain python source code and optionnaly outputs - the result of the execution
The cells are meant to be read and executed in chronological order from top to bottom.
You receive this instruction in the following context :
- the user just executed a "prompt" cell
- in this cell, he typed a user instruction which will be provided below
- you must execute this user instruction TAKING INTO ACCOUNT all the previous cells in the notebook
- the content of the previous cells is also provided below, you must interpret it as a CONVERSATION HISTORY
- your answer will be rendered as an output of the prompt cell in markdown format
The "code" cells of the notebook are backed by a python kernel which maintains state with functions and variables. 
The user can optionnaly provide a set of python functions for you to use as tools:
- if the user mentions the function with this very specific syntax `&myfunc` somewhere in the notebook (backticks mandatory)
- then the description of the function myfunc and its parameters will be provided to you as a tool you can call
Be careful to always check first if you could call these tools to better ground your answer, instead of trying to guess based on your pretraining knwoledge.
Only if you were not provided with the right tool, you can also generate and display code blocks in your response, that the user will be able to review and copy in a new "code" cell to execute it.
The user can optionnaly provide some of the python variables values as information for you:
- if the user mentions the variable with this very specific syntax `$myvar` somewhere in the notebook (backticks mandatory)
- then the value of the variable myvar will be provided to you as information to use to execute the user instruction
Same remark: prioritize using this information to ground your answer to the user question instead of trying to guess.
</execution_context_info>

<prompt_format_spec>
According to this execution context information, the prompt will be structured with XML tags as follows
- conversation_history -> list of cells of the notebook, which represents the history of the conversation
  - note -> markdown cell, directly contains a note from the user
  - prompt -> prompt sent to you during a previous turn of the conversation
    - user -> instruction written by the user
    - assistant -> answer previously generated by you
      - out -> part of the answer which was displayed
  - code -> python cell executed by the user
    - source -> python code written and executed by the user
    - outputs -> results of the python kernel execution
      - out -> part of the execution result which was displayed     
- referenced_variables -> optional
  - var name=... -> value of python variable referenced by the user (may be truncated if too long)
- user_instruction -> the very last user instruction written in a prompt cell you need to execute 
The definitions of the tools you can call are injected through the API.
</prompt_format_spec>

<output_format_spec>
Make sure you always give the FINAL answer in the same language as the user instruction.
For example, if the user_instruction below is written in french, answer in french, if it is written in german, answer in german.
But off course, for the intermediate turns in an agentic loop, this rule doesn't apply: you can generate as many tool calls as needed before generating the FINAL answer.
Try to be concise: just provide the answer, don't explain the obvious unless explicity prompted to do so.
Assume the user is intelligent and has no time to waste reading too long answers.
When you generate code, make sure to do it inside markdown fenced blocks.
</system_instructions>

<conversation_history>

{notebook_context}

</conversation_history>

<referenced_variables>

{referenced_variables}

</referenced_variables>

<user_instruction>

{user_instruction}

</user_instruction>
"""

# %% ../nbs/02_notebook.ipynb 95
FUNC_RE = re.compile(r"`&([a-zA-Z_][a-zA-Z0-9_]*)`")
VAR_RE  = re.compile(r"`\$([a-zA-Z_][a-zA-Z0-9_]*)`")

# %% ../nbs/02_notebook.ipynb 112
class ChatTurn:
    def __init__(self):
        self.thinking_chunks = []
        self.content_chunks = []
        self.tool_calls = {}
        self.turn_finished = False

    def append_thinking(self, chunk:str):
        self.thinking_chunks.append(chunk)

    @property
    def thinking(self):
        return "".join(self.thinking_chunks)
    
    def append_content(self, chunk:str):
        self.content_chunks.append(chunk)

    @property
    def content(self):
        return "".join(self.content_chunks)
    
    def append_tool_call(self, tool_name:str, params:dict):
        self.tool_calls[tool_name] = { "params": params }

    def start_tool_call(self, tool_name:str):
        self.tool_calls[tool_name]["start_time"] = time.time()

    def end_tool_call(self, tool_name:str, result: object):
        self.tool_calls[tool_name]["end_time"] = time.time()
        self.tool_calls[tool_name]["result"] = str(result)

    def end_turn(self):
        self.turn_finished = True
    
    def to_markdown(self, hide_thinking:bool=True, hide_tool_calls:bool=True):
        output = ""
        if len(self.thinking_chunks) > 0:
            if not hide_thinking:
                output += "> [Thinking]\n>\n> "
                output += ("".join(self.thinking_chunks)).replace("\n\n","\n>\n> ").replace("\n- ","\n - ") + "\n\n"
            else:
                output += f"> [Thinking] ... thought in {sum(s.count(' ') + s.count('\n') for s in self.thinking_chunks)} words\n\n"
        if len(self.content_chunks) > 0:
            output += "".join(self.content_chunks) + "\n\n"
        if len(self.tool_calls.keys()) > 0:
            if not hide_tool_calls:
                for tool_name in self.tool_calls.keys():
                    output += "> [Tool call]\n"
                    tool_call = self.tool_calls[tool_name]
                    if "params" in tool_call:
                        output += f"> - model wants to call `{tool_name}` with parameters `{tool_call["params"]}`\n"
                    if "start_time" in tool_call:
                        output += f"> - agent called {tool_name} at {datetime.fromtimestamp(tool_call["start_time"]).strftime("%H:%M:%S")}\n"
                    if "end_time" in tool_call:
                        result = tool_call["result"]
                        output += f"> - {tool_name} returned `{result if len(result)<=100 else result[:97]+'...'}` in {(tool_call["end_time"]-tool_call["start_time"]):.3f} sec\n"
                    output += "\n"
            else:
                for tool_name in self.tool_calls.keys():
                    tool_call = self.tool_calls[tool_name]
                    if "end_time" in tool_call:
                        result = self.tool_calls[tool_name]["result"]
                        output += f"> [Tool call] ... `{tool_name}` returned `{result if len(result)<=50 else result[:47]+'...'}`\n\n"
                    elif "start_time" in tool_call:                        
                        output += f"> [Tool call] ... agent is calling `{tool_name}`\n\n"
                    elif "params" in tool_call:                        
                        output += f"> [Tool call] ... model wants to call `{tool_name}`\n\n"
        return output       

# %% ../nbs/02_notebook.ipynb 116
def _refresh_display(chat_turns):
    clear_output(wait=True)
    output = ""
    for turn in chat_turns:
        output += turn.to_markdown()
    display(Markdown(output))

@patch
def chat(self: WordslabNotebook, user_instruction: str):    
    # Get notebook context
    notebook_context = self.get_context_for_llm()
    # Extract referenced tools and variables
    funcs_names = FUNC_RE.findall(notebook_context)
    vars_names = VAR_RE.findall(notebook_context)
    # Get tools schemas 
    tools = self.get_tools_schemas_and_functions(funcs_names)
    tools_schemas = [t[0] for t in tools.values()]
    # Get variables values
    vars_values = self.get_variables_values(vars_names)
    referenced_variables = "\n".join(L([Var(value, name=name) for name,value in vars_values.items()]).map(to_xml))
    
    # Format the prompt
    prompt = prompt_template.format(notebook_context=notebook_context,
                                    referenced_variables=referenced_variables,
                                    user_instruction=user_instruction)

    # Immediate user feedback
    display(Markdown(f"Processing {prompt.count(' ') + prompt.count('\n')} words (prompt+context) with `{self.chat_model}` ..."))
    
    # ollama agentic loop

    chat_turns = []
    messages = [{'role': 'user', 'content': prompt}]
    while True:
      stream = chat(
         model=self.chat_model,
         messages=messages,
         tools=tools_schemas,
         stream=True,
         think=self.chat_thinking,
      )
      
      chat_turn = ChatTurn()
      chat_turns.append(chat_turn)
        
      tool_calls = []

      # accumulate the partial fields
      for chunk in stream:
        if chunk.message.thinking:
          chat_turn.append_thinking(chunk.message.thinking)
          _refresh_display(chat_turns)
        if chunk.message.content:
          chat_turn.append_content(chunk.message.content)
          _refresh_display(chat_turns)
        if chunk.message.tool_calls:
          tool_calls.extend(chunk.message.tool_calls)            
          for tc in chunk.message.tool_calls:
            chat_turn.append_tool_call(tc.function.name, tc.function.arguments)
          _refresh_display(chat_turns)
    
      # append accumulated fields to the messages
      if chat_turn.thinking or chat_turn.content or tool_calls:
        messages.append({'role': 'assistant', 'thinking': chat_turn.thinking, 'content': chat_turn.content, 'tool_calls': tool_calls})

      # end the loop if there is no more tool calls
      if not tool_calls:    
        chat_turn.end_turn()
        break      
        
      # execute tool calls  
      if tool_calls:    
          for tc in tool_calls:
             if tc.function.name in tools:
                chat_turn.start_tool_call(tc.function.name)
                _refresh_display(chat_turns)
                result = tools[tc.function.name][1](**tc.function.arguments)
                chat_turn.end_tool_call(tc.function.name, result)
                _refresh_display(chat_turns)
             else:
               result = 'Unknown tool'
    
             # append tool call result to the messages 
             messages.append({'role': 'tool', 'tool_name': tc.function.name, 'content': str(result)})
              
    chat_turn.end_turn()
