Backfill missing chunk token counts in /api/get_chunks

Chunk cards in the source viewer read `metadata.token_count` and print a
bare "-" whenever it is absent. Ingestion records the count, but chunks
indexed before it was written -- and any path that rebuilds a chunk's
metadata from scratch -- reach the UI without it, so the whole file shows
"-" with no way to recover the number short of a re-ingest.

Fill the key in on the way out when it is missing or unusable (empty,
non-numeric, zero or negative), counting only the page being returned so
the cost is bounded by `per_page`. A count that is already recorded is
left untouched, including a numeric string, since stores round-trip
metadata types differently.

The fallback counts in cl100k rather than the embedding model's
tokenizer: loading that tokenizer can reach for a Hugging Face download,
which does not belong in a request path.

docsgpt/api/user/sources/chunks.py

@@ -43,6 +43,57 @@ def _resolve_source_for_write(doc_id: str, user: str):
         return SourcesRepository(conn).get_any(doc_id, owner)
 
 
+def _has_usable_token_count(metadata: dict) -> bool:
+    """Whether ``metadata`` already carries a count worth showing.
+
+    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.
+
+    Args:
+        metadata: A chunk's metadata mapping.
+
+    Returns:
+        True when ``token_count`` holds a 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
+    except (TypeError, ValueError):
+        return False
+
+
+def _with_token_counts(chunks: list) -> list:
+    """Fill in ``metadata.token_count`` for chunks that were stored without one.
+
+    Ingestion records a per-chunk count in the embedding model's tokenizer,
+    but chunks indexed before that was written -- and any path that rebuilt a
+    chunk's metadata from scratch -- reach the UI without the key, which then
+    renders a bare "-". The count recomputed here is cl100k rather than the
+    embedding model's tokenizer: it is a display fallback, and loading the
+    model's tokenizer would put a Hugging Face download in the request path.
+
+    Only the page being returned is counted, so the cost is bounded by
+    ``per_page`` rather than by the size of the index.
+
+    Args:
+        chunks: The chunk dicts about to be serialised.
+
+    Returns:
+        The same list, with each chunk's metadata normalised to a dict that
+        carries a ``token_count``.
+    """
+    for chunk in chunks:
+        metadata = chunk.get("metadata") or {}
+        if not _has_usable_token_count(metadata):
+            metadata["token_count"] = num_tokens_from_string(chunk.get("text") or "")
+        chunk["metadata"] = metadata
+    return chunks
+
+
 @sources_chunks_ns.route("/get_chunks")
 class GetChunks(Resource):
     @api.doc(
@@ -106,7 +157,7 @@ class GetChunks(Resource):
             total_chunks = len(chunks)
             start = (page - 1) * per_page
             end = start + per_page
-            paginated_chunks = chunks[start:end]
+            paginated_chunks = _with_token_counts(chunks[start:end])
 
             return make_response(
                 jsonify(

tests/api/user/sources/test_chunks.py

@@ -170,6 +170,112 @@ class TestGetChunks:
         assert response.status_code == 200
         assert response.json["total"] == 1
 
+    def test_backfills_missing_token_count(self, app, pg_conn):
+        # Chunks indexed before token_count was recorded (and any ingest path
+        # that dropped it) came back without the key, so the UI printed "-".
+        from docsgpt.api.user.sources.chunks import GetChunks
+
+        user = "u-tokens"
+        src = _seed_source(pg_conn, user=user)
+
+        fake_store = MagicMock()
+        fake_store.get_chunks.return_value = [
+            {"doc_id": "a", "text": "hello world", "metadata": {}},
+            {"doc_id": "b", "text": "second chunk", "metadata": None},
+            {"doc_id": "c", "text": "third chunk", "metadata": {"token_count": 42}},
+        ]
+
+        with _patch_db(pg_conn), patch(
+            "docsgpt.api.user.sources.chunks.get_vector_store",
+            return_value=fake_store,
+        ), app.test_request_context(f"/api/get_chunks?id={src['id']}"):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = GetChunks().get()
+
+        assert response.status_code == 200
+        chunks = response.json["chunks"]
+        assert all(c["metadata"]["token_count"] > 0 for c in chunks)
+        # 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}])
+    def test_recomputes_unusable_token_count(self, app, pg_conn, stored):
+        from docsgpt.api.user.sources.chunks import GetChunks
+
+        user = f"u-tok-{str(stored)[:6]}"
+        src = _seed_source(pg_conn, user=user)
+
+        fake_store = MagicMock()
+        fake_store.get_chunks.return_value = [
+            {"doc_id": "a", "text": "some text here", "metadata": {"token_count": stored}}
+        ]
+
+        with _patch_db(pg_conn), patch(
+            "docsgpt.api.user.sources.chunks.get_vector_store",
+            return_value=fake_store,
+        ), app.test_request_context(f"/api/get_chunks?id={src['id']}"):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = GetChunks().get()
+
+        assert response.status_code == 200
+        assert response.json["chunks"][0]["metadata"]["token_count"] > 0
+
+    def test_keeps_numeric_string_token_count(self, app, pg_conn):
+        # Some stores round-trip metadata values as strings; a usable count
+        # there is still a count and must not be recomputed.
+        from docsgpt.api.user.sources.chunks import GetChunks
+
+        user = "u-tok-str"
+        src = _seed_source(pg_conn, user=user)
+
+        fake_store = MagicMock()
+        fake_store.get_chunks.return_value = [
+            {"doc_id": "a", "text": "some text here", "metadata": {"token_count": "17"}}
+        ]
+
+        with _patch_db(pg_conn), patch(
+            "docsgpt.api.user.sources.chunks.get_vector_store",
+            return_value=fake_store,
+        ), app.test_request_context(f"/api/get_chunks?id={src['id']}"):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = GetChunks().get()
+
+        assert response.json["chunks"][0]["metadata"]["token_count"] == "17"
+
+    def test_backfills_only_the_requested_page(self, app, pg_conn):
+        # Counting is per-response work, so it must not run over chunks the
+        # caller never sees.
+        from docsgpt.api.user.sources.chunks import GetChunks
+
+        user = "u-tok-page"
+        src = _seed_source(pg_conn, user=user)
+
+        fake_store = MagicMock()
+        fake_store.get_chunks.return_value = [
+            {"doc_id": str(i), "text": f"chunk number {i}", "metadata": {}}
+            for i in range(10)
+        ]
+
+        with _patch_db(pg_conn), patch(
+            "docsgpt.api.user.sources.chunks.get_vector_store",
+            return_value=fake_store,
+        ), patch(
+            "docsgpt.api.user.sources.chunks.num_tokens_from_string",
+            return_value=7,
+        ) as counter, app.test_request_context(
+            f"/api/get_chunks?id={src['id']}&per_page=2&page=1"
+        ):
+            from flask import request
+            request.decoded_token = {"sub": user}
+            response = GetChunks().get()
+
+        assert response.status_code == 200
+        assert counter.call_count == 2
+        assert [c["metadata"]["token_count"] for c in response.json["chunks"]] == [7, 7]
+
     def test_returns_500_on_vector_store_error(self, app, pg_conn):
         from docsgpt.api.user.sources.chunks import GetChunks