mirror of
https://github.com/qdrant/qdrant-client.git
synced 2026-07-23 11:11:01 -05:00
264 lines
10 KiB
Python
264 lines
10 KiB
Python
from asyncio import get_event_loop
|
|
from functools import lru_cache
|
|
from typing import Any, Awaitable, Callable, Dict, Generic, Type, TypeVar, overload
|
|
from urllib.parse import urljoin
|
|
|
|
from httpx import AsyncClient, Client, Request, Response
|
|
from pydantic import ValidationError
|
|
from qdrant_client.common.client_exceptions import ResourceExhaustedResponse
|
|
from qdrant_client.http.api.aliases_api import AsyncAliasesApi, SyncAliasesApi
|
|
from qdrant_client.http.api.beta_api import AsyncBetaApi, SyncBetaApi
|
|
from qdrant_client.http.api.collections_api import AsyncCollectionsApi, SyncCollectionsApi
|
|
from qdrant_client.http.api.distributed_api import AsyncDistributedApi, SyncDistributedApi
|
|
from qdrant_client.http.api.indexes_api import AsyncIndexesApi, SyncIndexesApi
|
|
from qdrant_client.http.api.points_api import AsyncPointsApi, SyncPointsApi
|
|
from qdrant_client.http.api.search_api import AsyncSearchApi, SyncSearchApi
|
|
from qdrant_client.http.api.service_api import AsyncServiceApi, SyncServiceApi
|
|
from qdrant_client.http.api.snapshots_api import AsyncSnapshotsApi, SyncSnapshotsApi
|
|
from qdrant_client.http.exceptions import ResponseHandlingException, UnexpectedResponse
|
|
|
|
ClientT = TypeVar("ClientT", bound="ApiClient")
|
|
AsyncClientT = TypeVar("AsyncClientT", bound="AsyncApiClient")
|
|
|
|
|
|
class AsyncApis(Generic[AsyncClientT]):
|
|
def __init__(self, host: str, **kwargs: Any):
|
|
self.client = AsyncApiClient(host, **kwargs)
|
|
|
|
self.aliases_api = AsyncAliasesApi(self.client)
|
|
self.beta_api = AsyncBetaApi(self.client)
|
|
self.collections_api = AsyncCollectionsApi(self.client)
|
|
self.distributed_api = AsyncDistributedApi(self.client)
|
|
self.indexes_api = AsyncIndexesApi(self.client)
|
|
self.points_api = AsyncPointsApi(self.client)
|
|
self.search_api = AsyncSearchApi(self.client)
|
|
self.service_api = AsyncServiceApi(self.client)
|
|
self.snapshots_api = AsyncSnapshotsApi(self.client)
|
|
|
|
async def aclose(self) -> None:
|
|
await self.client.aclose()
|
|
|
|
|
|
class SyncApis(Generic[ClientT]):
|
|
def __init__(self, host: str, **kwargs: Any):
|
|
self.client = ApiClient(host, **kwargs)
|
|
|
|
self.aliases_api = SyncAliasesApi(self.client)
|
|
self.beta_api = SyncBetaApi(self.client)
|
|
self.collections_api = SyncCollectionsApi(self.client)
|
|
self.distributed_api = SyncDistributedApi(self.client)
|
|
self.indexes_api = SyncIndexesApi(self.client)
|
|
self.points_api = SyncPointsApi(self.client)
|
|
self.search_api = SyncSearchApi(self.client)
|
|
self.service_api = SyncServiceApi(self.client)
|
|
self.snapshots_api = SyncSnapshotsApi(self.client)
|
|
|
|
def close(self) -> None:
|
|
self.client.close()
|
|
|
|
|
|
T = TypeVar("T")
|
|
Send = Callable[[Request], Response]
|
|
SendAsync = Callable[[Request], Awaitable[Response]]
|
|
MiddlewareT = Callable[[Request, Send], Response]
|
|
AsyncMiddlewareT = Callable[[Request, SendAsync], Awaitable[Response]]
|
|
|
|
|
|
class ApiClient:
|
|
def __init__(self, host: str, **kwargs: Any) -> None:
|
|
self.host = host
|
|
self.middleware: MiddlewareT = BaseMiddleware()
|
|
self._client = Client(**kwargs)
|
|
|
|
@overload
|
|
def request(self, *, type_: Type[T], method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any) -> T:
|
|
...
|
|
|
|
@overload # noqa F811
|
|
def request(self, *, type_: None, method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any) -> None:
|
|
...
|
|
|
|
def request( # noqa F811
|
|
self, *, type_: Any, method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any
|
|
) -> Any:
|
|
if path_params is None:
|
|
path_params = {}
|
|
|
|
host = self.host if self.host.endswith("/") else self.host + "/"
|
|
url = url[1:] if url.startswith("/") else url
|
|
# in order to do a correct join, url join requires base_url to end with /, and url to not start with /,
|
|
# since url is treated as an absolute path and might truncate prefix in base_url
|
|
url = urljoin(host, url.format(**path_params))
|
|
if "params" in kwargs and "timeout" in kwargs["params"]:
|
|
kwargs["timeout"] = int(kwargs["params"]["timeout"])
|
|
request = self._client.build_request(method, url, **kwargs)
|
|
return self.send(request, type_)
|
|
|
|
@overload
|
|
def request_sync(self, *, type_: Type[T], **kwargs: Any) -> T:
|
|
...
|
|
|
|
@overload # noqa F811
|
|
def request_sync(self, *, type_: None, **kwargs: Any) -> None:
|
|
...
|
|
|
|
def request_sync(self, *, type_: Any, **kwargs: Any) -> Any: # noqa F811
|
|
"""
|
|
This method is not used by the generated apis, but is included for convenience
|
|
"""
|
|
return get_event_loop().run_until_complete(self.request(type_=type_, **kwargs))
|
|
|
|
def send(self, request: Request, type_: Type[T]) -> T:
|
|
response = self.middleware(request, self.send_inner)
|
|
|
|
if response.status_code == 429:
|
|
retry_after_s = response.headers.get("Retry-After", None)
|
|
try:
|
|
resp = response.json()
|
|
message = resp["status"]["error"] if resp["status"] and resp["status"]["error"] else ""
|
|
except Exception:
|
|
message = ""
|
|
|
|
if retry_after_s:
|
|
raise ResourceExhaustedResponse(message, retry_after_s)
|
|
|
|
if response.status_code in [200, 201, 202]:
|
|
try:
|
|
return parse_as_type(response.json(), type_)
|
|
except ValidationError as e:
|
|
raise ResponseHandlingException(e)
|
|
raise UnexpectedResponse.for_response(response)
|
|
|
|
def send_inner(self, request: Request) -> Response:
|
|
try:
|
|
response = self._client.send(request)
|
|
except Exception as e:
|
|
raise ResponseHandlingException(e)
|
|
return response
|
|
|
|
def close(self) -> None:
|
|
self._client.close()
|
|
|
|
def add_middleware(self, middleware: MiddlewareT) -> None:
|
|
current_middleware = self.middleware
|
|
|
|
def new_middleware(request: Request, call_next: Send) -> Response:
|
|
def inner_send(request: Request) -> Response:
|
|
return current_middleware(request, call_next)
|
|
|
|
return middleware(request, inner_send)
|
|
|
|
self.middleware = new_middleware
|
|
|
|
|
|
class AsyncApiClient:
|
|
def __init__(self, host: str = None, **kwargs: Any) -> None:
|
|
self.host = host
|
|
self.middleware: AsyncMiddlewareT = BaseAsyncMiddleware()
|
|
self._async_client = AsyncClient(**kwargs)
|
|
|
|
@overload
|
|
async def request(
|
|
self, *, type_: Type[T], method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any
|
|
) -> T:
|
|
...
|
|
|
|
@overload # noqa F811
|
|
async def request(
|
|
self, *, type_: None, method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any
|
|
) -> None:
|
|
...
|
|
|
|
async def request( # noqa F811
|
|
self, *, type_: Any, method: str, url: str, path_params: Dict[str, Any] = None, **kwargs: Any
|
|
) -> Any:
|
|
if path_params is None:
|
|
path_params = {}
|
|
|
|
host = self.host if self.host.endswith("/") else self.host + "/"
|
|
url = url[1:] if url.startswith("/") else url
|
|
# in order to do a correct join, url join requires base_url to end with /, and url to not start with /,
|
|
# since url is treated as an absolute path and might truncate prefix in base_url
|
|
url = urljoin(host, url.format(**path_params))
|
|
request = self._async_client.build_request(method, url, **kwargs)
|
|
return await self.send(request, type_)
|
|
|
|
@overload
|
|
def request_sync(self, *, type_: Type[T], **kwargs: Any) -> T:
|
|
...
|
|
|
|
@overload # noqa F811
|
|
def request_sync(self, *, type_: None, **kwargs: Any) -> None:
|
|
...
|
|
|
|
def request_sync(self, *, type_: Any, **kwargs: Any) -> Any: # noqa F811
|
|
"""
|
|
This method is not used by the generated apis, but is included for convenience
|
|
"""
|
|
return get_event_loop().run_until_complete(self.request(type_=type_, **kwargs))
|
|
|
|
async def send(self, request: Request, type_: Type[T]) -> T:
|
|
response = await self.middleware(request, self.send_inner)
|
|
|
|
if response.status_code == 429:
|
|
retry_after_s = response.headers.get("Retry-After", None)
|
|
try:
|
|
resp = response.json()
|
|
message = resp["status"]["error"] if resp["status"] and resp["status"]["error"] else ""
|
|
except Exception:
|
|
message = ""
|
|
|
|
if retry_after_s:
|
|
raise ResourceExhaustedResponse(message, retry_after_s)
|
|
|
|
if response.status_code in [200, 201, 202]:
|
|
try:
|
|
return parse_as_type(response.json(), type_)
|
|
except ValidationError as e:
|
|
raise ResponseHandlingException(e)
|
|
raise UnexpectedResponse.for_response(response)
|
|
|
|
async def send_inner(self, request: Request) -> Response:
|
|
try:
|
|
response = await self._async_client.send(request)
|
|
except Exception as e:
|
|
raise ResponseHandlingException(e)
|
|
return response
|
|
|
|
async def aclose(self) -> None:
|
|
await self._async_client.aclose()
|
|
|
|
def add_middleware(self, middleware: AsyncMiddlewareT) -> None:
|
|
current_middleware = self.middleware
|
|
|
|
async def new_middleware(request: Request, call_next: SendAsync) -> Response:
|
|
async def inner_send(request: Request) -> Response:
|
|
return await current_middleware(request, call_next)
|
|
|
|
return await middleware(request, inner_send)
|
|
|
|
self.middleware = new_middleware
|
|
|
|
|
|
class BaseAsyncMiddleware:
|
|
async def __call__(self, request: Request, call_next: SendAsync) -> Response:
|
|
return await call_next(request)
|
|
|
|
|
|
class BaseMiddleware:
|
|
def __call__(self, request: Request, call_next: Send) -> Response:
|
|
return call_next(request)
|
|
|
|
|
|
@lru_cache(maxsize=None)
|
|
def _get_parsing_type(type_: Any, source: str) -> Any:
|
|
from pydantic.main import create_model
|
|
|
|
type_name = getattr(type_, "__name__", str(type_))
|
|
return create_model(f"ParsingModel[{type_name}] (for {source})", obj=(type_, ...))
|
|
|
|
|
|
def parse_as_type(obj: Any, type_: Type[T]) -> T:
|
|
model_type = _get_parsing_type(type_, source=parse_as_type.__name__)
|
|
return model_type(obj=obj).obj
|