Genkit MCP
Official
by firebase
/**
* 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 {
ChatCompletionRequest,
ChatCompletionResponse,
CompletionChunk,
} from '@mistralai/mistralai-gcp/models/components';
import * as assert from 'assert';
import { GenerateRequest, GenerateResponseData } from 'genkit';
import { describe, it } from 'node:test';
import {
MistralConfigSchema,
fromMistralCompletionChunk,
fromMistralResponse,
toMistralRequest,
} from '../../src/modelgarden/mistral';
const MODEL_ID = 'mistral-large-2411';
describe('toMistralRequest', () => {
const testCases: {
should: string;
input: GenerateRequest<typeof MistralConfigSchema>;
expectedOutput: ChatCompletionRequest;
}[] = [
{
should: 'should transform genkit message (text content) correctly',
input: {
messages: [
{
role: 'user' as const,
content: [{ text: 'Tell a joke about dogs.' }],
},
],
},
expectedOutput: {
model: MODEL_ID,
messages: [
{
role: 'user',
content: 'Tell a joke about dogs.',
},
],
maxTokens: 1024,
temperature: 0.7,
},
},
{
should: 'should transform system message',
input: {
messages: [
{
role: 'system' as const,
content: [{ text: 'Talk like a pirate.' }],
},
{
role: 'user' as const,
content: [{ text: 'Tell a joke about dogs.' }],
},
],
},
expectedOutput: {
model: MODEL_ID,
messages: [
{
role: 'system',
content: 'Talk like a pirate.',
},
{
role: 'user',
content: 'Tell a joke about dogs.',
},
],
maxTokens: 1024,
temperature: 0.7,
},
},
{
should: 'should transform tool request correctly',
input: {
messages: [
{
role: 'user' as const,
content: [{ text: "What's the weather like today?" }],
},
],
tools: [
{
name: 'get_weather',
description: 'Get the weather for a location.',
inputSchema: {
type: 'object',
properties: {
location: {
type: 'string',
description: 'The city and state, e.g. San Francisco, CA',
},
},
required: ['location'],
},
},
],
},
expectedOutput: {
model: MODEL_ID,
messages: [
{
role: 'user',
content: "What's the weather like today?",
},
],
maxTokens: 1024,
temperature: 0.7,
tools: [
{
type: 'function',
function: {
name: 'get_weather',
description: 'Get the weather for a location.',
parameters: {
type: 'object',
properties: {
location: {
type: 'string',
description: 'The city and state, e.g. San Francisco, CA',
},
},
required: ['location'],
},
},
},
],
},
},
];
for (const test of testCases) {
it(test.should, () => {
assert.deepEqual(
toMistralRequest(MODEL_ID, test.input),
test.expectedOutput
);
});
}
});
describe('fromMistralResponse', () => {
const testCases: {
should: string;
input: GenerateRequest<typeof MistralConfigSchema>;
response: ChatCompletionResponse;
expectedOutput: GenerateResponseData;
}[] = [
{
should: 'should transform mistral message (text content) correctly',
input: {
messages: [
{
role: 'user' as const,
content: [{ text: 'Tell a joke about dogs.' }],
},
],
},
response: {
id: 'abcd1234',
object: 'chat.completion',
model: MODEL_ID,
choices: [
{
index: 0,
message: {
role: 'assistant',
content:
'Why do dogs make terrible comedians? Their jokes are too ruff!',
},
finishReason: 'stop',
},
],
usage: {
promptTokens: 123,
completionTokens: 234,
totalTokens: 357,
},
},
expectedOutput: {
message: {
role: 'model' as const,
content: [
{
text: 'Why do dogs make terrible comedians? Their jokes are too ruff!',
},
],
},
finishReason: 'stop',
usage: {
inputTokens: 123,
outputTokens: 234,
},
custom: {
id: 'abcd1234',
model: MODEL_ID,
created: undefined,
},
raw: {
id: 'abcd1234',
object: 'chat.completion',
model: MODEL_ID,
choices: [
{
index: 0,
message: {
role: 'assistant',
content:
'Why do dogs make terrible comedians? Their jokes are too ruff!',
},
finishReason: 'stop',
},
],
usage: {
promptTokens: 123,
completionTokens: 234,
totalTokens: 357,
},
},
},
},
{
should: 'should transform tool calls correctly',
input: {
messages: [
{
role: 'user' as const,
content: [{ text: "What's the weather like today?" }],
},
],
},
response: {
id: 'abcd1234',
object: 'chat.completion',
model: MODEL_ID,
choices: [
{
index: 0,
message: {
role: 'assistant',
content: null,
toolCalls: [
{
id: 'call_abc123',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"location":"San Francisco, CA"}',
},
},
],
},
finishReason: 'tool_calls',
},
],
usage: {
promptTokens: 123,
completionTokens: 234,
totalTokens: 357,
},
},
expectedOutput: {
message: {
role: 'model' as const,
content: [
{
toolRequest: {
ref: 'call_abc123',
name: 'get_weather',
input: { location: 'San Francisco, CA' },
},
},
],
},
finishReason: 'stop',
usage: {
inputTokens: 123,
outputTokens: 234,
},
custom: {
id: 'abcd1234',
model: MODEL_ID,
created: undefined,
},
raw: {
id: 'abcd1234',
object: 'chat.completion',
model: MODEL_ID,
choices: [
{
index: 0,
message: {
role: 'assistant',
content: null,
toolCalls: [
{
id: 'call_abc123',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"location":"San Francisco, CA"}',
},
},
],
},
finishReason: 'tool_calls',
},
],
usage: {
promptTokens: 123,
completionTokens: 234,
totalTokens: 357,
},
},
},
},
];
for (const test of testCases) {
it(test.should, () => {
assert.deepEqual(
fromMistralResponse(test.input, test.response),
test.expectedOutput
);
});
}
});
describe('validateToolSequence', () => {
it('should handle valid tool call and response sequence', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: "What's the weather?" }],
},
{
role: 'model' as const,
content: [
{
toolRequest: {
ref: 'call_123',
name: 'get_weather',
input: { location: 'San Francisco' },
},
},
],
},
{
role: 'tool' as const,
content: [
{
toolResponse: {
ref: 'call_123',
name: 'get_weather',
output: { temperature: 72, condition: 'sunny' },
},
},
],
},
],
};
// Should not throw an error
assert.doesNotThrow(() => toMistralRequest(MODEL_ID, input));
});
it('should throw error when tool response is missing', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: "What's the weather?" }],
},
{
role: 'model' as const,
content: [
{
toolRequest: {
ref: 'call_123',
name: 'get_weather',
input: { location: 'San Francisco' },
},
},
],
},
],
};
assert.throws(
() => toMistralRequest(MODEL_ID, input),
/Mismatch between tool calls/
);
});
it('should throw error when tool response id does not match call', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: "What's the weather?" }],
},
{
role: 'model' as const,
content: [
{
toolRequest: {
ref: 'call_123',
name: 'get_weather',
input: { location: 'San Francisco' },
},
},
],
},
{
role: 'tool' as const,
content: [
{
toolResponse: {
ref: 'wrong_id',
name: 'get_weather',
output: { temperature: 72, condition: 'sunny' },
},
},
],
},
],
};
assert.throws(
() => toMistralRequest(MODEL_ID, input),
/Tool response with ID wrong_id has no matching call/
);
});
});
describe('edge cases', () => {
it('should handle empty message content', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [],
},
],
};
const result = toMistralRequest(MODEL_ID, input);
assert.equal(result.messages[0].content, '');
});
it('should handle multiple text parts in content', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: 'Hello' }, { text: ' ' }, { text: 'World' }],
},
],
};
const result = toMistralRequest(MODEL_ID, input);
assert.equal(result.messages[0].content, 'Hello World');
});
it('should handle custom configuration options', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: 'test' }],
},
],
config: {
temperature: 0.5,
maxOutputTokens: 500,
topP: 0.9,
stopSequences: ['END'],
},
};
const result = toMistralRequest(MODEL_ID, input);
assert.equal(result.temperature, 0.5);
assert.equal(result.maxTokens, 500);
assert.equal(result.topP, 0.9);
assert.deepEqual(result.stop, ['END']);
});
});
describe('fromMistralResponse error handling', () => {
it('should handle response with no choices', () => {
const input = {
messages: [
{
role: 'user' as const,
content: [{ text: 'test' }],
},
],
};
const response = {
id: 'test',
object: 'chat.completion',
model: MODEL_ID,
choices: [],
usage: {
promptTokens: 0,
completionTokens: 0,
totalTokens: 0,
},
};
const result = fromMistralResponse(input, response);
assert.deepEqual(result.message!.content, []);
});
});
describe('fromMistralCompletionChunk', () => {
it('should handle text content chunk', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [
{
index: 0,
delta: {
role: 'assistant',
content: 'Hello world',
},
},
],
};
const parts = fromMistralCompletionChunk(
chunk as unknown as CompletionChunk
);
assert.deepEqual(parts, [{ text: 'Hello world' }]);
});
it('should handle tool call chunk', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [
{
index: 0,
delta: {
role: 'assistant',
toolCalls: [
{
id: 'call_123',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"location":"San Francisco"}',
},
},
],
},
},
],
};
const parts = fromMistralCompletionChunk(
chunk as unknown as CompletionChunk
);
assert.deepEqual(parts, [
{
toolRequest: {
ref: 'call_123',
name: 'get_weather',
input: { location: 'San Francisco' },
},
},
]);
});
it('should handle empty chunk', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [],
};
const parts = fromMistralCompletionChunk(chunk);
assert.deepEqual(parts, []);
});
it('should handle chunk with empty delta', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [
{
index: 0,
delta: {},
},
],
};
const parts = fromMistralCompletionChunk(
chunk as unknown as CompletionChunk
);
assert.deepEqual(parts, []);
});
it('should handle chunk with invalid tool call', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [
{
index: 0,
delta: {
toolCalls: [
{
id: 'call_123',
type: 'function',
// Missing function property
},
],
},
},
],
};
const parts = fromMistralCompletionChunk(
chunk as unknown as CompletionChunk
);
assert.deepEqual(parts, []); // Should skip invalid tool call
});
it('should handle mixed content chunk', () => {
const chunk = {
id: 'chunk_1',
model: MODEL_ID,
choices: [
{
index: 0,
delta: {
content: 'Some text',
toolCalls: [
{
id: 'call_123',
type: 'function',
function: {
name: 'get_weather',
arguments: '{"location":"San Francisco"}',
},
},
],
},
},
],
};
const parts = fromMistralCompletionChunk(
chunk as unknown as CompletionChunk
);
assert.deepEqual(parts, [
{ text: 'Some text' },
{
toolRequest: {
ref: 'call_123',
name: 'get_weather',
input: { location: 'San Francisco' },
},
},
]);
});
});