Use async client for ChromaDB
Feature description
This is a subissue to #567
Value and/or benefit
No response
Anything else?
No response
Currently, the only option to create an async client for ChromaDB is to use AsyncHttpClient. This cannot be used in a Python class initializer because Python class initializers must necessarily be synchronous.
Furthermore, we do not want to use an HTTP client unless we are necessarily making requests over HTTP because this will incur a performance hit.
@pmeier and I discussed adding an AsyncHttpClient side by side with a Client so that we can take advantage of asynchronous programming when we are making expensive HTTP requests. The finer technical details still need to be ironed out, particularly programmatically choosing the correct client at runtime. For example, we only want single store or retrieve methods, so the class should be able to decide which Chroma client to use internally.
For connecting, we can use the same pattern I layed out in https://github.com/Quansight/ragna/issues/572#issuecomment-2788832802.
As for having Client and AsyncHttpClient side-by-side: my proposal is to implement all methods in terms of the AsyncHttpClient and a thin wrapper class around Client to make it behave like an async one. For example:
import functools
from typing import Any, TypeVar
from ragna._utils import as_awaitable
T = TypeVar("T")
class _AsyncWrapper:
def __init__(self, o: Any) -> None:
self.__wrapped_object__ = o
def __getattribute__(self, item: str) -> Any:
o = object.__getattribute__(self, "__dict__")["__wrapped_object__"]
attr = getattr(o, item)
if callable(attr):
attr = functools.partial(as_awaitable, attr)
return attr
class Foo:
def bar(self, baz: str) -> None:
print(baz)
async def main() -> None:
foo = Foo()
foo.bar("sync")
afoo = _AsyncWrapper(foo)
await afoo.bar("async")
import asyncio
asyncio.run(main())
With this we could do something like
class Chroma:
def __init__(self):
if condition_for_using_http_client():
self._client = chromadb.AsyncHttpClient(...)
else:
self._client = _AsyncWrapper(chromadb.Client(...))
This way we can implement all other methods completely agnostic to the fact whether or not we are using a native sync or async client.