Merge pull request #94 from tharropoulos/nl-language

feat: add natural language search

Author: kishorenc

src/typesense/client.py

@@ -46,6 +46,7 @@
 from typesense.keys import Keys
 from typesense.metrics import Metrics
 from typesense.multi_search import MultiSearch
+from typesense.nl_search_models import NLSearchModels
 from typesense.operations import Operations
 from typesense.stemming import Stemming
 from typesense.stopwords import Stopwords
@@ -107,6 +108,7 @@ def __init__(self, config_dict: ConfigDict) -> None:
         self.stopwords = Stopwords(self.api_call)
         self.metrics = Metrics(self.api_call)
         self.conversations_models = ConversationsModels(self.api_call)
+        self.nl_search_models = NLSearchModels(self.api_call)
 
     def typed_collection(
         self,

src/typesense/nl_search_model.py

@@ -0,0 +1,108 @@
+"""
+This module provides functionality for managing individual NL search models in Typesense.
+
+Classes:
+    - NLSearchModel: Handles operations related to a specific NL search model.
+
+Methods:
+    - __init__: Initializes the NLSearchModel object.
+    - _endpoint_path: Constructs the API endpoint path for this specific NL search model.
+    - retrieve: Retrieves the details of this specific NL search model.
+    - update: Updates this specific NL search model.
+    - delete: Deletes this specific NL search model.
+
+The NLSearchModel class interacts with the Typesense API to manage operations on a
+specific NL search model. It provides methods to retrieve, update,
+and delete individual models.
+
+This module uses type hinting and is compatible with Python 3.11+ as well as earlier
+versions through the use of the typing_extensions library.
+"""
+
+from typesense.api_call import ApiCall
+from typesense.types.nl_search_model import (
+    NLSearchModelDeleteSchema,
+    NLSearchModelSchema,
+    NLSearchModelUpdateSchema,
+)
+
+
+class NLSearchModel:
+    """
+    Class for managing individual NL search models in Typesense.
+
+    This class provides methods to interact with a specific NL search model,
+    including retrieving, updating, and deleting it.
+
+    Attributes:
+        model_id (str): The ID of the NL search model.
+        api_call (ApiCall): The API call object for making requests.
+    """
+
+    def __init__(self, api_call: ApiCall, model_id: str) -> None:
+        """
+        Initialize the NLSearchModel object.
+
+        Args:
+            api_call (ApiCall): The API call object for making requests.
+            model_id (str): The ID of the NL search model.
+        """
+        self.model_id = model_id
+        self.api_call = api_call
+
+    def retrieve(self) -> NLSearchModelSchema:
+        """
+        Retrieve this specific NL search model.
+
+        Returns:
+            NLSearchModelSchema: The schema containing the NL search model details.
+        """
+        response = self.api_call.get(
+            self._endpoint_path,
+            as_json=True,
+            entity_type=NLSearchModelSchema,
+        )
+        return response
+
+    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 = self.api_call.put(
+            self._endpoint_path,
+            body=model,
+            entity_type=NLSearchModelSchema,
+        )
+        return response
+
+    def delete(self) -> NLSearchModelDeleteSchema:
+        """
+        Delete this specific NL search model.
+
+        Returns:
+            NLSearchModelDeleteSchema: The schema containing the deletion response.
+        """
+        response: NLSearchModelDeleteSchema = 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.nl_search_models import NLSearchModels
+
+        return "/".join([NLSearchModels.resource_path, self.model_id])

src/typesense/nl_search_models.py

@@ -0,0 +1,117 @@
+"""
+This module provides functionality for managing NL search models in Typesense.
+
+Classes:
+    - NLSearchModels: Handles operations related to NL search models.
+
+Methods:
+    - __init__: Initializes the NLSearchModels object.
+    - __getitem__: Retrieves or creates an NLSearchModel object for a given model_id.
+    - create: Creates a new NL search model.
+    - retrieve: Retrieves all NL search models.
+
+Attributes:
+    - resource_path: The API resource path for NL search models operations.
+
+The NLSearchModels class interacts with the Typesense API to manage
+NL search model operations.
+
+It provides methods to create and retrieve NL search models, as well as access
+individual NLSearchModel objects.
+
+This module uses type hinting and is compatible with Python 3.11+ as well as earlier
+versions through the use of the typing_extensions library.
+"""
+
+import sys
+
+from typesense.api_call import ApiCall
+from typesense.types.nl_search_model import (
+    NLSearchModelCreateSchema,
+    NLSearchModelSchema,
+    NLSearchModelsRetrieveSchema,
+)
+
+if sys.version_info > (3, 11):
+    import typing
+else:
+    import typing_extensions as typing
+
+from typesense.nl_search_model import NLSearchModel
+
+
+class NLSearchModels(object):
+    """
+    Class for managing NL search models in Typesense.
+
+    This class provides methods to interact with NL search models, including
+    creating, retrieving, and accessing individual models.
+
+    Attributes:
+        resource_path (str): The API resource path for NL search models operations.
+        api_call (ApiCall): The API call object for making requests.
+        nl_search_models (Dict[str, NLSearchModel]):
+            A dictionary of NLSearchModel objects.
+    """
+
+    resource_path: typing.Final[str] = "/nl_search_models"
+
+    def __init__(self, api_call: ApiCall) -> None:
+        """
+        Initialize the NLSearchModels object.
+
+        Args:
+            api_call (ApiCall): The API call object for making requests.
+        """
+        self.api_call = api_call
+        self.nl_search_models: typing.Dict[str, NLSearchModel] = {}
+
+    def __getitem__(self, model_id: str) -> NLSearchModel:
+        """
+        Get or create an NLSearchModel object for a given model_id.
+
+        Args:
+            model_id (str): The ID of the NL search model.
+
+        Returns:
+            NLSearchModel: The NLSearchModel object for the given ID.
+        """
+        if model_id not in self.nl_search_models:
+            self.nl_search_models[model_id] = NLSearchModel(
+                self.api_call,
+                model_id,
+            )
+        return self.nl_search_models[model_id]
+
+    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.
+        """
+        response = self.api_call.post(
+            endpoint=NLSearchModels.resource_path,
+            entity_type=NLSearchModelSchema,
+            as_json=True,
+            body=model,
+        )
+        return response
+
+    def retrieve(self) -> NLSearchModelsRetrieveSchema:
+        """
+        Retrieve all NL search models.
+
+        Returns:
+            NLSearchModelsRetrieveSchema: A list of all NL search models.
+        """
+        response: NLSearchModelsRetrieveSchema = self.api_call.get(
+            endpoint=NLSearchModels.resource_path,
+            entity_type=NLSearchModelsRetrieveSchema,
+            as_json=True,
+        )
+        return response

src/typesense/types/document.py

@@ -569,6 +569,23 @@ class CachingParameters(typing.TypedDict):
     cache_ttl: typing.NotRequired[int]
 
 
+class NLLanguageParameters(typing.TypedDict):
+    """
+    Parameters regarding [caching search results](https://typesense.org/docs/26.0/api/search.html#caching-parameters).
+
+    Attributes:
+      nl_query_prompt_cache_ttl (int): The duration (in seconds) that determines how long the schema prompts are cached.
+      nl_query (bool): Whether to use natural language in the query or not.
+      nl_model_id (str): The ID of the natural language model to use for the query.
+      nl_query_debug (bool): Whether to return the raw LLM response or not.
+    """
+
+    nl_query_prompt_cache_ttl: typing.NotRequired[int]
+    nl_query: typing.NotRequired[bool]
+    nl_model_id: typing.NotRequired[str]
+    nl_query_debug: typing.NotRequired[bool]
+
+
 class SearchParameters(
     RequiredSearchParameters,
     QueryParameters,
@@ -580,6 +597,7 @@ class SearchParameters(
     ResultsParameters,
     TypoToleranceParameters,
     CachingParameters,
+    NLLanguageParameters,
 ):
     """Parameters for searching documents."""
 
@@ -823,6 +841,38 @@ class Conversation(typing.TypedDict):
     query: str
 
 
+class LLMResponse(typing.TypedDict):
+    """
+    Schema for a raw LLM response.
+
+    Attributes:
+      content (str): Content of the LLM response.
+      extraction_method (str): Extraction method of the LLM response (e.g. "regex").
+      model (str): Model used to generate the response.
+    """
+
+    content: str
+    extraction_method: str
+    model: str
+
+
+class ParsedNLQuery(typing.TypedDict):
+    """
+    Schema for a parsed natural language query.
+
+    Attributes:
+      parse_time_ms (int): Parse time in milliseconds.
+      generated_params (SearchParameters): Generated parameters.
+      augmented_params (SearchParameters): Augmented parameters.
+      llm_response (LLMResponse): Raw LLM response.
+    """
+
+    parse_time_ms: int
+    generated_params: SearchParameters
+    augmented_params: SearchParameters
+    llm_response: typing.NotRequired[LLMResponse]
+
+
 class SearchResponse(typing.Generic[TDoc], typing.TypedDict):
     """
     Schema for a search response.
@@ -838,6 +888,7 @@ class SearchResponse(typing.Generic[TDoc], typing.TypedDict):
       hits (list[Hit[TDoc]]): List of hits in the search results.
       grouped_hits (list[GroupedHit[TDoc]]): List of grouped hits in the search results.
       conversation (Conversation): Conversation in the search results.
+      parsed_nl_query (ParsedNLQuery): Information about the natural language query
     """
 
     facet_counts: typing.List[SearchResponseFacetCountSchema]
@@ -850,6 +901,7 @@ class SearchResponse(typing.Generic[TDoc], typing.TypedDict):
     hits: typing.List[Hit[TDoc]]
     grouped_hits: typing.NotRequired[typing.List[GroupedHit[TDoc]]]
     conversation: typing.NotRequired[Conversation]
+    parsed_nl_query: typing.NotRequired[ParsedNLQuery]
 
 
 class DeleteSingleDocumentParameters(typing.TypedDict):