From 0aecdc32c927e4b776a7c0adddbb627cf4b52a7e Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:43:45 +0100 Subject: [PATCH] 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 | 53 ++++++++++++- tests/api/user/sources/test_chunks.py | 106 ++++++++++++++++++++++++++ 2 files changed, 158 insertions(+), 1 deletion(-) diff --git a/docsgpt/api/user/sources/chunks.py b/docsgpt/api/user/sources/chunks.py index 841510b8..2ee41393 100644 --- a/docsgpt/api/user/sources/chunks.py +++ b/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( diff --git a/tests/api/user/sources/test_chunks.py b/tests/api/user/sources/test_chunks.py index b4c5f878..a85113e3 100644 --- a/tests/api/user/sources/test_chunks.py +++ b/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