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

# %% auto 0
__all__ = []

# %% ../01_toolloop.ipynb
from .core import *
from fastcore.utils import *
from fastcore.meta import delegates
from fastcore.xtras import save_iter
from functools import wraps

from anthropic.types import TextBlock, Message, ToolUseBlock

# %% ../01_toolloop.ipynb
_final_prompt = "You have no more tool uses. Please summarize your findings. If you did not complete your goal please tell the user what further work needs to be done so they can choose how best to proceed."

# %% ../01_toolloop.ipynb
@patch
@delegates(Chat.__call__)
def toolloop(self:Chat,
             pr, # Prompt to pass to Claude
             max_steps=10, # Maximum number of tool requests to loop through
             cont_func:callable=noop, # Function that stops loop if returns False
             final_prompt=_final_prompt, # Prompt to add if last message is a tool call
             **kwargs):
    "Add prompt `pr` to dialog and get a response from Claude, automatically following up with `tool_use` messages"
    @save_iter
    def _f(o):
        init_n = len(self.h)
        r = self(pr, **kwargs)
        yield r
        if len(self.last)>1: yield self.last[1]
        for i in range(max_steps-1):
            if self.c.stop_reason!='tool_use': break
            r = self(final_prompt if i==max_steps-2 else None, **kwargs)
            yield r
            if len(self.last)>1: yield self.last[1]
            if not cont_func(*self.h[-3:]): break
        o.value = self.h[init_n+1:]
    return _f()

# %% ../01_toolloop.ipynb
from .asink import AsyncChat

# %% ../01_toolloop.ipynb
@patch
@delegates(AsyncChat.__call__)
def toolloop(
    self: AsyncChat,
    pr, # Prompt to pass to Claude
    max_steps=10, # Maximum number of tool requests to loop through
    cont_func: callable = noop, # Function that stops loop if returns False
    final_prompt = _final_prompt, # Prompt to add if last message is a tool call
    **kwargs
):
    "Add prompt `pr` to dialog and get a response from Claude, automatically following up with `tool_use` messages"
    @save_iter
    async def _f(o):
        init_n = len(self.h)
        r = await self(pr, **kwargs)
        yield r
        if len(self.last)>1: yield self.last[1]
        for i in range(max_steps-1):
            if self.c.stop_reason != 'tool_use': break
            r = await self(final_prompt if i==max_steps-2 else None, **kwargs)
            yield r
            if len(self.last)>1: yield self.last[1]
            if not cont_func(*self.h[-3:]): break
        o.value = self.h[init_n+1:]
    return _f()
