Replace ORM Collection/utility with MilvusClient API (pymilvus 2.5.x/2.6.x compatibility)

[sbright83]

Summary

• Replaces all legacy ORM-based operations (Collection, utility.*) with the modern MilvusClient API, compatible with both pymilvus 2.5.x and 2.6.x
• Migrates create_collection, drop_collection, has_collection, insert, delete, search, query, hybrid_search, describe_collection, create_index, load_collection, and get_load_state to MilvusClient equivalents
• Fixes partition_names not being forwarded to search, query, hybrid_search, and count_documents operations

Key changes

• self.col remains Optional[Collection] for backward compatibility, now implemented as a cached @property (same approach as langchain-milvus); all internal operations use MilvusClient. Since pymilvus 2.6.x no longer registers the client connection with the ORM connections registry, the store registers the client's gRPC handler itself so the ORM Collection can resolve the alias — this works on both 2.5.x and 2.6.x
• self.alias is kept as a public attribute
• Index creation uses prepare_index_params() / add_index() / create_index()
• _get_index keeps the ORM col.indexes iteration: MilvusClient.describe_index raises when the field has no index yet, which would break the first-time index creation path
• replica_number is forwarded to load_collection / load_partitions (no longer a silent no-op)
• Dropping a collection also clears pymilvus's connection-level schema cache (see [Bug]: pymilvus#3058), so drop_old=True cannot read back stale schema
• Search parameter key renamed from param= (ORM) to search_params= (MilvusClient); filter parameter renamed from expr= to filter= for search/query (AnnSearchRequest still uses expr=, correct per its API); hybrid_search ranker kwarg renamed from rerank= to ranker=
• _parse_search_result updated: dict(res.fields) replaces the ORM entity field iteration
• CI now runs the test suite against both pymilvus 2.5.x and 2.6.x

Notes

• from pymilvus.orm.types import infer_dtype_bydata remains as an ORM-module import; this function is a pure utility (no ORM state) that is not exported from the top-level pymilvus namespace and cannot be eliminated without duplicating its logic
• Full test suite passes locally against Milvus 2.6.14 with both pymilvus 2.5.18 and 2.6.11

[zc277584121]

Thanks for tackling this migration — the direction is right and the partition_names forwarding fix is a nice catch. Before merging, a few concerns around backward compatibility and a couple of potential regressions. The sister project langchain-ai/langchain-milvus went through the same migration recently (PR #104) and solved most of these cleanly — worth borrowing from.

🔴 Must fix

1. self.col type change breaks users

Changing self.col from Optional[Collection] to bool is a public-API break for anyone who has been reaching into store.col (it has no leading underscore, so it's effectively semi-public). langchain-milvus kept self.col as a real Collection via a cached @property, while internally switching all the checks to self.client.has_collection(name):

@property
def col(self) -> Optional[Collection]:
    current_key = f"{self.collection_name}:{self.alias}"
    if self._cache_key == current_key and self._col_cache is not None:
        return self._col_cache
    if self.client.has_collection(self.collection_name):
        self._col_cache = Collection(self.collection_name, using=self.alias)
        self._cache_key = current_key
        return self._col_cache
    return None

This preserves backward compatibility and still lets the rest of the code rely on the MilvusClient API.

2. Keep self.alias

self.alias = self.client._using was removed. Same concern as above — it's a public attribute, and langchain-milvus kept it for exactly this reason (also needed to construct the ORM Collection with using= in the property above).

3. _get_index regression with describe_index

client.describe_index(collection, field) raises when the field has no index, whereas the previous for x in self.col.indexes simply returned nothing. The check if self.col and self._get_index() is None in _create_index will now hit an exception on the first-time index creation path instead of falling through cleanly.

langchain-milvus deliberately kept the ORM self.col.indexes iteration here and left a comment noting that list_indexes + describe_index works but is less efficient and has this exact edge case. Recommend the same: keep the ORM path for _get_index only, using the cached Collection from the property above.

🟠 Strongly recommended

4. replica_number should not silently become a no-op

MilvusClient.load_collection and load_partitions both accept replica_number natively (pymilvus 2.5+). Please forward it instead of dropping it — silent behavioral degradation is worse than an API break, especially for users on Zilliz Cloud / multi-replica setups:

if partition_names:
    self.client.load_partitions(
        self.collection_name,
        partition_names=partition_names,
        replica_number=self.replica_number,
        timeout=timeout,
    )
else:
    self.client.load_collection(
        self.collection_name,
        replica_number=self.replica_number,
        timeout=timeout,
    )

5. Clear pymilvus schema cache on drop_collection

See pymilvus#3058 — after client.drop_collection(name), pymilvus internally caches the old schema, so recreating a collection with the same name can read back stale schema. langchain-milvus handles it like this:

self.client.drop_collection(self.collection_name)
conn = self.client._get_connection()
if hasattr(conn, "schema_cache"):
    conn.schema_cache.pop(self.collection_name, None)
# plus reset self._col_cache / self._cache_key

Relevant here because drop_old=True is a common path.

🟡 Discussion

6. pymilvus version matrix

The PR title advertises 2.5.x/2.6.x compatibility, but I don't see a CI matrix covering both. Two options:

• Add a CI matrix test against pymilvus 2.5.x and 2.6.x, or
• Follow langchain-milvus's approach and simply raise the floor (pymilvus>=2.6.3), dropping the dual-version goal.

Either is fine, but the current state ("claims to support both, tested against neither in CI") is risky.

---

Overall this is a valuable migration and I'd like to see it land. The items above (especially 1–3) are about avoiding silent breakage for existing users; 4 avoids a behavioral regression; 5 is a latent bug that the ORM path happened to mask. Happy to help review the next revision.

[sbright83]

Sorry for the delayed response and thank you for the thorough review! I tried to incorporate all 6 items in the current commit. Going through them point-by-point:

1. self.col type change — Restored as Optional[Collection] via a cached @property, following the langchain-milvus pattern you linked. Internal operations all stay on MilvusClient; the property exists purely for backward compatibility.

One gotcha worth flagging: the langchain-milvus snippet as written fails on recent pymilvus 2.6.x with ConnectionNotExistException. Since the ConnectionManager refactor, MilvusClient no longer registers its connection with the ORM connections registry — client._using is just a legacy-compat string (cm-<id>), so Collection(name, using=alias) can't resolve the alias. I work around it by registering the client's handler at init:

if hasattr(connections, "_alias_handlers"):
    connections._alias_handlers.setdefault(self.alias, self.client._get_connection())

This is a no-op on 2.5.x (where the alias is already registered) and makes the ORM Collection work on 2.6.x. langchain-milvus may have the same latent issue depending on which 2.6.x versions they test against.

2. self.alias — Restored (self.client._using), also needed for the property above.

3. _get_index regression — Reverted to the ORM col.indexes iteration with a comment explaining why (MilvusClient.describe_index raises when the field has no index, breaking first-time index creation). _create_search_params reads index["index_param"][...] again to match the ORM dict shape.

4. replica_number — Now forwarded, using your suggested shape: load_partitions(..., replica_number=...) when partition_names is set, load_collection(..., replica_number=...) otherwise.

5. Schema cache on drop — Added a _drop_collection helper that drops via the client, pops the collection from the connection's schema_cache (pymilvus#3058), and resets the cached Collection / cache key. Used on the drop_old=True path.

6. Version matrix — Kept the dual-version goal and added a CI matrix: the test workflow now runs against both pymilvus==2.5.* and pymilvus==2.6.* (with fail-fast: false).

Verified locally against Milvus 2.6.14: full suite passes with both pymilvus 2.5.18 and 2.6.11, and lint/mypy are clean.

[mergify]

Tick the box to add this pull request to the merge queue (same as @mergifyio queue).

• Queue this pull request