feat(rag): add Milvus provider#6503
Open
zc277584121 wants to merge 3 commits into
Open
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughAdded Milvus as a RAG vector-store provider, including configuration, optional dependency handling, factory dispatch, synchronous and asynchronous client operations, schema and filtering utilities, tests, and English/Arabic documentation. ChangesMilvus RAG integration
Sequence Diagram(s)sequenceDiagram
participant RAGConfig
participant RAGFactory
participant MilvusFactory
participant MilvusClient
participant Milvus
RAGConfig->>RAGFactory: select provider "milvus"
RAGFactory->>MilvusFactory: require Milvus factory
MilvusFactory->>Milvus: create pymilvus client
MilvusFactory->>MilvusClient: wrap client and embedding function
MilvusClient->>Milvus: create collection and upsert documents
MilvusClient->>Milvus: search vectors with filters
Milvus-->>MilvusClient: return search hits
MilvusClient-->>RAGFactory: return SearchResult values
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
zc277584121
force-pushed
the
add-milvus-rag-client
branch
from
July 10, 2026 01:02
8c3c2e6 to
4b2b4c3
Compare
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
lib/crewai/src/crewai/rag/milvus/types.py (1)
35-43: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
any_schema()accepts non-callable values for the embedding function field.
MilvusEmbeddingFunctionWrapper.__get_pydantic_core_schema__returnscore_schema.any_schema(), which means Pydantic will accept any value — including non-callables like42orNone— for theembedding_functionconfig field. This creates a gap between the type annotation (which expects anEmbeddingFunction) and runtime validation. If config is loaded from a file (YAML/JSON), invalid values would pass validation and only fail later with a confusingTypeErrorwhen the client invokes the function.Consider using
core_schema.callable_schema()or a custom validator to at least ensure the value is callable.♻️ Proposed refactor
class MilvusEmbeddingFunctionWrapper(EmbeddingFunction): """Base class for Milvus embedding functions used by Pydantic validation.""" `@classmethod` def __get_pydantic_core_schema__( cls, _source_type: Any, _handler: GetCoreSchemaHandler ) -> CoreSchema: """Generate Pydantic core schema for Milvus embedding functions.""" - return core_schema.any_schema() + return core_schema.callable_schema()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/crewai/src/crewai/rag/milvus/types.py` around lines 35 - 43, Update MilvusEmbeddingFunctionWrapper.__get_pydantic_core_schema__ to return a schema that validates values are callable, such as core_schema.callable_schema(), or add an equivalent custom validator; ensure invalid values like None or numeric values are rejected during Pydantic validation while valid embedding functions remain accepted.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@lib/crewai/src/crewai/rag/milvus/types.py`:
- Around line 35-43: Update
MilvusEmbeddingFunctionWrapper.__get_pydantic_core_schema__ to return a schema
that validates values are callable, such as core_schema.callable_schema(), or
add an equivalent custom validator; ensure invalid values like None or numeric
values are rejected during Pydantic validation while valid embedding functions
remain accepted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: c65fbc91-0076-4b2b-a7e6-686a7b87cb95
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (21)
docs/edge/ar/concepts/knowledge.mdxdocs/edge/en/concepts/knowledge.mdxlib/crewai/pyproject.tomllib/crewai/src/crewai/rag/config/optional_imports/base.pylib/crewai/src/crewai/rag/config/optional_imports/protocols.pylib/crewai/src/crewai/rag/config/optional_imports/providers.pylib/crewai/src/crewai/rag/config/optional_imports/types.pylib/crewai/src/crewai/rag/config/types.pylib/crewai/src/crewai/rag/factory.pylib/crewai/src/crewai/rag/milvus/__init__.pylib/crewai/src/crewai/rag/milvus/client.pylib/crewai/src/crewai/rag/milvus/config.pylib/crewai/src/crewai/rag/milvus/constants.pylib/crewai/src/crewai/rag/milvus/factory.pylib/crewai/src/crewai/rag/milvus/types.pylib/crewai/src/crewai/rag/milvus/utils.pylib/crewai/tests/rag/config/test_factory.pylib/crewai/tests/rag/config/test_optional_imports.pylib/crewai/tests/rag/milvus/__init__.pylib/crewai/tests/rag/milvus/test_client.pylib/crewai/tests/rag/milvus/test_config.py
Signed-off-by: Cheney Zhang <chen.zhang@zilliz.com>
zc277584121
force-pushed
the
add-milvus-rag-client
branch
from
July 15, 2026 00:32
4b2b4c3 to
07148f8
Compare
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
lib/crewai/src/crewai/rag/milvus/utils.py (1)
351-358: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winCache the Milvus Lite version lookup instead of resolving it per search hit.
_normalize_milvus_score(called once per hit in_process_search_results's loop) invokes this function on every COSINE-metric result, and each call re-runsimportlib.metadata.version("milvus-lite"), which scans installed package metadata. For searches returning many hits this adds avoidable per-result I/O; the result is invariant for the process lifetime and should be memoized.♻️ Proposed fix: memoize the version check
+import functools + + +@functools.lru_cache(maxsize=1) def _milvus_lite_uses_cosine_distance() -> bool: """Return whether the installed Milvus Lite version reports cosine distance.""" try: version = importlib.metadata.version("milvus-lite") except importlib.metadata.PackageNotFoundError: return False base_version = version.split("+", maxsplit=1)[0] return base_version in MILVUS_LITE_COSINE_DISTANCE_VERSIONS🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lib/crewai/src/crewai/rag/milvus/utils.py` around lines 351 - 358, Memoize the result of _milvus_lite_uses_cosine_distance so importlib.metadata.version("milvus-lite") is resolved only once per process. Preserve the existing PackageNotFoundError handling and version normalization, using the project’s established caching approach or a suitable function-level memoization mechanism.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@lib/crewai/src/crewai/rag/milvus/utils.py`:
- Around line 361-374: Update _normalize_milvus_score to handle the
VALID_METRIC_TYPES entry IP explicitly instead of using the cosine-style
fallback mapping. Normalize inner-product embeddings before scoring, or define
and document an IP-specific normalization that preserves meaningful score
ordering and threshold behavior, while keeping COSINE and L2 handling unchanged.
---
Nitpick comments:
In `@lib/crewai/src/crewai/rag/milvus/utils.py`:
- Around line 351-358: Memoize the result of _milvus_lite_uses_cosine_distance
so importlib.metadata.version("milvus-lite") is resolved only once per process.
Preserve the existing PackageNotFoundError handling and version normalization,
using the project’s established caching approach or a suitable function-level
memoization mechanism.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 2abebe57-53ef-41ef-be12-e1b78b6c31a1
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (21)
docs/edge/ar/concepts/knowledge.mdxdocs/edge/en/concepts/knowledge.mdxlib/crewai/pyproject.tomllib/crewai/src/crewai/rag/config/optional_imports/base.pylib/crewai/src/crewai/rag/config/optional_imports/protocols.pylib/crewai/src/crewai/rag/config/optional_imports/providers.pylib/crewai/src/crewai/rag/config/optional_imports/types.pylib/crewai/src/crewai/rag/config/types.pylib/crewai/src/crewai/rag/factory.pylib/crewai/src/crewai/rag/milvus/__init__.pylib/crewai/src/crewai/rag/milvus/client.pylib/crewai/src/crewai/rag/milvus/config.pylib/crewai/src/crewai/rag/milvus/constants.pylib/crewai/src/crewai/rag/milvus/factory.pylib/crewai/src/crewai/rag/milvus/types.pylib/crewai/src/crewai/rag/milvus/utils.pylib/crewai/tests/rag/config/test_factory.pylib/crewai/tests/rag/config/test_optional_imports.pylib/crewai/tests/rag/milvus/__init__.pylib/crewai/tests/rag/milvus/test_client.pylib/crewai/tests/rag/milvus/test_config.py
🚧 Files skipped from review as they are similar to previous changes (18)
- lib/crewai/pyproject.toml
- lib/crewai/src/crewai/rag/config/optional_imports/types.py
- lib/crewai/src/crewai/rag/milvus/factory.py
- lib/crewai/tests/rag/milvus/init.py
- lib/crewai/src/crewai/rag/milvus/init.py
- lib/crewai/tests/rag/milvus/test_config.py
- lib/crewai/src/crewai/rag/milvus/constants.py
- lib/crewai/src/crewai/rag/milvus/config.py
- lib/crewai/src/crewai/rag/factory.py
- lib/crewai/tests/rag/config/test_optional_imports.py
- lib/crewai/src/crewai/rag/config/optional_imports/providers.py
- lib/crewai/tests/rag/config/test_factory.py
- docs/edge/en/concepts/knowledge.mdx
- lib/crewai/src/crewai/rag/config/types.py
- lib/crewai/src/crewai/rag/config/optional_imports/protocols.py
- lib/crewai/src/crewai/rag/milvus/types.py
- lib/crewai/src/crewai/rag/milvus/client.py
- lib/crewai/tests/rag/milvus/test_client.py
Author
|
Hi @lorenzejay @joaomdmoura, this Milvus RAG provider PR is ready for maintainer review when you have a chance. It adds an optional Milvus provider following the existing RAG provider pattern, with Milvus Lite defaults, remote Milvus / Zilliz Cloud configuration, docs, and focused tests. Happy to adjust the scope or API shape to better match the project conventions. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Tests