549 lines
18 KiB
TypeScript
549 lines
18 KiB
TypeScript
/**
|
||
* Regression tests for graph_memory.ts response parsing (issue #4248).
|
||
*
|
||
* Exercises the three json_object call sites in MemoryGraph with a mocked LLM:
|
||
* 1. _retrieveNodesFromData → entity extraction
|
||
* 2. _establishNodesRelationsFromData → relation extraction
|
||
* 3. _getDeleteEntitiesFromSearchOutput → deletion identification
|
||
*
|
||
* Covers: malformed LLM responses, missing fields, bad JSON in toolCalls,
|
||
* string-only responses, empty tool calls, and prompt construction.
|
||
*
|
||
* See: https://github.com/mem0ai/mem0/issues/4248
|
||
*/
|
||
|
||
import { MemoryGraph } from "../src/memory/graph_memory";
|
||
import {
|
||
EXTRACT_RELATIONS_PROMPT,
|
||
getDeleteMessages,
|
||
} from "../src/graphs/utils";
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Mocks – we replace heavy dependencies so tests run without Neo4j / OpenAI
|
||
// ---------------------------------------------------------------------------
|
||
|
||
// Mock neo4j-driver: provides a fake Driver with a no-op session
|
||
jest.mock("neo4j-driver", () => ({
|
||
__esModule: true,
|
||
default: {
|
||
driver: jest.fn(() => ({
|
||
session: () => ({
|
||
run: jest.fn().mockResolvedValue({ records: [] }),
|
||
close: jest.fn(),
|
||
}),
|
||
})),
|
||
auth: { basic: jest.fn() },
|
||
},
|
||
}));
|
||
|
||
// Mock factory so constructor doesn't try to instantiate real LLMs / embedders
|
||
const mockGenerateResponse = jest.fn();
|
||
const mockGenerateChat = jest.fn();
|
||
const mockEmbed = jest.fn().mockResolvedValue([0.1, 0.2, 0.3]);
|
||
|
||
jest.mock("../src/utils/factory", () => ({
|
||
LLMFactory: {
|
||
create: jest.fn(() => ({
|
||
generateResponse: mockGenerateResponse,
|
||
generateChat: mockGenerateChat,
|
||
})),
|
||
},
|
||
EmbedderFactory: {
|
||
create: jest.fn(() => ({
|
||
embed: mockEmbed,
|
||
})),
|
||
},
|
||
}));
|
||
|
||
// Minimal config that satisfies the MemoryGraph constructor
|
||
function makeConfig(overrides: Record<string, any> = {}) {
|
||
return {
|
||
graphStore: {
|
||
config: {
|
||
url: "bolt://localhost:7687",
|
||
username: "neo4j",
|
||
password: "test",
|
||
},
|
||
...overrides,
|
||
},
|
||
embedder: { provider: "openai", config: {} },
|
||
llm: { provider: "openai", config: {} },
|
||
} as any;
|
||
}
|
||
|
||
// Helper to access private methods via `any` cast
|
||
function graph(overrides: Record<string, any> = {}): any {
|
||
return new MemoryGraph(makeConfig(overrides));
|
||
}
|
||
|
||
const FILTERS = { userId: "test-user" };
|
||
|
||
beforeEach(() => {
|
||
jest.clearAllMocks();
|
||
});
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// 1. _retrieveNodesFromData – entity extraction
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
describe("_retrieveNodesFromData", () => {
|
||
it("parses a well-formed extract_entities tool call", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "extract_entities",
|
||
arguments: JSON.stringify({
|
||
entities: [
|
||
{ entity: "Alice", entity_type: "person" },
|
||
{ entity: "Pizza", entity_type: "food" },
|
||
],
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData(
|
||
"Alice likes pizza",
|
||
FILTERS,
|
||
);
|
||
|
||
expect(result).toEqual({ alice: "person", pizza: "food" });
|
||
});
|
||
|
||
it("returns empty map when LLM returns a plain string", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce("I am a string, not an object");
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("returns empty map when toolCalls is undefined", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({});
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("returns empty map when toolCalls is an empty array", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("handles malformed JSON in tool call arguments gracefully", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{ name: "extract_entities", arguments: "NOT VALID JSON {{{" },
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
// Should not throw — the catch block in the source logs the error
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("handles missing entities array in arguments", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "extract_entities",
|
||
arguments: JSON.stringify({ wrong_key: [] }),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
// args.entities is undefined → for..of on undefined throws → caught
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("skips tool calls with unrelated names", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "some_other_tool",
|
||
arguments: JSON.stringify({
|
||
entities: [{ entity: "X", entity_type: "Y" }],
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
expect(result).toEqual({});
|
||
});
|
||
|
||
it("normalises entity names to lowercase with underscores", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "extract_entities",
|
||
arguments: JSON.stringify({
|
||
entities: [{ entity: "New York City", entity_type: "City Name" }],
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._retrieveNodesFromData("anything", FILTERS);
|
||
expect(result).toEqual({ new_york_city: "city_name" });
|
||
});
|
||
|
||
it("passes json_object response format and the correct system prompt", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
await mg._retrieveNodesFromData("test data", FILTERS);
|
||
|
||
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
|
||
expect(responseFormat).toEqual({ type: "json_object" });
|
||
|
||
const systemMsg = messages[0].content as string;
|
||
expect(systemMsg.toLowerCase()).toContain("json");
|
||
expect(systemMsg).toContain("test-user");
|
||
});
|
||
});
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// 2. _establishNodesRelationsFromData – relation extraction
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
describe("_establishNodesRelationsFromData", () => {
|
||
it("parses a well-formed establish_relationships tool call", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "establish_relationships",
|
||
arguments: JSON.stringify({
|
||
entities: [
|
||
{ source: "Alice", relationship: "likes", destination: "Pizza" },
|
||
],
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._establishNodesRelationsFromData(
|
||
"Alice likes pizza",
|
||
FILTERS,
|
||
{ alice: "person", pizza: "food" },
|
||
);
|
||
|
||
expect(result).toEqual([
|
||
{ source: "alice", relationship: "likes", destination: "pizza" },
|
||
]);
|
||
});
|
||
|
||
it("returns empty array when LLM returns a string", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce("just a string");
|
||
|
||
const mg = graph();
|
||
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("returns empty array when toolCalls is empty", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("returns empty array when entities key is missing from arguments", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "establish_relationships",
|
||
arguments: JSON.stringify({ not_entities: [] }),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._establishNodesRelationsFromData("x", FILTERS, {});
|
||
// args.entities is undefined → falls back to []
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("throws on malformed JSON in tool call arguments (no try/catch in source)", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [{ name: "establish_relationships", arguments: "<<BROKEN>>" }],
|
||
});
|
||
|
||
const mg = graph();
|
||
// _establishNodesRelationsFromData does JSON.parse without try/catch
|
||
await expect(
|
||
mg._establishNodesRelationsFromData("x", FILTERS, {}),
|
||
).rejects.toThrow();
|
||
});
|
||
|
||
it("appends JSON format suffix to system prompt (no custom prompt)", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
await mg._establishNodesRelationsFromData("data", FILTERS, { a: "b" });
|
||
|
||
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
|
||
expect(responseFormat).toEqual({ type: "json_object" });
|
||
|
||
const systemContent = messages[0].content as string;
|
||
expect(systemContent.toLowerCase()).toContain("json");
|
||
expect(systemContent).toContain("test-user");
|
||
expect(systemContent).not.toContain("USER_ID");
|
||
// CUSTOM_PROMPT placeholder stays when no custom prompt is configured
|
||
// (only replaced when config.graphStore.customPrompt is set)
|
||
});
|
||
|
||
it("appends JSON format suffix and custom prompt when configured", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph({ customPrompt: "Focus on food relationships only." });
|
||
await mg._establishNodesRelationsFromData("data", FILTERS, {});
|
||
|
||
const [messages] = mockGenerateResponse.mock.calls[0];
|
||
const systemContent = messages[0].content as string;
|
||
expect(systemContent.toLowerCase()).toContain("json");
|
||
expect(systemContent).toContain("Focus on food relationships only.");
|
||
});
|
||
});
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// 3. _getDeleteEntitiesFromSearchOutput – deletion identification
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
describe("_getDeleteEntitiesFromSearchOutput", () => {
|
||
const SEARCH_OUTPUT = [
|
||
{
|
||
source: "alice",
|
||
source_id: "1",
|
||
relationship: "likes",
|
||
relation_id: "r1",
|
||
destination: "pizza",
|
||
destination_id: "2",
|
||
similarity: 0.95,
|
||
},
|
||
];
|
||
|
||
it("parses a well-formed delete_graph_memory tool call", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "delete_graph_memory",
|
||
arguments: JSON.stringify({
|
||
source: "Alice",
|
||
relationship: "likes",
|
||
destination: "Pizza",
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
SEARCH_OUTPUT,
|
||
"Alice hates pizza",
|
||
FILTERS,
|
||
);
|
||
|
||
expect(result).toEqual([
|
||
{ source: "alice", relationship: "likes", destination: "pizza" },
|
||
]);
|
||
});
|
||
|
||
it("returns empty array when LLM returns a string", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce("string response");
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
SEARCH_OUTPUT,
|
||
"x",
|
||
FILTERS,
|
||
);
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("returns empty array when no tool calls are present", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
SEARCH_OUTPUT,
|
||
"x",
|
||
FILTERS,
|
||
);
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("skips non-delete_graph_memory tool calls", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "noop",
|
||
arguments: JSON.stringify({}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
SEARCH_OUTPUT,
|
||
"x",
|
||
FILTERS,
|
||
);
|
||
expect(result).toEqual([]);
|
||
});
|
||
|
||
it("collects multiple delete tool calls", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "delete_graph_memory",
|
||
arguments: JSON.stringify({
|
||
source: "A",
|
||
relationship: "r1",
|
||
destination: "B",
|
||
}),
|
||
},
|
||
{
|
||
name: "delete_graph_memory",
|
||
arguments: JSON.stringify({
|
||
source: "C",
|
||
relationship: "r2",
|
||
destination: "D",
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
SEARCH_OUTPUT,
|
||
"x",
|
||
FILTERS,
|
||
);
|
||
expect(result).toHaveLength(2);
|
||
expect(result[0].source).toBe("a");
|
||
expect(result[1].source).toBe("c");
|
||
});
|
||
|
||
it("passes json_object format and includes 'json' in system prompt", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
await mg._getDeleteEntitiesFromSearchOutput(SEARCH_OUTPUT, "data", FILTERS);
|
||
|
||
const [messages, responseFormat] = mockGenerateResponse.mock.calls[0];
|
||
expect(responseFormat).toEqual({ type: "json_object" });
|
||
|
||
const systemContent = messages[0].content as string;
|
||
expect(systemContent.toLowerCase()).toContain("json");
|
||
expect(systemContent).toContain("test-user");
|
||
expect(systemContent).not.toContain("USER_ID");
|
||
});
|
||
|
||
it("handles empty searchOutput array", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
|
||
const mg = graph();
|
||
const result = await mg._getDeleteEntitiesFromSearchOutput(
|
||
[],
|
||
"data",
|
||
FILTERS,
|
||
);
|
||
expect(result).toEqual([]);
|
||
});
|
||
});
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// 4. Prompt construction — JSON keyword present in every json_object site
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
describe("Prompt construction — all json_object sites include 'json'", () => {
|
||
it("_retrieveNodesFromData system message includes 'json' for any userId", async () => {
|
||
for (const userId of ["", "user-1", "special<>chars", "ユーザー"]) {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
const mg = graph();
|
||
await mg._retrieveNodesFromData("test", { userId });
|
||
|
||
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
|
||
expect(systemMsg.toLowerCase()).toContain("json");
|
||
}
|
||
});
|
||
|
||
it("_establishNodesRelationsFromData system message includes 'json' for any userId", async () => {
|
||
for (const userId of ["", "user-1", "special<>chars"]) {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
const mg = graph();
|
||
await mg._establishNodesRelationsFromData("test", { userId }, {});
|
||
|
||
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
|
||
expect(systemMsg.toLowerCase()).toContain("json");
|
||
}
|
||
});
|
||
|
||
it("_getDeleteEntitiesFromSearchOutput system message includes 'json' for any userId", async () => {
|
||
for (const userId of ["", "user-1", "special<>chars"]) {
|
||
mockGenerateResponse.mockResolvedValueOnce({ toolCalls: [] });
|
||
const mg = graph();
|
||
await mg._getDeleteEntitiesFromSearchOutput([], "test", { userId });
|
||
|
||
const systemMsg = mockGenerateResponse.mock.calls.at(-1)![0][0].content;
|
||
expect(systemMsg.toLowerCase()).toContain("json");
|
||
}
|
||
});
|
||
});
|
||
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
// 5. Edge cases – malformed entity fields in _removeSpacesFromEntities
|
||
// ═══════════════════════════════════════════════════════════════════════════
|
||
|
||
describe("_removeSpacesFromEntities (via _establishNodesRelationsFromData)", () => {
|
||
it("normalises spaces and case in entity source/relationship/destination", async () => {
|
||
mockGenerateResponse.mockResolvedValueOnce({
|
||
toolCalls: [
|
||
{
|
||
name: "establish_relationships",
|
||
arguments: JSON.stringify({
|
||
entities: [
|
||
{
|
||
source: "New York",
|
||
relationship: "Capital Of",
|
||
destination: "United States",
|
||
},
|
||
],
|
||
}),
|
||
},
|
||
],
|
||
});
|
||
|
||
const mg = graph();
|
||
const result = await mg._establishNodesRelationsFromData(
|
||
"test",
|
||
FILTERS,
|
||
{},
|
||
);
|
||
|
||
expect(result).toEqual([
|
||
{
|
||
source: "new_york",
|
||
relationship: "capital_of",
|
||
destination: "united_states",
|
||
},
|
||
]);
|
||
});
|
||
});
|