MCP Terminal Server
by dillip285
/**
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { GenkitError, z } from '@genkit-ai/core';
import { Registry } from '@genkit-ai/core/registry';
import * as assert from 'assert';
import { beforeEach, describe, it } from 'node:test';
import { defineReranker, rerank } from '../../src/reranker';
import { Document } from '../../src/retriever';
describe('reranker', () => {
describe('defineReranker()', () => {
let registry: Registry;
beforeEach(() => {
registry = new Registry();
});
it('reranks documents based on custom logic', async () => {
const customReranker = defineReranker(
registry,
{
name: 'reranker',
configSchema: z.object({
k: z.number().optional(),
}),
},
async (query, documents, options) => {
// Custom reranking logic: score based on string length similarity to query
const queryLength = query.text.length;
const rerankedDocs = documents.map((doc) => {
const score = Math.abs(queryLength - doc.text.length);
return {
...doc,
metadata: { ...doc.metadata, score },
};
});
return {
documents: rerankedDocs
.sort((a, b) => a.metadata.score - b.metadata.score)
.slice(0, options.k || 3),
};
}
);
// Sample documents for testing
const documents = [
Document.fromText('short'),
Document.fromText('a bit longer'),
Document.fromText('this is a very long document'),
];
const query = Document.fromText('medium length');
const rerankedDocuments = await rerank(registry, {
reranker: customReranker,
query,
documents,
options: { k: 2 },
});
// Validate the reranked results
assert.equal(rerankedDocuments.length, 2);
assert.ok(rerankedDocuments[0].text.includes('a bit longer'));
assert.ok(rerankedDocuments[1].text.includes('short'));
});
it('handles missing options gracefully', async () => {
const customReranker = defineReranker(
registry,
{
name: 'reranker',
configSchema: z.object({
k: z.number().optional(),
}),
},
async (query, documents, options) => {
const rerankedDocs = documents.map((doc) => {
const score = Math.random(); // Simplified scoring for testing
return {
...doc,
metadata: { ...doc.metadata, score },
};
});
return {
documents: rerankedDocs.sort(
(a, b) => b.metadata.score - a.metadata.score
),
};
}
);
const documents = [Document.fromText('doc1'), Document.fromText('doc2')];
const query = Document.fromText('test query');
const rerankedDocuments = await rerank(registry, {
reranker: customReranker,
query,
documents,
options: { k: 2 },
});
assert.equal(rerankedDocuments.length, 2);
assert.ok(typeof rerankedDocuments[0].metadata.score === 'number');
});
it('validates config schema and throws error on invalid input', async () => {
const customReranker = defineReranker(
registry,
{
name: 'reranker',
configSchema: z.object({
k: z.number().min(1),
}),
},
async (query, documents, options) => {
// Simplified scoring for testing
const rerankedDocs = documents.map((doc) => ({
...doc,
metadata: { score: Math.random() },
}));
return {
documents: rerankedDocs.sort(
(a, b) => b.metadata.score - a.metadata.score
),
};
}
);
const documents = [Document.fromText('doc1')];
const query = Document.fromText('test query');
try {
await rerank(registry, {
reranker: customReranker,
query,
documents,
options: { k: 0 }, // Invalid input: k must be at least 1
});
assert.fail('Expected validation error');
} catch (err) {
assert.ok(err instanceof GenkitError);
assert.equal(err.status, 'INVALID_ARGUMENT');
}
});
it('preserves document metadata after reranking', async () => {
const customReranker = defineReranker(
registry,
{
name: 'reranker',
},
async (query, documents) => {
const rerankedDocs = documents.map((doc, i) => ({
...doc,
metadata: { ...doc.metadata, score: 2 - i },
}));
return {
documents: rerankedDocs.sort(
(a, b) => b.metadata.score - a.metadata.score
),
};
}
);
const documents = [
new Document({ content: [], metadata: { originalField: 'test1' } }),
new Document({ content: [], metadata: { originalField: 'test2' } }),
];
const query = Document.fromText('test query');
const rerankedDocuments = await rerank(registry, {
reranker: customReranker,
query,
documents,
});
assert.equal(rerankedDocuments[0].metadata.originalField, 'test1');
assert.equal(rerankedDocuments[1].metadata.originalField, 'test2');
});
it('handles errors thrown by the reranker', async () => {
const customReranker = defineReranker(
registry,
{
name: 'reranker',
},
async (query, documents) => {
// Simulate an error in the reranker logic
throw new GenkitError({
message: 'Something went wrong during reranking',
status: 'INTERNAL',
});
}
);
const documents = [Document.fromText('doc1'), Document.fromText('doc2')];
const query = Document.fromText('test query');
try {
await rerank(registry, {
reranker: customReranker,
query,
documents,
});
assert.fail('Expected an error to be thrown');
} catch (err) {
assert.ok(err instanceof GenkitError);
assert.equal(err.status, 'INTERNAL');
assert.equal(
err.message,
'INTERNAL: Something went wrong during reranking'
);
}
});
});
});