Skip to content

feat: add natural language search #94

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions src/typesense/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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,
Expand Down
108 changes: 108 additions & 0 deletions src/typesense/nl_search_model.py
Original file line number Diff line number Diff line change
@@ -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])
117 changes: 117 additions & 0 deletions src/typesense/nl_search_models.py
Original file line number Diff line number Diff line change
@@ -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
52 changes: 52 additions & 0 deletions src/typesense/types/document.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -580,6 +597,7 @@ class SearchParameters(
ResultsParameters,
TypoToleranceParameters,
CachingParameters,
NLLanguageParameters,
):
"""Parameters for searching documents."""

Expand Down Expand Up @@ -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.
Expand All @@ -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]
Expand All @@ -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):
Expand Down
Loading