diff --git a/docs/my-website/docs/providers/lm_studio.md b/docs/my-website/docs/providers/lm_studio.md index 45c546ada68f..0cf9acff33db 100644 --- a/docs/my-website/docs/providers/lm_studio.md +++ b/docs/my-website/docs/providers/lm_studio.md @@ -153,3 +153,26 @@ response = embedding( ) print(response) ``` + + +## Structured Output + +LM Studio supports structured outputs via JSON Schema. You can pass a pydantic model or a raw schema using `response_format`. +LiteLLM sends the schema as `{ "type": "json_schema", "json_schema": {"schema": } }`. + +```python +from pydantic import BaseModel +from litellm import completion + +class Book(BaseModel): + title: str + author: str + year: int + +response = completion( + model="lm_studio/llama-3-8b-instruct", + messages=[{"role": "user", "content": "Tell me about The Hobbit"}], + response_format=Book, +) +print(response.choices[0].message.content) +``` \ No newline at end of file diff --git a/litellm/llms/lm_studio/chat/transformation.py b/litellm/llms/lm_studio/chat/transformation.py index 147e8e923f2b..f7a2cc0f28a5 100644 --- a/litellm/llms/lm_studio/chat/transformation.py +++ b/litellm/llms/lm_studio/chat/transformation.py @@ -18,3 +18,32 @@ def _get_openai_compatible_provider_info( api_key or get_secret_str("LM_STUDIO_API_KEY") or " " ) # vllm does not require an api key return api_base, dynamic_api_key + + def map_openai_params( + self, + non_default_params: dict, + optional_params: dict, + model: str, + drop_params: bool, + ) -> dict: + for param, value in list(non_default_params.items()): + if param == "response_format" and isinstance(value, dict): + if value.get("type") == "json_schema": + if "json_schema" not in value and "schema" in value: + optional_params["response_format"] = { + "type": "json_schema", + "json_schema": {"schema": value.get("schema")}, + } + else: + optional_params["response_format"] = value + non_default_params.pop(param, None) + elif value.get("type") == "json_object": + optional_params["response_format"] = value + non_default_params.pop(param, None) + + return super().map_openai_params( + non_default_params=non_default_params, + optional_params=optional_params, + model=model, + drop_params=drop_params, + ) \ No newline at end of file diff --git a/litellm/utils.py b/litellm/utils.py index fbae31d41f8f..8519070a37fd 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1824,6 +1824,7 @@ def supports_response_schema( PROVIDERS_GLOBALLY_SUPPORT_RESPONSE_SCHEMA = [ litellm.LlmProviders.PREDIBASE, litellm.LlmProviders.FIREWORKS_AI, + litellm.LlmProviders.LM_STUDIO, ] if custom_llm_provider in PROVIDERS_GLOBALLY_SUPPORT_RESPONSE_SCHEMA: diff --git a/tests/litellm/llms/lm_studio/test_lm_studio_chat_transformation.py b/tests/litellm/llms/lm_studio/test_lm_studio_chat_transformation.py new file mode 100644 index 000000000000..718142ad1751 --- /dev/null +++ b/tests/litellm/llms/lm_studio/test_lm_studio_chat_transformation.py @@ -0,0 +1,51 @@ +import os +import sys + +from pydantic import BaseModel + +sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../.."))) + +from litellm.llms.lm_studio.chat.transformation import LMStudioChatConfig +from litellm.utils import get_optional_params + + +class Book(BaseModel): + title: str + author: str + year: int + + +class TestLMStudioChatConfigResponseFormat: + def test_get_optional_params_with_pydantic_model(self): + optional_params = get_optional_params( + model="lm_studio/test-model", + response_format=Book, + custom_llm_provider="lm_studio", + ) + + assert "response_format" in optional_params + transformed = optional_params["response_format"] + assert transformed.get("type") == "json_schema" + schema = transformed.get("json_schema", {}).get("schema") + assert schema["properties"] == Book.model_json_schema()["properties"] + + def test_map_openai_params_with_dict_json_schema(self): + config = LMStudioChatConfig() + schema = Book.model_json_schema() + response_format_dict = { + "type": "json_schema", + "json_schema": {"schema": schema}, + } + + non_default_params = {"response_format": response_format_dict} + optional_params = get_optional_params( + model="lm_studio/test-model", + response_format=response_format_dict, + custom_llm_provider="lm_studio", + ) + + mapped = config.map_openai_params(non_default_params, {}, "lm_studio/test-model", False) + mapped_schema = mapped["response_format"]["json_schema"]["schema"] + assert mapped_schema["properties"] == schema["properties"] + opt_schema = optional_params["response_format"]["json_schema"]["schema"] + assert opt_schema["properties"] == schema["properties"]