diff --git a/docsgpt/api/user/sources/chunks.py b/docsgpt/api/user/sources/chunks.py index 2ee41393..ae289aa6 100644 --- a/docsgpt/api/user/sources/chunks.py +++ b/docsgpt/api/user/sources/chunks.py @@ -1,5 +1,7 @@ """Source document management chunk management.""" +import math + from flask import current_app, jsonify, make_response, request from flask_restx import fields, Namespace, Resource @@ -48,22 +50,24 @@ def _has_usable_token_count(metadata: dict) -> bool: Stores round-trip metadata differently -- pgvector keeps JSON types, the Mongo backend can hand back strings -- so a numeric string counts as - recorded. Anything else (missing, empty, non-numeric, zero or negative) - does not. + recorded. Anything else (missing, empty, non-numeric, zero, negative, or + non-finite) does not: ``float("inf")`` is greater than zero but is not a + number of tokens, and it reaches ``toLocaleString`` in the UI as "∞". Args: metadata: A chunk's metadata mapping. Returns: - True when ``token_count`` holds a positive number. + True when ``token_count`` holds a finite positive number. """ raw = metadata.get("token_count") if isinstance(raw, bool) or not isinstance(raw, (int, float, str)): return False try: - return float(raw) > 0 + value = float(raw) except (TypeError, ValueError): return False + return math.isfinite(value) and value > 0 def _with_token_counts(chunks: list) -> list: diff --git a/tests/api/user/sources/test_chunks.py b/tests/api/user/sources/test_chunks.py index a85113e3..83829b78 100644 --- a/tests/api/user/sources/test_chunks.py +++ b/tests/api/user/sources/test_chunks.py @@ -199,7 +199,10 @@ class TestGetChunks: # An already-recorded count is left alone, whatever tokenizer produced it. assert chunks[2]["metadata"]["token_count"] == 42 - @pytest.mark.parametrize("stored", ["", 0, -1, "abc", None, {"n": 1}]) + @pytest.mark.parametrize( + "stored", + ["", 0, -1, "abc", None, {"n": 1}, "inf", "-inf", "nan", float("inf")], + ) def test_recomputes_unusable_token_count(self, app, pg_conn, stored): from docsgpt.api.user.sources.chunks import GetChunks