feat: Add Spell Correction Spec to Google Cloud Enterprise Search connector (#8705)

This commit is contained in:
Holt Skinner 2023-08-03 12:38:45 -05:00 committed by GitHub
parent b2b71b0d35
commit 8a8917e0d9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -56,6 +56,16 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever):
SearchResponse.total_size is zero. SearchResponse.total_size is zero.
2 - Automatic query expansion built by the Search API. 2 - Automatic query expansion built by the Search API.
""" """
spell_correction_mode: int = Field(default=2, ge=0, le=2)
"""Specification to determine under which conditions query expansion should occur.
0 - Unspecified spell correction mode. In this case, server behavior defaults
to auto.
1 - Suggestion only. Search API will try to find a spell suggestion if there is any
and put in the `SearchResponse.corrected_query`.
The spell suggestion will not be used as the search query.
2 - Automatic spell correction built by the Search API.
Search will be based on the corrected query if found.
"""
credentials: Any = None credentials: Any = None
"""The default custom credentials (google.auth.credentials.Credentials) to use """The default custom credentials (google.auth.credentials.Credentials) to use
when making API calls. If not provided, credentials will be ascertained from when making API calls. If not provided, credentials will be ascertained from
@ -134,6 +144,10 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever):
condition=self.query_expansion_condition, condition=self.query_expansion_condition,
) )
spell_correction_spec = SearchRequest.SpellCorrectionSpec(
mode=self.spell_correction_mode
)
if self.get_extractive_answers: if self.get_extractive_answers:
extractive_content_spec = ( extractive_content_spec = (
SearchRequest.ContentSearchSpec.ExtractiveContentSpec( SearchRequest.ContentSearchSpec.ExtractiveContentSpec(
@ -158,6 +172,7 @@ class GoogleCloudEnterpriseSearchRetriever(BaseRetriever):
page_size=self.max_documents, page_size=self.max_documents,
content_search_spec=content_search_spec, content_search_spec=content_search_spec,
query_expansion_spec=query_expansion_spec, query_expansion_spec=query_expansion_spec,
spell_correction_spec=spell_correction_spec,
) )
def _get_relevant_documents( def _get_relevant_documents(