Skip to content

Commit

Permalink
test(llm): add ollama e2e test
Browse files Browse the repository at this point in the history
Signed-off-by: Tomas Dvorak <[email protected]>
  • Loading branch information
Tomas2D committed Oct 7, 2024
1 parent f0995ea commit e4ec120
Showing 1 changed file with 54 additions and 0 deletions.
54 changes: 54 additions & 0 deletions tests/e2e/adapters/ollama/chat.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,54 @@
/**
* Copyright 2024 IBM Corp.
*
* 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 { BaseMessage, Role } from "@/llms/primitives/message.js";
import { OllamaChatLLM } from "@/adapters/ollama/chat.js";
import { Ollama } from "ollama";

const host = process.env.OLLAMA_HOST;

describe.runIf(Boolean(host))("Ollama Chat LLM", () => {
const createChatLLM = () => {
return new OllamaChatLLM({
modelId: "llama3.1",
parameters: {
temperature: 0,
num_predict: 5,
},
client: new Ollama({
host,
}),
});
};

it("Generates", async () => {
const conversation = [
BaseMessage.of({
role: Role.SYSTEM,
text: `You are a helpful and respectful and honest assistant. Your name is Bee.`,
}),
];
const llm = createChatLLM();
const response = await llm.generate([
...conversation,
BaseMessage.of({
role: "user",
text: "What is your name? Just output the name without any additional comment.",
}),
]);
expect(response.getTextContent()).includes("Bee");
});
});

0 comments on commit e4ec120

Please sign in to comment.