feat(wren-ai-service): add litellm embedder #1247
Conversation
cyyeh
added
module/ai-service
WalkthroughThis pull request updates multiple configuration files to switch embedding providers from the legacy ones to Changes
Sequence Diagram(s)sequenceDiagram
participant U as User
participant P as LitellmEmbedderProvider
participant T as AsyncTextEmbedder
participant API as Embedding API
U->>P: Request text embedder instance
P->>T: Return AsyncTextEmbedder
U->>T: Invoke run(text)
T->>API: Send text embedding request
API-->>T: Return embedding result
T-->>U: Deliver embedding data with metadata
sequenceDiagram
participant U as User
participant P as LitellmEmbedderProvider
participant D as AsyncDocumentEmbedder
participant API as Embedding API
U->>P: Request document embedder instance
P->>D: Return AsyncDocumentEmbedder
U->>D: Invoke run(documents, batch_size)
loop Batch processing
D->>API: Send batch embedding request
API-->>D: Return batch embedding result
end
D-->>U: Provide complete document embeddings
Possibly related PRs
Suggested reviewers
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (Invoked using PR comments)
Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (10)
wren-ai-service/src/providers/embedder/litellm.py (8)
1-3: Consider consolidating or reorganizing imports.
Imports from standard libraries (logging,os) and third-party libraries (openai,backoff,haystack,litellm, etc.) can be grouped by convention, improving readability.
5-9: Ensure consistent error-handling strategy.
Usingbackoffis great for resiliency, but confirm that other exceptions (e.g., network interruptions, HTTP 4xx errors) are also addressed. If not, you may consider broadening your exception handling strategy or logging additional diagnostic info.
18-30: Add optional multi-field metadata embedding support.
At present,_prepare_texts_to_embedonly concatenatesdoc.content. Consider if you’d like to embed additional fields from the metadata in the same pass. This might enhance the embedding by providing more context.
33-48: Clarify or document default parameter values in constructor.
The constructor forAsyncTextEmbedderhas parametersapi_key,model, etc., which can beNone. It might be helpful to include docstrings or default values in the signature to clarify usage, especially if they are required.
79-94: Update docstring for improved clarity.
AsyncDocumentEmbedderdoes not have a docstring explaining its usage or requirements. Adding a brief summary of its purpose and parameters helps future contributors.
95-125: Guard against missing usage keys in_embed_batch.
When appending usage statistics, the code assumesprompt_tokensandtotal_tokensalways exist. Future changes to the provider or partial usage returns could break this. Consider a safe retrieval pattern, e.g.meta["usage"].get("prompt_tokens", 0).- meta["usage"]["prompt_tokens"] += response.usage.prompt_tokens - meta["usage"]["total_tokens"] += response.usage.total_tokens + meta["usage"]["prompt_tokens"] = meta["usage"].get("prompt_tokens", 0) + getattr(response.usage, "prompt_tokens", 0) + meta["usage"]["total_tokens"] = meta["usage"].get("total_tokens", 0) + getattr(response.usage, "total_tokens", 0)
127-154: Optionally allow chunking for large documents.
Embedding large documents, even in batches, can be expensive or lead to hitting model limits. Consider offering an optional chunk-based approach to further split very large documents.
156-197: Consider lazy loading or caching for scaled usage.
Initial logging in the provider is helpful, but if usage grows significantly (e.g., frequent model changes, dynamic environment changes), consider advanced patterns like caching or lazy loading. This can optimize repeated instantiations.wren-ai-service/docs/config_examples/config.ollama.yaml (2)
17-21: Enhance configuration documentation for better user guidance.Consider adding comments to explain:
- The reason for using the
openai/prefix in the model name- Alternative
api_baseconfigurations for non-Docker environments- The implications of the 120-second timeout value
112-112: Remove trailing space.There's a trailing space after
nomic-embed-textthat should be removed.- embedder: litellm_embedder.openai/nomic-embed-text + embedder: litellm_embedder.openai/nomic-embed-text🧰 Tools
🪛 YAMLlint (1.35.1)
[error] 112-112: trailing spaces
(trailing-spaces)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (4)
wren-ai-service/docs/config_examples/config.ollama.yaml(4 hunks)wren-ai-service/src/providers/__init__.py(1 hunks)wren-ai-service/src/providers/embedder/litellm.py(1 hunks)wren-ai-service/tests/pytest/providers/test_loader.py(1 hunks)
🧰 Additional context used
🪛 YAMLlint (1.35.1)
wren-ai-service/docs/config_examples/config.ollama.yaml
[error] 112-112: trailing spaces
(trailing-spaces)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
- GitHub Check: build-image (linux_arm64_runner, linux/arm64)
🔇 Additional comments (4)
wren-ai-service/src/providers/embedder/litellm.py (1)
58-77: Confirmresponse.usageaccess safety.
While the code checksif response.usage, it later references usage keys (e.g.,prompt_tokens) in other parts of the file. Ensure all code paths handle partial or missing usage data gracefully.wren-ai-service/tests/pytest/providers/test_loader.py (1)
35-36: Great addition to test coverage.
Verifying"litellm_embedder"ensures that the new provider is properly loaded. This test helps confirm that the provider is wired up.wren-ai-service/docs/config_examples/config.ollama.yaml (2)
43-43: LGTM! Consistent embedder configuration across all pipelines.The embedder references have been correctly and consistently updated to use the new LiteLLM format throughout all pipeline configurations.
Also applies to: 46-46, 49-49, 53-53, 56-56, 92-92, 103-103, 109-109, 112-112, 115-115
Line range hint
1-141: Verify security considerations in the example configuration.Please ensure:
- The example configuration doesn't expose any sensitive information
- Add a comment warning users about securing their API endpoints
- Consider adding a note about the implications of having
development: truein production environments🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 42-42: wrong indentation: expected 0 but found 2
(indentation)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (6)
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml (2)
16-21: Embedder Section Update & Indentation Warning:
The embedder configuration correctly updates the provider tolitellm_embedderand the model togemini/text-embedding-004. However, YAMLlint reports an indentation warning on line 19 (wrong indentation: expected 0 but found 2). Please review the list-item indentation undermodels:to ensure it meets your YAML style guidelines.🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 19-19: wrong indentation: expected 0 but found 2
(indentation)
Line range hint
43-115: Pipeline Embedder References & Trailing Space Cleanup:
Pipeline entries have been consistently updated to uselitellm_embedder.text-embedding-004. Additionally, static analysis detected trailing spaces on line 112. Removing these extra spaces will clean up the YAML formatting.For example, update line 112 as follows:
- embedder: litellm_embedder.text-embedding-004 + embedder: litellm_embedder.text-embedding-004🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 42-42: wrong indentation: expected 0 but found 2
(indentation)
wren-ai-service/docs/config_examples/config.groq.yaml (2)
17-23: Embedder Block Indentation Issue:
The embedder section is updated correctly with the provider set tolitellm_embedderand modeltext-embedding-3-large. However, YAMLlint reports an indentation warning on line 21. Please adjust the indentation of the list item to conform to YAML standards (e.g., remove unnecessary leading spaces).🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 21-21: wrong indentation: expected 0 but found 2
(indentation)
Line range hint
45-117: Pipeline Embedder References & Trailing Space Removal:
Pipeline entries now correctly referencelitellm_embedder.text-embedding-3-large. Note that a trailing space was detected on line 114; removing it will improve the YAML formatting.For instance, update line 114 as follows:
- embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.text-embedding-3-large🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 44-44: wrong indentation: expected 0 but found 2
(indentation)
wren-ai-service/docs/config_examples/config.deepseek.yaml (2)
33-39: Embedder Configuration Indentation Warning:
The embedder block now useslitellm_embedderwith modeltext-embedding-3-large, which is correct. However, YAMLlint indicates an indentation issue on line 37 for the list item. Please adjust the indentation to meet YAML formatting standards.🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 37-37: wrong indentation: expected 0 but found 2
(indentation)
Line range hint
61-133: Pipeline Embedder Updates & Trailing Space Cleanup:
Pipeline entries consistently uselitellm_embedder.text-embedding-3-large, maintaining alignment with the new provider changes. Additionally, there is a trailing space on line 130. Removing that trailing space will improve file cleanliness.For example, update line 130 as shown:
- embedder: litellm_embedder.text-embedding-3-large + embedder: litellm_embedder.text-embedding-3-large🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 60-60: wrong indentation: expected 0 but found 2
(indentation)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (5)
wren-ai-service/docs/config_examples/config.deepseek.yaml(4 hunks)wren-ai-service/docs/config_examples/config.google_ai_studio.yaml(4 hunks)wren-ai-service/docs/config_examples/config.groq.yaml(4 hunks)wren-ai-service/docs/config_examples/config.ollama.yaml(4 hunks)wren-ai-service/tests/pytest/providers/test_loader.py(2 hunks)
🚧 Files skipped from review as they are similar to previous changes (1)
- wren-ai-service/tests/pytest/providers/test_loader.py
🧰 Additional context used
🪛 YAMLlint (1.35.1)
wren-ai-service/docs/config_examples/config.google_ai_studio.yaml
[warning] 19-19: wrong indentation: expected 0 but found 2
(indentation)
[error] 112-112: trailing spaces
(trailing-spaces)
wren-ai-service/docs/config_examples/config.groq.yaml
[warning] 21-21: wrong indentation: expected 0 but found 2
(indentation)
[error] 114-114: trailing spaces
(trailing-spaces)
wren-ai-service/docs/config_examples/config.deepseek.yaml
[warning] 37-37: wrong indentation: expected 0 but found 2
(indentation)
[error] 130-130: trailing spaces
(trailing-spaces)
wren-ai-service/docs/config_examples/config.ollama.yaml
[error] 112-112: trailing spaces
(trailing-spaces)
⏰ Context from checks skipped due to timeout of 90000ms (1)
- GitHub Check: pytest
🔇 Additional comments (2)
wren-ai-service/docs/config_examples/config.ollama.yaml (2)
17-21: Embedder Configuration Update:
The embedder block now useslitellm_embedderwith the updated model valueopenai/nomic-embed-text, which aligns with the new provider requirements.
Line range hint
43-115: Pipeline Embedder References Update:
All pipeline entries have been updated to referencelitellm_embedder.openai/nomic-embed-text, ensuring consistency across configurations. Please verify that the associatedllmand other service references are still aligned with your overall service architecture.🧰 Tools
🪛 YAMLlint (1.35.1)
[warning] 42-42: wrong indentation: expected 0 but found 2
(indentation)
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (2)
wren-ai-service/src/providers/document_store/qdrant.py (2)
244-247: Consider improving the error message specificity.The error message could be more helpful by including the actual value that caused the error.
- msg = f"DocumentStore.write_documents() expects a list of Documents but got an element of {type(doc)}." + msg = f"DocumentStore.write_documents() expects a list of Documents but got an element of type {type(doc)}: {repr(doc)}"
274-289: Consider handling the last batch size in progress bar update.The progress bar update uses a fixed
write_batch_sizewhich might not be accurate for the last batch.- progress_bar.update(self.write_batch_size) + progress_bar.update(len(document_batch))
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/src/providers/document_store/qdrant.py(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: pytest
- GitHub Check: Analyze (javascript-typescript)
- GitHub Check: Analyze (go)
🔇 Additional comments (1)
wren-ai-service/src/providers/document_store/qdrant.py (1)
357-362: LGTM! The addition follows existing patterns.The new line follows the established pattern of initializing document stores for different datasets. The change is well-integrated with the existing code structure.
Let's verify the usage of this new document store:
✅ Verification successful
The "sql_pairs" document store initialization is properly integrated.
The shell script output confirms that the "sql_pairs" dataset is referenced in multiple parts of the codebase (including tests, pipelines, and configuration files), thereby validating the new initialization in the Qdrant provider.
- Confirmed consistent usage across tests (e.g.,
test_sql_pairs_preparation.py,test_sql_pairs_deletion.py) and pipelines.- Configuration files and other modules (e.g., in
src/config.pyandsrc/web/v1/services/sql_pairs_preparation.py) reference "sql_pairs" appropriately.🏁 Scripts executed
The following scripts were executed for the analysis:
Script:
#!/bin/bash # Description: Check for references to the new sql_pairs dataset # Expected: Find references to confirm this dataset is used elsewhere in the codebase # Search for direct references to sql_pairs rg -l "sql_pairs" # Search for any SQL-related dataset configurations rg -l "dataset.*sql"Length of output: 1979
cyyeh
force-pushed
the
feat/ai-service/litellm-embedder
branch
from
ef91a18 to
1f0d299
Compare
issue #1078
Summary by CodeRabbit
New Features
Tests
Chores