feat: adds support for pinning rules

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<PinningRule[]> {
+    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<string[]> {
+    return this.client.request<string[]>({
+      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<void> {
@@ -799,7 +869,7 @@ export class Index {
     })
   }
 
-  public async insertDocuments(documents: AnyObject | AnyObject[], init?: ClientRequestInit): Promise<void> {
+  public async insertDocuments<T = AnyObject | AnyObject[]>(documents: T, init?: ClientRequestInit): Promise<void> {
     await this.oramaInterface.request<void>({
       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<void> {
+  public async upsertDocuments<T = AnyObject[]>(documents: T, init?: ClientRequestInit): Promise<void> {
     await this.oramaInterface.request<void>({
       path: `/v1/collections/${this.collectionID}/indexes/${this.indexID}/upsert`,
-      body: documents,
+      body: documents as AnyObject[],
       method: 'POST',
       init,
       apiKeyPosition: 'header',

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<PinningRule, 'id'> & {
+  id?: string
+}

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)
+})