typellm 0.2.0__tar.gz → 0.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.0
+Version: 0.2.1
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/

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

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

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

@@ -83,6 +83,30 @@ def width(payload):
     return 1 if isinstance(payload["text"], str) else len(payload["text"])
 
 
+class PrefixWarmupTests(unittest.TestCase):
+    MIXED = {"a": {"type": "integer"}, "b": {"type": "number"}, "n": {"type": "string"},
+             "k": {"type": "string", "enum": ["x", "y"]}}
+
+    def test_the_context_is_warmed_before_any_batch_that_shares_it(self):
+        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)
+                warm = client.sglang.requests("count")
+                self.assertEqual(len(warm), 1)
+                # First, before the thinking, number, text and scoring batches.
+                self.assertIs(client.sglang.payloads[0], warm[0])
+                self.assertIn("Receipt", warm[0]["text"])
+                self.assertNotIn('{"a":', warm[0]["text"])
+
+    def test_open_fields_alone_are_not_warmed(self):
+        client = TypeLLMClient(model="fake")
+        client.sglang = FakeServer()
+        client.generate(context="Receipt", questions={"a": {"type": "integer"}, "n": {"type": "string"}})
+        self.assertEqual(client.sglang.requests("count"), [])
+
+
 class BatchedThinkingTests(unittest.TestCase):
     QUESTIONS = {
         "flag": {"type": "boolean"},

{typellm-0.2.0 → typellm-0.2.1}/tests/test_grammar_fields.py

@@ -5,7 +5,7 @@ import unittest
 from typellm import SGLangError, TypeLLMClient
 from typellm.runtime import _numeric_text_is_complete, numeric_pattern
 
-from tests.test_batching import FakeServer, width
+from tests.test_batching import FakeServer, is_number_pattern, width
 from tests.test_nullable import NullServer
 
 
@@ -72,6 +72,20 @@ class GrammarNumberTests(unittest.TestCase):
         self.assertEqual(numbers["sampling_params"][0]["temperature"], 0)
         self.assertIn('{"a": 7}', client.last_prompts[0])
 
+    def test_numbers_and_strings_of_a_layer_share_one_request(self):
+        client = TypeLLMClient(model="fake")
+        client.sglang = FakeServer()
+        result = client.generate(context="Receipt", questions={
+            "a": {"type": "integer"}, "n": {"type": "string"},
+            "b": {"type": "number"}, "m": {"type": ["string", "null"]}})
+        self.assertEqual(result, {"a": 7, "n": "blue", "b": 7, "m": "blue"})
+        [request] = regex_requests(client.sglang)
+        self.assertEqual(width(request), 4)
+        patterns = [p["regex"] for p in request["sampling_params"]]
+        self.assertEqual([is_number_pattern(p) for p in patterns], [True, True, False, False])
+        self.assertTrue(all(t.endswith(k) for t, k in zip(request["text"], ('{"a":', '{"b":', '{"n":', '{"m":'))))
+        self.assertIn("null", patterns[3])
+
     def test_sampling_passes_the_temperature_and_no_truncation(self):
         client = TypeLLMClient(model="fake", mode="sample", temperature=0.7, seed=1)
         client.sglang = FakeServer()

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

@@ -60,6 +60,12 @@ class FakeSGLang:
         self.number_patterns.extend(patterns)
         return [next(self.numbers, " 7}") for _ in prefixes]
 
+    def generate_fields(self, number_prefixes, patterns, number_max_tokens, text_prefixes, max_lengths,
+                        *, temperature=0, number_seed=0, text_seed=0, after_key=False, nullable=None):
+        numbers = self.generate_numbers(number_prefixes, patterns, number_max_tokens)
+        texts = self.generate_texts(text_prefixes, max_lengths) if text_prefixes else []
+        return numbers, texts
+
     def score_candidates(self, prefix, candidate_ids):
         self.prompts.append(prefix)
         self.candidate_sets.append(list(candidate_ids))

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

@@ -510,23 +510,10 @@ def numeric_pattern(numeric_type: str, max_digits: int, nullable: bool = False)
     return " ?" + value + r"\}?"
 
 
-def _generate_numeric_batch(
-    client: SGLangClient,
-    items: Sequence[tuple[str, Choice]],
-    mode: str,
-    temperature: float,
-    rng: random.Random,
-    max_digits: int,
+def _parse_numbers(
+    items: Sequence[tuple[str, Choice]], texts: Sequence[str],
 ) -> list[tuple[int | float | None, str, str]]:
-    """Decode prefilled numeric fields in one grammar-constrained batch request."""
-    patterns = [numeric_pattern(d.numeric_type or "", max_digits, d.nullable) for _, d in items]
-    texts = client.generate_numbers(
-        [prompt for prompt, _ in items], patterns,
-        # Digits, sign, point and the closing brace; tokens hold one or more characters.
-        max_new_tokens=max_digits + 4,
-        temperature=0 if mode == "argmax" else temperature,
-        seed=rng.randrange(2**31),
-    )
+    """Read the grammar-generated numbers: (value, closed prompt, value text)."""
     outputs = []
     for (prompt, decision), raw in zip(items, texts):
         text = raw.strip()
@@ -758,6 +745,14 @@ def _execute_batch_decisions(
             {label: token_id for label, (token_id, _) in label_tokens.items()},
         )
 
+    # Warm the exact common prefix once, before any batch that shares it: prompts
+    # prefilled in the same batch cannot reuse each other's cache, so thinking,
+    # number and text batches would each prefill the context once per prompt.
+    # SGLang then forks the cached state; TypeLLM never reads or moves KV tensors.
+    if finite_indexes and not incremental:
+        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
     prompts = [ready[decision_slots[index]] for index in finite_indexes]
     scoring_prompts = [ready[slot] + prefill for slot, prefill in zip(scoring_slots, scoring_prefills)]
@@ -771,37 +766,36 @@ def _execute_batch_decisions(
     numeric_pending = [item for item in open_pending if item[1].numeric_type is not None]
     text_pending = [item for item in open_pending if item[1].text_type]
 
-    if numeric_pending:
-        decoded = _generate_numeric_batch(
-            client,
-            [(prompt + decision.answer_prefill, decision) for _, decision, _, prompt in numeric_pending],
-            mode, temperature, rng, numeric_max_digits,
-        )
-        for (index, decision, messages, prompt), (value, _completed, generated_text) in zip(numeric_pending, decoded):
-            open_results[index] = (open_row(decision, value),
-                                   complete(prompt, messages, _closed_answer(decision, generated_text)))
-
-    if text_pending:
-        values = client.generate_texts(
+    # A layer's numbers and strings decode side by side in one request.
+    if numeric_pending or text_pending:
+        number_items = [(prompt + decision.answer_prefill, decision)
+                        for _, decision, _, prompt in numeric_pending]
+        number_seed = rng.randrange(2**31) if numeric_pending else 0
+        text_seed = rng.randrange(2**31) if text_pending else 0
+        raw_numbers, values = client.generate_fields(
+            [prompt for prompt, _ in number_items],
+            [numeric_pattern(d.numeric_type or "", numeric_max_digits, d.nullable) for _, d in number_items],
+            # Digits, sign, point and the closing brace; tokens hold one or more characters.
+            numeric_max_digits + 4,
             # From {"name": the model writes the string's first token, quote included.
             [prompt + decision.answer_prefill for _, decision, _, prompt in text_pending],
             [decision.max_length for _, decision, _, _ in text_pending],
             temperature=0 if mode == "argmax" else temperature,
-            seed=rng.randrange(2**31),
+            number_seed=number_seed,
+            text_seed=text_seed,
             after_key=all(decision.answer_prefill for _, decision, _, _ in text_pending),
             # A nullable string writes null or its text in the same request.
             nullable=[decision.nullable for _, decision, _, _ in text_pending],
         )
+        for (index, decision, messages, prompt), (value, _completed, generated_text) in zip(
+                numeric_pending, _parse_numbers(number_items, raw_numbers)):
+            open_results[index] = (open_row(decision, value),
+                                   complete(prompt, messages, _closed_answer(decision, generated_text)))
         for (index, decision, messages, prompt), value in zip(text_pending, values):
             completed = complete(prompt, messages, _closed_answer(decision, json.dumps(value, ensure_ascii=False)))
             open_results[index] = (open_row(decision, value), completed)
 
-    # Warm the exact common prefix once, then let SGLang fork the cached state
-    # across the K batched prompts. TypeLLM never reads or moves KV tensors.
     if prompts:
-        if not incremental:
-            cache_meta = client.cache_prefix(shared_prefix)
-            LOG.info("batch_shared_prefix_cached_tokens=%s", cache_meta.get("cached_tokens"))
         scored, elapsed = client.score_candidates_batch(scoring_prompts, scoring_ids)
         grouped_scores = []
         offset = 0

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

@@ -744,35 +744,26 @@ class SGLangClient:
             results.append((scores, meta if isinstance(meta, Mapping) else {}))
         return results, elapsed
 
-    def generate_numbers(
-        self,
-        prefixes: Sequence[str],
-        patterns: Sequence[str],
-        max_new_tokens: int,
-        *,
-        temperature: float = 0,
-        seed: int = 0,
-    ) -> list[str]:
-        """Generate one regex-constrained value per prompt in a native batch.
-
-        SGLang masks every step to the pattern, so a number takes one request
-        instead of one per token. Returns the raw generated text.
-        """
-        if len(prefixes) != len(patterns):
-            raise ValueError("prefixes and patterns must have the same length")
-        if not prefixes:
-            return []
-        params = [{"max_new_tokens": max_new_tokens, "temperature": temperature,
-                   "top_p": 1.0, "top_k": -1, "min_p": 0.0,
-                   "sampling_seed": seed + i, "regex": pattern}
-                  for i, pattern in enumerate(patterns)]
-        response = self._generate({"text": list(prefixes), "sampling_params": params})
+    def _batch(self, prefixes: Sequence[str], params: Sequence[Mapping[str, Any]], what: str) -> list[Any]:
+        """One /generate request with one sampling-params dict per prompt."""
+        response = self._generate({"text": list(prefixes), "sampling_params": list(params)})
         if isinstance(response, Mapping) and len(prefixes) == 1:
             response = [response]
         if not isinstance(response, list) or len(response) != len(prefixes):
-            raise SGLangError("Unexpected number batch response shape")
+            raise SGLangError(f"Unexpected {what} batch response shape")
+        return response
+
+    @staticmethod
+    def _number_params(patterns: Sequence[str], max_new_tokens: int, temperature: float, seed: int) -> list[dict]:
+        return [{"max_new_tokens": max_new_tokens, "temperature": temperature,
+                 "top_p": 1.0, "top_k": -1, "min_p": 0.0,
+                 "sampling_seed": seed + i, "regex": pattern}
+                for i, pattern in enumerate(patterns)]
+
+    @staticmethod
+    def _read_numbers(items: Sequence[Any]) -> list[str]:
         texts = []
-        for item in response:
+        for item in items:
             meta = item.get("meta_info", {}) if isinstance(item, Mapping) else {}
             finish = meta.get("finish_reason", {}) if isinstance(meta, Mapping) else {}
             kind = finish.get("type") if isinstance(finish, Mapping) else finish
@@ -782,34 +773,10 @@ class SGLangClient:
             texts.append(text)
         return texts
 
-    def generate_texts(
-        self,
-        prefixes: Sequence[str],
-        max_lengths: Sequence[int | None],
-        *,
-        temperature: float = 0,
-        seed: int = 0,
-        after_key: bool = False,
-        nullable: Sequence[bool] | None = None,
-    ) -> list[str | None]:
-        """Generate JSON strings in a native batch, then validate every value.
-
-        With after_key, every prompt ends with '{"name":', and the model writes the
-        opening quote, the characters and the closing '"}'. Writing the quote
-        itself lets it start with a merged token such as ' "$', which keeps the
-        first character. A max length then truncates: generation is capped near
-        that many tokens and the string is cut to that many characters, as a
-        length-bounded grammar would.
-        """
-        if len(prefixes) != len(max_lengths):
-            raise ValueError("prefixes and max_lengths must have the same length")
-        nullable = [False] * len(prefixes) if nullable is None else list(nullable)
-        if len(nullable) != len(prefixes):
-            raise ValueError("prefixes and nullable must have the same length")
+    def _text_params(self, max_lengths: Sequence[int | None], nullable: Sequence[bool],
+                     after_key: bool, temperature: float, seed: int) -> list[dict]:
         if any(nullable) and not after_key:
             raise ValueError("nullable text needs the prefilled '{\"name\":' prompt")
-        if not prefixes:
-            return []
         params = []
         for limit, can_be_null in zip(max_lengths, nullable):
             budget = self.text_max_tokens
@@ -832,13 +799,13 @@ class SGLangClient:
                 constraint = {"json_schema": json.dumps(schema)}
             params.append({"max_new_tokens": budget,
                            "temperature": temperature, "sampling_seed": seed, **constraint})
-        response = self._generate({"text": list(prefixes), "sampling_params": params})
-        if isinstance(response, Mapping) and len(prefixes) == 1:
-            response = [response]
-        if not isinstance(response, list) or len(response) != len(prefixes):
-            raise SGLangError("Unexpected text batch response shape")
+        return params
+
+    @staticmethod
+    def _read_texts(items: Sequence[Any], max_lengths: Sequence[int | None],
+                    nullable: Sequence[bool], after_key: bool) -> list[str | None]:
         values: list[str | None] = []
-        for item, limit, can_be_null in zip(response, max_lengths, nullable):
+        for item, limit, can_be_null in zip(items, max_lengths, nullable):
             if not isinstance(item, Mapping):
                 raise SGLangError("Invalid text response")
             meta = item.get("meta_info", {})
@@ -874,6 +841,93 @@ class SGLangClient:
             values.append(value)
         return values
 
+    @staticmethod
+    def _check_lengths(nullable, prefixes, max_lengths):
+        if len(prefixes) != len(max_lengths):
+            raise ValueError("prefixes and max_lengths must have the same length")
+        nullable = [False] * len(prefixes) if nullable is None else list(nullable)
+        if len(nullable) != len(prefixes):
+            raise ValueError("prefixes and nullable must have the same length")
+        return nullable
+
+    def generate_numbers(
+        self,
+        prefixes: Sequence[str],
+        patterns: Sequence[str],
+        max_new_tokens: int,
+        *,
+        temperature: float = 0,
+        seed: int = 0,
+    ) -> list[str]:
+        """Generate one regex-constrained value per prompt in a native batch.
+
+        SGLang masks every step to the pattern, so a number takes one request
+        instead of one per token. Returns the raw generated text.
+        """
+        if len(prefixes) != len(patterns):
+            raise ValueError("prefixes and patterns must have the same length")
+        if not prefixes:
+            return []
+        params = self._number_params(patterns, max_new_tokens, temperature, seed)
+        return self._read_numbers(self._batch(prefixes, params, "number"))
+
+    def generate_texts(
+        self,
+        prefixes: Sequence[str],
+        max_lengths: Sequence[int | None],
+        *,
+        temperature: float = 0,
+        seed: int = 0,
+        after_key: bool = False,
+        nullable: Sequence[bool] | None = None,
+    ) -> list[str | None]:
+        """Generate JSON strings in a native batch, then validate every value.
+
+        With after_key, every prompt ends with '{"name":', and the model writes the
+        opening quote, the characters and the closing '"}'. Writing the quote
+        itself lets it start with a merged token such as ' "$', which keeps the
+        first character. A max length then truncates: generation is capped near
+        that many tokens and the string is cut to that many characters, as a
+        length-bounded grammar would.
+        """
+        nullable = self._check_lengths(nullable, prefixes, max_lengths)
+        if not prefixes:
+            return []
+        params = self._text_params(max_lengths, nullable, after_key, temperature, seed)
+        return self._read_texts(self._batch(prefixes, params, "text"), max_lengths, nullable, after_key)
+
+    def generate_fields(
+        self,
+        number_prefixes: Sequence[str],
+        patterns: Sequence[str],
+        number_max_tokens: int,
+        text_prefixes: Sequence[str],
+        max_lengths: Sequence[int | None],
+        *,
+        temperature: float = 0,
+        number_seed: int = 0,
+        text_seed: int = 0,
+        after_key: bool = False,
+        nullable: Sequence[bool] | None = None,
+    ) -> tuple[list[str], list[str | None]]:
+        """Numbers and strings of one layer in a single request, decoded side by side.
+
+        Returns what generate_numbers and generate_texts would for each part.
+        """
+        if len(number_prefixes) != len(patterns):
+            raise ValueError("prefixes and patterns must have the same length")
+        nullable = self._check_lengths(nullable, text_prefixes, max_lengths)
+        params = (self._number_params(patterns, number_max_tokens, temperature, number_seed)
+                  + (self._text_params(max_lengths, nullable, after_key, temperature, text_seed)
+                     if text_prefixes else []))
+        prefixes = list(number_prefixes) + list(text_prefixes)
+        if not prefixes:
+            return [], []
+        items = self._batch(prefixes, params, "field")
+        split = len(number_prefixes)
+        return (self._read_numbers(items[:split]),
+                self._read_texts(items[split:], max_lengths, nullable, after_key))
+
     def flush_cache(self) -> None:
         reply = self._request("/flush_cache", {}, allow_text=True)
         if isinstance(reply, str) and not reply.startswith("Cache flushed."):

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.0
+Version: 0.2.1
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/