typellm 0.2.1__tar.gz → 0.2.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (33) hide show
  1. {typellm-0.2.1 → typellm-0.2.3}/PKG-INFO +27 -16
  2. {typellm-0.2.1 → typellm-0.2.3}/README.md +26 -15
  3. {typellm-0.2.1 → typellm-0.2.3}/pyproject.toml +1 -1
  4. {typellm-0.2.1 → typellm-0.2.3}/tests/test_batching.py +95 -11
  5. {typellm-0.2.1 → typellm-0.2.3}/tests/test_concurrency.py +6 -1
  6. {typellm-0.2.1 → typellm-0.2.3}/tests/test_dependencies.py +1 -1
  7. {typellm-0.2.1 → typellm-0.2.3}/tests/test_images.py +15 -1
  8. {typellm-0.2.1 → typellm-0.2.3}/tests/test_protocols.py +10 -10
  9. {typellm-0.2.1 → typellm-0.2.3}/tests/test_text.py +2 -2
  10. {typellm-0.2.1 → typellm-0.2.3}/tests/test_typellm.py +12 -14
  11. {typellm-0.2.1 → typellm-0.2.3}/typellm/runtime.py +51 -14
  12. {typellm-0.2.1 → typellm-0.2.3}/typellm/schema.py +10 -1
  13. {typellm-0.2.1 → typellm-0.2.3}/typellm/sglang.py +116 -28
  14. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/PKG-INFO +27 -16
  15. {typellm-0.2.1 → typellm-0.2.3}/LICENSE +0 -0
  16. {typellm-0.2.1 → typellm-0.2.3}/setup.cfg +0 -0
  17. {typellm-0.2.1 → typellm-0.2.3}/tests/test_grammar_fields.py +0 -0
  18. {typellm-0.2.1 → typellm-0.2.3}/tests/test_nullable.py +0 -0
  19. {typellm-0.2.1 → typellm-0.2.3}/tests/test_permutations.py +0 -0
  20. {typellm-0.2.1 → typellm-0.2.3}/tests/test_prefill.py +0 -0
  21. {typellm-0.2.1 → typellm-0.2.3}/tests/test_probabilities.py +0 -0
  22. {typellm-0.2.1 → typellm-0.2.3}/tests/test_prompts.py +0 -0
  23. {typellm-0.2.1 → typellm-0.2.3}/typellm/__init__.py +0 -0
  24. {typellm-0.2.1 → typellm-0.2.3}/typellm/__main__.py +0 -0
  25. {typellm-0.2.1 → typellm-0.2.3}/typellm/benchmark.py +0 -0
  26. {typellm-0.2.1 → typellm-0.2.3}/typellm/cli.py +0 -0
  27. {typellm-0.2.1 → typellm-0.2.3}/typellm/images.py +0 -0
  28. {typellm-0.2.1 → typellm-0.2.3}/typellm/protocol.py +0 -0
  29. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/SOURCES.txt +0 -0
  30. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/dependency_links.txt +0 -0
  31. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/entry_points.txt +0 -0
  32. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/requires.txt +0 -0
  33. {typellm-0.2.1 → typellm-0.2.3}/typellm.egg-info/top_level.txt +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: typellm
3
- Version: 0.2.1
3
+ Version: 0.2.3
4
4
  Summary: Type-Safe Decoding for Autoregressive LLMs
5
5
  License-Expression: Apache-2.0
6
6
  Project-URL: Homepage, https://typellm.ai/
@@ -168,6 +168,9 @@ result = client.generate(
168
168
  )
169
169
  ```
170
170
 
171
+ Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
172
+ longer answers, or set `maxLength` on a field to cap its characters.
173
+
171
174
  Ask for a numeric answer without enumerating every possible value:
172
175
 
173
176
  ```python
@@ -229,23 +232,28 @@ result = client.generate(
229
232
 
230
233
  ## Thinking mode
231
234
 
232
- Thinking is off by default. Enable it when constructing the client:
235
+ Thinking is off by default. Turn it on for the fields that need it; the others
236
+ answer at once, and the fields that think reason side by side:
233
237
 
234
238
  ```python
235
- client = TypeLLMClient(
236
- "http://127.0.0.1:30000",
237
- model="Qwen/Qwen3.8-27B",
238
- thinking=True,
239
- )
240
- result = client.generate(context=context, questions=questions)
239
+ result = client.generate(context=context, questions={
240
+ "total": {"type": "number"},
241
+ "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
242
+ "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
243
+ "thinking": True, "thinking_budget": 1024},
244
+ })
241
245
  ```
242
246
 
243
- Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
244
- default. When reasoning reaches the budget, TypeLLM closes it and moves on to
245
- the typed answer.
247
+ `thinking_budget` caps a field's reasoning; there is no budget by default.
248
+ `TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
249
+ without its own. When reasoning reaches the budget, TypeLLM closes it and moves
250
+ on to the typed answer.
251
+
252
+ After a call, `client.last_thinking` maps each field that thought to its
253
+ reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
246
254
 
247
- Models with always-on thinking still reason with `thinking=False`;
248
- `thinking_budget` applies to them too. See [Supported models](#supported-models).
255
+ Models with always-on thinking reason on every field; `thinking_budget` applies
256
+ to them too. See [Supported models](#supported-models).
249
257
 
250
258
  ## Image input
251
259
 
@@ -416,11 +424,14 @@ result = client.generate(
416
424
  cancel=cancel, # set it from another thread; raises GenerationCancelled
417
425
  )
418
426
  print(client.last_usage)
419
- # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
427
+ # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
420
428
  ```
421
429
 
422
- `last_usage` reports the SGLang requests and tokens of the last call made in the
423
- current thread, including a call that failed partway. A timeout or cancel stops
430
+ `last_usage` reports the tokens of the last call made in the current thread,
431
+ including a call that failed partway. `input_tokens` is what you sent, each part
432
+ counted once: the context, the questions as JSON, and the images. The other
433
+ counts are SGLang's for the call's requests, where every prompt carries the
434
+ shared context. A timeout or cancel stops
424
435
  the call before its next SGLang request.
425
436
 
426
437
  ## Cost analysis
@@ -153,6 +153,9 @@ result = client.generate(
153
153
  )
154
154
  ```
155
155
 
156
+ Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
157
+ longer answers, or set `maxLength` on a field to cap its characters.
158
+
156
159
  Ask for a numeric answer without enumerating every possible value:
157
160
 
158
161
  ```python
@@ -214,23 +217,28 @@ result = client.generate(
214
217
 
215
218
  ## Thinking mode
216
219
 
217
- Thinking is off by default. Enable it when constructing the client:
220
+ Thinking is off by default. Turn it on for the fields that need it; the others
221
+ answer at once, and the fields that think reason side by side:
218
222
 
219
223
  ```python
220
- client = TypeLLMClient(
221
- "http://127.0.0.1:30000",
222
- model="Qwen/Qwen3.8-27B",
223
- thinking=True,
224
- )
225
- result = client.generate(context=context, questions=questions)
224
+ result = client.generate(context=context, questions={
225
+ "total": {"type": "number"},
226
+ "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
227
+ "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
228
+ "thinking": True, "thinking_budget": 1024},
229
+ })
226
230
  ```
227
231
 
228
- Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
229
- default. When reasoning reaches the budget, TypeLLM closes it and moves on to
230
- the typed answer.
232
+ `thinking_budget` caps a field's reasoning; there is no budget by default.
233
+ `TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
234
+ without its own. When reasoning reaches the budget, TypeLLM closes it and moves
235
+ on to the typed answer.
236
+
237
+ After a call, `client.last_thinking` maps each field that thought to its
238
+ reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
231
239
 
232
- Models with always-on thinking still reason with `thinking=False`;
233
- `thinking_budget` applies to them too. See [Supported models](#supported-models).
240
+ Models with always-on thinking reason on every field; `thinking_budget` applies
241
+ to them too. See [Supported models](#supported-models).
234
242
 
235
243
  ## Image input
236
244
 
@@ -401,11 +409,14 @@ result = client.generate(
401
409
  cancel=cancel, # set it from another thread; raises GenerationCancelled
402
410
  )
403
411
  print(client.last_usage)
404
- # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
412
+ # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
405
413
  ```
406
414
 
407
- `last_usage` reports the SGLang requests and tokens of the last call made in the
408
- current thread, including a call that failed partway. A timeout or cancel stops
415
+ `last_usage` reports the tokens of the last call made in the current thread,
416
+ including a call that failed partway. `input_tokens` is what you sent, each part
417
+ counted once: the context, the questions as JSON, and the images. The other
418
+ counts are SGLang's for the call's requests, where every prompt carries the
419
+ shared context. A timeout or cancel stops
409
420
  the call before its next SGLang request.
410
421
 
411
422
  ## Cost analysis
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "typellm"
7
- version = "0.2.1"
7
+ version = "0.2.3"
8
8
  description = "Type-Safe Decoding for Autoregressive LLMs"
9
9
  readme = "README.md"
10
10
  license = "Apache-2.0"
@@ -26,8 +26,8 @@ class FakeServer(SGLangClient):
26
26
  Numeric fields decode to 7: the server prefers the end token, then "7".
27
27
  """
28
28
 
29
- def __init__(self, thinking=False):
30
- super().__init__(model="fake", thinking=thinking)
29
+ def __init__(self):
30
+ super().__init__(model="fake")
31
31
  self._chat_tokenizer = ThinkingTokenizer()
32
32
  self._context_length_cache = 100_000
33
33
  self.payloads = []
@@ -83,6 +83,10 @@ def width(payload):
83
83
  return 1 if isinstance(payload["text"], str) else len(payload["text"])
84
84
 
85
85
 
86
+ def thinks(questions):
87
+ return {name: {**field, "thinking": True} for name, field in questions.items()}
88
+
89
+
86
90
  class PrefixWarmupTests(unittest.TestCase):
87
91
  MIXED = {"a": {"type": "integer"}, "b": {"type": "number"}, "n": {"type": "string"},
88
92
  "k": {"type": "string", "enum": ["x", "y"]}}
@@ -91,8 +95,8 @@ class PrefixWarmupTests(unittest.TestCase):
91
95
  for thinking in (False, True):
92
96
  with self.subTest(thinking=thinking):
93
97
  client = TypeLLMClient(model="fake")
94
- client.sglang = FakeServer(thinking=thinking)
95
- client.generate(context="Receipt", questions=self.MIXED)
98
+ client.sglang = FakeServer()
99
+ client.generate(context="Receipt", questions=thinks(self.MIXED) if thinking else self.MIXED)
96
100
  warm = client.sglang.requests("count")
97
101
  self.assertEqual(len(warm), 1)
98
102
  # First, before the thinking, number, text and scoring batches.
@@ -107,6 +111,86 @@ class PrefixWarmupTests(unittest.TestCase):
107
111
  self.assertEqual(client.sglang.requests("count"), [])
108
112
 
109
113
 
114
+ class PerFieldThinkingTests(unittest.TestCase):
115
+ def thinking_prompts(self, server):
116
+ return [t for p in server.requests("think")
117
+ for t in ([p["text"]] if isinstance(p["text"], str) else p["text"])]
118
+
119
+ def test_only_fields_that_ask_for_it_think(self):
120
+ client = TypeLLMClient(model="fake")
121
+ client.sglang = FakeServer()
122
+ result = client.generate(context="Receipt", questions={
123
+ "hard": {"type": "boolean", "instructions": "Hard?", "thinking": True},
124
+ "easy": {"type": "boolean", "instructions": "Easy?"},
125
+ "count": {"type": "integer"},
126
+ })
127
+ self.assertEqual(result, {"hard": True, "easy": True, "count": 7})
128
+ [prompt] = self.thinking_prompts(client.sglang) # one batch, one prompt
129
+ self.assertIn("Hard?", prompt)
130
+ self.assertTrue(prompt.endswith("<think>\n"))
131
+ answers = [t for p in client.sglang.payloads if p not in client.sglang.requests("think")
132
+ for t in ([p["text"]] if isinstance(p["text"], str) else p["text"])]
133
+ self.assertTrue(any("Easy?" in t and "<think>" not in t for t in answers))
134
+
135
+ def test_thinking_false_is_the_same_as_leaving_it_out(self):
136
+ client = TypeLLMClient(model="fake")
137
+ client.sglang = FakeServer()
138
+ client.generate(context="Receipt", questions={
139
+ "a": {"type": "boolean", "instructions": "A?", "thinking": True},
140
+ "b": {"type": "boolean", "instructions": "B?", "thinking": True},
141
+ "skip": {"type": "boolean", "instructions": "Skip?", "thinking": False},
142
+ })
143
+ prompts = self.thinking_prompts(client.sglang)
144
+ self.assertEqual(len(client.sglang.requests("think")), 1)
145
+ self.assertEqual(len(prompts), 2)
146
+ self.assertFalse(any("Skip?" in t for t in prompts))
147
+
148
+ def test_each_field_gets_its_own_budget(self):
149
+ client = TypeLLMClient(model="fake")
150
+ client.sglang = FakeServer()
151
+ client.sglang.thinking_budget = 500 # the client's budget, for fields without their own
152
+ client.generate(context="Receipt", questions={
153
+ "short": {"type": "boolean", "instructions": "Short?", "thinking": True, "thinking_budget": 100},
154
+ "long": {"type": "boolean", "instructions": "Long?", "thinking": True, "thinking_budget": 300},
155
+ "default": {"type": "boolean", "instructions": "Default?", "thinking": True},
156
+ })
157
+ [request] = client.sglang.requests("think")
158
+ budgets = {next(q for q in ("Short?", "Long?", "Default?") if q in t): params["max_new_tokens"]
159
+ for t, params in zip(request["text"], request["sampling_params"])}
160
+ self.assertEqual(budgets, {"Short?": 100, "Long?": 300, "Default?": 500})
161
+
162
+ def test_reasoning_and_its_tokens_are_reported_per_field(self):
163
+ class Metered(FakeServer):
164
+ def _request(self, path, payload=None, *, allow_text=False):
165
+ response = super()._request(path, payload, allow_text=allow_text)
166
+ if path == "/generate":
167
+ params = payload["sampling_params"]
168
+ thinking = (params[0] if isinstance(params, list) else params).get("stop") == [THINK_STOP]
169
+ for item in response if isinstance(response, list) else [response]:
170
+ item["meta_info"]["completion_tokens"] = 40 if thinking else 1
171
+ return response
172
+
173
+ client = TypeLLMClient(model="fake")
174
+ client.sglang = Metered()
175
+ client.generate(context="Receipt", questions={
176
+ "hard": {"type": "boolean", "instructions": "Hard?", "thinking": True},
177
+ "easy": {"type": "boolean", "instructions": "Easy?"},
178
+ })
179
+ self.assertEqual(client.last_thinking, {"hard": "Reasoned."})
180
+ self.assertEqual(client.last_usage.thinking_tokens, 40)
181
+ self.assertGreater(client.last_usage.completion_tokens, 40) # answers count too
182
+ with self.assertRaises(ValueError):
183
+ client.generate(questions={"x": {"type": "boolean"}})
184
+ self.assertEqual(client.last_thinking, {})
185
+
186
+ def test_invalid_settings_are_schema_errors(self):
187
+ from typellm import SchemaError
188
+ for field in ({"type": "boolean", "thinking": "yes"}, {"type": "boolean", "thinking_budget": 0},
189
+ {"type": "boolean", "thinking_budget": 2.5}):
190
+ with self.subTest(field=field), self.assertRaises(SchemaError):
191
+ TypeLLMClient(model="fake").compile_schema({"type": "object", "properties": {"x": field}})
192
+
193
+
110
194
  class BatchedThinkingTests(unittest.TestCase):
111
195
  QUESTIONS = {
112
196
  "flag": {"type": "boolean"},
@@ -117,8 +201,8 @@ class BatchedThinkingTests(unittest.TestCase):
117
201
 
118
202
  def test_one_thinking_request_covers_every_prompt_in_a_layer(self):
119
203
  client = TypeLLMClient(model="fake")
120
- client.sglang = FakeServer(thinking=True)
121
- result = client.generate(context="Receipt", questions=self.QUESTIONS)
204
+ client.sglang = FakeServer()
205
+ result = client.generate(context="Receipt", questions=thinks(self.QUESTIONS))
122
206
  self.assertEqual(result, {"flag": True, "count": 7, "name": "blue", "pick": "x"})
123
207
  [think] = client.sglang.requests("think")
124
208
  # Four fields plus the reversed ordering of "pick".
@@ -130,19 +214,19 @@ class BatchedThinkingTests(unittest.TestCase):
130
214
 
131
215
  def test_dag_thinks_once_per_layer(self):
132
216
  client = TypeLLMClient(model="fake")
133
- client.sglang = FakeServer(thinking=True)
134
- client.generate(context="Receipt", questions={
217
+ client.sglang = FakeServer()
218
+ client.generate(context="Receipt", questions=thinks({
135
219
  "a": {"type": "boolean"}, "b": {"type": "integer"},
136
220
  "c": {"type": "boolean", "depends_on": ["a", "b"]},
137
221
  "d": {"type": "string", "depends_on": ["a"]},
138
- })
222
+ }))
139
223
  self.assertEqual([width(p) for p in client.sglang.requests("think")], [2, 2])
140
224
 
141
225
  def test_images_count_prompt_tokens_in_one_batch(self):
142
226
  client = TypeLLMClient(model="fake")
143
- client.sglang = FakeServer(thinking=True)
227
+ client.sglang = FakeServer()
144
228
  client.generate(context="Receipt", images=[PNG],
145
- questions={"a": {"type": "boolean"}, "b": {"type": "boolean"}})
229
+ questions=thinks({"a": {"type": "boolean"}, "b": {"type": "boolean"}}))
146
230
  counts = [p for p in client.sglang.requests("count") if width(p) > 1]
147
231
  self.assertEqual([width(p) for p in counts], [2])
148
232
  [think] = client.sglang.requests("think")
@@ -1,3 +1,4 @@
1
+ import json
1
2
  import pickle
2
3
  import sys
3
4
  import threading
@@ -104,10 +105,13 @@ class UsageTests(unittest.TestCase):
104
105
  self.assertIsNone(client.last_usage)
105
106
  client.generate(context="Receipt", questions=self.QUESTIONS)
106
107
  n = prompts_sent(client.sglang)
108
+ # The input counts once, though every prompt carries it.
109
+ sent = client.sglang.count_tokens("Receipt") + client.sglang.count_tokens(json.dumps(self.QUESTIONS))
107
110
  self.assertEqual(client.last_usage, Usage(
108
111
  requests=len(client.sglang.payloads),
109
- prompt_tokens=10 * n, cached_tokens=4 * n, completion_tokens=n,
112
+ prompt_tokens=10 * n, cached_tokens=4 * n, completion_tokens=n, input_tokens=sent,
110
113
  ))
114
+ self.assertGreater(client.last_usage.requests, 1)
111
115
  client.sglang.payloads.clear()
112
116
  client.generate(context="Receipt", questions={"b": {"type": "boolean"}})
113
117
  self.assertEqual(client.last_usage.requests, len(client.sglang.payloads))
@@ -136,6 +140,7 @@ class UsageTests(unittest.TestCase):
136
140
  with self.assertRaisesRegex(SGLangError, "went away"):
137
141
  client.generate(context="Receipt", questions=self.QUESTIONS)
138
142
  self.assertEqual(client.last_usage.requests, 2)
143
+ self.assertGreater(client.last_usage.input_tokens, 0)
139
144
  with self.assertRaises(ValueError):
140
145
  client.generate(questions=self.QUESTIONS)
141
146
  self.assertIsNone(client.last_usage)
@@ -102,7 +102,7 @@ class DependencyTests(unittest.TestCase):
102
102
  self.assertTrue(client.last_prompts[1].startswith(client.last_prompts[0]))
103
103
 
104
104
  def test_always_thinking_template_continuation_runs_once(self):
105
- client = SGLangClient(thinking=False)
105
+ client = SGLangClient()
106
106
 
107
107
  class Tokenizer:
108
108
  chat_template = 'test'
@@ -255,7 +255,6 @@ class ImageRequestTests(unittest.TestCase):
255
255
 
256
256
  def test_thinking_budget_counts_image_tokens_on_the_server(self):
257
257
  client = FakeServerClient()
258
- client.thinking = True
259
258
  client.thinking_budget = None
260
259
  client._context_length_cache = 1000
261
260
  prefix = VisionTokenizer().apply_chat_template(
@@ -272,6 +271,21 @@ class ImageRequestTests(unittest.TestCase):
272
271
  self.assertEqual(rest, [])
273
272
 
274
273
 
274
+ class ImageInputTokenTests(unittest.TestCase):
275
+ def test_images_count_as_the_server_expands_them(self):
276
+ client = TypeLLMClient(model="fake-vl")
277
+ client.sglang = FakeServerClient()
278
+ questions = {"a": {"type": "boolean"}}
279
+ client.generate(context="Receipt", images=[PNG], questions=questions)
280
+ first = client.sglang.generate_payloads[0]["text"]
281
+ first = first if isinstance(first, str) else first[0]
282
+ count = client.sglang.count_tokens
283
+ # The server reported 900 prompt tokens, with the placeholder expanded.
284
+ image = 900 - count(first) + count(VISION)
285
+ self.assertEqual(client.last_usage.input_tokens,
286
+ count("Receipt") + count('{"a": {"type": "boolean"}}') + image)
287
+
288
+
275
289
  class RuntimeContentTests(unittest.TestCase):
276
290
  def test_fake_backends_receive_image_parts(self):
277
291
  class Recording(FakeSGLang):
@@ -39,8 +39,8 @@ class ProtocolTokenizer(FakeChatTokenizer):
39
39
 
40
40
 
41
41
  class ProtocolTests(unittest.TestCase):
42
- def client(self, family, thinking=False, response=None):
43
- client = SGLangClient(thinking=thinking, thinking_budget=64)
42
+ def client(self, family, response=None):
43
+ client = SGLangClient(thinking_budget=64)
44
44
  client._chat_tokenizer = ProtocolTokenizer(family)
45
45
  client._context_length_cache = 8192
46
46
  end = '</think>'
@@ -59,9 +59,9 @@ class ProtocolTests(unittest.TestCase):
59
59
 
60
60
  def test_gemma_is_rejected_before_inference(self):
61
61
  for thinking in (False, True):
62
- client = self.client('gemma', thinking)
62
+ client = self.client('gemma')
63
63
  with self.assertRaisesRegex(SGLangError, 'Gemma 4.*not supported'):
64
- client.render_chat([], add_generation_prompt=True)
64
+ client.render_chat([], add_generation_prompt=True, thinking=thinking)
65
65
  client._request.assert_not_called()
66
66
 
67
67
  def test_native_turn_stop_recovers_preserved_and_filtered_terminators(self):
@@ -69,7 +69,7 @@ class ProtocolTests(unittest.TestCase):
69
69
  for matched in (ending, token):
70
70
  for suffix in ('', ending + '\n'):
71
71
  with self.subTest(family=family, matched=matched, suffix=suffix):
72
- client = self.client(family, True, {
72
+ client = self.client(family, {
73
73
  'text': 'Keep reasoning' + suffix,
74
74
  'meta_info': {'finish_reason': {'type': 'stop', 'matched': matched}},
75
75
  })
@@ -91,15 +91,15 @@ class ProtocolTests(unittest.TestCase):
91
91
  ('Partial<|im_end|>', {'type': 'error', 'matched': 3}),
92
92
  ]:
93
93
  with self.subTest(text=text, finish=finish):
94
- client = self.client('minicpm5', True, {'text': text, 'meta_info': {'finish_reason': finish}})
94
+ client = self.client('minicpm5', {'text': text, 'meta_info': {'finish_reason': finish}})
95
95
  with self.assertRaises(SGLangError):
96
96
  client._finish_thinking('<think>\n')
97
97
 
98
98
  def test_minicpm5_switch_and_ring_always_thinking(self):
99
99
  for family in ('minicpm5', 'ring'):
100
100
  for thinking in (False, True):
101
- client = self.client(family, thinking)
102
- prompt = client.render_chat([], add_generation_prompt=True)
101
+ client = self.client(family)
102
+ prompt = client.render_chat([], add_generation_prompt=True, thinking=thinking)
103
103
  expected = thinking or family == 'ring'
104
104
  self.assertEqual(client._request.called, expected)
105
105
  parent = client.complete_chat_prefix(prompt, 'A')
@@ -108,9 +108,9 @@ class ProtocolTests(unittest.TestCase):
108
108
 
109
109
  def test_non_thinking_templates_reject_requested_thinking(self):
110
110
  for family in ('ling',):
111
- client = self.client(family, True)
111
+ client = self.client(family)
112
112
  with self.assertRaisesRegex(SGLangError, 'may not support thinking'):
113
- client.render_chat([], add_generation_prompt=True)
113
+ client.render_chat([], add_generation_prompt=True, thinking=True)
114
114
  client._request.assert_not_called()
115
115
 
116
116
 
@@ -63,13 +63,13 @@ class TextTests(unittest.TestCase):
63
63
  TypeLLMClient(text_max_tokens=budget)
64
64
 
65
65
  def test_thinking_then_constrained_text(self):
66
- client=TypeLLMClient(thinking=True)
66
+ client=TypeLLMClient()
67
67
  client.sglang._context_length_cache=8192
68
68
  client.sglang._chat_tokenizer=FakeChatTokenizer()
69
69
  client.sglang._chat_tokenizer.apply_chat_template=lambda *args,**kw: "assistant\n<think>\n" if kw["enable_thinking"] else "completed"
70
70
  responses=[{'text':'brief</think>'},[{'text':' "done"}','meta_info':{'finish_reason':{'type':'stop'}}}]]
71
71
  with patch.object(client.sglang,'_request',side_effect=responses) as request:
72
- self.assertEqual(client.generate(state='x',questions={'t':{'type':'string'}}),{'t':'done'})
72
+ self.assertEqual(client.generate(state='x',questions={'t':{'type':'string','thinking':True}}),{'t':'done'})
73
73
  self.assertIn('</think>',request.call_args.args[1]['text'][0])
74
74
  self.assertTrue(request.call_args.args[1]['text'][0].endswith('{"t":'))
75
75
  self.assertIn('regex',request.call_args.args[1]['sampling_params'][0])
@@ -427,7 +427,7 @@ class ThinkingTests(unittest.TestCase):
427
427
  from typellm import run_schema
428
428
  self.assertIsNone(SGLangClient().thinking_budget)
429
429
  self.assertIsNone(TypeLLMClient().sglang.thinking_budget)
430
- client = SGLangClient(thinking=True)
430
+ client = SGLangClient()
431
431
  client._chat_tokenizer = FakeChatTokenizer()
432
432
  client._context_length_cache = 8192
433
433
  client._request = Mock(return_value={"text": "Done.</think>"})
@@ -443,7 +443,7 @@ class ThinkingTests(unittest.TestCase):
443
443
 
444
444
  def make_client(self, response=None):
445
445
  from unittest.mock import Mock
446
- client = SGLangClient(thinking=True, thinking_budget=128)
446
+ client = SGLangClient(thinking_budget=128)
447
447
  client._context_length_cache = 8192
448
448
  tokenizer = FakeChatTokenizer()
449
449
  tokenizer.apply_chat_template = Mock(return_value="assistant\n<think>\n")
@@ -453,7 +453,7 @@ class ThinkingTests(unittest.TestCase):
453
453
 
454
454
  def test_thinking_stops_before_constrained_answer(self):
455
455
  client = self.make_client()
456
- prompt = client.render_chat([], add_generation_prompt=True)
456
+ prompt = client.render_chat([], add_generation_prompt=True, thinking=True)
457
457
  self.assertEqual(prompt, "assistant\n<think>\nWork done. </think>\n\n")
458
458
  self.assertNotIn("illegal answer", prompt)
459
459
  params = client._request.call_args.args[1]["sampling_params"]
@@ -464,7 +464,7 @@ class ThinkingTests(unittest.TestCase):
464
464
 
465
465
  def test_length_stop_forces_closure_and_retains_reasoning(self):
466
466
  client = self.make_client({"text": "Partial reasoning", "meta_info": {"finish_reason": {"type": "length"}}})
467
- prompt = client.render_chat([], add_generation_prompt=True)
467
+ prompt = client.render_chat([], add_generation_prompt=True, thinking=True)
468
468
  self.assertIn("Partial reasoning", prompt)
469
469
  self.assertIn("I will now give the final answer.", prompt)
470
470
  self.assertTrue(prompt.endswith("</think>\n\n"))
@@ -499,7 +499,7 @@ class ThinkingTests(unittest.TestCase):
499
499
  for text in ("Partial", "Partial</think>"):
500
500
  client = self.make_client({"text": text, "meta_info": {"finish_reason": {"type": "abort"}}})
501
501
  with self.assertRaisesRegex(SGLangError, "aborted"):
502
- client.render_chat([], add_generation_prompt=True)
502
+ client.render_chat([], add_generation_prompt=True, thinking=True)
503
503
 
504
504
  def test_forced_thinking_keeps_all_final_decoders(self):
505
505
  stops = [{"type": "length"}, {"type": "stop", "matched": "<|im_end|>"}]
@@ -528,20 +528,19 @@ class ThinkingTests(unittest.TestCase):
528
528
  client = self.make_client()
529
529
  client._request.return_value = response
530
530
  with self.subTest(response=response), self.assertRaises(SGLangError):
531
- client.render_chat([], add_generation_prompt=True)
531
+ client.render_chat([], add_generation_prompt=True, thinking=True)
532
532
 
533
533
  def test_unsupported_template_does_not_generate(self):
534
534
  from typellm import SGLangError
535
535
  client = self.make_client()
536
536
  client._chat_tokenizer.apply_chat_template.return_value = "<think></think>"
537
537
  with self.assertRaisesRegex(SGLangError, "native chat template"):
538
- client.render_chat([], add_generation_prompt=True)
538
+ client.render_chat([], add_generation_prompt=True, thinking=True)
539
539
  client._request.assert_not_called()
540
540
 
541
541
  def test_always_thinking_template_reasons_even_when_thinking_is_off(self):
542
542
  client = self.make_client()
543
- client.thinking = False
544
- prompt = client.render_chat([], add_generation_prompt=True)
543
+ prompt = client.render_chat([], add_generation_prompt=True) # thinking not asked for
545
544
  self.assertEqual(prompt, "assistant\n<think>\nWork done. </think>\n\n")
546
545
  self.assertFalse(client._chat_tokenizer.apply_chat_template.call_args.kwargs["enable_thinking"])
547
546
 
@@ -552,12 +551,11 @@ class ThinkingTests(unittest.TestCase):
552
551
  self.assertFalse(client._chat_tokenizer.apply_chat_template.call_args.kwargs["enable_thinking"])
553
552
 
554
553
  def test_public_configuration_and_validation(self):
555
- self.assertFalse(TypeLLMClient().sglang.thinking)
556
- client = TypeLLMClient(thinking=True, thinking_budget=256)
557
- self.assertTrue(client.sglang.thinking)
554
+ client = TypeLLMClient(thinking_budget=256)
558
555
  self.assertEqual(client.sglang.thinking_budget, 256)
559
- for kwargs in [{"thinking": "false"}, {"thinking": 1}, {"thinking_budget": 0},
560
- {"thinking_budget": True}, {"thinking_budget": 1.5}]:
556
+ with self.assertRaises(TypeError):
557
+ TypeLLMClient(thinking=True) # thinking is chosen per field
558
+ for kwargs in [{"thinking_budget": 0}, {"thinking_budget": True}, {"thinking_budget": 1.5}]:
561
559
  with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
562
560
  TypeLLMClient(**kwargs)
563
561
 
@@ -63,6 +63,8 @@ class Choice:
63
63
  return_probabilities: bool = False
64
64
  depends_on: tuple[str, ...] | None = None
65
65
  nullable: bool = False
66
+ thinking: bool | None = None
67
+ thinking_budget: int | None = None
66
68
 
67
69
  def __post_init__(self) -> None:
68
70
  if not self.choices and self.numeric_type is None and not self.text_type:
@@ -154,9 +156,8 @@ class TypeLLMClient:
154
156
  label_pool: Sequence[str] | None = None,
155
157
  numeric_max_digits: int = 32,
156
158
  tokenizer: str | None = None,
157
- thinking: bool = False,
158
159
  thinking_budget: int | None = None,
159
- text_max_tokens: int = 512,
160
+ text_max_tokens: int = 128,
160
161
  ) -> None:
161
162
  _validate_decoding(mode, temperature)
162
163
  if type(numeric_max_digits) is not int or numeric_max_digits <= 0:
@@ -166,7 +167,6 @@ class TypeLLMClient:
166
167
  model,
167
168
  timeout,
168
169
  tokenizer=tokenizer,
169
- thinking=thinking,
170
170
  thinking_budget=thinking_budget,
171
171
  text_max_tokens=text_max_tokens,
172
172
  answer_reserve_tokens=numeric_max_digits + 3,
@@ -184,23 +184,33 @@ class TypeLLMClient:
184
184
  self._last_usage: ContextVar[Usage | None] = ContextVar(
185
185
  f"typellm_last_usage_{id(self)}", default=None
186
186
  )
187
+ self._last_thinking: ContextVar[dict[str, str]] = ContextVar(
188
+ f"typellm_last_thinking_{id(self)}", default={}
189
+ )
187
190
 
188
191
  def __getstate__(self) -> dict[str, Any]:
189
192
  state = self.__dict__.copy()
190
193
  del state["_last_prompts"]
191
194
  del state["_last_usage"]
195
+ del state["_last_thinking"]
192
196
  return state
193
197
 
194
198
  def __setstate__(self, state: dict[str, Any]) -> None:
195
199
  self.__dict__.update(state)
196
200
  self._last_prompts = ContextVar(f"typellm_last_prompts_{id(self)}", default=[])
197
201
  self._last_usage = ContextVar(f"typellm_last_usage_{id(self)}", default=None)
202
+ self._last_thinking = ContextVar(f"typellm_last_thinking_{id(self)}", default={})
198
203
 
199
204
  @property
200
205
  def last_prompts(self) -> list[str]:
201
206
  """Final prompts of the last generate() call made in this thread or task."""
202
207
  return self._last_prompts.get()
203
208
 
209
+ @property
210
+ def last_thinking(self) -> dict[str, str]:
211
+ """Reasoning of the fields that thought in the last call in this thread or task."""
212
+ return self._last_thinking.get()
213
+
204
214
  @property
205
215
  def last_usage(self) -> Usage | None:
206
216
  """Requests and tokens of the last generate() call made in this thread or task.
@@ -256,6 +266,8 @@ class TypeLLMClient:
256
266
  return_probabilities=item.return_probabilities,
257
267
  depends_on=item.depends_on,
258
268
  nullable=item.nullable,
269
+ thinking=item.thinking,
270
+ thinking_budget=item.thinking_budget,
259
271
  )
260
272
  for item in decisions
261
273
  ]
@@ -362,6 +374,7 @@ class TypeLLMClient:
362
374
  Both stop the call before its next request to SGLang.
363
375
  """
364
376
  self._last_usage.set(None)
377
+ self._last_thinking.set({})
365
378
  if (context is None) == (state is None):
366
379
  raise ValueError("provide exactly one of context or state")
367
380
  context = state if state is not None else context
@@ -382,6 +395,11 @@ class TypeLLMClient:
382
395
  with call_scope(timeout, cancel) as scope:
383
396
  try:
384
397
  decisions = self.compile_schema(schema)
398
+ count = getattr(self.sglang, "count_tokens", None)
399
+ if callable(count):
400
+ scope.usage.input_tokens = sum(map(count, (
401
+ context, json.dumps(questions if questions is not None else schema, ensure_ascii=False))))
402
+ scope.unmeasured_images = len(encoded_images)
385
403
  # Independent fields run together; depends_on turns the fields into a
386
404
  # graph whose layers run in order.
387
405
  run = (_execute_dependency_decisions if any(d.depends_on is not None for d in decisions)
@@ -396,6 +414,8 @@ class TypeLLMClient:
396
414
  finally:
397
415
  self._last_usage.set(scope.usage)
398
416
  self._last_prompts.set(prompts)
417
+ self._last_thinking.set({decision.name: row["thinking"] for decision, row in zip(decisions, rows)
418
+ if row.get("thinking")})
399
419
 
400
420
  output: dict[str, Any] = {}
401
421
  for decision, row in zip(decisions, rows):
@@ -694,14 +714,24 @@ def _execute_batch_decisions(
694
714
  defer = callable(getattr(client, "prepare_answer_prefixes", None))
695
715
  raw_prompts: list[str] = []
696
716
 
697
- def generation_prompt(parent, content, messages):
717
+ raw_thinking: list[bool | None] = []
718
+ raw_budgets: list[int | None] = []
719
+
720
+ def generation_prompt(parent, content, messages, decision):
721
+ # A field's own thinking settings are passed only when it has them, so
722
+ # clients without per-prompt thinking keep working.
723
+ own = {key: value for key, value in (("thinking", decision.thinking),
724
+ ("thinking_budget", decision.thinking_budget))
725
+ if value is not None}
698
726
  if parent is not None:
699
- prompt = (client.extend_chat_prefix(parent, content, finish_thinking=False)
700
- if defer else client.extend_chat_prefix(parent, content))
727
+ prompt = (client.extend_chat_prefix(parent, content, finish_thinking=False, **own)
728
+ if defer else client.extend_chat_prefix(parent, content, **own))
701
729
  else:
702
- prompt = (client.render_chat(messages, add_generation_prompt=True, finish_thinking=False)
703
- if defer else client.render_chat(messages, add_generation_prompt=True))
730
+ prompt = (client.render_chat(messages, add_generation_prompt=True, finish_thinking=False, **own)
731
+ if defer else client.render_chat(messages, add_generation_prompt=True, **own))
704
732
  raw_prompts.append(prompt)
733
+ raw_thinking.append(decision.thinking)
734
+ raw_budgets.append(decision.thinking_budget)
705
735
  return len(raw_prompts) - 1
706
736
 
707
737
  decision_slots = {}
@@ -712,7 +742,7 @@ def _execute_batch_decisions(
712
742
  {"role": "user", "content": question_content(decision)}
713
743
  ]
714
744
  parent = (parent_prefixes or {}).get(decision.name)
715
- decision_slots[index] = generation_prompt(parent, question_content(decision), messages)
745
+ decision_slots[index] = generation_prompt(parent, question_content(decision), messages, decision)
716
746
  if decision.text_type:
717
747
  text_pending.append((index, decision, messages))
718
748
  continue
@@ -734,7 +764,7 @@ def _execute_batch_decisions(
734
764
  scoring_slots.append(
735
765
  decision_slots[index] if variant.choices == decision.choices else
736
766
  generation_prompt(parent, variant_content,
737
- shared_messages + [{"role": "user", "content": variant_content}]))
767
+ shared_messages + [{"role": "user", "content": variant_content}], variant))
738
768
  scoring_ids.append(ids)
739
769
  scoring_prefills.append(decision.label_prefill)
740
770
  finite_indexes.append(index)
@@ -753,7 +783,10 @@ def _execute_batch_decisions(
753
783
  cache_meta = client.cache_prefix(shared_prefix)
754
784
  LOG.info("batch_shared_prefix_cached_tokens=%s", cache_meta.get("cached_tokens"))
755
785
 
756
- ready = client.prepare_answer_prefixes(raw_prompts) if defer else raw_prompts
786
+ per_field = any(t is not None for t in raw_thinking) or any(b is not None for b in raw_budgets)
787
+ ready = (client.prepare_answer_prefixes(raw_prompts, **({"thinking": raw_thinking, "budgets": raw_budgets}
788
+ if per_field else {}))
789
+ if defer else raw_prompts)
757
790
  prompts = [ready[decision_slots[index]] for index in finite_indexes]
758
791
  scoring_prompts = [ready[slot] + prefill for slot, prefill in zip(scoring_slots, scoring_prefills)]
759
792
  # Open fields continue from {"name": ; history gets the closed object.
@@ -860,6 +893,7 @@ def _execute_batch_decisions(
860
893
 
861
894
  if open_results:
862
895
  finite_rows = iter(zip(results, completed_prompts))
896
+
863
897
  merged_results: list[dict] = []
864
898
  merged_prompts: list[str] = []
865
899
  for index in range(len(decisions)):
@@ -870,6 +904,11 @@ def _execute_batch_decisions(
870
904
  merged_results.append(row)
871
905
  merged_prompts.append(prompt)
872
906
  results, completed_prompts = merged_results, merged_prompts
907
+ # Each field's reasoning, when it thought: what finishing thinking added.
908
+ read = getattr(client, "thinking_text", None)
909
+ for index, row in enumerate(results):
910
+ slot = decision_slots[index]
911
+ row["thinking"] = read(raw_prompts[slot], ready[slot]) if callable(read) and defer else None
873
912
  return results, completed_prompts
874
913
 
875
914
 
@@ -894,9 +933,8 @@ def run_schema(
894
933
  seed: int | None = None,
895
934
  numeric_max_digits: int = 32,
896
935
  tokenizer: str | None = None,
897
- thinking: bool = False,
898
936
  thinking_budget: int | None = None,
899
- text_max_tokens: int = 512,
937
+ text_max_tokens: int = 128,
900
938
  print_final_prompt: bool = False,
901
939
  ) -> dict[str, Any]:
902
940
  client = TypeLLMClient(
@@ -907,7 +945,6 @@ def run_schema(
907
945
  seed=seed,
908
946
  numeric_max_digits=numeric_max_digits,
909
947
  tokenizer=tokenizer,
910
- thinking=thinking,
911
948
  thinking_budget=thinking_budget,
912
949
  text_max_tokens=text_max_tokens,
913
950
  )
@@ -32,6 +32,9 @@ class Decision:
32
32
  return_probabilities: bool = False
33
33
  depends_on: tuple[str, ...] | None = None
34
34
  nullable: bool = False
35
+ # None follows the client's thinking and thinking_budget.
36
+ thinking: bool | None = None
37
+ thinking_budget: int | None = None
35
38
 
36
39
 
37
40
  def _has_duplicates(values: Sequence[Any]) -> bool:
@@ -231,7 +234,13 @@ def compile_json_schema(schema: Mapping[str, Any]) -> list[Decision]:
231
234
  if len(set(dependencies)) != len(dependencies):
232
235
  raise SchemaError(f"depends_on for {decision.name!r} contains duplicates")
233
236
  dependencies = tuple(dependencies)
234
- compiled.append(replace(decision, depends_on=dependencies))
237
+ thinking = field.get("thinking")
238
+ if thinking is not None and type(thinking) is not bool:
239
+ raise SchemaError(f"thinking for {decision.name!r} must be a boolean")
240
+ budget = field.get("thinking_budget")
241
+ if budget is not None and (type(budget) is not int or budget <= 0):
242
+ raise SchemaError(f"thinking_budget for {decision.name!r} must be a positive integer")
243
+ compiled.append(replace(decision, depends_on=dependencies, thinking=thinking, thinking_budget=budget))
235
244
  dependency_layers(compiled)
236
245
  return compiled
237
246
 
@@ -16,6 +16,11 @@ import httpx
16
16
  from .protocol import detect_protocol
17
17
 
18
18
 
19
+ def _overrides(thinking: bool | None, budget: int | None) -> tuple:
20
+ """A prompt's own thinking settings as extra arguments; none when it has none."""
21
+ return () if thinking is None and budget is None else (thinking, budget)
22
+
23
+
19
24
  class SGLangError(RuntimeError):
20
25
  def __init__(self, message: str = "", *, status: int | None = None) -> None:
21
26
  super().__init__(message)
@@ -32,12 +37,21 @@ class GenerationCancelled(RuntimeError):
32
37
 
33
38
  @dataclass
34
39
  class Usage:
35
- """What SGLang reported for the /generate requests of one generate() call."""
40
+ """The tokens of one generate() call.
41
+
42
+ input_tokens is what the caller sent, each part counted once: the context,
43
+ the questions or schema as JSON, and the images as the server expands them.
44
+ The other counts are what SGLang reported for the call's /generate requests,
45
+ where the shared context is part of every prompt.
46
+ """
36
47
 
37
48
  requests: int = 0
38
49
  prompt_tokens: int = 0
39
50
  cached_tokens: int = 0
40
51
  completion_tokens: int = 0
52
+ # The part of completion_tokens that is reasoning, not typed answers.
53
+ thinking_tokens: int = 0
54
+ input_tokens: int = 0
41
55
 
42
56
  def _add(self, response: Any) -> None:
43
57
  self.requests += 1
@@ -51,11 +65,25 @@ class Usage:
51
65
  setattr(self, name, getattr(self, name) + value)
52
66
 
53
67
 
68
+ def _count_thinking(response: Any) -> None:
69
+ """Add a thinking request's generated tokens to this call's usage."""
70
+ scope = _call_scope.get()
71
+ if scope is None:
72
+ return
73
+ for item in response if isinstance(response, list) else [response]:
74
+ meta = item.get("meta_info") if isinstance(item, Mapping) else None
75
+ tokens = meta.get("completion_tokens") if isinstance(meta, Mapping) else None
76
+ if type(tokens) is int:
77
+ scope.usage.thinking_tokens += tokens
78
+
79
+
54
80
  @dataclass
55
81
  class _CallScope:
56
82
  usage: Usage = field(default_factory=Usage)
57
83
  deadline: float | None = None # time.monotonic() when the call runs out of time
58
84
  cancel: threading.Event | None = None
85
+ # Images whose tokens input_tokens still lacks; the first response measures them.
86
+ unmeasured_images: int = 0
59
87
 
60
88
  def expired(self) -> bool:
61
89
  return self.deadline is not None and time.monotonic() >= self.deadline
@@ -140,13 +168,10 @@ class SGLangClient:
140
168
  timeout: float = 120.0,
141
169
  tokenizer: str | None = None,
142
170
  *,
143
- thinking: bool = False,
144
171
  thinking_budget: int | None = None,
145
- text_max_tokens: int = 512,
172
+ text_max_tokens: int = 128,
146
173
  answer_reserve_tokens: int = 64,
147
174
  ) -> None:
148
- if type(thinking) is not bool:
149
- raise ValueError("thinking must be a boolean")
150
175
  if thinking_budget is not None and (type(thinking_budget) is not int or thinking_budget <= 0):
151
176
  raise ValueError("thinking_budget must be a positive integer or None")
152
177
  if type(text_max_tokens) is not int or text_max_tokens <= 0:
@@ -156,7 +181,7 @@ class SGLangClient:
156
181
  self.answer_reserve_tokens = max(answer_reserve_tokens, text_max_tokens)
157
182
  self._context_length_cache: int | None = None
158
183
  self.text_max_tokens = text_max_tokens
159
- self.thinking = thinking
184
+ # The budget of fields that think without their own.
160
185
  self.thinking_budget = thinking_budget
161
186
  self.base_url = base_url.rstrip("/")
162
187
  self.model = model
@@ -335,8 +360,27 @@ class SGLangClient:
335
360
  raise GenerationTimeout("generate() ran out of time") from exc
336
361
  raise
337
362
  scope.usage._add(response)
363
+ if scope.unmeasured_images:
364
+ self._measure_images(scope, payload, response)
338
365
  return response
339
366
 
367
+ def count_tokens(self, text: str) -> int:
368
+ return len(self._get_chat_tokenizer().encode(text, add_special_tokens=False))
369
+
370
+ def _measure_images(self, scope: _CallScope, payload: Mapping[str, Any], response: Any) -> None:
371
+ """Add the images' tokens to input_tokens: the server's prompt count less the local one."""
372
+ text = payload.get("text")
373
+ prompt = text if isinstance(text, str) else text[0] if isinstance(text, list) and text else None
374
+ item = response[0] if isinstance(response, list) and response else response
375
+ meta = item.get("meta_info") if isinstance(item, Mapping) else None
376
+ served = meta.get("prompt_tokens") if isinstance(meta, Mapping) else None
377
+ if not isinstance(prompt, str) or type(served) is not int:
378
+ return
379
+ # The server replaces each placeholder with the image's tokens.
380
+ placeholders = scope.unmeasured_images * self.count_tokens(self.image_placeholder())
381
+ scope.usage.input_tokens += max(0, served - self.count_tokens(prompt) + placeholders)
382
+ scope.unmeasured_images = 0
383
+
340
384
  def _with_images(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
341
385
  images = self._active_images.get()
342
386
  if not images:
@@ -421,8 +465,14 @@ class SGLangClient:
421
465
  *,
422
466
  add_generation_prompt: bool,
423
467
  finish_thinking: bool = True,
468
+ thinking: bool | None = None,
469
+ thinking_budget: int | None = None,
424
470
  ) -> str:
425
- """Render history; optionally finish thinking before constrained decoding."""
471
+ """Render history; optionally finish thinking before constrained decoding.
472
+
473
+ thinking asks this prompt to reason; thinking_budget overrides the client's.
474
+ """
475
+ think = bool(thinking)
426
476
  tokenizer = self._get_chat_tokenizer()
427
477
  if not getattr(tokenizer, "chat_template", None):
428
478
  raise SGLangError("The served tokenizer does not define a chat template")
@@ -435,7 +485,7 @@ class SGLangClient:
435
485
  list(messages),
436
486
  tokenize=False,
437
487
  add_generation_prompt=add_generation_prompt,
438
- enable_thinking=self.thinking and add_generation_prompt,
488
+ enable_thinking=think and add_generation_prompt,
439
489
  )
440
490
  except Exception as exc:
441
491
  raise SGLangError(
@@ -450,16 +500,19 @@ class SGLangClient:
450
500
  if bos and rendered.startswith(bos) and tokenizer.encode("x")[0] == tokenizer.bos_token_id:
451
501
  rendered = rendered[len(bos):]
452
502
  if add_generation_prompt and finish_thinking:
453
- return self._prepare_answer_prefix(rendered)
503
+ return self._prepare_answer_prefix(rendered, *_overrides(thinking, thinking_budget))
454
504
  return rendered
455
505
 
456
- def _prepare_answer_prefix(self, prefix: str) -> str:
506
+ def _prepare_answer_prefix(self, prefix: str, thinking: bool | None = None,
507
+ budget: int | None = None) -> str:
457
508
  protocol = detect_protocol(self._get_chat_tokenizer())
458
- if self.thinking or protocol.has_open_thinking(prefix):
459
- return self._finish_thinking(prefix)
509
+ # A template that always opens thinking reasons whatever the setting.
510
+ if thinking or protocol.has_open_thinking(prefix):
511
+ return self._finish_thinking(prefix, *(() if budget is None else (budget,)))
460
512
  return prefix
461
513
 
462
- def _continuation_parts(self, question: str | None = None) -> tuple[str, str]:
514
+ def _continuation_parts(self, question: str | None = None,
515
+ thinking: bool | None = None) -> tuple[str, str]:
463
516
  # Derive turn delimiters from the actual tokenizer, never hardcode a
464
517
  # model's chat tokens. The marker is confined to this local template probe.
465
518
  marker = "TYPELLM_ASSISTANT_BOUNDARY_8b46c9"
@@ -474,6 +527,7 @@ class SGLangClient:
474
527
  extended = self.render_chat(
475
528
  messages + [{"role": "user", "content": question}],
476
529
  add_generation_prompt=True, finish_thinking=False,
530
+ **({} if thinking is None else {"thinking": thinking}),
477
531
  )
478
532
  if extended.count(marker) != 1:
479
533
  raise SGLangError("Chat template cannot preserve assistant content for KV continuation")
@@ -486,20 +540,47 @@ class SGLangClient:
486
540
  closing, _ = self._continuation_parts()
487
541
  return prompt + answer + closing
488
542
 
489
- def extend_chat_prefix(self, prefix: str, question: str, *, finish_thinking: bool = True) -> str:
490
- _, suffix = self._continuation_parts(question)
543
+ def extend_chat_prefix(self, prefix: str, question: str, *, finish_thinking: bool = True,
544
+ thinking: bool | None = None, thinking_budget: int | None = None) -> str:
545
+ _, suffix = self._continuation_parts(question, thinking)
491
546
  prompt = prefix + suffix
492
- return self._prepare_answer_prefix(prompt) if finish_thinking else prompt
547
+ return (self._prepare_answer_prefix(prompt, *_overrides(thinking, thinking_budget))
548
+ if finish_thinking else prompt)
549
+
550
+ def thinking_text(self, prompt: str, finished: str) -> str | None:
551
+ """The reasoning that finishing thinking added to prompt, or None."""
552
+ if finished == prompt or not finished.startswith(prompt):
553
+ return None
554
+ added = finished[len(prompt):]
555
+ close = detect_protocol(self._get_chat_tokenizer()).thinking_close
556
+ if close and close in added:
557
+ added = added[:added.index(close)]
558
+ return added.strip() or None
559
+
560
+ def prepare_answer_prefixes(
561
+ self,
562
+ prefixes: Sequence[str],
563
+ thinking: Sequence[bool | None] | None = None,
564
+ budgets: Sequence[int | None] | None = None,
565
+ ) -> list[str]:
566
+ """Finish thinking for many generation prompts in one batched request.
493
567
 
494
- def prepare_answer_prefixes(self, prefixes: Sequence[str]) -> list[str]:
495
- """Finish thinking for many generation prompts in one batched request."""
568
+ thinking says which prompts reason and budgets their own budgets (None: the
569
+ client's); only the prompts that think join the batch.
570
+ """
496
571
  protocol = detect_protocol(self._get_chat_tokenizer())
497
- pending = [i for i, p in enumerate(prefixes) if self.thinking or protocol.has_open_thinking(p)]
572
+ wants = [bool(t) for t in (thinking or [None] * len(prefixes))]
573
+ budgets = list(budgets or [None] * len(prefixes))
574
+ pending = [i for i, p in enumerate(prefixes) if wants[i] or protocol.has_open_thinking(p)]
498
575
  finished = list(prefixes)
576
+ # Budgets go only to prompts that have their own, as before for the rest.
577
+ own = [budgets[i] for i in pending]
499
578
  if len(pending) == 1:
500
- finished[pending[0]] = self._finish_thinking(prefixes[pending[0]])
579
+ finished[pending[0]] = self._finish_thinking(prefixes[pending[0]], *([] if own[0] is None else own))
501
580
  elif pending:
502
- for i, value in zip(pending, self._finish_thinking_batch([prefixes[i] for i in pending])):
581
+ done = self._finish_thinking_batch([prefixes[i] for i in pending],
582
+ *(() if all(b is None for b in own) else (own,)))
583
+ for i, value in zip(pending, done):
503
584
  finished[i] = value
504
585
  return finished
505
586
 
@@ -527,12 +608,15 @@ class SGLangClient:
527
608
  self._context_length_cache = min(limits)
528
609
  return self._context_length_cache
529
610
 
530
- def _finish_thinking(self, prefix: str) -> str:
531
- params, image_tokens = self._thinking_params(prefix)
611
+ def _finish_thinking(self, prefix: str, budget: int | None = None) -> str:
612
+ params, image_tokens = self._thinking_params(prefix, budget=budget)
532
613
  response = self._generate({"text": prefix, "sampling_params": params})
614
+ _count_thinking(response)
533
615
  return self._complete_thinking(prefix, response, image_tokens)
534
616
 
535
- def _finish_thinking_batch(self, prefixes: Sequence[str]) -> list[str]:
617
+ def _finish_thinking_batch(self, prefixes: Sequence[str],
618
+ budgets: Sequence[int | None] | None = None) -> list[str]:
619
+ budgets = list(budgets or [None] * len(prefixes))
536
620
  served: list[int | None] = [None] * len(prefixes)
537
621
  if self._active_images.get():
538
622
  response = self._generate({
@@ -543,23 +627,26 @@ class SGLangClient:
543
627
  raise SGLangError("Unexpected prompt-count batch response shape")
544
628
  served = [item.get("meta_info", {}).get("prompt_tokens") if isinstance(item, Mapping) else None
545
629
  for item in response]
546
- planned = [self._thinking_params(prefix, count) for prefix, count in zip(prefixes, served)]
630
+ planned = [self._thinking_params(prefix, count, budget)
631
+ for prefix, count, budget in zip(prefixes, served, budgets)]
547
632
  response = self._generate({
548
633
  "text": list(prefixes),
549
634
  "sampling_params": [params for params, _ in planned],
550
635
  })
551
636
  if not isinstance(response, list) or len(response) != len(prefixes):
552
637
  raise SGLangError("Unexpected thinking batch response shape")
638
+ _count_thinking(response)
553
639
  return [self._complete_thinking(prefix, item, image_tokens)
554
640
  for prefix, item, (_, image_tokens) in zip(prefixes, response, planned)]
555
641
 
556
- def _thinking_params(self, prefix: str, served_tokens: int | None = None) -> tuple[dict[str, Any], int]:
642
+ def _thinking_params(self, prefix: str, served_tokens: int | None = None,
643
+ budget: int | None = None) -> tuple[dict[str, Any], int]:
557
644
  """Return sampling params for one thinking request and its image-token count."""
558
645
  tokenizer = self._get_chat_tokenizer()
559
646
  protocol = detect_protocol(tokenizer)
560
647
  if not protocol.has_open_thinking(prefix):
561
648
  raise SGLangError(
562
- f"thinking=True requires a native chat template ending in an open "
649
+ f"thinking requires a native chat template ending in an open "
563
650
  f"{protocol.thinking_open} block; this template may not support thinking"
564
651
  )
565
652
  forced_end = protocol.forced_close()
@@ -579,7 +666,8 @@ class SGLangClient:
579
666
  available = self._context_length() - prefix_tokens - self.answer_reserve_tokens - closing_tokens - 16
580
667
  if available <= 0:
581
668
  raise SGLangError("Input leaves no room for thinking and the final constrained answer")
582
- limit = available if self.thinking_budget is None else min(available, self.thinking_budget)
669
+ budget = self.thinking_budget if budget is None else budget
670
+ limit = available if budget is None else min(available, budget)
583
671
  return {
584
672
  "max_new_tokens": limit,
585
673
  "temperature": 0.6, "top_p": 0.95, "top_k": 20,
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: typellm
3
- Version: 0.2.1
3
+ Version: 0.2.3
4
4
  Summary: Type-Safe Decoding for Autoregressive LLMs
5
5
  License-Expression: Apache-2.0
6
6
  Project-URL: Homepage, https://typellm.ai/
@@ -168,6 +168,9 @@ result = client.generate(
168
168
  )
169
169
  ```
170
170
 
171
+ Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
172
+ longer answers, or set `maxLength` on a field to cap its characters.
173
+
171
174
  Ask for a numeric answer without enumerating every possible value:
172
175
 
173
176
  ```python
@@ -229,23 +232,28 @@ result = client.generate(
229
232
 
230
233
  ## Thinking mode
231
234
 
232
- Thinking is off by default. Enable it when constructing the client:
235
+ Thinking is off by default. Turn it on for the fields that need it; the others
236
+ answer at once, and the fields that think reason side by side:
233
237
 
234
238
  ```python
235
- client = TypeLLMClient(
236
- "http://127.0.0.1:30000",
237
- model="Qwen/Qwen3.8-27B",
238
- thinking=True,
239
- )
240
- result = client.generate(context=context, questions=questions)
239
+ result = client.generate(context=context, questions={
240
+ "total": {"type": "number"},
241
+ "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
242
+ "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
243
+ "thinking": True, "thinking_budget": 1024},
244
+ })
241
245
  ```
242
246
 
243
- Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
244
- default. When reasoning reaches the budget, TypeLLM closes it and moves on to
245
- the typed answer.
247
+ `thinking_budget` caps a field's reasoning; there is no budget by default.
248
+ `TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
249
+ without its own. When reasoning reaches the budget, TypeLLM closes it and moves
250
+ on to the typed answer.
251
+
252
+ After a call, `client.last_thinking` maps each field that thought to its
253
+ reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
246
254
 
247
- Models with always-on thinking still reason with `thinking=False`;
248
- `thinking_budget` applies to them too. See [Supported models](#supported-models).
255
+ Models with always-on thinking reason on every field; `thinking_budget` applies
256
+ to them too. See [Supported models](#supported-models).
249
257
 
250
258
  ## Image input
251
259
 
@@ -416,11 +424,14 @@ result = client.generate(
416
424
  cancel=cancel, # set it from another thread; raises GenerationCancelled
417
425
  )
418
426
  print(client.last_usage)
419
- # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
427
+ # Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
420
428
  ```
421
429
 
422
- `last_usage` reports the SGLang requests and tokens of the last call made in the
423
- current thread, including a call that failed partway. A timeout or cancel stops
430
+ `last_usage` reports the tokens of the last call made in the current thread,
431
+ including a call that failed partway. `input_tokens` is what you sent, each part
432
+ counted once: the context, the questions as JSON, and the images. The other
433
+ counts are SGLang's for the call's requests, where every prompt carries the
434
+ shared context. A timeout or cancel stops
424
435
  the call before its next SGLang request.
425
436
 
426
437
  ## Cost analysis
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes