Audrey/add recency decay

src/inference/__tests__/inference.test.ts

@@ -214,3 +214,67 @@ describe('Inference Class: rerank', () => {
     }
   });
 });
+
+describe('Inference Class: rerank, additive recency decay', () => {
+  let inference: Inference;
+
+  beforeEach(() => {
+    const config: PineconeConfiguration = { apiKey: 'test-api-key' };
+    const infApi = inferenceOperationsBuilder(config);
+    inference = new Inference(infApi);
+  });
+
+  test('Additive recency decay, confirm error is thrown if timestamp is invalid', async () => {
+    const invalidTimestamp = 'invalid-timestamp';
+    const rerankOriginalResponse = {
+      model: 'bge-reranker-v2-m3',
+      data: [
+        {
+          index: 0,
+          score: 0.09586197932098765,
+          document: {
+            text: "This is today's document",
+            timestamp: invalidTimestamp, // Should throw error
+          },
+        },
+      ],
+      usage: { rerankUnits: 1 },
+    };
+
+    try {
+      inference.addAdditiveDecay(rerankOriginalResponse, { decay: true });
+    } catch (error) {
+      expect(error).toEqual(
+        new Error(
+          `Invalid date format: ${rerankOriginalResponse.data[0].document['timestamp']}`
+        )
+      );
+    }
+  });
+
+  test('Additive recency decay, confirm error is thrown if doc does not have `timestamp` filed', async () => {
+    const rerankOriginalResponse = {
+      model: 'bge-reranker-v2-m3',
+      data: [
+        {
+          index: 0,
+          score: 0.09586197932098765,
+          document: {
+            text: "This is today's document",
+          },
+        },
+      ],
+      usage: { rerankUnits: 1 },
+    };
+
+    try {
+      inference.addAdditiveDecay(rerankOriginalResponse, { decay: true });
+    } catch (error) {
+      expect(error).toEqual(
+        new Error(
+          `Document ${rerankOriginalResponse.data[0].index} does not have a \`timestamp\` field`
+        )
+      );
+    }
+  });
+});

src/inference/inference.ts

@@ -13,6 +13,9 @@
   topN?: number;
   returnDocuments?: boolean;
   rankFields?: Array<string>;
+  decay?: boolean;
+  decayThreshold?: number;
+  decayWeight?: number;
   parameters?: { [key: string]: string };
 }
 
@@ -58,6 +61,86 @@
     return new EmbeddingsList(response.model, response.data, response.usage);
   }
 
+  /* Add an additive recency decay to a ranked results list from the /embed endpoint.
+   *
+   * Additive decay means we *add* a decay factor to the original score (vs multiplying it, or taking the log of it,
+   *  etc.).
+   *
+   * The factors that contribute to the final score are:
+   * - Document timestamp: The timestamp of the document, provided in string form, including ms (e.g. "2013-09-05
+   *  15:34:00").
+   * - Today's date: Today's timestamp.
+   * - decayThreshold (default 30 days): Time period (in days) after which the decay starts significantly affecting.
+   *  If a document is within the threshold, the decay will scale based on how old the document is. If it is older
+   *  than the threshold, the document is treated as fully decayed (normalized decay of 1).
+   *   - Increasing this value:
+   *     - Effect: Recency decay is more gradual; documents remain relevant for a longer time.
+   *     - Use case: When freshness/recency is _less_ important (e.g. product reviews)
+   *   - Decreasing this value:
+   *     - Effect: Recency decay is more abrupt; documents lose relevance faster.
+   *     - Use case: When freshness/recency is _more_ important (e.g. news articles).
+   * - decayWeight (default 0.5): The magnitude of the decay's impact on document scores.
+   *   - Increasing this value:
+   *     - Effect: Decay has a stronger impact on document scores; older docs are heavily penalized.
+   *     - Use case: You want to more strongly prioritize recency.
+   *   - Decreasing this value:
+   *     - Effect: Decay has a weaker impact on document scores; older documents have a better chance at
+   *  retaining their original score/ranking.
+   *     - Use case: You want to prioritize recency less.
+   *
+   * @param response - The original response object from the /embed endpoint.
+   * @param options - The original options object passed to the /embed endpoint.
+   * */
+  addAdditiveDecay(
+    response: RerankResult,
+    options: RerankOptions
+  ): RerankResult {
+    if (options.rankFields) {
+      options.rankFields = ['timestamp'];
+    }
+
+    const convertDaysToSeconds = (decayThreshold: number) => {
+      return decayThreshold * 24 * 60 * 60;
+    };
+
+    const {
+      decayThreshold = convertDaysToSeconds(30),
+      decayWeight = 0.5,
+      ...previousOptions
+    } = options;
+
+    for (const doc of response.data) {
+      if (doc.document && doc.document['timestamp']) {
+        // Convert timestamp (e.g. "2013-09-05 15:34:00") to milliseconds
+        const timestamp = new Date(doc.document['timestamp']).getTime();
+        console.log('timestamp', timestamp);
+        if (isNaN(timestamp)) {
+          throw new Error(`Invalid date format: ${doc.document['timestamp']}`);
+        }
+
+        const now = new Date().getTime(); // Calculate current time (ms)
+
+        // Calculate time decay in seconds (more manageable than ms)
+        const decay = (now - timestamp) / 1000;
+
+        // Normalize decay by n-days (s); docs > threshold have same decay so ancient docs don't get penalized more
+        const normalizedDecay = Math.min(decay / decayThreshold, 1); // Cap at 1 for documents > decayThreshold
+
+        // Apply decay to the original score, scaling new score to a manageable range; ^ decayWeight, ^ impact decay
+        // has on ranking
+        doc.score = doc.score - normalizedDecay * decayWeight; // Additive part is here
+      } else {
+        throw new Error(
+          `Document ${doc.index} does not have a \`timestamp\` field`
+        );
+      }
+    }
+    // Reorder response.data according to new scores
+    response.data.sort((a, b) => b.score - a.score);
+
+    return response;
+  }
+
   /** Rerank documents against a query with a reranking model. Each document is ranked in descending relevance order
    *  against the query provided.
    *
@@ -189,6 +272,12 @@
       },
     };
 
-    return await this._inferenceApi.rerank(req);
+    const response = await this._inferenceApi.rerank(req);
+
+    if (options.decay) {
+      return this.addAdditiveDecay(response, options);
+    }
+
+    return response;
   }
 }

src/integration/inference/rerank.test.ts

@@ -4,6 +4,8 @@ describe('Integration Test: Pinecone Inference API rerank endpoint', () => {
   let model: string;
   let query: string;
   let documents: Array<string>;
+  let decayDocs: Array<{ text: string; timestamp: string }>;
+  let decayQuery: string;
   let pinecone: Pinecone;
 
   beforeAll(() => {
@@ -12,6 +14,23 @@ describe('Integration Test: Pinecone Inference API rerank endpoint', () => {
       'document content 1 yay I am about turkey',
       'document content 2',
     ];
+    decayDocs = [
+      { text: "This is today's doc", timestamp: '2024-09-30 13:02:00' },
+      {
+        text: "This is today's doc, but at an earlier timestamp",
+        timestamp: '2024-09-30 11:02:00',
+      },
+      {
+        text: "This is yesterday's document",
+        timestamp: '2024-09-29 11:02:00',
+      },
+      {
+        text: "Now, this is two days' ago's document",
+        timestamp: '2024-09-28 11:02:00',
+      },
+      { text: 'This is a super old doc', timestamp: '2020-01-01 11:02:00' },
+    ];
+    decayQuery = 'what was yesterday like?';
     model = 'bge-reranker-v2-m3';
     const apiKey = process.env.PINECONE_API_KEY || '';
     pinecone = new Pinecone({ apiKey });
@@ -34,4 +53,135 @@ describe('Integration Test: Pinecone Inference API rerank endpoint', () => {
     // (Just ignoring the fact that technically doc.document['text'] could be undefined)
     expect(response.data.map((doc) => doc.document['text'])).toBeDefined();
   });
+
+  test('Additive recency decay, confirm default expectations', async () => {
+    const rerankOptions = {
+      decay: true,
+    };
+    const nonDecayedResponse = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs
+    );
+    const rankOrderOfNonDecayedResponse = nonDecayedResponse.data.map(
+      (doc) => doc.index
+    );
+
+    const decayedResponse = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs,
+      rerankOptions
+    );
+    const rankOrderOfDecayedResponse = decayedResponse.data.map(
+      (doc) => doc.index
+    );
+
+    if (
+      nonDecayedResponse.data[0].document &&
+      decayedResponse.data[0].document
+    ) {
+      // Non-decayed response should put doc whose text is "This is yesterday's document" in position 0, since query is
+      // "what was yesterday like?"
+      expect(nonDecayedResponse.data[0].document.text).toContain(
+        decayDocs[2].text
+      );
+      // Decayed response should put doc whose text is "This is today's doc" in position 0, since it's the most
+      // recent doc, and we want to rank by recency first and foremost
+      expect(decayedResponse.data[0].document.text).toContain(
+        decayDocs[0].text
+      );
+    }
+
+    // General assertion that the rank order of the non-decayed response is different from the rank order of the decayed
+    expect(rankOrderOfNonDecayedResponse).not.toEqual(
+      rankOrderOfDecayedResponse
+    );
+  });
+
+  test('Additive recency decay, super small decayWeight', async () => {
+    // This test should yield a rank order extremely similar to the rank order of a non-decayed response
+    const rerankOptionsSmallDecay = {
+      decay: true,
+      decayWeight: 0.000000001,
+    };
+    const responseNoDecay = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs
+    );
+    const responseSmallDecay = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs,
+      rerankOptionsSmallDecay
+    );
+    const responseDefaultDecay = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs,
+      { decay: true }
+    );
+
+    expect(responseNoDecay.data.map((doc) => doc.index)).toEqual(
+      responseSmallDecay.data.map((doc) => doc.index)
+    );
+    expect(responseSmallDecay.data.map((doc) => doc.index)).not.toEqual(
+      responseDefaultDecay.data.map((doc) => doc.index)
+    );
+  });
+
+  test('Additive recency decay, super small decayThreshold', async () => {
+    // With a super small decayThreshold, we expect only docs whose timestamps fall into the threshold of now-0.25
+    // days to get a decay added to them; since no docs in `decayDocs` meet this expectation, we expect the rank
+    // order to be the same as if no decay were applied
+    const rerankOptionsSmallThreshold = {
+      decay: true,
+      decayThreshold: 0.25, // 1/4 day
+    };
+    const responseNoDecay = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs
+    );
+    const responseHighThreshold = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs,
+      rerankOptionsSmallThreshold
+    );
+
+    expect(responseNoDecay.data.map((doc) => doc.index)).toEqual(
+      responseHighThreshold.data.map((doc) => doc.index)
+    );
+  });
+
+  test('Additive recency decay, large decayThreshold', async () => {
+    // With a large decayThreshold, we expect all docs to get a decay added to them, and this decay will be gradual.
+    // Since decay is applied uniformally across all docs, their rank order is the same as if no decay had been applied,
+    // but their scores will be lower.
+    const rerankOptionsLargeThreshold = {
+      decay: true,
+      decayThreshold: 300, // 1 year
+    };
+    const responseNoDecay = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs
+    );
+    const responseLargeThreshold = await pinecone.inference.rerank(
+      model,
+      decayQuery,
+      decayDocs,
+      rerankOptionsLargeThreshold
+    );
+
+    let noDecaySum: number = 0;
+    responseNoDecay.data.forEach((a) => (noDecaySum += a.score));
+
+    let largeThresholdSum: number = 0;
+    responseLargeThreshold.data.forEach((a) => (largeThresholdSum += a.score));
+
+    expect(noDecaySum).toBeGreaterThan(largeThresholdSum);
+  });
 });