[Enhancement]: Type-Safety for Async Milvus Client

[shubh3ai]

Is there an existing issue for this?

• I have searched the existing issues

What would you like to be added?

The Async Milvus Client right now returns dict objects for all the API requests that it makes. This makes the code much harder to manage. For example, the documents on the official page are at times not in sync with the Python library. Or it makes the code harder to debug. ATM I have to go into the Go Source code for the milvus repo to make sure that the values that I'm looking for in the request body actually makes sense.

Adding Pydantic Datatypes or any other datatype to enforce a structure on the return responses will help with:

1. easier to debug the code
2. makes it easier to manage the pymilvus library itself

I'm willing to work on this issue and add type safety myself if the maintainers are onboard with the proposal.

Why is this needed?

No response

Anything else?

No response

[shubh3ai]

For example, here is a document on fetching the load state for a collection: https://milvus.io/api-reference/pymilvus/v2.6.x/MilvusClient/Management/get_load_state.md

Now, if I want to ensure the correctness of my code, I need to know the possible return values of this function. Are there just types of load states i.e., Loaded, NotLoad? Or are there other possible states eg, LoadFailed, LoadProgressing?

And what if down the line, if there is an API update and the names of these variables change. Should I again look at the source code to find out what the correct string values for Load State could be?

To resolve issues like these and actually be able to make this library usable in production, I would strongly recommend adding type safety to the client library.

[XuanYang-cn]

Yeah, I'm totally with you. There's one more thing that I'm concern: compatibility.

Let's make this an official proposal:
We need type safety for MilvusClient and AsyncMilvusClient's return results for some hard-using APIs results while keeping compatible with previous output, so that code written in older PyMilvus are still able to run in higher version of PyMilvus.

Here's how I plan to do this: for every API return type changes, we need the following steps: open an issue and doc the Behavior before, expected Behavior after, and how we're gonna do it with compatibility keep in mind.

1.get_load_state returns a Dictionary before

ret = {"state": "Loaded"}

# in the previous code, here's how we use it
if ret.get("state") == "Loaded":
    pass

# Or define lots of consts to avoid misspell 
# (which is the most common mistakes yet the hardest one to debug in Python)
STATE_LOADED = "Loaded"
STATE_KEY = "state"
if ret.get(STATE_KEY) == STATE_LOADED:
    pass

if ret[STATE_KEY] == STATE_LOADED:
    pass

2. New return types and usage after:

from enum import StrEnum

class LoadState(StrEnum):
    Loaded = "Loaded"
    LoadFail = "LoadFail"

    # The following code is for compatibility
    #  Maybe we need to mark it as deprecated to print a deprecated warning and plan to remove version.
    def __getitem__(self, key: str) -> str:
        if key == "state":
            return self.value
        return {}[key]

    def get(self, key: str, default: str | None = None) -> str:
        if key == "state":
            return self.value
        return {}.get(key, default)


if __name__ == "__main__":
    a = LoadState.Loaded
    print(f"a = {a}")
    print(f"a['state'] = {a['state']}")
    print(f"a.get('state') = {a.get('state')}")

# Output
# a = Loaded
# a['state'] = Loaded
# a.get('state') = Loaded

@shubh3ai What do you think?

[XuanYang-cn]

@shubh3ai We can started by creating sub-issues in this issue. And I cannot wait working with you to make PyMilvus better and easier to use.

[XuanYang-cn]

But I don't encourage to change all APIs in one big pr. cuz big api changes may not convincing our tests and engneering team that it's safe to use. And we also need proper tests for compatibility.