Support LMStudio

This commit is contained in:
Yiorgis Gozadinos 2025-11-25 15:27:54 +02:00
parent 7876830eca
commit 0303e6a8cd
No known key found for this signature in database
9 changed files with 138 additions and 2 deletions

View file

@ -1,6 +1,13 @@
# Changelog
## [Unreleased]
### Added
- **LM Studio Provider**: Added support for LM Studio as a provider for embeddings and QA/research models
- Configure with `provider: lm_studio` in embeddings, QA, or research model settings
- Supports thinking control for reasoning models (gpt-oss, etc.)
- Default base URL: `http://localhost:1234`
### Fixed
- **Configuration**: Fixed `init-config` command generating invalid configuration files (#165)

View file

@ -7,8 +7,8 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
## Features
- **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
- **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
- **Multiple QA providers**: Any provider/model supported by Pydantic AI
- **Multiple embedding providers**: Ollama, LM Studio, VoyageAI, OpenAI, vLLM
- **Multiple QA providers**: Any provider/model supported by Pydantic AI (Ollama, LM Studio, OpenAI, Anthropic, etc.)
- **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
- **Reranking**: Default search result reranking with MixedBread AI, Cohere, Zero Entropy, or vLLM
- **Question answering**: Built-in QA agents on your documents

View file

@ -59,6 +59,7 @@ See the [Pydantic AI thinking documentation](https://ai.pydantic.dev/thinking/)
- **Bedrock**: Claude, OpenAI, and Qwen models
- **Ollama**: Models supporting reasoning (gpt-oss, etc.)
- **vLLM**: Models supporting reasoning (gpt-oss, etc.)
- **LM Studio**: Models supporting reasoning (gpt-oss, etc.)
**When to use:**
- Disable for simple queries, RAG workflows, speed-critical applications
@ -148,6 +149,23 @@ providers:
**Note:** You need to run a vLLM server separately with an embedding model loaded.
### LM Studio
[LM Studio](https://lmstudio.ai/) provides a local OpenAI-compatible API server for running models:
```yaml
embeddings:
provider: lm_studio
model: text-embedding-qwen3-embedding-4b
vector_dim: 2560
providers:
lm_studio:
base_url: http://localhost:1234
```
**Note:** LM Studio must be running with an embedding model loaded. The default URL is `http://localhost:1234`.
## Question Answering Providers
Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used.
@ -226,6 +244,29 @@ providers:
**Note:** You need to run a vLLM server separately with a model that supports tool calling loaded. Consult the specific model's documentation for proper vLLM serving configuration.
### LM Studio
Use LM Studio for local question answering and research:
```yaml
qa:
model:
provider: lm_studio
name: openai/gpt-oss-20b
enable_thinking: false
research:
model:
provider: lm_studio
name: openai/gpt-oss-20b
providers:
lm_studio:
base_url: http://localhost:1234
```
**Note:** LM Studio must be running with a chat model that supports tool calling loaded.
### Other Providers
Any provider supported by Pydantic AI can be used. Examples:

View file

@ -11,6 +11,7 @@ from haiku.rag.config.models import (
ConversionOptions,
EmbeddingsConfig,
LanceDBConfig,
LMStudioConfig,
MonitorConfig,
OllamaConfig,
ProcessingConfig,
@ -36,6 +37,7 @@ __all__ = [
"ResearchConfig",
"ProcessingConfig",
"OllamaConfig",
"LMStudioConfig",
"VLLMConfig",
"ProvidersConfig",
"find_config_file",

View file

@ -135,9 +135,14 @@ class DoclingServeConfig(BaseModel):
timeout: int = 300
class LMStudioConfig(BaseModel):
base_url: str = "http://localhost:1234"
class ProvidersConfig(BaseModel):
ollama: OllamaConfig = Field(default_factory=OllamaConfig)
vllm: VLLMConfig = Field(default_factory=VLLMConfig)
lm_studio: LMStudioConfig = Field(default_factory=LMStudioConfig)
docling_serve: DoclingServeConfig = Field(default_factory=DoclingServeConfig)

View file

@ -46,4 +46,11 @@ def get_embedder(config: AppConfig = Config) -> EmbedderBase:
config.embeddings.model, config.embeddings.vector_dim, config
)
if config.embeddings.provider == "lm_studio":
from haiku.rag.embeddings.lm_studio import Embedder as LMStudioEmbedder
return LMStudioEmbedder(
config.embeddings.model, config.embeddings.vector_dim, config
)
raise ValueError(f"Unsupported embedding provider: {config.embeddings.provider}")

View file

@ -0,0 +1,28 @@
from typing import overload
from openai import AsyncOpenAI
from haiku.rag.embeddings.base import EmbedderBase
class Embedder(EmbedderBase):
@overload
async def embed(self, text: str) -> list[float]: ...
@overload
async def embed(self, text: list[str]) -> list[list[float]]: ...
async def embed(self, text: str | list[str]) -> list[float] | list[list[float]]:
client = AsyncOpenAI(
base_url=f"{self._config.providers.lm_studio.base_url}/v1", api_key="dummy"
)
if not text:
return []
response = await client.embeddings.create(
model=self._model,
input=text,
)
if isinstance(text, str):
return response.data[0].embedding
else:
return [item.embedding for item in response.data]

View file

@ -233,6 +233,29 @@ def get_model(
settings=vllm_settings,
)
elif provider == "lm_studio":
model_settings = None
# Apply thinking control for gpt-oss
if model == "gpt-oss" and model_config.enable_thinking is not None:
if model_config.enable_thinking is False:
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="low")
else:
model_settings = OpenAIChatModelSettings(openai_reasoning_effort="high")
model_settings = apply_common_settings(
model_settings, OpenAIChatModelSettings, model_config
)
return OpenAIChatModel(
model_name=model,
provider=OpenAIProvider(
base_url=f"{app_config.providers.lm_studio.base_url}/v1",
api_key="dummy",
),
settings=model_settings,
)
else:
# For any other provider, use string format and let Pydantic AI handle it
return f"{provider}:{model}"

View file

@ -3,6 +3,7 @@ import pytest
from haiku.rag.config import (
AppConfig,
EmbeddingsConfig,
LMStudioConfig,
OllamaConfig,
ProvidersConfig,
VLLMConfig,
@ -71,6 +72,28 @@ def test_openai_embedder_uses_config():
assert embedder._config == custom_config
def test_lm_studio_embedder_uses_config():
"""Test that lm_studio embedder uses the config passed to get_embedder."""
custom_config = AppConfig(
embeddings=EmbeddingsConfig(
provider="lm_studio",
model="custom-lm-studio-model",
vector_dim=1024,
),
providers=ProvidersConfig(
lm_studio=LMStudioConfig(base_url="http://custom-lmstudio:5678"),
),
)
embedder = get_embedder(custom_config)
assert embedder._model == "custom-lm-studio-model"
assert embedder._vector_dim == 1024
assert (
embedder._config.providers.lm_studio.base_url == "http://custom-lmstudio:5678"
)
@pytest.mark.skipif(
True, reason="VoyageAI is an optional dependency, may not be installed"
)