Repository: run-llama/llama_index
Pull request: Feat/lmstudio integration
# Description

Add Support for LM Studio
Resolves #13556

## New Package?

Did I fill in the `tool.llamahub` section in the `pyproject.toml` and provide a detailed README.md for my new integration or package?

- [x] Yes
- [ ] No

## Version Bump?

Did I bump the version in the `pyproject.toml` file of the package I am updating? (Except for the `llama-index-core` package)

- [x] Yes
- [ ] No

## Type of Change

Please delete options that are not relevant.

- [x] New feature (non-breaking change which adds functionality)

## How Has This Been Tested?

Please describe the tests that you ran to verify your changes. Provide instructions so we can reproduce. Please also list any relevant details for your test configuration

- [x] Added new notebook (that tests end-to-end)

## Suggested Checklist:

- [x] I have performed a self-review of my own code
- [x] I have commented my code, particularly in hard-to-understand areas
- [x] I have made corresponding changes to the documentation
- [x] I have added Google Colab support for the newly added notebooks.
- [x] My changes generate no new warnings
- [x] New and existing unit tests pass locally with my changes
- [x] I ran `make format; make lint` to appease the lint gods


Problem statement:
[Feature Request]: Add LMStudio support
### Feature Description

Add a way to connect LMStudio with LlamaIndex.

### Reason

Add direct implementation for LMStudio

### Value of Feature

This features enables loading models on LMStudio and allows LlamaIndex index to send request to the API provided by LMStudio.

Hints from issue discussion:

Commit message: Merge branch 'run-llama:main' into feat/lmstudio-integration

Patch to review:
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/.gitignore b/llama-index-integrations/llms/llama-index-llms-lmstudio/.gitignore
new file mode 100644
index 0000000000..990c18de22
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/.gitignore
@@ -0,0 +1,153 @@
+llama_index/_static
+.DS_Store
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+bin/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+etc/
+include/
+lib/
+lib64/
+parts/
+sdist/
+share/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+.ruff_cache
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+notebooks/
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+pyvenv.cfg
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Jetbrains
+.idea
+modules/
+*.swp
+
+# VsCode
+.vscode
+
+# pipenv
+Pipfile
+Pipfile.lock
+
+# pyright
+pyrightconfig.json
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/BUILD b/llama-index-integrations/llms/llama-index-llms-lmstudio/BUILD
new file mode 100644
index 0000000000..0896ca890d
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/BUILD
@@ -0,0 +1,3 @@
+poetry_requirements(
+    name="poetry",
+)
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/Makefile b/llama-index-integrations/llms/llama-index-llms-lmstudio/Makefile
new file mode 100644
index 0000000000..b9eab05aa3
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/Makefile
@@ -0,0 +1,17 @@
+GIT_ROOT ?= $(shell git rev-parse --show-toplevel)
+
+help:	## Show all Makefile targets.
+	@grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf "\033[33m%-30s\033[0m %s\n", $$1, $$2}'
+
+format:	## Run code autoformatters (black).
+	pre-commit install
+	git ls-files | xargs pre-commit run black --files
+
+lint:	## Run linters: pre-commit (black, ruff, codespell) and mypy
+	pre-commit install && git ls-files | xargs pre-commit run --show-diff-on-failure --files
+
+test:	## Run tests via pytest.
+	pytest tests
+
+watch-docs:	## Build and watch documentation.
+	sphinx-autobuild docs/ docs/_build/html --open-browser --watch $(GIT_ROOT)/llama_index/
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/README.md b/llama-index-integrations/llms/llama-index-llms-lmstudio/README.md
new file mode 100644
index 0000000000..d35182fb06
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/README.md
@@ -0,0 +1,32 @@
+# LlamaIndex Llms Integration: Lmstudio
+
+## Usage Steps
+
+1. Open LM Studio App and go to the Local Server Tab
+2. In the Configuration settings, enable Apply Prompt Formatting
+3. Load the model of your choice
+4. Start your server
+
+```python
+from llama_index.llms.lmstudio import LMStudio
+
+llm = LMStudio(
+    model_name="Hermes-2-Pro-Llama-3-8B",
+    base_url="http://localhost:1234/v1",
+    temperature=0.7,
+)
+
+messages = [
+    ChatMessage(
+        role=MessageRole.SYSTEM,
+        content="You an expert AI assistant. Help User with their queries.",
+    ),
+    ChatMessage(
+        role=MessageRole.USER,
+        content="What is the significance of the number 42?",
+    ),
+]
+
+response = llm.chat(messages=messages)
+print(str(response))
+```
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/BUILD b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/BUILD
new file mode 100644
index 0000000000..db46e8d6c9
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/BUILD
@@ -0,0 +1 @@
+python_sources()
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/__init__.py b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/__init__.py
new file mode 100644
index 0000000000..30f1f67d0e
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/__init__.py
@@ -0,0 +1,3 @@
+from llama_index.llms.lmstudio.base import LMStudio
+
+__all__ = ["LMStudio"]
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/base.py b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/base.py
new file mode 100644
index 0000000000..96ab119fcf
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/base.py
@@ -0,0 +1,245 @@
+import logging
+import json
+import httpx
+from httpx import Timeout
+from typing import Any, Dict, Sequence, Tuple
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
+from llama_index.core.llms.custom import CustomLLM
+
+from llama_index.core.base.llms.generic_utils import (
+    stream_chat_to_completion_decorator,
+    chat_to_completion_decorator,
+    achat_to_completion_decorator,
+    astream_chat_to_completion_decorator,
+)
+from llama_index.core.constants import (
+    DEFAULT_TEMPERATURE,
+    DEFAULT_CONTEXT_WINDOW,
+    DEFAULT_NUM_OUTPUTS,
+)
+from llama_index.core.base.llms.types import (
+    ChatMessage,
+    ChatResponse,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseAsyncGen,
+    CompletionResponseGen,
+    LLMMetadata,
+    MessageRole,
+)
+
+
+def get_additional_kwargs(
+    response: Dict[str, Any], exclude: Tuple[str, ...]
+) -> Dict[str, Any]:
+    return {k: v for k, v in response.items() if k not in exclude}
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REQUEST_TIMEOUT = 30.0
+
+
+class LMStudio(CustomLLM):
+    base_url: str = Field(
+        default="http://localhost:1234/v1",
+        description="Base url the model is hosted under.",
+    )
+
+    context_window: int = Field(
+        default=DEFAULT_CONTEXT_WINDOW,
+        description="The maximum number of context tokens for the model.",
+        gt=0,
+    )
+
+    model_name: str = Field(description="The model to use.")
+
+    request_timeout: float = Field(
+        default=DEFAULT_REQUEST_TIMEOUT,
+        description="The timeout for making http request in seconds to LM Studio API server.",
+    )
+    num_output: int = Field(
+        default=DEFAULT_NUM_OUTPUTS,
+        description=LLMMetadata.__fields__["num_output"].field_info.description,
+    )
+
+    is_chat_model: bool = Field(
+        default=True,
+        description=(
+            "LM Studio API supports chat."
+            + LLMMetadata.__fields__["is_chat_model"].field_info.description
+        ),
+    )
+
+    temperature: float = Field(
+        default=DEFAULT_TEMPERATURE,
+        description=("The temperature to use for sampling."),
+        gte=0.0,
+        lte=1.0,
+    )
+
+    timeout: float = Field(
+        default=120, description=("The timeout to use in seconds."), gte=0
+    )
+
+    additional_kwargs: Dict[str, Any] = Field(
+        default_factory=dict, description=("Additional kwargs to pass to the model.")
+    )
+
+    def __create_payload_from_messages(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> Dict[str, Any]:
+        return {
+            "model": self.model_name,
+            "messages": [
+                {
+                    "role": message.role.value,
+                    "content": message.content,
+                    **(
+                        message.additional_kwargs
+                        if message.additional_kwargs is not None
+                        else {}
+                    ),
+                }
+                for message in messages
+            ],
+            "options": self._model_kwargs,
+            "stream": False,
+            **kwargs,
+        }
+
+    def __create_chat_response_from_http_response(
+        self, response: httpx.Response
+    ) -> ChatResponse:
+        raw = response.json()
+        message = raw["choices"][0]["message"]
+        return ChatResponse(
+            message=ChatMessage(
+                content=message.get("content"),
+                role=MessageRole(message.get("role")),
+                additional_kwargs=get_additional_kwargs(message, ("content", "role")),
+            ),
+            raw=raw,
+            additional_kwargs=get_additional_kwargs(raw, ("choices",)),
+        )
+
+    @llm_chat_callback()
+    def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse:
+        payload = self.__create_payload_from_messages(messages, **kwargs)
+        with httpx.Client(timeout=Timeout(self.request_timeout)) as client:
+            response = client.post(
+                url=f"{self.base_url}/chat/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            return self.__create_chat_response_from_http_response(response)
+
+    @llm_chat_callback()
+    async def achat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponse:
+        payload = self.__create_payload_from_messages(messages, **kwargs)
+        async with httpx.AsyncClient(timeout=Timeout(self.request_timeout)) as client:
+            response = await client.post(
+                url=f"{self.base_url}/chat/completions",
+                json=payload,
+            )
+            response.raise_for_status()
+            return self.__create_chat_response_from_http_response(response)
+
+    @llm_completion_callback()
+    def complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
+        complete_fn = chat_to_completion_decorator(self.chat)
+        return complete_fn(prompt, **kwargs)
+
+    @llm_completion_callback()
+    async def acomplete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponse:
+        acomplete_fn = achat_to_completion_decorator(self.achat)
+        return await acomplete_fn(prompt, **kwargs)
+
+    @llm_chat_callback()
+    def stream_chat(
+        self, messages: Sequence[ChatMessage], **kwargs: Any
+    ) -> ChatResponseGen:
+        payload = self.__create_payload_from_messages(messages, stream=True, **kwargs)
+
+        with httpx.Client(timeout=Timeout(self.request_timeout)) as client:
+            with client.stream(
+                method="POST",
+                url=f"{self.base_url}/chat/completions",
+                json=payload,
+            ) as response:
+                response.raise_for_status()
+                text = ""
+                for line in response.iter_lines():
+                    if line:
+                        line = line.strip()
+                        if isinstance(line, bytes):
+                            line = line.decode("utf-8")
+
+                        if line.startswith("data: [DONE]"):
+                            break
+
+                        # Slice the line to remove the "data: " prefix
+                        chunk = json.loads(line[5:])
+
+                        delta = chunk["choices"][0].get("delta")
+
+                        role = delta.get("role") or MessageRole.ASSISTANT
+                        content_delta = delta.get("content") or ""
+                        text += content_delta
+
+                        yield ChatResponse(
+                            message=ChatMessage(
+                                content=text,
+                                role=MessageRole(role),
+                                additional_kwargs=get_additional_kwargs(
+                                    chunk, ("choices",)
+                                ),
+                            ),
+                            delta=content_delta,
+                            raw=chunk,
+                            additional_kwargs=get_additional_kwargs(
+                                chunk, ("choices",)
+                            ),
+                        )
+
+    @llm_completion_callback()
+    def stream_complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponseGen:
+        stream_complete_fn = stream_chat_to_completion_decorator(self.stream_chat)
+        return stream_complete_fn(prompt, **kwargs)
+
+    @llm_completion_callback()
+    async def astream_complete(
+        self, prompt: str, formatted: bool = False, **kwargs: Any
+    ) -> CompletionResponseAsyncGen:
+        astream_complete_fn = astream_chat_to_completion_decorator(self.astream_chat)
+        return await astream_complete_fn(prompt, **kwargs)
+
+    @property
+    def metadata(self) -> LLMMetadata:
+        """LLM metadata."""
+        return LLMMetadata(
+            context_window=self.context_window,
+            num_output=self.num_output,
+            model_name=self.model_name,
+            is_chat_model=self.is_chat_model,
+        )
+
+    @property
+    def _model_kwargs(self) -> Dict[str, Any]:
+        base_kwargs = {
+            "temperature": self.temperature,
+            "num_ctx": self.context_window,
+        }
+        return {
+            **base_kwargs,
+            **self.additional_kwargs,
+        }
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-lmstudio/pyproject.toml
new file mode 100644
index 0000000000..d6c4be89a6
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/pyproject.toml
@@ -0,0 +1,53 @@
+[build-system]
+build-backend = "poetry.core.masonry.api"
+requires = ["poetry-core"]
+
+[tool.codespell]
+check-filenames = true
+check-hidden = true
+skip = "*.csv,*.html,*.json,*.jsonl,*.pdf,*.txt,*.ipynb"
+
+[tool.llamahub]
+contains_example = false
+import_path = "llama_index.llms.lmstudio"
+
+[tool.llamahub.class_authors]
+LMStudio = "shivamklr"
+
+[tool.mypy]
+disallow_untyped_defs = true
+exclude = ["_static", "build", "examples", "notebooks", "venv"]
+ignore_missing_imports = true
+python_version = "3.8"
+
+[tool.poetry]
+authors = ["Shivam Kalra <dev@kaltq.com>"]
+description = "llama-index llms lmstudio integration"
+license = "MIT"
+name = "llama-index-llms-lmstudio"
+packages = [{include = "llama_index/"}]
+readme = "README.md"
+version = "0.1.0"
+
+[tool.poetry.dependencies]
+python = ">=3.8.1,<4.0"
+llama-index-core = "^0.10.0"
+
+[tool.poetry.group.dev.dependencies]
+black = {extras = ["jupyter"], version = "<=23.9.1,>=23.7.0"}
+codespell = {extras = ["toml"], version = ">=v2.2.6"}
+ipython = "8.10.0"
+jupyter = "^1.0.0"
+mypy = "0.991"
+pre-commit = "3.2.0"
+pylint = "2.15.10"
+pytest = "7.2.1"
+pytest-mock = "3.11.1"
+ruff = "0.0.292"
+tree-sitter-languages = "^1.8.0"
+types-Deprecated = ">=0.1.0"
+types-PyYAML = "^6.0.12.12"
+types-protobuf = "^4.24.0.4"
+types-redis = "4.5.5.0"
+types-requests = "2.28.11.8"  # TODO: unpin when mypy>0.991
+types-setuptools = "67.1.0.0"
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/BUILD b/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/BUILD
new file mode 100644
index 0000000000..dabf212d7e
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/BUILD
@@ -0,0 +1 @@
+python_tests()
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/__init__.py b/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/test_llms_lmstudio.py b/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/test_llms_lmstudio.py
new file mode 100644
index 0000000000..46897f21ec
--- /dev/null
+++ b/llama-index-integrations/llms/llama-index-llms-lmstudio/tests/test_llms_lmstudio.py
@@ -0,0 +1,7 @@
+from llama_index.core.base.llms.base import BaseLLM
+from llama_index.llms.lmstudio import LMStudio
+
+
+def test_embedding_class():
+    names_of_base_classes = [b.__name__ for b in LMStudio.__mro__]
+    assert BaseLLM.__name__ in names_of_base_classes


Code review comments:
diff_hunk: @@ -0,0 +1,245 @@
+import logging
+import json
+import httpx
+from httpx import Timeout
+from typing import Any, Dict, Sequence, Tuple
+from llama_index.core.bridge.pydantic import Field
+from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback
+from llama_index.core.llms.custom import CustomLLM
+
+from llama_index.core.base.llms.generic_utils import (
+    stream_chat_to_completion_decorator,
+    chat_to_completion_decorator,
+    achat_to_completion_decorator,
+    astream_chat_to_completion_decorator,
+)
+from llama_index.core.constants import (
+    DEFAULT_TEMPERATURE,
+    DEFAULT_CONTEXT_WINDOW,
+    DEFAULT_NUM_OUTPUTS,
+)
+from llama_index.core.base.llms.types import (
+    ChatMessage,
+    ChatResponse,
+    ChatResponseGen,
+    CompletionResponse,
+    CompletionResponseAsyncGen,
+    CompletionResponseGen,
+    LLMMetadata,
+    MessageRole,
+)
+
+
+def get_additional_kwargs(
+    response: Dict[str, Any], exclude: Tuple[str, ...]
+) -> Dict[str, Any]:
+    return {k: v for k, v in response.items() if k not in exclude}
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_REQUEST_TIMEOUT = 30.0
+
+
+class LMStudio(CustomLLM):
+    base_url: str = Field(
+        default="http://localhost:1234/v1",
+        description="Base url the model is hosted under.",
+    )
+
+    context_window: int = Field(
+        default=DEFAULT_CONTEXT_WINDOW,
+        description="The maximum number of context tokens for the model.",
+        gt=0,
+    )
+
+    model_name: str = Field(description="The model to use.")
+
+    request_timeout: float = Field(
+        default=DEFAULT_REQUEST_TIMEOUT,
+        description="The timeout for making http request in seconds to LM Studio API server.",
+    )
+    num_output: int = Field(
+        default=DEFAULT_NUM_OUTPUTS,
+        description=LLMMetadata.__fields__["num_output"].field_info.description,
+    )
+
+    is_chat_model: bool = Field(
+        default=True,
+        description=(
+            "LM Studio API supports chat."
+            + LLMMetadata.__fields__["is_chat_model"].field_info.description
+        ),
+    )
+
+    temperature: float = Field(
+        default=DEFAULT_TEMPERATURE,
+        description=("The temperature to use for sampling."),
+        gte=0.0,
+        lte=1.0,
+    )
+
+    timeout: float = Field(
+        default=120, description=("The timeout to use in seconds."), gte=0
+    )
+
+    additional_kwargs: Dict[str, Any] = Field(
+        default_factory=dict, description=("Additional kwargs to pass to the model.")
+    )
+
+    def __create_payload_from_messages(
path: llama-index-integrations/llms/llama-index-llms-lmstudio/llama_index/llms/lmstudio/base.py
text: @user1:
nit: lets just use single under-score for private methods

@author:
👍 Fixed the typo