feat(LM Studio): Add response_format param for LM Studio to config (mem0ai#2502)

Author: siroa

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": {}}},
         }
     }
 }

mem0/configs/llms/base.py

@@ -43,6 +43,7 @@ def __init__(
         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 @@ def __init__(
         :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 @@ def __init__(
 
         # 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

mem0/llms/lmstudio.py

@@ -46,6 +46,8 @@ def generate_response(
         }
         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

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!"