diff --git a/CHANGELOG.md b/CHANGELOG.md index 1341c7ec..408159f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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) diff --git a/README.md b/README.md index 4325c308..0de7c536 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index 00f6e7c3..1f2a396c 100644 --- a/docs/configuration/providers.md +++ b/docs/configuration/providers.md @@ -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: diff --git a/haiku_rag_slim/haiku/rag/config/__init__.py b/haiku_rag_slim/haiku/rag/config/__init__.py index 40c18ee0..5caf60fe 100644 --- a/haiku_rag_slim/haiku/rag/config/__init__.py +++ b/haiku_rag_slim/haiku/rag/config/__init__.py @@ -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", diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index dbc3ce11..5ed0663a 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index c8e00fce..70582b49 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/haiku_rag_slim/haiku/rag/embeddings/__init__.py @@ -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}") diff --git a/haiku_rag_slim/haiku/rag/embeddings/lm_studio.py b/haiku_rag_slim/haiku/rag/embeddings/lm_studio.py new file mode 100644 index 00000000..fcde0606 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/embeddings/lm_studio.py @@ -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] diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index eac06a34..8ce14b54 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -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}" diff --git a/tests/test_embedder_config.py b/tests/test_embedder_config.py index 6b0d31fb..64a1be96 100644 --- a/tests/test_embedder_config.py +++ b/tests/test_embedder_config.py @@ -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" )