my-project-public

repository

loading code, commits, and activity

repositories

loading repo index

my-project-public — gitlawb

#1	from unittest.mock import Mock
#2
#3	import pytest
#4
#5	from mem0.configs.llms.base import BaseLlmConfig
#6	from mem0.llms.langchain import LangchainLLM
#7
#8	# Add the import for BaseChatModel
#9	try:
#10	from langchain.chat_models.base import BaseChatModel
#11	except ImportError:
#12	from unittest.mock import MagicMock
#13
#14	BaseChatModel = MagicMock
#15
#16
#17	@pytest.fixture
#18	def mock_langchain_model():
#19	"""Mock a Langchain model for testing."""
#20	mock_model = Mock(spec=BaseChatModel)
#21	mock_model.invoke.return_value = Mock(content="This is a test response")
#22	return mock_model
#23
#24
#25	def test_langchain_initialization(mock_langchain_model):
#26	"""Test that LangchainLLM initializes correctly with a valid model."""
#27	# Create a config with the model instance directly
#28	config = BaseLlmConfig(model=mock_langchain_model, temperature=0.7, max_tokens=100, api_key="test-api-key")
#29
#30	# Initialize the LangchainLLM
#31	llm = LangchainLLM(config)
#32
#33	# Verify the model was correctly assigned
#34	assert llm.langchain_model == mock_langchain_model
#35
#36
#37	def test_generate_response(mock_langchain_model):
#38	"""Test that generate_response correctly processes messages and returns a response."""
#39	# Create a config with the model instance
#40	config = BaseLlmConfig(model=mock_langchain_model, temperature=0.7, max_tokens=100, api_key="test-api-key")
#41
#42	# Initialize the LangchainLLM
#43	llm = LangchainLLM(config)
#44
#45	# Create test messages
#46	messages = [
#47	{"role": "system", "content": "You are a helpful assistant."},
#48	{"role": "user", "content": "Hello, how are you?"},
#49	{"role": "assistant", "content": "I'm doing well! How can I help you?"},
#50	{"role": "user", "content": "Tell me a joke."},
#51	]
#52
#53	# Get response
#54	response = llm.generate_response(messages)
#55
#56	# Verify the correct message format was passed to the model
#57	expected_langchain_messages = [
#58	("system", "You are a helpful assistant."),
#59	("human", "Hello, how are you?"),
#60	("ai", "I'm doing well! How can I help you?"),
#61	("human", "Tell me a joke."),
#62	]
#63
#64	mock_langchain_model.invoke.assert_called_once()
#65	# Extract the first argument of the first call
#66	actual_messages = mock_langchain_model.invoke.call_args[0][0]
#67	assert actual_messages == expected_langchain_messages
#68	assert response == "This is a test response"
#69
#70
#71	def test_generate_response_with_tools(mock_langchain_model):
#72	config = BaseLlmConfig(model=mock_langchain_model, temperature=0.7, max_tokens=100, api_key="test-api-key")
#73	llm = LangchainLLM(config)
#74
#75	messages = [
#76	{"role": "system", "content": "You are a helpful assistant."},
#77	{"role": "user", "content": "Add a new memory: Today is a sunny day."},
#78	]
#79	tools = [
#80	{
#81	"type": "function",
#82	"function": {
#83	"name": "add_memory",
#84	"description": "Add a memory",
#85	"parameters": {
#86	"type": "object",
#87	"properties": {"data": {"type": "string", "description": "Data to add to memory"}},
#88	"required": ["data"],
#89	},
#90	},
#91	}
#92	]
#93
#94	mock_response = Mock()
#95	mock_response.content = "I've added the memory for you."
#96
#97	mock_tool_call = Mock()
#98	mock_tool_call.__getitem__ = Mock(
#99	side_effect={"name": "add_memory", "args": {"data": "Today is a sunny day."}}.__getitem__
#100	)
#101
#102	mock_response.tool_calls = [mock_tool_call]
#103	mock_langchain_model.invoke.return_value = mock_response
#104	mock_langchain_model.bind_tools.return_value = mock_langchain_model
#105
#106	response = llm.generate_response(messages, tools=tools)
#107
#108	mock_langchain_model.invoke.assert_called_once()
#109
#110	assert response["content"] == "I've added the memory for you."
#111	assert len(response["tool_calls"]) == 1
#112	assert response["tool_calls"][0]["name"] == "add_memory"
#113	assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
#114
#115
#116	def test_invalid_model():
#117	"""Test that LangchainLLM raises an error with an invalid model."""
#118	config = BaseLlmConfig(model="not-a-valid-model-instance", temperature=0.7, max_tokens=100, api_key="test-api-key")
#119
#120	with pytest.raises(ValueError, match="`model` must be an instance of BaseChatModel"):
#121	LangchainLLM(config)
#122
#123
#124	def test_missing_model():
#125	"""Test that LangchainLLM raises an error when model is None."""
#126	config = BaseLlmConfig(model=None, temperature=0.7, max_tokens=100, api_key="test-api-key")
#127
#128	with pytest.raises(ValueError, match="`model` parameter is required"):
#129	LangchainLLM(config)
#130

z6Mkq5mY3JWtxoxUobWcfNHm7AkRubgSWEZTkBVqZXJviFZ5/my-project-public