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

# %% auto #0
__all__ = ['mk_schema', 'Agent', 'call_llm_tool', 'mk_transfer_doc', 'with_model_nm', 'Runner', 'msg_fragment',
           'mk_narrative_cast', 'filter_out_history', 'assistant_response', 'tool_response', 'mk_cb_msg']

# %% ../nbs/00_core.ipynb #82f94d48
import asyncio, inspect, json, uuid
from copy import deepcopy
from functools import partial
from itertools import groupby
from typing import Callable
from fastcore.basics import patch, Self
from fastcore.utils import L, store_attr
from litellm import ModelResponse, completion
from toolslm.funccall import get_schema
try: from IPython.display import display
except ImportError: display = print

# %% ../nbs/00_core.ipynb #b3bf967f
def mk_schema(fn, strip=("tool_ctx",)):
    sig = inspect.signature(fn)
    fn.__signature__ = sig.replace(parameters=[v for k,v in sig.parameters.items() if k not in strip])
    schema = {"type": "function", "function": get_schema(fn, pname="parameters")}
    fn.__signature__ = sig
    return schema

# %% ../nbs/00_core.ipynb #f7a38a90
class Agent:
    def __init__(self, 
        model:str,           # model str in litellm format
        name="",             # name of the agent
        desc="",             # what the agent does (used by parent agents to decide when to transfer)
        sp=None,             # the system prompt
        tools=None,          # a list of functions to provide to the agent
        subagents=None,      # a list subagents that this agent can transfer to
        narrative_cast=True, # to narrative cast or not
        include_hist=True,   # if False: will remove all history that doesn't belong to this agent on call
        before_cb=None,      # a callback func(state:dict, msgs:list). Any non-None return will skip the llm
        after_cb=None        # a callback func(state:dict, response) that can edit and modify the llm response
    ):
        store_attr(but="tools,subagents")
        self.tool_map = {}
        self.subagents = subagents or []
        if tools: [self.add_tool(tool) for tool in tools]
    
    @property
    def tool_schemas(self): return [schema for _,schema in self.tool_map.values()]
    def add_tool(self, tool:Callable): self.tool_map[tool.__name__] = (tool, mk_schema(tool))

# %% ../nbs/00_core.ipynb #2d6002e1
def call_llm_tool(tc_func:dict, tool_map:dict, state:dict):
    try: 
        fn,_ = tool_map[tc_func.name]
        kwargs = json.loads(tc_func.arguments)
        if "tool_ctx" in inspect.signature(fn).parameters: kwargs["tool_ctx"] = state
        result = fn(**kwargs)
    except Exception as e: 
        result = f"Error: {e}"
    return result if isinstance(result, str) else json.dumps(result)

# %% ../nbs/00_core.ipynb #3926d1c2
def mk_transfer_doc(agent, parent):
    doc = "Transfer runner to a different agent by providing an agent name.\nSupported agents:"
    if parent: doc += f"\n - (Parent) '{parent.name}': {parent.desc}"
    for sub in agent.subagents: doc += f"\n - '{sub.name}': {sub.desc}"
    return doc

# %% ../nbs/00_core.ipynb #357a806a
def with_model_nm(llm_output, nm:str):
    msg = llm_output.choices[0].message.content
    if not msg: return llm_output
    out = deepcopy(llm_output)
    out.choices[0].message.content = f"[{nm}]: {msg}"
    return out

# %% ../nbs/00_core.ipynb #2ec86448
class Runner:
    def __init__(self, agent:Agent):
        self.history = []
        self.state = {}
        self.agent_map = {}
        self.active_agent = agent
        if agent.subagents: self.register_subagents(agent)
    
    def register_subagents(self, agent:Agent, parent=None):
        def transfer_to_agent(agent_name:str): self.active_agent = self.agent_map[agent_name]
        transfer_to_agent.__doc__ = mk_transfer_doc(agent, parent)
        self.agent_map[agent.name] = agent
        agent.add_tool(transfer_to_agent)
        [self.register_subagents(sub, parent=agent) for sub in agent.subagents]
        
    def __call__(self, msg:list|str, max_calls=20, do_display=True, **kwargs):
        if isinstance(msg, str): self.history.append({"content": msg, "role": "user", "agent": self.active_agent.name})
        else:                    self.history = list(msg)

        for _ in range(max_calls):
            nm = self.active_agent.name
            llm_output = self.active_agent(self.history, state=self.state, **kwargs)
            if do_display: display(with_model_nm(llm_output, nm))
            choice = llm_output.choices[0]
            self.history.append({**choice.message.model_dump(), "agent": nm})
            if choice.finish_reason != "tool_calls": break
            for tool_call in choice.message.tool_calls:
                tool_res = call_llm_tool(tool_call.function, self.active_agent.tool_map, self.state)
                self.history.append({"role": "tool", "tool_call_id": tool_call.id, "content": tool_res, "agent": nm})
        return choice


# %% ../nbs/00_core.ipynb #2566e5f5
def msg_fragment(d:dict, results:dict):
    c = d.get("content") or ""
    for tc in d.get("tool_calls") or []:
        kwargs = ",".join([f"{k}={v}" for k,v in json.loads(tc["function"]["arguments"]).items()])
        output = results.get(tc["id"])
        c += f" {tc['function']['name']}({kwargs}) -> {output}"
    return c

def mk_narrative_cast(msg_list:list, active_agent:str):
    narrative_msg_list = []
    for key, group in groupby(msg_list, key=lambda x: "user" if x["role"] == "user" else x.get("agent")):
        if key in ("user", active_agent): narrative_msg_list.extend(group)
        else:
            group = list(group)
            results = {t["tool_call_id"]: t["content"] for t in group if t["role"] == "tool"}
            casted_msg = "\n".join(msg_fragment(d, results) for d in group if d["role"] != "tool")
            f = {"role": "user", "content": f"[Agent '{key}'] {casted_msg}"}
            narrative_msg_list.append(f)
    return narrative_msg_list

# %% ../nbs/00_core.ipynb #ab42f4db
def filter_out_history(msgs:list, active_agent:str):
    end = L(msgs).argfirst(lambda o: o["role"] == "assistant" and o.get("agent") == active_agent)
    valid_starts = L(msgs)[:end].argwhere(lambda o: o["role"] == "user" or o["role"] == "assistant" and o.get("agent") != active_agent)
    return msgs[valid_starts[-1]:] if valid_starts else msgs

# %% ../nbs/00_core.ipynb #1ddfa145
@patch
def preproc_msg(self:Agent, msgs:list|str, state=None):
    if isinstance(msgs, str): msgs = [{"content": msgs, "role": "user"}]
    if not self.include_hist: msgs = filter_out_history(msgs, active_agent=self.name)
    if self.narrative_cast:   msgs = mk_narrative_cast(msgs, active_agent=self.name)
    if self.sp:               msgs = [{"content": self.sp.format_map(state or {}), "role": "system"}, *msgs]
    return msgs

# %% ../nbs/00_core.ipynb #a8ea4152
def assistant_response(content:str): return ModelResponse(choices=[{"message": {"role": "assistant", "content": content}}])

def tool_response(name:str, arguments:dict):
    return ModelResponse(choices=[{"message": {"role": "assistant", "content": "", "tool_calls": [
            {"id": f"toolu_{uuid.uuid4().hex[:24]}", "type": "function", "function": {"name": name, "arguments": json.dumps(arguments)}}
        ]}, "finish_reason": "tool_calls"}])

def mk_cb_msg(msg:str|partial|ModelResponse):
    match msg:
        case str()           as o: return assistant_response(o)
        case partial()       as o: return tool_response(o.func.__name__, o.keywords)
        case ModelResponse() as o: return o

# %% ../nbs/00_core.ipynb #b6ca0d34
@patch
def run_callback(self:Agent, fn, msgs:list, state:dict): 
    if callable(fn): return mk_cb_msg(fn(state or {}, msgs))

@patch
def __call__(self:Agent, msgs:list|str, state=None, **kwargs):
    msgs = self.preproc_msg(msgs, state)
    if res := self.run_callback(self.before_cb, msgs, state): return res
    response = completion(model=self.model, messages=msgs, tools=self.tool_schemas, **kwargs)
    return self.run_callback(self.after_cb, response, state) or response 
