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

Show token counts on source chunks

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