Skip to content

Commit cb27c74

Browse files
tharropoulostharropoulos
committed
feat(client): add nl search models functionality
- add `NLSearchModel` class for individual model operations (retrieve, update, delete) - add `NLSearchModels` class for collection operations (create, retrieve, __getitem__) - add type definitions in `types/nl_search_model.py` with support for multiple llm providers - integrate `nl_search_models` into main `Client` class with proper imports - add test fixtures and tests for both individual and collection operations
1 parent f33c6a3 commit cb27c74

File tree

7 files changed

+661
-0
lines changed

7 files changed

+661
-0
lines changed

‎src/typesense/client.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,7 @@
4646
from typesense.keys import Keys
4747
from typesense.metrics import Metrics
4848
from typesense.multi_search import MultiSearch
49+
from typesense.nl_search_models import NLSearchModels
4950
from typesense.operations import Operations
5051
from typesense.stemming import Stemming
5152
from typesense.stopwords import Stopwords
@@ -107,6 +108,7 @@ def __init__(self, config_dict: ConfigDict) -> None:
107108
self.stopwords = Stopwords(self.api_call)
108109
self.metrics = Metrics(self.api_call)
109110
self.conversations_models = ConversationsModels(self.api_call)
111+
self.nl_search_models = NLSearchModels(self.api_call)
110112

111113
def typed_collection(
112114
self,

‎src/typesense/nl_search_model.py‎

Lines changed: 108 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,108 @@
1+
"""
2+
This module provides functionality for managing individual NL search models in Typesense.
3+
4+
Classes:
5+
- NLSearchModel: Handles operations related to a specific NL search model.
6+
7+
Methods:
8+
- __init__: Initializes the NLSearchModel object.
9+
- _endpoint_path: Constructs the API endpoint path for this specific NL search model.
10+
- retrieve: Retrieves the details of this specific NL search model.
11+
- update: Updates this specific NL search model.
12+
- delete: Deletes this specific NL search model.
13+
14+
The NLSearchModel class interacts with the Typesense API to manage operations on a
15+
specific NL search model. It provides methods to retrieve, update,
16+
and delete individual models.
17+
18+
This module uses type hinting and is compatible with Python 3.11+ as well as earlier
19+
versions through the use of the typing_extensions library.
20+
"""
21+
22+
from typesense.api_call import ApiCall
23+
from typesense.types.nl_search_model import (
24+
NLSearchModelDeleteSchema,
25+
NLSearchModelSchema,
26+
NLSearchModelUpdateSchema,
27+
)
28+
29+
30+
class NLSearchModel:
31+
"""
32+
Class for managing individual NL search models in Typesense.
33+
34+
This class provides methods to interact with a specific NL search model,
35+
including retrieving, updating, and deleting it.
36+
37+
Attributes:
38+
model_id (str): The ID of the NL search model.
39+
api_call (ApiCall): The API call object for making requests.
40+
"""
41+
42+
def __init__(self, api_call: ApiCall, model_id: str) -> None:
43+
"""
44+
Initialize the NLSearchModel object.
45+
46+
Args:
47+
api_call (ApiCall): The API call object for making requests.
48+
model_id (str): The ID of the NL search model.
49+
"""
50+
self.model_id = model_id
51+
self.api_call = api_call
52+
53+
def retrieve(self) -> NLSearchModelSchema:
54+
"""
55+
Retrieve this specific NL search model.
56+
57+
Returns:
58+
NLSearchModelSchema: The schema containing the NL search model details.
59+
"""
60+
response = self.api_call.get(
61+
self._endpoint_path,
62+
as_json=True,
63+
entity_type=NLSearchModelSchema,
64+
)
65+
return response
66+
67+
def update(self, model: NLSearchModelUpdateSchema) -> NLSearchModelSchema:
68+
"""
69+
Update this specific NL search model.
70+
71+
Args:
72+
model (NLSearchModelUpdateSchema):
73+
The schema containing the updated model details.
74+
75+
Returns:
76+
NLSearchModelSchema: The schema containing the updated NL search model.
77+
"""
78+
response: NLSearchModelSchema = self.api_call.put(
79+
self._endpoint_path,
80+
body=model,
81+
entity_type=NLSearchModelSchema,
82+
)
83+
return response
84+
85+
def delete(self) -> NLSearchModelDeleteSchema:
86+
"""
87+
Delete this specific NL search model.
88+
89+
Returns:
90+
NLSearchModelDeleteSchema: The schema containing the deletion response.
91+
"""
92+
response: NLSearchModelDeleteSchema = self.api_call.delete(
93+
self._endpoint_path,
94+
entity_type=NLSearchModelDeleteSchema,
95+
)
96+
return response
97+
98+
@property
99+
def _endpoint_path(self) -> str:
100+
"""
101+
Construct the API endpoint path for this specific NL search model.
102+
103+
Returns:
104+
str: The constructed endpoint path.
105+
"""
106+
from typesense.nl_search_models import NLSearchModels
107+
108+
return "/".join([NLSearchModels.resource_path, self.model_id])

‎src/typesense/nl_search_models.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
"""
2+
This module provides functionality for managing NL search models in Typesense.
3+
4+
Classes:
5+
- NLSearchModels: Handles operations related to NL search models.
6+
7+
Methods:
8+
- __init__: Initializes the NLSearchModels object.
9+
- __getitem__: Retrieves or creates an NLSearchModel object for a given model_id.
10+
- create: Creates a new NL search model.
11+
- retrieve: Retrieves all NL search models.
12+
13+
Attributes:
14+
- resource_path: The API resource path for NL search models operations.
15+
16+
The NLSearchModels class interacts with the Typesense API to manage
17+
NL search model operations.
18+
19+
It provides methods to create and retrieve NL search models, as well as access
20+
individual NLSearchModel objects.
21+
22+
This module uses type hinting and is compatible with Python 3.11+ as well as earlier
23+
versions through the use of the typing_extensions library.
24+
"""
25+
26+
import sys
27+
28+
from typesense.api_call import ApiCall
29+
from typesense.types.nl_search_model import (
30+
NLSearchModelCreateSchema,
31+
NLSearchModelSchema,
32+
NLSearchModelsRetrieveSchema,
33+
)
34+
35+
if sys.version_info > (3, 11):
36+
import typing
37+
else:
38+
import typing_extensions as typing
39+
40+
from typesense.nl_search_model import NLSearchModel
41+
42+
43+
class NLSearchModels(object):
44+
"""
45+
Class for managing NL search models in Typesense.
46+
47+
This class provides methods to interact with NL search models, including
48+
creating, retrieving, and accessing individual models.
49+
50+
Attributes:
51+
resource_path (str): The API resource path for NL search models operations.
52+
api_call (ApiCall): The API call object for making requests.
53+
nl_search_models (Dict[str, NLSearchModel]):
54+
A dictionary of NLSearchModel objects.
55+
"""
56+
57+
resource_path: typing.Final[str] = "/nl_search_models"
58+
59+
def __init__(self, api_call: ApiCall) -> None:
60+
"""
61+
Initialize the NLSearchModels object.
62+
63+
Args:
64+
api_call (ApiCall): The API call object for making requests.
65+
"""
66+
self.api_call = api_call
67+
self.nl_search_models: typing.Dict[str, NLSearchModel] = {}
68+
69+
def __getitem__(self, model_id: str) -> NLSearchModel:
70+
"""
71+
Get or create an NLSearchModel object for a given model_id.
72+
73+
Args:
74+
model_id (str): The ID of the NL search model.
75+
76+
Returns:
77+
NLSearchModel: The NLSearchModel object for the given ID.
78+
"""
79+
if model_id not in self.nl_search_models:
80+
self.nl_search_models[model_id] = NLSearchModel(
81+
self.api_call,
82+
model_id,
83+
)
84+
return self.nl_search_models[model_id]
85+
86+
def create(self, model: NLSearchModelCreateSchema) -> NLSearchModelSchema:
87+
"""
88+
Create a new NL search model.
89+
90+
Args:
91+
model (NLSearchModelCreateSchema):
92+
The schema for creating the NL search model.
93+
94+
Returns:
95+
NLSearchModelSchema: The created NL search model.
96+
"""
97+
response = self.api_call.post(
98+
endpoint=NLSearchModels.resource_path,
99+
entity_type=NLSearchModelSchema,
100+
as_json=True,
101+
body=model,
102+
)
103+
return response
104+
105+
def retrieve(self) -> NLSearchModelsRetrieveSchema:
106+
"""
107+
Retrieve all NL search models.
108+
109+
Returns:
110+
NLSearchModelsRetrieveSchema: A list of all NL search models.
111+
"""
112+
response: NLSearchModelsRetrieveSchema = self.api_call.get(
113+
endpoint=NLSearchModels.resource_path,
114+
entity_type=NLSearchModelsRetrieveSchema,
115+
as_json=True,
116+
)
117+
return response

‎src/typesense/types/nl_search_model.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
"""NLSearchModel types for Typesense Python Client."""
2+
3+
import sys
4+
5+
if sys.version_info >= (3, 11):
6+
import typing
7+
else:
8+
import typing_extensions as typing
9+
10+
11+
class NLSearchModelBase(typing.TypedDict):
12+
"""
13+
Base schema with all possible fields for NL search models.
14+
15+
Attributes:
16+
model_name (str): Name of the LLM model.
17+
api_key (str): The LLM service's API Key.
18+
api_url (str): The API URL for the LLM service.
19+
max_bytes (int): The maximum number of bytes to send to the LLM.
20+
temperature (float): The temperature parameter for the LLM.
21+
system_prompt (str): The system prompt for the LLM.
22+
top_p (float): The top_p parameter (Google-specific).
23+
top_k (int): The top_k parameter (Google-specific).
24+
stop_sequences (list[str]): Stop sequences for the LLM (Google-specific).
25+
api_version (str): API version (Google-specific).
26+
project_id (str): GCP project ID (GCP Vertex AI specific).
27+
access_token (str): Access token for GCP (GCP Vertex AI specific).
28+
refresh_token (str): Refresh token for GCP (GCP Vertex AI specific).
29+
client_id (str): Client ID for GCP (GCP Vertex AI specific).
30+
client_secret (str): Client secret for GCP (GCP Vertex AI specific).
31+
region (str): Region for GCP (GCP Vertex AI specific).
32+
max_output_tokens (int): Maximum output tokens (GCP Vertex AI specific).
33+
account_id (str): Account ID (Cloudflare specific).
34+
"""
35+
36+
model_name: str
37+
api_key: typing.NotRequired[str]
38+
api_url: typing.NotRequired[str]
39+
max_bytes: typing.NotRequired[int]
40+
temperature: typing.NotRequired[float]
41+
system_prompt: typing.NotRequired[str]
42+
# Google-specific parameters
43+
top_p: typing.NotRequired[float]
44+
top_k: typing.NotRequired[int]
45+
stop_sequences: typing.NotRequired[typing.List[str]]
46+
api_version: typing.NotRequired[str]
47+
# GCP Vertex AI specific
48+
project_id: typing.NotRequired[str]
49+
access_token: typing.NotRequired[str]
50+
refresh_token: typing.NotRequired[str]
51+
client_id: typing.NotRequired[str]
52+
client_secret: typing.NotRequired[str]
53+
region: typing.NotRequired[str]
54+
max_output_tokens: typing.NotRequired[int]
55+
# Cloudflare specific
56+
account_id: typing.NotRequired[str]
57+
58+
59+
class NLSearchModelCreateSchema(NLSearchModelBase):
60+
"""
61+
Schema for creating a new NL search model.
62+
63+
Attributes:
64+
id (str): The custom ID of the model.
65+
"""
66+
67+
id: typing.NotRequired[str]
68+
69+
70+
class NLSearchModelUpdateSchema(typing.TypedDict):
71+
"""
72+
Base schema with all possible fields for NL search models.
73+
74+
Attributes:
75+
model_name (str): Name of the LLM model.
76+
api_key (str): The LLM service's API Key.
77+
api_url (str): The API URL for the LLM service.
78+
max_bytes (int): The maximum number of bytes to send to the LLM.
79+
temperature (float): The temperature parameter for the LLM.
80+
system_prompt (str): The system prompt for the LLM.
81+
top_p (float): The top_p parameter (Google-specific).
82+
top_k (int): The top_k parameter (Google-specific).
83+
stop_sequences (list[str]): Stop sequences for the LLM (Google-specific).
84+
api_version (str): API version (Google-specific).
85+
project_id (str): GCP project ID (GCP Vertex AI specific).
86+
access_token (str): Access token for GCP (GCP Vertex AI specific).
87+
refresh_token (str): Refresh token for GCP (GCP Vertex AI specific).
88+
client_id (str): Client ID for GCP (GCP Vertex AI specific).
89+
client_secret (str): Client secret for GCP (GCP Vertex AI specific).
90+
region (str): Region for GCP (GCP Vertex AI specific).
91+
max_output_tokens (int): Maximum output tokens (GCP Vertex AI specific).
92+
account_id (str): Account ID (Cloudflare specific).
93+
"""
94+
95+
model_name: typing.NotRequired[str]
96+
api_key: typing.NotRequired[str]
97+
api_url: typing.NotRequired[str]
98+
max_bytes: typing.NotRequired[int]
99+
temperature: typing.NotRequired[float]
100+
system_prompt: typing.NotRequired[str]
101+
# Google-specific parameters
102+
top_p: typing.NotRequired[float]
103+
top_k: typing.NotRequired[int]
104+
stop_sequences: typing.NotRequired[typing.List[str]]
105+
api_version: typing.NotRequired[str]
106+
# GCP Vertex AI specific
107+
project_id: typing.NotRequired[str]
108+
access_token: typing.NotRequired[str]
109+
refresh_token: typing.NotRequired[str]
110+
client_id: typing.NotRequired[str]
111+
client_secret: typing.NotRequired[str]
112+
region: typing.NotRequired[str]
113+
max_output_tokens: typing.NotRequired[int]
114+
# Cloudflare specific
115+
account_id: typing.NotRequired[str]
116+
117+
118+
class NLSearchModelDeleteSchema(typing.TypedDict):
119+
"""
120+
Schema for deleting an NL search model.
121+
122+
Attributes:
123+
id (str): The ID of the model.
124+
"""
125+
126+
id: str
127+
128+
129+
class NLSearchModelSchema(NLSearchModelBase):
130+
"""
131+
Schema for an NL search model.
132+
133+
Attributes:
134+
id (str): The ID of the model.
135+
"""
136+
137+
id: str
138+
139+
140+
NLSearchModelsRetrieveSchema = typing.List[NLSearchModelSchema]

0 commit comments

Comments
 (0)