Skip to content

Commit

Permalink
feat(viewer): add similarity search
Browse files Browse the repository at this point in the history
- Introduces a new submodule `utils/search` that provides the search
  features. So far, it provides a function `searchSimilarMumblings` that
  performs similarity search on mumblings.

issue codemonger-io#28
  • Loading branch information
kikuomax committed Oct 12, 2023
1 parent 69362bc commit e80be82
Showing 1 changed file with 50 additions and 0 deletions.
50 changes: 50 additions & 0 deletions cdk/viewer/src/utils/search.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
import * as util from 'node:util';
import { InvokeCommand, LambdaClient } from '@aws-sdk/client-lambda';

// Shared Lambda client.
let lambda: LambdaClient | null = null;

/** Mumbling in similarity search results. */
export interface SimilarMumbling {
/** ID (URL) of the mumbling fragment. */
id: string;
/** Approximate squared distance. */
distance: number;
}

/**
* Searches embeddings similar to a given query.
*
* @remarks
*
* You have to specify the name of the Lambda function that performs actual
* search to the environment variable `SEARCH_SIMILAR_MUMBLINGS_FUNCTION_NAME`.
*/
export async function searchSimilarMumblings(
embedding: [number],
): Promise<[SimilarMumbling]> {
if (process.env.SEARCH_SIMILAR_MUMBLINGS_FUNCTION_NAME == null) {
throw new Error('SEARCH_SIMILAR_MUMBLINGS_FUNCTION_NAME is not set');
}
if (lambda == null) {
lambda = new LambdaClient({});
}
const res = await lambda.send(new InvokeCommand({
FunctionName: process.env.SEARCH_SIMILAR_MUMBLINGS_FUNCTION_NAME,
InvocationType: 'RequestResponse',
Payload: JSON.stringify(embedding),
}));
if (res.StatusCode !== 200){
throw new Error(`similarity search function failed with ${res.StatusCode}`);
}
if (res.Payload == null) {
throw new Error('similarity search function returned nothing');
}
const data = new util.TextDecoder().decode(res.Payload);
const results = JSON.parse(data);
if (!Array.isArray(results)) {
throw new Error('similarity search function returned non-array');
}
// TODO: verify results
return results as [SimilarMumbling];
}

0 comments on commit e80be82

Please sign in to comment.