feat(ai-usage): track embedding usage separately with usage_type column#981
feat(ai-usage): track embedding usage separately with usage_type column#981Abh1shxkk wants to merge 11 commits intoInsForge:mainInsForge/InsForge:mainfrom Abh1shxkk:feat/embeddings-usage-type-979Abh1shxkk/InsForge:feat/embeddings-usage-type-979Copy head branch name to clipboard
Conversation
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughAdds a non-null Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
backend/src/services/ai/embedding.service.ts (1)
32-39:⚠️ Potential issue | 🔴 CriticalBuild-breaking type error:
sendRequest()does not return{ result, source }.The pipeline failure confirms this issue. Per
openrouter.provider.ts:340-388,sendRequest<T>()returnsPromise<T>directly (the rawCreateEmbeddingResponse), not an object withresultandsourceproperties.Either:
- Modify
OpenRouterProvider.sendRequest()to return{ result: T, source: string }, or- Revert to the previous pattern and find another way to detect BYOK usage
🔧 Option 2: Revert to direct response handling
- const { result: response, source } = await this.openRouterProvider.sendRequest((client) => + const response = await this.openRouterProvider.sendRequest((client) => client.embeddings.create({ model: options.model, input: options.input, encoding_format: options.encoding_format || 'float', dimensions: options.dimensions, }) );This would also require updating the BYOK check at line 57 to use a different mechanism.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/src/services/ai/embedding.service.ts` around lines 32 - 39, The call to openRouterProvider.sendRequest currently returns the raw CreateEmbeddingResponse (Promise<T>), not an object with { result, source }, so update the embedding call in embedding.service.ts to accept the direct response from sendRequest (e.g., assign the returned value to a variable like response or embeddingResult) and stop destructuring { result, source }; then adjust the BYOK detection logic (the check referencing source/BYOK) to use a different signal available on the raw CreateEmbeddingResponse or another provider method (e.g., inspect fields on the returned response or add a helper on OpenRouterProvider) so BYOK detection no longer relies on a non-existent source property.backend/src/services/ai/ai-usage.service.ts (1)
167-171:⚠️ Potential issue | 🟡 MinorInclude
usage_typein SELECT statement to return accurate data.The
getUsageByConfig()method omits theusage_typecolumn from its SELECT query. While Zod validation won't fail (the schema hasusageType: usageTypeSchema.default('chat')), this causes all returned records to haveusageType: 'chat'regardless of their actual database value, resulting in data inaccuracy. Other methods in the same file correctly select this column (e.g., line 268). The database column exists and stores distinct values: 'chat', 'embedding', and 'image_generation'.🛠️ Proposed fix to include usageType
let query = ` SELECT id, config_id as "configId", input_tokens as "inputTokens", output_tokens as "outputTokens", image_count as "imageCount", - image_resolution as "imageResolution", created_at as "createdAt" + image_resolution as "imageResolution", created_at as "createdAt", + usage_type as "usageType" FROM ai.usage WHERE config_id = $1 `;🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/src/services/ai/ai-usage.service.ts` around lines 167 - 171, The SELECT in getUsageByConfig is missing the usage_type column so returned rows get the Zod default ('chat') instead of the DB value; update the SQL in getUsageByConfig to include usage_type as "usageType" (matching other methods' mapping) so the query returns the actual usage_type values for downstream validation and mapping.
🧹 Nitpick comments (1)
backend/src/infra/database/migrations/025_add-usage-type-to-ai-usage.sql (1)
2-2: Consider adding an index onusage_typefor query performance.The
getUsageSummary()method usesFILTER (WHERE usage_type = 'embedding')clauses. As the table grows, an index onusage_typecould improve aggregation performance.📊 Optional index suggestion
CREATE INDEX IF NOT EXISTS idx_ai_usage_usage_type ON ai.usage (usage_type);🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@backend/src/infra/database/migrations/025_add-usage-type-to-ai-usage.sql` at line 2, Add a non-blocking index creation to the migration to improve query performance on the new usage_type column used by getUsageSummary(); specifically, in the migration that adds the usage_type column on table ai.usage add a CREATE INDEX IF NOT EXISTS (e.g., idx_ai_usage_usage_type on ai.usage (usage_type)) so aggregates and FILTER (WHERE usage_type = 'embedding') clauses run faster as the table grows. Ensure the index creation is included in the same migration flow (or a follow-up migration) and uses IF NOT EXISTS to be idempotent.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@backend/src/services/ai/embedding.service.ts`:
- Around line 56-63: The BYOK check using `source !== 'byok'` is ineffective
because `source` is not returned by `sendRequest()`; update the embedding usage
tracking in the method that calls `this.aiUsageService.trackEmbeddingUsage` to
use an explicit BYOK flag instead: add and read a boolean like `options.isByok`
(or `options.source`) passed into the provider call, or detect BYOK from
`aiConfig` (e.g., `aiConfig.isByok`) before invoking tracking; then change the
condition to `if (aiConfig?.id && tokenUsage && !options.isByok)` (or
`!aiConfig.isByok`) so BYOK requests are correctly excluded when calling
`trackEmbeddingUsage`.
---
Outside diff comments:
In `@backend/src/services/ai/ai-usage.service.ts`:
- Around line 167-171: The SELECT in getUsageByConfig is missing the usage_type
column so returned rows get the Zod default ('chat') instead of the DB value;
update the SQL in getUsageByConfig to include usage_type as "usageType"
(matching other methods' mapping) so the query returns the actual usage_type
values for downstream validation and mapping.
In `@backend/src/services/ai/embedding.service.ts`:
- Around line 32-39: The call to openRouterProvider.sendRequest currently
returns the raw CreateEmbeddingResponse (Promise<T>), not an object with {
result, source }, so update the embedding call in embedding.service.ts to accept
the direct response from sendRequest (e.g., assign the returned value to a
variable like response or embeddingResult) and stop destructuring { result,
source }; then adjust the BYOK detection logic (the check referencing
source/BYOK) to use a different signal available on the raw
CreateEmbeddingResponse or another provider method (e.g., inspect fields on the
returned response or add a helper on OpenRouterProvider) so BYOK detection no
longer relies on a non-existent source property.
---
Nitpick comments:
In `@backend/src/infra/database/migrations/025_add-usage-type-to-ai-usage.sql`:
- Line 2: Add a non-blocking index creation to the migration to improve query
performance on the new usage_type column used by getUsageSummary();
specifically, in the migration that adds the usage_type column on table ai.usage
add a CREATE INDEX IF NOT EXISTS (e.g., idx_ai_usage_usage_type on ai.usage
(usage_type)) so aggregates and FILTER (WHERE usage_type = 'embedding') clauses
run faster as the table grows. Ensure the index creation is included in the same
migration flow (or a follow-up migration) and uses IF NOT EXISTS to be
idempotent.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 915f2379-8ef7-452f-b889-e5bec0d4e549
📒 Files selected for processing (4)
backend/src/infra/database/migrations/025_add-usage-type-to-ai-usage.sqlbackend/src/services/ai/ai-usage.service.tsbackend/src/services/ai/embedding.service.tsshared-schemas/src/ai.schema.ts
|
on top of this we also want the correspoding ui to display embedding usages |
Hey @tonychang04, quick clarification — when we discussed surfacing embedding usage in the UI, I ended up opening a separate issue #978 (AI Usage Dashboard) and built it in a separate PR #982. The backend fixes (usage_type column, trackEmbeddingUsage) are in this PR, and the frontend (type badges, embedding token counts in summary) is in #982 which depends on this merging first. Two options going forward: Keep them separate — merge this first, then #982 brings the UI |
Hey @tonychang04, @Fermionic-Lyu just to clarify — I had already opened issue #978 (AI Usage Dashboard) before you raised #979, so I split the work into two PRs. The backend tracking is here in #981 and the frontend UI is in #982. Both are linked and #982 depends on this merging first. Let me know if you'd prefer I consolidate them into one PR! |
|
@Abh1shxkk I prefer to keep it seperate, for the focus in this repo, can we follow somthing like the old ui? for the other branch. you can give us the view for the new ui you propsed(that doesn't need mebedding) |
Hey @tonychang04, got it — keeping them separate. For #982, should I update the UI to follow the existing AIPage.tsx style (inline stats/table within the same page layout) rather than a separate dashboard page? And should I hold off on showing embedding-specific fields until #981 merges, or include them with a fallback? |
Hey @tonychang04, the UI in PR #982 already follows the same pattern as AIPage.tsx — same layout, same card/table structure, same styling conventions. The usageType field (chat/embedding/image) is there but gracefully handled until #981 merges. Would love your feedback on it! |
…#981 - Add AIUsagePage with stats cards: Total Requests, Input/Output Tokens, Images Generated, Embedding Requests, Embedding Tokens - Add date range filter (This week / This month / All time) - Add usage records table with TypeBadge (chat/embedding/image_generation) - Add /dashboard/ai/usage route in AppRoutes - Add Models/Usage secondary menu under Model Gateway - Support exact route matching in SecondaryMenu - Update shared-schemas: usageTypeSchema, embeddingRequests, embeddingTokens Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
frontend/src/features/ai/pages/AIUsagePage.tsx (1)
280-287: Use theme tokens for the usage badges.These
bg-blue-100/dark:bg-blue-900/30style pairs hard-code palette values inside a themed page, so they can drift from the rest of the app when theme tokens change. Please switch the badge variants to CSS-variable-backed styles instead.As per coding guidelines, "frontend/src/**/*.{tsx,jsx,css}: Maintain consistent color scheme using CSS variables".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@frontend/src/features/ai/pages/AIUsagePage.tsx` around lines 280 - 287, The badge color classes in the styles map (styles, keys: chat, embedding, image_generation) use hard-coded Tailwind palette classes; replace them with CSS-variable–backed tokens: add theme CSS variables for each badge (e.g., --badge-chat-bg, --badge-chat-text, and dark-mode counterparts) in your global/theme CSS, update the styles map to reference token-backed utility classes or neutral classes combined with inline style variables (so the className for the span uses those variables instead of bg-blue-100/dark:bg-blue-900/30), and ensure dark-mode variants are provided by the theme CSS (using .dark or prefers-color-scheme rules). Update AIUsagePage.tsx to use the new token-backed styles for the three keys (chat, embedding, image_generation) so badges follow the app theme tokens.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@frontend/src/features/ai/pages/AIUsagePage.tsx`:
- Around line 28-40: The filter labels say "This week"/"This month" but
getDateRange(range) computes rolling windows; change to calendar-anchored
ranges: for 'week' set startDate to the beginning of the current calendar week
(e.g., set time to 00:00:00 of the week's first day, choose Sunday or Monday
consistent with UI) and for 'month' set startDate to the first day of the
current month at 00:00:00 (use start.setDate(1) instead of subtracting a month
to avoid end-of-month rollover issues), keep endDate as now (or end of today if
preferred); update the same logic used in the other occurrence noted (lines
94-105) so both places use the anchored calendar-week and calendar-month
calculations in getDateRange and its duplicate.
---
Nitpick comments:
In `@frontend/src/features/ai/pages/AIUsagePage.tsx`:
- Around line 280-287: The badge color classes in the styles map (styles, keys:
chat, embedding, image_generation) use hard-coded Tailwind palette classes;
replace them with CSS-variable–backed tokens: add theme CSS variables for each
badge (e.g., --badge-chat-bg, --badge-chat-text, and dark-mode counterparts) in
your global/theme CSS, update the styles map to reference token-backed utility
classes or neutral classes combined with inline style variables (so the
className for the span uses those variables instead of
bg-blue-100/dark:bg-blue-900/30), and ensure dark-mode variants are provided by
the theme CSS (using .dark or prefers-color-scheme rules). Update
AIUsagePage.tsx to use the new token-backed styles for the three keys (chat,
embedding, image_generation) so badges follow the app theme tokens.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 76bd1543-6be2-4815-a80d-da228b581361
📒 Files selected for processing (5)
frontend/src/components/layout/SecondaryMenu.tsxfrontend/src/features/ai/pages/AIUsagePage.tsxfrontend/src/lib/routing/AppRoutes.tsxfrontend/src/lib/utils/menuItems.tsshared-schemas/src/ai.schema.ts
🚧 Files skipped from review as they are similar to previous changes (1)
- shared-schemas/src/ai.schema.ts
|
Hey @tonychang04, here's the frontend embedding UI running locally on the feat/embeddings-usage-type-979 branch: AI Usage page at /dashboard/ai/usage with 6 stats cards: Total Requests, Input Tokens, Output Tokens, Images Generated, Embedding Requests, and Embedding Tokens
|
|
can you rebase the changes @Abh1shxkk and i can take a look at the embeddings ui once it is exposed? |
429c362 to
7abad0b
Compare
…#981 - Add AIUsagePage with stats cards: Total Requests, Input/Output Tokens, Images Generated, Embedding Requests, Embedding Tokens - Add date range filter (This week / This month / All time) - Add usage records table with TypeBadge (chat/embedding/image_generation) - Add /dashboard/ai/usage route in AppRoutes - Add Models/Usage secondary menu under Model Gateway - Support exact route matching in SecondaryMenu - Update shared-schemas: usageTypeSchema, embeddingRequests, embeddingTokens Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
42245ad to
4f20715
Compare
|
Hey @tonychang04 Rebased the branch on latest main and resolved all the merge conflicts. Here's what changed: The usage tracking is now integrated directly into the Model Gateway page instead of being a separate page. There's a Usage button next to Gateway credentials that opens a clean modal showing all the stats. The modal has date filters (Today, 7 days, 30 days, All time) so you can drill into specific time periods. It shows total requests, input/output tokens, images generated, embedding requests and embedding tokens at a glance. Below the stats there's a scrollable activity table showing each request with the date, time, model name, type (Chat, Embed, Image) and token counts. Also fixed the migration numbering conflict that came up during the rebase. Our usage_type migration is now 029 and the admin email fix is 030 so they run after all the upstream migrations. The sidebar we had before (Models/Usage) is removed since everything lives on one page now. All files pass ESLint and Prettier checks. 
|
|
please let me know if there is any issue i am free today @tonychang04 |
|
howevery for this pr embeddings can only be tracked when click usage... in the dashboard there is no place to enable and disable embeddings right? |
@tonychang04 just to clarify, clicking the Usage button does not control whether embeddings are tracked or not, it only opens the usage view. In this PR, embedding requests are tracked automatically on the backend whenever the embeddings endpoint is used, and they show up separately in the usage modal with their own type and counts. You’re right though that there isn’t a separate place on the main Model Gateway dashboard to explicitly enable/disable embeddings as a distinct UI flow in this branch. This PR is mainly about fixing the tracking so embeddings are no longer counted as chat. If you want, I can handle the dashboard-side UX for embeddings separately in a follow-up |
|
@Abh1shxkk yah if you can have 1 pr. that
That is the goal of the issue |
…age_generation requests - Migration 025: adds usage_type column (default 'chat'), backfills image_generation rows - shared-schemas: adds usageTypeSchema, usageType field to aiUsageDataSchema/aiUsageRecordSchema, embeddingRequests/embeddingTokens to aiUsageSummarySchema - ai-usage.service: new trackEmbeddingUsage(), trackChatUsage/trackImageGenerationUsage write explicit usage_type, getUsageSummary returns embedding breakdown, getAllUsage includes usage_type in SELECT - embedding.service: use trackEmbeddingUsage() instead of trackChatUsage(), also thread source from sendRequest to skip tracking for BYOK Closes InsForge#979 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Cast sendRequest return type explicitly to fix TS2339 errors on result/source destructuring. Add OpenAI.Embedding type annotation to fix implicit any on map callback. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…rom sendRequest
sendRequest() on this branch returns T not { result, source }, so source was
always undefined and the BYOK guard was ineffective. Switch to isByokActive()
which correctly detects BYOK from the provider.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…#981 - Add AIUsagePage with stats cards: Total Requests, Input/Output Tokens, Images Generated, Embedding Requests, Embedding Tokens - Add date range filter (This week / This month / All time) - Add usage records table with TypeBadge (chat/embedding/image_generation) - Add /dashboard/ai/usage route in AppRoutes - Add Models/Usage secondary menu under Model Gateway - Support exact route matching in SecondaryMenu - Update shared-schemas: usageTypeSchema, embeddingRequests, embeddingTokens Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…te ranges - Add usageTypeSchema enum and usageType field to aiUsageDataSchema so ai-usage.service.ts trackUsage() typechecks correctly - Fix getDateRange: 'week' now anchors to start of current calendar week (Sunday 00:00:00), 'month' anchors to first day of current month 00:00:00 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
- Show last 10 usage records in AIPage with type badges (chat/embedding/image_generation) - Remove separate AIUsagePage, route, and secondary menu — fits existing UI Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ture Update @/components alias to relative import path to match the packages/dashboard package convention. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
4f20715 to
e9cd568
Compare
- AILayout now renders AISidebar + Outlet (Models / Usage tabs) - Register /dashboard/ai/usage route pointing to AIUsagePage - Remove UsageDialog modal and Usage button from AIPage — sidebar nav replaces it - Fix packages/dashboard tsconfig: add rootDir to silence TS warning Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Summary
Closes #979
025— addsusage_type VARCHAR(20) NOT NULL DEFAULT 'chat'toai.usage, backfills existing image generation rows to'image_generation'shared-schemas— addsusageTypeSchema('chat' | 'embedding' | 'image_generation'),usageTypefield toaiUsageDataSchemaandaiUsageRecordSchema, andembeddingRequests/embeddingTokensbreakdown toaiUsageSummarySchemaai-usage.service— newtrackEmbeddingUsage()method;trackChatUsageandtrackImageGenerationUsagenow write explicitusage_type;getUsageSummaryreturns per-type breakdown;getAllUsageincludesusage_typein SELECTembedding.service— callstrackEmbeddingUsage()instead oftrackChatUsage(), threadssourcefromsendRequestto skip tracking for BYOKHow did you test this change?
docker exec -it insforge-postgres psql -U postgres -d insforge -c "SELECT column_name, data_type, column_default FROM information_schema.columns WHERE table_schema='ai' AND table_name='usage' AND column_name='usage_type';"
→ Column added with DEFAULT 'chat' ✅
curl -X POST http://localhost:7130/api/ai/chat/completion
-H "Authorization: Bearer ik_local_dev_access_key_..."
-H "Content-Type: application/json"
-d '{"model":"meta-llama/llama-3.2-1b-instruct","messages":[{"role":"user","content":"say hi"}]}'
DB result:
usage_type | input_tokens | output_tokens
-----------+--------------+--------------
chat | 13 | 3
✅
GET /api/ai/usage?limit=5
Response includes "usageType":"chat" on each record ✅
GET /api/ai/usage/summary
Response:
{
"totalInputTokens": 13,
"totalOutputTokens": 3,
"totalTokens": 16,
"totalImageCount": 0,
"totalRequests": 1,
"embeddingRequests": 0,
"embeddingTokens": 0
}
embeddingRequests and embeddingTokens fields present ✅
Summary by CodeRabbit
New Features
Bug Fixes
Chores
Note
Track embedding AI usage separately with a
usage_typecolumnusage_typeVARCHAR(20) column (default'chat') toai.usagevia migration, backfilling rows withimage_count > 0as'image_generation'.trackEmbeddingUsagetoAIUsageServiceand updatesEmbeddingServiceto call it instead oftrackChatUsage, so embedding requests are now recorded withusage_type='embedding'.getUsageSummaryandgetAllUsageto returnembeddingRequests,embeddingTokens, andusageTypeper record./dashboard/ai/usagewith date filters, pagination, and per-recordusageTypebadges, plus a sidebar for AI section navigation.Macroscope summarized f4d61e9.