Commit 1ef9cba
authored
Feature/prompt templates and lmstudio sdk (#171)
* Add prompt template support and LM Studio SDK integration
Features:
- Prompt template support for embedding models (via --embedding-prompt-template)
- LM Studio SDK integration for automatic context length detection
- Hybrid token limit discovery (Ollama → LM Studio → Registry → Default)
- Client-side token truncation to prevent silent failures
- Automatic persistence of embedding_options to .meta.json
Implementation:
- Added _query_lmstudio_context_limit() with Node.js subprocess bridge
- Modified compute_embeddings_openai() to apply prompt templates before truncation
- Extended CLI with --embedding-prompt-template flag for build and search
- URL detection for LM Studio (port 1234 or lmstudio/lm.studio keywords)
- HTTP→WebSocket URL conversion for SDK compatibility
Tests:
- 60 passing tests across 5 test files
- Comprehensive coverage of prompt templates, LM Studio integration, and token handling
- Parametrized tests for maintainability and clarity
* Add integration tests and fix LM Studio SDK bridge
Features:
- End-to-end integration tests for prompt template with EmbeddingGemma
- Integration tests for hybrid token limit discovery mechanism
- Tests verify real-world functionality with live services (LM Studio, Ollama)
Fixes:
- LM Studio SDK bridge now uses client.embedding.load() for embedding models
- Fixed NODE_PATH resolution to include npm global modules
- Fixed integration test to use WebSocket URL (ws://) for SDK bridge
Tests:
- test_prompt_template_e2e.py: 8 integration tests covering:
- Prompt template prepending with LM Studio (EmbeddingGemma)
- LM Studio SDK bridge for context length detection
- Ollama dynamic token limit detection
- Hybrid discovery fallback mechanism (registry, default)
- All tests marked with @pytest.mark.integration for selective execution
- Tests gracefully skip when services unavailable
Documentation:
- Updated tests/README.md with integration test section
- Added prerequisites and running instructions
- Documented that prompt templates are ONLY for EmbeddingGemma
- Added integration marker to pyproject.toml
Test Results:
- All 8 integration tests passing with live services
- Confirmed prompt templates work correctly with EmbeddingGemma
- Verified LM Studio SDK bridge auto-detects context length (2048)
- Validated hybrid token limit discovery across all backends
* Add prompt template support to Ollama mode
Extends prompt template functionality from OpenAI mode to Ollama for backend consistency.
Changes:
- Add provider_options parameter to compute_embeddings_ollama()
- Apply prompt template before token truncation (lines 1005-1011)
- Pass provider_options through compute_embeddings() call chain
Tests:
- test_ollama_embedding_with_prompt_template: Verifies templates work with Ollama
- test_ollama_prompt_template_affects_embeddings: Confirms embeddings differ with/without template
- Both tests pass with live Ollama service (2/2 passing)
Usage:
leann build --embedding-mode ollama --embedding-prompt-template "query: " ...
* Fix LM Studio SDK bridge to respect JIT auto-evict settings
Problem: SDK bridge called client.embedding.load() which loaded models into
LM Studio memory and bypassed JIT auto-evict settings, causing duplicate
model instances to accumulate.
Root cause analysis (from Perplexity research):
- Explicit SDK load() commands are treated as "pinned" models
- JIT auto-evict only applies to models loaded reactively via API requests
- SDK-loaded models remain in memory until explicitly unloaded
Solutions implemented:
1. Add model.unload() after metadata query (line 243)
- Load model temporarily to get context length
- Unload immediately to hand control back to JIT system
- Subsequent API requests trigger JIT load with auto-evict
2. Add token limit caching to prevent repeated SDK calls
- Cache discovered limits in _token_limit_cache dict (line 48)
- Key: (model_name, base_url), Value: token_limit
- Prevents duplicate load/unload cycles within same process
- Cache shared across all discovery methods (Ollama, SDK, registry)
Tests:
- TestTokenLimitCaching: 5 tests for cache behavior (integrated into test_token_truncation.py)
- Manual testing confirmed no duplicate models in LM Studio after fix
- All existing tests pass
Impact:
- Respects user's LM Studio JIT and auto-evict settings
- Reduces model memory footprint
- Faster subsequent builds (cached limits)
* Document prompt template and LM Studio SDK features
Added comprehensive documentation for new optional embedding features:
Configuration Guide (docs/configuration-guide.md):
- New section: "Optional Embedding Features"
- Task-Specific Prompt Templates subsection:
- Explains EmbeddingGemma use case with document/query prompts
- CLI and Python API examples
- Clear warnings about compatible vs incompatible models
- References to GitHub issue #155 and HuggingFace blog
- LM Studio Auto-Detection subsection:
- Prerequisites (Node.js + @lmstudio/sdk)
- How auto-detection works (4-step process)
- Benefits and optional nature clearly stated
FAQ (docs/faq.md):
- FAQ #2: When should I use prompt templates?
- DO/DON'T guidance with examples
- Links to detailed configuration guide
- FAQ #3: Why is LM Studio loading multiple copies?
- Explains the JIT auto-evict fix
- Troubleshooting steps if still seeing issues
- FAQ #4: Do I need Node.js and @lmstudio/sdk?
- Clarifies it's completely optional
- Lists benefits if installed
- Installation instructions
Cross-references between documents for easy navigation between quick reference and detailed guides.
* Add separate build/query template support for task-specific models
Task-specific models like EmbeddingGemma require different templates for indexing vs searching. Store both templates at build time and auto-apply query template during search with backward compatibility.
* Consolidate prompt template tests from 44 to 37 tests
Merged redundant no-op tests, removed low-value implementation tests, consolidated parameterized CLI tests, and removed hanging over-mocked test. All tests pass with improved focus on behavioral testing.
* Fix query template application in compute_query_embedding
Query templates were only applied in the fallback code path, not when using the embedding server (default path). This meant stored query templates in index metadata were ignored during MCP and CLI searches.
Changes:
- Move template application to before any computation path (searcher_base.py:109-110)
- Add comprehensive tests for both server and fallback paths
- Consolidate tests into test_prompt_template_persistence.py
Tests verify:
- Template applied when using embedding server
- Template applied in fallback path
- Consistent behavior between both paths
* Apply ruff formatting and fix linting issues
- Remove unused imports
- Fix import ordering
- Remove unused variables
- Apply code formatting
* Fix CI test failures: mock OPENAI_API_KEY in tests
Tests were failing in CI because compute_embeddings_openai() checks for OPENAI_API_KEY before using the mocked client. Added monkeypatch to set fake API key in test fixture.1 parent a635509 commit 1ef9cba
File tree
15 files changed
+3095
-15
lines changed- docs
- packages/leann-core/src/leann
- tests
15 files changed
+3095
-15
lines changed| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
158 | 158 | | |
159 | 159 | | |
160 | 160 | | |
| 161 | + | |
| 162 | + | |
| 163 | + | |
| 164 | + | |
| 165 | + | |
| 166 | + | |
| 167 | + | |
| 168 | + | |
| 169 | + | |
| 170 | + | |
| 171 | + | |
| 172 | + | |
| 173 | + | |
| 174 | + | |
| 175 | + | |
| 176 | + | |
| 177 | + | |
| 178 | + | |
| 179 | + | |
| 180 | + | |
| 181 | + | |
| 182 | + | |
| 183 | + | |
| 184 | + | |
| 185 | + | |
| 186 | + | |
| 187 | + | |
| 188 | + | |
| 189 | + | |
| 190 | + | |
| 191 | + | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
| 195 | + | |
| 196 | + | |
| 197 | + | |
| 198 | + | |
| 199 | + | |
| 200 | + | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
| 209 | + | |
| 210 | + | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
| 216 | + | |
| 217 | + | |
| 218 | + | |
| 219 | + | |
| 220 | + | |
| 221 | + | |
| 222 | + | |
| 223 | + | |
| 224 | + | |
| 225 | + | |
| 226 | + | |
| 227 | + | |
| 228 | + | |
| 229 | + | |
| 230 | + | |
| 231 | + | |
| 232 | + | |
| 233 | + | |
| 234 | + | |
| 235 | + | |
| 236 | + | |
| 237 | + | |
| 238 | + | |
| 239 | + | |
| 240 | + | |
| 241 | + | |
| 242 | + | |
| 243 | + | |
| 244 | + | |
| 245 | + | |
| 246 | + | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
161 | 250 | | |
162 | 251 | | |
163 | 252 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
8 | 8 | | |
9 | 9 | | |
10 | 10 | | |
| 11 | + | |
| 12 | + | |
| 13 | + | |
| 14 | + | |
| 15 | + | |
| 16 | + | |
| 17 | + | |
| 18 | + | |
| 19 | + | |
| 20 | + | |
| 21 | + | |
| 22 | + | |
| 23 | + | |
| 24 | + | |
| 25 | + | |
| 26 | + | |
| 27 | + | |
| 28 | + | |
| 29 | + | |
| 30 | + | |
| 31 | + | |
| 32 | + | |
| 33 | + | |
| 34 | + | |
| 35 | + | |
| 36 | + | |
| 37 | + | |
| 38 | + | |
| 39 | + | |
| 40 | + | |
| 41 | + | |
| 42 | + | |
| 43 | + | |
| 44 | + | |
| 45 | + | |
| 46 | + | |
| 47 | + | |
| 48 | + | |
| 49 | + | |
| 50 | + | |
| 51 | + | |
| 52 | + | |
| 53 | + | |
| 54 | + | |
| 55 | + | |
| 56 | + | |
| 57 | + | |
| 58 | + | |
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
916 | 916 | | |
917 | 917 | | |
918 | 918 | | |
| 919 | + | |
919 | 920 | | |
920 | 921 | | |
921 | 922 | | |
| |||
979 | 980 | | |
980 | 981 | | |
981 | 982 | | |
| 983 | + | |
| 984 | + | |
| 985 | + | |
| 986 | + | |
| 987 | + | |
| 988 | + | |
| 989 | + | |
| 990 | + | |
| 991 | + | |
| 992 | + | |
| 993 | + | |
| 994 | + | |
| 995 | + | |
982 | 996 | | |
983 | 997 | | |
984 | 998 | | |
985 | 999 | | |
| 1000 | + | |
986 | 1001 | | |
987 | 1002 | | |
988 | 1003 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
144 | 144 | | |
145 | 145 | | |
146 | 146 | | |
| 147 | + | |
| 148 | + | |
| 149 | + | |
| 150 | + | |
| 151 | + | |
| 152 | + | |
| 153 | + | |
| 154 | + | |
| 155 | + | |
| 156 | + | |
| 157 | + | |
| 158 | + | |
147 | 159 | | |
148 | 160 | | |
149 | 161 | | |
| |||
260 | 272 | | |
261 | 273 | | |
262 | 274 | | |
| 275 | + | |
| 276 | + | |
| 277 | + | |
| 278 | + | |
| 279 | + | |
| 280 | + | |
263 | 281 | | |
264 | 282 | | |
265 | 283 | | |
| |||
1398 | 1416 | | |
1399 | 1417 | | |
1400 | 1418 | | |
| 1419 | + | |
| 1420 | + | |
| 1421 | + | |
| 1422 | + | |
| 1423 | + | |
| 1424 | + | |
| 1425 | + | |
| 1426 | + | |
1401 | 1427 | | |
1402 | 1428 | | |
1403 | 1429 | | |
| |||
0 commit comments