feat(nl-search): add async support for nl search model operations

tharropoulos · tharropoulos · commit 6a6e4d374b39 · 2025-12-09T11:05:54.000+02:00
- add AsyncNLSearchModel class for async individual model operations
- add AsyncNLSearchModels class for async mod
diff --git a/src/typesense/async_nl_search_model.py b/src/typesense/async_nl_search_model.py
@@ -0,0 +1,102 @@
+"""
+This module provides async functionality for managing individual NL search models in Typesense.
+
+It contains the AsyncNLSearchModel class, which allows for retrieving, updating, and deleting
+NL search models asynchronously.
+
+Classes:
+    AsyncNLSearchModel: Manages async operations on a single NL search model in the Typesense API.
+
+Dependencies:
+    - typesense.async_api_call: Provides the AsyncApiCall class for making async API requests.
+    - typesense.types.nl_search_model: Provides NLSearchModelDeleteSchema, NLSearchModelSchema, and NLSearchModelUpdateSchema types.
+
+Note: This module uses conditional imports to support both Python 3.11+ and earlier versions.
+"""
+
+from typesense.async_api_call import AsyncApiCall
+from typesense.types.nl_search_model import (
+    NLSearchModelDeleteSchema,
+    NLSearchModelSchema,
+    NLSearchModelUpdateSchema,
+)
+
+
+class AsyncNLSearchModel:
+    """
+    Manages async operations on a single NL search model in the Typesense API.
+
+    This class provides async methods to retrieve, update, and delete an NL search model.
+
+    Attributes:
+        model_id (str): The ID of the NL search model.
+        api_call (AsyncApiCall): The AsyncApiCall instance for making async API requests.
+    """
+
+    def __init__(self, api_call: AsyncApiCall, model_id: str) -> None:
+        """
+        Initialize the AsyncNLSearchModel instance.
+
+        Args:
+            api_call (AsyncApiCall): The AsyncApiCall instance for making async API requests.
+            model_id (str): The ID of the NL search model.
+        """
+        self.model_id = model_id
+        self.api_call = api_call
+
+    async def retrieve(self) -> NLSearchModelSchema:
+        """
+        Retrieve this specific NL search model.
+
+        Returns:
+            NLSearchModelSchema: The schema containing the NL search model details.
+        """
+        response: NLSearchModelSchema = await self.api_call.get(
+            self._endpoint_path,
+            as_json=True,
+            entity_type=NLSearchModelSchema,
+        )
+        return response
+
+    async def update(self, model: NLSearchModelUpdateSchema) -> NLSearchModelSchema:
+        """
+        Update this specific NL search model.
+
+        Args:
+            model (NLSearchModelUpdateSchema):
+              The schema containing the updated model details.
+
+        Returns:
+            NLSearchModelSchema: The schema containing the updated NL search model.
+        """
+        response: NLSearchModelSchema = await self.api_call.put(
+            self._endpoint_path,
+            body=model,
+            entity_type=NLSearchModelSchema,
+        )
+        return response
+
+    async def delete(self) -> NLSearchModelDeleteSchema:
+        """
+        Delete this specific NL search model.
+
+        Returns:
+            NLSearchModelDeleteSchema: The schema containing the deletion response.
+        """
+        response: NLSearchModelDeleteSchema = await self.api_call.delete(
+            self._endpoint_path,
+            entity_type=NLSearchModelDeleteSchema,
+        )
+        return response
+
+    @property
+    def _endpoint_path(self) -> str:
+        """
+        Construct the API endpoint path for this specific NL search model.
+
+        Returns:
+            str: The constructed endpoint path.
+        """
+        from typesense.async_nl_search_models import AsyncNLSearchModels
+
+        return "/".join([AsyncNLSearchModels.resource_path, self.model_id])
diff --git a/src/typesense/async_nl_search_models.py b/src/typesense/async_nl_search_models.py
@@ -0,0 +1,130 @@
+"""
+This module provides async functionality for managing NL search models in Typesense.
+
+It contains the AsyncNLSearchModels class, which allows for creating, retrieving, and
+accessing individual NL search models asynchronously.
+
+Classes:
+    AsyncNLSearchModels: Manages NL search models in the Typesense API (async).
+
+Dependencies:
+    - typesense.async_api_call: Provides the AsyncApiCall class for making async API requests.
+    - typesense.async_nl_search_model: Provides the AsyncNLSearchModel class for individual NL search model operations.
+    - typesense.types.nl_search_model: Provides NLSearchModelCreateSchema, NLSearchModelSchema, and NLSearchModelsRetrieveSchema types.
+
+Note: This module uses conditional imports to support both Python 3.11+ and earlier versions.
+"""
+
+import sys
+
+from typesense.async_api_call import AsyncApiCall
+from typesense.async_nl_search_model import AsyncNLSearchModel
+from typesense.types.nl_search_model import (
+    NLSearchModelCreateSchema,
+    NLSearchModelSchema,
+    NLSearchModelsRetrieveSchema,
+)
+
+if sys.version_info >= (3, 11):
+    import typing
+else:
+    import typing_extensions as typing
+
+
+class AsyncNLSearchModels:
+    """
+    Manages NL search models in the Typesense API (async).
+
+    This class provides async methods to create, retrieve, and access individual NL search models.
+
+    Attributes:
+        resource_path (str): The API endpoint path for NL search models operations.
+        api_call (AsyncApiCall): The AsyncApiCall instance for making async API requests.
+        nl_search_models (Dict[str, AsyncNLSearchModel]):
+            A dictionary of AsyncNLSearchModel instances, keyed by model ID.
+    """
+
+    resource_path: typing.Final[str] = "/nl_search_models"
+
+    def __init__(self, api_call: AsyncApiCall) -> None:
+        """
+        Initialize the AsyncNLSearchModels instance.
+
+        Args:
+            api_call (AsyncApiCall): The AsyncApiCall instance for making async API requests.
+        """
+        self.api_call = api_call
+        self.nl_search_models: typing.Dict[str, AsyncNLSearchModel] = {}
+
+    def __getitem__(self, model_id: str) -> AsyncNLSearchModel:
+        """
+        Get or create an AsyncNLSearchModel instance for a given model ID.
+
+        This method allows accessing NL search models using dictionary-like syntax.
+        If the AsyncNLSearchModel instance doesn't exist, it creates a new one.
+
+        Args:
+            model_id (str): The ID of the NL search model.
+
+        Returns:
+            AsyncNLSearchModel: The AsyncNLSearchModel instance for the specified model ID.
+
+        Example:
+            >>> nl_search_models = AsyncNLSearchModels(async_api_call)
+            >>> model = nl_search_models["model_id"]
+        """
+        if model_id not in self.nl_search_models:
+            self.nl_search_models[model_id] = AsyncNLSearchModel(
+                self.api_call,
+                model_id,
+            )
+        return self.nl_search_models[model_id]
+
+    async def create(self, model: NLSearchModelCreateSchema) -> NLSearchModelSchema:
+        """
+        Create a new NL search model.
+
+        Args:
+            model (NLSearchModelCreateSchema):
+                The schema for creating the NL search model.
+
+        Returns:
+            NLSearchModelSchema: The created NL search model.
+
+        Example:
+            >>> nl_search_models = AsyncNLSearchModels(async_api_call)
+            >>> model = await nl_search_models.create(
+            ...     {
+            ...         "api_key": "key",
+            ...         "model_name": "openai/gpt-3.5-turbo",
+            ...         "system_prompt": "System prompt",
+            ...     }
+            ... )
+        """
+        response: NLSearchModelSchema = await self.api_call.post(
+            endpoint=AsyncNLSearchModels.resource_path,
+            entity_type=NLSearchModelSchema,
+            as_json=True,
+            body=model,
+        )
+        return response
+
+    async def retrieve(self) -> NLSearchModelsRetrieveSchema:
+        """
+        Retrieve all NL search models.
+
+        Returns:
+            NLSearchModelsRetrieveSchema: A list of all NL search models.
+
+        Example:
+            >>> nl_search_models = AsyncNLSearchModels(async_api_call)
+            >>> all_models = await nl_search_models.retrieve()
+            >>> for model in all_models:
+            ...     print(model["id"])
+        """
+        response: NLSearchModelsRetrieveSchema = await self.api_call.get(
+            endpoint=AsyncNLSearchModels.resource_path,
+            entity_type=NLSearchModelsRetrieveSchema,
+            as_json=True,
+        )
+        return response
diff --git a/tests/fixtures/nl_search_model_fixtures.py b/tests/fixtures/nl_search_model_fixtures.py
@@ -7,6 +7,9 @@
 from dotenv import load_dotenv
 
 from typesense.api_call import ApiCall
+from typesense.async_api_call import AsyncApiCall
+from typesense.async_nl_search_model import AsyncNLSearchModel
+from typesense.async_nl_search_models import AsyncNLSearchModels
 from typesense.nl_search_model import NLSearchModel
 from typesense.nl_search_models import NLSearchModels
 
@@ -76,3 +79,27 @@ def actual_nl_search_models_fixture(
 ) -> NLSearchModels:
     """Return an NLSearchModels object using a real API."""
     return NLSearchModels(actual_api_call)
+
+
+@pytest.fixture(scope="function", name="actual_async_nl_search_models")
+def actual_async_nl_search_models_fixture(
+    actual_async_api_call: AsyncApiCall,
+) -> AsyncNLSearchModels:
+    """Return a AsyncNLSearchModels object using a real API."""
+    return AsyncNLSearchModels(actual_async_api_call)
+
+
+@pytest.fixture(scope="function", name="fake_async_nl_search_models")
+def fake_async_nl_search_models_fixture(
+    fake_async_api_call: AsyncApiCall,
+) -> AsyncNLSearchModels:
+    """Return a AsyncNLSearchModels object with test values."""
+    return AsyncNLSearchModels(fake_async_api_call)
+
+
+@pytest.fixture(scope="function", name="fake_async_nl_search_model")
+def fake_async_nl_search_model_fixture(
+    fake_async_api_call: AsyncApiCall,
+) -> AsyncNLSearchModel:
+    """Return a AsyncNLSearchModel object with test values."""
+    return AsyncNLSearchModel(fake_async_api_call, "nl_search_model_id")
diff --git a/tests/nl_search_model_test.py b/tests/nl_search_model_test.py
@@ -11,6 +11,9 @@
     assert_to_contain_keys,
 )
 from typesense.api_call import ApiCall
+from typesense.async_api_call import AsyncApiCall
+from typesense.async_nl_search_model import AsyncNLSearchModel
+from typesense.async_nl_search_models import AsyncNLSearchModels
 from typesense.nl_search_model import NLSearchModel
 from typesense.nl_search_models import NLSearchModels
 
@@ -40,6 +43,29 @@ def test_init(fake_api_call: ApiCall) -> None:
     )
 
 
+def test_init_async(fake_async_api_call: AsyncApiCall) -> None:
+    """Test that the AsyncNLSearchModel object is initialized correctly."""
+    nl_search_model = AsyncNLSearchModel(
+        fake_async_api_call,
+        "nl_search_model_id",
+    )
+
+    assert nl_search_model.model_id == "nl_search_model_id"
+    assert_match_object(nl_search_model.api_call, fake_async_api_call)
+    assert_object_lists_match(
+        nl_search_model.api_call.node_manager.nodes,
+        fake_async_api_call.node_manager.nodes,
+    )
+    assert_match_object(
+        nl_search_model.api_call.config.nearest_node,
+        fake_async_api_call.config.nearest_node,
+    )
+    assert (
+        nl_search_model._endpoint_path  # noqa: WPS437
+        == "/nl_search_models/nl_search_model_id"
+    )
+
+
 @pytest.mark.open_ai
 def test_actual_retrieve(
     actual_nl_search_models: NLSearchModels,
@@ -97,3 +123,62 @@ def test_actual_delete(
     )
 
     assert response.get("id") == create_nl_search_model
+
+
+@pytest.mark.open_ai
+async def test_actual_retrieve_async(
+    actual_async_nl_search_models: AsyncNLSearchModels,
+    delete_all_nl_search_models: None,
+    create_nl_search_model: str,
+) -> None:
+    """Test it can retrieve an NL search model from Typesense Server."""
+    response = await actual_async_nl_search_models[create_nl_search_model].retrieve()
+
+    assert_to_contain_keys(
+        response,
+        ["id", "model_name", "system_prompt", "max_bytes", "api_key"],
+    )
+    assert response.get("id") == create_nl_search_model
+
+
+@pytest.mark.open_ai
+async def test_actual_update_async(
+    actual_async_nl_search_models: AsyncNLSearchModels,
+    delete_all_nl_search_models: None,
+    create_nl_search_model: str,
+) -> None:
+    """Test that it can update an NL search model from Typesense Server."""
+    response = await actual_async_nl_search_models[create_nl_search_model].update(
+        {"system_prompt": "This is a new system prompt for NL search"},
+    )
+
+    assert_to_contain_keys(
+        response,
+        [
+            "id",
+            "model_name",
+            "system_prompt",
+            "max_bytes",
+            "api_key",
+        ],
+    )
+
+    assert response.get("system_prompt") == "This is a new system prompt for NL search"
+    assert response.get("id") == create_nl_search_model
+
+
+@pytest.mark.open_ai
+async def test_actual_delete_async(
+    actual_async_nl_search_models: AsyncNLSearchModels,
+    delete_all_nl_search_models: None,
+    create_nl_search_model: str,
+) -> None:
+    """Test that it can delete an NL search model from Typesense Server."""
+    response = await actual_async_nl_search_models[create_nl_search_model].delete()
+
+    assert_to_contain_keys(
+        response,
+        ["id"],
+    )
+
+    assert response.get("id") == create_nl_search_model
diff --git a/tests/nl_search_models_test.py b/tests/nl_search_models_test.py
