expose decoder

packages/cli/src/parse.ts

@@ -98,7 +98,7 @@ export async function parseJinja2(
     }
 ) {
     let src = await readFile(file, { encoding: "utf-8" })
-    if (PROMPTY_REGEX.test(file)) src = promptyParse(src).content
+    if (PROMPTY_REGEX.test(file)) src = promptyParse(file, src).content
     else if (MD_REGEX.test(file)) src = splitMarkdown(src).content
 
     const vars: Record<string, any> = parseOptionsVars(
@@ -188,7 +188,7 @@ export async function parseTokens(
     options: { excludedFiles: string[]; model: string }
 ) {
     const { model = DEFAULT_MODEL } = options || {}
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
 
     const files = await expandFiles(filesGlobs, options?.excludedFiles)
     console.log(`parsing ${files.length} files`)
@@ -222,7 +222,7 @@ export async function prompty2genaiscript(
             : replaceExt(f, ".genai.mts")
         console.log(`${f} -> ${gf}`)
         const content = await readText(f)
-        const doc = promptyParse(content)
+        const doc = promptyParse(f, content)
         const script = promptyToGenAIScript(doc)
         await writeText(gf, script)
     }

packages/cli/src/run.ts

@@ -272,7 +272,7 @@ export async function runScript(
                 DOCS_CONFIGURATION_URL
             )
         }
-        trace.options.encoder = await resolveTokenEncoder(info.model)
+        trace.options.encoder = (await resolveTokenEncoder(info.model)).encode
         await runtimeHost.models.pullModel(info.model)
 
         let tokenColor = 0

packages/core/src/anthropic.ts

@@ -197,7 +197,7 @@ export const AnthropicChatCompletion: ChatCompletionHandler = async (
     const { requestOptions, partialCb, cancellationToken, inner } = options
     const { headers } = requestOptions || {}
     const { model } = parseModelIdentifier(req.model)
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
 
     const anthropic = new Anthropic({
         baseURL: cfg.base,

packages/core/src/chat.ts

@@ -140,7 +140,7 @@ async function runToolCalls(
 ) {
     const projFolder = host.projectFolder()
     const { cancellationToken, trace, model } = options || {}
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
     assert(!!trace)
     let edits: Edits[] = []
 

packages/core/src/encoders.test.ts

@@ -6,27 +6,27 @@ import { encode as defaultEncode } from "gpt-tokenizer"
 describe("resolveTokenEncoder", () => {
     test("gpt-3.5-turbo", async () => {
         const encoder = await resolveTokenEncoder("gpt-3.5-turbo")
-        const result = encoder("test line")
+        const result = encoder.encode("test line")
         assert.deepEqual(result, [1985, 1584])
     })
     test("gpt-4", async () => {
         const encoder = await resolveTokenEncoder("gpt-4")
-        const result = encoder("test line")
+        const result = encoder.encode("test line")
         assert.deepEqual(result, [1985, 1584])
     })
     test("gpt-4o", async () => {
         const encoder = await resolveTokenEncoder("gpt-4o")
-        const result = encoder("test line")
+        const result = encoder.encode("test line")
         assert.deepEqual(result, [3190, 2543])
     })
     test("gpt-4o-mini", async () => {
         const encoder = await resolveTokenEncoder("gpt-4o-mini")
-        const result = encoder("test line")
+        const result = encoder.encode("test line")
         assert.deepEqual(result, [3190, 2543])
     })
     test("gpt-4o forbidden", async () => {
         const encoder = await resolveTokenEncoder("gpt-4o")
-        const result = encoder("<|im_end|>")
+        const result = encoder.encode("<|im_end|>")
         assert.deepEqual(result, [27, 91, 321, 13707, 91, 29])
     })
 })

packages/core/src/encoders.ts

@@ -1,26 +1,37 @@
 // Import the function to parse model identifiers
 import { parseModelIdentifier } from "./models"
 
+export interface TokenEncoders {
+    encode: TokenEncoder
+    decode: TokenDecoder
+}
+
 /**
  * Resolves the appropriate token encoder based on the given model ID.
  * @param modelId - The identifier for the model to resolve the encoder for.
  * @returns A Promise that resolves to a TokenEncoder function.
  */
 export async function resolveTokenEncoder(
     modelId: string
-): Promise<TokenEncoder> {
+): Promise<TokenEncoders> {
     // Parse the model identifier to extract the model information
     const { model } = parseModelIdentifier(modelId)
     const module = model // Assign model to module for dynamic import path
 
     const options = { disallowedSpecial: new Set<string>() }
     try {
         // Attempt to dynamically import the encoder module for the specified model
-        const mod = await import(`gpt-tokenizer/model/${module}`)
-        return (line) => mod.encode(line, options) // Return the encoder function
+        const { encode, decode } = await import(`gpt-tokenizer/model/${module}`)
+        return Object.freeze<TokenEncoders>({
+            encode: (line) => encode(line, options), // Return the default encoder function
+            decode,
+        })
     } catch (e) {
         // If the specific model encoder is not found, default to gpt-4o encoder
-        const { encode } = await import("gpt-tokenizer")
-        return (line) => encode(line, options) // Return the default encoder function
+        const { encode, decode } = await import("gpt-tokenizer")
+        return Object.freeze<TokenEncoders>({
+            encode: (line) => encode(line, options), // Return the default encoder function
+            decode,
+        })
     }
 }

packages/core/src/git.ts

@@ -282,7 +282,7 @@ export class GitClient implements Git {
         }
         if (!nameOnly && llmify) {
             res = llmifyDiff(res)
-            const encoder = await resolveTokenEncoder(
+            const { encode: encoder } = await resolveTokenEncoder(
                 runtimeHost.defaultModelOptions.model || DEFAULT_MODEL
             )
             const tokens = estimateTokens(res, encoder)

packages/core/src/globals.ts

@@ -123,14 +123,14 @@ export function installGlobals() {
 
     glb.tokenizers = Object.freeze<Tokenizers>({
         count: async (text, options) => {
-            const encoder = await resolveTokenEncoder(
+            const { encode: encoder } = await resolveTokenEncoder(
                 options?.model || runtimeHost.defaultModelOptions.model
             )
             const c = await estimateTokens(text, encoder)
             return c
         },
         truncate: async (text, maxTokens, options) => {
-            const encoder = await resolveTokenEncoder(
+            const { encode: encoder } = await resolveTokenEncoder(
                 options?.model || runtimeHost.defaultModelOptions.model
             )
             return await truncateTextToTokens(text, maxTokens, encoder, options)

packages/core/src/openai.ts

@@ -73,7 +73,7 @@ export const OpenAIChatCompletion: ChatCompletionHandler = async (
     const { headers = {}, ...rest } = requestOptions || {}
     const { token, source, ...cfgNoToken } = cfg
     const { model } = parseModelIdentifier(req.model)
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
 
     const cache = !!cacheOrName || !!cacheName
     const cacheStore = getChatCompletionCache(

packages/core/src/parsers.ts

@@ -34,7 +34,7 @@ export async function createParsers(options: {
     model: string
 }): Promise<Parsers> {
     const { trace, model } = options
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
     return Object.freeze<Parsers>({
         JSON5: (text, options) =>
             JSON5TryParse(filenameOrFileToContent(text), options?.defaultValue),
@@ -120,6 +120,6 @@ export async function createParsers(options: {
         },
         diff: (f1, f2) => llmifyDiff(createDiff(f1, f2)),
         tidyData: (rows, options) => tidyData(rows, options),
-        hash: async (text, options) => await hash(text, options)
+        hash: async (text, options) => await hash(text, options),
     })
 }

packages/core/src/promptdom.ts

@@ -568,7 +568,7 @@ async function resolvePromptNode(
     model: string,
     root: PromptNode
 ): Promise<{ errors: number }> {
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
     let err = 0
     const names = new Set<string>()
     const uniqueName = (n_: string) => {
@@ -742,7 +742,7 @@ async function resolveImportPrompty(
     args: Record<string, string | number | boolean>,
     options: ImportTemplateOptions
 ) {
-    const { messages } = promptyParse(f.content)
+    const { messages } = promptyParse(f.filename, f.content)
     for (const message of messages) {
         const txt = jinjaRenderChatMessage(message, args)
         if (message.role === "assistant")
@@ -761,7 +761,7 @@ async function truncatePromptNode(
     options?: TraceOptions
 ): Promise<boolean> {
     const { trace } = options || {}
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
     let truncated = false
 
     const cap = (n: {
@@ -923,7 +923,7 @@ export async function renderPromptNode(
 ): Promise<PromptNodeRender> {
     const { trace, flexTokens } = options || {}
     const { model } = parseModelIdentifier(modelId)
-    const encoder = await resolveTokenEncoder(model)
+    const { encode: encoder } = await resolveTokenEncoder(model)
 
     await resolvePromptNode(model, node)
     await tracePromptNode(trace, node)

packages/core/src/prompty.test.ts

@@ -4,7 +4,7 @@ import assert from "node:assert/strict"
 
 describe("promptyParse", () => {
     test("correctly parses an empty markdown string", () => {
-        const result = promptyParse("")
+        const result = promptyParse(undefined, "")
         assert.deepStrictEqual(result, {
             meta: {},
             frontmatter: {},
@@ -15,7 +15,7 @@ describe("promptyParse", () => {
 
     test("correctly parses a markdown string without frontmatter", () => {
         const content = "This is a sample content without frontmatter."
-        const result = promptyParse(content)
+        const result = promptyParse(undefined, "")
         assert.deepStrictEqual(result, {
             meta: {},
             frontmatter: {},
@@ -40,7 +40,7 @@ sample:
 ---
 # Heading
 Content below heading.`
-        const result = promptyParse(markdownString)
+        const result = promptyParse(undefined, markdownString)
         assert.deepStrictEqual(result.frontmatter, {
             name: "Test",
             description: "A test description",
@@ -59,7 +59,7 @@ assistant:
 Assistant's reply
 user:
 Another message from the user`
-        const result = promptyParse(markdownContent)
+        const result = promptyParse(undefined, markdownContent)
         assert.deepStrictEqual(result.messages, [
             { role: "user", content: "User's message" },
             { role: "assistant", content: "Assistant's reply" },
@@ -69,7 +69,7 @@ Another message from the user`
 
     test("correctly handles a markdown string with content but without roles", () => {
         const markdownContent = `Just some content without specifying roles.`
-        const result = promptyParse(markdownContent)
+        const result = promptyParse(undefined, markdownContent)
         assert.deepStrictEqual(result.messages, [
             { role: "system", content: markdownContent },
         ])

packages/core/src/types/prompt_template.d.ts

@@ -1083,6 +1083,7 @@ interface ParseZipOptions {
 }
 
 type TokenEncoder = (text: string) => number[]
+type TokenDecoder = (lines: Iterable<number>) => string
 
 interface CSVParseOptions {
     delimiter?: string