Skip to content

fix: add tool support to Anthropic LLM for graph memory compatibility #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nickrahman merged 9 commits into from
Mar 15, 2026
Merged

fix: add tool support to Anthropic LLM for graph memory compatibility #1

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
5c2d718
fix: add tool support to Anthropic LLM for graph memory compatibility
nickrahman Mar 15, 2026
76843b4
fix: complete Anthropic provider tool_choice, top_p, and response_for…
nickrahman Mar 15, 2026
b0e9779
fix(anthropic): prevent message double-pass, guard empty response, fi…
nickrahman Mar 15, 2026
4e5d09b
fix(graph_memory): guard JSON.parse calls and prevent Cypher injection
nickrahman Mar 15, 2026
56f47f5
refactor: simplify top_p handling, add runtime guards, clean up tests
nickrahman Mar 15, 2026
69b1995
fix: normalize relationship types before validation, add response typ…
nickrahman Mar 15, 2026
5e9cc8a
fix: skip anthropic tests when anthropic package not installed
nickrahman Mar 15, 2026
b6f5ca7
fix: add noqa E402 for imports after importorskip
nickrahman Mar 15, 2026
362a5ba
fix: add anthropic to llms optional dependency group
nickrahman Mar 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 82 additions & 15 deletions mem0-ts/src/oss/src/llms/anthropic.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,37 +18,104 @@ export class AnthropicLLM implements LLM {
async generateResponse(
messages: Message[],
responseFormat?: { type: string },
): Promise<string> {
// Extract system message if present
tools?: any[],
toolChoice: string = "auto",
Comment on lines +21 to +22

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

New tools and toolChoice parameters enable graph memory to use Anthropic for entity extraction via function calling. Return type widens from Promise<string> to Promise<string | LLMResponse> — callers that don't pass tools still get a plain string.

): Promise<string | LLMResponse> {
const systemMessage = messages.find((msg) => msg.role === "system");
const otherMessages = messages.filter((msg) => msg.role !== "system");

const response = await this.client.messages.create({
model: this.model,
messages: otherMessages.map((msg) => ({
let mappedMessages = messages
.filter((msg) => msg.role !== "system")
.map((msg) => ({
role: msg.role as "user" | "assistant",
content:
typeof msg.content === "string"
? msg.content
: msg.content.image_url.url,
})),
system:
typeof systemMessage?.content === "string"
? systemMessage.content
: undefined,
}));

if (responseFormat?.type === "json_object" && !tools) {
const last = mappedMessages[mappedMessages.length - 1];
if (last && last.role === "user") {
mappedMessages = [
...mappedMessages.slice(0, -1),
{ ...last, content: last.content + "\n\nYou must respond with valid JSON only." },
];
}
}

const systemContent =
typeof systemMessage?.content === "string" ? systemMessage.content : undefined;

const params: Record<string, any> = {
model: this.model,
messages: mappedMessages,
system: systemContent,
max_tokens: 4096,
});
};

if (tools) {
params.tools = this._convertTools(tools);
const mapped = this._mapToolChoice(toolChoice);
if (mapped !== null) {
params.tool_choice = mapped;
}
}

const response = await this.client.messages.create(params as any);

if (tools) {
let content = "";
const toolCalls: Array<{ name: string; arguments: string }> = [];
for (const block of response.content) {
if (block.type === "text") {
content = block.text;
} else if (block.type === "tool_use") {
toolCalls.push({
name: block.name,
arguments: JSON.stringify(block.input),
});

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: arguments are JSON-stringified here (matching the OpenAI tool_call format that graph_memory.ts expects to JSON.parse). The Python implementation returns the raw dict instead — this is an intentional divergence since Python callers expect a dict while TS callers expect a string.

}
}
return { content, role: "assistant", toolCalls };
}

const firstBlock = response.content[0];
if (!firstBlock) {
throw new Error("Empty response from Anthropic API");
}
if (firstBlock.type === "text") {
return firstBlock.text;
} else {
throw new Error("Unexpected response type from Anthropic API");
}
throw new Error("Unexpected response type from Anthropic API");
}

private _mapToolChoice(toolChoice: string): Record<string, string> | null {
if (toolChoice === "auto") return { type: "auto" };
if (toolChoice === "required") return { type: "any" };
if (toolChoice === "none") return null;
return { type: "tool", name: toolChoice };
}

private _convertTools(tools: any[]): any[] {
Comment on lines +94 to +98

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maps OpenAI-style tool_choice values to Anthropic's format. Key differences:

  • "required" → {type: "any"} (Anthropic's equivalent meaning "must use a tool")
  • "none" → null (caller omits the param entirely; Anthropic has no "none" equivalent — we just don't send tools)
  • A specific tool name → {type: "tool", name: "..."} (force a particular tool)

// Validate structure before mapping to catch malformed tool definitions early
return tools.map((tool, i) => {
if (!tool.function) {
throw new Error(`Tool at index ${i} is missing required key 'function'`);
}
const { name, description, parameters } = tool.function;
if (!name || !description || !parameters) {
throw new Error(
`Tool at index ${i} is missing required function keys (name, description, parameters)`,
);
}
return { name, description, input_schema: parameters };
});
}

async generateChat(messages: Message[]): Promise<LLMResponse> {
const response = await this.generateResponse(messages);
if (typeof response !== "string") {
throw new Error("generateChat received a non-string response; use generateResponse with tools instead");
}
return {
content: response,
role: "assistant",
Comment on lines 119 to 121

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Runtime guard replacing an unsafe as string cast. generateChat calls generateResponse without tools, so it should always get a string — but if the contract ever changes, this throws a clear error instead of silently producing [object Object] as the content.

Expand Down
64 changes: 35 additions & 29 deletions mem0-ts/src/oss/src/memory/graph_memory.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,15 +23,6 @@ interface SearchOutput {
similarity: number;
}

interface ToolCall {

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

These local ToolCall and LLMResponse interfaces were duplicates of the canonical definitions already exported from ../llms/base.ts. Removed to avoid divergence — if base.ts adds a field (e.g., role), this local copy would silently fall out of sync.

name: string;
arguments: string;
}

interface LLMResponse {
toolCalls?: ToolCall[];
}

interface Tool {
type: string;
function: {
Expand All @@ -52,7 +43,7 @@ export class MemoryGraph {
private graph: Driver;
private embeddingModel: Embedder;

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

structuredLlm (removed) was a second LLM instance created with the identical config as this.llm. Now that generateResponse accepts tools directly, there's no need for a separate instance. All three call sites (_retrieveNodesFromData, _establishNodesRelationsFromData, _getDeleteEntitiesFromSearchOutput) now use this.llm instead.

private llm: LLM;
private structuredLlm: LLM;

private llmProvider: string;
private threshold: number;

Expand All @@ -79,19 +70,10 @@ export class MemoryGraph {
this.config.embedder.config,
);
Comment on lines 70 to 71

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Simplified from a three-step if/if override pattern to nullish coalescing. Same precedence: graphStore.llm.provider > llm.provider > "openai" default.


this.llmProvider = "openai";
if (this.config.llm?.provider) {
this.llmProvider = this.config.llm.provider;
}
if (this.config.graphStore?.llm?.provider) {
this.llmProvider = this.config.graphStore.llm.provider;
}
this.llmProvider =
this.config.graphStore?.llm?.provider ?? this.config.llm?.provider ?? "openai";

this.llm = LLMFactory.create(this.llmProvider, this.config.llm.config);
this.structuredLlm = LLMFactory.create(
this.llmProvider,
this.config.llm.config,
);
this.threshold = 0.7;
}

Expand Down Expand Up @@ -208,7 +190,7 @@ export class MemoryGraph {
filters: Record<string, any>,
) {
const tools = [EXTRACT_ENTITIES_TOOL] as Tool[];
const searchResults = await this.structuredLlm.generateResponse(
const searchResults = await this.llm.generateResponse(
[
{
role: "system",
Expand Down Expand Up @@ -284,7 +266,7 @@ export class MemoryGraph {
}

const tools = [RELATIONS_TOOL] as Tool[];
const extractedEntities = await this.structuredLlm.generateResponse(
const extractedEntities = await this.llm.generateResponse(
messages,
{ type: "json_object" },
tools,
Expand All @@ -294,8 +276,12 @@ export class MemoryGraph {
if (typeof extractedEntities !== "string" && extractedEntities.toolCalls) {
const toolCall = extractedEntities.toolCalls[0];
if (toolCall && toolCall.arguments) {
const args = JSON.parse(toolCall.arguments);
entities = args.entities || [];
try {
const args = JSON.parse(toolCall.arguments);
entities = args.entities || [];
} catch (e) {
logger.error(`Failed to parse relation tool arguments: ${e}`);
}
Comment on lines 276 to +284

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added try/catch around JSON.parse of LLM tool call arguments. Previously, malformed JSON from the LLM would throw an uncaught SyntaxError that crashed the entire add() operation. Now it logs and falls back to an empty entity list — consistent with the existing catch block in _retrieveNodesFromData above.

}
}

Expand Down Expand Up @@ -385,7 +371,7 @@ export class MemoryGraph {
);

const tools = [DELETE_MEMORY_TOOL_GRAPH] as Tool[];
const memoryUpdates = await this.structuredLlm.generateResponse(
const memoryUpdates = await this.llm.generateResponse(
[
{ role: "system", content: systemPrompt },
{ role: "user", content: userPrompt },
Expand All @@ -398,7 +384,11 @@ export class MemoryGraph {
if (typeof memoryUpdates !== "string" && memoryUpdates.toolCalls) {
for (const item of memoryUpdates.toolCalls) {
if (item.name === "delete_graph_memory") {
toBeDeleted.push(JSON.parse(item.arguments));
try {
toBeDeleted.push(JSON.parse(item.arguments));
} catch (e) {
logger.error(`Failed to parse delete tool arguments: ${e}`);
}
}
}
}
Expand All @@ -416,7 +406,8 @@ export class MemoryGraph {

try {
Comment on lines 406 to 407

@nickrahman nickrahman Mar 15, 2026

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Cypher injection prevention. relationship comes from LLM output and is interpolated directly into Cypher queries (Neo4j doesn't support parameterized relationship types). _sanitizeRelationshipType normalizes to lowercase, replaces spaces with underscores, then validates against [a-z0-9_]+ — rejecting any value with special characters that could alter the query.

for (const item of toBeDeleted) {
const { source, destination, relationship } = item;
const { source, destination } = item;
const relationship = this._sanitizeRelationshipType(item.relationship);

const cypher = `
MATCH (n {name: $source_name, user_id: $user_id})
Expand Down Expand Up @@ -454,7 +445,8 @@ export class MemoryGraph {

try {
for (const item of toBeAdded) {
const { source, destination, relationship } = item;
const { source, destination } = item;
const relationship = this._sanitizeRelationshipType(item.relationship);
const sourceType = entityTypeMap[source] || "unknown";
const destinationType = entityTypeMap[destination] || "unknown";

Expand Down Expand Up @@ -581,6 +573,20 @@ export class MemoryGraph {
}));
}

/**
* Validate that a relationship type contains only characters safe for Cypher
* interpolation (alphanumerics and underscores). LLM-controlled values must
* be validated before being embedded in query strings because Neo4j does not
* support parameterized relationship types.
*/
private _sanitizeRelationshipType(relationship: string): string {
const normalized = relationship.toLowerCase().replace(/\s+/g, "_");
if (!/^[a-z0-9_]+$/.test(normalized)) {
throw new Error(`Unsafe relationship type rejected: "${relationship}"`);
}
return normalized;
}

private async _searchSourceNode(
sourceEmbedding: number[],
userId: string,
Expand Down
Loading
Loading