feat: adding tests for knowledge.ts

Author: ai16z-demirix

packages/core/src/tests/knowledge.test.ts

@@ -0,0 +1,103 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { AgentRuntime } from "../runtime";
+import { KnowledgeItem, Memory, UUID } from "../types";
+import knowledge from "../knowledge";
+
+// Mock dependencies
+vi.mock("../embedding", () => ({
+    embed: vi.fn().mockResolvedValue([0.1, 0.2, 0.3]),
+}));
+
+vi.mock("../generation", () => ({
+    splitChunks: vi.fn().mockResolvedValue(["chunk1", "chunk2"]),
+}));
+
+vi.mock("../uuid", () => ({
+    stringToUuid: vi.fn().mockImplementation((str) => str as UUID),
+}));
+
+describe("Knowledge", () => {
+    let mockRuntime: AgentRuntime;
+    const mockAgentId = "agent-123" as UUID;
+    const mockMemoryId = "memory-123" as UUID;
+
+    beforeEach(() => {
+        mockRuntime = {
+            agentId: mockAgentId,
+            knowledgeManager: {
+                searchMemoriesByEmbedding: vi.fn().mockResolvedValue([
+                    {
+                        content: {
+                            text: "Test fragment",
+                            source: mockMemoryId,
+                        },
+                        similarity: 0.9,
+                    },
+                ]),
+                createMemory: vi.fn().mockResolvedValue(undefined),
+            },
+            documentsManager: {
+                getMemoryById: vi.fn().mockResolvedValue({
+                    content: { text: "Test document" },
+                }),
+                createMemory: vi.fn().mockResolvedValue(undefined),
+            },
+        } as unknown as AgentRuntime;
+    });
+
+    describe("get", () => {
+        it("should retrieve knowledge based on message content", async () => {
+            const message: Memory = {
+                agentId: mockAgentId,
+                content: { text: "test query" },
+            } as Memory;
+
+            const result = await knowledge.get(mockRuntime, message);
+
+            expect(result).toEqual(["Test document"]);
+            expect(mockRuntime.knowledgeManager.searchMemoriesByEmbedding).toHaveBeenCalled();
+            expect(mockRuntime.documentsManager.getMemoryById).toHaveBeenCalledWith(mockMemoryId);
+        });
+    });
+
+    describe("set", () => {
+        it("should store knowledge item and its fragments", async () => {
+            const knowledgeItem: KnowledgeItem = {
+                id: "test-id" as UUID,
+                content: { text: "test content" },
+            };
+
+            await knowledge.set(mockRuntime, knowledgeItem);
+
+            expect(mockRuntime.documentsManager.createMemory).toHaveBeenCalled();
+            expect(mockRuntime.knowledgeManager.createMemory).toHaveBeenCalledTimes(2); // Once for each chunk
+        });
+    });
+
+    describe("preprocess", () => {
+        it("should clean markdown and code blocks from text", () => {
+            const input = `
+# Header
+\`\`\`
+code block
+\`\`\`
+\`inline code\`
+![image alt](image.jpg)
+[link text](http://example.com)
+<div>html</div>
+---
+`;
+            const expected = "header image alt link text html";
+            expect(knowledge.preprocess(input)).toBe(expected);
+        });
+
+        it("should handle empty input", () => {
+            expect(knowledge.preprocess("")).toBe("");
+        });
+
+        it("should convert text to lowercase", () => {
+            const plainText = "This is Just Plain TEXT.";
+            expect(knowledge.preprocess(plainText)).toBe("this is just plain text");
+        });
+    });
+});