From e4ec120df3b79818b35b8e3cdeed95b208d2e8e6 Mon Sep 17 00:00:00 2001
From: Tomas Dvorak <toomas2d@gmail.com>
Date: Mon, 7 Oct 2024 21:58:02 +0200
Subject: [PATCH] test(llm): add ollama e2e test

Signed-off-by: Tomas Dvorak <toomas2d@gmail.com>
---
 tests/e2e/adapters/ollama/chat.test.ts | 54 ++++++++++++++++++++++++++
 1 file changed, 54 insertions(+)
 create mode 100644 tests/e2e/adapters/ollama/chat.test.ts

diff --git a/tests/e2e/adapters/ollama/chat.test.ts b/tests/e2e/adapters/ollama/chat.test.ts
new file mode 100644
index 00000000..1ae73a94
--- /dev/null
+++ b/tests/e2e/adapters/ollama/chat.test.ts
@@ -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");
+  });
+});