diff --git a/server/src/controllers/vectorDB/hardSyncVectorDB.ts b/server/src/controllers/vectorDB/hardSyncVectorDB.ts new file mode 100644 index 000000000..0bbf8c29a --- /dev/null +++ b/server/src/controllers/vectorDB/hardSyncVectorDB.ts @@ -0,0 +1,14 @@ +import { hardSyncVectorDB } from 'src/services/vectorDB/hardSyncVectorDBService'; +import { publicProcedure } from '../../trpc'; + +export function hardSyncVectorDBRoute() { + return publicProcedure.mutation(async () => { + try { + await hardSyncVectorDB(); + return { message: 'Vectors synced successfully!' }; + } catch (error) { + console.error(`Failed to sync vectors: ${error.message}`, error); + throw error; + } + }); +} diff --git a/server/src/controllers/vectorDB/index.ts b/server/src/controllers/vectorDB/index.ts new file mode 100644 index 000000000..b8f2974c6 --- /dev/null +++ b/server/src/controllers/vectorDB/index.ts @@ -0,0 +1 @@ +export * from './hardSyncVectorDB'; diff --git a/server/src/integrations/ai/client.ts b/server/src/integrations/ai/client.ts index 4f9a4c7d2..9f8e08338 100644 --- a/server/src/integrations/ai/client.ts +++ b/server/src/integrations/ai/client.ts @@ -13,6 +13,13 @@ class AiClient { this.EXECUTE_AI_MODEL_URL = `https://api.cloudflare.com/client/v4/accounts/${this.accountId}/ai/run/${this.MODEL_NAME}`; } + /** + * @returns The name of the model to use to generate embeddings. + */ + public static get modelName(): string { + return AiClient.instance.MODEL_NAME; + } + public static get instance(): AiClient { if (!AiClient._instance) { throw new Error('AiClient instance not initialized.'); @@ -33,13 +40,15 @@ class AiClient { } public async run(text: string | string[]) { + const body = JSON.stringify({ text }); + console.log({ bodySizeKb: body.length / 1024 }); const response = await fetch(this.EXECUTE_AI_MODEL_URL, { method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${this.apiKey}`, }, - body: JSON.stringify({ text }), + body, }); if (!response.ok) { @@ -61,7 +70,7 @@ class AiClient { contentList: string[], transform: (embedding: number[], index: number) => T, ): Promise { - const MAX_BATCH_SIZE = 100; // REF: https://developers.cloudflare.com/workers-ai/models/bge-base-en-v1.5/ + const MAX_BATCH_SIZE = 25; // 0; // REF: https://developers.cloudflare.com/workers-ai/models/bge-base-en-v1.5/ const MAX_ROUND = Math.ceil(contentList.length / MAX_BATCH_SIZE); const result: T[] = []; @@ -72,10 +81,13 @@ class AiClient { (round + 1) * MAX_BATCH_SIZE, ); + console.log(`[${round + 1}/${MAX_ROUND}] Batch size: ${batch.length}`); const { result: { data }, } = await AiClient.instance.run(batch); + // await new Promise((resolve) => setTimeout(resolve, 1000)); + // Flatten the result for (let idx = 0; idx < batch.length; idx++) { const embedding = data[idx]; diff --git a/server/src/routes/trpcRouter.ts b/server/src/routes/trpcRouter.ts index e6b16f804..a42f10f22 100644 --- a/server/src/routes/trpcRouter.ts +++ b/server/src/routes/trpcRouter.ts @@ -110,6 +110,7 @@ import { addPackTemplateRoute, importPackTemplatesRoute, } from '../controllers/packTemplates'; +import { hardSyncVectorDBRoute } from '../controllers/vectorDB'; import { router as trpcRouter } from '../trpc'; @@ -218,6 +219,8 @@ export const appRouter = trpcRouter({ addToFavorite: addToFavoriteRoute(), getUserFavorites: getUserFavoritesRoute(), getFavoritePacksByUser: getFavoritePacksByUserRoute(), + // Vectorize routes + hardSyncVectorDB: hardSyncVectorDBRoute(), }); export type AppRouter = typeof appRouter; diff --git a/server/src/services/vectorDB/hardSyncVectorDBService.ts b/server/src/services/vectorDB/hardSyncVectorDBService.ts new file mode 100644 index 000000000..ba8c590e5 --- /dev/null +++ b/server/src/services/vectorDB/hardSyncVectorDBService.ts @@ -0,0 +1,28 @@ +import { VectorClient } from 'src/vector/client'; +import * as ItemService from 'src/services/item/item.service'; +import { summarizeItem } from 'src/utils/item'; + +export const hardSyncVectorDB = async () => { + const deleteVectorDBResponse = await VectorClient.instance.deleteVectorDB(); + console.log('deleteVectorDBResponse', deleteVectorDBResponse); + + const createVectorDBResponse = await VectorClient.instance.createVectorDB(); + console.log('createVectorDBResponse', createVectorDBResponse); + + const { items } = await ItemService.getItemsGloballyService(0, 0); + console.log('itemsCount', items.length); + + const records = []; + for (const item of items) { + records.push({ + id: item.id, + content: summarizeItem(item), + namespace: 'item', + metadata: { + id: item.id, + global: item.global, + }, + }); + } + await VectorClient.instance.syncRecords(records); +}; diff --git a/server/src/vector/client.ts b/server/src/vector/client.ts index 59738a1cf..267a5e308 100644 --- a/server/src/vector/client.ts +++ b/server/src/vector/client.ts @@ -19,13 +19,15 @@ class VectorClient { private readonly apiKey: string; private readonly indexName: string; private readonly accountId: string; + private readonly VECTORIZE_INDEX_BASE_URL: string; private readonly VECTORIZE_INDEX_URL: string; private constructor(apiKey: string, indexName: string, accountId: string) { this.apiKey = apiKey; this.indexName = indexName; this.accountId = accountId; - this.VECTORIZE_INDEX_URL = `https://api.cloudflare.com/client/v4/accounts/${this.accountId}/vectorize/v2/indexes/${this.indexName}`; + this.VECTORIZE_INDEX_BASE_URL = `https://api.cloudflare.com/client/v4/accounts/${this.accountId}/vectorize/v2/indexes`; + this.VECTORIZE_INDEX_URL = `${this.VECTORIZE_INDEX_BASE_URL}/${this.indexName}`; } public static get instance(): VectorClient { @@ -107,6 +109,9 @@ class VectorClient { metadata: Metadata; }>, ) { + if (vectors.length === 0) { + throw new Error('No vectors provided for upsert.'); + } let ndjsonBody = ''; for (const vector of vectors) { ndjsonBody += `${JSON.stringify(vector)}\n`; @@ -133,6 +138,80 @@ class VectorClient { return await response.json(); } + /** + * Creates a new vector database based on the environment index name. + * @returns A promise that resolves with the result of the vector database creation. + */ + public async createVectorDB() { + const response = await fetch(this.VECTORIZE_INDEX_BASE_URL, { + method: 'POST', + headers: { + 'Content-Type': 'application/json', + Authorization: `Bearer ${this.apiKey}`, + }, + body: JSON.stringify({ + name: this.indexName, + description: 'Vector database for packrat', + config: { preset: AiClient.modelName }, + }), + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `Failed to create vector: ${response.statusText} - ${errorText}`, + ); + } + + const responseData: { + success: boolean; + result: Record; + errors: Array<{ code: number; message: string }>; + messages: Array<{ code: number; message: string }>; + } = await response.json(); + if (!responseData.success) { + throw new Error( + `Failed to create vector: ${JSON.stringify({ errors: responseData.errors })}`, + ); + } + + return { result: responseData.result, messages: responseData.messages }; + } + + /** + * Deletes the vector database. + * @returns A promise that resolves when the vector database is deleted. + */ + public async deleteVectorDB(): Promise> { + const response = await fetch(this.VECTORIZE_INDEX_URL, { + method: 'DELETE', + headers: { + Authorization: `Bearer ${this.apiKey}`, + }, + }); + + if (!response.ok) { + const errorText = await response.text(); + throw new Error( + `Failed to delete vector: ${response.statusText} - ${errorText}`, + ); + } + + const responseData: { + success: boolean; + result: Record; + errors: Array<{ code: number; message: string }>; + messages: Array<{ code: number; message: string }>; + } = await response.json(); + if (!responseData.success) { + throw new Error( + `Failed to delete vector: ${JSON.stringify({ errors: responseData.errors })}`, + ); + } + + return responseData.result; + } + public async delete(id: string) { const url = `${this.VECTORIZE_INDEX_URL}/delete-by-ids`;