From 25c3f5dc25717028243cd7a2332a8f1e467a85c2 Mon Sep 17 00:00:00 2001
From: Pabloader <pabloader@pabloader.ru>
Date: Sun, 3 Nov 2024 16:48:00 +0000
Subject: [PATCH] AI story: tool calling check

---
 bun.lockb                                 | Bin 31377 -> 31409 bytes
 package.json                              |   1 +
 src/games/ai/components/header/header.tsx |   4 +-
 src/games/ai/contexts/llm.tsx             |  55 +++---
 src/games/ai/contexts/state.tsx           |   6 +-
 src/games/ai/huggingface.ts               | 226 ++++++++++++++++++++--
 6 files changed, 241 insertions(+), 51 deletions(-)

diff --git a/bun.lockb b/bun.lockb
index 597e5c52cb15522c32ca7bfe92a904b08f5b46ee..d9cef444a408f13ab3f110010778c6c78fb73765 100755
GIT binary patch
delta 2564
zcmeHJO-x)>6u$S-8RiW`ff@>LhDs<{3fKV#D9|$Vm}=8dS`cvfwft4V!A*;@s7;tL
zja6$)phq`NH+5ASjGBq2){3B}YC@An)3`8%bYn797mZ>~TaDj&?~SscHZiV@C;9Go
z&(A&g|Gqc1BHvq)lT9*jsyP_BB7|@X5xs2j3^*P7Z7D<=_$Jr|KB)+yf)}B$2RoC5
zNC7WE&j8O__JVb0oEyw}FIhYa<~)ryA+o^rV6=}!hvDEvm3ARg!N<VdFwL^t!QAm3
zbjS&o!92hvYXJ9>P5V{o8(=?!2{50qCZeJ->qnsTKo5aA-w@a-ViVFQ<x*^3DNK^r
zC~nggyQE@!K<=WjU6&D>w!@^*GEDoZC|Q?pQW)hmnnsyPag=RT<j^Hi809vaMmbJ#
zhpv8(q1T}+8{KkLE5b%2YCs*t#M@y^qFc$;>P2V)Xb#roFI42z)gV@27i<bQlc#9f
zsmre@?$m9`nClMu(YaFwDV(CKFF_9`a>-9Ao}w!b2Nk;lvVp=bT~5$6%BvJd>4_Dk
z-gV_+4R{?4S&MBHa_2^{n<CqaTFuxY@d$v$7Fz+1D*{-U+5RXikk!m_+W=-S!16at
zKe_WF98(5xn{tbNV3xHQiy>+augTnQ2f)13vaiML$2J+ljFZlLV)V}yLzvl758!x;
zy#UL<$v7^e3E)L+5p5`q%7d9_&}xPm%maPW(pR%W2XeBEz3sBR%-r2m0B1P>urPC$
zK7j500PBMQ>nsmENr>5KM4HU)aP{dqxBCD2`tP1EKkxHDK3}og+;@DLiHGS4q~dgZ
z^xX7Y%h3ir{l=l>4!VIJfRA_x;48qlrUl^psssE`Q~;Gg6|fzs2Fhr@^`X2>)cSz|
zfRB79z&E-P*az&VmDa+H2Gsar@B?ljhsxWscITtU$GZdITiXoqbHmRfKapDiegcbu
z5`dpS4^RrcYhCU-dabQl=EknK9hB-Jo&#i@f}xVwvmu9U#5zQKzWp`zv4iKoP<Rmx
zIe`-E5Zyg9J2HRTTd2q~Z+Wd3Z{-<yHo!BHKJfj~O&`3)o(gX{tTS}3U6YsTBa~O@
zdV7T@@x{InisoO~Q~I(Zdy&J71zVs`IyAXRev~&U*s00ebQtxM^m>P;IFo3yGh40-
z@+7{!O?!_2ob~YfR-`IJDy;burFLqb#5a5B_|&pz<>N8<E8xd~&GdJur&G&GeCap7
z(dBynl@}HeRp~9~D6`9|#4o{YYxpze_irZu3P}8IY}$Mye@@-}Aw1jB0VZBRi+_%P
zm@>LFIZ7p6n&&<((Ag#G>(Z+3<3~p8&n5c9)7_!Py1B>ESB*KDMtU2W6F)M8hYS4m
zDQPwE`^>Q%qZ{b|9^FN`L@B+RqN-HZovj=ZvAx}MNv6>Zpg`|g=->73_ay#wrb@d{
zZF%2ctH^O5=H~O(7)jJG`H_aaO3(BK>3rXu^J2eY5Za40+CQ^yczA3mmND?PgEGTg
KV>u&Z>3;wr_+9A$

delta 2666
zcmeHJT}&KR6ux(spPhxaR75%~SW9aVkY#uI+d|8fDlxbM0`e<aSbu;e+Mq2a4OYOI
zHpHZr9!*RzG3kqLThwM!qqbVGnixY+A2s%&F>SJ~rfI@jo2J$8%+5p}mc$1idz0^e
zd(OG%%>6m{>^&8ivvK)Wi<Gmi4U?yZ5H2C6PFOq#&Vqbd3b6tF3pf)zp$MUZmmseP
z--Mh2J_C6pc;3=yrzFOyV2&HLI1J`EwW&hng3H11pPcG~fdg$#6Jj0sWiUJ5aU}F{
zFn>4(<^Vr~d4kK9{c3tbe-`pR(9d82%=@i{{N+r@K?LE6jzQ1yhQKai9F~45myJcG
z$RW>BELB(Hf=bfrrH>+My4*w2G^iw9fvS#*({=d<MbMw4DEbtNq2Eo#PF=oE5%l>K
zMgJnjoVsdYX*=+#RJ!i0QLd%Xs9LYKVc}g+I_P?OjrswkdPq)|<RvP0>8ckyup2ss
z-Q=qjb?Ndmin;XEKVhk*KU{UvOOXs+9mCoi>{#*?#WHl|Mmm*b*2@4zGIe=?qUg_3
z486-JUUw_ghrQ$dGg~soghY;IXO`5Mw|SBkI5s{Zu?@gxiy{E0E*=87F|)pu1+>-7
z;~oZ>%K&cwg?YS32>LB=%OwByu_IbC#%73sX-zPX-wrU>TKZ(n_BzYXjLa50EqN7B
zng$&Pq^D@I2K=9lbP>$}Z{m|y{D0XzgJg!p3fn9PW?n&uC9h`IcLKa(<e<sj7WaU;
zG4lkT26$XQz^(r&bQ~ZAuwlTG!<Gy(6DwI#*|xjMbvL<^Cf6opNYdmowr~E}pJV3~
z-v-V-9;(spwuvUrMo!sYpby}5=S<~m&<?N%Ut$zlvuOCi6+k6W1ylp&Ko%{wXK&%0
z9t1e)`C{<xejEq_PXIc_+lz9zNT99@FQ5T#pn%FcavQi{^2Op?+yZd<;WEi3^Fe@1
zVlhwxa5>BaP~OBR)-4UtOh>EnYsa8eNBE+m^;2VKsj;s!MFz18Q;jFT8eKf<DN=X~
z%r=9b|L_iHg)dJOe0ofgUQdOG3*;QYKDoxNkNPGasZ^w&EtvKRy3nP`({!V&Le5iZ
zw<gcfMf6KF(5+S3|M|5;$M;fEXjGA*N;qQc>{mg<^vo+OV{e{^#qaU4WrF6rHMjjL
zC_OSK<_q6pyU*kG;ESu`yY0W~md6gR<leu&9iuSDgCkj{@*b_q{)-<vID5q%|8gAm
zibRyPXBMLaJzBN>Sh(-)-psv+_ASA)67$3N*PhuO2QBq%k}A#hYBHPt#st=0ru|xJ
zDgXMI_Wk*L%@3Fh3{!Kj=C+?b1%(&4E~tf{!(N@3dX%OS-~5W(di+fN{#;+@g&hlz
zVpJe8s+NA})vD~b&@(S>+qo-aLjZO^>^M{Vb@Wru@gJL3-oK>C*L+wYw(cqw^l5UL
z%0gPc{pR^1RP&Agz_lp+;qCYM13}ezp>M&F7$;318VtP$`F5xoeQ>~vp?=7B0d0g1
z4KyIx=(~Xh*SmwJPTbs2M+Z+Dt>JH-8+VP24-ZH73=f5$eZG8nczlRtB;W9kPG|iC
DcPxB^

diff --git a/package.json b/package.json
index 4960a8d..b2190b0 100644
--- a/package.json
+++ b/package.json
@@ -8,6 +8,7 @@
     "bake": "bun build/build.ts"
   },
   "dependencies": {
+    "@huggingface/gguf": "0.1.12",
     "@huggingface/hub": "0.19.0",
     "@huggingface/jinja": "0.3.1",
     "@inquirer/select": "2.3.10",
diff --git a/src/games/ai/components/header/header.tsx b/src/games/ai/components/header/header.tsx
index 8031bc8..bf60fda 100644
--- a/src/games/ai/components/header/header.tsx
+++ b/src/games/ai/components/header/header.tsx
@@ -60,8 +60,8 @@ export const Header = () => {
                     onBlur={handleBlurUrl}
                     class={blockConnection.value ? '' : urlValid ? styles.valid : styles.invalid}
                 />
-                <select value={instruct} onChange={setInstruct}>
-                    {modelName && modelTemplate && <option value={modelTemplate}>{modelName}</option>}
+                <select value={instruct} onChange={setInstruct} title='Instruct template'>
+                    {modelName && modelTemplate && <option value={modelTemplate} title='Native for model'>{modelName}</option>}
                     {Object.entries(Instruct).map(([label, value]) => (
                         <option value={value} key={value}>
                             {label.toLowerCase()}
diff --git a/src/games/ai/contexts/llm.tsx b/src/games/ai/contexts/llm.tsx
index c8b6852..eb9c92e 100644
--- a/src/games/ai/contexts/llm.tsx
+++ b/src/games/ai/contexts/llm.tsx
@@ -3,17 +3,11 @@ import SSE from "@common/sse";
 import { createContext } from "preact";
 import { useCallback, useContext, useEffect, useMemo, useState } from "preact/hooks";
 import { MessageTools, type IMessage } from "../messages";
-import { StateContext } from "./state";
+import { Instruct, StateContext } from "./state";
 import { useBool } from "@common/hooks/useBool";
 import { Template } from "@huggingface/jinja";
 import { Huggingface } from "../huggingface";
 
-
-interface ITemplateMessage {
-    role: 'user' | 'assistant' | 'system';
-    content: string;
-}
-
 interface ICompileArgs {
     keepUsers?: number;
 }
@@ -29,6 +23,7 @@ interface IContext {
     blockConnection: ReturnType<typeof useBool>;
     modelName: string;
     modelTemplate: string;
+    hasToolCalls: boolean;
     promptTokens: number;
     contextLength: number;
 }
@@ -65,7 +60,7 @@ export const normalizeModel = (model: string) => {
         normalizedModel = currentModel;
 
         currentModel = currentModel
-            .replace(/[ ._-]\d(\d*k|\d+)(-context|$)/i, '') // remove context length, i.e. -32k
+            .replace(/[ ._-]\d+(k$|-context)/i, '') // remove context length, i.e. -32k
             .replace(/[ ._-](gptq|awq|exl2?|imat|i\d)/i, '') // remove quant name
             .replace(/([ ._-]?gg(uf|ml)[ ._-]?(v[ ._-]?\d)?)/i, '') // remove gguf-v3/ggml/etc
             .replace(/[ ._-]i?q([ ._-]?\d[ ._-]?(k?[ ._-]?x*[ ._-]?[lms]?)?)+/i, '') // remove quant size
@@ -82,19 +77,6 @@ export const normalizeModel = (model: string) => {
         .trim();
 }
 
-export const applyChatTemplate = (messages: ITemplateMessage[], templateString: string) => {
-    const template = new Template(templateString);
-
-    console.log(`Applying template:\n${templateString}`, messages);
-
-    const prompt = template.render({
-        messages,
-        add_generation_prompt: true,
-    });
-
-    return prompt;
-};
-
 export const LLMContext = createContext<ILLMContext>({} as ILLMContext);
 
 export const LLMContextProvider = ({ children }: { children?: any }) => {
@@ -109,6 +91,7 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
     const [contextLength, setContextLength] = useState(0);
     const [modelName, setModelName] = useState('');
     const [modelTemplate, setModelTemplate] = useState('');
+    const [hasToolCalls, setHasToolCalls] = useState(false);
 
     const userPromptTemplate = useMemo(() => {
         try {
@@ -162,17 +145,17 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
             const isRegen = isAssistantLast && !MessageTools.getSwipe(lastMessage)?.content;
             const isContinue = isAssistantLast && !isRegen;
 
-            const userMessages = promptMessages.filter(m => m.role === 'user');
-            const lastUserMessage = userMessages.at(-1);
-            const firstUserMessage = userMessages.at(0);
-
             if (isContinue) {
                 promptMessages.push(MessageTools.create(userPromptTemplate.render({})));
             }
 
+            const userMessages = promptMessages.filter(m => m.role === 'user');
+            const lastUserMessage = userMessages.at(-1);
+            const firstUserMessage = userMessages.at(0);
+
             const system = `${systemPrompt}\n\n${lore}`.trim();
 
-            const templateMessages: ITemplateMessage[] = [
+            const templateMessages: Huggingface.ITemplateMessage[] = [
                 { role: 'system', content: system },
             ];
 
@@ -229,7 +212,7 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
                 });
             }
 
-            const prompt = applyChatTemplate(templateMessages, instruct);
+            const prompt = Huggingface.applyChatTemplate(instruct, templateMessages);
             return {
                 prompt,
                 isContinue,
@@ -342,16 +325,14 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
 
             for await (const chunk of actions.generate(prompt)) {
                 text += chunk;
-                setPromptTokens(tokens + 1);
+                setPromptTokens(tokens + Math.round(text.length * 0.25));
                 editMessage(messageId, text);
             }
 
             text = MessageTools.trimSentence(text);
             editMessage(messageId, text);
 
-            const generatedTokens = await actions.countTokens(text);
-
-            setPromptTokens(tokens + generatedTokens);
+            setPromptTokens(0); // trigger calculation
 
             MessageTools.playReady();
         }
@@ -381,6 +362,8 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
                     if (template) {
                         setModelTemplate(template);
                         setInstruct(template);
+                    } else {
+                        setInstruct(Instruct.CHATML);
                     }
                 });
         }
@@ -395,11 +378,21 @@ export const LLMContextProvider = ({ children }: { children?: any }) => {
         }
     }, [actions, promptTokens, messages, blockConnection.value]);
 
+    useEffect(() => {
+        try {
+            const hasTools = Huggingface.testToolCalls(instruct);
+            setHasToolCalls(hasTools);
+        } catch {
+            setHasToolCalls(false);
+        }
+    }, [instruct]);
+
     const rawContext: IContext = {
         generating: generating.value,
         blockConnection,
         modelName,
         modelTemplate,
+        hasToolCalls,
         promptTokens,
         contextLength,
     };
diff --git a/src/games/ai/contexts/state.tsx b/src/games/ai/contexts/state.tsx
index 98d7809..0945e94 100644
--- a/src/games/ai/contexts/state.tsx
+++ b/src/games/ai/contexts/state.tsx
@@ -38,12 +38,12 @@ interface IActions {
 const SAVE_KEY = 'ai_game_save_state';
 
 export enum Instruct {
+    CHATML = `{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n\n' }}{% endif %}`,
+
     LLAMA = `{% for message in messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] | trim + '<|eot_id|>' %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>\n\n' }}{% endif %}`,
 
     MISTRAL = `{%- if messages[0]['role'] == 'system' %}{%- set system_message = messages[0]['content'] %}{%- set loop_messages = messages[1:] %}{%- else %}{%- set loop_messages = messages %}{%- endif %}{%- for message in loop_messages %}{%- if message['role'] == 'user' %}{%- if loop.first and system_message is defined %}{{- ' [INST] ' + system_message + '\n\n' + message['content'] + ' [/INST]' }}{%- else %}{{- ' [INST] ' + message['content'] + ' [/INST]' }}{%- endif %}{%- elif message['role'] == 'assistant' %}{{- ' ' + message['content'] + '</s>'}}{%- endif %}{%- endfor %}`,
 
-    CHATML = `{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n\n' }}{% endif %}`,
-
     ALPACA = `{% for message in messages %}{% if message['role'] == 'system' and message['content'] %}{{ message['content'] + '\n\n'}}{% elif message['role'] == 'user' %}{{'### Instruction:\n\n' + message['content'] + '\n\n'}}{% elif message['role'] == 'assistant' %}{{'### Response:\n\n' + message['content'] + '\n\n'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '### Response:\n\n' }}{% endif %}`,
 };
 
@@ -58,7 +58,7 @@ export const loadContext = (): IContext => {
     const defaultContext: IContext = {
         connectionUrl: 'http://localhost:5001',
         input: '',
-        instruct: Instruct.LLAMA,
+        instruct: Instruct.CHATML,
         systemPrompt: 'You are creative writer. Write a story based on the world description below.',
         lore: '',
         userPrompt: `{% if prompt %}{% if isStart %}Start{% else %}Continue{% endif %} this story, taking information into account: {{ prompt | trim }}
diff --git a/src/games/ai/huggingface.ts b/src/games/ai/huggingface.ts
index e083908..1666297 100644
--- a/src/games/ai/huggingface.ts
+++ b/src/games/ai/huggingface.ts
@@ -1,27 +1,104 @@
+import { gguf } from '@huggingface/gguf';
 import * as hub from '@huggingface/hub';
+import { Template } from '@huggingface/jinja';
 
 export namespace Huggingface {
+    export interface ITemplateMessage {
+        role: 'user' | 'assistant' | 'system';
+        content: string;
+    }
+
+    interface INumberParameter {
+        type: 'number';
+        enum?: number[];
+        description?: string;
+    }
+
+    interface IStringParameter {
+        type: 'string';
+        enum?: string[];
+        description?: string;
+    }
+
+    interface IArrayParameter {
+        type: 'array';
+        description?: string;
+        items: IParameter;
+    }
+
+    interface IObjectParameter {
+        type: 'object';
+        description?: string;
+        properties: Record<string, IParameter>;
+        required?: string[];
+    }
+
+    type IParameter = INumberParameter | IStringParameter | IArrayParameter | IObjectParameter;
+
+    interface ITool {
+        type: 'function',
+        function: {
+            name: string;
+            description?: string;
+            parameters?: IObjectParameter;
+        }
+    }
+
+    export interface IFunction {
+        name: string;
+        description?: string;
+        parameters?: Record<string, IParameter>;
+    }
+
     interface TokenizerConfig {
         chat_template: string;
-        eos_token: string;
+        bos_token?: string;
+        eos_token?: string;
     }
 
+    const TEMPLATE_CACHE_KEY = 'ai_game_template_cache';
+
+    const loadCache = (): Record<string, string> => {
+        const json = localStorage.getItem(TEMPLATE_CACHE_KEY);
+
+        try {
+            if (json) {
+                const cache = JSON.parse(json);
+                if (cache && typeof cache === 'object') {
+                    return cache
+                }
+            }
+        } catch { }
+
+        return {};
+    };
+
+    const saveCache = (cache: Record<string, string>) => {
+        const json = JSON.stringify(cache);
+        localStorage.setItem(TEMPLATE_CACHE_KEY, json);
+    };
+
+    const templateCache: Record<string, string> = loadCache();
+
     const hasField = <T extends string>(obj: unknown, field: T): obj is Record<T, unknown> => (
         obj != null && typeof obj === 'object' && (field in obj)
     );
     const isTokenizerConfig = (obj: unknown): obj is TokenizerConfig => (
         hasField(obj, 'chat_template') && typeof obj.chat_template === 'string'
-        && hasField(obj, 'eos_token') && typeof obj.eos_token === 'string'
+        && (!hasField(obj, 'eos_token') || !obj.eos_token || typeof obj.eos_token === 'string')
+        && (!hasField(obj, 'bos_token') || !obj.bos_token || typeof obj.bos_token === 'string')
     );
 
     const loadHuggingfaceTokenizerConfig = async (model: string): Promise<TokenizerConfig | null> => {
-        console.log(`Searching for model '${model}'`);
+        console.log(`[huggingface] searching config for '${model}'`);
 
         const models = hub.listModels({ search: { query: model }, additionalFields: ['config'] });
+        const recheckModels: hub.ModelEntry[] = [];
 
         let tokenizerConfig: TokenizerConfig | null = null;
 
         for await (const model of models) {
+            recheckModels.push(model);
             const { config } = model;
             if (hasField(config, 'tokenizer_config') && isTokenizerConfig(config.tokenizer_config)) {
                 tokenizerConfig = config.tokenizer_config;
@@ -29,6 +106,7 @@ export namespace Huggingface {
             }
 
             try {
+                console.log(`[huggingface] searching config in '${model.name}/tokenizer_config.json'`);
                 const fileResponse = await hub.downloadFile({ repo: model.name, path: 'tokenizer_config.json' });
                 if (fileResponse?.ok) {
                     const maybeConfig = await fileResponse.json();
@@ -38,31 +116,149 @@ export namespace Huggingface {
                     }
                 }
             } catch { }
+
+        }
+        if (!tokenizerConfig) {
+            for (const model of recheckModels.slice(0, 10)) {
+                try {
+                    for await (const file of hub.listFiles({ repo: model.name, recursive: true })) {
+                        if (file.type !== 'file' || !file.path.endsWith('.gguf')) continue;
+                        try {
+                            console.log(`[huggingface] searching config in '${model.name}/${file.path}'`);
+                            const fileInfo = await hub.fileDownloadInfo({ repo: model.name, path: file.path });
+                            if (fileInfo?.downloadLink) {
+                                const { metadata } = await gguf(fileInfo.downloadLink);
+                                if ('tokenizer.chat_template' in metadata) {
+                                    const chat_template = metadata['tokenizer.chat_template'];
+                                    const tokens = metadata['tokenizer.ggml.tokens'];
+                                    const bos_token = tokens[metadata['tokenizer.ggml.bos_token_id']];
+                                    const eos_token = tokens[metadata['tokenizer.ggml.eos_token_id']];
+
+                                    const maybeConfig = {
+                                        chat_template,
+                                        bos_token,
+                                        eos_token,
+                                    }
+
+                                    if (isTokenizerConfig(maybeConfig)) {
+                                        tokenizerConfig = maybeConfig;
+                                        break;
+                                    }
+                                } else if ('tokenizer.ggml.model' in metadata) {
+                                    break; // no reason to touch different quants
+                                }
+                            }
+                        } catch { }
+                    }
+
+                } catch { }
+
+                if (tokenizerConfig) {
+                    break;
+                }
+            }
         }
 
         if (tokenizerConfig) {
-            console.log(`Huggingface config for '${model}' found.`);
+            console.log(`[huggingface] found config for '${model}'`);
             return {
                 chat_template: tokenizerConfig.chat_template,
                 eos_token: tokenizerConfig.eos_token,
+                bos_token: tokenizerConfig.bos_token,
             };
         }
 
-        console.log(`Huggingface config for '${model}' not found.`);
+        console.log(`[huggingface] not found config for '${model}'`);
         return null;
     };
 
-    export const findModelTemplate = async (modelName: string): Promise<string | null> => {
-        const config = await loadHuggingfaceTokenizerConfig(modelName);
-
-        if (config?.chat_template?.trim()) {
-            const template = config.chat_template.trim()
-                .replace('eos_token', `'${config.eos_token}'`)
-                .replace('bos_token', `''`);
-
-            return template;
+    function updateRequired<T extends IParameter>(param: T): T {
+        if ('items' in param) {
+            updateRequired(param.items);
+        } else if ('properties' in param) {
+            for (const prop of Object.values(param.properties)) {
+                updateRequired(prop);
+            }
+            param.required = Object.keys(param.properties);
         }
 
-        return null;
+        return param;
     }
+
+    const convertFunctionToTool = (fn: IFunction): ITool => ({
+        type: 'function',
+        function: {
+            name: fn.name,
+            description: fn.description,
+            parameters: updateRequired({
+                type: 'object',
+                properties: fn.parameters ?? {},
+            })
+        }
+    })
+
+    export const testToolCalls = (template: string): boolean => {
+        const history: ITemplateMessage[] = [
+            { role: 'system', content: 'You are calculator.' },
+            { role: 'user', content: 'Calculate 2 + 2.' },
+        ];
+
+        const needle = '___AWOORWA_NEEDLE__';
+
+        const tools: IFunction[] = [{
+            name: 'add',
+            description: 'Test function',
+            parameters: {
+                a: { type: 'number' },
+                b: { type: 'number' },
+                c: { type: 'array', items: { type: 'number' } },
+                d: { type: 'object', properties: { inside: { type: 'number', description: needle } } },
+            }
+        }];
+
+        const text = applyChatTemplate(template, history, tools);
+
+        console.log(text);
+
+        return text.includes(needle);
+    }
+
+    export const findModelTemplate = async (modelName: string): Promise<string | null> => {
+        let template = templateCache[modelName] ?? null;
+
+        if (template) {
+            console.log(`[huggingface] found cached template for '${modelName}'`);
+        } else {
+            const config = await loadHuggingfaceTokenizerConfig(modelName);
+
+            if (config?.chat_template?.trim()) {
+                template = config.chat_template.trim()
+                    .replaceAll('eos_token', `'${config.eos_token ?? ''}'`)
+                    .replaceAll('bos_token', `''`);
+
+                if (config.bos_token) {
+                    template = template
+                        .replaceAll(config.bos_token, '')
+                        .replace(/\{\{ ?(''|"") ?\}\}/g, '');
+                }
+            }
+        }
+
+        templateCache[modelName] = template;
+        saveCache(templateCache);
+
+        return template;
+    }
+
+    export const applyChatTemplate = (templateString: string, messages: ITemplateMessage[], functions?: IFunction[]) => {
+        const template = new Template(templateString);
+
+        const prompt = template.render({
+            messages,
+            add_generation_prompt: true,
+            tools: functions?.map(convertFunctionToTool),
+        });
+
+        return prompt;
+    };
 }
\ No newline at end of file