Merge pull request #2823 from arc53-machine/fix-chunk-token-counts

Show token counts on source chunks

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
 
@@ -43,6 +45,59 @@ 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, 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 finite positive number.
+    """
+    raw = metadata.get("token_count")
+    if isinstance(raw, bool) or not isinstance(raw, (int, float, str)):
+        return False
+    try:
+        value = float(raw)
+    except (TypeError, ValueError):
+        return False
+    return math.isfinite(value) and value > 0
+
+
+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 +161,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(

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:

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<ChunksProps> = ({
                         <div className="w-full">
                           <div className="border-border bg-muted dark:border-border dark:bg-card flex w-full items-center justify-between border-b px-4 py-3">
                             <div className="dark:text-muted-foreground text-sm text-[#59636E]">
-                              {chunk.metadata.token_count
-                                ? chunk.metadata.token_count.toLocaleString()
-                                : '-'}{' '}
+                              {formatChunkTokens(chunk.metadata)}{' '}
                               {t('settings.sources.tokensUnit')}
                             </div>
                           </div>
@@ -663,9 +662,7 @@ const Chunks: React.FC<ChunksProps> = ({
                 <div className="border-border dark:border-border relative flex w-full flex-col overflow-hidden rounded-md border">
                   <div className="border-border bg-muted dark:border-border dark:bg-card flex w-full items-center justify-between border-b px-4 py-3">
                     <div className="dark:text-muted-foreground text-sm text-[#59636E]">
-                      {editingChunk.metadata.token_count
-                        ? editingChunk.metadata.token_count.toLocaleString()
-                        : '-'}{' '}
+                      {formatChunkTokens(editingChunk.metadata)}{' '}
                       {t('settings.sources.tokensUnit')}
                     </div>
                   </div>

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<typeof formatChunkTokens>[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,
+    );
+  });
+});

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();
+}

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 =

tests/api/user/sources/test_chunks.py

@@ -170,6 +170,115 @@ 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}, "inf", "-inf", "nan", float("inf")],
+    )
+    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
 

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
     ):