chore(wren-ai-service): improve embedding performance and allow setting batch size for embedding #1814
Conversation
cyyeh
added
module/ai-service
WalkthroughThe Changes
Poem
✨ Finishing Touches
🧪 Generate Unit Tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. 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: 2
🧹 Nitpick comments (1)
wren-ai-service/src/providers/embedder/litellm.py (1)
110-113: Fix whitespace formatting issue.Static analysis detected a formatting issue with whitespace before the colon.
Apply this diff to fix the formatting:
- batches = [ - texts_to_embed[i : i + batch_size] - for i in range(0, len(texts_to_embed), batch_size) - ] + batches = [ + texts_to_embed[i:i + batch_size] + for i in range(0, len(texts_to_embed), batch_size) + ]
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
wren-ai-service/src/providers/embedder/litellm.py(4 hunks)
🧰 Additional context used
🪛 Flake8 (7.2.0)
wren-ai-service/src/providers/embedder/litellm.py
[error] 111-111: whitespace before ':'
(E203)
⏰ Context from checks skipped due to timeout of 90000ms (4)
- GitHub Check: pytest
- GitHub Check: pytest
- GitHub Check: Analyze (go)
- GitHub Check: Analyze (javascript-typescript)
🔇 Additional comments (5)
wren-ai-service/src/providers/embedder/litellm.py (5)
1-1: LGTM: Import addition supports concurrent processing.The asyncio import is necessary for the
asyncio.gatherfunctionality introduced in the_embed_batchmethod.
84-84: LGTM: Configurable batch size improves flexibility.The addition of the
batch_sizeparameter with a sensible default of 32 allows users to tune performance based on their needs, and storing it as an instance attribute_batch_sizefollows good encapsulation practices.Also applies to: 92-92
100-117: Excellent performance improvement with concurrent processing.The refactoring to use
asyncio.gatherfor concurrent batch processing is a significant performance improvement over the previous sequential approach. The implementation correctly:
- Creates batches of the specified size
- Processes all batches concurrently
- Maintains proper error handling through the existing backoff decorator
119-137: LGTM: Proper metadata aggregation for concurrent responses.The metadata aggregation logic correctly handles:
- Extending embeddings from all batch responses
- Accumulating usage statistics (prompt_tokens, total_tokens)
- Preserving model information from the first response
141-141: AsyncDocumentEmbedder.run signature change—no callers found using removed parametersA full-repo search for
AsyncDocumentEmbedder.run,batch_size, andprogress_barshowed no existing calls passing the now-removed parameters. This breaking change has no impact on current callers.
| embeddings, meta = await self._embed_batch( | ||
| texts_to_embed=texts_to_embed, | ||
| batch_size=batch_size, | ||
| progress_bar=progress_bar, | ||
| batch_size=self._batch_size, | ||
| ) |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Update method call to match refactored signature.
The call to _embed_batch should be updated to remove the batch_size parameter since it's no longer used by the method.
Apply this diff to match the refactored method signature:
- embeddings, meta = await self._embed_batch(
- texts_to_embed=texts_to_embed,
- batch_size=self._batch_size,
- )
+ embeddings, meta = await self._embed_batch(
+ texts_to_embed=texts_to_embed,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| embeddings, meta = await self._embed_batch( | |
| texts_to_embed=texts_to_embed, | |
| batch_size=batch_size, | |
| progress_bar=progress_bar, | |
| batch_size=self._batch_size, | |
| ) | |
| embeddings, meta = await self._embed_batch( | |
| texts_to_embed=texts_to_embed, | |
| ) |
🤖 Prompt for AI Agents
In wren-ai-service/src/providers/embedder/litellm.py around lines 154 to 157,
the call to the _embed_batch method incorrectly includes the batch_size
parameter, which has been removed in the refactored method signature. Remove the
batch_size argument from the method call so that only texts_to_embed is passed
to _embed_batch.
| self._api_base_url = api_base_url | ||
| self._timeout = timeout | ||
| self._kwargs = kwargs | ||
|
|
||
| async def _embed_batch( | ||
| self, texts_to_embed: List[str], batch_size: int, progress_bar: bool = True | ||
| self, texts_to_embed: List[str], batch_size: int |
There was a problem hiding this comment.
🛠️ Refactor suggestion
Remove unused parameter from method signature.
The batch_size parameter in the _embed_batch method signature is no longer used since the method now uses self._batch_size. This creates inconsistency and potential confusion.
Apply this diff to clean up the method signature:
- async def _embed_batch(
- self, texts_to_embed: List[str], batch_size: int
- ) -> Tuple[List[List[float]], Dict[str, Any]]:
+ async def _embed_batch(
+ self, texts_to_embed: List[str]
+ ) -> Tuple[List[List[float]], Dict[str, Any]]:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| self, texts_to_embed: List[str], batch_size: int | |
| async def _embed_batch( | |
| self, texts_to_embed: List[str] | |
| ) -> Tuple[List[List[float]], Dict[str, Any]]: |
🤖 Prompt for AI Agents
In wren-ai-service/src/providers/embedder/litellm.py at line 98, remove the
unused parameter batch_size from the _embed_batch method signature since the
method uses self._batch_size internally. This will clean up the method signature
and avoid confusion about the parameter's usage.
Summary by CodeRabbit
New Features
Refactor