typellm 0.2.2__tar.gz → 0.2.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.2
+Version: 0.2.4
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/
@@ -424,11 +424,14 @@ 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, thinking_tokens=0)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
 ```
 
-`last_usage` reports the SGLang requests and tokens of the last call made in the
-current thread, including a call that failed partway. A timeout or cancel stops
+`last_usage` reports the tokens of the last call made in the current thread,
+including a call that failed partway. `input_tokens` is what you sent, each part
+counted once: the context, the questions as JSON, and the images. The other
+counts are SGLang's for the call's requests, where every prompt carries the
+shared context. A timeout or cancel stops
 the call before its next SGLang request.
 
 ## Cost analysis

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

@@ -409,11 +409,14 @@ 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, thinking_tokens=0)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
 ```
 
-`last_usage` reports the SGLang requests and tokens of the last call made in the
-current thread, including a call that failed partway. A timeout or cancel stops
+`last_usage` reports the tokens of the last call made in the current thread,
+including a call that failed partway. `input_tokens` is what you sent, each part
+counted once: the context, the questions as JSON, and the images. The other
+counts are SGLang's for the call's requests, where every prompt carries the
+shared context. A timeout or cancel stops
 the call before its next SGLang request.
 
 ## Cost analysis

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

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

{typellm-0.2.2 → typellm-0.2.4}/tests/test_concurrency.py

@@ -1,3 +1,4 @@
+import json
 import pickle
 import sys
 import threading
@@ -104,10 +105,13 @@ class UsageTests(unittest.TestCase):
         self.assertIsNone(client.last_usage)
         client.generate(context="Receipt", questions=self.QUESTIONS)
         n = prompts_sent(client.sglang)
+        # The input counts once, though every prompt carries it.
+        sent = client.sglang.count_tokens("Receipt") + client.sglang.count_tokens(json.dumps(self.QUESTIONS))
         self.assertEqual(client.last_usage, Usage(
             requests=len(client.sglang.payloads),
-            prompt_tokens=10 * n, cached_tokens=4 * n, completion_tokens=n,
+            prompt_tokens=10 * n, cached_tokens=4 * n, completion_tokens=n, input_tokens=sent,
         ))
+        self.assertGreater(client.last_usage.requests, 1)
         client.sglang.payloads.clear()
         client.generate(context="Receipt", questions={"b": {"type": "boolean"}})
         self.assertEqual(client.last_usage.requests, len(client.sglang.payloads))
@@ -136,6 +140,7 @@ class UsageTests(unittest.TestCase):
         with self.assertRaisesRegex(SGLangError, "went away"):
             client.generate(context="Receipt", questions=self.QUESTIONS)
         self.assertEqual(client.last_usage.requests, 2)
+        self.assertGreater(client.last_usage.input_tokens, 0)
         with self.assertRaises(ValueError):
             client.generate(questions=self.QUESTIONS)
         self.assertIsNone(client.last_usage)
@@ -207,6 +212,38 @@ class DeadlineTests(unittest.TestCase):
             client._generate({"text": "x"})
         self.assertNotIsInstance(caught.exception, GenerationTimeout)
 
+    def dropping(self, error, times):
+        """A client whose server drops the first `times` requests with `error`."""
+        seen = []
+
+        def handler(request):
+            seen.append(request)
+            if len(seen) <= times:
+                raise error("connection dropped", request=request)
+            return httpx.Response(200, json={"text": "ok"})
+
+        client = SGLangClient()
+        client._http_client = httpx.Client(transport=httpx.MockTransport(handler))
+        return client, seen
+
+    def test_a_dropped_connection_is_retried_once(self):
+        for error in (httpx.ReadError, httpx.RemoteProtocolError, httpx.WriteError):
+            client, seen = self.dropping(error, 1)
+            with self.subTest(error=error.__name__):
+                self.assertEqual(client._request("/generate", {"text": "x"}), {"text": "ok"})
+                self.assertEqual(len(seen), 2)
+        client, seen = self.dropping(httpx.ReadError, 2)
+        with self.assertRaisesRegex(SGLangError, "Could not reach"):
+            client._request("/generate", {"text": "x"})
+        self.assertEqual(len(seen), 2)
+
+    def test_a_refused_connection_is_not_retried(self):
+        for error in (httpx.ConnectError, httpx.ReadTimeout):
+            client, seen = self.dropping(error, 1)
+            with self.subTest(error=error.__name__), self.assertRaises(SGLangError):
+                client._request("/generate", {"text": "x"})
+            self.assertEqual(len(seen), 1)
+
     def test_invalid_budgets_are_rejected(self):
         class Recording(FakeServer):
             def _request(self, path, payload=None, *, allow_text=False):

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

@@ -271,6 +271,21 @@ class ImageRequestTests(unittest.TestCase):
         self.assertEqual(rest, [])
 
 
+class ImageInputTokenTests(unittest.TestCase):
+    def test_images_count_as_the_server_expands_them(self):
+        client = TypeLLMClient(model="fake-vl")
+        client.sglang = FakeServerClient()
+        questions = {"a": {"type": "boolean"}}
+        client.generate(context="Receipt", images=[PNG], questions=questions)
+        first = client.sglang.generate_payloads[0]["text"]
+        first = first if isinstance(first, str) else first[0]
+        count = client.sglang.count_tokens
+        # The server reported 900 prompt tokens, with the placeholder expanded.
+        image = 900 - count(first) + count(VISION)
+        self.assertEqual(client.last_usage.input_tokens,
+                         count("Receipt") + count('{"a": {"type": "boolean"}}') + image)
+
+
 class RuntimeContentTests(unittest.TestCase):
     def test_fake_backends_receive_image_parts(self):
         class Recording(FakeSGLang):

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

@@ -395,6 +395,11 @@ class TypeLLMClient:
         with call_scope(timeout, cancel) as scope:
             try:
                 decisions = self.compile_schema(schema)
+                count = getattr(self.sglang, "count_tokens", None)
+                if callable(count):
+                    scope.usage.input_tokens = sum(map(count, (
+                        context, json.dumps(questions if questions is not None else schema, ensure_ascii=False))))
+                    scope.unmeasured_images = len(encoded_images)
                 # Independent fields run together; depends_on turns the fields into a
                 # graph whose layers run in order.
                 run = (_execute_dependency_decisions if any(d.depends_on is not None for d in decisions)

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

@@ -37,7 +37,13 @@ class GenerationCancelled(RuntimeError):
 
 @dataclass
 class Usage:
-    """What SGLang reported for the /generate requests of one generate() call."""
+    """The tokens of one generate() call.
+
+    input_tokens is what the caller sent, each part counted once: the context,
+    the questions or schema as JSON, and the images as the server expands them.
+    The other counts are what SGLang reported for the call's /generate requests,
+    where the shared context is part of every prompt.
+    """
 
     requests: int = 0
     prompt_tokens: int = 0
@@ -45,6 +51,7 @@ class Usage:
     completion_tokens: int = 0
     # The part of completion_tokens that is reasoning, not typed answers.
     thinking_tokens: int = 0
+    input_tokens: int = 0
 
     def _add(self, response: Any) -> None:
         self.requests += 1
@@ -75,6 +82,8 @@ class _CallScope:
     usage: Usage = field(default_factory=Usage)
     deadline: float | None = None  # time.monotonic() when the call runs out of time
     cancel: threading.Event | None = None
+    # Images whose tokens input_tokens still lacks; the first response measures them.
+    unmeasured_images: int = 0
 
     def expired(self) -> bool:
         return self.deadline is not None and time.monotonic() >= self.deadline
@@ -243,9 +252,11 @@ class SGLangClient:
     def _http(self) -> httpx.Client:
         with self._load_lock:
             if self._http_client is None:
-                # Keep-alive connections: one call makes many small requests.
+                # Keep-alive connections: one call makes many small requests. Idle
+                # ones are dropped before SGLang closes them (5 s), so a request
+                # rarely goes out on a connection the server has just closed.
                 self._http_client = httpx.Client(limits=httpx.Limits(
-                    max_connections=None, max_keepalive_connections=64,
+                    max_connections=None, max_keepalive_connections=64, keepalive_expiry=2.0,
                 ))
             return self._http_client
 
@@ -260,24 +271,15 @@ class SGLangClient:
         scope = _call_scope.get()
         timeout = self.timeout if scope is None else scope.socket_timeout(self.timeout)
         try:
-            # Streamed so a failed body read still reports the status it follows.
-            with self._http().stream(
-                "GET" if payload is None else "POST",
-                self.base_url + path,
-                content=body,
-                headers={"Content-Type": "application/json"},
-                timeout=timeout,
-            ) as response:
-                if response.status_code >= 400:
-                    try:
-                        detail = response.read().decode("utf-8", errors="replace")
-                    except httpx.HTTPError as read_error:
-                        detail = f"Could not read error response: {read_error}"
-                    raise SGLangError(
-                        f"SGLang {path} returned HTTP {response.status_code}: {detail}",
-                        status=response.status_code,
-                    )
-                raw = response.read().decode("utf-8")
+            try:
+                raw = self._send(path, payload is None, body, timeout)
+            except (httpx.ReadError, httpx.WriteError, httpx.RemoteProtocolError):
+                # The connection dropped before an answer, most often a pooled one
+                # the server had just closed. SGLang is fine: try once on a new one.
+                if scope is not None:
+                    scope.check()
+                    timeout = scope.socket_timeout(self.timeout)
+                raw = self._send(path, payload is None, body, timeout)
         except httpx.HTTPError as exc:
             # Connection failures, timeouts, resets and truncated responses.
             raise SGLangError(
@@ -292,6 +294,26 @@ class SGLangClient:
                 f"SGLang {path} returned non-JSON data: {raw[:500]}"
             ) from exc
 
+    def _send(self, path: str, get: bool, body: bytes | None, timeout: float) -> str:
+        # Streamed so a failed body read still reports the status it follows.
+        with self._http().stream(
+            "GET" if get else "POST",
+            self.base_url + path,
+            content=body,
+            headers={"Content-Type": "application/json"},
+            timeout=timeout,
+        ) as response:
+            if response.status_code >= 400:
+                try:
+                    detail = response.read().decode("utf-8", errors="replace")
+                except httpx.HTTPError as read_error:
+                    detail = f"Could not read error response: {read_error}"
+                raise SGLangError(
+                    f"SGLang {path} returned HTTP {response.status_code}: {detail}",
+                    status=response.status_code,
+                )
+            return response.read().decode("utf-8")
+
     @contextmanager
     def images(self, images: Sequence[str]) -> Iterator[None]:
         """Attach encoded images to every /generate request made in this block."""
@@ -351,8 +373,27 @@ class SGLangClient:
                 raise GenerationTimeout("generate() ran out of time") from exc
             raise
         scope.usage._add(response)
+        if scope.unmeasured_images:
+            self._measure_images(scope, payload, response)
         return response
 
+    def count_tokens(self, text: str) -> int:
+        return len(self._get_chat_tokenizer().encode(text, add_special_tokens=False))
+
+    def _measure_images(self, scope: _CallScope, payload: Mapping[str, Any], response: Any) -> None:
+        """Add the images' tokens to input_tokens: the server's prompt count less the local one."""
+        text = payload.get("text")
+        prompt = text if isinstance(text, str) else text[0] if isinstance(text, list) and text else None
+        item = response[0] if isinstance(response, list) and response else response
+        meta = item.get("meta_info") if isinstance(item, Mapping) else None
+        served = meta.get("prompt_tokens") if isinstance(meta, Mapping) else None
+        if not isinstance(prompt, str) or type(served) is not int:
+            return
+        # The server replaces each placeholder with the image's tokens.
+        placeholders = scope.unmeasured_images * self.count_tokens(self.image_placeholder())
+        scope.usage.input_tokens += max(0, served - self.count_tokens(prompt) + placeholders)
+        scope.unmeasured_images = 0
+
     def _with_images(self, payload: Mapping[str, Any]) -> Mapping[str, Any]:
         images = self._active_images.get()
         if not images:

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: typellm
-Version: 0.2.2
+Version: 0.2.4
 Summary: Type-Safe Decoding for Autoregressive LLMs
 License-Expression: Apache-2.0
 Project-URL: Homepage, https://typellm.ai/
@@ -424,11 +424,14 @@ 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, thinking_tokens=0)
+# Usage(requests=4, prompt_tokens=1830, cached_tokens=1504, completion_tokens=7, thinking_tokens=0, input_tokens=410)
 ```
 
-`last_usage` reports the SGLang requests and tokens of the last call made in the
-current thread, including a call that failed partway. A timeout or cancel stops
+`last_usage` reports the tokens of the last call made in the current thread,
+including a call that failed partway. `input_tokens` is what you sent, each part
+counted once: the context, the questions as JSON, and the images. The other
+counts are SGLang's for the call's requests, where every prompt carries the
+shared context. A timeout or cancel stops
 the call before its next SGLang request.
 
 ## Cost analysis