mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 02:14:23 +00:00
fix(wiki): maintain sources.tokens for wikis (sum of page token_counts)
sources.tokens was only set at ingest, so wikis showed no/stale token count on the card (blank wikis showed nothing; converted ones showed the stale original count). rebuild_wiki_directory_structure (called after every wiki mutation) now also sets sources.tokens to the sum of wiki_pages.token_count; blank wiki create sets tokens=0. So the card reflects live wiki content.
This commit is contained in:
1 parent
2616ec5d1b
commit
f809f660f4
4 files changed
+42
-1
No files matched your search
@@ -763,6 +763,7 @@ class CreateWikiSource(Resource):
|
||||
type="wiki",
|
||||
config={"kind": "wiki"},
|
||||
directory_structure={},
|
||||
tokens=0,
|
||||
)
|
||||
if initial_content:
|
||||
WikiPagesRepository(conn).upsert(
|
||||
|
||||
@@ -249,7 +249,8 @@ def rebuild_wiki_directory_structure(
|
||||
|
||||
pages = WikiPagesRepository(conn).list_for_source(source_id)
|
||||
tree = build_wiki_directory_structure(pages)
|
||||
total_tokens = sum(int(page.get("token_count") or 0) for page in pages)
|
||||
SourcesRepository(conn).update(
|
||||
source_id, owner_id, {"directory_structure": tree}
|
||||
source_id, owner_id, {"directory_structure": tree, "tokens": total_tokens}
|
||||
)
|
||||
return tree
|
||||
@@ -94,6 +94,7 @@ class TestCreateWikiSource:
|
||||
row = SourcesRepository(pg_conn).get_any(source_id, user)
|
||||
assert row["type"] == "wiki"
|
||||
assert row["config"]["kind"] == "wiki"
|
||||
assert int(row["tokens"]) == 0
|
||||
# No seed content → no re-embed, and never any ingest/reingest task.
|
||||
mock_reembed.assert_not_called()
|
||||
mock_ingest.assert_not_called()
|
||||
@@ -105,6 +106,7 @@ class TestCreateWikiSource:
|
||||
WikiPage,
|
||||
WIKI_INDEX_PATH,
|
||||
)
|
||||
from application.storage.db.repositories.sources import SourcesRepository
|
||||
|
||||
user = "u-wiki-seed"
|
||||
with _patch_db(pg_conn), patch(
|
||||
@@ -124,6 +126,8 @@ class TestCreateWikiSource:
|
||||
|
||||
assert response.status_code == 200
|
||||
source_id = response.json["source_id"]
|
||||
row = SourcesRepository(pg_conn).get_any(source_id, user)
|
||||
assert int(row["tokens"]) > 0
|
||||
# Re-embed fires once, only for the seed page; no ingest/reingest ever.
|
||||
mock_reembed.assert_called_once()
|
||||
assert mock_reembed.call_args.args[0] == source_id
|
||||
|
||||
@@ -239,6 +239,41 @@ class TestDirectoryStructure:
|
||||
stored = SourcesRepository(pg_conn).get(sid, "rebuild-user")
|
||||
assert "a.md" in (stored["directory_structure"] or {})
|
||||
|
||||
def test_rebuild_sets_source_tokens_to_sum(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
owner = "rebuild-tokens-user"
|
||||
sid = str(
|
||||
SourcesRepository(pg_conn).create(
|
||||
"rebuild-tokens-src", user_id=owner, type="wiki"
|
||||
)["id"]
|
||||
)
|
||||
page_a = repo.upsert(sid, "/a.md", "alpha beta gamma")
|
||||
page_b = repo.upsert(sid, "/dir/b.md", "delta epsilon")
|
||||
expected = page_a["token_count"] + page_b["token_count"]
|
||||
|
||||
rebuild_wiki_directory_structure(pg_conn, sid, owner)
|
||||
|
||||
stored = SourcesRepository(pg_conn).get(sid, owner)
|
||||
assert int(stored["tokens"]) == expected
|
||||
|
||||
def test_rebuild_after_delete_lowers_source_tokens(self, pg_conn):
|
||||
repo = _repo(pg_conn)
|
||||
owner = "rebuild-delete-user"
|
||||
sid = str(
|
||||
SourcesRepository(pg_conn).create(
|
||||
"rebuild-delete-src", user_id=owner, type="wiki"
|
||||
)["id"]
|
||||
)
|
||||
page_a = repo.upsert(sid, "/a.md", "alpha beta gamma")
|
||||
repo.upsert(sid, "/b.md", "delta epsilon")
|
||||
rebuild_wiki_directory_structure(pg_conn, sid, owner)
|
||||
|
||||
repo.delete_by_path(sid, "/b.md")
|
||||
rebuild_wiki_directory_structure(pg_conn, sid, owner)
|
||||
|
||||
stored = SourcesRepository(pg_conn).get(sid, owner)
|
||||
assert int(stored["tokens"]) == page_a["token_count"]
|
||||
|
||||
|
||||
class TestSetEmbedStatus:
|
||||
def test_set_embed_status(self, pg_conn):
|
||||
|
||||
Reference in new issue
Block a user