typellm 0.2.1__tar.gz → 0.2.2__tar.gz

{typellm-0.2.1 → typellm-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.1
+Version: 0.2.2
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/
@@ -168,6 +168,9 @@ result = client.generate(
 )
 ```
 
+Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
+longer answers, or set `maxLength` on a field to cap its characters.
+
 Ask for a numeric answer without enumerating every possible value:
 
 ```python
@@ -229,23 +232,28 @@ result = client.generate(
 
 ## Thinking mode
 
-Thinking is off by default. Enable it when constructing the client:
+Thinking is off by default. Turn it on for the fields that need it; the others
+answer at once, and the fields that think reason side by side:
 
 ```python
-client = TypeLLMClient(
-    "http://127.0.0.1:30000",
-    model="Qwen/Qwen3.8-27B",
-    thinking=True,
-)
-result = client.generate(context=context, questions=questions)
+result = client.generate(context=context, questions={
+    "total": {"type": "number"},
+    "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
+    "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
+                  "thinking": True, "thinking_budget": 1024},
+})
 ```
 
-Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
-default. When reasoning reaches the budget, TypeLLM closes it and moves on to
-the typed answer.
+`thinking_budget` caps a field's reasoning; there is no budget by default.
+`TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
+without its own. When reasoning reaches the budget, TypeLLM closes it and moves
+on to the typed answer.
+
+After a call, `client.last_thinking` maps each field that thought to its
+reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
 
-Models with always-on thinking still reason with `thinking=False`;
-`thinking_budget` applies to them too. See [Supported models](#supported-models).
+Models with always-on thinking reason on every field; `thinking_budget` applies
+to them too. See [Supported models](#supported-models).
 
 ## Image input
 
@@ -416,7 +424,7 @@ result = client.generate(
     cancel=cancel,   # set it from another thread; raises GenerationCancelled
 )
 print(client.last_usage)
-# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0)
 ```
 
 `last_usage` reports the SGLang requests and tokens of the last call made in the

{typellm-0.2.1 → typellm-0.2.2}/README.md

@@ -153,6 +153,9 @@ result = client.generate(
 )
 ```
 
+Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
+longer answers, or set `maxLength` on a field to cap its characters.
+
 Ask for a numeric answer without enumerating every possible value:
 
 ```python
@@ -214,23 +217,28 @@ result = client.generate(
 
 ## Thinking mode
 
-Thinking is off by default. Enable it when constructing the client:
+Thinking is off by default. Turn it on for the fields that need it; the others
+answer at once, and the fields that think reason side by side:
 
 ```python
-client = TypeLLMClient(
-    "http://127.0.0.1:30000",
-    model="Qwen/Qwen3.8-27B",
-    thinking=True,
-)
-result = client.generate(context=context, questions=questions)
+result = client.generate(context=context, questions={
+    "total": {"type": "number"},
+    "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
+    "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
+                  "thinking": True, "thinking_budget": 1024},
+})
 ```
 
-Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
-default. When reasoning reaches the budget, TypeLLM closes it and moves on to
-the typed answer.
+`thinking_budget` caps a field's reasoning; there is no budget by default.
+`TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
+without its own. When reasoning reaches the budget, TypeLLM closes it and moves
+on to the typed answer.
+
+After a call, `client.last_thinking` maps each field that thought to its
+reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
 
-Models with always-on thinking still reason with `thinking=False`;
-`thinking_budget` applies to them too. See [Supported models](#supported-models).
+Models with always-on thinking reason on every field; `thinking_budget` applies
+to them too. See [Supported models](#supported-models).
 
 ## Image input
 
@@ -401,7 +409,7 @@ result = client.generate(
     cancel=cancel,   # set it from another thread; raises GenerationCancelled
 )
 print(client.last_usage)
-# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0)
 ```
 
 `last_usage` reports the SGLang requests and tokens of the last call made in the

{typellm-0.2.1 → typellm-0.2.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "typellm"
-version = "0.2.1"
+version = "0.2.2"
 description = "Type-Safe Decoding for Autoregressive LLMs"
 readme = "README.md"
 license = "Apache-2.0"

{typellm-0.2.1 → typellm-0.2.2}/tests/test_batching.py

@@ -26,8 +26,8 @@ class FakeServer(SGLangClient):
     Numeric fields decode to 7: the server prefers the end token, then "7".
     """
 
-    def __init__(self, thinking=False):
-        super().__init__(model="fake", thinking=thinking)
+    def __init__(self):
+        super().__init__(model="fake")
         self._chat_tokenizer = ThinkingTokenizer()
         self._context_length_cache = 100_000
         self.payloads = []
@@ -83,6 +83,10 @@ def width(payload):
     return 1 if isinstance(payload["text"], str) else len(payload["text"])
 
 
+def thinks(questions):
+    return {name: {**field, "thinking": True} for name, field in questions.items()}
+
+
 class PrefixWarmupTests(unittest.TestCase):
     MIXED = {"a": {"type": "integer"}, "b": {"type": "number"}, "n": {"type": "string"},
              "k": {"type": "string", "enum": ["x", "y"]}}
@@ -91,8 +95,8 @@ class PrefixWarmupTests(unittest.TestCase):
         for thinking in (False, True):
             with self.subTest(thinking=thinking):
                 client = TypeLLMClient(model="fake")
-                client.sglang = FakeServer(thinking=thinking)
-                client.generate(context="Receipt", questions=self.MIXED)
+                client.sglang = FakeServer()
+                client.generate(context="Receipt", questions=thinks(self.MIXED) if thinking else self.MIXED)
                 warm = client.sglang.requests("count")
                 self.assertEqual(len(warm), 1)
                 # First, before the thinking, number, text and scoring batches.
@@ -107,6 +111,86 @@ class PrefixWarmupTests(unittest.TestCase):
         self.assertEqual(client.sglang.requests("count"), [])
 
 
+class PerFieldThinkingTests(unittest.TestCase):
+    def thinking_prompts(self, server):
+        return [t for p in server.requests("think")
+                for t in ([p["text"]] if isinstance(p["text"], str) else p["text"])]
+
+    def test_only_fields_that_ask_for_it_think(self):
+        client = TypeLLMClient(model="fake")
+        client.sglang = FakeServer()
+        result = client.generate(context="Receipt", questions={
+            "hard": {"type": "boolean", "instructions": "Hard?", "thinking": True},
+            "easy": {"type": "boolean", "instructions": "Easy?"},
+            "count": {"type": "integer"},
+        })
+        self.assertEqual(result, {"hard": True, "easy": True, "count": 7})
+        [prompt] = self.thinking_prompts(client.sglang)  # one batch, one prompt
+        self.assertIn("Hard?", prompt)
+        self.assertTrue(prompt.endswith("<think>\n"))
+        answers = [t for p in client.sglang.payloads if p not in client.sglang.requests("think")
+                   for t in ([p["text"]] if isinstance(p["text"], str) else p["text"])]
+        self.assertTrue(any("Easy?" in t and "<think>" not in t for t in answers))
+
+    def test_thinking_false_is_the_same_as_leaving_it_out(self):
+        client = TypeLLMClient(model="fake")
+        client.sglang = FakeServer()
+        client.generate(context="Receipt", questions={
+            "a": {"type": "boolean", "instructions": "A?", "thinking": True},
+            "b": {"type": "boolean", "instructions": "B?", "thinking": True},
+            "skip": {"type": "boolean", "instructions": "Skip?", "thinking": False},
+        })
+        prompts = self.thinking_prompts(client.sglang)
+        self.assertEqual(len(client.sglang.requests("think")), 1)
+        self.assertEqual(len(prompts), 2)
+        self.assertFalse(any("Skip?" in t for t in prompts))
+
+    def test_each_field_gets_its_own_budget(self):
+        client = TypeLLMClient(model="fake")
+        client.sglang = FakeServer()
+        client.sglang.thinking_budget = 500  # the client's budget, for fields without their own
+        client.generate(context="Receipt", questions={
+            "short": {"type": "boolean", "instructions": "Short?", "thinking": True, "thinking_budget": 100},
+            "long": {"type": "boolean", "instructions": "Long?", "thinking": True, "thinking_budget": 300},
+            "default": {"type": "boolean", "instructions": "Default?", "thinking": True},
+        })
+        [request] = client.sglang.requests("think")
+        budgets = {next(q for q in ("Short?", "Long?", "Default?") if q in t): params["max_new_tokens"]
+                   for t, params in zip(request["text"], request["sampling_params"])}
+        self.assertEqual(budgets, {"Short?": 100, "Long?": 300, "Default?": 500})
+
+    def test_reasoning_and_its_tokens_are_reported_per_field(self):
+        class Metered(FakeServer):
+            def _request(self, path, payload=None, *, allow_text=False):
+                response = super()._request(path, payload, allow_text=allow_text)
+                if path == "/generate":
+                    params = payload["sampling_params"]
+                    thinking = (params[0] if isinstance(params, list) else params).get("stop") == [THINK_STOP]
+                    for item in response if isinstance(response, list) else [response]:
+                        item["meta_info"]["completion_tokens"] = 40 if thinking else 1
+                return response
+
+        client = TypeLLMClient(model="fake")
+        client.sglang = Metered()
+        client.generate(context="Receipt", questions={
+            "hard": {"type": "boolean", "instructions": "Hard?", "thinking": True},
+            "easy": {"type": "boolean", "instructions": "Easy?"},
+        })
+        self.assertEqual(client.last_thinking, {"hard": "Reasoned."})
+        self.assertEqual(client.last_usage.thinking_tokens, 40)
+        self.assertGreater(client.last_usage.completion_tokens, 40)  # answers count too
+        with self.assertRaises(ValueError):
+            client.generate(questions={"x": {"type": "boolean"}})
+        self.assertEqual(client.last_thinking, {})
+
+    def test_invalid_settings_are_schema_errors(self):
+        from typellm import SchemaError
+        for field in ({"type": "boolean", "thinking": "yes"}, {"type": "boolean", "thinking_budget": 0},
+                      {"type": "boolean", "thinking_budget": 2.5}):
+            with self.subTest(field=field), self.assertRaises(SchemaError):
+                TypeLLMClient(model="fake").compile_schema({"type": "object", "properties": {"x": field}})
+
+
 class BatchedThinkingTests(unittest.TestCase):
     QUESTIONS = {
         "flag": {"type": "boolean"},
@@ -117,8 +201,8 @@ class BatchedThinkingTests(unittest.TestCase):
 
     def test_one_thinking_request_covers_every_prompt_in_a_layer(self):
         client = TypeLLMClient(model="fake")
-        client.sglang = FakeServer(thinking=True)
-        result = client.generate(context="Receipt", questions=self.QUESTIONS)
+        client.sglang = FakeServer()
+        result = client.generate(context="Receipt", questions=thinks(self.QUESTIONS))
         self.assertEqual(result, {"flag": True, "count": 7, "name": "blue", "pick": "x"})
         [think] = client.sglang.requests("think")
         # Four fields plus the reversed ordering of "pick".
@@ -130,19 +214,19 @@ class BatchedThinkingTests(unittest.TestCase):
 
     def test_dag_thinks_once_per_layer(self):
         client = TypeLLMClient(model="fake")
-        client.sglang = FakeServer(thinking=True)
-        client.generate(context="Receipt", questions={
+        client.sglang = FakeServer()
+        client.generate(context="Receipt", questions=thinks({
             "a": {"type": "boolean"}, "b": {"type": "integer"},
             "c": {"type": "boolean", "depends_on": ["a", "b"]},
             "d": {"type": "string", "depends_on": ["a"]},
-        })
+        }))
         self.assertEqual([width(p) for p in client.sglang.requests("think")], [2, 2])
 
     def test_images_count_prompt_tokens_in_one_batch(self):
         client = TypeLLMClient(model="fake")
-        client.sglang = FakeServer(thinking=True)
+        client.sglang = FakeServer()
         client.generate(context="Receipt", images=[PNG],
-                        questions={"a": {"type": "boolean"}, "b": {"type": "boolean"}})
+                        questions=thinks({"a": {"type": "boolean"}, "b": {"type": "boolean"}}))
         counts = [p for p in client.sglang.requests("count") if width(p) > 1]
         self.assertEqual([width(p) for p in counts], [2])
         [think] = client.sglang.requests("think")

{typellm-0.2.1 → typellm-0.2.2}/tests/test_dependencies.py

@@ -102,7 +102,7 @@ class DependencyTests(unittest.TestCase):
         self.assertTrue(client.last_prompts[1].startswith(client.last_prompts[0]))
 
     def test_always_thinking_template_continuation_runs_once(self):
-        client = SGLangClient(thinking=False)
+        client = SGLangClient()
 
         class Tokenizer:
             chat_template = 'test'

{typellm-0.2.1 → typellm-0.2.2}/tests/test_images.py

@@ -255,7 +255,6 @@ class ImageRequestTests(unittest.TestCase):
 
     def test_thinking_budget_counts_image_tokens_on_the_server(self):
         client = FakeServerClient()
-        client.thinking = True
         client.thinking_budget = None
         client._context_length_cache = 1000
         prefix = VisionTokenizer().apply_chat_template(

{typellm-0.2.1 → typellm-0.2.2}/tests/test_protocols.py

@@ -39,8 +39,8 @@ class ProtocolTokenizer(FakeChatTokenizer):
 
 
 class ProtocolTests(unittest.TestCase):
-    def client(self, family, thinking=False, response=None):
-        client = SGLangClient(thinking=thinking, thinking_budget=64)
+    def client(self, family, response=None):
+        client = SGLangClient(thinking_budget=64)
         client._chat_tokenizer = ProtocolTokenizer(family)
         client._context_length_cache = 8192
         end = '</think>'
@@ -59,9 +59,9 @@ class ProtocolTests(unittest.TestCase):
 
     def test_gemma_is_rejected_before_inference(self):
         for thinking in (False, True):
-            client = self.client('gemma', thinking)
+            client = self.client('gemma')
             with self.assertRaisesRegex(SGLangError, 'Gemma 4.*not supported'):
-                client.render_chat([], add_generation_prompt=True)
+                client.render_chat([], add_generation_prompt=True, thinking=thinking)
             client._request.assert_not_called()
 
     def test_native_turn_stop_recovers_preserved_and_filtered_terminators(self):
@@ -69,7 +69,7 @@ class ProtocolTests(unittest.TestCase):
             for matched in (ending, token):
                 for suffix in ('', ending + '\n'):
                     with self.subTest(family=family, matched=matched, suffix=suffix):
-                        client = self.client(family, True, {
+                        client = self.client(family, {
                             'text': 'Keep reasoning' + suffix,
                             'meta_info': {'finish_reason': {'type': 'stop', 'matched': matched}},
                         })
@@ -91,15 +91,15 @@ class ProtocolTests(unittest.TestCase):
             ('Partial<|im_end|>', {'type': 'error', 'matched': 3}),
         ]:
             with self.subTest(text=text, finish=finish):
-                client = self.client('minicpm5', True, {'text': text, 'meta_info': {'finish_reason': finish}})
+                client = self.client('minicpm5', {'text': text, 'meta_info': {'finish_reason': finish}})
                 with self.assertRaises(SGLangError):
                     client._finish_thinking('<think>\n')
 
     def test_minicpm5_switch_and_ring_always_thinking(self):
         for family in ('minicpm5', 'ring'):
             for thinking in (False, True):
-                client = self.client(family, thinking)
-                prompt = client.render_chat([], add_generation_prompt=True)
+                client = self.client(family)
+                prompt = client.render_chat([], add_generation_prompt=True, thinking=thinking)
                 expected = thinking or family == 'ring'
                 self.assertEqual(client._request.called, expected)
                 parent = client.complete_chat_prefix(prompt, 'A')
@@ -108,9 +108,9 @@ class ProtocolTests(unittest.TestCase):
 
     def test_non_thinking_templates_reject_requested_thinking(self):
         for family in ('ling',):
-            client = self.client(family, True)
+            client = self.client(family)
             with self.assertRaisesRegex(SGLangError, 'may not support thinking'):
-                client.render_chat([], add_generation_prompt=True)
+                client.render_chat([], add_generation_prompt=True, thinking=True)
             client._request.assert_not_called()
 
 

{typellm-0.2.1 → typellm-0.2.2}/tests/test_text.py

@@ -63,13 +63,13 @@ class TextTests(unittest.TestCase):
                 TypeLLMClient(text_max_tokens=budget)
 
     def test_thinking_then_constrained_text(self):
-        client=TypeLLMClient(thinking=True)
+        client=TypeLLMClient()
         client.sglang._context_length_cache=8192
         client.sglang._chat_tokenizer=FakeChatTokenizer()
         client.sglang._chat_tokenizer.apply_chat_template=lambda *args,**kw: "assistant\n<think>\n" if kw["enable_thinking"] else "completed"
         responses=[{'text':'brief</think>'},[{'text':' "done"}','meta_info':{'finish_reason':{'type':'stop'}}}]]
         with patch.object(client.sglang,'_request',side_effect=responses) as request:
-            self.assertEqual(client.generate(state='x',questions={'t':{'type':'string'}}),{'t':'done'})
+            self.assertEqual(client.generate(state='x',questions={'t':{'type':'string','thinking':True}}),{'t':'done'})
             self.assertIn('</think>',request.call_args.args[1]['text'][0])
             self.assertTrue(request.call_args.args[1]['text'][0].endswith('{"t":'))
             self.assertIn('regex',request.call_args.args[1]['sampling_params'][0])

{typellm-0.2.1 → typellm-0.2.2}/tests/test_typellm.py

@@ -427,7 +427,7 @@ class ThinkingTests(unittest.TestCase):
         from typellm import run_schema
         self.assertIsNone(SGLangClient().thinking_budget)
         self.assertIsNone(TypeLLMClient().sglang.thinking_budget)
-        client = SGLangClient(thinking=True)
+        client = SGLangClient()
         client._chat_tokenizer = FakeChatTokenizer()
         client._context_length_cache = 8192
         client._request = Mock(return_value={"text": "Done.</think>"})
@@ -443,7 +443,7 @@ class ThinkingTests(unittest.TestCase):
 
     def make_client(self, response=None):
         from unittest.mock import Mock
-        client = SGLangClient(thinking=True, thinking_budget=128)
+        client = SGLangClient(thinking_budget=128)
         client._context_length_cache = 8192
         tokenizer = FakeChatTokenizer()
         tokenizer.apply_chat_template = Mock(return_value="assistant\n<think>\n")
@@ -453,7 +453,7 @@ class ThinkingTests(unittest.TestCase):
 
     def test_thinking_stops_before_constrained_answer(self):
         client = self.make_client()
-        prompt = client.render_chat([], add_generation_prompt=True)
+        prompt = client.render_chat([], add_generation_prompt=True, thinking=True)
         self.assertEqual(prompt, "assistant\n<think>\nWork done. </think>\n\n")
         self.assertNotIn("illegal answer", prompt)
         params = client._request.call_args.args[1]["sampling_params"]
@@ -464,7 +464,7 @@ class ThinkingTests(unittest.TestCase):
 
     def test_length_stop_forces_closure_and_retains_reasoning(self):
         client = self.make_client({"text": "Partial reasoning", "meta_info": {"finish_reason": {"type": "length"}}})
-        prompt = client.render_chat([], add_generation_prompt=True)
+        prompt = client.render_chat([], add_generation_prompt=True, thinking=True)
         self.assertIn("Partial reasoning", prompt)
         self.assertIn("I will now give the final answer.", prompt)
         self.assertTrue(prompt.endswith("</think>\n\n"))
@@ -499,7 +499,7 @@ class ThinkingTests(unittest.TestCase):
         for text in ("Partial", "Partial</think>"):
             client = self.make_client({"text": text, "meta_info": {"finish_reason": {"type": "abort"}}})
             with self.assertRaisesRegex(SGLangError, "aborted"):
-                client.render_chat([], add_generation_prompt=True)
+                client.render_chat([], add_generation_prompt=True, thinking=True)
 
     def test_forced_thinking_keeps_all_final_decoders(self):
         stops = [{"type": "length"}, {"type": "stop", "matched": "<|im_end|>"}]
@@ -528,20 +528,19 @@ class ThinkingTests(unittest.TestCase):
             client = self.make_client()
             client._request.return_value = response
             with self.subTest(response=response), self.assertRaises(SGLangError):
-                client.render_chat([], add_generation_prompt=True)
+                client.render_chat([], add_generation_prompt=True, thinking=True)
 
     def test_unsupported_template_does_not_generate(self):
         from typellm import SGLangError
         client = self.make_client()
         client._chat_tokenizer.apply_chat_template.return_value = "<think></think>"
         with self.assertRaisesRegex(SGLangError, "native chat template"):
-            client.render_chat([], add_generation_prompt=True)
+            client.render_chat([], add_generation_prompt=True, thinking=True)
         client._request.assert_not_called()
 
     def test_always_thinking_template_reasons_even_when_thinking_is_off(self):
         client = self.make_client()
-        client.thinking = False
-        prompt = client.render_chat([], add_generation_prompt=True)
+        prompt = client.render_chat([], add_generation_prompt=True)  # thinking not asked for
         self.assertEqual(prompt, "assistant\n<think>\nWork done. </think>\n\n")
         self.assertFalse(client._chat_tokenizer.apply_chat_template.call_args.kwargs["enable_thinking"])
 
@@ -552,12 +551,11 @@ class ThinkingTests(unittest.TestCase):
         self.assertFalse(client._chat_tokenizer.apply_chat_template.call_args.kwargs["enable_thinking"])
 
     def test_public_configuration_and_validation(self):
-        self.assertFalse(TypeLLMClient().sglang.thinking)
-        client = TypeLLMClient(thinking=True, thinking_budget=256)
-        self.assertTrue(client.sglang.thinking)
+        client = TypeLLMClient(thinking_budget=256)
         self.assertEqual(client.sglang.thinking_budget, 256)
-        for kwargs in [{"thinking": "false"}, {"thinking": 1}, {"thinking_budget": 0},
-                       {"thinking_budget": True}, {"thinking_budget": 1.5}]:
+        with self.assertRaises(TypeError):
+            TypeLLMClient(thinking=True)  # thinking is chosen per field
+        for kwargs in [{"thinking_budget": 0}, {"thinking_budget": True}, {"thinking_budget": 1.5}]:
             with self.subTest(kwargs=kwargs), self.assertRaises(ValueError):
                 TypeLLMClient(**kwargs)
 

{typellm-0.2.1 → typellm-0.2.2}/typellm/runtime.py

@@ -63,6 +63,8 @@ class Choice:
     return_probabilities: bool = False
     depends_on: tuple[str, ...] | None = None
     nullable: bool = False
+    thinking: bool | None = None
+    thinking_budget: int | None = None
 
     def __post_init__(self) -> None:
         if not self.choices and self.numeric_type is None and not self.text_type:
@@ -154,9 +156,8 @@ class TypeLLMClient:
         label_pool: Sequence[str] | None = None,
         numeric_max_digits: int = 32,
         tokenizer: str | None = None,
-        thinking: bool = False,
         thinking_budget: int | None = None,
-        text_max_tokens: int = 512,
+        text_max_tokens: int = 128,
     ) -> None:
         _validate_decoding(mode, temperature)
         if type(numeric_max_digits) is not int or numeric_max_digits <= 0:
@@ -166,7 +167,6 @@ class TypeLLMClient:
             model,
             timeout,
             tokenizer=tokenizer,
-            thinking=thinking,
             thinking_budget=thinking_budget,
             text_max_tokens=text_max_tokens,
             answer_reserve_tokens=numeric_max_digits + 3,
@@ -184,23 +184,33 @@ class TypeLLMClient:
         self._last_usage: ContextVar[Usage | None] = ContextVar(
             f"typellm_last_usage_{id(self)}", default=None
         )
+        self._last_thinking: ContextVar[dict[str, str]] = ContextVar(
+            f"typellm_last_thinking_{id(self)}", default={}
+        )
 
     def __getstate__(self) -> dict[str, Any]:
         state = self.__dict__.copy()
         del state["_last_prompts"]
         del state["_last_usage"]
+        del state["_last_thinking"]
         return state
 
     def __setstate__(self, state: dict[str, Any]) -> None:
         self.__dict__.update(state)
         self._last_prompts = ContextVar(f"typellm_last_prompts_{id(self)}", default=[])
         self._last_usage = ContextVar(f"typellm_last_usage_{id(self)}", default=None)
+        self._last_thinking = ContextVar(f"typellm_last_thinking_{id(self)}", default={})
 
     @property
     def last_prompts(self) -> list[str]:
         """Final prompts of the last generate() call made in this thread or task."""
         return self._last_prompts.get()
 
+    @property
+    def last_thinking(self) -> dict[str, str]:
+        """Reasoning of the fields that thought in the last call in this thread or task."""
+        return self._last_thinking.get()
+
     @property
     def last_usage(self) -> Usage | None:
         """Requests and tokens of the last generate() call made in this thread or task.
@@ -256,6 +266,8 @@ class TypeLLMClient:
                     return_probabilities=item.return_probabilities,
                     depends_on=item.depends_on,
                     nullable=item.nullable,
+                    thinking=item.thinking,
+                    thinking_budget=item.thinking_budget,
                 )
                 for item in decisions
             ]
@@ -362,6 +374,7 @@ class TypeLLMClient:
         Both stop the call before its next request to SGLang.
         """
         self._last_usage.set(None)
+        self._last_thinking.set({})
         if (context is None) == (state is None):
             raise ValueError("provide exactly one of context or state")
         context = state if state is not None else context
@@ -396,6 +409,8 @@ class TypeLLMClient:
             finally:
                 self._last_usage.set(scope.usage)
         self._last_prompts.set(prompts)
+        self._last_thinking.set({decision.name: row["thinking"] for decision, row in zip(decisions, rows)
+                                 if row.get("thinking")})
 
         output: dict[str, Any] = {}
         for decision, row in zip(decisions, rows):
@@ -694,14 +709,24 @@ def _execute_batch_decisions(
     defer = callable(getattr(client, "prepare_answer_prefixes", None))
     raw_prompts: list[str] = []
 
-    def generation_prompt(parent, content, messages):
+    raw_thinking: list[bool | None] = []
+    raw_budgets: list[int | None] = []
+
+    def generation_prompt(parent, content, messages, decision):
+        # A field's own thinking settings are passed only when it has them, so
+        # clients without per-prompt thinking keep working.
+        own = {key: value for key, value in (("thinking", decision.thinking),
+                                             ("thinking_budget", decision.thinking_budget))
+               if value is not None}
         if parent is not None:
-            prompt = (client.extend_chat_prefix(parent, content, finish_thinking=False)
-                      if defer else client.extend_chat_prefix(parent, content))
+            prompt = (client.extend_chat_prefix(parent, content, finish_thinking=False, **own)
+                      if defer else client.extend_chat_prefix(parent, content, **own))
         else:
-            prompt = (client.render_chat(messages, add_generation_prompt=True, finish_thinking=False)
-                      if defer else client.render_chat(messages, add_generation_prompt=True))
+            prompt = (client.render_chat(messages, add_generation_prompt=True, finish_thinking=False, **own)
+                      if defer else client.render_chat(messages, add_generation_prompt=True, **own))
         raw_prompts.append(prompt)
+        raw_thinking.append(decision.thinking)
+        raw_budgets.append(decision.thinking_budget)
         return len(raw_prompts) - 1
 
     decision_slots = {}
@@ -712,7 +737,7 @@ def _execute_batch_decisions(
             {"role": "user", "content": question_content(decision)}
         ]
         parent = (parent_prefixes or {}).get(decision.name)
-        decision_slots[index] = generation_prompt(parent, question_content(decision), messages)
+        decision_slots[index] = generation_prompt(parent, question_content(decision), messages, decision)
         if decision.text_type:
             text_pending.append((index, decision, messages))
             continue
@@ -734,7 +759,7 @@ def _execute_batch_decisions(
             scoring_slots.append(
                 decision_slots[index] if variant.choices == decision.choices else
                 generation_prompt(parent, variant_content,
-                                  shared_messages + [{"role": "user", "content": variant_content}]))
+                                  shared_messages + [{"role": "user", "content": variant_content}], variant))
             scoring_ids.append(ids)
             scoring_prefills.append(decision.label_prefill)
         finite_indexes.append(index)
@@ -753,7 +778,10 @@ def _execute_batch_decisions(
         cache_meta = client.cache_prefix(shared_prefix)
         LOG.info("batch_shared_prefix_cached_tokens=%s", cache_meta.get("cached_tokens"))
 
-    ready = client.prepare_answer_prefixes(raw_prompts) if defer else raw_prompts
+    per_field = any(t is not None for t in raw_thinking) or any(b is not None for b in raw_budgets)
+    ready = (client.prepare_answer_prefixes(raw_prompts, **({"thinking": raw_thinking, "budgets": raw_budgets}
+                                                           if per_field else {}))
+             if defer else raw_prompts)
     prompts = [ready[decision_slots[index]] for index in finite_indexes]
     scoring_prompts = [ready[slot] + prefill for slot, prefill in zip(scoring_slots, scoring_prefills)]
     # Open fields continue from {"name": ; history gets the closed object.
@@ -860,6 +888,7 @@ def _execute_batch_decisions(
 
     if open_results:
         finite_rows = iter(zip(results, completed_prompts))
+
         merged_results: list[dict] = []
         merged_prompts: list[str] = []
         for index in range(len(decisions)):
@@ -870,6 +899,11 @@ def _execute_batch_decisions(
             merged_results.append(row)
             merged_prompts.append(prompt)
         results, completed_prompts = merged_results, merged_prompts
+    # Each field's reasoning, when it thought: what finishing thinking added.
+    read = getattr(client, "thinking_text", None)
+    for index, row in enumerate(results):
+        slot = decision_slots[index]
+        row["thinking"] = read(raw_prompts[slot], ready[slot]) if callable(read) and defer else None
     return results, completed_prompts
 
 
@@ -894,9 +928,8 @@ def run_schema(
     seed: int | None = None,
     numeric_max_digits: int = 32,
     tokenizer: str | None = None,
-    thinking: bool = False,
     thinking_budget: int | None = None,
-    text_max_tokens: int = 512,
+    text_max_tokens: int = 128,
     print_final_prompt: bool = False,
 ) -> dict[str, Any]:
     client = TypeLLMClient(
@@ -907,7 +940,6 @@ def run_schema(
         seed=seed,
         numeric_max_digits=numeric_max_digits,
         tokenizer=tokenizer,
-        thinking=thinking,
         thinking_budget=thinking_budget,
         text_max_tokens=text_max_tokens,
     )

{typellm-0.2.1 → typellm-0.2.2}/typellm/schema.py

@@ -32,6 +32,9 @@ class Decision:
     return_probabilities: bool = False
     depends_on: tuple[str, ...] | None = None
     nullable: bool = False
+    # None follows the client's thinking and thinking_budget.
+    thinking: bool | None = None
+    thinking_budget: int | None = None
 
 
 def _has_duplicates(values: Sequence[Any]) -> bool:
@@ -231,7 +234,13 @@ def compile_json_schema(schema: Mapping[str, Any]) -> list[Decision]:
             if len(set(dependencies)) != len(dependencies):
                 raise SchemaError(f"depends_on for {decision.name!r} contains duplicates")
             dependencies = tuple(dependencies)
-        compiled.append(replace(decision, depends_on=dependencies))
+        thinking = field.get("thinking")
+        if thinking is not None and type(thinking) is not bool:
+            raise SchemaError(f"thinking for {decision.name!r} must be a boolean")
+        budget = field.get("thinking_budget")
+        if budget is not None and (type(budget) is not int or budget <= 0):
+            raise SchemaError(f"thinking_budget for {decision.name!r} must be a positive integer")
+        compiled.append(replace(decision, depends_on=dependencies, thinking=thinking, thinking_budget=budget))
     dependency_layers(compiled)
     return compiled
 

{typellm-0.2.1 → typellm-0.2.2}/typellm/sglang.py

@@ -16,6 +16,11 @@ import httpx
 from .protocol import detect_protocol
 
 
+def _overrides(thinking: bool | None, budget: int | None) -> tuple:
+    """A prompt's own thinking settings as extra arguments; none when it has none."""
+    return () if thinking is None and budget is None else (thinking, budget)
+
+
 class SGLangError(RuntimeError):
     def __init__(self, message: str = "", *, status: int | None = None) -> None:
         super().__init__(message)
@@ -38,6 +43,8 @@ class Usage:
     prompt_tokens: int = 0
     cached_tokens: int = 0
     completion_tokens: int = 0
+    # The part of completion_tokens that is reasoning, not typed answers.
+    thinking_tokens: int = 0
 
     def _add(self, response: Any) -> None:
         self.requests += 1
@@ -51,6 +58,18 @@ class Usage:
                     setattr(self, name, getattr(self, name) + value)
 
 
+def _count_thinking(response: Any) -> None:
+    """Add a thinking request's generated tokens to this call's usage."""
+    scope = _call_scope.get()
+    if scope is None:
+        return
+    for item in response if isinstance(response, list) else [response]:
+        meta = item.get("meta_info") if isinstance(item, Mapping) else None
+        tokens = meta.get("completion_tokens") if isinstance(meta, Mapping) else None
+        if type(tokens) is int:
+            scope.usage.thinking_tokens += tokens
+
+
 @dataclass
 class _CallScope:
     usage: Usage = field(default_factory=Usage)
@@ -140,13 +159,10 @@ class SGLangClient:
         timeout: float = 120.0,
         tokenizer: str | None = None,
         *,
-        thinking: bool = False,
         thinking_budget: int | None = None,
-        text_max_tokens: int = 512,
+        text_max_tokens: int = 128,
         answer_reserve_tokens: int = 64,
     ) -> None:
-        if type(thinking) is not bool:
-            raise ValueError("thinking must be a boolean")
         if thinking_budget is not None and (type(thinking_budget) is not int or thinking_budget <= 0):
             raise ValueError("thinking_budget must be a positive integer or None")
         if type(text_max_tokens) is not int or text_max_tokens <= 0:
@@ -156,7 +172,7 @@ class SGLangClient:
         self.answer_reserve_tokens = max(answer_reserve_tokens, text_max_tokens)
         self._context_length_cache: int | None = None
         self.text_max_tokens = text_max_tokens
-        self.thinking = thinking
+        # The budget of fields that think without their own.
         self.thinking_budget = thinking_budget
         self.base_url = base_url.rstrip("/")
         self.model = model
@@ -421,8 +437,14 @@ class SGLangClient:
         *,
         add_generation_prompt: bool,
         finish_thinking: bool = True,
+        thinking: bool | None = None,
+        thinking_budget: int | None = None,
     ) -> str:
-        """Render history; optionally finish thinking before constrained decoding."""
+        """Render history; optionally finish thinking before constrained decoding.
+
+        thinking asks this prompt to reason; thinking_budget overrides the client's.
+        """
+        think = bool(thinking)
         tokenizer = self._get_chat_tokenizer()
         if not getattr(tokenizer, "chat_template", None):
             raise SGLangError("The served tokenizer does not define a chat template")
@@ -435,7 +457,7 @@ class SGLangClient:
                 list(messages),
                 tokenize=False,
                 add_generation_prompt=add_generation_prompt,
-                enable_thinking=self.thinking and add_generation_prompt,
+                enable_thinking=think and add_generation_prompt,
             )
         except Exception as exc:
             raise SGLangError(
@@ -450,16 +472,19 @@ class SGLangClient:
         if bos and rendered.startswith(bos) and tokenizer.encode("x")[0] == tokenizer.bos_token_id:
             rendered = rendered[len(bos):]
         if add_generation_prompt and finish_thinking:
-            return self._prepare_answer_prefix(rendered)
+            return self._prepare_answer_prefix(rendered, *_overrides(thinking, thinking_budget))
         return rendered
 
-    def _prepare_answer_prefix(self, prefix: str) -> str:
+    def _prepare_answer_prefix(self, prefix: str, thinking: bool | None = None,
+                               budget: int | None = None) -> str:
         protocol = detect_protocol(self._get_chat_tokenizer())
-        if self.thinking or protocol.has_open_thinking(prefix):
-            return self._finish_thinking(prefix)
+        # A template that always opens thinking reasons whatever the setting.
+        if thinking or protocol.has_open_thinking(prefix):
+            return self._finish_thinking(prefix, *(() if budget is None else (budget,)))
         return prefix
 
-    def _continuation_parts(self, question: str | None = None) -> tuple[str, str]:
+    def _continuation_parts(self, question: str | None = None,
+                            thinking: bool | None = None) -> tuple[str, str]:
         # Derive turn delimiters from the actual tokenizer, never hardcode a
         # model's chat tokens. The marker is confined to this local template probe.
         marker = "TYPELLM_ASSISTANT_BOUNDARY_8b46c9"
@@ -474,6 +499,7 @@ class SGLangClient:
         extended = self.render_chat(
             messages + [{"role": "user", "content": question}],
             add_generation_prompt=True, finish_thinking=False,
+            **({} if thinking is None else {"thinking": thinking}),
         )
         if extended.count(marker) != 1:
             raise SGLangError("Chat template cannot preserve assistant content for KV continuation")
@@ -486,20 +512,47 @@ class SGLangClient:
         closing, _ = self._continuation_parts()
         return prompt + answer + closing
 
-    def extend_chat_prefix(self, prefix: str, question: str, *, finish_thinking: bool = True) -> str:
-        _, suffix = self._continuation_parts(question)
+    def extend_chat_prefix(self, prefix: str, question: str, *, finish_thinking: bool = True,
+                           thinking: bool | None = None, thinking_budget: int | None = None) -> str:
+        _, suffix = self._continuation_parts(question, thinking)
         prompt = prefix + suffix
-        return self._prepare_answer_prefix(prompt) if finish_thinking else prompt
+        return (self._prepare_answer_prefix(prompt, *_overrides(thinking, thinking_budget))
+                if finish_thinking else prompt)
+
+    def thinking_text(self, prompt: str, finished: str) -> str | None:
+        """The reasoning that finishing thinking added to prompt, or None."""
+        if finished == prompt or not finished.startswith(prompt):
+            return None
+        added = finished[len(prompt):]
+        close = detect_protocol(self._get_chat_tokenizer()).thinking_close
+        if close and close in added:
+            added = added[:added.index(close)]
+        return added.strip() or None
+
+    def prepare_answer_prefixes(
+        self,
+        prefixes: Sequence[str],
+        thinking: Sequence[bool | None] | None = None,
+        budgets: Sequence[int | None] | None = None,
+    ) -> list[str]:
+        """Finish thinking for many generation prompts in one batched request.
 
-    def prepare_answer_prefixes(self, prefixes: Sequence[str]) -> list[str]:
-        """Finish thinking for many generation prompts in one batched request."""
+        thinking says which prompts reason and budgets their own budgets (None: the
+        client's); only the prompts that think join the batch.
+        """
         protocol = detect_protocol(self._get_chat_tokenizer())
-        pending = [i for i, p in enumerate(prefixes) if self.thinking or protocol.has_open_thinking(p)]
+        wants = [bool(t) for t in (thinking or [None] * len(prefixes))]
+        budgets = list(budgets or [None] * len(prefixes))
+        pending = [i for i, p in enumerate(prefixes) if wants[i] or protocol.has_open_thinking(p)]
         finished = list(prefixes)
+        # Budgets go only to prompts that have their own, as before for the rest.
+        own = [budgets[i] for i in pending]
         if len(pending) == 1:
-            finished[pending[0]] = self._finish_thinking(prefixes[pending[0]])
+            finished[pending[0]] = self._finish_thinking(prefixes[pending[0]], *([] if own[0] is None else own))
         elif pending:
-            for i, value in zip(pending, self._finish_thinking_batch([prefixes[i] for i in pending])):
+            done = self._finish_thinking_batch([prefixes[i] for i in pending],
+                                               *(() if all(b is None for b in own) else (own,)))
+            for i, value in zip(pending, done):
                 finished[i] = value
         return finished
 
@@ -527,12 +580,15 @@ class SGLangClient:
             self._context_length_cache = min(limits)
         return self._context_length_cache
 
-    def _finish_thinking(self, prefix: str) -> str:
-        params, image_tokens = self._thinking_params(prefix)
+    def _finish_thinking(self, prefix: str, budget: int | None = None) -> str:
+        params, image_tokens = self._thinking_params(prefix, budget=budget)
         response = self._generate({"text": prefix, "sampling_params": params})
+        _count_thinking(response)
         return self._complete_thinking(prefix, response, image_tokens)
 
-    def _finish_thinking_batch(self, prefixes: Sequence[str]) -> list[str]:
+    def _finish_thinking_batch(self, prefixes: Sequence[str],
+                               budgets: Sequence[int | None] | None = None) -> list[str]:
+        budgets = list(budgets or [None] * len(prefixes))
         served: list[int | None] = [None] * len(prefixes)
         if self._active_images.get():
             response = self._generate({
@@ -543,23 +599,26 @@ class SGLangClient:
                 raise SGLangError("Unexpected prompt-count batch response shape")
             served = [item.get("meta_info", {}).get("prompt_tokens") if isinstance(item, Mapping) else None
                       for item in response]
-        planned = [self._thinking_params(prefix, count) for prefix, count in zip(prefixes, served)]
+        planned = [self._thinking_params(prefix, count, budget)
+                   for prefix, count, budget in zip(prefixes, served, budgets)]
         response = self._generate({
             "text": list(prefixes),
             "sampling_params": [params for params, _ in planned],
         })
         if not isinstance(response, list) or len(response) != len(prefixes):
             raise SGLangError("Unexpected thinking batch response shape")
+        _count_thinking(response)
         return [self._complete_thinking(prefix, item, image_tokens)
                 for prefix, item, (_, image_tokens) in zip(prefixes, response, planned)]
 
-    def _thinking_params(self, prefix: str, served_tokens: int | None = None) -> tuple[dict[str, Any], int]:
+    def _thinking_params(self, prefix: str, served_tokens: int | None = None,
+                         budget: int | None = None) -> tuple[dict[str, Any], int]:
         """Return sampling params for one thinking request and its image-token count."""
         tokenizer = self._get_chat_tokenizer()
         protocol = detect_protocol(tokenizer)
         if not protocol.has_open_thinking(prefix):
             raise SGLangError(
-                f"thinking=True requires a native chat template ending in an open "
+                f"thinking requires a native chat template ending in an open "
                 f"{protocol.thinking_open} block; this template may not support thinking"
             )
         forced_end = protocol.forced_close()
@@ -579,7 +638,8 @@ class SGLangClient:
         available = self._context_length() - prefix_tokens - self.answer_reserve_tokens - closing_tokens - 16
         if available <= 0:
             raise SGLangError("Input leaves no room for thinking and the final constrained answer")
-        limit = available if self.thinking_budget is None else min(available, self.thinking_budget)
+        budget = self.thinking_budget if budget is None else budget
+        limit = available if budget is None else min(available, budget)
         return {
             "max_new_tokens": limit,
             "temperature": 0.6, "top_p": 0.95, "top_k": 20,

{typellm-0.2.1 → typellm-0.2.2}/typellm.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.1
+Version: 0.2.2
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/
@@ -168,6 +168,9 @@ result = client.generate(
 )
 ```
 
+Free text stops after 128 tokens. Pass `text_max_tokens=` to the client for
+longer answers, or set `maxLength` on a field to cap its characters.
+
 Ask for a numeric answer without enumerating every possible value:
 
 ```python
@@ -229,23 +232,28 @@ result = client.generate(
 
 ## Thinking mode
 
-Thinking is off by default. Enable it when constructing the client:
+Thinking is off by default. Turn it on for the fields that need it; the others
+answer at once, and the fields that think reason side by side:
 
 ```python
-client = TypeLLMClient(
-    "http://127.0.0.1:30000",
-    model="Qwen/Qwen3.8-27B",
-    thinking=True,
-)
-result = client.generate(context=context, questions=questions)
+result = client.generate(context=context, questions={
+    "total": {"type": "number"},
+    "category": {"type": "string", "enum": ["meal", "travel", "equipment"]},
+    "policy_ok": {"type": "boolean", "instructions": "Does it meet the travel policy?",
+                  "thinking": True, "thinking_budget": 1024},
+})
 ```
 
-Set `thinking_budget=2048` to cap reasoning per field; there is no budget by
-default. When reasoning reaches the budget, TypeLLM closes it and moves on to
-the typed answer.
+`thinking_budget` caps a field's reasoning; there is no budget by default.
+`TypeLLMClient(..., thinking_budget=2048)` sets one for every field that thinks
+without its own. When reasoning reaches the budget, TypeLLM closes it and moves
+on to the typed answer.
+
+After a call, `client.last_thinking` maps each field that thought to its
+reasoning, and `client.last_usage.thinking_tokens` counts the reasoning tokens.
 
-Models with always-on thinking still reason with `thinking=False`;
-`thinking_budget` applies to them too. See [Supported models](#supported-models).
+Models with always-on thinking reason on every field; `thinking_budget` applies
+to them too. See [Supported models](#supported-models).
 
 ## Image input
 
@@ -416,7 +424,7 @@ result = client.generate(
     cancel=cancel,   # set it from another thread; raises GenerationCancelled
 )
 print(client.last_usage)
-# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0)
 ```
 
 `last_usage` reports the SGLang requests and tokens of the last call made in the