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

# %% auto 0
__all__ = ['empty', 'models', 'text_only_models', 'has_streaming_models', 'has_sp_models', 'has_temp_models', 'models_azure',
           'can_stream', 'can_set_sp', 'can_set_temp', 'usage', 'wrap_latex', 'Client', 'mk_openai_func',
           'mk_tool_choice', 'get_stream', 'call_func_openai', 'mk_toolres', 'Chat', 'mk_msg', 'mk_msgs', 'Response',
           'Responses', 'ResponseUsage', 'ResponseCompletedEvent', 'ResponseTextDeltaEvent', 'ResponseCreatedEvent',
           'ResponseInProgressEvent', 'ResponseOutputItemAddedEvent', 'ResponseContentPartAddedEvent',
           'ResponseTextDoneEvent', 'ResponseContentPartDoneEvent', 'ResponseOutputItemDoneEvent',
           'ResponseFunctionToolCall']

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

import inspect, typing, mimetypes, base64, json, msglm
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 toolslm.funccall import *

from openai import types
from openai import OpenAI,NOT_GIVEN,AzureOpenAI
from openai.resources import chat
from openai.types.responses.response import Response
from openai.resources.responses.responses import Responses
from openai.types.responses.response_usage import ResponseUsage

from openai.types.responses import (
    ResponseCompletedEvent, ResponseTextDeltaEvent, ResponseCreatedEvent, ResponseInProgressEvent,
    ResponseOutputItemAddedEvent, ResponseContentPartAddedEvent, ResponseTextDoneEvent, 
    ResponseContentPartDoneEvent, ResponseOutputItemDoneEvent, ResponseCompletedEvent,
    ResponseFunctionToolCall)

# %% ../00_core.ipynb
_all_ = ['mk_msg', 'mk_msgs', 'Response', 'Responses', 'ResponseUsage', 'ResponseCompletedEvent', 'ResponseTextDeltaEvent', 'ResponseCreatedEvent', 'ResponseInProgressEvent', 'ResponseOutputItemAddedEvent', 'ResponseContentPartAddedEvent', 'ResponseTextDoneEvent', 'ResponseContentPartDoneEvent', 'ResponseOutputItemDoneEvent', 'ResponseCompletedEvent', 'ResponseFunctionToolCall']

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

# %% ../00_core.ipynb
models = 'gpt-5.2', 'gpt-5.2-pro', 'gpt-5.2-chat-latest', 'gpt-5.1-codex', 'gpt-5-mini', 'gpt-5-nano', '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_sp_models = set(models) - set(('o1-mini', 'o3-mini'))
has_temp_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_sp(m): return m in has_sp_models
def can_set_temp(m): return m in has_temp_models

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

<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 `ResponseUsage`."
    return ResponseUsage(input_tokens=inp, output_tokens=out, total_tokens=inp+out, input_tokens_details={'cached_tokens':0}, output_tokens_details={'cached_tokens':0, 'reasoning_tokens':0})

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

# %% ../00_core.ipynb
@patch
def __add__(self:ResponseUsage, b):
    "Add together each of `input_tokens` and `output_tokens`"
    return usage(self.input_tokens+b.input_tokens, self.output_tokens+b.output_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, api_key_env=None, base_url=None):
        "Basic LLM messages client."
        self.model,self.use = model,usage(0,0)
        self.text_only = model in text_only_models
        if not cli:
            cli = OpenAI(api_key=os.getenv(api_key_env or "OPENAI_API_KEY"), base_url=base_url )
        self.c = cli.responses

# %% ../00_core.ipynb
@patch
def _r(self:Client, r):
    "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 mk_openai_func(f): 
    if isinstance(f, dict): return f
    sc = get_schema(f, 'parameters')
    if 'parameters' in sc: sc['parameters'].pop('title', None)
    return dict(type='function', **sc)

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

# %% ../00_core.ipynb
@save_iter
def get_stream(o, r, cli, cb=None):
    if not hasattr(o, 'events'): o.events = []
    for x in r:
        o.events.append(x)
        if isinstance(x, ResponseTextDeltaEvent): yield x.delta
        elif isinstance(x, ResponseCompletedEvent):
            o.value = x.response
            cli.use += x.response.usage
    if cb: cb(o.value)

# %% ../00_core.ipynb
@patch
@delegates(Responses.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?
             tools:Optional[list]=None, # List of tools to make available
             tool_choice:Optional[str]=None, # Forced tool choice
             cb:callable=None, # Callback after completion
             **kwargs):
    "Make a call to LLM."
    if tools: 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."
    tools = [mk_openai_func(o) for o in listify(tools)]
    r = self.c.create(
        model=self.model, input=msgs, max_output_tokens=maxtok, stream=stream, instructions=sp,
        tools=tools, tool_choice=mk_tool_choice(tool_choice), **kwargs)
    if stream: return get_stream(r, self, cb=cb)
    else:
        res = self._r(r)
        if cb: cb(res)
        return res

# %% ../00_core.ipynb
def call_func_openai(func, ns:Optional[abc.Mapping]=None):
    return call_func(func.name, json.loads(func.arguments), ns, raise_on_err=False)

# %% ../00_core.ipynb
def _toolres(r, ns):
    "Create a result dict from `tcs`."
    tcs = [o for o in getattr(r, 'output', []) if isinstance(o, ResponseFunctionToolCall)]
    if ns is None: ns = globals()
    return { tc.call_id: call_func_openai(tc, ns=mk_ns(ns)) for tc in tcs }

# %% ../00_core.ipynb
def mk_toolres(
    r:abc.Mapping, # Response containing tool use request
    ns:Optional[abc.Mapping]=None # Namespace to search for tools
    ):
    "Create a `tool_result` message from response `r`."
    tr = _toolres(r, ns)
    r = mk_msg(r)
    res = [r] if isinstance(r, dict) else listify(r)
    for k,v in tr.items(): res.append(dict(type="function_call_output", call_id=k, output=str(v)))
    return res

# %% ../00_core.ipynb
@patch
@delegates(Client.__call__)
def structured(self:Client,
               msgs: list, # Prompt
               tools:Optional[list]=None, # List of tools to make available to OpenAI model
               ns:Optional[abc.Mapping]=None, # Namespace to search for tools
               **kwargs):
    "Return the value of all tool calls (generally used for structured outputs)"
    if ns is None: ns = mk_ns(tools)
    r = self(msgs, tools=tools, tool_choice='required', **kwargs)
    return first(_toolres(r, ns).values())

# %% ../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
                 **kw):
        "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,self.kw = sp,tools,tool_choice,ns,kw
    
    @property
    def use(self): return self.c.use

# %% ../00_core.ipynb
@patch
@delegates(Responses.create)
def __call__(self:Chat,
             pr=None,  # Prompt / message
             stream:bool=False, # Stream response?
             tools=None, # Tools to use
             tool_choice=None, # Required tools to use
             **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 not tools: tools = self.tools
    if not tool_choice: tool_choice = self.tool_choice
    kw = self.kw | kwargs
    def _cb(v):
        self.last = mk_toolres(v, ns=self.ns)
        self.h += self.last
    res = self.c(self.h, sp=self.sp, stream=stream, cb=_cb, tools=tools, **kw)
    return res

# %% ../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'
