feat: specify text embedding dim

[aaron-ang]

Changes Made

Related Issues

Closes #5555.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR implements support for specifying custom embedding dimensions for both LM Studio and Transformers text embedding providers, addressing issue #5555.

Key Changes:

• Removed deprecation warnings about ignored dimensions parameter
• Added dimensions field to both LMStudioTextEmbedderDescriptor and TransformersTextEmbedderDescriptor
• Implemented validation in __post_init__() to ensure dimensions don't exceed model capabilities
• LM Studio: validates against OpenAI model profiles and sets supports_overriding_dimensions flag
• Transformers: validates dimensions against model's hidden_size using AutoConfig.from_pretrained()
• Updated get_dimensions() to return specified dimensions when provided
• Transformers embedder now uses truncate_dim parameter in model.encode()
• LM Studio passes dimensions to OpenAI API using omit type when not supported

Issues Found:

• Edge case in LM Studio where None could be passed instead of omit if supports_overriding_dimensions is manually set

Confidence Score: 4/5

• This PR is mostly safe to merge with one edge case to address
• The implementation is well-structured and handles the main use cases correctly. Validation logic prevents invalid dimensions from being set. However, there's an edge case in LM Studio's instantiate() method where None could be passed instead of omit if embed_options is manually configured.
• Pay attention to daft/ai/lm_studio/protocols/text_embedder.py for the None vs omit edge case

Important Files Changed

Filename	Overview
daft/ai/lm_studio/protocols/text_embedder.py	Added dimensions parameter support with validation and OpenAI API integration, edge case with None handling in instantiate()
daft/ai/transformers/protocols/text_embedder.py	Added dimensions parameter with validation using AutoConfig, passes truncate_dim to model.encode()

Sequence Diagram

sequenceDiagram
    participant User
    participant Provider
    participant Descriptor
    participant Embedder
    participant API

    User->>Provider: get_text_embedder(model, dimensions)
    Provider->>Descriptor: Create descriptor with dimensions
    Descriptor->>Descriptor: __post_init__() - validate dimensions
    alt LM Studio
        Descriptor->>API: Check if model in _models
        alt Model supports overriding
            Descriptor->>Descriptor: Set supports_overriding_dimensions=True
        else Model doesn't support
            Descriptor-->>User: ValueError
        end
    else Transformers
        Descriptor->>API: AutoConfig.from_pretrained()
        API-->>Descriptor: hidden_size
        alt dimensions > hidden_size
            Descriptor-->>User: ValueError
        end
    end
    User->>Descriptor: get_dimensions()
    alt dimensions specified
        Descriptor-->>User: Return specified dimensions
    else dimensions not specified
        alt LM Studio
            Descriptor->>API: embeddings.create("dimension probe")
            API-->>Descriptor: embedding vector
            Descriptor-->>User: Return detected dimensions
        else Transformers
            Descriptor->>API: AutoConfig.from_pretrained()
            API-->>Descriptor: hidden_size
            Descriptor-->>User: Return model dimensions
        end
    end
    User->>Descriptor: instantiate()
    Descriptor->>Embedder: Create embedder
    alt LM Studio
        Embedder->>API: Pass dimensions if supports_overriding_dimensions else omit
    else Transformers
        Embedder->>API: Pass truncate_dim=dimensions
    end

Loading

[greptile-apps]

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

Model download happens in __post_init__ during validation

Calling AutoConfig.from_pretrained() in __post_init__ will download the model config every time a descriptor is created, even when validation isn't needed (e.g., when dimensions is None). This causes unnecessary network calls and delays during normal usage.

Move the model config fetch to get_dimensions() method where it's already called, or cache it after first use.

Prompt To Fix With AI

This is a comment left during a code review.
Path: daft/ai/transformers/protocols/text_embedder.py
Line: 36:36

Comment:
Model download happens in `__post_init__` during validation

Calling `AutoConfig.from_pretrained()` in `__post_init__` will download the model config every time a descriptor is created, even when validation isn't needed (e.g., when `dimensions` is `None`). This causes unnecessary network calls and delays during normal usage.

Move the model config fetch to `get_dimensions()` method where it's already called, or cache it after first use.

How can I resolve this? If you propose a fix, please make it concise.

[aaron-ang]

AutoConfig.from_pretrained() only runs when dimensions is not None. The early return in __post_init__ prevents any model config lookup when no validation is needed.

[codecov]

Codecov Report

❌ Patch coverage is 58.62069% with 12 lines in your changes missing coverage. Please review.
✅ Project coverage is 43.41%. Comparing base (aa8add2) to head (5deae60).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
daft/ai/transformers/protocols/text_embedder.py	53.33%	7 Missing ⚠️
daft/ai/lm_studio/protocols/text_embedder.py	61.53%	5 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #6097       +/-   ##
===========================================
- Coverage   72.91%   43.41%   -29.50%     
===========================================
  Files         973      909       -64     
  Lines      126196   112757    -13439     
===========================================
- Hits        92016    48956    -43060     
- Misses      34180    63801    +29621     

Files with missing lines	Coverage Δ
daft/ai/lm_studio/provider.py	85.00% <ø> (+2.39%)	⬆️
daft/ai/transformers/provider.py	84.44% <100.00%> (+1.11%)	⬆️
daft/ai/lm_studio/protocols/text_embedder.py	76.47% <61.53%> (-5.59%)	⬇️
daft/ai/transformers/protocols/text_embedder.py	75.86% <53.33%> (-6.75%)	⬇️

... and 653 files with indirect coverage changes

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[aaron-ang]

@greptile-apps re-review.

[greptile-apps]

_{2 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}