diff --git a/docs/components/llms/models/lmstudio.mdx b/docs/components/llms/models/lmstudio.mdx index 9c5195611..fa3b23474 100644 --- a/docs/components/llms/models/lmstudio.mdx +++ b/docs/components/llms/models/lmstudio.mdx @@ -23,6 +23,7 @@ config = { "temperature": 0.2, "max_tokens": 2000, "lmstudio_base_url": "http://localhost:1234/v1", # default LM Studio API URL + "lmstudio_response_format": {"type": "json_schema", "json_schema": {"type": "object", "schema": {}}}, } } } diff --git a/mem0/configs/llms/base.py b/mem0/configs/llms/base.py index 271de2ef4..900556832 100644 --- a/mem0/configs/llms/base.py +++ b/mem0/configs/llms/base.py @@ -43,6 +43,7 @@ class BaseLlmConfig(ABC): sarvam_base_url: Optional[str] = "https://api.sarvam.ai/v1", # LM Studio specific lmstudio_base_url: Optional[str] = "http://localhost:1234/v1", + lmstudio_response_format: dict = None, # AWS Bedrock specific aws_access_key_id: Optional[str] = None, aws_secret_access_key: Optional[str] = None, @@ -95,6 +96,8 @@ class BaseLlmConfig(ABC): :type sarvam_base_url: Optional[str], optional :param lmstudio_base_url: LM Studio base URL to be use, defaults to "http://localhost:1234/v1" :type lmstudio_base_url: Optional[str], optional + :param lmstudio_response_format: LM Studio response format to be use, defaults to None + :type lmstudio_response_format: Optional[Dict], optional """ self.model = model @@ -134,6 +137,7 @@ class BaseLlmConfig(ABC): # LM Studio specific self.lmstudio_base_url = lmstudio_base_url + self.lmstudio_response_format = lmstudio_response_format # AWS Bedrock specific self.aws_access_key_id = aws_access_key_id diff --git a/mem0/llms/lmstudio.py b/mem0/llms/lmstudio.py index 7e003eaf6..cdcafe555 100644 --- a/mem0/llms/lmstudio.py +++ b/mem0/llms/lmstudio.py @@ -46,6 +46,8 @@ class LMStudioLLM(LLMBase): } if response_format: params["response_format"] = response_format + if self.config.lmstudio_response_format is not None: + params["response_format"] = self.config.lmstudio_response_format response = self.client.chat.completions.create(**params) return response.choices[0].message.content diff --git a/tests/llms/test_lm_studio.py b/tests/llms/test_lm_studio.py index bed1b98b6..37033e024 100644 --- a/tests/llms/test_lm_studio.py +++ b/tests/llms/test_lm_studio.py @@ -42,3 +42,30 @@ def test_generate_response_without_tools(mock_lm_studio_client): ) assert response == "I'm doing well, thank you for asking!" + +def test_generate_response_specifying_response_format(mock_lm_studio_client): + config = BaseLlmConfig( + model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", + temperature=0.7, + max_tokens=100, + top_p=1.0, + lmstudio_response_format={"type": "json_schema"}, # Specifying the response format in config + ) + llm = LMStudioLLM(config) + messages = [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello, how are you?"}, + ] + + response = llm.generate_response(messages) + + mock_lm_studio_client.chat.completions.create.assert_called_once_with( + model="lmstudio-community/Meta-Llama-3.1-8B-Instruct-GGUF/Meta-Llama-3.1-8B-Instruct-Q4_K_M.gguf", + messages=messages, + temperature=0.7, + max_tokens=100, + top_p=1.0, + response_format={"type": "json_schema"}, + ) + + assert response == "I'm doing well, thank you for asking!" \ No newline at end of file