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

# %% auto 0
__all__ = ['MISTRAL_API_KEY', 'model_types', 'all_models', 'vision_models', 'embed_models', 'ocr_models', 'tool_models',
           'text_only_models', 'has_streaming_models', 'has_system_prompt_models', 'has_temperature_models', 'models',
           'can_stream', 'can_set_system_prompt', 'can_set_temperature', 'can_use_image', 'find_block', 'contents',
           'usage', 'Client', 'get_stream', 'mk_mistralai_func', 'mk_tool_choice', 'call_func_mistralai', 'mk_toolres',
           'mock_tooluse', 'Chat', 'mk_msg', 'mk_msgs']

# %% ../00_core.ipynb 3
import os
from collections import abc
from fastcore.utils import *
from fastcore.meta import delegates
import ast, json, base64
from random import choices
from string import ascii_letters,digits

# Mistral AI messages are compatible with OpenAI's messages
from msglm import mk_msg_openai as mk_msg, mk_msgs_openai as mk_msgs
from toolslm.funccall import *

import mistralai
from mistralai import Mistral
from mistralai.models import ChatCompletionChoice, ChatCompletionResponse, UsageInfo, CompletionEvent
from mistralai.models.functioncall import FunctionCall
from mistralai.types import BaseModel

from typing import Optional

try: from IPython import display
except: display=None

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

# %% ../00_core.ipynb 5
MISTRAL_API_KEY = os.environ.get("MISTRAL_API_KEY")

# %% ../00_core.ipynb 6
model_types = {
    # Premier models
    'codestral-2501': 'codestral-latest', # code generation model
    'mistral-large-2411': 'mistral-large-latest', # top-tier reasoning model for high-complexity tasks
    'pixtral-large-2411': 'pixtral-large-latest', # frontier-class multimodal model
    'mistral-saba-2502': 'mistral-saba-latest', # model for languages from the Middle East and South Asia
    'ministral-3b-2410': 'ministral-3b-latest', # edge model
    'ministral-8b-2410': 'ministral-8b-latest', # edge model with high performance/price ratio
    'mistral-embed-2312': 'mistral-embed', # embedding model
    'mistral-moderation-2411': 'mistral-moderation-latest', # moderation service to detect harmful text content
    'mistral-ocr-2503': 'mistral-ocr-latest', # OCR model to extract interleaved text and images
    
    # Free models (with weight availability)
    'mistral-small-2503': 'mistral-small-latest', # small model with image understanding capabilities
    
    # Research models
    'open-mistral-nemo-2407': 'open-mistral-nemo', # multilingual open source model
}

# %% ../00_core.ipynb 7
all_models = list(model_types)

# %% ../00_core.ipynb 8
vision_models = ['pixtral-large-2411', 'mistral-small-2503', 'mistral-ocr-2503']

# %% ../00_core.ipynb 9
embed_models = ['mistral-embed-2312']

# %% ../00_core.ipynb 10
ocr_models = ['mistral-ocr-2503']

# %% ../00_core.ipynb 11
tool_models = ['codestral-2501', 'mistral-large-2411', 
               'mistral-small-2503', 'ministral-3b-2410', 'ministral-8b-2410',
               'pixtral-large-2411', 'open-mistral-nemo-2407']

# %% ../00_core.ipynb 12
text_only_models = set(all_models) - set(vision_models) - set(embed_models) - set(ocr_models)

# %% ../00_core.ipynb 13
has_streaming_models = set(all_models) - set(embed_models) - set(ocr_models)
has_system_prompt_models = set(all_models) - set(embed_models) - set(ocr_models)
has_temperature_models = set(all_models) - set(embed_models) - set(ocr_models)

# %% ../00_core.ipynb 14
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
def can_use_image(m): return m in vision_models

# %% ../00_core.ipynb 16
models = all_models

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

# %% ../00_core.ipynb 32
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 34
@patch
def _repr_markdown_(self:ChatCompletionResponse):
    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 37
def usage(inp=0, # input tokens
          out=0,  # Output tokens
         ):
    "Slightly more concise version of `UsageInfo`."
    return UsageInfo(prompt_tokens=inp, completion_tokens=out, total_tokens=inp+out)

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

# %% ../00_core.ipynb 41
@patch
def __add__(self:UsageInfo, 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 44
@patch(as_prop=True)
def total(self:UsageInfo): return self.total_tokens

# %% ../00_core.ipynb 57
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.can_use_tools = model in tool_models
        self.c = (cli or Mistral(api_key=os.environ.get("MISTRAL_API_KEY")))

# %% ../00_core.ipynb 60
@patch
def _r(self:Client, r:ChatCompletionResponse):
    "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 62
def get_stream(r):
    for o in r:
        o = contents(o)
        if o and isinstance(o, str): yield(o)

# %% ../00_core.ipynb 64
@patch
@delegates(mistralai.Chat.complete)
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 self.can_use_tools, "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 sp and self.model in has_system_prompt_models: msgs = [mk_msg(sp, 'system')] + list(msgs)
    chat_args = dict(model=self.model, messages=msgs, max_tokens=maxtok, **kwargs)
    r = self.c.chat.stream(**chat_args) if stream else self.c.chat.complete(**chat_args)
    return self._r(r) if not stream else get_stream(map(self._r, r))

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

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

# %% ../00_core.ipynb 84
def call_func_mistralai(func:FunctionCall, ns:Optional[abc.Mapping]=None): 
    return call_func(func.name, ast.literal_eval(func.arguments), ns)

# %% ../00_core.ipynb 88
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_mistralai(func, ns=ns))
        res.append(mk_msg(str(cts), 'tool', tool_call_id=tc.id, name=func.name))
    return res

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

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 103
@patch
@delegates(Client.__call__)
def structured(self:Client,
               msgs: list, # Prompt
               tools:Optional[list]=None, # List of tools to make available to OpenAI model
               obj:Optional=None, # Class to search for tools
               ns:Optional[abc.Mapping]=None, # Namespace to search for tools
               **kwargs):
    "Return the value of all tool calls (generally used for structured outputs)"
    tools = listify(tools)
    if ns is None: ns=mk_ns(*tools)
    tools = [mk_mistralai_func(o) for o in tools]
    if obj is not None: ns = mk_ns(obj)
    res = self(msgs, tools=tools, tool_choice='required', **kwargs)
    cts = getattr(res, 'choices', [])
    tcs = [call_func_mistralai(t.function, ns=ns) for o in cts for t in (o.message.tool_calls or [])]
    return tcs

# %% ../00_core.ipynb 107
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
                 tool_choice:Optional[str]=None): # Forced tool choice
        "OpenAI chat client."
        assert model or cli
        self.c = (cli or Client(model))
        self.h,self.sp,self.tools,self.tool_choice = [],sp,tools,tool_choice
    
    @property
    def use(self): return self.c.use

# %% ../00_core.ipynb 109
@patch
@delegates(mistralai.Chat.complete)
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_mistralai_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.h += mk_toolres(res, ns=self.tools)
    return res
