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

# %% auto 0
__all__ = ['empty', 'models', 'text_only_models', 'has_streaming_models', 'has_system_prompt_models', 'has_temperature_models',
           'models_azure', 'can_stream', 'can_set_system_prompt', 'can_set_temperature', 'find_block', 'contents',
           'usage', 'wrap_latex', 'Client', 'get_stream', 'mk_openai_func', 'mk_tool_choice', 'call_func_openai',
           'mk_toolres', 'mock_tooluse', 'Chat', 'mk_msg', 'mk_msgs']

# %% ../00_core.ipynb
from fastcore import imghdr
from fastcore.utils import *
from fastcore.meta import delegates

import inspect, typing, mimetypes, base64, json, ast
from collections import abc
from random import choices
from string import ascii_letters,digits

from msglm import mk_msg_openai as mk_msg, mk_msgs_openai as mk_msgs

from openai import types
from openai import Completion,OpenAI,NOT_GIVEN,AzureOpenAI
from openai.resources import chat
from openai.resources.chat import Completions
from openai.types.chat.chat_completion import ChatCompletion, ChatCompletionMessage
from openai.types.completion_usage import CompletionUsage

from toolslm.funccall import *

# %% ../00_core.ipynb
_all_ = ['mk_msg', 'mk_msgs']

# %% ../00_core.ipynb
empty = inspect.Parameter.empty

# %% ../00_core.ipynb
models = 'o1-preview', 'o1-mini', 'gpt-4o', 'gpt-4o-mini', 'gpt-4-turbo', 'gpt-4', 'gpt-4-32k', 'gpt-3.5-turbo', 'gpt-3.5-turbo-instruct', 'o1', 'o3-mini', 'chatgpt-4o-latest', 'o1-pro', 'o3', 'o4-mini', 'gpt-4.1', 'gpt-4.1-mini', 'gpt-4.1-nano'

# %% ../00_core.ipynb
text_only_models = 'o1-preview', 'o1-mini', 'o3-mini'

# %% ../00_core.ipynb
has_streaming_models = set(models) - set(('o1-mini', 'o3-mini'))
has_system_prompt_models = set(models) - set(('o1-mini', 'o3-mini'))
has_temperature_models = set(models) - set(('o1', 'o1-mini', 'o3-mini'))

# %% ../00_core.ipynb
def can_stream(m): return m in has_streaming_models
def can_set_system_prompt(m): return m in has_system_prompt_models
def can_set_temperature(m): return m in has_temperature_models

# %% ../00_core.ipynb
def find_block(r:abc.Mapping, # The message to look in
              ):
    "Find the message in `r`."
    m = nested_idx(r, 'choices', 0)
    if not m: return m
    if hasattr(m, 'message'): return m.message
    return m.delta

# %% ../00_core.ipynb
def contents(r):
    "Helper to get the contents from response `r`."
    blk = find_block(r)
    if not blk: return r
    if hasattr(blk, 'content'): return getattr(blk,'content')
    return blk

# %% ../00_core.ipynb
@patch
def _repr_markdown_(self:ChatCompletion):
    det = '\n- '.join(f'{k}: {v}' for k,v in dict(self).items())
    res = contents(self)
    if not res: return f"- {det}"
    return f"""{contents(self)}

<details>

- {det}

</details>"""

# %% ../00_core.ipynb
def usage(inp=0, # Number of prompt tokens
          out=0  # Number of completion tokens
         ):
    "Slightly more concise version of `CompletionUsage`."
    return CompletionUsage(prompt_tokens=inp, completion_tokens=out, total_tokens=inp+out)

# %% ../00_core.ipynb
@patch
def __repr__(self:CompletionUsage): return f'In: {self.prompt_tokens}; Out: {self.completion_tokens}; Total: {self.total_tokens}'

# %% ../00_core.ipynb
@patch
def __add__(self:CompletionUsage, b):
    "Add together each of `input_tokens` and `output_tokens`"
    return usage(self.prompt_tokens+b.prompt_tokens, self.completion_tokens+b.completion_tokens)

# %% ../00_core.ipynb
def wrap_latex(text):
    "Replace OpenAI LaTeX codes with markdown-compatible ones"
    text = re.sub(r"\\\((.*?)\\\)", lambda o: f"${o.group(1)}$", text)
    res = re.sub(r"\\\[(.*?)\\\]", lambda o: f"$${o.group(1)}$$", text, flags=re.DOTALL)
    return res

# %% ../00_core.ipynb
class Client:
    def __init__(self, model, cli=None):
        "Basic LLM messages client."
        self.model,self.use = model,usage(0,0)
        self.text_only = model in text_only_models
        self.c = (cli or OpenAI()).chat.completions

# %% ../00_core.ipynb
@patch
def _r(self:Client, r:ChatCompletion):
    "Store the result of the message and accrue total usage."
    self.result = r
    if getattr(r,'usage',None): self.use += r.usage
    return r

# %% ../00_core.ipynb
def get_stream(r):
    for o in r:
        o = contents(o)
        if o and isinstance(o, str): yield(o)

# %% ../00_core.ipynb
@patch
@delegates(Completions.create)
def __call__(self:Client,
             msgs:list, # List of messages in the dialog
             sp:str='', # System prompt
             maxtok=4096, # Maximum tokens
             stream:bool=False, # Stream response?
             **kwargs):
    "Make a call to LLM."
    if 'tools' in kwargs: assert not self.text_only, "Tool use is not supported by the current model type."
    if any(c['type'] == 'image_url' for msg in msgs if isinstance(msg, dict) and isinstance(msg.get('content'), list) for c in msg['content']): assert not self.text_only, "Images are not supported by the current model type."
    if stream: kwargs['stream_options'] = {"include_usage": True}
    if sp and self.model in has_system_prompt_models:
        msgs = [mk_msg(sp, 'system')] + list(msgs)

    r = self.c.create(
        model=self.model, messages=msgs, max_completion_tokens=maxtok, stream=stream, **kwargs)
    if not stream: return self._r(r)
    else: return get_stream(map(self._r, r))

# %% ../00_core.ipynb
def mk_openai_func(f): 
    sc = get_schema(f, 'parameters')
    if 'parameters' in sc: sc['parameters'].pop('title', None)
    return dict(type='function', function=sc)

# %% ../00_core.ipynb
def mk_tool_choice(f): return dict(type='function', function={'name':f})

# %% ../00_core.ipynb
def call_func_openai(func:types.chat.chat_completion_message_tool_call.Function, ns:Optional[abc.Mapping]=None):
    return call_func(func.name, ast.literal_eval(func.arguments), ns, raise_on_err=False)

# %% ../00_core.ipynb
def mk_toolres(
    r:abc.Mapping, # Tool use request response
    ns:Optional[abc.Mapping]=None, # Namespace to search for tools
    obj:Optional=None # Class to search for tools
    ):
    "Create a `tool_result` message from response `r`."
    r = mk_msg(r)
    tcs = getattr(r, 'tool_calls', [])
    res = [r]
    if ns is None: ns = globals()
    if obj is not None: ns = mk_ns(obj)
    for tc in (tcs or []):
        func = tc.function
        cts = str(call_func_openai(func, ns=ns))
        res.append(mk_msg(str(cts), 'tool', tool_call_id=tc.id, name=func.name))
    return res

# %% ../00_core.ipynb
def _mock_id(): return 'call_' + ''.join(choices(ascii_letters+digits, k=24))

def mock_tooluse(name:str, # The name of the called function
                 res,  # The result of calling the function
                 **kwargs): # The arguments to the function
    ""
    id = _mock_id()
    func = dict(arguments=json.dumps(kwargs), name=name)
    tc = dict(id=id, function=func, type='function')
    req = dict(content=None, role='assistant', tool_calls=[tc])
    resp = mk_msg('' if res is None else str(res), 'tool', tool_call_id=id, name=name)
    return [req,resp]

# %% ../00_core.ipynb
class Chat:
    def __init__(self,
                 model:Optional[str]=None, # Model to use (leave empty if passing `cli`)
                 cli:Optional[Client]=None, # Client to use (leave empty if passing `model`)
                 sp='', # Optional system prompt
                 tools:Optional[list]=None, # List of tools to make available
                 hist: list = None,  # Initialize history
                 tool_choice:Optional[str]=None, # Forced tool choice
                 ns:Optional[abc.Mapping]=None): # Namespace to search for tools
        "OpenAI chat client."
        assert model or cli
        self.c = (cli or Client(model))
        self.h = hist if hist else []
        if ns is None: ns=tools
        self.sp,self.tools,self.tool_choice,self.ns = sp,listify(tools),tool_choice,ns
    
    @property
    def use(self): return self.c.use

# %% ../00_core.ipynb
@patch
@delegates(Completions.create)
def __call__(self:Chat,
             pr=None,  # Prompt / message
             stream:bool=False, # Stream response?
             **kwargs):
    "Add prompt `pr` to dialog and get a response"
    if isinstance(pr,str): pr = pr.strip()
    if pr: self.h.append(mk_msg(pr))
    if self.tools: kwargs['tools'] = [mk_openai_func(o) for o in self.tools]
    if self.tool_choice: kwargs['tool_choice'] = mk_tool_choice(self.tool_choice)
    res = self.c(self.h, sp=self.sp, stream=stream, **kwargs)
    self.last = mk_toolres(res, ns=self.ns)
    self.h += self.last
    return res

# %% ../00_core.ipynb
@patch(as_prop=True)
def structured(self:Chat): return self.h[-1]['content']

# %% ../00_core.ipynb
models_azure = 'o1-preview', 'o1-mini', 'gpt-4o', 'gpt-4o-mini', 'gpt-4-turbo', 'gpt-4', 'gpt-4-32k', 'gpt-3.5-turbo', 'gpt-3.5-turbo-instruct', 'o1', 'o3-mini', 'chatgpt-4o-latest', 'o1-pro', 'o3', 'o4-mini', 'gpt-4.1', 'gpt-4.1-mini', 'gpt-4.1-nano'
