Migrated get_embeddings from litellm service to kairon #1844
Conversation
…esponding test cases
WalkthroughThis pull request introduces new methods in the Changes
Possibly related PRs
Suggested reviewers
Poem
✨ Finishing Touches
🪧 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: 0
🧹 Nitpick comments (5)
kairon/shared/llm/processor.py (3)
188-189: Preserve original traceback when re-raising exceptions.
Useraise Exception(...) from eto keep context of the original error:- except Exception as e: - raise Exception(f"Failed to fetch embeddings: {str(e)}") + except Exception as e: + raise Exception(f"Failed to fetch embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
189-189: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
532-532: Preserve original traceback for sparse embedding errors.
Useraise ... from ewhen re-raising:- except Exception as e: - raise Exception(f"Error processing sparse embeddings: {str(e)}") + except Exception as e: + raise Exception(f"Error processing sparse embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
532-532: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
548-548: Preserve original traceback for rerank embedding errors.
Useraise ... from ewhen re-raising:- except Exception as e: - raise Exception(f"Error processing rerank embeddings: {str(e)}") + except Exception as e: + raise Exception(f"Error processing rerank embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
548-548: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
tests/unit_test/llm_test.py (2)
21-22: Remove unused imports.
BotSecretTypeandBotSecretsare imported but never used. Consider dropping them to maintain a tidy import list.-from kairon.shared.admin.constants import BotSecretType -from kairon.shared.admin.data_objects import BotSecrets🧰 Tools
🪛 Ruff (0.8.2)
21-21:
kairon.shared.admin.constants.BotSecretTypeimported but unusedRemove unused import:
kairon.shared.admin.constants.BotSecretType(F401)
22-22:
kairon.shared.admin.data_objects.BotSecretsimported but unusedRemove unused import:
kairon.shared.admin.data_objects.BotSecrets(F401)
2650-2650: Remove assignments to the unusedprocessorvariable.
Following static analysis recommendations, remove these unused assignments:- processor = LLMProcessor(bot="test_bot", llm_type="openai") ... - processor = LLMProcessor(bot="test_bot", llm_type="openai") ... - processor = LLMProcessor(bot="test_bot", llm_type="openai") ... - processor = LLMProcessor(bot="test_bot", llm_type="openai") ... - processor = LLMProcessor(bot="test_bot", llm_type="openai") ... - processor = LLMProcessor(bot="test_bot", llm_type="openai")Also applies to: 2681-2681, 2713-2713, 2746-2746, 2777-2777, 2809-2809
🧰 Tools
🪛 Ruff (0.8.2)
2650-2650: Local variable
processoris assigned to but never usedRemove assignment to unused variable
processor(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (2)
kairon/shared/llm/processor.py(4 hunks)tests/unit_test/llm_test.py(10 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
kairon/shared/llm/processor.py
189-189: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
532-532: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
548-548: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
tests/unit_test/llm_test.py
21-21: kairon.shared.admin.constants.BotSecretType imported but unused
Remove unused import: kairon.shared.admin.constants.BotSecretType
(F401)
22-22: kairon.shared.admin.data_objects.BotSecrets imported but unused
Remove unused import: kairon.shared.admin.data_objects.BotSecrets
(F401)
2650-2650: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
2681-2681: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
2713-2713: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
2746-2746: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
2777-2777: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
2809-2809: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: Analyze (python)
- GitHub Check: Python CI
- GitHub Check: Codacy Static Code Analysis
🔇 Additional comments (4)
kairon/shared/llm/processor.py (4)
57-58: Initialize embeddings at construction.
These calls toself.load_sparse_embedding_model()andself.load_rerank_embedding_model()ensure that the embeddings are available as soon as theLLMProcessoris instantiated. This is a good design choice for on-demand functionality.
148-159: Token truncation logic looks good.
Thetruncate_textmethod correctly slices the token list toEMBEDDING_CTX_LENGTHand re-derives the string. This helps avoid oversize embeddings.
490-491: Surface error messages to guide troubleshooting.
Raising a descriptive exception here is appropriate to help diagnose configuration failures. No further change required.
492-513: Sparse and Re-rank embedding model loading.
These methods implement a robust fallback for caching directories. The design ensures a single initialization. This is a sensible approach.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (1)
requirements/prod.txt (1)
72-72: Consider pinning the 'nltk' dependency version.
The re-includednltkdependency is currently unversioned. For consistency across environments and to avoid potential compatibility issues with future releases, it is advisable to specify a tested version (e.g.,nltk==3.8.0).
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (1)
requirements/prod.txt(1 hunks)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: Codacy Static Code Analysis
- GitHub Check: Python CI
- GitHub Check: Analyze (python)
🔇 Additional comments (1)
requirements/prod.txt (1)
73-73: Validate the 'fastembed==0.5.1' dependency for feature consistency.
The new dependencyfastembed==0.5.1appears to support the new embedding functionalities introduced in the LLMProcessor. Please confirm that this version meets the API requirements and does not introduce conflicts with other dependencies in the project.
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (7)
kairon/shared/llm/processor.py (5)
157-186: Improve exception handling with proper chaining.The method effectively handles both single string and list inputs, properly truncates texts, and collects embeddings from multiple models. However, the exception handling could be improved.
except Exception as e: - raise Exception(f"Failed to fetch embeddings: {str(e)}") + raise Exception(f"Failed to fetch embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
186-186: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
487-488: Improve exception handling with proper chaining.Use exception chaining to preserve the original traceback information.
- raise Exception(f"Failed to fetch vector configs: {http_response.get('message', 'Unknown error')}") + raise Exception(f"Failed to fetch vector configs: {http_response.get('message', 'Unknown error')}") from None
500-510: DRY principle violation - consider refactoring model loading methods.Both
load_sparse_embedding_modelandload_rerank_embedding_modelhave identical cache directory logic. Consider extracting this common functionality.+ @classmethod + def _get_cache_dir(cls): + hf_cache_dir = os.path.expanduser("~/.cache/huggingface/hub") + kairon_cache_dir = "./kairon/pre-trained-models/" + return hf_cache_dir if os.path.exists(hf_cache_dir) else kairon_cache_dir @classmethod def load_sparse_embedding_model(cls): - hf_cache_dir = os.path.expanduser("~/.cache/huggingface/hub") - kairon_cache_dir = "./kairon/pre-trained-models/" - - cache_dir = hf_cache_dir if os.path.exists(hf_cache_dir) else kairon_cache_dir + cache_dir = cls._get_cache_dir() if cls._sparse_embedding is None: cls._sparse_embedding = SparseTextEmbedding("Qdrant/bm25", cache_dir=cache_dir) logging.info("SPARSE MODEL LOADED") @classmethod def load_rerank_embedding_model(cls): - hf_cache_dir = os.path.expanduser("~/.cache/huggingface/hub") - kairon_cache_dir = "./kairon/pre-trained-models/" - - cache_dir = hf_cache_dir if os.path.exists(hf_cache_dir) else kairon_cache_dir + cache_dir = cls._get_cache_dir() if cls._rerank_embedding is None: cls._rerank_embedding = LateInteractionTextEmbedding("colbert-ir/colbertv2.0", cache_dir=cache_dir) logging.info("RERANK MODEL LOADED")
511-530: Improve exception handling with proper chaining in get_sparse_embedding.The method correctly generates sparse embeddings, but could benefit from proper exception chaining.
try: embeddings = list(self._sparse_embedding.passage_embed(sentences)) return [ {"values": emb.values.tolist(), "indices": emb.indices.tolist()} for emb in embeddings ] except Exception as e: - raise Exception(f"Error processing sparse embeddings: {str(e)}") + raise Exception(f"Error processing sparse embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
529-529: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
531-545: Improve exception handling with proper chaining in get_rerank_embedding.The method correctly generates rerank embeddings, but could benefit from proper exception chaining.
try: embeddings = list(self._rerank_embedding.passage_embed(sentences)) return [emb.tolist() for emb in embeddings] except Exception as e: - raise Exception(f"Error processing rerank embeddings: {str(e)}") + raise Exception(f"Error processing rerank embeddings: {str(e)}") from e🧰 Tools
🪛 Ruff (0.8.2)
545-545: Within an
exceptclause, raise exceptions withraise ... from errorraise ... from Noneto distinguish them from errors in exception handling(B904)
tests/unit_test/llm_test.py (2)
21-22: Remove unused imports.These imports are not used in the test file and should be removed.
-from kairon.shared.admin.constants import BotSecretType -from kairon.shared.admin.data_objects import BotSecrets, LLMSecret +from kairon.shared.admin.data_objects import LLMSecret🧰 Tools
🪛 Ruff (0.8.2)
21-21:
kairon.shared.admin.constants.BotSecretTypeimported but unusedRemove unused import:
kairon.shared.admin.constants.BotSecretType(F401)
22-22:
kairon.shared.admin.data_objects.BotSecretsimported but unusedRemove unused import:
kairon.shared.admin.data_objects.BotSecrets(F401)
2628-2659: Test for model initialization could be improved.The test correctly checks that the model doesn't reload if already initialized, but the unused variable
processorcould be removed.logger.add(log_output, format="{message}") - processor = LLMProcessor(bot="test_bot", llm_type="openai") + # Initialize to verify no additional loading occurs + LLMProcessor(bot="test_bot", llm_type="openai") logger.remove()🧰 Tools
🪛 Ruff (0.8.2)
2650-2650: Local variable
processoris assigned to but never usedRemove assignment to unused variable
processor(F841)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
📒 Files selected for processing (7)
kairon/actions/definitions/prompt.py(1 hunks)kairon/shared/cognition/processor.py(1 hunks)kairon/shared/llm/processor.py(3 hunks)kairon/shared/vector_embeddings/db/qdrant.py(1 hunks)kairon/train.py(1 hunks)tests/integration_test/action_service_test.py(8 hunks)tests/unit_test/llm_test.py(10 hunks)
🧰 Additional context used
🪛 Ruff (0.8.2)
kairon/shared/llm/processor.py
186-186: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
529-529: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
545-545: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
tests/unit_test/llm_test.py
21-21: kairon.shared.admin.constants.BotSecretType imported but unused
Remove unused import: kairon.shared.admin.constants.BotSecretType
(F401)
22-22: kairon.shared.admin.data_objects.BotSecrets imported but unused
Remove unused import: kairon.shared.admin.data_objects.BotSecrets
(F401)
2650-2650: Local variable processor is assigned to but never used
Remove assignment to unused variable processor
(F841)
⏰ Context from checks skipped due to timeout of 90000ms (3)
- GitHub Check: Codacy Static Code Analysis
- GitHub Check: Python CI
- GitHub Check: Analyze (python)
🔇 Additional comments (22)
kairon/train.py (1)
20-21: Model loading added at module levelThese two method calls ensure that sparse and rerank embedding models are loaded during module initialization. This aligns with the PR objective of migrating embedding functionality from litellm service to kairon.
kairon/actions/definitions/prompt.py (1)
17-19: Model loading added at module levelAdding these embedding model loading methods ensures that the models are ready when
ActionPromptneeds to process user queries. This is consistent with the changes in other files and supports the migration from litellm service.kairon/shared/vector_embeddings/db/qdrant.py (1)
11-13: Import and model loading addedThe LLMProcessor import has been uncommented and the embedding model loading methods have been added. This is particularly important for the Qdrant class which directly interacts with vector embeddings and uses
llm.get_embedding()for vector operations.kairon/shared/cognition/processor.py (1)
588-590: Model loading added in upsert_data methodUnlike the other files where model loading happens at module level, here the models are loaded inside the
upsert_datamethod just before creating the LLMProcessor instance. This ensures the models are available when generating embeddings for data being upserted to Qdrant.tests/integration_test/action_service_test.py (8)
3808-3810: Good addition of passthru URLs for embedding modelsThese additions allow the test to access Hugging Face embedding models (BM25 and ColBERT) directly, which is consistent with the migration of embedding functionality from litellm to Kairon. This ensures the tests can properly validate the new internal embedding implementation.
3901-3903: Good addition of passthru URLs for embedding modelsProperly added the same Hugging Face model URLs to allow this test to access the embedding models directly. This is consistent with the embedding functionality migration.
3993-3995: Good addition of passthru URLs for embedding modelsThese passthru URLs correctly allow real HTTP requests to the Hugging Face models (sparse BM25 and reranking ColBERT) needed for the embedding functionality that's been migrated to Kairon.
4088-4090: Good addition of passthru URLs for embedding modelsThe addition of these Hugging Face model URLs is consistent with the embedding implementation migration. Since this test specifically mocks the
get_embeddingmethod, allowing passthru to embedding model URLs ensures the test correctly simulates the integration between components.
4225-4227: Good addition of passthru URLs for embedding modelsThe addition of passthru URLs for the BM25 and ColBERT models allows proper testing of the payload search functionality using the migrated embedding implementation.
4316-4318: Good addition of passthru URLs for embedding modelsThese passthru URLs appropriately support testing error handling scenarios with the new embedding implementation, ensuring the tests properly validate error cases with the migrated functionality.
4379-4381: Good addition of passthru URLs for embedding modelsProperly added passthru URLs for the embedding models while testing slot-based embedding search. The mock of
get_embeddingcombined with these passthru URLs ensures thorough testing of the migrated functionality.
4519-4521: Good addition of passthru URLs for embedding modelsThe addition of these passthru URLs completes the testing coverage for the no-response-dispatch scenario with the new embedding implementation. This ensures proper validation of all edge cases with the migrated functionality.
kairon/shared/llm/processor.py (2)
145-155: Well-structured implementation for text truncation.The
truncate_textmethod correctly handles the tokenization of input texts and ensures they don't exceed the embedding context length limit of 8191 tokens, which is essential for OpenAI's embeddings.
489-499: Good implementation for model loading with fallback mechanism.The method properly checks for the Hugging Face cache directory and falls back to a local cache directory if needed. The singleton pattern ensures the model is only loaded once.
tests/unit_test/llm_test.py (8)
4-4: Proper imports for testing.Good use of MagicMock and AsyncMock for testing asynchronous methods.
7-9: New imports for testing logging and stream handling.Added appropriate imports for capturing log output during tests.
13-13: Import for embedding models.Good addition of the SparseTextEmbedding and LateInteractionTextEmbedding imports to test the new functionality.
1950-2046: Well-structured test for collection hybrid query.The test thoroughly verifies the hybrid query functionality, including checking the request payload and validating the response format.
2047-2081: Good test for handling request failures.Proper testing of error handling when a request fails.
2119-2146: Comprehensive test for sparse embedding with single sentence.The test properly mocks the sparse embedding model and verifies the output format.
2365-2414: Complete test coverage for single text embedding.The test thoroughly covers the entire embedding pipeline for a single text input, including dense, sparse, and rerank embeddings.
2417-2478: Thorough test for multiple text embedding.Good test coverage for processing multiple texts through the embedding pipeline.
Summary by CodeRabbit
New Features
Tests
Chores
fastembedand re-addednltk.