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 1/4] 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 From d787f54c2d44fd189c6bd64c95f88ce0c60a2057 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:44:24 +0100 Subject: [PATCH 2/4] Keep the chunker's metadata when re-embedding a wiki page `reembed_wiki_page_worker` built each chunk's metadata from scratch as `{source, title, filename}`, discarding everything the chunker attached to the document -- including the per-chunk `token_count` every other ingest path records. Wiki chunks therefore reached the source viewer with no count at all. Carry `extra_info` over and let the page's own path and title win on top, so a chunk that inherited a stale source from an earlier conversion is still keyed and filtered correctly. --- docsgpt/worker.py | 10 ++-- tests/worker/test_reembed_wiki_page.py | 74 ++++++++++++++++++++++++++ 2 files changed, 80 insertions(+), 4 deletions(-) diff --git a/docsgpt/worker.py b/docsgpt/worker.py index 947f59c2..9fbf57f3 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -2615,10 +2615,12 @@ def reembed_wiki_page_worker(self, source_id, path, content_hash, user): added = 0 for chunk in chunks: - store.add_chunk( - chunk.text, - metadata={"source": path, "title": title, "filename": path}, - ) + # Start from what the chunker produced -- ``token_count`` above + # all, which the source viewer reads per chunk -- and let the + # page's own identity win over anything stale it inherited. + metadata = dict(chunk.extra_info or {}) + metadata.update({"source": path, "title": title, "filename": path}) + store.add_chunk(chunk.text, metadata=metadata) added += 1 with db_session() as conn: diff --git a/tests/worker/test_reembed_wiki_page.py b/tests/worker/test_reembed_wiki_page.py index 8447f8a3..04e10043 100644 --- a/tests/worker/test_reembed_wiki_page.py +++ b/tests/worker/test_reembed_wiki_page.py @@ -100,6 +100,80 @@ class TestReembedWikiPageWorker: ) assert result == {"status": "embedded", "added": 2, "deleted": 3} + def test_chunk_metadata_survives_the_rewrite( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """The chunker's own metadata -- ``token_count`` above all -- has to + reach the store, or the source viewer shows "-" for every wiki chunk. + """ + from docsgpt import worker + + source_id = _seed_source(pg_conn) + + store = MagicMock(name="vector_store") + store.delete_chunks_by_source_path.return_value = 0 + _patch_store(monkeypatch, store) + + _patch_repo(monkeypatch, {"content": "page body", "title": "Page Title"}) + _patch_chunker( + monkeypatch, + [ + Document(text="c1", extra_info={"token_count": 11, "custom": "keep"}), + Document(text="c2", extra_info={"token_count": 7}), + ], + ) + + worker.reembed_wiki_page_worker( + task_self, source_id, "guide/intro.md", "hash-1", "alice" + ) + + counts = [ + call.kwargs["metadata"]["token_count"] + for call in store.add_chunk.call_args_list + ] + assert counts == [11, 7] + assert store.add_chunk.call_args_list[0].kwargs["metadata"]["custom"] == "keep" + + def test_page_identity_wins_over_stale_chunk_metadata( + self, pg_conn, patch_worker_db, task_self, monkeypatch + ): + """Whatever the chunker inherited, the page's own path and title are + what the store is keyed and filtered on. + """ + from docsgpt import worker + + source_id = _seed_source(pg_conn) + + store = MagicMock(name="vector_store") + store.delete_chunks_by_source_path.return_value = 0 + _patch_store(monkeypatch, store) + + _patch_repo(monkeypatch, {"content": "body", "title": "Page Title"}) + _patch_chunker( + monkeypatch, + [ + Document( + text="c1", + extra_info={ + "source": "stale/path.md", + "title": "Stale", + "filename": "stale/path.md", + "token_count": 5, + }, + ) + ], + ) + + worker.reembed_wiki_page_worker( + task_self, source_id, "guide/intro.md", "hash-1", "alice" + ) + + metadata = store.add_chunk.call_args_list[0].kwargs["metadata"] + assert metadata["source"] == "guide/intro.md" + assert metadata["filename"] == "guide/intro.md" + assert metadata["title"] == "Page Title" + assert metadata["token_count"] == 5 + def test_a_reembed_stamps_the_model_on_the_source( self, pg_conn, patch_worker_db, task_self, monkeypatch ): From 3954de9b9ea6954bf4fa9bb5aab48ffbfd922c3d Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 22 Sep 2026 13:45:40 +0100 Subject: [PATCH 3/4] Render chunk token counts through one tested helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both chunk headers inlined the same `token_count ? … : '-'` expression, and `ChunkType` declared every metadata value a string even though the JSON-typed stores return the count as a number. Move the formatting into `formatChunkTokens`, which accepts either shape and only falls back to a dash when the value is genuinely unusable, and type `ChunkType.metadata` to match what the API actually returns. --- frontend/src/components/Chunks.tsx | 9 +++---- frontend/src/components/chunkUtils.test.ts | 31 ++++++++++++++++++++++ frontend/src/components/chunkUtils.ts | 27 +++++++++++++++++++ frontend/src/settings/types/index.ts | 12 ++++++++- 4 files changed, 72 insertions(+), 7 deletions(-) create mode 100644 frontend/src/components/chunkUtils.test.ts create mode 100644 frontend/src/components/chunkUtils.ts diff --git a/frontend/src/components/Chunks.tsx b/frontend/src/components/Chunks.tsx index e01ead13..8efa7f64 100644 --- a/frontend/src/components/Chunks.tsx +++ b/frontend/src/components/Chunks.tsx @@ -20,6 +20,7 @@ import ConfirmationModal from '../modals/ConfirmationModal'; import { ActiveState } from '../models/misc'; import { selectToken } from '../preferences/preferenceSlice'; import { ChunkType } from '../settings/types'; +import { formatChunkTokens } from './chunkUtils'; import Pagination from './DocumentPagination'; import SkeletonLoader from './SkeletonLoader'; import { Button } from './ui/button'; @@ -628,9 +629,7 @@ const Chunks: React.FC = ({
- {chunk.metadata.token_count - ? chunk.metadata.token_count.toLocaleString() - : '-'}{' '} + {formatChunkTokens(chunk.metadata)}{' '} {t('settings.sources.tokensUnit')}
@@ -663,9 +662,7 @@ const Chunks: React.FC = ({
- {editingChunk.metadata.token_count - ? editingChunk.metadata.token_count.toLocaleString() - : '-'}{' '} + {formatChunkTokens(editingChunk.metadata)}{' '} {t('settings.sources.tokensUnit')}
diff --git a/frontend/src/components/chunkUtils.test.ts b/frontend/src/components/chunkUtils.test.ts new file mode 100644 index 00000000..6ff36b9a --- /dev/null +++ b/frontend/src/components/chunkUtils.test.ts @@ -0,0 +1,31 @@ +import { describe, expect, it } from 'vitest'; + +import { formatChunkTokens, UNKNOWN_TOKEN_COUNT } from './chunkUtils'; + +type ChunkMetadata = Parameters[0]; + +describe('formatChunkTokens', () => { + it('formats a numeric count with separators', () => { + expect(formatChunkTokens({ token_count: 1234 })).toBe( + (1234).toLocaleString(), + ); + }); + + it('formats a count a store handed back as a string', () => { + expect(formatChunkTokens({ token_count: '1234' })).toBe( + (1234).toLocaleString(), + ); + }); + + it('falls back to a dash when the count is unusable', () => { + for (const token_count of [undefined, 0, -1, 'abc', '']) { + expect(formatChunkTokens({ token_count })).toBe(UNKNOWN_TOKEN_COUNT); + } + }); + + it('tolerates metadata the store returned as null', () => { + expect(formatChunkTokens(null as unknown as ChunkMetadata)).toBe( + UNKNOWN_TOKEN_COUNT, + ); + }); +}); diff --git a/frontend/src/components/chunkUtils.ts b/frontend/src/components/chunkUtils.ts new file mode 100644 index 00000000..56853106 --- /dev/null +++ b/frontend/src/components/chunkUtils.ts @@ -0,0 +1,27 @@ +/** + * Pure helpers for the chunk cards in the source viewer. Kept free of React so + * the formatting rules are unit-testable in isolation. + */ + +import { ChunkType } from '../settings/types'; + +/** What a chunk card shows when the token count is genuinely unknown. */ +export const UNKNOWN_TOKEN_COUNT = '-'; + +/** + * Render a chunk's token count for display. + * + * The backend fills the count in for chunks that were indexed without one, but + * stores round-trip metadata types differently: pgvector keeps JSON numbers + * while other backends hand the same value back as a string. Both are counts, + * so both get thousands separators; only a missing or unusable value falls + * back to a dash. + */ +export function formatChunkTokens(metadata: ChunkType['metadata']): string { + const raw = metadata?.token_count; + const count = typeof raw === 'string' ? Number(raw.trim()) : raw; + if (typeof count !== 'number' || !Number.isFinite(count) || count <= 0) { + return UNKNOWN_TOKEN_COUNT; + } + return count.toLocaleString(); +} diff --git a/frontend/src/settings/types/index.ts b/frontend/src/settings/types/index.ts index fe662e7b..81e557a8 100644 --- a/frontend/src/settings/types/index.ts +++ b/frontend/src/settings/types/index.ts @@ -3,7 +3,17 @@ import { ConfigRequirements } from '../../modals/types'; export type ChunkType = { doc_id: string; text: string; - metadata: { [key: string]: string }; + /** + * Chunk metadata as the vector store recorded it. Values are strings for + * most backends, but numbers (`token_count`) survive the round trip on the + * JSON-typed ones, so consumers must handle both. + */ + metadata: { + title?: string; + source?: string; + token_count?: number | string; + [key: string]: number | string | undefined; + }; }; export type LogEventType = From db8db62c937af7228782a3e6c3199ab952534276 Mon Sep 17 00:00:00 2001 From: arc53-machine <232052973+arc53-machine@users.noreply.github.com> Date: Tue, 22 Sep 2026 14:01:27 +0100 Subject: [PATCH 4/4] Treat a non-finite token count as unrecorded MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `float("inf") > 0` is true, so a stored `"inf"` was preserved as if it were a real count and reached the UI as "∞". Require a finite positive value, which also covers `nan` explicitly rather than by accident. --- docsgpt/api/user/sources/chunks.py | 12 ++++++++---- tests/api/user/sources/test_chunks.py | 5 ++++- 2 files changed, 12 insertions(+), 5 deletions(-) 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