diff --git a/mem0/llms/langchain.py b/mem0/llms/langchain.py index 9833cd5bd..7cc9afa7f 100644 --- a/mem0/llms/langchain.py +++ b/mem0/llms/langchain.py @@ -57,6 +57,7 @@ class LangchainLLM(LLMBase): response_format=None, tools: Optional[List[Dict]] = None, tool_choice: str = "auto", + **kwargs, ): """ Generate a response based on the given messages using langchain_community. @@ -66,6 +67,8 @@ class LangchainLLM(LLMBase): response_format (str or object, optional): Format of the response. Not used in Langchain. tools (list, optional): List of tools that the model can call. tool_choice (str, optional): Tool choice method. + **kwargs: Additional model parameters forwarded to the underlying LangChain + model's ``invoke`` (matches the ``LLMBase.generate_response`` contract). Returns: str: The generated response. @@ -90,5 +93,5 @@ class LangchainLLM(LLMBase): if tools: langchain_model = langchain_model.bind_tools(tools=tools, tool_choice=tool_choice) - response: AIMessage = langchain_model.invoke(langchain_messages) + response: AIMessage = langchain_model.invoke(langchain_messages, **kwargs) return self._parse_response(response, tools) diff --git a/mem0/llms/sarvam.py b/mem0/llms/sarvam.py index 6ef836ed6..f963d57cb 100644 --- a/mem0/llms/sarvam.py +++ b/mem0/llms/sarvam.py @@ -28,7 +28,7 @@ class SarvamLLM(LLMBase): getattr(self.config, "sarvam_base_url", None) or os.getenv("SARVAM_API_BASE") or "https://api.sarvam.ai/v1" ) - def generate_response(self, messages: List[Dict[str, str]], response_format=None) -> str: + def generate_response(self, messages: List[Dict[str, str]], response_format=None, **kwargs) -> str: """ Generate a response based on the given messages using Sarvam-M. @@ -36,6 +36,8 @@ class SarvamLLM(LLMBase): messages (list): List of message dicts containing 'role' and 'content'. response_format (str or object, optional): Format of the response. Currently not used by Sarvam API. + **kwargs: Additional provider-specific parameters forwarded to the Sarvam + request payload (matches the ``LLMBase.generate_response`` contract). Returns: str: The generated response. @@ -72,6 +74,9 @@ class SarvamLLM(LLMBase): if param in self.config.model: params[param] = self.config.model[param] + # Forward any per-call provider-specific parameters (LLMBase contract). + params.update(kwargs) + try: response = requests.post(url, headers=headers, json=params, timeout=30) response.raise_for_status() diff --git a/mem0/llms/together.py b/mem0/llms/together.py index d2af10c18..0ec08f00f 100644 --- a/mem0/llms/together.py +++ b/mem0/llms/together.py @@ -58,6 +58,7 @@ class TogetherLLM(LLMBase): response_format=None, tools: Optional[List[Dict]] = None, tool_choice: str = "auto", + **kwargs, ): """ Generate a response based on the given messages using TogetherAI. @@ -67,6 +68,8 @@ class TogetherLLM(LLMBase): response_format (str or object, optional): Format of the response. Defaults to "text". tools (list, optional): List of tools that the model can call. Defaults to None. tool_choice (str, optional): Tool choice method. Defaults to "auto". + **kwargs: Additional provider-specific parameters forwarded to the Together + client (matches the ``LLMBase.generate_response`` contract). Returns: str: The generated response. @@ -78,6 +81,7 @@ class TogetherLLM(LLMBase): "max_tokens": self.config.max_tokens, "top_p": self.config.top_p, } + params.update(kwargs) if response_format: params["response_format"] = response_format if tools: # TODO: Remove tools if no issues found with new memory addition logic diff --git a/tests/llms/test_langchain.py b/tests/llms/test_langchain.py index 1b5bf4c18..bd2f3eb4c 100644 --- a/tests/llms/test_langchain.py +++ b/tests/llms/test_langchain.py @@ -115,6 +115,20 @@ def test_generate_response_with_tools(mock_langchain_model): assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."} +def test_generate_response_forwards_extra_kwargs(mock_langchain_model): + """Per the LLMBase contract, extra model kwargs must be accepted and forwarded to + the underlying LangChain model's ``invoke``.""" + config = BaseLlmConfig(model=mock_langchain_model, temperature=0.7, max_tokens=100, api_key="test-api-key") + llm = LangchainLLM(config) + messages = [{"role": "user", "content": "Hello"}] + + response = llm.generate_response(messages, frequency_penalty=0.5) + + mock_langchain_model.invoke.assert_called_once() + assert mock_langchain_model.invoke.call_args.kwargs["frequency_penalty"] == 0.5 + assert response == "This is a test response" + + def test_invalid_model(): """Test that LangchainLLM raises an error with an invalid model.""" config = BaseLlmConfig(model="not-a-valid-model-instance", temperature=0.7, max_tokens=100, api_key="test-api-key") diff --git a/tests/llms/test_sarvam.py b/tests/llms/test_sarvam.py new file mode 100644 index 000000000..765a51796 --- /dev/null +++ b/tests/llms/test_sarvam.py @@ -0,0 +1,46 @@ +from unittest.mock import Mock, patch + +import pytest + +from mem0.configs.llms.base import BaseLlmConfig +from mem0.llms.sarvam import SarvamLLM + + +@pytest.fixture +def sarvam_llm(): + config = BaseLlmConfig(model="sarvam-m", temperature=0.7, max_tokens=100, top_p=1.0, api_key="test-api-key") + return SarvamLLM(config) + + +def _mock_post(content="Hello there!"): + mock_response = Mock() + mock_response.raise_for_status.return_value = None + mock_response.json.return_value = {"choices": [{"message": {"content": content}}]} + return mock_response + + +def test_generate_response_returns_content(sarvam_llm): + with patch("mem0.llms.sarvam.requests.post", return_value=_mock_post("Hi!")) as mock_post: + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello, how are you?"}, + ] + response = sarvam_llm.generate_response(messages) + + assert response == "Hi!" + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["model"] == "sarvam-m" + assert sent_payload["messages"] == messages + assert sent_payload["temperature"] == 0.7 + + +def test_generate_response_forwards_extra_kwargs(sarvam_llm): + """Per the LLMBase contract, extra provider-specific kwargs must be accepted and + forwarded into the Sarvam request payload.""" + with patch("mem0.llms.sarvam.requests.post", return_value=_mock_post("Hi!")) as mock_post: + messages = [{"role": "user", "content": "Hello"}] + response = sarvam_llm.generate_response(messages, frequency_penalty=0.5) + + assert response == "Hi!" + sent_payload = mock_post.call_args.kwargs["json"] + assert sent_payload["frequency_penalty"] == 0.5 diff --git a/tests/llms/test_together.py b/tests/llms/test_together.py index 7c59ee419..63428baad 100644 --- a/tests/llms/test_together.py +++ b/tests/llms/test_together.py @@ -84,3 +84,21 @@ def test_generate_response_with_tools(mock_together_client): assert len(response["tool_calls"]) == 1 assert response["tool_calls"][0]["name"] == "add_memory" assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."} + + +def test_generate_response_forwards_extra_kwargs(mock_together_client): + """Per the LLMBase contract, extra provider-specific kwargs must be accepted and + forwarded to the Together client (matching openai/deepseek/vllm/xai behavior).""" + config = BaseLlmConfig(model="mistralai/Mixtral-8x7B-Instruct-v0.1", temperature=0.7, max_tokens=100, top_p=1.0) + llm = TogetherLLM(config) + messages = [{"role": "user", "content": "Hello"}] + + mock_response = Mock() + mock_response.choices = [Mock(message=Mock(content="Hi"))] + mock_together_client.chat.completions.create.return_value = mock_response + + response = llm.generate_response(messages, frequency_penalty=0.5) + + assert response == "Hi" + call_kwargs = mock_together_client.chat.completions.create.call_args.kwargs + assert call_kwargs["frequency_penalty"] == 0.5