Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 27 additions & 2 deletions apps/sim/app/api/v1/knowledge/route.integration.ts
Original file line number Diff line number Diff line change
@@ -1,18 +1,21 @@
/**
* Knowledge-base document totals against real PostgreSQL: the public v1 list and detail count
* only the documents their caller can read, and the internal list reads no document at all
* unless the caller asks for totals.
* unless the caller asks for totals. A request without the flag is counted, since a page loaded
* before the flag existed requires both totals on every row.
*/
import type { Principal } from '@sim/auth/principal'
import { db } from '@sim/db'
import { document, organization, user, workspace } from '@sim/db/schema'
import { createMockRequest } from '@sim/testing'
import { authMock, authMockFns, createMockRequest } from '@sim/testing'
import { generateId } from '@sim/utils/id'
import { eq, inArray } from 'drizzle-orm'
import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'

const caller = vi.hoisted(() => ({ userId: '' }))

vi.mock('@/lib/auth', () => authMock)

vi.mock('@/app/api/v1/middleware', async (importOriginal) => ({
...(await importOriginal<typeof import('@/app/api/v1/middleware')>()),
authenticateRequest: async () => ({
Expand All @@ -34,6 +37,7 @@ import {
seedKnowledgeAclFixture,
} from '@/lib/knowledge/__integration__/seed-source-access-fixture'
import { listInternalKnowledgeBases } from '@/lib/knowledge/application/knowledge-bases'
import { GET as listInternalKnowledgeBasesRoute } from '@/app/api/knowledge/route'
import { GET as getKnowledgeBase } from '@/app/api/v1/knowledge/[id]/route'
import { GET as listKnowledgeBases } from '@/app/api/v1/knowledge/route'

Expand All @@ -46,6 +50,10 @@ describe('knowledge-base document totals in PostgreSQL', () => {
throw new Error('Unexpected provider request in knowledge-base count tests')
})
caller.userId = ids.bobId
authMockFns.mockGetSession.mockResolvedValue({
user: { id: ids.bobId },
session: { id: 'fixture-reader' },
})
await seedKnowledgeAclFixture(ids, { connectorType: 'google_drive' })
await db.insert(document).values([
{
Expand Down Expand Up @@ -134,4 +142,21 @@ describe('knowledge-base document totals in PostgreSQL', () => {
).knowledgeBases
expect(counted).toMatchObject({ docCount: 1, tokenCount: 10 })
})

it('counts an internal list request that omits the flag', async () => {
const list = async (query: string) => {
const response = await listInternalKnowledgeBasesRoute(
createMockRequest('GET', undefined, {}, `http://localhost/api/knowledge?${query}`),
{ params: Promise.resolve({}) }
)
expect(response.status).toBe(200)
return (await response.json()).data
}
const workspaceQuery = `workspaceId=${ids.workspaceId}&scope=active`
expect(await list(workspaceQuery)).toEqual([
expect.objectContaining({ id: ids.knowledgeBaseId, docCount: 1, tokenCount: 10 }),
])
const [uncounted] = await list(`${workspaceQuery}&includeCounts=false`)
expect(uncounted).not.toHaveProperty('tokenCount')
})
})
Original file line number Diff line number Diff line change
Expand Up @@ -250,7 +250,7 @@ export function useMentionData(props: UseMentionDataProps): MentionDataReturn {
try {
setIsLoadingKnowledge(true)
const result = await requestJson(listKnowledgeBasesContract, {
query: { workspaceId },
query: { workspaceId, includeCounts: false },
})
const items = result.data
const sorted = [...items].sort((a, b) => {
Expand Down
8 changes: 6 additions & 2 deletions apps/sim/lib/api/contracts/knowledge/base.ts
Original file line number Diff line number Diff line change
Expand Up @@ -26,8 +26,12 @@ export type KnowledgeScope = z.output<typeof knowledgeScopeSchema>
export const listKnowledgeBasesQuerySchema = z.object({
workspaceId: z.string().min(1).optional(),
scope: knowledgeScopeSchema.default('active'),
/** Adds `docCount` and `tokenCount` for the documents the caller can see; costs a document scan. */
includeCounts: booleanQueryFlagSchema.optional().default(false),
/**
* Adds `docCount` and `tokenCount` for the documents the caller can see; costs a document scan.
* Absent means counted: a page loaded before this flag existed requires both totals on every
* row, and current clients always send it.
*/
includeCounts: booleanQueryFlagSchema.optional().default(true),
Comment thread
waleedlatif1 marked this conversation as resolved.
})

/**
Expand Down
Loading