# This file was auto-generated by Fern from our API Definition.

import typing
import urllib.parse
from json.decoder import JSONDecodeError

from ...core.api_error import ApiError
from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ...core.jsonable_encoder import jsonable_encoder
from ...errors.unprocessable_entity_error import UnprocessableEntityError
from ...types.http_validation_error import HttpValidationError
from ...types.retrieve_results import RetrieveResults

try:
    import pydantic.v1 as pydantic  # type: ignore
except ImportError:
    import pydantic  # type: ignore

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)


class RetrievalClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._client_wrapper = client_wrapper

    def run_search(
        self,
        *,
        dense_similarity_top_k: typing.Optional[int] = OMIT,
        sparse_similarity_top_k: typing.Optional[int] = OMIT,
        enable_reranking: typing.Optional[bool] = OMIT,
        rerank_top_n: typing.Optional[int] = OMIT,
        alpha: typing.Optional[float] = OMIT,
        search_filters: typing.Optional[typing.Dict[str, typing.List[typing.Any]]] = OMIT,
        query: str,
        pipeline_id: str,
        class_name: typing.Optional[str] = OMIT,
    ) -> RetrieveResults:
        """
        Get retrieval results for a query

        Parameters:
            - dense_similarity_top_k: typing.Optional[int]. Number of nodes for dense retrieval.

            - sparse_similarity_top_k: typing.Optional[int]. Number of nodes for sparse retrieval.

            - enable_reranking: typing.Optional[bool]. Enable reranking for retrieval

            - rerank_top_n: typing.Optional[int]. Number of reranked nodes for returning.

            - alpha: typing.Optional[float]. Alpha value for hybrid retrieval to determine the weights between dense and sparse retrieval. 0 is sparse retrieval and 1 is dense retrieval.

            - search_filters: typing.Optional[typing.Dict[str, typing.List[typing.Any]]]. Search filters for retrieval. the format of search_filters is a dict of {key: (operator, value)}

            - query: str. The query to retrieve against.

            - pipeline_id: str. The ID of the managed pipeline that the query was retrieved against.

            - class_name: typing.Optional[str].
        """
        _request: typing.Dict[str, typing.Any] = {"query": query, "pipeline_id": pipeline_id}
        if dense_similarity_top_k is not OMIT:
            _request["dense_similarity_top_k"] = dense_similarity_top_k
        if sparse_similarity_top_k is not OMIT:
            _request["sparse_similarity_top_k"] = sparse_similarity_top_k
        if enable_reranking is not OMIT:
            _request["enable_reranking"] = enable_reranking
        if rerank_top_n is not OMIT:
            _request["rerank_top_n"] = rerank_top_n
        if alpha is not OMIT:
            _request["alpha"] = alpha
        if search_filters is not OMIT:
            _request["search_filters"] = search_filters
        if class_name is not OMIT:
            _request["class_name"] = class_name
        _response = self._client_wrapper.httpx_client.request(
            "POST",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/retrieval/retrieve"),
            json=jsonable_encoder(_request),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(RetrieveResults, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)


class AsyncRetrievalClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._client_wrapper = client_wrapper

    async def run_search(
        self,
        *,
        dense_similarity_top_k: typing.Optional[int] = OMIT,
        sparse_similarity_top_k: typing.Optional[int] = OMIT,
        enable_reranking: typing.Optional[bool] = OMIT,
        rerank_top_n: typing.Optional[int] = OMIT,
        alpha: typing.Optional[float] = OMIT,
        search_filters: typing.Optional[typing.Dict[str, typing.List[typing.Any]]] = OMIT,
        query: str,
        pipeline_id: str,
        class_name: typing.Optional[str] = OMIT,
    ) -> RetrieveResults:
        """
        Get retrieval results for a query

        Parameters:
            - dense_similarity_top_k: typing.Optional[int]. Number of nodes for dense retrieval.

            - sparse_similarity_top_k: typing.Optional[int]. Number of nodes for sparse retrieval.

            - enable_reranking: typing.Optional[bool]. Enable reranking for retrieval

            - rerank_top_n: typing.Optional[int]. Number of reranked nodes for returning.

            - alpha: typing.Optional[float]. Alpha value for hybrid retrieval to determine the weights between dense and sparse retrieval. 0 is sparse retrieval and 1 is dense retrieval.

            - search_filters: typing.Optional[typing.Dict[str, typing.List[typing.Any]]]. Search filters for retrieval. the format of search_filters is a dict of {key: (operator, value)}

            - query: str. The query to retrieve against.

            - pipeline_id: str. The ID of the managed pipeline that the query was retrieved against.

            - class_name: typing.Optional[str].
        """
        _request: typing.Dict[str, typing.Any] = {"query": query, "pipeline_id": pipeline_id}
        if dense_similarity_top_k is not OMIT:
            _request["dense_similarity_top_k"] = dense_similarity_top_k
        if sparse_similarity_top_k is not OMIT:
            _request["sparse_similarity_top_k"] = sparse_similarity_top_k
        if enable_reranking is not OMIT:
            _request["enable_reranking"] = enable_reranking
        if rerank_top_n is not OMIT:
            _request["rerank_top_n"] = rerank_top_n
        if alpha is not OMIT:
            _request["alpha"] = alpha
        if search_filters is not OMIT:
            _request["search_filters"] = search_filters
        if class_name is not OMIT:
            _request["class_name"] = class_name
        _response = await self._client_wrapper.httpx_client.request(
            "POST",
            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/retrieval/retrieve"),
            json=jsonable_encoder(_request),
            headers=self._client_wrapper.get_headers(),
            timeout=60,
        )
        if 200 <= _response.status_code < 300:
            return pydantic.parse_obj_as(RetrieveResults, _response.json())  # type: ignore
        if _response.status_code == 422:
            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
        try:
            _response_json = _response.json()
        except JSONDecodeError:
            raise ApiError(status_code=_response.status_code, body=_response.text)
        raise ApiError(status_code=_response.status_code, body=_response_json)
