Fix: Correct Graph LLM Initialization and Tool Schema

[EnzoFanAccount]

Description

This PR addresses two critical bugs that prevented the Graph Memory feature from functioning correctly when using a specific, but valid, configuration. The primary issue occurred when Mem0 was initialized with a main llm config containing a non-serializable object (e.g., a Langchain BaseChatModel instance), while a separate, serializable dictionary was provided for the graph_store.llm.

This led to two sequential failures:

1. A TypeError during LLM request serialization because the incorrect LLM config was used.
2. An invalid_request_error from the OpenAI API due to a malformed tool schema.

These fixes ensure that the graph memory module correctly respects its dedicated configuration and uses a valid tool schema, making it robust and usable in complex integration scenarios.

1. memory/graph_memory.py

• Problem: The MemoryGraph class was always using the main llm_config (self.config.llm.config), ignoring the specific configuration provided in self.config.graph_store.llm.config. This caused a crash when the main config was non-serializable.
• Solution: The __init__ method has been updated to correctly prioritize the graph-specific LLM configuration. It now checks for the presence of config.graph_store.llm and uses its config dictionary if available. This ensures the correct, serializable configuration is passed to the LlmFactory for graph operations.

2. graphs/tools.py

• Problem: The RELATIONS_STRUCT_TOOL had typos in its schema. The properties keys were source_entity, relatationship, and destination_entity, and the required array was also inconsistent. This caused a 400 Bad Request from the OpenAI API due to an invalid function schema.
• Solution: The keys in both the properties object and the required array have been corrected to source, relationship, and destination. This aligns the tool schema with OpenAI's requirements and the expectations of the downstream parsing logic.

Fixes # (issue)

Type of change

Please delete options that are not relevant.

• [x] Bug fix (non-breaking change which fixes an issue)

How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

Please delete options that are not relevant.

• [x] Test Script (please provide) This bug was found when testing this PR over at Browser Use. The way to run into the bug in the most loyal way to how I did would be to clone Browser Use, make the changes proposed in my PR, and run the example script provided:

import asyncio
import os
from dotenv import load_dotenv

from browser_use import Agent
from browser_use.agent.memory import MemoryConfig
from langchain_openai import ChatOpenAI

load_dotenv()
NEO4J_URI = os.getenv('NEO4J_URI')
NEO4J_USERNAME = os.getenv('NEO4J_USERNAME')
NEO4J_PASSWORD = os.getenv('NEO4J_PASSWORD')
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY')

print(f'Attempting direct connect to URI: {NEO4J_URI}, User: {NEO4J_USERNAME}')

# --- 1. Initialize LLMs ---
agent_llm = ChatOpenAI(model='gpt-4o-mini', temperature=0.0, api_key=OPENAI_API_KEY)

# --- 2. Configure Mem0 Graph Memory via Browser Use's MemoryConfig ---
# Mem0 uses agent_id as user_id for graph operations context.
AGENT_ID_FOR_MEM0 = 'user_browser_task_graph_v1'

graph_memory_config = MemoryConfig(
	# Main LLM instance for Mem0
	llm_instance=agent_llm,
	agent_id=AGENT_ID_FOR_MEM0,
	memory_interval=2,
	# --- Graph Store Configuration ---
	graph_store_provider='neo4j',
	graph_store_config_override={
		'url': NEO4J_URI,
		'username': NEO4J_USERNAME,
		'password': NEO4J_PASSWORD,
	},
	# Optional: Override LLM specifically for graph operations within Mem0.
	# This tells Mem0 to use `mem0_internal_graph_llm` for its graph-related LLM calls.
	graph_store_llm_config_override={
		'provider': 'openai',
		'config': {'model': 'gpt-4o-mini', 'api_key': OPENAI_API_KEY, 'temperature': 0.0},
	},
	# Optional: Custom prompt for Mem0's graph entity/relationship extraction.
	graph_store_custom_prompt=(
		'From the conversation and browser interactions, extract entities like URLs, search queries, '
		'page titles, section titles, key facts extracted, user intents, and agent actions. '
		"Use 'USER_ID' for the user ('user_browser_task_graph_v1' in this case) and 'AGENT' for the assistant. "
		'Capture relationships such as: '
		"'AGENT NAVIGATED_TO URL', 'AGENT SEARCHED_FOR Query', 'URL HAS_TITLE PageTitle', "
		"'PAGE_CONTENT_CONTAINS Fact', 'AGENT_EXTRACTED Fact FROM URL', "
		"'USER_ID TASKED_AGENT_WITH UserTask', 'AGENT_ACTION_PERFORMED ActionDetail'. "
		'Focus on the sequence of web operations and information flow.'
	),
)

TASK_DESCRIPTION = (
	'Go to en.wikipedia.org. Search for "Persistent Memory". '
	'Navigate to the "Persistent Memory" Wikipedia page. '
	'Find and remember the title of the first main section of content on that page. '
	'Then, state what you found.'
)

# --- 3. Initialize Browser Use Agent ---

agent = Agent(
	task=TASK_DESCRIPTION,
	llm=agent_llm,
	enable_memory=True,  # This is True by default if memory_config is provided
	memory_config=graph_memory_config,
)


async def main():
	print(f"Starting Browser Use Agent for task: '{TASK_DESCRIPTION}'")
	print(f'Mem0 is configured with Neo4j for graph memory. Agent ID: {AGENT_ID_FOR_MEM0}')
	print(f'Ensure Neo4j is running at: {NEO4J_URI}')

	try:
		history = await agent.run()
		print('\n--- Agent Run Completed ---')

		final_result = history.final_result()
		if final_result:
			print(f"Agent's final output: {final_result}")
		else:
			print('Agent did not produce a final output within max_steps.')

		print(f'\nTotal steps taken: {history.number_of_steps()}')
		print(f'Total duration: {history.total_duration_seconds():.2f}s')

		if history.has_errors():
			print('\nErrors occurred during the run:')
			for i, error_msg in enumerate(history.errors()):
				if error_msg:
					print(f'  Step {i + 1}: {error_msg}')
		else:
			print('\nNo errors reported during the run.')

	except Exception as e:
		print(f'An error occurred during agent execution: {e}')

	print('\n--- Verification ---')
	print(f"Check your Neo4j instance for graph data related to user_id '{AGENT_ID_FOR_MEM0}'.")
	print('Example Cypher query for Neo4j Browser:')
	print(
		f"MATCH (n)-[r]-(m) WHERE n.user_id = '{AGENT_ID_FOR_MEM0}' OR m.user_id = '{AGENT_ID_FOR_MEM0}' RETURN n, r, m LIMIT 50"
	)
	print("You should see nodes and relationships reflecting the agent's actions (navigation, search, extraction) and findings.")


if __name__ == '__main__':
	asyncio.run(main())

Additionally, to help spot the bug more clearly, a debug section was added to llms/openai.py:

        if response_format:
            params["response_format"] = response_format
        if tools:  # TODO: Remove tools if no issues found with new memory addition logic
            params["tools"] = tools
            params["tool_choice"] = tool_choice

        import logging  # Make sure logging is imported if not already
        logger_openai_llm = logging.getLogger(__name__) # Get logger for this module
        try:
            params_json_str = json.dumps(params, indent=2)
            logger_openai_llm.error(f"MEM0_OPENAILLM_DEBUG: Params to OpenAI SDK: {params_json_str}")
        except TypeError as te:
            logger_openai_llm.error(f"MEM0_OPENAILLM_DEBUG: Params to OpenAI SDK (contains non-serializable, printing repr): {repr(params)}")
            logger_openai_llm.error(f"MEM0_OPENAILLM_DEBUG: Non-serializable error: {te}")

        response = self.client.chat.completions.create(**params)
        return self._parse_response(response, tools)

Checklist:

• [x] My code follows the style guidelines of this project
• [x] I have performed a self-review of my own code
• [x] I have commented my code, particularly in hard-to-understand areas
• [x] I have made corresponding changes to the documentation
• [x] My changes generate no new warnings
• [x] I have added tests that prove my fix is effective or that my feature works
• [x] New and existing unit tests pass locally with my changes
• [x] Any dependent changes have been merged and published in downstream modules
• [x] I have checked my code and corrected any misspellings

Maintainer Checklist

• [ ] closes #xxxx (Replace xxxx with the GitHub issue number)
• [ ] Made sure Checks passed

[CLAassistant]

All committers have signed the CLA.

[EnzoFanAccount]

@Dev-Khant Could you take a look at this when you have the time? I love the graph memory feature but it's not fully functional without this patch. It'd be nice to have it fixed before the next release. Thanks in advance!