Expose count_tokens() and add refresh_token_counts() helper#22
Expose count_tokens() and add refresh_token_counts() helper#22syedkazim110 wants to merge 4 commits into
Conversation
|
Warning Review limit reached
More reviews will be available in 56 minutes and 57 seconds. Learn how PR review limits work. Your organization has run out of usage credits. Purchase more in the billing tab. ⌛ How to resolve this issue?After more reviews become available, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available. Please see our Fair Usage Limits Policy for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Pro Run ID: ⛔ Files ignored due to path filters (1)
📒 Files selected for processing (5)
📝 WalkthroughWalkthroughThis PR adds token counting functionality to the pgedge_vectorizer extension. It introduces 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Up to standards ✅🟢 Issues
|
| Metric | Results |
|---|---|
| Duplication | 0 |
NEW Get contextual insights on your PRs based on Codacy's metrics, along with PR and Jira context, without leaving GitHub. Enable AI reviewer
TIP This summary will be updated as you push new changes.
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
test/sql/count_tokens.sql (1)
72-75: ⚡ Quick winAssert exact updated-row count instead of
>= 1.These checks are too weak for the stated contract (“returns number of rows updated”). Please assert equality with table row count so partial updates cannot pass silently.
Proposed test tightening
-SELECT pgedge_vectorizer.refresh_token_counts( - 'count_tokens_refresh_test_content_chunks'::regclass -) >= 1 AS refresh_returned_count; +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_refresh_test_content_chunks'::regclass +) = ( + SELECT COUNT(*) FROM count_tokens_refresh_test_content_chunks +) AS refresh_returned_count; ... -SELECT pgedge_vectorizer.refresh_token_counts( - 'count_tokens_ns_test.ns_chunks'::regclass -) >= 1 AS refresh_returned_count; +SELECT pgedge_vectorizer.refresh_token_counts( + 'count_tokens_ns_test.ns_chunks'::regclass +) = ( + SELECT COUNT(*) FROM count_tokens_ns_test.ns_chunks +) AS refresh_returned_count;Also applies to: 108-110
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@test/sql/count_tokens.sql` around lines 72 - 75, Replace the weak >= 1 assertion with an exact equality check: call pgedge_vectorizer.refresh_token_counts('count_tokens_refresh_test_content_chunks'::regclass) and assert it equals the actual row count of the target table (e.g. compare to (SELECT COUNT(*) FROM count_tokens_refresh_test_content_chunks)), so the test enforces the exact number of rows updated; apply the same change for the other occurrence around the pgedge_vectorizer.refresh_token_counts call on lines ~108-110.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@sql/pgedge_vectorizer--1.0--1.1.sql`:
- Around line 994-1007: The refresh_token_counts function uses
pgedge_vectorizer.count_tokens(content) but the normal chunk write paths still
set token_count = length(chunk_text) / 4; update all chunk INSERT/UPDATE write
paths to compute token_count with pgedge_vectorizer.count_tokens(content) (i.e.,
replace occurrences assigning token_count from length(chunk_text)/4) so
inserted/updated rows use pgedge_vectorizer.count_tokens(content) consistently;
apply this change to every chunk write path in this migration (the INSERT/UPDATE
statements around the earlier referenced token_count assignments) and mirror the
same edits in the corresponding sql/pgedge_vectorizer--1.1.sql file.
---
Nitpick comments:
In `@test/sql/count_tokens.sql`:
- Around line 72-75: Replace the weak >= 1 assertion with an exact equality
check: call
pgedge_vectorizer.refresh_token_counts('count_tokens_refresh_test_content_chunks'::regclass)
and assert it equals the actual row count of the target table (e.g. compare to
(SELECT COUNT(*) FROM count_tokens_refresh_test_content_chunks)), so the test
enforces the exact number of rows updated; apply the same change for the other
occurrence around the pgedge_vectorizer.refresh_token_counts call on lines
~108-110.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 954868f4-7d14-46ba-b7bc-78efe6b85103
⛔ Files ignored due to path filters (1)
test/expected/count_tokens.outis excluded by!**/*.out
📒 Files selected for processing (4)
Makefilesql/pgedge_vectorizer--1.0--1.1.sqlsql/pgedge_vectorizer--1.1.sqltest/sql/count_tokens.sql
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
There was a problem hiding this comment.
🧹 Nitpick comments (1)
sql/pgedge_vectorizer--1.1.sql (1)
1000-1004: ⚡ Quick winSkip no-op rewrites in
refresh_token_counts().The
UPDATErecomputes and rewritestoken_countfor every row, even those already correct. On a large chunk table this produces unnecessary dead tuples/WAL. Filtering on a real change also makes the returnedROW_COUNTreflect rows that actually changed, matching the "number of rows updated" wording in the comment.♻️ Proposed change to avoid redundant writes
EXECUTE format( - 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content)', + 'UPDATE %s SET token_count = pgedge_vectorizer.count_tokens(content) + WHERE token_count IS DISTINCT FROM pgedge_vectorizer.count_tokens(content)', p_chunk_table );Note: this changes the return value to count only rows whose
token_countactually changed. The Test 4 assertion (>= 1) still holds since all rows are forced to0beforehand. If callers rely on the count equaling total row count, keep the unconditional form.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@sql/pgedge_vectorizer--1.1.sql` around lines 1000 - 1004, In refresh_token_counts(), avoid rewriting unchanged rows by changing the EXECUTE UPDATE that sets token_count via pgedge_vectorizer.count_tokens(content) to only update rows where the computed value differs from the stored one (use "IS DISTINCT FROM" or equivalent comparison against the computed token count); build the dynamic SQL using p_chunk_table and the same pgedge_vectorizer.count_tokens(content) expression in a WHERE clause (or JOIN to a subquery that computes new_tokens) so GET DIAGNOSTICS rows_updated will reflect only rows that actually changed.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@sql/pgedge_vectorizer--1.1.sql`:
- Around line 1000-1004: In refresh_token_counts(), avoid rewriting unchanged
rows by changing the EXECUTE UPDATE that sets token_count via
pgedge_vectorizer.count_tokens(content) to only update rows where the computed
value differs from the stored one (use "IS DISTINCT FROM" or equivalent
comparison against the computed token count); build the dynamic SQL using
p_chunk_table and the same pgedge_vectorizer.count_tokens(content) expression in
a WHERE clause (or JOIN to a subquery that computes new_tokens) so GET
DIAGNOSTICS rows_updated will reflect only rows that actually changed.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: 06364c6c-0043-499d-a071-2233c9072810
📒 Files selected for processing (2)
sql/pgedge_vectorizer--1.0--1.1.sqlsql/pgedge_vectorizer--1.1.sql
🚧 Files skipped from review as they are similar to previous changes (1)
- sql/pgedge_vectorizer--1.0--1.1.sql
|
@davepage here's the count_tokens() SQL exposure + refresh_token_counts() helper you mentioned keeping from the closed tiktoken PR. No plpython3u, no GUC, no tiktoken dependency — just the two functions and 5 regression tests. Let me know if anything needs adjusting. |
dpage
left a comment
dpage
left a comment
There was a problem hiding this comment.
Review summary
The refresh_token_counts() design and test coverage here are good, but there's a critical, merge-blocking defect: the PR declares a C function symbol that does not exist, which breaks not only the new functions but the core chunking pipeline after upgrade.
⚠️ Note: this wasn't caught by the green checks because the GitHub Actions CI job that runsmake installcheck(.github/workflows/ci.yml:123) is inaction_requiredstate for this fork PR and never ran — only Codacy and CodeRabbit executed, neither of which builds/runs the extension (CodeRabbit also ignores.outfiles via the!**/*.outfilter).
🔴 Must Fix — count_tokens() references a non-existent C symbol
Both sql/pgedge_vectorizer--1.1.sql:136 and sql/pgedge_vectorizer--1.0--1.1.sql:136 declare:
CREATE FUNCTION pgedge_vectorizer.count_tokens(p_text TEXT) RETURNS INT
AS 'MODULE_PATHNAME', 'pgedge_vectorizer_count_tokens_sql'
LANGUAGE C IMMUTABLE STRICT;The PR description calls this "a SQL declaration for the existing C count_tokens() function" — but the existing count_tokens in src/tokenizer.c:32 is an internal C function (int count_tokens(const char *text, const char *model)), not a PG_FUNCTION_INFO_V1 SQL wrapper. The symbol pgedge_vectorizer_count_tokens_sql is defined nowhere — it appears only in these two SQL files. This PR changes no C files.
I verified this by building the branch locally (PG 17.4) and inspecting the compiled library:
$ nm -gU pgedge_vectorizer.dylib | grep -i count_tokens
(no matches)
$ nm -gU pgedge_vectorizer.dylib | grep -i bm25_tokenize # a real SQL fn, for contrast
_pg_finfo_pgedge_vectorizer_bm25_tokenize
_pgedge_vectorizer_bm25_tokenize
Consequences
- Any call to
pgedge_vectorizer.count_tokens(...)fails at runtime withERROR: could not find function "pgedge_vectorizer_count_tokens_sql" in file ".../pgedge_vectorizer.so". - Worse: the PR rewrites all three chunk-write paths (
recreate_chunks()and the trigger/existing-row insert paths — lines 358, 675, 970 in--1.1.sql) to callpgedge_vectorizer.count_tokens(chunk_text). So after upgrading to 1.1, every chunking operation errors out —enable_vectorization(), the insert trigger, andrecreate_chunks()all break. This is a severe regression, not just a broken new helper. refresh_token_counts()is dead on arrival for the same reason.
Suggested fix
Add the SQL-callable wrapper in src/tokenizer.c:
PG_FUNCTION_INFO_V1(pgedge_vectorizer_count_tokens_sql);
Datum
pgedge_vectorizer_count_tokens_sql(PG_FUNCTION_ARGS)
{
text *txt = PG_GETARG_TEXT_PP(0);
char *s = text_to_cstring(txt);
PG_RETURN_INT32(count_tokens(s, pgedge_vectorizer_model));
}(plus the prototype in src/pgedge_vectorizer.h). Then regenerate test/expected/count_tokens.out from an actual run — the committed .out cannot have come from a passing installcheck, since the symbol does not resolve.
🟡 Should Fix — get the fork CI run approved
The installcheck regression suite never executed for this PR (fork workflow pending approval). Once the C fix is in, please have a maintainer approve the workflow run so the new count_tokens regression test actually runs before merge. The absent regression run should be treated as a blocker in its own right.
ℹ️ Notes (non-blocking)
- Semantics change in the chunk-write paths. The old inline
length(chunk_text) / 4used byte length with floor division;count_tokens()uses UTF-8 character count with ceiling division(n+3)/4. This is an improvement (it makes the plpgsql paths consistent with the C chunker, which already usescount_tokens) and is correctly reflected in the tests ('你好世界'→ 1). Worth a changelog line since storedtoken_countvalues will change for multibyte/short text. refresh_token_counts()evaluatescount_tokens(content)twice (inSETand theIS DISTINCT FROMguard). Minor; acceptable since the function isIMMUTABLEand the no-op-skip guard is a nice addition.format('UPDATE %s ...', p_chunk_table)is injection-safe becausep_chunk_tableisregclass(output is a validated, properly-quoted, schema-qualified identifier). Consistent with existing code. 👍
Once the C wrapper lands and the regression suite passes for real, this is a nice addition. Thanks @syedkazim110!
Add a SQL declaration for the existing C count_tokens() function so users can call pgedge_vectorizer.count_tokens(text) directly to inspect approximate token counts (~4 chars/token, UTF-8 aware). Add refresh_token_counts(chunk_table regclass) to back-fill the token_count column on existing chunk rows without a full recreate_chunks() rebuild. Returns the number of rows updated. Includes regression tests covering approximation values, NULL/UTF-8 handling, the refresh path, and schema-qualified chunk tables.
The three PLpgSQL insert paths (enable_vectorization, vectorization_trigger, recreate_chunks) were using length(chunk_text) / 4 (floor division) while count_tokens() and refresh_token_counts() use ceiling division ((char_count + 3) / 4). This meant refresh_token_counts() would silently change values the trigger had just written for any chunk whose length is not divisible by 4. Replace all six occurrences (three per SQL file) with pgedge_vectorizer.count_tokens(chunk_text) so every write path uses the same formula.
…ions Add WHERE token_count IS DISTINCT FROM count_tokens(content) so the UPDATE only touches rows whose stored value actually differs. On large chunk tables this avoids unnecessary WAL writes and dead tuples. The return value now reflects rows that changed, not rows scanned. Tighten the two refresh test assertions from >= 1 to = COUNT(*) so a partial update cannot silently pass the regression test.
The SQL migration files declared pgedge_vectorizer_count_tokens_sql as a C symbol but the wrapper was never added to src/tokenizer.c, causing every call to pgedge_vectorizer.count_tokens() (and all three chunk-write paths that call it) to fail with "could not find function" at runtime. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
syedkazim110
force-pushed
the
feature/count-tokens-expose
branch
from
0dd7f16 to
74f9073
Compare
2 participants
Add a SQL declaration for the existing C count_tokens() function so users can call pgedge_vectorizer.count_tokens(text) directly to inspect approximate token counts (~4 chars/token, UTF-8 aware).
Add refresh_token_counts(chunk_table regclass) to back-fill the token_count column on existing chunk rows without a full recreate_chunks() rebuild. Returns the number of rows updated.
Includes regression tests covering approximation values, NULL/UTF-8 handling, the refresh path, and schema-qualified chunk tables.