feat(embed_text): Support LM Studio as a provider #5103
Conversation
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
There was a problem hiding this comment.
Greptile Summary
This PR adds LM Studio as a new text embedding provider to the Daft AI framework. LM Studio is a local AI server that provides an OpenAI-compatible API for running embedding models locally, making it an attractive option for users who want to avoid external API calls or run custom models.
The implementation cleverly extends the existing OpenAI provider architecture since LM Studio maintains API compatibility with OpenAI. The key changes include:
- New LMStudioProvider class in
daft/ai/openai/__init__.pythat inherits fromOpenAIProviderbut configures defaults for local usage (localhost:1234/v1) and sets a dummy API key since LM Studio doesn't require authentication - LMStudioTextEmbedderDescriptor that dynamically discovers embedding dimensions by making a probe request to the server, which is necessary because LM Studio can load arbitrary models with varying dimensions
- Provider registration in
daft/ai/provider.pythat adds 'lm_studio' to the available providers list - Comprehensive test coverage in
tests/ai/test_lm_studio.pythat validates the provider functionality with proper mocking
The base_url handling includes smart logic to automatically append '/v1' if not present, ensuring compatibility with LM Studio's expected endpoint format. This implementation maintains the same interface as other providers while accommodating LM Studio's unique characteristics of local deployment and variable model dimensions.
Confidence score: 3/5
- This PR introduces architectural complexity with the dynamic dimension discovery mechanism that could fail in various scenarios
- Score reflects concerns about network calls during descriptor initialization and potential reliability issues with the probe request approach
- Pay close attention to the dimension discovery logic in
LMStudioTextEmbedderDescriptor.get_dimensions()method
4 files reviewed, 3 comments
| def get_dimensions(self) -> EmbeddingDimensions: | ||
| try: | ||
| client = OpenAI(**self.provider_options) | ||
| response = client.embeddings.create( | ||
| input="dimension probe", | ||
| model=self.model_name, | ||
| encoding_format="float", | ||
| ) | ||
| size = len(response.data[0].embedding) | ||
| return EmbeddingDimensions(size=size, dtype=DataType.float32()) | ||
| except Exception as ex: | ||
| raise ValueError("Failed to determine embedding dimensions from LM Studio.") from ex |
There was a problem hiding this comment.
style: The dimension probing creates a new OpenAI client and makes a network request during descriptor creation. This could be expensive if called repeatedly and may fail if the LM Studio server is temporarily unavailable. Consider caching the dimensions or moving this logic to instantiation time.
There was a problem hiding this comment.
This is a local network request
| except Exception as ex: | ||
| raise ValueError("Failed to determine embedding dimensions from LM Studio.") from ex |
There was a problem hiding this comment.
style: The broad exception catch could mask specific connection errors. Consider catching more specific exceptions like OpenAIError or connection-related exceptions to provide better error messages.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #5103 +/- ##
==========================================
- Coverage 76.10% 76.07% -0.04%
==========================================
Files 950 953 +3
Lines 130488 130628 +140
==========================================
+ Hits 99310 99377 +67
- Misses 31178 31251 +73
🚀 New features to boost your workflow:
|
desmondcheongzx
changed the title
## Changes Made
Add LM Studio as a text embedding provider.
For example
```
import daft
from daft.ai.provider import load_provider
from daft.functions.ai import embed_text
provider = load_provider("lm_studio", base_url="http://127.0.0.1:1234") # This base_url parameter is optional if you're using the defaults for LM Studio. You can modify this as needed.
model = "text-embedding-nomic-embed-text-v1.5" # Select a text embedding model that you've loaded into LM Studio.
(
daft.read_huggingface("Open-Orca/OpenOrca")
.with_column("embedding", embed_text(daft.col("response"), provider=provider, model=model))
.show()
)
```
Changes Made
Add LM Studio as a text embedding provider.
For example