mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-14 08:17:34 +00:00
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.
This commit is contained in:
1 parent
f5b02ed374
commit
0aecdc32c9
2 files changed
+158
-1
No files matched your search
@@ -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(
|
||||
|
||||
Reference in new issue
Block a user