1/**
2 * Service for model-specific token counting and prompt truncation. Tokenization
3 * depends on the target provider, model, and encoding rules, so this module
4 * leaves the actual tokenization function to the service implementation.
5 *
6 * The `Tokenizer` service can count tokens for raw prompt input and shorten a
7 * prompt to a token limit by keeping the newest messages that fit. This module
8 * defines the service tag, the service interface, and a `make` constructor that
9 * builds a full tokenizer service from a token-counting function.
10 *
11 * @stability unstable
12 * @since 4.0.0
13 */
14import * as Context from "../Context.ts"
15import * as Effect from "../Effect.ts"
16import * as Predicate from "../Predicate.ts"
17import type * as AiError from "./AiError.ts"
18import * as Prompt from "./Prompt.ts"
19
20/**
21 * Service tag for model tokenization services.
22 *
23 * **When to use**
24 *
25 * Use to access or provide model-specific token counting and prompt truncation
26 * operations.
27 *
28 * **Details**
29 *
30 * This tag provides access to tokenization functionality throughout your
31 * application, enabling token counting and prompt truncation capabilities.
32 *
33 * **Example** (Accessing the Tokenizer service)
34 *
35 * ```ts import.meta.vitest
36 * import { Effect } from "effect"
37 * import { Tokenizer } from "effect/ai"
38 *
39 * const useTokenizer = Effect.gen(function*() {
40 * const tokenizer = yield* Tokenizer.Tokenizer
41 * const tokens = yield* tokenizer.tokenize("Hello, world!")
42 * return tokens.length
43 * })
44 *
45 * const tokenizer = Tokenizer.make({
46 * tokenize: (prompt) => Effect.succeed(prompt.content.map((_, index) => index))
47 * })
48 * const result = useTokenizer.pipe(Effect.provideService(Tokenizer.Tokenizer, tokenizer))
49 * await Effect.runPromise(result) // => 1
50 * ```
51 *
52 * @stability unstable
53 * @category services
54 * @since 4.0.0
55 */
56export class Tokenizer extends Context.Service<Tokenizer, Service>()(
57 "effect/ai/Tokenizer"
58) {}
59
60/**
61 * Tokenizer service interface providing text tokenization and truncation
62 * operations.
63 *
64 * **Details**
65 *
66 * This interface defines the core operations for converting text to tokens and
67 * managing content length within token limits for AI model compatibility.
68 *
69 * **Example** (Implementing a custom tokenizer)
70 *
71 * ```ts import.meta.vitest
72 * import { Effect } from "effect"
73 * import { Prompt } from "effect/ai"
74 * import type { Tokenizer } from "effect/ai"
75 *
76 * const customTokenizer: Tokenizer.Service = {
77 * tokenize: (input) =>
78 * Effect.succeed(input.toString().split(" ").map((_, i) => i)),
79 * truncate: (input, maxTokens) =>
80 * Effect.succeed(Prompt.make(input.toString().slice(0, maxTokens * 5)))
81 * }
82 *
83 * const tokenCount = (await Effect.runPromise(customTokenizer.tokenize("one two three"))).length // => 3
84 * const messageCount = (await Effect.runPromise(customTokenizer.truncate("hello world", 1))).content.length // => 1
85 * ```
86 *
87 * @stability unstable
88 * @category models
89 * @since 4.0.0
90 */
91export interface Service {
92 /**
93 * Converts text input into an array of token numbers.
94 */
95 readonly tokenize: (
96 /**
97 * The text input to tokenize.
98 */
99 input: Prompt.RawInput
100 ) => Effect.Effect<Array<number>, AiError.AiError>
101 /**
102 * Truncates text input to fit within the specified token limit.
103 */
104 readonly truncate: (
105 /**
106 * The text input to truncate.
107 */
108 input: Prompt.RawInput,
109 /**
110 * Maximum number of tokens to retain.
111 */
112 tokens: number
113 ) => Effect.Effect<Prompt.Prompt, AiError.AiError>
114}
115
116/**
117 * Creates a Tokenizer service implementation from tokenization options.
118 *
119 * **Details**
120 *
121 * This function constructs a complete Tokenizer service by providing a
122 * tokenization function. The service handles both tokenization and
123 * truncation operations using the provided tokenizer.
124 *
125 * **Example** (Creating a word tokenizer)
126 *
127 * ```ts import.meta.vitest
128 * import { Effect } from "effect"
129 * import { Tokenizer } from "effect/ai"
130 *
131 * // Simple word-based tokenizer
132 * const wordTokenizer = Tokenizer.make({
133 * tokenize: (prompt) =>
134 * Effect.succeed(
135 * prompt.content
136 * .flatMap((msg) =>
137 * typeof msg.content === "string"
138 * ? msg.content.split(" ")
139 * : msg.content.flatMap((part) =>
140 * part.type === "text" ? part.text.split(" ") : []
141 * )
142 * )
143 * .map((_, index) => index)
144 * )
145 * })
146 *
147 * await Effect.runPromise(wordTokenizer.tokenize("hello effect world")) // => [0, 1, 2]
148 * ```
149 *
150 * @stability unstable
151 * @category constructors
152 * @since 4.0.0
153 */
154export const make = (options: {
155 readonly tokenize: (content: Prompt.Prompt) => Effect.Effect<Array<number>, AiError.AiError>
156}): Service =>
157 Tokenizer.of({
158 tokenize(input) {
159 return options.tokenize(Prompt.make(input))
160 },
161 truncate(input, tokens) {
162 return truncate(Prompt.make(input), options.tokenize, tokens)
163 }
164 })
165
166const truncate = (
167 self: Prompt.Prompt,
168 tokenize: (input: Prompt.Prompt) => Effect.Effect<Array<number>, AiError.AiError>,
169 maxTokens: number
170): Effect.Effect<Prompt.Prompt, AiError.AiError> =>
171 Effect.suspend(() => {
172 let inputMessages = self.content
173 let outputMessages: Array<Prompt.Message> = []
174 const loop: Effect.Effect<Prompt.Prompt, AiError.AiError> = Effect.suspend(() => {
175 const message = inputMessages[inputMessages.length - 1]
176 if (Predicate.isUndefined(message)) {
177 return Effect.succeed(Prompt.fromMessages(outputMessages))
178 }
179 inputMessages = inputMessages.slice(0, inputMessages.length - 1)
180 const candidate = [message, ...outputMessages]
181 return Effect.flatMap(tokenize(Prompt.fromMessages(candidate)), (tokens) => {
182 if (tokens.length > maxTokens) {
183 return Effect.succeed(Prompt.fromMessages(outputMessages))
184 }
185 outputMessages = candidate
186 return loop
187 })
188 })
189 return loop
190 })