diff --git a/src/collection.ts b/src/collection.ts index 39fe78f..6584e3d 100644 --- a/src/collection.ts +++ b/src/collection.ts @@ -7,6 +7,8 @@ import type { NLPSearchStreamResult, NLPSearchStreamStatus, Nullable, + PinningRule, + PinningRuleInsertObject, SearchParams, SearchResult, TrainingSetInsertParameters, @@ -431,6 +433,72 @@ class HooksNamespace { } } +class PinningRulesNamespace { + private client: Client + private collectionID: string + private indexID: string + + constructor(client: Client, collectionID: string, indexID: string) { + this.client = client + this.collectionID = collectionID + this.indexID = indexID + } + + public insert(rule: PinningRuleInsertObject): Promise<{ success: boolean }> { + if (!rule.id) { + rule.id = createRandomString(32) + } + + return this.client.request<{ success: true }>({ + path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/pin_rules/insert`, + body: rule, + method: 'POST', + apiKeyPosition: 'header', + target: 'writer', + }) + } + + public update(rule: PinningRuleInsertObject): Promise<{ success: boolean }> { + if (!rule.id) { + rule.id = createRandomString(32) + } + + return this.insert(rule) + } + + public async list(): Promise { + const results = await this.client.request<{ data: PinningRule[] }>({ + path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/pin_rules/list`, + method: 'GET', + apiKeyPosition: 'header', + target: 'writer', + }) + + return results.data + } + + public listIDs(): Promise { + return this.client.request({ + path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/pin_rules/ids`, + method: 'GET', + apiKeyPosition: 'query-params', + target: 'reader', + }) + } + + public delete(id: string): Promise<{ success: boolean }> { + return this.client.request<{ success: true }>({ + path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/pin_rules/delete`, + method: 'POST', + body: { + pin_rule_id_to_delete: id, + }, + apiKeyPosition: 'header', + target: 'writer', + }) + } +} + class LogsNamespace { private client: Client private collectionID: string @@ -781,12 +849,14 @@ export class Index { private collectionID: string private oramaInterface: Client public transaction: Transaction + public pinningRules: PinningRulesNamespace constructor(oramaInterface: Client, collectionID: string, indexID: string) { this.indexID = indexID this.collectionID = collectionID this.oramaInterface = oramaInterface this.transaction = new Transaction(oramaInterface, collectionID, indexID) + this.pinningRules = new PinningRulesNamespace(oramaInterface, collectionID, indexID) } public async reindex(init?: ClientRequestInit): Promise { @@ -799,7 +869,7 @@ export class Index { }) } - public async insertDocuments(documents: AnyObject | AnyObject[], init?: ClientRequestInit): Promise { + public async insertDocuments(documents: T, init?: ClientRequestInit): Promise { await this.oramaInterface.request({ path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/insert`, body: Array.isArray(documents) ? documents : [documents], @@ -821,10 +891,10 @@ export class Index { }) } - public async upsertDocuments(documents: AnyObject[], init?: ClientRequestInit): Promise { + public async upsertDocuments(documents: T, init?: ClientRequestInit): Promise { await this.oramaInterface.request({ path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/upsert`, - body: documents, + body: documents as AnyObject[], method: 'POST', init, apiKeyPosition: 'header', diff --git a/src/lib/types.ts b/src/lib/types.ts index 70e3f27..b5d6c28 100644 --- a/src/lib/types.ts +++ b/src/lib/types.ts @@ -71,7 +71,7 @@ export type SearchParams = { facets?: AnyObject indexes?: string[] datasourceIDs?: string[] - boost: { [key: string]: number } + boost?: { [key: string]: number } exact?: boolean threshold?: number tolerance?: number @@ -306,3 +306,27 @@ export type TrainingSetQueryOptimizer = { } export type TrainingSetInsertParameters = TrainingSetQueryOptimizer['queries'] + +export type PinningRuleAnchoringType = 'is' + +export type PinningRuleCondition = { + anchoring: PinningRuleAnchoringType + pattern: string +} + +export type PinningRuleConsequencePromote = { + doc_id: string + position: number +} + +export type PinningRule = { + id: string + conditions: PinningRuleCondition[] + consequence: { + promote?: PinningRuleConsequencePromote[] + } +} + +export type PinningRuleInsertObject = Omit & { + id?: string +} diff --git a/tests/orama.collection.e2e.test.ts b/tests/orama.collection.e2e.test.ts index 55ea993..2f0bc0f 100644 --- a/tests/orama.collection.e2e.test.ts +++ b/tests/orama.collection.e2e.test.ts @@ -1,6 +1,8 @@ +import type { PinningRuleInsertObject } from '../src/lib/types.ts' + import { z } from 'npm:zod@3.24.3' import { assert, assertEquals, assertFalse, assertNotEquals } from 'jsr:@std/assert' -import { CollectionManager, OramaCloud, OramaCoreManager } from '../src/index.ts' +import { CollectionManager, OramaCoreManager } from '../src/index.ts' import { createRandomString } from '../src/lib/utils.ts' const manager = new OramaCoreManager({ @@ -288,7 +290,7 @@ export default { beforeAnswer }; assertEquals(hooksAfterAfter.BeforeRetrieval, null) }) -Deno.test('CollectionManager: stream logs', async () => { +Deno.test.ignore('CollectionManager: stream logs', async () => { await collectionManager.hooks.insert({ name: 'BeforeRetrieval', code: ` @@ -371,3 +373,58 @@ Deno.test('CollectionManager: can handle transaction', async () => { assertEquals(countAfter, 1) assertEquals(docsAfter.hits[0].document.id, '3') }) + +Deno.test('CollectionManager: can handle pinning rules', async () => { + const newIndexId = createRandomString(32) + + await collectionManager.index.create({ + id: newIndexId, + }) + + const index = collectionManager.index.set(newIndexId) + + await index.insertDocuments([ + { id: '1', name: 'Blue Jeans' }, + { id: '2', name: 'Red T-Shirt' }, + { id: '3', name: 'Green Hoodie' }, + { id: '4', name: 'Yellow Socks' }, + ]) + + const pinningRule: PinningRuleInsertObject = { + id: 'test_rule', + conditions: [ + { + anchoring: 'is', + pattern: 'Blue Jeans', + }, + ], + consequence: { + promote: [ + { + doc_id: '2', + position: 1, + }, + ], + }, + } + + await index.pinningRules.insert(pinningRule) + + const rules = await index.pinningRules.list() + assertEquals(rules.length, 1) + assertEquals(rules[0].id, 'test_rule') + + const result = await collectionManager.search({ + term: 'Blue Jeans', + indexes: [newIndexId], + }) + + assertEquals(result.hits.length, 2) + assertEquals(result.hits[0].document.id, '1') + assertEquals(result.hits[1].document.id, '2') + + await index.pinningRules.delete('test_rule') + + const newRules = await index.pinningRules.list() + assertEquals(newRules.length, 0) +})