diff --git a/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/rank_service.rst b/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/rank_service.rst new file mode 100644 index 000000000000..649695ae0c96 --- /dev/null +++ b/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/rank_service.rst @@ -0,0 +1,6 @@ +RankService +----------------------------- + +.. automodule:: google.cloud.discoveryengine_v1alpha.services.rank_service + :members: + :inherited-members: diff --git a/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/services_.rst b/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/services_.rst index aa2f4b4bb8c9..7607a16f88c4 100644 --- a/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/services_.rst +++ b/packages/google-cloud-discoveryengine/docs/discoveryengine_v1alpha/services_.rst @@ -11,6 +11,7 @@ Services for Google Cloud Discoveryengine v1alpha API document_service engine_service estimate_billing_service + rank_service recommendation_service schema_service search_service diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/__init__.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/__init__.py index dffac0fccc78..ef04e9e9e8c4 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/__init__.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/__init__.py @@ -41,6 +41,7 @@ EstimateBillingServiceAsyncClient, EstimateBillingServiceClient, ) +from .services.rank_service import RankServiceAsyncClient, RankServiceClient from .services.recommendation_service import ( RecommendationServiceAsyncClient, RecommendationServiceClient, @@ -114,12 +115,13 @@ UpdateDataStoreRequest, UpdateDocumentProcessingConfigRequest, ) -from .types.document import Document +from .types.document import Document, ProcessedDocument from .types.document_processing_config import DocumentProcessingConfig from .types.document_service import ( CreateDocumentRequest, DeleteDocumentRequest, GetDocumentRequest, + GetProcessedDocumentRequest, ListDocumentsRequest, ListDocumentsResponse, UpdateDocumentRequest, @@ -147,7 +149,14 @@ ) from .types.import_config import ( BigQuerySource, + BigtableOptions, + BigtableSource, + CloudSqlSource, + FhirStoreSource, + FirestoreSource, GcsSource, + ImportCompletionSuggestionsMetadata, + ImportCompletionSuggestionsResponse, ImportDocumentsMetadata, ImportDocumentsRequest, ImportDocumentsResponse, @@ -158,6 +167,7 @@ ImportUserEventsMetadata, ImportUserEventsRequest, ImportUserEventsResponse, + SpannerSource, ) from .types.purge_config import ( PurgeDocumentsMetadata, @@ -171,6 +181,7 @@ PurgeUserEventsRequest, PurgeUserEventsResponse, ) +from .types.rank_service import RankingRecord, RankRequest, RankResponse from .types.recommendation_service import RecommendRequest, RecommendResponse from .types.schema import FieldConfig, Schema from .types.schema_service import ( @@ -248,6 +259,7 @@ "DocumentServiceAsyncClient", "EngineServiceAsyncClient", "EstimateBillingServiceAsyncClient", + "RankServiceAsyncClient", "RecommendationServiceAsyncClient", "SchemaServiceAsyncClient", "SearchServiceAsyncClient", @@ -264,8 +276,11 @@ "BatchVerifyTargetSitesRequest", "BatchVerifyTargetSitesResponse", "BigQuerySource", + "BigtableOptions", + "BigtableSource", "Chunk", "ChunkServiceClient", + "CloudSqlSource", "CollectUserEventRequest", "CompleteQueryRequest", "CompleteQueryResponse", @@ -321,7 +336,9 @@ "EstimateDataSizeResponse", "FetchDomainVerificationStatusRequest", "FetchDomainVerificationStatusResponse", + "FhirStoreSource", "FieldConfig", + "FirestoreSource", "GcsSource", "GetAclConfigRequest", "GetChunkRequest", @@ -330,12 +347,15 @@ "GetDocumentProcessingConfigRequest", "GetDocumentRequest", "GetEngineRequest", + "GetProcessedDocumentRequest", "GetSchemaRequest", "GetServingConfigRequest", "GetSiteSearchEngineRequest", "GetTargetSiteRequest", "GuidedSearchSpec", "IdpConfig", + "ImportCompletionSuggestionsMetadata", + "ImportCompletionSuggestionsResponse", "ImportDocumentsMetadata", "ImportDocumentsRequest", "ImportDocumentsResponse", @@ -369,6 +389,7 @@ "PanelInfo", "PauseEngineRequest", "Principal", + "ProcessedDocument", "PurgeDocumentsMetadata", "PurgeDocumentsRequest", "PurgeDocumentsResponse", @@ -379,6 +400,10 @@ "PurgeUserEventsMetadata", "PurgeUserEventsRequest", "PurgeUserEventsResponse", + "RankRequest", + "RankResponse", + "RankServiceClient", + "RankingRecord", "RecommendRequest", "RecommendResponse", "RecommendationServiceClient", @@ -402,6 +427,7 @@ "SiteSearchEngineServiceClient", "SiteVerificationInfo", "SolutionType", + "SpannerSource", "SuggestionDenyListEntry", "TargetSite", "TextInput", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/gapic_metadata.json b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/gapic_metadata.json index 8c6560f8ad4d..2b2e7e2a394d 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/gapic_metadata.json +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/gapic_metadata.json @@ -420,6 +420,11 @@ "get_document" ] }, + "GetProcessedDocument": { + "methods": [ + "get_processed_document" + ] + }, "ImportDocuments": { "methods": [ "import_documents" @@ -460,6 +465,11 @@ "get_document" ] }, + "GetProcessedDocument": { + "methods": [ + "get_processed_document" + ] + }, "ImportDocuments": { "methods": [ "import_documents" @@ -500,6 +510,11 @@ "get_document" ] }, + "GetProcessedDocument": { + "methods": [ + "get_processed_document" + ] + }, "ImportDocuments": { "methods": [ "import_documents" @@ -697,6 +712,40 @@ } } }, + "RankService": { + "clients": { + "grpc": { + "libraryClient": "RankServiceClient", + "rpcs": { + "Rank": { + "methods": [ + "rank" + ] + } + } + }, + "grpc-async": { + "libraryClient": "RankServiceAsyncClient", + "rpcs": { + "Rank": { + "methods": [ + "rank" + ] + } + } + }, + "rest": { + "libraryClient": "RankServiceClient", + "rpcs": { + "Rank": { + "methods": [ + "rank" + ] + } + } + } + } + }, "RecommendationService": { "clients": { "grpc": { diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/acl_config_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/acl_config_service/transports/rest.py index 4f516e5b7554..38a3b6d8ef63 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/acl_config_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/acl_config_service/transports/rest.py @@ -551,6 +551,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/chunk_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/chunk_service/transports/rest.py index 2f791c760cba..41341880aa7d 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/chunk_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/chunk_service/transports/rest.py @@ -556,6 +556,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/completion_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/completion_service/transports/rest.py index 36b146494576..22ef00f28f3c 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/completion_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/completion_service/transports/rest.py @@ -382,6 +382,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -885,6 +889,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/conversational_search_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/conversational_search_service/transports/rest.py index e6b7f8459831..f1214ed848b3 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/conversational_search_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/conversational_search_service/transports/rest.py @@ -1163,6 +1163,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/data_store_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/data_store_service/transports/rest.py index d64ce6962b24..5543795f24f1 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/data_store_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/data_store_service/transports/rest.py @@ -511,6 +511,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -1450,6 +1454,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/async_client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/async_client.py index cf0b53c800b2..6cae338eb123 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/async_client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/async_client.py @@ -83,6 +83,8 @@ class DocumentServiceAsyncClient: parse_branch_path = staticmethod(DocumentServiceClient.parse_branch_path) document_path = staticmethod(DocumentServiceClient.document_path) parse_document_path = staticmethod(DocumentServiceClient.parse_document_path) + fhir_store_path = staticmethod(DocumentServiceClient.fhir_store_path) + parse_fhir_store_path = staticmethod(DocumentServiceClient.parse_fhir_store_path) common_billing_account_path = staticmethod( DocumentServiceClient.common_billing_account_path ) @@ -1145,6 +1147,128 @@ async def sample_purge_documents(): # Done; return the response. return response + async def get_processed_document( + self, + request: Optional[ + Union[document_service.GetProcessedDocumentRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> document.ProcessedDocument: + r"""Gets the parsed layout information for a + [Document][google.cloud.discoveryengine.v1alpha.Document]. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import discoveryengine_v1alpha + + async def sample_get_processed_document(): + # Create a client + client = discoveryengine_v1alpha.DocumentServiceAsyncClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.GetProcessedDocumentRequest( + name="name_value", + processed_document_type="CHUNKED_DOCUMENT", + ) + + # Make the request + response = await client.get_processed_document(request=request) + + # Handle the response + print(response) + + Args: + request (Optional[Union[google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest, dict]]): + The request object. Request message for + [DocumentService.GetDocument][google.cloud.discoveryengine.v1alpha.DocumentService.GetDocument] + method. + name (:class:`str`): + Required. Full resource name of + [Document][google.cloud.discoveryengine.v1alpha.Document], + such as + ``projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/branches/{branch}/documents/{document}``. + + If the caller does not have permission to access the + [Document][google.cloud.discoveryengine.v1alpha.Document], + regardless of whether or not it exists, a + ``PERMISSION_DENIED`` error is returned. + + If the requested + [Document][google.cloud.discoveryengine.v1alpha.Document] + does not exist, a ``NOT_FOUND`` error is returned. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.discoveryengine_v1alpha.types.ProcessedDocument: + Document captures all raw metadata + information of items to be recommended + or searched. + + """ + # Create or coerce a protobuf request object. + # Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + request = document_service.GetProcessedDocumentRequest(request) + + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.get_processed_document, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + async def list_operations( self, request: Optional[operations_pb2.ListOperationsRequest] = None, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/client.py index 208f80f7e439..441460593ca8 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/client.py @@ -249,6 +249,30 @@ def parse_document_path(path: str) -> Dict[str, str]: ) return m.groupdict() if m else {} + @staticmethod + def fhir_store_path( + project: str, + location: str, + dataset: str, + fhir_store: str, + ) -> str: + """Returns a fully-qualified fhir_store string.""" + return "projects/{project}/locations/{location}/datasets/{dataset}/fhirStores/{fhir_store}".format( + project=project, + location=location, + dataset=dataset, + fhir_store=fhir_store, + ) + + @staticmethod + def parse_fhir_store_path(path: str) -> Dict[str, str]: + """Parses a fhir_store path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/locations/(?P.+?)/datasets/(?P.+?)/fhirStores/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + @staticmethod def common_billing_account_path( billing_account: str, @@ -1577,6 +1601,128 @@ def sample_purge_documents(): # Done; return the response. return response + def get_processed_document( + self, + request: Optional[ + Union[document_service.GetProcessedDocumentRequest, dict] + ] = None, + *, + name: Optional[str] = None, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> document.ProcessedDocument: + r"""Gets the parsed layout information for a + [Document][google.cloud.discoveryengine.v1alpha.Document]. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import discoveryengine_v1alpha + + def sample_get_processed_document(): + # Create a client + client = discoveryengine_v1alpha.DocumentServiceClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.GetProcessedDocumentRequest( + name="name_value", + processed_document_type="CHUNKED_DOCUMENT", + ) + + # Make the request + response = client.get_processed_document(request=request) + + # Handle the response + print(response) + + Args: + request (Union[google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest, dict]): + The request object. Request message for + [DocumentService.GetDocument][google.cloud.discoveryengine.v1alpha.DocumentService.GetDocument] + method. + name (str): + Required. Full resource name of + [Document][google.cloud.discoveryengine.v1alpha.Document], + such as + ``projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/branches/{branch}/documents/{document}``. + + If the caller does not have permission to access the + [Document][google.cloud.discoveryengine.v1alpha.Document], + regardless of whether or not it exists, a + ``PERMISSION_DENIED`` error is returned. + + If the requested + [Document][google.cloud.discoveryengine.v1alpha.Document] + does not exist, a ``NOT_FOUND`` error is returned. + + This corresponds to the ``name`` field + on the ``request`` instance; if ``request`` is provided, this + should not be set. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.discoveryengine_v1alpha.types.ProcessedDocument: + Document captures all raw metadata + information of items to be recommended + or searched. + + """ + # Create or coerce a protobuf request object. + # Quick check: If we got a request object, we should *not* have + # gotten any keyword arguments that map to the request. + has_flattened_params = any([name]) + if request is not None and has_flattened_params: + raise ValueError( + "If the `request` argument is set, then none of " + "the individual field arguments should be set." + ) + + # Minor optimization to avoid making a copy if the user passes + # in a document_service.GetProcessedDocumentRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, document_service.GetProcessedDocumentRequest): + request = document_service.GetProcessedDocumentRequest(request) + # If we have keyword arguments corresponding to fields on the + # request, apply these. + if name is not None: + request.name = name + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.get_processed_document] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + def __enter__(self) -> "DocumentServiceClient": return self diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/base.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/base.py index 4b477f330d04..bb2471bfc1a0 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/base.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/base.py @@ -179,6 +179,11 @@ def _prep_wrapped_messages(self, client_info): default_timeout=None, client_info=client_info, ), + self.get_processed_document: gapic_v1.method.wrap_method( + self.get_processed_document, + default_timeout=None, + client_info=client_info, + ), } def close(self): @@ -261,6 +266,15 @@ def purge_documents( ]: raise NotImplementedError() + @property + def get_processed_document( + self, + ) -> Callable[ + [document_service.GetProcessedDocumentRequest], + Union[document.ProcessedDocument, Awaitable[document.ProcessedDocument]], + ]: + raise NotImplementedError() + @property def list_operations( self, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc.py index 2f5881ca1f31..fb4312c79d5c 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc.py @@ -468,6 +468,35 @@ def purge_documents( ) return self._stubs["purge_documents"] + @property + def get_processed_document( + self, + ) -> Callable[ + [document_service.GetProcessedDocumentRequest], document.ProcessedDocument + ]: + r"""Return a callable for the get processed document method over gRPC. + + Gets the parsed layout information for a + [Document][google.cloud.discoveryengine.v1alpha.Document]. + + Returns: + Callable[[~.GetProcessedDocumentRequest], + ~.ProcessedDocument]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_processed_document" not in self._stubs: + self._stubs["get_processed_document"] = self.grpc_channel.unary_unary( + "/google.cloud.discoveryengine.v1alpha.DocumentService/GetProcessedDocument", + request_serializer=document_service.GetProcessedDocumentRequest.serialize, + response_deserializer=document.ProcessedDocument.deserialize, + ) + return self._stubs["get_processed_document"] + def close(self): self.grpc_channel.close() diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc_asyncio.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc_asyncio.py index 1748888d465f..a7225aa7eb18 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc_asyncio.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/grpc_asyncio.py @@ -482,6 +482,36 @@ def purge_documents( ) return self._stubs["purge_documents"] + @property + def get_processed_document( + self, + ) -> Callable[ + [document_service.GetProcessedDocumentRequest], + Awaitable[document.ProcessedDocument], + ]: + r"""Return a callable for the get processed document method over gRPC. + + Gets the parsed layout information for a + [Document][google.cloud.discoveryengine.v1alpha.Document]. + + Returns: + Callable[[~.GetProcessedDocumentRequest], + Awaitable[~.ProcessedDocument]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_processed_document" not in self._stubs: + self._stubs["get_processed_document"] = self.grpc_channel.unary_unary( + "/google.cloud.discoveryengine.v1alpha.DocumentService/GetProcessedDocument", + request_serializer=document_service.GetProcessedDocumentRequest.serialize, + response_deserializer=document.ProcessedDocument.deserialize, + ) + return self._stubs["get_processed_document"] + def close(self): return self.grpc_channel.close() diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/rest.py index d5a665e1bae6..ac96353a83ad 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/document_service/transports/rest.py @@ -99,6 +99,14 @@ def post_get_document(self, response): logging.log(f"Received response: {response}") return response + def pre_get_processed_document(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_get_processed_document(self, response): + logging.log(f"Received response: {response}") + return response + def pre_import_documents(self, request, metadata): logging.log(f"Received request: {request}") return request, metadata @@ -193,6 +201,29 @@ def post_get_document(self, response: document.Document) -> document.Document: """ return response + def pre_get_processed_document( + self, + request: document_service.GetProcessedDocumentRequest, + metadata: Sequence[Tuple[str, str]], + ) -> Tuple[document_service.GetProcessedDocumentRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for get_processed_document + + Override in a subclass to manipulate the request or metadata + before they are sent to the DocumentService server. + """ + return request, metadata + + def post_get_processed_document( + self, response: document.ProcessedDocument + ) -> document.ProcessedDocument: + """Post-rpc interceptor for get_processed_document + + Override in a subclass to manipulate the response + after it is returned by the DocumentService server but before + it is returned to user code. + """ + return response + def pre_import_documents( self, request: import_config.ImportDocumentsRequest, @@ -490,6 +521,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -852,6 +887,104 @@ def __call__( resp = self._interceptor.post_get_document(resp) return resp + class _GetProcessedDocument(DocumentServiceRestStub): + def __hash__(self): + return hash("GetProcessedDocument") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = { + "processedDocumentType": {}, + } + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + def __call__( + self, + request: document_service.GetProcessedDocumentRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> document.ProcessedDocument: + r"""Call the get processed document method over HTTP. + + Args: + request (~.document_service.GetProcessedDocumentRequest): + The request object. Request message for + [DocumentService.GetDocument][google.cloud.discoveryengine.v1alpha.DocumentService.GetDocument] + method. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + ~.document.ProcessedDocument: + Document captures all raw metadata + information of items to be recommended + or searched. + + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/branches/*/documents/*}:getProcessedDocument", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/branches/*/documents/*}:getProcessedDocument", + }, + ] + request, metadata = self._interceptor.pre_get_processed_document( + request, metadata + ) + pb_request = document_service.GetProcessedDocumentRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update(self._get_unset_required_fields(query_params)) + + query_params["$alt"] = "json;enum-encoding=int" + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = document.ProcessedDocument() + pb_resp = document.ProcessedDocument.pb(resp) + + json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) + resp = self._interceptor.post_get_processed_document(resp) + return resp + class _ImportDocuments(DocumentServiceRestStub): def __hash__(self): return hash("ImportDocuments") @@ -1272,6 +1405,16 @@ def get_document( # In C++ this would require a dynamic_cast return self._GetDocument(self._session, self._host, self._interceptor) # type: ignore + @property + def get_processed_document( + self, + ) -> Callable[ + [document_service.GetProcessedDocumentRequest], document.ProcessedDocument + ]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._GetProcessedDocument(self._session, self._host, self._interceptor) # type: ignore + @property def import_documents( self, @@ -1383,6 +1526,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/async_client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/async_client.py index 6d0d94c10d6c..d4f10b9519d3 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/async_client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/async_client.py @@ -298,7 +298,7 @@ async def sample_create_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.CreateEngineRequest( parent="parent_value", @@ -591,7 +591,7 @@ async def sample_update_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateEngineRequest( engine=engine, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/client.py index 530179d20b6c..008db6ce6821 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/client.py @@ -733,7 +733,7 @@ def sample_create_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.CreateEngineRequest( parent="parent_value", @@ -1026,7 +1026,7 @@ def sample_update_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateEngineRequest( engine=engine, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/transports/rest.py index ef302e425d42..fa0ea17da226 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/engine_service/transports/rest.py @@ -524,6 +524,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -1500,6 +1504,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/estimate_billing_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/estimate_billing_service/transports/rest.py index d93088ce0bba..30d076310d4d 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/estimate_billing_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/estimate_billing_service/transports/rest.py @@ -314,6 +314,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -583,6 +587,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/__init__.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/__init__.py new file mode 100644 index 000000000000..f34036ed49c9 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/__init__.py @@ -0,0 +1,22 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from .async_client import RankServiceAsyncClient +from .client import RankServiceClient + +__all__ = ( + "RankServiceClient", + "RankServiceAsyncClient", +) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/async_client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/async_client.py new file mode 100644 index 000000000000..255c96741a2f --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/async_client.py @@ -0,0 +1,469 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from collections import OrderedDict +import functools +import re +from typing import ( + Dict, + Mapping, + MutableMapping, + MutableSequence, + Optional, + Sequence, + Tuple, + Type, + Union, +) + +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import retry_async as retries +from google.api_core.client_options import ClientOptions +from google.auth import credentials as ga_credentials # type: ignore +from google.oauth2 import service_account # type: ignore + +from google.cloud.discoveryengine_v1alpha import gapic_version as package_version + +try: + OptionalRetry = Union[retries.AsyncRetry, gapic_v1.method._MethodDefault, None] +except AttributeError: # pragma: NO COVER + OptionalRetry = Union[retries.AsyncRetry, object, None] # type: ignore + +from google.cloud.location import locations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore + +from google.cloud.discoveryengine_v1alpha.types import rank_service + +from .client import RankServiceClient +from .transports.base import DEFAULT_CLIENT_INFO, RankServiceTransport +from .transports.grpc_asyncio import RankServiceGrpcAsyncIOTransport + + +class RankServiceAsyncClient: + """Service for ranking text records.""" + + _client: RankServiceClient + + # Copy defaults from the synchronous client for use here. + # Note: DEFAULT_ENDPOINT is deprecated. Use _DEFAULT_ENDPOINT_TEMPLATE instead. + DEFAULT_ENDPOINT = RankServiceClient.DEFAULT_ENDPOINT + DEFAULT_MTLS_ENDPOINT = RankServiceClient.DEFAULT_MTLS_ENDPOINT + _DEFAULT_ENDPOINT_TEMPLATE = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE + _DEFAULT_UNIVERSE = RankServiceClient._DEFAULT_UNIVERSE + + ranking_config_path = staticmethod(RankServiceClient.ranking_config_path) + parse_ranking_config_path = staticmethod( + RankServiceClient.parse_ranking_config_path + ) + common_billing_account_path = staticmethod( + RankServiceClient.common_billing_account_path + ) + parse_common_billing_account_path = staticmethod( + RankServiceClient.parse_common_billing_account_path + ) + common_folder_path = staticmethod(RankServiceClient.common_folder_path) + parse_common_folder_path = staticmethod(RankServiceClient.parse_common_folder_path) + common_organization_path = staticmethod(RankServiceClient.common_organization_path) + parse_common_organization_path = staticmethod( + RankServiceClient.parse_common_organization_path + ) + common_project_path = staticmethod(RankServiceClient.common_project_path) + parse_common_project_path = staticmethod( + RankServiceClient.parse_common_project_path + ) + common_location_path = staticmethod(RankServiceClient.common_location_path) + parse_common_location_path = staticmethod( + RankServiceClient.parse_common_location_path + ) + + @classmethod + def from_service_account_info(cls, info: dict, *args, **kwargs): + """Creates an instance of this client using the provided credentials + info. + + Args: + info (dict): The service account private key info. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + RankServiceAsyncClient: The constructed client. + """ + return RankServiceClient.from_service_account_info.__func__(RankServiceAsyncClient, info, *args, **kwargs) # type: ignore + + @classmethod + def from_service_account_file(cls, filename: str, *args, **kwargs): + """Creates an instance of this client using the provided credentials + file. + + Args: + filename (str): The path to the service account private key json + file. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + RankServiceAsyncClient: The constructed client. + """ + return RankServiceClient.from_service_account_file.__func__(RankServiceAsyncClient, filename, *args, **kwargs) # type: ignore + + from_service_account_json = from_service_account_file + + @classmethod + def get_mtls_endpoint_and_cert_source( + cls, client_options: Optional[ClientOptions] = None + ): + """Return the API endpoint and client cert source for mutual TLS. + + The client cert source is determined in the following order: + (1) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is not "true", the + client cert source is None. + (2) if `client_options.client_cert_source` is provided, use the provided one; if the + default client cert source exists, use the default one; otherwise the client cert + source is None. + + The API endpoint is determined in the following order: + (1) if `client_options.api_endpoint` if provided, use the provided one. + (2) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is "always", use the + default mTLS endpoint; if the environment variable is "never", use the default API + endpoint; otherwise if client cert source exists, use the default mTLS endpoint, otherwise + use the default API endpoint. + + More details can be found at https://google.aip.dev/auth/4114. + + Args: + client_options (google.api_core.client_options.ClientOptions): Custom options for the + client. Only the `api_endpoint` and `client_cert_source` properties may be used + in this method. + + Returns: + Tuple[str, Callable[[], Tuple[bytes, bytes]]]: returns the API endpoint and the + client cert source to use. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If any errors happen. + """ + return RankServiceClient.get_mtls_endpoint_and_cert_source(client_options) # type: ignore + + @property + def transport(self) -> RankServiceTransport: + """Returns the transport used by the client instance. + + Returns: + RankServiceTransport: The transport used by the client instance. + """ + return self._client.transport + + @property + def api_endpoint(self): + """Return the API endpoint used by the client instance. + + Returns: + str: The API endpoint used by the client instance. + """ + return self._client._api_endpoint + + @property + def universe_domain(self) -> str: + """Return the universe domain used by the client instance. + + Returns: + str: The universe domain used + by the client instance. + """ + return self._client._universe_domain + + get_transport_class = functools.partial( + type(RankServiceClient).get_transport_class, type(RankServiceClient) + ) + + def __init__( + self, + *, + credentials: Optional[ga_credentials.Credentials] = None, + transport: Union[str, RankServiceTransport] = "grpc_asyncio", + client_options: Optional[ClientOptions] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiates the rank service async client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, ~.RankServiceTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]): + Custom options for the client. + + 1. The ``api_endpoint`` property can be used to override the + default endpoint provided by the client when ``transport`` is + not explicitly provided. Only if this property is not set and + ``transport`` was not explicitly provided, the endpoint is + determined by the GOOGLE_API_USE_MTLS_ENDPOINT environment + variable, which have one of the following values: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto-switch to the + default mTLS endpoint if client certificate is present; this is + the default value). + + 2. If the GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide a client certificate for mTLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + + 3. The ``universe_domain`` property can be used to override the + default "googleapis.com" universe. Note that ``api_endpoint`` + property still takes precedence; and ``universe_domain`` is + currently not supported for mTLS. + + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + """ + self._client = RankServiceClient( + credentials=credentials, + transport=transport, + client_options=client_options, + client_info=client_info, + ) + + async def rank( + self, + request: Optional[Union[rank_service.RankRequest, dict]] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> rank_service.RankResponse: + r"""Ranks a list of text records based on the given input + query. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import discoveryengine_v1alpha + + async def sample_rank(): + # Create a client + client = discoveryengine_v1alpha.RankServiceAsyncClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.RankRequest( + ranking_config="ranking_config_value", + ) + + # Make the request + response = await client.rank(request=request) + + # Handle the response + print(response) + + Args: + request (Optional[Union[google.cloud.discoveryengine_v1alpha.types.RankRequest, dict]]): + The request object. Request message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.discoveryengine_v1alpha.types.RankResponse: + Response message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + """ + # Create or coerce a protobuf request object. + request = rank_service.RankRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.rank, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("ranking_config", request.ranking_config),) + ), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + async def list_operations( + self, + request: Optional[operations_pb2.ListOperationsRequest] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.ListOperationsResponse: + r"""Lists operations that match the specified filter in the request. + + Args: + request (:class:`~.operations_pb2.ListOperationsRequest`): + The request object. Request message for + `ListOperations` method. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + Returns: + ~.operations_pb2.ListOperationsResponse: + Response message for ``ListOperations`` method. + """ + # Create or coerce a protobuf request object. + # The request isn't a proto-plus wrapped type, + # so it must be constructed via keyword expansion. + if isinstance(request, dict): + request = operations_pb2.ListOperationsRequest(**request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.list_operations, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + async def get_operation( + self, + request: Optional[operations_pb2.GetOperationRequest] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.Operation: + r"""Gets the latest state of a long-running operation. + + Args: + request (:class:`~.operations_pb2.GetOperationRequest`): + The request object. Request message for + `GetOperation` method. + retry (google.api_core.retry_async.AsyncRetry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + Returns: + ~.operations_pb2.Operation: + An ``Operation`` object. + """ + # Create or coerce a protobuf request object. + # The request isn't a proto-plus wrapped type, + # so it must be constructed via keyword expansion. + if isinstance(request, dict): + request = operations_pb2.GetOperationRequest(**request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method_async.wrap_method( + self._client._transport.get_operation, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._client._validate_universe_domain() + + # Send the request. + response = await rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + async def __aenter__(self) -> "RankServiceAsyncClient": + return self + + async def __aexit__(self, exc_type, exc, tb): + await self.transport.close() + + +DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=package_version.__version__ +) + + +__all__ = ("RankServiceAsyncClient",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/client.py new file mode 100644 index 000000000000..1d5ad3325014 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/client.py @@ -0,0 +1,887 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from collections import OrderedDict +import os +import re +from typing import ( + Dict, + Mapping, + MutableMapping, + MutableSequence, + Optional, + Sequence, + Tuple, + Type, + Union, + cast, +) +import warnings + +from google.api_core import client_options as client_options_lib +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials # type: ignore +from google.auth.exceptions import MutualTLSChannelError # type: ignore +from google.auth.transport import mtls # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.oauth2 import service_account # type: ignore + +from google.cloud.discoveryengine_v1alpha import gapic_version as package_version + +try: + OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None] +except AttributeError: # pragma: NO COVER + OptionalRetry = Union[retries.Retry, object, None] # type: ignore + +from google.cloud.location import locations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore + +from google.cloud.discoveryengine_v1alpha.types import rank_service + +from .transports.base import DEFAULT_CLIENT_INFO, RankServiceTransport +from .transports.grpc import RankServiceGrpcTransport +from .transports.grpc_asyncio import RankServiceGrpcAsyncIOTransport +from .transports.rest import RankServiceRestTransport + + +class RankServiceClientMeta(type): + """Metaclass for the RankService client. + + This provides class-level methods for building and retrieving + support objects (e.g. transport) without polluting the client instance + objects. + """ + + _transport_registry = OrderedDict() # type: Dict[str, Type[RankServiceTransport]] + _transport_registry["grpc"] = RankServiceGrpcTransport + _transport_registry["grpc_asyncio"] = RankServiceGrpcAsyncIOTransport + _transport_registry["rest"] = RankServiceRestTransport + + def get_transport_class( + cls, + label: Optional[str] = None, + ) -> Type[RankServiceTransport]: + """Returns an appropriate transport class. + + Args: + label: The name of the desired transport. If none is + provided, then the first transport in the registry is used. + + Returns: + The transport class to use. + """ + # If a specific transport is requested, return that one. + if label: + return cls._transport_registry[label] + + # No transport is requested; return the default (that is, the first one + # in the dictionary). + return next(iter(cls._transport_registry.values())) + + +class RankServiceClient(metaclass=RankServiceClientMeta): + """Service for ranking text records.""" + + @staticmethod + def _get_default_mtls_endpoint(api_endpoint): + """Converts api endpoint to mTLS endpoint. + + Convert "*.sandbox.googleapis.com" and "*.googleapis.com" to + "*.mtls.sandbox.googleapis.com" and "*.mtls.googleapis.com" respectively. + Args: + api_endpoint (Optional[str]): the api endpoint to convert. + Returns: + str: converted mTLS api endpoint. + """ + if not api_endpoint: + return api_endpoint + + mtls_endpoint_re = re.compile( + r"(?P[^.]+)(?P\.mtls)?(?P\.sandbox)?(?P\.googleapis\.com)?" + ) + + m = mtls_endpoint_re.match(api_endpoint) + name, mtls, sandbox, googledomain = m.groups() + if mtls or not googledomain: + return api_endpoint + + if sandbox: + return api_endpoint.replace( + "sandbox.googleapis.com", "mtls.sandbox.googleapis.com" + ) + + return api_endpoint.replace(".googleapis.com", ".mtls.googleapis.com") + + # Note: DEFAULT_ENDPOINT is deprecated. Use _DEFAULT_ENDPOINT_TEMPLATE instead. + DEFAULT_ENDPOINT = "discoveryengine.googleapis.com" + DEFAULT_MTLS_ENDPOINT = _get_default_mtls_endpoint.__func__( # type: ignore + DEFAULT_ENDPOINT + ) + + _DEFAULT_ENDPOINT_TEMPLATE = "discoveryengine.{UNIVERSE_DOMAIN}" + _DEFAULT_UNIVERSE = "googleapis.com" + + @classmethod + def from_service_account_info(cls, info: dict, *args, **kwargs): + """Creates an instance of this client using the provided credentials + info. + + Args: + info (dict): The service account private key info. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + RankServiceClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_info(info) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + @classmethod + def from_service_account_file(cls, filename: str, *args, **kwargs): + """Creates an instance of this client using the provided credentials + file. + + Args: + filename (str): The path to the service account private key json + file. + args: Additional arguments to pass to the constructor. + kwargs: Additional arguments to pass to the constructor. + + Returns: + RankServiceClient: The constructed client. + """ + credentials = service_account.Credentials.from_service_account_file(filename) + kwargs["credentials"] = credentials + return cls(*args, **kwargs) + + from_service_account_json = from_service_account_file + + @property + def transport(self) -> RankServiceTransport: + """Returns the transport used by the client instance. + + Returns: + RankServiceTransport: The transport used by the client + instance. + """ + return self._transport + + @staticmethod + def ranking_config_path( + project: str, + location: str, + ranking_config: str, + ) -> str: + """Returns a fully-qualified ranking_config string.""" + return "projects/{project}/locations/{location}/rankingConfigs/{ranking_config}".format( + project=project, + location=location, + ranking_config=ranking_config, + ) + + @staticmethod + def parse_ranking_config_path(path: str) -> Dict[str, str]: + """Parses a ranking_config path into its component segments.""" + m = re.match( + r"^projects/(?P.+?)/locations/(?P.+?)/rankingConfigs/(?P.+?)$", + path, + ) + return m.groupdict() if m else {} + + @staticmethod + def common_billing_account_path( + billing_account: str, + ) -> str: + """Returns a fully-qualified billing_account string.""" + return "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + + @staticmethod + def parse_common_billing_account_path(path: str) -> Dict[str, str]: + """Parse a billing_account path into its component segments.""" + m = re.match(r"^billingAccounts/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_folder_path( + folder: str, + ) -> str: + """Returns a fully-qualified folder string.""" + return "folders/{folder}".format( + folder=folder, + ) + + @staticmethod + def parse_common_folder_path(path: str) -> Dict[str, str]: + """Parse a folder path into its component segments.""" + m = re.match(r"^folders/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_organization_path( + organization: str, + ) -> str: + """Returns a fully-qualified organization string.""" + return "organizations/{organization}".format( + organization=organization, + ) + + @staticmethod + def parse_common_organization_path(path: str) -> Dict[str, str]: + """Parse a organization path into its component segments.""" + m = re.match(r"^organizations/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_project_path( + project: str, + ) -> str: + """Returns a fully-qualified project string.""" + return "projects/{project}".format( + project=project, + ) + + @staticmethod + def parse_common_project_path(path: str) -> Dict[str, str]: + """Parse a project path into its component segments.""" + m = re.match(r"^projects/(?P.+?)$", path) + return m.groupdict() if m else {} + + @staticmethod + def common_location_path( + project: str, + location: str, + ) -> str: + """Returns a fully-qualified location string.""" + return "projects/{project}/locations/{location}".format( + project=project, + location=location, + ) + + @staticmethod + def parse_common_location_path(path: str) -> Dict[str, str]: + """Parse a location path into its component segments.""" + m = re.match(r"^projects/(?P.+?)/locations/(?P.+?)$", path) + return m.groupdict() if m else {} + + @classmethod + def get_mtls_endpoint_and_cert_source( + cls, client_options: Optional[client_options_lib.ClientOptions] = None + ): + """Deprecated. Return the API endpoint and client cert source for mutual TLS. + + The client cert source is determined in the following order: + (1) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is not "true", the + client cert source is None. + (2) if `client_options.client_cert_source` is provided, use the provided one; if the + default client cert source exists, use the default one; otherwise the client cert + source is None. + + The API endpoint is determined in the following order: + (1) if `client_options.api_endpoint` if provided, use the provided one. + (2) if `GOOGLE_API_USE_CLIENT_CERTIFICATE` environment variable is "always", use the + default mTLS endpoint; if the environment variable is "never", use the default API + endpoint; otherwise if client cert source exists, use the default mTLS endpoint, otherwise + use the default API endpoint. + + More details can be found at https://google.aip.dev/auth/4114. + + Args: + client_options (google.api_core.client_options.ClientOptions): Custom options for the + client. Only the `api_endpoint` and `client_cert_source` properties may be used + in this method. + + Returns: + Tuple[str, Callable[[], Tuple[bytes, bytes]]]: returns the API endpoint and the + client cert source to use. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If any errors happen. + """ + + warnings.warn( + "get_mtls_endpoint_and_cert_source is deprecated. Use the api_endpoint property instead.", + DeprecationWarning, + ) + if client_options is None: + client_options = client_options_lib.ClientOptions() + use_client_cert = os.getenv("GOOGLE_API_USE_CLIENT_CERTIFICATE", "false") + use_mtls_endpoint = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto") + if use_client_cert not in ("true", "false"): + raise ValueError( + "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`" + ) + if use_mtls_endpoint not in ("auto", "never", "always"): + raise MutualTLSChannelError( + "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`" + ) + + # Figure out the client cert source to use. + client_cert_source = None + if use_client_cert == "true": + if client_options.client_cert_source: + client_cert_source = client_options.client_cert_source + elif mtls.has_default_client_cert_source(): + client_cert_source = mtls.default_client_cert_source() + + # Figure out which api endpoint to use. + if client_options.api_endpoint is not None: + api_endpoint = client_options.api_endpoint + elif use_mtls_endpoint == "always" or ( + use_mtls_endpoint == "auto" and client_cert_source + ): + api_endpoint = cls.DEFAULT_MTLS_ENDPOINT + else: + api_endpoint = cls.DEFAULT_ENDPOINT + + return api_endpoint, client_cert_source + + @staticmethod + def _read_environment_variables(): + """Returns the environment variables used by the client. + + Returns: + Tuple[bool, str, str]: returns the GOOGLE_API_USE_CLIENT_CERTIFICATE, + GOOGLE_API_USE_MTLS_ENDPOINT, and GOOGLE_CLOUD_UNIVERSE_DOMAIN environment variables. + + Raises: + ValueError: If GOOGLE_API_USE_CLIENT_CERTIFICATE is not + any of ["true", "false"]. + google.auth.exceptions.MutualTLSChannelError: If GOOGLE_API_USE_MTLS_ENDPOINT + is not any of ["auto", "never", "always"]. + """ + use_client_cert = os.getenv( + "GOOGLE_API_USE_CLIENT_CERTIFICATE", "false" + ).lower() + use_mtls_endpoint = os.getenv("GOOGLE_API_USE_MTLS_ENDPOINT", "auto").lower() + universe_domain_env = os.getenv("GOOGLE_CLOUD_UNIVERSE_DOMAIN") + if use_client_cert not in ("true", "false"): + raise ValueError( + "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`" + ) + if use_mtls_endpoint not in ("auto", "never", "always"): + raise MutualTLSChannelError( + "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`" + ) + return use_client_cert == "true", use_mtls_endpoint, universe_domain_env + + @staticmethod + def _get_client_cert_source(provided_cert_source, use_cert_flag): + """Return the client cert source to be used by the client. + + Args: + provided_cert_source (bytes): The client certificate source provided. + use_cert_flag (bool): A flag indicating whether to use the client certificate. + + Returns: + bytes or None: The client cert source to be used by the client. + """ + client_cert_source = None + if use_cert_flag: + if provided_cert_source: + client_cert_source = provided_cert_source + elif mtls.has_default_client_cert_source(): + client_cert_source = mtls.default_client_cert_source() + return client_cert_source + + @staticmethod + def _get_api_endpoint( + api_override, client_cert_source, universe_domain, use_mtls_endpoint + ): + """Return the API endpoint used by the client. + + Args: + api_override (str): The API endpoint override. If specified, this is always + the return value of this function and the other arguments are not used. + client_cert_source (bytes): The client certificate source used by the client. + universe_domain (str): The universe domain used by the client. + use_mtls_endpoint (str): How to use the mTLS endpoint, which depends also on the other parameters. + Possible values are "always", "auto", or "never". + + Returns: + str: The API endpoint to be used by the client. + """ + if api_override is not None: + api_endpoint = api_override + elif use_mtls_endpoint == "always" or ( + use_mtls_endpoint == "auto" and client_cert_source + ): + _default_universe = RankServiceClient._DEFAULT_UNIVERSE + if universe_domain != _default_universe: + raise MutualTLSChannelError( + f"mTLS is not supported in any universe other than {_default_universe}." + ) + api_endpoint = RankServiceClient.DEFAULT_MTLS_ENDPOINT + else: + api_endpoint = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=universe_domain + ) + return api_endpoint + + @staticmethod + def _get_universe_domain( + client_universe_domain: Optional[str], universe_domain_env: Optional[str] + ) -> str: + """Return the universe domain used by the client. + + Args: + client_universe_domain (Optional[str]): The universe domain configured via the client options. + universe_domain_env (Optional[str]): The universe domain configured via the "GOOGLE_CLOUD_UNIVERSE_DOMAIN" environment variable. + + Returns: + str: The universe domain to be used by the client. + + Raises: + ValueError: If the universe domain is an empty string. + """ + universe_domain = RankServiceClient._DEFAULT_UNIVERSE + if client_universe_domain is not None: + universe_domain = client_universe_domain + elif universe_domain_env is not None: + universe_domain = universe_domain_env + if len(universe_domain.strip()) == 0: + raise ValueError("Universe Domain cannot be an empty string.") + return universe_domain + + @staticmethod + def _compare_universes( + client_universe: str, credentials: ga_credentials.Credentials + ) -> bool: + """Returns True iff the universe domains used by the client and credentials match. + + Args: + client_universe (str): The universe domain configured via the client options. + credentials (ga_credentials.Credentials): The credentials being used in the client. + + Returns: + bool: True iff client_universe matches the universe in credentials. + + Raises: + ValueError: when client_universe does not match the universe in credentials. + """ + + default_universe = RankServiceClient._DEFAULT_UNIVERSE + credentials_universe = getattr(credentials, "universe_domain", default_universe) + + if client_universe != credentials_universe: + raise ValueError( + "The configured universe domain " + f"({client_universe}) does not match the universe domain " + f"found in the credentials ({credentials_universe}). " + "If you haven't configured the universe domain explicitly, " + f"`{default_universe}` is the default." + ) + return True + + def _validate_universe_domain(self): + """Validates client's and credentials' universe domains are consistent. + + Returns: + bool: True iff the configured universe domain is valid. + + Raises: + ValueError: If the configured universe domain is not valid. + """ + self._is_universe_domain_valid = ( + self._is_universe_domain_valid + or RankServiceClient._compare_universes( + self.universe_domain, self.transport._credentials + ) + ) + return self._is_universe_domain_valid + + @property + def api_endpoint(self): + """Return the API endpoint used by the client instance. + + Returns: + str: The API endpoint used by the client instance. + """ + return self._api_endpoint + + @property + def universe_domain(self) -> str: + """Return the universe domain used by the client instance. + + Returns: + str: The universe domain used by the client instance. + """ + return self._universe_domain + + def __init__( + self, + *, + credentials: Optional[ga_credentials.Credentials] = None, + transport: Optional[Union[str, RankServiceTransport]] = None, + client_options: Optional[Union[client_options_lib.ClientOptions, dict]] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + ) -> None: + """Instantiates the rank service client. + + Args: + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + transport (Union[str, RankServiceTransport]): The + transport to use. If set to None, a transport is chosen + automatically. + client_options (Optional[Union[google.api_core.client_options.ClientOptions, dict]]): + Custom options for the client. + + 1. The ``api_endpoint`` property can be used to override the + default endpoint provided by the client when ``transport`` is + not explicitly provided. Only if this property is not set and + ``transport`` was not explicitly provided, the endpoint is + determined by the GOOGLE_API_USE_MTLS_ENDPOINT environment + variable, which have one of the following values: + "always" (always use the default mTLS endpoint), "never" (always + use the default regular endpoint) and "auto" (auto-switch to the + default mTLS endpoint if client certificate is present; this is + the default value). + + 2. If the GOOGLE_API_USE_CLIENT_CERTIFICATE environment variable + is "true", then the ``client_cert_source`` property can be used + to provide a client certificate for mTLS transport. If + not provided, the default SSL client certificate will be used if + present. If GOOGLE_API_USE_CLIENT_CERTIFICATE is "false" or not + set, no client certificate will be used. + + 3. The ``universe_domain`` property can be used to override the + default "googleapis.com" universe. Note that the ``api_endpoint`` + property still takes precedence; and ``universe_domain`` is + currently not supported for mTLS. + + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + """ + self._client_options = client_options + if isinstance(self._client_options, dict): + self._client_options = client_options_lib.from_dict(self._client_options) + if self._client_options is None: + self._client_options = client_options_lib.ClientOptions() + self._client_options = cast( + client_options_lib.ClientOptions, self._client_options + ) + + universe_domain_opt = getattr(self._client_options, "universe_domain", None) + + ( + self._use_client_cert, + self._use_mtls_endpoint, + self._universe_domain_env, + ) = RankServiceClient._read_environment_variables() + self._client_cert_source = RankServiceClient._get_client_cert_source( + self._client_options.client_cert_source, self._use_client_cert + ) + self._universe_domain = RankServiceClient._get_universe_domain( + universe_domain_opt, self._universe_domain_env + ) + self._api_endpoint = None # updated below, depending on `transport` + + # Initialize the universe domain validation. + self._is_universe_domain_valid = False + + api_key_value = getattr(self._client_options, "api_key", None) + if api_key_value and credentials: + raise ValueError( + "client_options.api_key and credentials are mutually exclusive" + ) + + # Save or instantiate the transport. + # Ordinarily, we provide the transport, but allowing a custom transport + # instance provides an extensibility point for unusual situations. + transport_provided = isinstance(transport, RankServiceTransport) + if transport_provided: + # transport is a RankServiceTransport instance. + if credentials or self._client_options.credentials_file or api_key_value: + raise ValueError( + "When providing a transport instance, " + "provide its credentials directly." + ) + if self._client_options.scopes: + raise ValueError( + "When providing a transport instance, provide its scopes " + "directly." + ) + self._transport = cast(RankServiceTransport, transport) + self._api_endpoint = self._transport.host + + self._api_endpoint = self._api_endpoint or RankServiceClient._get_api_endpoint( + self._client_options.api_endpoint, + self._client_cert_source, + self._universe_domain, + self._use_mtls_endpoint, + ) + + if not transport_provided: + import google.auth._default # type: ignore + + if api_key_value and hasattr( + google.auth._default, "get_api_key_credentials" + ): + credentials = google.auth._default.get_api_key_credentials( + api_key_value + ) + + Transport = type(self).get_transport_class(cast(str, transport)) + self._transport = Transport( + credentials=credentials, + credentials_file=self._client_options.credentials_file, + host=self._api_endpoint, + scopes=self._client_options.scopes, + client_cert_source_for_mtls=self._client_cert_source, + quota_project_id=self._client_options.quota_project_id, + client_info=client_info, + always_use_jwt_access=True, + api_audience=self._client_options.api_audience, + ) + + def rank( + self, + request: Optional[Union[rank_service.RankRequest, dict]] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> rank_service.RankResponse: + r"""Ranks a list of text records based on the given input + query. + + .. code-block:: python + + # This snippet has been automatically generated and should be regarded as a + # code template only. + # It will require modifications to work: + # - It may require correct/in-range values for request initialization. + # - It may require specifying regional endpoints when creating the service + # client as shown in: + # https://googleapis.dev/python/google-api-core/latest/client_options.html + from google.cloud import discoveryengine_v1alpha + + def sample_rank(): + # Create a client + client = discoveryengine_v1alpha.RankServiceClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.RankRequest( + ranking_config="ranking_config_value", + ) + + # Make the request + response = client.rank(request=request) + + # Handle the response + print(response) + + Args: + request (Union[google.cloud.discoveryengine_v1alpha.types.RankRequest, dict]): + The request object. Request message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + google.cloud.discoveryengine_v1alpha.types.RankResponse: + Response message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + """ + # Create or coerce a protobuf request object. + # Minor optimization to avoid making a copy if the user passes + # in a rank_service.RankRequest. + # There's no risk of modifying the input as we've already verified + # there are no flattened fields. + if not isinstance(request, rank_service.RankRequest): + request = rank_service.RankRequest(request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = self._transport._wrapped_methods[self._transport.rank] + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata( + (("ranking_config", request.ranking_config),) + ), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + def __enter__(self) -> "RankServiceClient": + return self + + def __exit__(self, type, value, traceback): + """Releases underlying transport's resources. + + .. warning:: + ONLY use as a context manager if the transport is NOT shared + with other clients! Exiting the with block will CLOSE the transport + and may cause errors in other clients! + """ + self.transport.close() + + def list_operations( + self, + request: Optional[operations_pb2.ListOperationsRequest] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.ListOperationsResponse: + r"""Lists operations that match the specified filter in the request. + + Args: + request (:class:`~.operations_pb2.ListOperationsRequest`): + The request object. Request message for + `ListOperations` method. + retry (google.api_core.retry.Retry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + Returns: + ~.operations_pb2.ListOperationsResponse: + Response message for ``ListOperations`` method. + """ + # Create or coerce a protobuf request object. + # The request isn't a proto-plus wrapped type, + # so it must be constructed via keyword expansion. + if isinstance(request, dict): + request = operations_pb2.ListOperationsRequest(**request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method.wrap_method( + self._transport.list_operations, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + def get_operation( + self, + request: Optional[operations_pb2.GetOperationRequest] = None, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Union[float, object] = gapic_v1.method.DEFAULT, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.Operation: + r"""Gets the latest state of a long-running operation. + + Args: + request (:class:`~.operations_pb2.GetOperationRequest`): + The request object. Request message for + `GetOperation` method. + retry (google.api_core.retry.Retry): Designation of what errors, + if any, should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + Returns: + ~.operations_pb2.Operation: + An ``Operation`` object. + """ + # Create or coerce a protobuf request object. + # The request isn't a proto-plus wrapped type, + # so it must be constructed via keyword expansion. + if isinstance(request, dict): + request = operations_pb2.GetOperationRequest(**request) + + # Wrap the RPC method; this adds retry and timeout information, + # and friendly error handling. + rpc = gapic_v1.method.wrap_method( + self._transport.get_operation, + default_timeout=None, + client_info=DEFAULT_CLIENT_INFO, + ) + + # Certain fields should be provided within the metadata header; + # add these here. + metadata = tuple(metadata) + ( + gapic_v1.routing_header.to_grpc_metadata((("name", request.name),)), + ) + + # Validate the universe domain. + self._validate_universe_domain() + + # Send the request. + response = rpc( + request, + retry=retry, + timeout=timeout, + metadata=metadata, + ) + + # Done; return the response. + return response + + +DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=package_version.__version__ +) + + +__all__ = ("RankServiceClient",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/__init__.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/__init__.py new file mode 100644 index 000000000000..0084baa3bfcc --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/__init__.py @@ -0,0 +1,36 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from collections import OrderedDict +from typing import Dict, Type + +from .base import RankServiceTransport +from .grpc import RankServiceGrpcTransport +from .grpc_asyncio import RankServiceGrpcAsyncIOTransport +from .rest import RankServiceRestInterceptor, RankServiceRestTransport + +# Compile a registry of transports. +_transport_registry = OrderedDict() # type: Dict[str, Type[RankServiceTransport]] +_transport_registry["grpc"] = RankServiceGrpcTransport +_transport_registry["grpc_asyncio"] = RankServiceGrpcAsyncIOTransport +_transport_registry["rest"] = RankServiceRestTransport + +__all__ = ( + "RankServiceTransport", + "RankServiceGrpcTransport", + "RankServiceGrpcAsyncIOTransport", + "RankServiceRestTransport", + "RankServiceRestInterceptor", +) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/base.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/base.py new file mode 100644 index 000000000000..66a0514579f2 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/base.py @@ -0,0 +1,182 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import abc +from typing import Awaitable, Callable, Dict, Optional, Sequence, Union + +import google.api_core +from google.api_core import exceptions as core_exceptions +from google.api_core import gapic_v1 +from google.api_core import retry as retries +import google.auth # type: ignore +from google.auth import credentials as ga_credentials # type: ignore +from google.cloud.location import locations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account # type: ignore + +from google.cloud.discoveryengine_v1alpha import gapic_version as package_version +from google.cloud.discoveryengine_v1alpha.types import rank_service + +DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=package_version.__version__ +) + + +class RankServiceTransport(abc.ABC): + """Abstract transport class for RankService.""" + + AUTH_SCOPES = ("https://www.googleapis.com/auth/cloud-platform",) + + DEFAULT_HOST: str = "discoveryengine.googleapis.com" + + def __init__( + self, + *, + host: str = DEFAULT_HOST, + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + always_use_jwt_access: Optional[bool] = False, + api_audience: Optional[str] = None, + **kwargs, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): + The hostname to connect to (default: 'discoveryengine.googleapis.com'). + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scopes (Optional[Sequence[str]]): A list of scopes. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + always_use_jwt_access (Optional[bool]): Whether self signed JWT should + be used for service account credentials. + """ + + scopes_kwargs = {"scopes": scopes, "default_scopes": self.AUTH_SCOPES} + + # Save the scopes. + self._scopes = scopes + + # If no credentials are provided, then determine the appropriate + # defaults. + if credentials and credentials_file: + raise core_exceptions.DuplicateCredentialArgs( + "'credentials_file' and 'credentials' are mutually exclusive" + ) + + if credentials_file is not None: + credentials, _ = google.auth.load_credentials_from_file( + credentials_file, **scopes_kwargs, quota_project_id=quota_project_id + ) + elif credentials is None: + credentials, _ = google.auth.default( + **scopes_kwargs, quota_project_id=quota_project_id + ) + # Don't apply audience if the credentials file passed from user. + if hasattr(credentials, "with_gdch_audience"): + credentials = credentials.with_gdch_audience( + api_audience if api_audience else host + ) + + # If the credentials are service account credentials, then always try to use self signed JWT. + if ( + always_use_jwt_access + and isinstance(credentials, service_account.Credentials) + and hasattr(service_account.Credentials, "with_always_use_jwt_access") + ): + credentials = credentials.with_always_use_jwt_access(True) + + # Save the credentials. + self._credentials = credentials + + # Save the hostname. Default to port 443 (HTTPS) if none is specified. + if ":" not in host: + host += ":443" + self._host = host + + @property + def host(self): + return self._host + + def _prep_wrapped_messages(self, client_info): + # Precompute the wrapped methods. + self._wrapped_methods = { + self.rank: gapic_v1.method.wrap_method( + self.rank, + default_timeout=None, + client_info=client_info, + ), + } + + def close(self): + """Closes resources associated with the transport. + + .. warning:: + Only call this method if the transport is NOT shared + with other clients - this may cause errors in other clients! + """ + raise NotImplementedError() + + @property + def rank( + self, + ) -> Callable[ + [rank_service.RankRequest], + Union[rank_service.RankResponse, Awaitable[rank_service.RankResponse]], + ]: + raise NotImplementedError() + + @property + def list_operations( + self, + ) -> Callable[ + [operations_pb2.ListOperationsRequest], + Union[ + operations_pb2.ListOperationsResponse, + Awaitable[operations_pb2.ListOperationsResponse], + ], + ]: + raise NotImplementedError() + + @property + def get_operation( + self, + ) -> Callable[ + [operations_pb2.GetOperationRequest], + Union[operations_pb2.Operation, Awaitable[operations_pb2.Operation]], + ]: + raise NotImplementedError() + + @property + def kind(self) -> str: + raise NotImplementedError() + + +__all__ = ("RankServiceTransport",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc.py new file mode 100644 index 000000000000..d5b8acd58595 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc.py @@ -0,0 +1,302 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from typing import Callable, Dict, Optional, Sequence, Tuple, Union +import warnings + +from google.api_core import gapic_v1, grpc_helpers +import google.auth # type: ignore +from google.auth import credentials as ga_credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.cloud.location import locations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore +import grpc # type: ignore + +from google.cloud.discoveryengine_v1alpha.types import rank_service + +from .base import DEFAULT_CLIENT_INFO, RankServiceTransport + + +class RankServiceGrpcTransport(RankServiceTransport): + """gRPC backend transport for RankService. + + Service for ranking text records. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _stubs: Dict[str, Callable] + + def __init__( + self, + *, + host: str = "discoveryengine.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + channel: Optional[grpc.Channel] = None, + api_mtls_endpoint: Optional[str] = None, + client_cert_source: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + ssl_channel_credentials: Optional[grpc.ChannelCredentials] = None, + client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + always_use_jwt_access: Optional[bool] = False, + api_audience: Optional[str] = None, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): + The hostname to connect to (default: 'discoveryengine.googleapis.com'). + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional(Sequence[str])): A list of scopes. This argument is + ignored if ``channel`` is provided. + channel (Optional[grpc.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or application default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for the grpc channel. It is ignored if ``channel`` is provided. + client_cert_source_for_mtls (Optional[Callable[[], Tuple[bytes, bytes]]]): + A callback to provide client certificate bytes and private key bytes, + both in PEM format. It is used to configure a mutual TLS channel. It is + ignored if ``channel`` or ``ssl_channel_credentials`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + always_use_jwt_access (Optional[bool]): Whether self signed JWT should + be used for service account credentials. + + Raises: + google.auth.exceptions.MutualTLSChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._grpc_channel = None + self._ssl_channel_credentials = ssl_channel_credentials + self._stubs: Dict[str, Callable] = {} + + if api_mtls_endpoint: + warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning) + if client_cert_source: + warnings.warn("client_cert_source is deprecated", DeprecationWarning) + + if channel: + # Ignore credentials if a channel was passed. + credentials = False + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + + else: + if api_mtls_endpoint: + host = api_mtls_endpoint + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + self._ssl_channel_credentials = SslCredentials().ssl_credentials + + else: + if client_cert_source_for_mtls and not ssl_channel_credentials: + cert, key = client_cert_source_for_mtls() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + + # The base transport sets the host, credentials and scopes + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + client_info=client_info, + always_use_jwt_access=always_use_jwt_access, + api_audience=api_audience, + ) + + if not self._grpc_channel: + self._grpc_channel = type(self).create_channel( + self._host, + # use the credentials which are saved + credentials=self._credentials, + # Set ``credentials_file`` to ``None`` here as + # the credentials that we saved earlier should be used. + credentials_file=None, + scopes=self._scopes, + ssl_credentials=self._ssl_channel_credentials, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Wrap messages. This must be done after self._grpc_channel exists + self._prep_wrapped_messages(client_info) + + @classmethod + def create_channel( + cls, + host: str = "discoveryengine.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> grpc.Channel: + """Create and return a gRPC channel object. + Args: + host (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is mutually exclusive with credentials. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + grpc.Channel: A gRPC channel object. + + Raises: + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + + return grpc_helpers.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + quota_project_id=quota_project_id, + default_scopes=cls.AUTH_SCOPES, + scopes=scopes, + default_host=cls.DEFAULT_HOST, + **kwargs, + ) + + @property + def grpc_channel(self) -> grpc.Channel: + """Return the channel designed to connect to this service.""" + return self._grpc_channel + + @property + def rank(self) -> Callable[[rank_service.RankRequest], rank_service.RankResponse]: + r"""Return a callable for the rank method over gRPC. + + Ranks a list of text records based on the given input + query. + + Returns: + Callable[[~.RankRequest], + ~.RankResponse]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "rank" not in self._stubs: + self._stubs["rank"] = self.grpc_channel.unary_unary( + "/google.cloud.discoveryengine.v1alpha.RankService/Rank", + request_serializer=rank_service.RankRequest.serialize, + response_deserializer=rank_service.RankResponse.deserialize, + ) + return self._stubs["rank"] + + def close(self): + self.grpc_channel.close() + + @property + def get_operation( + self, + ) -> Callable[[operations_pb2.GetOperationRequest], operations_pb2.Operation]: + r"""Return a callable for the get_operation method over gRPC.""" + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_operation" not in self._stubs: + self._stubs["get_operation"] = self.grpc_channel.unary_unary( + "/google.longrunning.Operations/GetOperation", + request_serializer=operations_pb2.GetOperationRequest.SerializeToString, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["get_operation"] + + @property + def list_operations( + self, + ) -> Callable[ + [operations_pb2.ListOperationsRequest], operations_pb2.ListOperationsResponse + ]: + r"""Return a callable for the list_operations method over gRPC.""" + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_operations" not in self._stubs: + self._stubs["list_operations"] = self.grpc_channel.unary_unary( + "/google.longrunning.Operations/ListOperations", + request_serializer=operations_pb2.ListOperationsRequest.SerializeToString, + response_deserializer=operations_pb2.ListOperationsResponse.FromString, + ) + return self._stubs["list_operations"] + + @property + def kind(self) -> str: + return "grpc" + + +__all__ = ("RankServiceGrpcTransport",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc_asyncio.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc_asyncio.py new file mode 100644 index 000000000000..40aed2d9dc58 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/grpc_asyncio.py @@ -0,0 +1,303 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from typing import Awaitable, Callable, Dict, Optional, Sequence, Tuple, Union +import warnings + +from google.api_core import gapic_v1, grpc_helpers_async +from google.auth import credentials as ga_credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.cloud.location import locations_pb2 # type: ignore +from google.longrunning import operations_pb2 # type: ignore +import grpc # type: ignore +from grpc.experimental import aio # type: ignore + +from google.cloud.discoveryengine_v1alpha.types import rank_service + +from .base import DEFAULT_CLIENT_INFO, RankServiceTransport +from .grpc import RankServiceGrpcTransport + + +class RankServiceGrpcAsyncIOTransport(RankServiceTransport): + """gRPC AsyncIO backend transport for RankService. + + Service for ranking text records. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends protocol buffers over the wire using gRPC (which is built on + top of HTTP/2); the ``grpcio`` package must be installed. + """ + + _grpc_channel: aio.Channel + _stubs: Dict[str, Callable] = {} + + @classmethod + def create_channel( + cls, + host: str = "discoveryengine.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + quota_project_id: Optional[str] = None, + **kwargs, + ) -> aio.Channel: + """Create and return a gRPC AsyncIO channel object. + Args: + host (Optional[str]): The host for the channel to use. + credentials (Optional[~.Credentials]): The + authorization credentials to attach to requests. These + credentials identify this application to the service. If + none are specified, the client will attempt to ascertain + the credentials from the environment. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + kwargs (Optional[dict]): Keyword arguments, which are passed to the + channel creation. + Returns: + aio.Channel: A gRPC AsyncIO channel object. + """ + + return grpc_helpers_async.create_channel( + host, + credentials=credentials, + credentials_file=credentials_file, + quota_project_id=quota_project_id, + default_scopes=cls.AUTH_SCOPES, + scopes=scopes, + default_host=cls.DEFAULT_HOST, + **kwargs, + ) + + def __init__( + self, + *, + host: str = "discoveryengine.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + channel: Optional[aio.Channel] = None, + api_mtls_endpoint: Optional[str] = None, + client_cert_source: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + ssl_channel_credentials: Optional[grpc.ChannelCredentials] = None, + client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + always_use_jwt_access: Optional[bool] = False, + api_audience: Optional[str] = None, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): + The hostname to connect to (default: 'discoveryengine.googleapis.com'). + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + This argument is ignored if ``channel`` is provided. + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional[Sequence[str]]): A optional list of scopes needed for this + service. These are only used when credentials are not specified and + are passed to :func:`google.auth.default`. + channel (Optional[aio.Channel]): A ``Channel`` instance through + which to make calls. + api_mtls_endpoint (Optional[str]): Deprecated. The mutual TLS endpoint. + If provided, it overrides the ``host`` argument and tries to create + a mutual TLS channel with client SSL credentials from + ``client_cert_source`` or application default SSL credentials. + client_cert_source (Optional[Callable[[], Tuple[bytes, bytes]]]): + Deprecated. A callback to provide client SSL certificate bytes and + private key bytes, both in PEM format. It is ignored if + ``api_mtls_endpoint`` is None. + ssl_channel_credentials (grpc.ChannelCredentials): SSL credentials + for the grpc channel. It is ignored if ``channel`` is provided. + client_cert_source_for_mtls (Optional[Callable[[], Tuple[bytes, bytes]]]): + A callback to provide client certificate bytes and private key bytes, + both in PEM format. It is used to configure a mutual TLS channel. It is + ignored if ``channel`` or ``ssl_channel_credentials`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you're developing + your own client library. + always_use_jwt_access (Optional[bool]): Whether self signed JWT should + be used for service account credentials. + + Raises: + google.auth.exceptions.MutualTlsChannelError: If mutual TLS transport + creation failed for any reason. + google.api_core.exceptions.DuplicateCredentialArgs: If both ``credentials`` + and ``credentials_file`` are passed. + """ + self._grpc_channel = None + self._ssl_channel_credentials = ssl_channel_credentials + self._stubs: Dict[str, Callable] = {} + + if api_mtls_endpoint: + warnings.warn("api_mtls_endpoint is deprecated", DeprecationWarning) + if client_cert_source: + warnings.warn("client_cert_source is deprecated", DeprecationWarning) + + if channel: + # Ignore credentials if a channel was passed. + credentials = False + # If a channel was explicitly provided, set it. + self._grpc_channel = channel + self._ssl_channel_credentials = None + else: + if api_mtls_endpoint: + host = api_mtls_endpoint + + # Create SSL credentials with client_cert_source or application + # default SSL credentials. + if client_cert_source: + cert, key = client_cert_source() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + else: + self._ssl_channel_credentials = SslCredentials().ssl_credentials + + else: + if client_cert_source_for_mtls and not ssl_channel_credentials: + cert, key = client_cert_source_for_mtls() + self._ssl_channel_credentials = grpc.ssl_channel_credentials( + certificate_chain=cert, private_key=key + ) + + # The base transport sets the host, credentials and scopes + super().__init__( + host=host, + credentials=credentials, + credentials_file=credentials_file, + scopes=scopes, + quota_project_id=quota_project_id, + client_info=client_info, + always_use_jwt_access=always_use_jwt_access, + api_audience=api_audience, + ) + + if not self._grpc_channel: + self._grpc_channel = type(self).create_channel( + self._host, + # use the credentials which are saved + credentials=self._credentials, + # Set ``credentials_file`` to ``None`` here as + # the credentials that we saved earlier should be used. + credentials_file=None, + scopes=self._scopes, + ssl_credentials=self._ssl_channel_credentials, + quota_project_id=quota_project_id, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Wrap messages. This must be done after self._grpc_channel exists + self._prep_wrapped_messages(client_info) + + @property + def grpc_channel(self) -> aio.Channel: + """Create the channel designed to connect to this service. + + This property caches on the instance; repeated calls return + the same channel. + """ + # Return the channel from cache. + return self._grpc_channel + + @property + def rank( + self, + ) -> Callable[[rank_service.RankRequest], Awaitable[rank_service.RankResponse]]: + r"""Return a callable for the rank method over gRPC. + + Ranks a list of text records based on the given input + query. + + Returns: + Callable[[~.RankRequest], + Awaitable[~.RankResponse]]: + A function that, when called, will call the underlying RPC + on the server. + """ + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "rank" not in self._stubs: + self._stubs["rank"] = self.grpc_channel.unary_unary( + "/google.cloud.discoveryengine.v1alpha.RankService/Rank", + request_serializer=rank_service.RankRequest.serialize, + response_deserializer=rank_service.RankResponse.deserialize, + ) + return self._stubs["rank"] + + def close(self): + return self.grpc_channel.close() + + @property + def get_operation( + self, + ) -> Callable[[operations_pb2.GetOperationRequest], operations_pb2.Operation]: + r"""Return a callable for the get_operation method over gRPC.""" + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "get_operation" not in self._stubs: + self._stubs["get_operation"] = self.grpc_channel.unary_unary( + "/google.longrunning.Operations/GetOperation", + request_serializer=operations_pb2.GetOperationRequest.SerializeToString, + response_deserializer=operations_pb2.Operation.FromString, + ) + return self._stubs["get_operation"] + + @property + def list_operations( + self, + ) -> Callable[ + [operations_pb2.ListOperationsRequest], operations_pb2.ListOperationsResponse + ]: + r"""Return a callable for the list_operations method over gRPC.""" + # Generate a "stub function" on-the-fly which will actually make + # the request. + # gRPC handles serialization and deserialization, so we just need + # to pass in the functions for each. + if "list_operations" not in self._stubs: + self._stubs["list_operations"] = self.grpc_channel.unary_unary( + "/google.longrunning.Operations/ListOperations", + request_serializer=operations_pb2.ListOperationsRequest.SerializeToString, + response_deserializer=operations_pb2.ListOperationsResponse.FromString, + ) + return self._stubs["list_operations"] + + +__all__ = ("RankServiceGrpcAsyncIOTransport",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/rest.py new file mode 100644 index 000000000000..fb339d0730e0 --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/rank_service/transports/rest.py @@ -0,0 +1,598 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import dataclasses +import json # type: ignore +import re +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union +import warnings + +from google.api_core import gapic_v1, path_template, rest_helpers, rest_streaming +from google.api_core import exceptions as core_exceptions +from google.api_core import retry as retries +from google.auth import credentials as ga_credentials # type: ignore +from google.auth.transport.grpc import SslCredentials # type: ignore +from google.auth.transport.requests import AuthorizedSession # type: ignore +from google.cloud.location import locations_pb2 # type: ignore +from google.protobuf import json_format +import grpc # type: ignore +from requests import __version__ as requests_version + +try: + OptionalRetry = Union[retries.Retry, gapic_v1.method._MethodDefault, None] +except AttributeError: # pragma: NO COVER + OptionalRetry = Union[retries.Retry, object, None] # type: ignore + + +from google.longrunning import operations_pb2 # type: ignore + +from google.cloud.discoveryengine_v1alpha.types import rank_service + +from .base import DEFAULT_CLIENT_INFO as BASE_DEFAULT_CLIENT_INFO +from .base import RankServiceTransport + +DEFAULT_CLIENT_INFO = gapic_v1.client_info.ClientInfo( + gapic_version=BASE_DEFAULT_CLIENT_INFO.gapic_version, + grpc_version=None, + rest_version=requests_version, +) + + +class RankServiceRestInterceptor: + """Interceptor for RankService. + + Interceptors are used to manipulate requests, request metadata, and responses + in arbitrary ways. + Example use cases include: + * Logging + * Verifying requests according to service or custom semantics + * Stripping extraneous information from responses + + These use cases and more can be enabled by injecting an + instance of a custom subclass when constructing the RankServiceRestTransport. + + .. code-block:: python + class MyCustomRankServiceInterceptor(RankServiceRestInterceptor): + def pre_rank(self, request, metadata): + logging.log(f"Received request: {request}") + return request, metadata + + def post_rank(self, response): + logging.log(f"Received response: {response}") + return response + + transport = RankServiceRestTransport(interceptor=MyCustomRankServiceInterceptor()) + client = RankServiceClient(transport=transport) + + + """ + + def pre_rank( + self, request: rank_service.RankRequest, metadata: Sequence[Tuple[str, str]] + ) -> Tuple[rank_service.RankRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for rank + + Override in a subclass to manipulate the request or metadata + before they are sent to the RankService server. + """ + return request, metadata + + def post_rank( + self, response: rank_service.RankResponse + ) -> rank_service.RankResponse: + """Post-rpc interceptor for rank + + Override in a subclass to manipulate the response + after it is returned by the RankService server but before + it is returned to user code. + """ + return response + + def pre_get_operation( + self, + request: operations_pb2.GetOperationRequest, + metadata: Sequence[Tuple[str, str]], + ) -> Tuple[operations_pb2.GetOperationRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for get_operation + + Override in a subclass to manipulate the request or metadata + before they are sent to the RankService server. + """ + return request, metadata + + def post_get_operation( + self, response: operations_pb2.Operation + ) -> operations_pb2.Operation: + """Post-rpc interceptor for get_operation + + Override in a subclass to manipulate the response + after it is returned by the RankService server but before + it is returned to user code. + """ + return response + + def pre_list_operations( + self, + request: operations_pb2.ListOperationsRequest, + metadata: Sequence[Tuple[str, str]], + ) -> Tuple[operations_pb2.ListOperationsRequest, Sequence[Tuple[str, str]]]: + """Pre-rpc interceptor for list_operations + + Override in a subclass to manipulate the request or metadata + before they are sent to the RankService server. + """ + return request, metadata + + def post_list_operations( + self, response: operations_pb2.ListOperationsResponse + ) -> operations_pb2.ListOperationsResponse: + """Post-rpc interceptor for list_operations + + Override in a subclass to manipulate the response + after it is returned by the RankService server but before + it is returned to user code. + """ + return response + + +@dataclasses.dataclass +class RankServiceRestStub: + _session: AuthorizedSession + _host: str + _interceptor: RankServiceRestInterceptor + + +class RankServiceRestTransport(RankServiceTransport): + """REST backend transport for RankService. + + Service for ranking text records. + + This class defines the same methods as the primary client, so the + primary client can load the underlying transport implementation + and call it. + + It sends JSON representations of protocol buffers over HTTP/1.1 + + """ + + def __init__( + self, + *, + host: str = "discoveryengine.googleapis.com", + credentials: Optional[ga_credentials.Credentials] = None, + credentials_file: Optional[str] = None, + scopes: Optional[Sequence[str]] = None, + client_cert_source_for_mtls: Optional[Callable[[], Tuple[bytes, bytes]]] = None, + quota_project_id: Optional[str] = None, + client_info: gapic_v1.client_info.ClientInfo = DEFAULT_CLIENT_INFO, + always_use_jwt_access: Optional[bool] = False, + url_scheme: str = "https", + interceptor: Optional[RankServiceRestInterceptor] = None, + api_audience: Optional[str] = None, + ) -> None: + """Instantiate the transport. + + Args: + host (Optional[str]): + The hostname to connect to (default: 'discoveryengine.googleapis.com'). + credentials (Optional[google.auth.credentials.Credentials]): The + authorization credentials to attach to requests. These + credentials identify the application to the service; if none + are specified, the client will attempt to ascertain the + credentials from the environment. + + credentials_file (Optional[str]): A file with credentials that can + be loaded with :func:`google.auth.load_credentials_from_file`. + This argument is ignored if ``channel`` is provided. + scopes (Optional(Sequence[str])): A list of scopes. This argument is + ignored if ``channel`` is provided. + client_cert_source_for_mtls (Callable[[], Tuple[bytes, bytes]]): Client + certificate to configure mutual TLS HTTP channel. It is ignored + if ``channel`` is provided. + quota_project_id (Optional[str]): An optional project to use for billing + and quota. + client_info (google.api_core.gapic_v1.client_info.ClientInfo): + The client info used to send a user-agent string along with + API requests. If ``None``, then default info will be used. + Generally, you only need to set this if you are developing + your own client library. + always_use_jwt_access (Optional[bool]): Whether self signed JWT should + be used for service account credentials. + url_scheme: the protocol scheme for the API endpoint. Normally + "https", but for testing or local servers, + "http" can be specified. + """ + # Run the base constructor + # TODO(yon-mg): resolve other ctor params i.e. scopes, quota, etc. + # TODO: When custom host (api_endpoint) is set, `scopes` must *also* be set on the + # credentials object + maybe_url_match = re.match("^(?Phttp(?:s)?://)?(?P.*)$", host) + if maybe_url_match is None: + raise ValueError( + f"Unexpected hostname structure: {host}" + ) # pragma: NO COVER + + url_match_items = maybe_url_match.groupdict() + + host = f"{url_scheme}://{host}" if not url_match_items["scheme"] else host + + super().__init__( + host=host, + credentials=credentials, + client_info=client_info, + always_use_jwt_access=always_use_jwt_access, + api_audience=api_audience, + ) + self._session = AuthorizedSession( + self._credentials, default_host=self.DEFAULT_HOST + ) + if client_cert_source_for_mtls: + self._session.configure_mtls_channel(client_cert_source_for_mtls) + self._interceptor = interceptor or RankServiceRestInterceptor() + self._prep_wrapped_messages(client_info) + + class _Rank(RankServiceRestStub): + def __hash__(self): + return hash("Rank") + + __REQUIRED_FIELDS_DEFAULT_VALUES: Dict[str, Any] = {} + + @classmethod + def _get_unset_required_fields(cls, message_dict): + return { + k: v + for k, v in cls.__REQUIRED_FIELDS_DEFAULT_VALUES.items() + if k not in message_dict + } + + def __call__( + self, + request: rank_service.RankRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> rank_service.RankResponse: + r"""Call the rank method over HTTP. + + Args: + request (~.rank_service.RankRequest): + The request object. Request message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + ~.rank_service.RankResponse: + Response message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "post", + "uri": "/v1alpha/{ranking_config=projects/*/locations/*/rankingConfigs/*}:rank", + "body": "*", + }, + ] + request, metadata = self._interceptor.pre_rank(request, metadata) + pb_request = rank_service.RankRequest.pb(request) + transcoded_request = path_template.transcode(http_options, pb_request) + + # Jsonify the request body + + body = json_format.MessageToJson( + transcoded_request["body"], use_integers_for_enums=True + ) + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads( + json_format.MessageToJson( + transcoded_request["query_params"], + use_integers_for_enums=True, + ) + ) + query_params.update(self._get_unset_required_fields(query_params)) + + query_params["$alt"] = "json;enum-encoding=int" + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params, strict=True), + data=body, + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + # Return the response + resp = rank_service.RankResponse() + pb_resp = rank_service.RankResponse.pb(resp) + + json_format.Parse(response.content, pb_resp, ignore_unknown_fields=True) + resp = self._interceptor.post_rank(resp) + return resp + + @property + def rank(self) -> Callable[[rank_service.RankRequest], rank_service.RankResponse]: + # The return type is fine, but mypy isn't sophisticated enough to determine what's going on here. + # In C++ this would require a dynamic_cast + return self._Rank(self._session, self._host, self._interceptor) # type: ignore + + @property + def get_operation(self): + return self._GetOperation(self._session, self._host, self._interceptor) # type: ignore + + class _GetOperation(RankServiceRestStub): + def __call__( + self, + request: operations_pb2.GetOperationRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.Operation: + r"""Call the get operation method over HTTP. + + Args: + request (operations_pb2.GetOperationRequest): + The request object for GetOperation method. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + operations_pb2.Operation: Response from GetOperation method. + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataConnector/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/branches/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/schemas/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/siteSearchEngine/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/siteSearchEngine/targetSites/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/engines/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/branches/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/models/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/operations/*}", + }, + ] + + request, metadata = self._interceptor.pre_get_operation(request, metadata) + request_kwargs = json_format.MessageToDict(request) + transcoded_request = path_template.transcode(http_options, **request_kwargs) + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads(json.dumps(transcoded_request["query_params"])) + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params), + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + resp = operations_pb2.Operation() + resp = json_format.Parse(response.content.decode("utf-8"), resp) + resp = self._interceptor.post_get_operation(resp) + return resp + + @property + def list_operations(self): + return self._ListOperations(self._session, self._host, self._interceptor) # type: ignore + + class _ListOperations(RankServiceRestStub): + def __call__( + self, + request: operations_pb2.ListOperationsRequest, + *, + retry: OptionalRetry = gapic_v1.method.DEFAULT, + timeout: Optional[float] = None, + metadata: Sequence[Tuple[str, str]] = (), + ) -> operations_pb2.ListOperationsResponse: + r"""Call the list operations method over HTTP. + + Args: + request (operations_pb2.ListOperationsRequest): + The request object for ListOperations method. + retry (google.api_core.retry.Retry): Designation of what errors, if any, + should be retried. + timeout (float): The timeout for this request. + metadata (Sequence[Tuple[str, str]]): Strings which should be + sent along with the request as metadata. + + Returns: + operations_pb2.ListOperationsResponse: Response from ListOperations method. + """ + + http_options: List[Dict[str, str]] = [ + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataConnector}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/branches/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/schemas/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/siteSearchEngine/targetSites}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*/siteSearchEngine}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/dataStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*/engines/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/collections/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/branches/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/models/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*}/operations", + }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*}/operations", + }, + ] + + request, metadata = self._interceptor.pre_list_operations(request, metadata) + request_kwargs = json_format.MessageToDict(request) + transcoded_request = path_template.transcode(http_options, **request_kwargs) + + uri = transcoded_request["uri"] + method = transcoded_request["method"] + + # Jsonify the query params + query_params = json.loads(json.dumps(transcoded_request["query_params"])) + + # Send the request + headers = dict(metadata) + headers["Content-Type"] = "application/json" + + response = getattr(self._session, method)( + "{host}{uri}".format(host=self._host, uri=uri), + timeout=timeout, + headers=headers, + params=rest_helpers.flatten_query_params(query_params), + ) + + # In case of error, raise the appropriate core_exceptions.GoogleAPICallError exception + # subclass. + if response.status_code >= 400: + raise core_exceptions.from_http_response(response) + + resp = operations_pb2.ListOperationsResponse() + resp = json_format.Parse(response.content.decode("utf-8"), resp) + resp = self._interceptor.post_list_operations(resp) + return resp + + @property + def kind(self) -> str: + return "rest" + + def close(self): + self._session.close() + + +__all__ = ("RankServiceRestTransport",) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/recommendation_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/recommendation_service/transports/rest.py index d173c1a79509..b3eae0c09863 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/recommendation_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/recommendation_service/transports/rest.py @@ -438,6 +438,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/schema_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/schema_service/transports/rest.py index 8ff9d41ed798..f64580bf219e 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/schema_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/schema_service/transports/rest.py @@ -435,6 +435,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -1118,6 +1122,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_service/transports/rest.py index 0af3a99d2950..9c5be5135af1 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_service/transports/rest.py @@ -436,6 +436,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_tuning_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_tuning_service/transports/rest.py index 7e71bd4789a4..71e5ae9679d5 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_tuning_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/search_tuning_service/transports/rest.py @@ -314,6 +314,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -583,6 +587,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/async_client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/async_client.py index d646f0b5417a..7891ec8832cf 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/async_client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/async_client.py @@ -61,7 +61,9 @@ class ServingConfigServiceAsyncClient: - """Service for operations related to [ServingConfig][].""" + """Service for operations related to + [ServingConfig][google.cloud.discoveryengine.v1alpha.ServingConfig]. + """ _client: ServingConfigServiceClient @@ -300,7 +302,7 @@ async def sample_update_serving_config(): serving_config = discoveryengine_v1alpha.ServingConfig() serving_config.media_config.content_watched_percentage_threshold = 0.3811 serving_config.display_name = "display_name_value" - serving_config.solution_type = "SOLUTION_TYPE_CHAT" + serving_config.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateServingConfigRequest( serving_config=serving_config, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/client.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/client.py index 537e7ac8f2f8..4ecafa4a418f 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/client.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/client.py @@ -104,7 +104,9 @@ def get_transport_class( class ServingConfigServiceClient(metaclass=ServingConfigServiceClientMeta): - """Service for operations related to [ServingConfig][].""" + """Service for operations related to + [ServingConfig][google.cloud.discoveryengine.v1alpha.ServingConfig]. + """ @staticmethod def _get_default_mtls_endpoint(api_endpoint): @@ -712,7 +714,7 @@ def sample_update_serving_config(): serving_config = discoveryengine_v1alpha.ServingConfig() serving_config.media_config.content_watched_percentage_threshold = 0.3811 serving_config.display_name = "display_name_value" - serving_config.solution_type = "SOLUTION_TYPE_CHAT" + serving_config.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateServingConfigRequest( serving_config=serving_config, diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc.py index 92acf5ade2f1..7510879650fa 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc.py @@ -36,7 +36,8 @@ class ServingConfigServiceGrpcTransport(ServingConfigServiceTransport): """gRPC backend transport for ServingConfigService. - Service for operations related to [ServingConfig][]. + Service for operations related to + [ServingConfig][google.cloud.discoveryengine.v1alpha.ServingConfig]. This class defines the same methods as the primary client, so the primary client can load the underlying transport implementation diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc_asyncio.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc_asyncio.py index 1b6f1a481c23..ac0370f4d030 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc_asyncio.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/grpc_asyncio.py @@ -37,7 +37,8 @@ class ServingConfigServiceGrpcAsyncIOTransport(ServingConfigServiceTransport): """gRPC AsyncIO backend transport for ServingConfigService. - Service for operations related to [ServingConfig][]. + Service for operations related to + [ServingConfig][google.cloud.discoveryengine.v1alpha.ServingConfig]. This class defines the same methods as the primary client, so the primary client can load the underlying transport implementation diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/rest.py index 2ee2059e4829..c1818764bf5b 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/serving_config_service/transports/rest.py @@ -232,7 +232,8 @@ class ServingConfigServiceRestStub: class ServingConfigServiceRestTransport(ServingConfigServiceTransport): """REST backend transport for ServingConfigService. - Service for operations related to [ServingConfig][]. + Service for operations related to + [ServingConfig][google.cloud.discoveryengine.v1alpha.ServingConfig]. This class defines the same methods as the primary client, so the primary client can load the underlying transport implementation @@ -737,6 +738,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/site_search_engine_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/site_search_engine_service/transports/rest.py index eecb33082934..1a3f8294ff31 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/site_search_engine_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/site_search_engine_service/transports/rest.py @@ -685,6 +685,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -2174,6 +2178,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/user_event_service/transports/rest.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/user_event_service/transports/rest.py index 1462db92b57e..2be9f70f7f63 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/user_event_service/transports/rest.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/services/user_event_service/transports/rest.py @@ -412,6 +412,10 @@ def operations_client(self) -> operations_v1.AbstractOperationsClient: "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", @@ -1055,6 +1059,10 @@ def __call__( "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/dataStores/*/operations/*}", }, + { + "method": "get", + "uri": "/v1alpha/{name=projects/*/locations/*/evaluations/*/operations/*}", + }, { "method": "get", "uri": "/v1alpha/{name=projects/*/locations/*/operations/*}", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/__init__.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/__init__.py index 413667215f6e..c09ad2d93af2 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/__init__.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/__init__.py @@ -64,12 +64,13 @@ UpdateDataStoreRequest, UpdateDocumentProcessingConfigRequest, ) -from .document import Document +from .document import Document, ProcessedDocument from .document_processing_config import DocumentProcessingConfig from .document_service import ( CreateDocumentRequest, DeleteDocumentRequest, GetDocumentRequest, + GetProcessedDocumentRequest, ListDocumentsRequest, ListDocumentsResponse, UpdateDocumentRequest, @@ -97,7 +98,14 @@ ) from .import_config import ( BigQuerySource, + BigtableOptions, + BigtableSource, + CloudSqlSource, + FhirStoreSource, + FirestoreSource, GcsSource, + ImportCompletionSuggestionsMetadata, + ImportCompletionSuggestionsResponse, ImportDocumentsMetadata, ImportDocumentsRequest, ImportDocumentsResponse, @@ -108,6 +116,7 @@ ImportUserEventsMetadata, ImportUserEventsRequest, ImportUserEventsResponse, + SpannerSource, ) from .purge_config import ( PurgeDocumentsMetadata, @@ -121,6 +130,7 @@ PurgeUserEventsRequest, PurgeUserEventsResponse, ) +from .rank_service import RankingRecord, RankRequest, RankResponse from .recommendation_service import RecommendRequest, RecommendResponse from .schema import FieldConfig, Schema from .schema_service import ( @@ -238,10 +248,12 @@ "UpdateDataStoreRequest", "UpdateDocumentProcessingConfigRequest", "Document", + "ProcessedDocument", "DocumentProcessingConfig", "CreateDocumentRequest", "DeleteDocumentRequest", "GetDocumentRequest", + "GetProcessedDocumentRequest", "ListDocumentsRequest", "ListDocumentsResponse", "UpdateDocumentRequest", @@ -263,7 +275,14 @@ "EstimateDataSizeRequest", "EstimateDataSizeResponse", "BigQuerySource", + "BigtableOptions", + "BigtableSource", + "CloudSqlSource", + "FhirStoreSource", + "FirestoreSource", "GcsSource", + "ImportCompletionSuggestionsMetadata", + "ImportCompletionSuggestionsResponse", "ImportDocumentsMetadata", "ImportDocumentsRequest", "ImportDocumentsResponse", @@ -274,6 +293,7 @@ "ImportUserEventsMetadata", "ImportUserEventsRequest", "ImportUserEventsResponse", + "SpannerSource", "PurgeDocumentsMetadata", "PurgeDocumentsRequest", "PurgeDocumentsResponse", @@ -284,6 +304,9 @@ "PurgeUserEventsMetadata", "PurgeUserEventsRequest", "PurgeUserEventsResponse", + "RankingRecord", + "RankRequest", + "RankResponse", "RecommendRequest", "RecommendResponse", "FieldConfig", diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/chunk.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/chunk.py index c7dc880260f0..949a111bbb1a 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/chunk.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/chunk.py @@ -40,7 +40,7 @@ class Chunk(proto.Message): This field must be a UTF-8 encoded string with a length limit of 1024 characters. id (str): - Unique chunk id of the current chunk. + Unique chunk ID of the current chunk. content (str): Content is a string from a document (parsed content). @@ -50,6 +50,10 @@ class Chunk(proto.Message): derived_struct_data (google.protobuf.struct_pb2.Struct): Output only. This field is OUTPUT_ONLY. It contains derived data that are not in the original input document. + page_span (google.cloud.discoveryengine_v1alpha.types.Chunk.PageSpan): + Page span of the chunk. + chunk_metadata (google.cloud.discoveryengine_v1alpha.types.Chunk.ChunkMetadata): + Output only. Metadata of the current chunk. """ class DocumentMetadata(proto.Message): @@ -72,6 +76,58 @@ class DocumentMetadata(proto.Message): number=2, ) + class PageSpan(proto.Message): + r"""Page span of the chunk. + + Attributes: + page_start (int): + The start page of the chunk. + page_end (int): + The end page of the chunk. + """ + + page_start: int = proto.Field( + proto.INT32, + number=1, + ) + page_end: int = proto.Field( + proto.INT32, + number=2, + ) + + class ChunkMetadata(proto.Message): + r"""Metadata of the current chunk. This field is only populated on + [SearchService.Search][google.cloud.discoveryengine.v1alpha.SearchService.Search] + API. + + Attributes: + previous_chunks (MutableSequence[google.cloud.discoveryengine_v1alpha.types.Chunk]): + The previous chunks of the current chunk. The number is + controlled by + [SearchRequest.ContentSearchSpec.ChunkSpec.num_previous_chunks][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.ChunkSpec.num_previous_chunks]. + This field is only populated on + [SearchService.Search][google.cloud.discoveryengine.v1alpha.SearchService.Search] + API. + next_chunks (MutableSequence[google.cloud.discoveryengine_v1alpha.types.Chunk]): + The next chunks of the current chunk. The number is + controlled by + [SearchRequest.ContentSearchSpec.ChunkSpec.num_next_chunks][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.ChunkSpec.num_next_chunks]. + This field is only populated on + [SearchService.Search][google.cloud.discoveryengine.v1alpha.SearchService.Search] + API. + """ + + previous_chunks: MutableSequence["Chunk"] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message="Chunk", + ) + next_chunks: MutableSequence["Chunk"] = proto.RepeatedField( + proto.MESSAGE, + number=2, + message="Chunk", + ) + name: str = proto.Field( proto.STRING, number=1, @@ -94,6 +150,16 @@ class DocumentMetadata(proto.Message): number=4, message=struct_pb2.Struct, ) + page_span: PageSpan = proto.Field( + proto.MESSAGE, + number=6, + message=PageSpan, + ) + chunk_metadata: ChunkMetadata = proto.Field( + proto.MESSAGE, + number=7, + message=ChunkMetadata, + ) __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/common.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/common.py index d749dbe770bd..b694a83f6f44 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/common.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/common.py @@ -51,10 +51,13 @@ class IndustryVertical(proto.Enum): not specific to any industry vertical. MEDIA (2): The media industry vertical. + HEALTHCARE_FHIR (7): + The healthcare FHIR vertical. """ INDUSTRY_VERTICAL_UNSPECIFIED = 0 GENERIC = 1 MEDIA = 2 + HEALTHCARE_FHIR = 7 class SolutionType(proto.Enum): @@ -70,11 +73,17 @@ class SolutionType(proto.Enum): SOLUTION_TYPE_CHAT (3): Used for use cases related to the Generative AI agent. + SOLUTION_TYPE_GENERATIVE_CHAT (4): + Used for use cases related to the Generative Chat agent. + It's used for Generative chat engine only, the associated + data stores must enrolled with ``SOLUTION_TYPE_CHAT`` + solution. """ SOLUTION_TYPE_UNSPECIFIED = 0 SOLUTION_TYPE_RECOMMENDATION = 1 SOLUTION_TYPE_SEARCH = 2 SOLUTION_TYPE_CHAT = 3 + SOLUTION_TYPE_GENERATIVE_CHAT = 4 class SearchTier(proto.Enum): diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document.py index 36f8cc395f55..0134792178ce 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document.py @@ -27,6 +27,7 @@ package="google.cloud.discoveryengine.v1alpha", manifest={ "Document", + "ProcessedDocument", }, ) @@ -253,4 +254,34 @@ class AccessRestriction(proto.Message): ) +class ProcessedDocument(proto.Message): + r"""Document captures all raw metadata information of items to be + recommended or searched. + + + .. _oneof: https://proto-plus-python.readthedocs.io/en/stable/fields.html#oneofs-mutually-exclusive-fields + + Attributes: + json_data (str): + The JSON string representation of the + processed document. + + This field is a member of `oneof`_ ``processed_data_format``. + document (str): + Required. Full resource name of the referenced document, in + the format + ``projects/*/locations/*/collections/*/dataStores/*/branches/*/documents/*``. + """ + + json_data: str = proto.Field( + proto.STRING, + number=2, + oneof="processed_data_format", + ) + document: str = proto.Field( + proto.STRING, + number=1, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document_service.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document_service.py index c01f4e45e1a8..b4cae08db273 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document_service.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/document_service.py @@ -31,6 +31,7 @@ "CreateDocumentRequest", "UpdateDocumentRequest", "DeleteDocumentRequest", + "GetProcessedDocumentRequest", }, ) @@ -272,4 +273,77 @@ class DeleteDocumentRequest(proto.Message): ) +class GetProcessedDocumentRequest(proto.Message): + r"""Request message for + [DocumentService.GetDocument][google.cloud.discoveryengine.v1alpha.DocumentService.GetDocument] + method. + + Attributes: + name (str): + Required. Full resource name of + [Document][google.cloud.discoveryengine.v1alpha.Document], + such as + ``projects/{project}/locations/{location}/collections/{collection}/dataStores/{data_store}/branches/{branch}/documents/{document}``. + + If the caller does not have permission to access the + [Document][google.cloud.discoveryengine.v1alpha.Document], + regardless of whether or not it exists, a + ``PERMISSION_DENIED`` error is returned. + + If the requested + [Document][google.cloud.discoveryengine.v1alpha.Document] + does not exist, a ``NOT_FOUND`` error is returned. + processed_document_type (google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest.ProcessedDocumentType): + Required. What type of processing to return. + processed_document_format (google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest.ProcessedDocumentFormat): + What format output should be. If + unspecified, defaults to JSON. + """ + + class ProcessedDocumentType(proto.Enum): + r"""The type of processing to return in the response. + + Values: + PROCESSED_DOCUMENT_TYPE_UNSPECIFIED (0): + Default value. + PARSED_DOCUMENT (1): + Available for all data store parsing configs. + CHUNKED_DOCUMENT (2): + Only available if ChunkingConfig is enabeld + on the data store. + """ + PROCESSED_DOCUMENT_TYPE_UNSPECIFIED = 0 + PARSED_DOCUMENT = 1 + CHUNKED_DOCUMENT = 2 + + class ProcessedDocumentFormat(proto.Enum): + r"""The format of the returned processed document. If + unspecified, defaults to JSON. + + Values: + PROCESSED_DOCUMENT_FORMAT_UNSPECIFIED (0): + Default value. + JSON (1): + output format will be a JSON string + representation of processed document. + """ + PROCESSED_DOCUMENT_FORMAT_UNSPECIFIED = 0 + JSON = 1 + + name: str = proto.Field( + proto.STRING, + number=1, + ) + processed_document_type: ProcessedDocumentType = proto.Field( + proto.ENUM, + number=2, + enum=ProcessedDocumentType, + ) + processed_document_format: ProcessedDocumentFormat = proto.Field( + proto.ENUM, + number=3, + enum=ProcessedDocumentFormat, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/import_config.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/import_config.py index 268492cc53b0..6a15ffbe540d 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/import_config.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/import_config.py @@ -17,6 +17,7 @@ from typing import MutableMapping, MutableSequence +from google.protobuf import field_mask_pb2 # type: ignore from google.protobuf import timestamp_pb2 # type: ignore from google.rpc import status_pb2 # type: ignore from google.type import date_pb2 # type: ignore @@ -29,6 +30,12 @@ manifest={ "GcsSource", "BigQuerySource", + "SpannerSource", + "BigtableOptions", + "BigtableSource", + "FhirStoreSource", + "CloudSqlSource", + "FirestoreSource", "ImportErrorConfig", "ImportUserEventsRequest", "ImportUserEventsResponse", @@ -39,6 +46,8 @@ "ImportSuggestionDenyListEntriesRequest", "ImportSuggestionDenyListEntriesResponse", "ImportSuggestionDenyListEntriesMetadata", + "ImportCompletionSuggestionsResponse", + "ImportCompletionSuggestionsMetadata", }, ) @@ -178,6 +187,404 @@ class BigQuerySource(proto.Message): ) +class SpannerSource(proto.Message): + r"""The Spanner source for importing data + + Attributes: + project_id (str): + The project ID that the Spanner source is in + with a length limit of 128 characters. If not + specified, inherits the project ID from the + parent request. + instance_id (str): + Required. The instance ID of the source + Spanner table. + database_id (str): + Required. The database ID of the source + Spanner table. + table_id (str): + Required. The table name of the Spanner + database that needs to be imported. + enable_data_boost (bool): + Whether to apply data boost on Spanner export. Enabling this + option will incur additional cost. More info can be found + `here `__. + """ + + project_id: str = proto.Field( + proto.STRING, + number=1, + ) + instance_id: str = proto.Field( + proto.STRING, + number=2, + ) + database_id: str = proto.Field( + proto.STRING, + number=3, + ) + table_id: str = proto.Field( + proto.STRING, + number=4, + ) + enable_data_boost: bool = proto.Field( + proto.BOOL, + number=5, + ) + + +class BigtableOptions(proto.Message): + r"""The Bigtable Options object that contains information to + support the import. + + Attributes: + key_field_name (str): + The field name used for saving row key value in the + document. The name has to match the pattern + ``[a-zA-Z0-9][a-zA-Z0-9-_]*``. + families (MutableMapping[str, google.cloud.discoveryengine_v1alpha.types.BigtableOptions.BigtableColumnFamily]): + The mapping from family names to an object + that contains column families level information + for the given column family. If a family is not + present in this map it will be ignored. + """ + + class Type(proto.Enum): + r"""The type of values in a Bigtable column or column family. The values + are expected to be encoded using `HBase + Bytes.toBytes `__ + function when the encoding value is set to ``BINARY``. + + Values: + TYPE_UNSPECIFIED (0): + The type is unspecified. + STRING (1): + String type. + NUMBER (2): + Numerical type. + INTEGER (3): + Integer type. + VAR_INTEGER (4): + Variable length integer type. + BIG_NUMERIC (5): + BigDecimal type. + BOOLEAN (6): + Boolean type. + JSON (7): + JSON type. + """ + TYPE_UNSPECIFIED = 0 + STRING = 1 + NUMBER = 2 + INTEGER = 3 + VAR_INTEGER = 4 + BIG_NUMERIC = 5 + BOOLEAN = 6 + JSON = 7 + + class Encoding(proto.Enum): + r"""The encoding mode of a Bigtable column or column family. + + Values: + ENCODING_UNSPECIFIED (0): + The encoding is unspecified. + TEXT (1): + Text encoding. + BINARY (2): + Binary encoding. + """ + ENCODING_UNSPECIFIED = 0 + TEXT = 1 + BINARY = 2 + + class BigtableColumnFamily(proto.Message): + r"""The column family of the Bigtable. + + Attributes: + field_name (str): + The field name to use for this column family in the + document. The name has to match the pattern + ``[a-zA-Z0-9][a-zA-Z0-9-_]*``. If not set, it is parsed from + the family name with best effort. However, due to different + naming patterns, field name collisions could happen, where + parsing behavior is undefined. + encoding (google.cloud.discoveryengine_v1alpha.types.BigtableOptions.Encoding): + The encoding mode of the values when the type is not STRING. + Acceptable encoding values are: + + - ``TEXT``: indicates values are alphanumeric text strings. + - ``BINARY``: indicates values are encoded using + ``HBase Bytes.toBytes`` family of functions. This can be + overridden for a specific column by listing that column + in ``columns`` and specifying an encoding for it. + type_ (google.cloud.discoveryengine_v1alpha.types.BigtableOptions.Type): + The type of values in this column family. The values are + expected to be encoded using ``HBase Bytes.toBytes`` + function when the encoding value is set to ``BINARY``. + columns (MutableSequence[google.cloud.discoveryengine_v1alpha.types.BigtableOptions.BigtableColumn]): + The list of objects that contains column + level information for each column. If a column + is not present in this list it will be ignored. + """ + + field_name: str = proto.Field( + proto.STRING, + number=1, + ) + encoding: "BigtableOptions.Encoding" = proto.Field( + proto.ENUM, + number=2, + enum="BigtableOptions.Encoding", + ) + type_: "BigtableOptions.Type" = proto.Field( + proto.ENUM, + number=3, + enum="BigtableOptions.Type", + ) + columns: MutableSequence[ + "BigtableOptions.BigtableColumn" + ] = proto.RepeatedField( + proto.MESSAGE, + number=4, + message="BigtableOptions.BigtableColumn", + ) + + class BigtableColumn(proto.Message): + r"""The column of the Bigtable. + + Attributes: + qualifier (bytes): + Required. Qualifier of the column. If it + cannot be decoded with utf-8, use a base-64 + encoded string instead. + field_name (str): + The field name to use for this column in the document. The + name has to match the pattern ``[a-zA-Z0-9][a-zA-Z0-9-_]*``. + If not set, it is parsed from the qualifier bytes with best + effort. However, due to different naming patterns, field + name collisions could happen, where parsing behavior is + undefined. + encoding (google.cloud.discoveryengine_v1alpha.types.BigtableOptions.Encoding): + The encoding mode of the values when the type is not + ``STRING``. Acceptable encoding values are: + + - ``TEXT``: indicates values are alphanumeric text strings. + - ``BINARY``: indicates values are encoded using + ``HBase Bytes.toBytes`` family of functions. This can be + overridden for a specific column by listing that column + in ``columns`` and specifying an encoding for it. + type_ (google.cloud.discoveryengine_v1alpha.types.BigtableOptions.Type): + The type of values in this column family. The values are + expected to be encoded using ``HBase Bytes.toBytes`` + function when the encoding value is set to ``BINARY``. + """ + + qualifier: bytes = proto.Field( + proto.BYTES, + number=1, + ) + field_name: str = proto.Field( + proto.STRING, + number=2, + ) + encoding: "BigtableOptions.Encoding" = proto.Field( + proto.ENUM, + number=3, + enum="BigtableOptions.Encoding", + ) + type_: "BigtableOptions.Type" = proto.Field( + proto.ENUM, + number=4, + enum="BigtableOptions.Type", + ) + + key_field_name: str = proto.Field( + proto.STRING, + number=1, + ) + families: MutableMapping[str, BigtableColumnFamily] = proto.MapField( + proto.STRING, + proto.MESSAGE, + number=2, + message=BigtableColumnFamily, + ) + + +class BigtableSource(proto.Message): + r"""The Cloud Bigtable source for importing data. + + Attributes: + project_id (str): + The project ID that the Bigtable source is in + with a length limit of 128 characters. If not + specified, inherits the project ID from the + parent request. + instance_id (str): + Required. The instance ID of the Cloud + Bigtable that needs to be imported. + table_id (str): + Required. The table ID of the Cloud Bigtable + that needs to be imported. + bigtable_options (google.cloud.discoveryengine_v1alpha.types.BigtableOptions): + Required. Bigtable options that contains + information needed when parsing data into typed + structures. For example, column type + annotations. + """ + + project_id: str = proto.Field( + proto.STRING, + number=1, + ) + instance_id: str = proto.Field( + proto.STRING, + number=2, + ) + table_id: str = proto.Field( + proto.STRING, + number=3, + ) + bigtable_options: "BigtableOptions" = proto.Field( + proto.MESSAGE, + number=4, + message="BigtableOptions", + ) + + +class FhirStoreSource(proto.Message): + r"""Cloud FhirStore source import data from. + + Attributes: + fhir_store (str): + Required. The full resource name of the FHIR store to import + data from, in the format of + ``projects/{project}/locations/{location}/datasets/{dataset}/fhirStores/{fhir_store}``. + gcs_staging_dir (str): + Intermediate Cloud Storage directory used for + the import with a length limit of 2,000 + characters. Can be specified if one wants to + have the FhirStore export to a specific Cloud + Storage directory. + """ + + fhir_store: str = proto.Field( + proto.STRING, + number=1, + ) + gcs_staging_dir: str = proto.Field( + proto.STRING, + number=2, + ) + + +class CloudSqlSource(proto.Message): + r"""Cloud SQL source import data from. + + Attributes: + project_id (str): + The project ID that the Cloud SQL source is + in with a length limit of 128 characters. If not + specified, inherits the project ID from the + parent request. + instance_id (str): + Required. The Cloud SQL instance to copy the + data from with a length limit of 256 characters. + database_id (str): + Required. The Cloud SQL database to copy the + data from with a length limit of 256 characters. + table_id (str): + Required. The Cloud SQL table to copy the + data from with a length limit of 256 characters. + gcs_staging_dir (str): + Intermediate Cloud Storage directory used for + the import with a length limit of 2,000 + characters. Can be specified if one wants to + have the Cloud SQL export to a specific Cloud + Storage directory. + + Please ensure that the Cloud SQL service account + has the necessary Cloud Storage Admin + permissions to access the specified Cloud + Storage directory. + offload (bool): + Option for serverless export. Enabling this option will + incur additional cost. More info can be found + `here `__. + """ + + project_id: str = proto.Field( + proto.STRING, + number=1, + ) + instance_id: str = proto.Field( + proto.STRING, + number=2, + ) + database_id: str = proto.Field( + proto.STRING, + number=3, + ) + table_id: str = proto.Field( + proto.STRING, + number=4, + ) + gcs_staging_dir: str = proto.Field( + proto.STRING, + number=5, + ) + offload: bool = proto.Field( + proto.BOOL, + number=6, + ) + + +class FirestoreSource(proto.Message): + r"""Firestore source import data from. + + Attributes: + project_id (str): + The project ID that the Cloud SQL source is + in with a length limit of 128 characters. If not + specified, inherits the project ID from the + parent request. + database_id (str): + Required. The Firestore database to copy the + data from with a length limit of 256 characters. + collection_id (str): + Required. The Firestore collection to copy + the data from with a length limit of 1,500 + characters. + gcs_staging_dir (str): + Intermediate Cloud Storage directory used for + the import with a length limit of 2,000 + characters. Can be specified if one wants to + have the Firestore export to a specific Cloud + Storage directory. + + Please ensure that the Firestore service account + has the necessary Cloud Storage Admin + permissions to access the specified Cloud + Storage directory. + """ + + project_id: str = proto.Field( + proto.STRING, + number=1, + ) + database_id: str = proto.Field( + proto.STRING, + number=2, + ) + collection_id: str = proto.Field( + proto.STRING, + number=3, + ) + gcs_staging_dir: str = proto.Field( + proto.STRING, + number=4, + ) + + class ImportErrorConfig(proto.Message): r"""Configuration of destination for Import related errors. @@ -376,6 +783,8 @@ class ImportDocumentsMetadata(proto.Message): failure_count (int): Count of entries that encountered errors while processing. + total_count (int): + Total count of entries that were processed. """ create_time: timestamp_pb2.Timestamp = proto.Field( @@ -396,6 +805,10 @@ class ImportDocumentsMetadata(proto.Message): proto.INT64, number=4, ) + total_count: int = proto.Field( + proto.INT64, + number=5, + ) class ImportDocumentsRequest(proto.Message): @@ -421,6 +834,26 @@ class ImportDocumentsRequest(proto.Message): bigquery_source (google.cloud.discoveryengine_v1alpha.types.BigQuerySource): BigQuery input source. + This field is a member of `oneof`_ ``source``. + fhir_store_source (google.cloud.discoveryengine_v1alpha.types.FhirStoreSource): + FhirStore input source. + + This field is a member of `oneof`_ ``source``. + spanner_source (google.cloud.discoveryengine_v1alpha.types.SpannerSource): + Spanner input source. + + This field is a member of `oneof`_ ``source``. + cloud_sql_source (google.cloud.discoveryengine_v1alpha.types.CloudSqlSource): + Cloud SQL input source. + + This field is a member of `oneof`_ ``source``. + firestore_source (google.cloud.discoveryengine_v1alpha.types.FirestoreSource): + Firestore input source. + + This field is a member of `oneof`_ ``source``. + bigtable_source (google.cloud.discoveryengine_v1alpha.types.BigtableSource): + Cloud Bigtable input source. + This field is a member of `oneof`_ ``source``. parent (str): Required. The parent branch resource name, such as @@ -433,6 +866,10 @@ class ImportDocumentsRequest(proto.Message): The mode of reconciliation between existing documents and the documents to be imported. Defaults to [ReconciliationMode.INCREMENTAL][google.cloud.discoveryengine.v1alpha.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL]. + update_mask (google.protobuf.field_mask_pb2.FieldMask): + Indicates which fields in the provided + imported documents to update. If not set, the + default is to update all fields. auto_generate_ids (bool): Whether to automatically generate IDs for the documents if absent. @@ -450,52 +887,61 @@ class ImportDocumentsRequest(proto.Message): [id_field][google.cloud.discoveryengine.v1alpha.ImportDocumentsRequest.id_field], otherwise, documents without IDs fail to be imported. - Only set this field when using - [GcsSource][google.cloud.discoveryengine.v1alpha.GcsSource] - or - [BigQuerySource][google.cloud.discoveryengine.v1alpha.BigQuerySource], - and when - [GcsSource.data_schema][google.cloud.discoveryengine.v1alpha.GcsSource.data_schema] - or - [BigQuerySource.data_schema][google.cloud.discoveryengine.v1alpha.BigQuerySource.data_schema] - is ``custom`` or ``csv``. Otherwise, an INVALID_ARGUMENT - error is thrown. + Supported data sources: + + - [GcsSource][google.cloud.discoveryengine.v1alpha.GcsSource]. + [GcsSource.data_schema][google.cloud.discoveryengine.v1alpha.GcsSource.data_schema] + must be ``custom`` or ``csv``. Otherwise, an + INVALID_ARGUMENT error is thrown. + - [BigQuerySource][google.cloud.discoveryengine.v1alpha.BigQuerySource]. + [BigQuerySource.data_schema][google.cloud.discoveryengine.v1alpha.BigQuerySource.data_schema] + must be ``custom`` or ``csv``. Otherwise, an + INVALID_ARGUMENT error is thrown. + - [SpannerSource][google.cloud.discoveryengine.v1alpha.SpannerSource]. + - [CloudSqlSource][google.cloud.discoveryengine.v1alpha.CloudSqlSource]. + - [FirestoreSource][google.cloud.discoveryengine.v1alpha.FirestoreSource]. + - [BigtableSource][google.cloud.discoveryengine.v1alpha.BigtableSource]. id_field (str): - The field in the Cloud Storage and BigQuery sources that - indicates the unique IDs of the documents. + The field indicates the ID field or column to be used as + unique IDs of the documents. For [GcsSource][google.cloud.discoveryengine.v1alpha.GcsSource] it is the key of the JSON field. For instance, ``my_id`` for - JSON ``{"my_id": "some_uuid"}``. For - [BigQuerySource][google.cloud.discoveryengine.v1alpha.BigQuerySource] - it is the column name of the BigQuery table where the unique - ids are stored. + JSON ``{"my_id": "some_uuid"}``. For others, it may be the + column name of the table where the unique ids are stored. - The values of the JSON field or the BigQuery column are used - as the + The values of the JSON field or the table column are used as + the [Document.id][google.cloud.discoveryengine.v1alpha.Document.id]s. - The JSON field or the BigQuery column must be of string - type, and the values must be set as valid strings conform to + The JSON field or the table column must be of string type, + and the values must be set as valid strings conform to `RFC-1034 `__ with 1-63 characters. Otherwise, documents without valid IDs fail to be imported. - Only set this field when using - [GcsSource][google.cloud.discoveryengine.v1alpha.GcsSource] - or - [BigQuerySource][google.cloud.discoveryengine.v1alpha.BigQuerySource], - and when - [GcsSource.data_schema][google.cloud.discoveryengine.v1alpha.GcsSource.data_schema] - or - [BigQuerySource.data_schema][google.cloud.discoveryengine.v1alpha.BigQuerySource.data_schema] - is ``custom``. And only set this field when + Only set this field when [auto_generate_ids][google.cloud.discoveryengine.v1alpha.ImportDocumentsRequest.auto_generate_ids] is unset or set as ``false``. Otherwise, an INVALID_ARGUMENT error is thrown. If it is unset, a default value ``_id`` is used when importing from the allowed data sources. + + Supported data sources: + + - [GcsSource][google.cloud.discoveryengine.v1alpha.GcsSource]. + [GcsSource.data_schema][google.cloud.discoveryengine.v1alpha.GcsSource.data_schema] + must be ``custom`` or ``csv``. Otherwise, an + INVALID_ARGUMENT error is thrown. + - [BigQuerySource][google.cloud.discoveryengine.v1alpha.BigQuerySource]. + [BigQuerySource.data_schema][google.cloud.discoveryengine.v1alpha.BigQuerySource.data_schema] + must be ``custom`` or ``csv``. Otherwise, an + INVALID_ARGUMENT error is thrown. + - [SpannerSource][google.cloud.discoveryengine.v1alpha.SpannerSource]. + - [CloudSqlSource][google.cloud.discoveryengine.v1alpha.CloudSqlSource]. + - [FirestoreSource][google.cloud.discoveryengine.v1alpha.FirestoreSource]. + - [BigtableSource][google.cloud.discoveryengine.v1alpha.BigtableSource]. """ class ReconciliationMode(proto.Enum): @@ -554,6 +1000,36 @@ class InlineSource(proto.Message): oneof="source", message="BigQuerySource", ) + fhir_store_source: "FhirStoreSource" = proto.Field( + proto.MESSAGE, + number=10, + oneof="source", + message="FhirStoreSource", + ) + spanner_source: "SpannerSource" = proto.Field( + proto.MESSAGE, + number=11, + oneof="source", + message="SpannerSource", + ) + cloud_sql_source: "CloudSqlSource" = proto.Field( + proto.MESSAGE, + number=12, + oneof="source", + message="CloudSqlSource", + ) + firestore_source: "FirestoreSource" = proto.Field( + proto.MESSAGE, + number=13, + oneof="source", + message="FirestoreSource", + ) + bigtable_source: "BigtableSource" = proto.Field( + proto.MESSAGE, + number=15, + oneof="source", + message="BigtableSource", + ) parent: str = proto.Field( proto.STRING, number=1, @@ -568,6 +1044,11 @@ class InlineSource(proto.Message): number=6, enum=ReconciliationMode, ) + update_mask: field_mask_pb2.FieldMask = proto.Field( + proto.MESSAGE, + number=7, + message=field_mask_pb2.FieldMask, + ) auto_generate_ids: bool = proto.Field( proto.BOOL, number=8, @@ -733,4 +1214,65 @@ class ImportSuggestionDenyListEntriesMetadata(proto.Message): ) +class ImportCompletionSuggestionsResponse(proto.Message): + r"""Response of the + [CompletionService.ImportCompletionSuggestions][google.cloud.discoveryengine.v1alpha.CompletionService.ImportCompletionSuggestions] + method. If the long running operation is done, this message is + returned by the google.longrunning.Operations.response field if the + operation is successful. + + Attributes: + error_samples (MutableSequence[google.rpc.status_pb2.Status]): + A sample of errors encountered while + processing the request. + success_count (int): + Count of + [CompletionSuggestion][google.cloud.discoveryengine.v1alpha.CompletionSuggestion]s + successfully imported. + failure_count (int): + Count of + [CompletionSuggestion][google.cloud.discoveryengine.v1alpha.CompletionSuggestion]s + that failed to be imported. + """ + + error_samples: MutableSequence[status_pb2.Status] = proto.RepeatedField( + proto.MESSAGE, + number=1, + message=status_pb2.Status, + ) + success_count: int = proto.Field( + proto.INT64, + number=2, + ) + failure_count: int = proto.Field( + proto.INT64, + number=3, + ) + + +class ImportCompletionSuggestionsMetadata(proto.Message): + r"""Metadata related to the progress of the + ImportCompletionSuggestions operation. This will be returned by + the google.longrunning.Operation.metadata field. + + Attributes: + create_time (google.protobuf.timestamp_pb2.Timestamp): + Operation create time. + update_time (google.protobuf.timestamp_pb2.Timestamp): + Operation last update time. If the operation + is done, this is also the finish time. + """ + + create_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=1, + message=timestamp_pb2.Timestamp, + ) + update_time: timestamp_pb2.Timestamp = proto.Field( + proto.MESSAGE, + number=2, + message=timestamp_pb2.Timestamp, + ) + + __all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/rank_service.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/rank_service.py new file mode 100644 index 000000000000..f2cc0ac2ac0f --- /dev/null +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/rank_service.py @@ -0,0 +1,152 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +from __future__ import annotations + +from typing import MutableMapping, MutableSequence + +import proto # type: ignore + +__protobuf__ = proto.module( + package="google.cloud.discoveryengine.v1alpha", + manifest={ + "RankingRecord", + "RankRequest", + "RankResponse", + }, +) + + +class RankingRecord(proto.Message): + r"""Record message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + Attributes: + id (str): + The unique ID to represent the record. + title (str): + The title of the record. Empty by default. At least one of + [title][google.cloud.discoveryengine.v1alpha.RankingRecord.title] + or + [content][google.cloud.discoveryengine.v1alpha.RankingRecord.content] + should be set otherwise an INVALID_ARGUMENT error is thrown. + content (str): + The content of the record. Empty by default. At least one of + [title][google.cloud.discoveryengine.v1alpha.RankingRecord.title] + or + [content][google.cloud.discoveryengine.v1alpha.RankingRecord.content] + should be set otherwise an INVALID_ARGUMENT error is thrown. + score (float): + The score of this record based on the given + query and selected model. + """ + + id: str = proto.Field( + proto.STRING, + number=1, + ) + title: str = proto.Field( + proto.STRING, + number=2, + ) + content: str = proto.Field( + proto.STRING, + number=3, + ) + score: float = proto.Field( + proto.FLOAT, + number=4, + ) + + +class RankRequest(proto.Message): + r"""Request message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + Attributes: + ranking_config (str): + Required. The resource name of the rank service config, such + as + ``projects/{project_num}/locations/{location_id}/rankingConfigs/default_ranking_config``. + model (str): + The identifier of the model to use. It is one of: + + - ``semantic-ranker-512@latest``: Semantic ranking model + with maxiumn input token size 512. + + It is set to ``semantic-ranker-512@latest`` by default if + unspecified. + top_n (int): + The number of results to return. If this is + unset or no bigger than zero, returns all + results. + query (str): + The query to use. + records (MutableSequence[google.cloud.discoveryengine_v1alpha.types.RankingRecord]): + Required. A list of records to rank. At most + 200 records to rank. + ignore_record_details_in_response (bool): + If true, the response will contain only + record ID and score. By default, it is false, + the response will contain record details. + """ + + ranking_config: str = proto.Field( + proto.STRING, + number=1, + ) + model: str = proto.Field( + proto.STRING, + number=2, + ) + top_n: int = proto.Field( + proto.INT32, + number=3, + ) + query: str = proto.Field( + proto.STRING, + number=4, + ) + records: MutableSequence["RankingRecord"] = proto.RepeatedField( + proto.MESSAGE, + number=5, + message="RankingRecord", + ) + ignore_record_details_in_response: bool = proto.Field( + proto.BOOL, + number=6, + ) + + +class RankResponse(proto.Message): + r"""Response message for + [RankService.Rank][google.cloud.discoveryengine.v1alpha.RankService.Rank] + method. + + Attributes: + records (MutableSequence[google.cloud.discoveryengine_v1alpha.types.RankingRecord]): + A list of records sorted by descending score. + """ + + records: MutableSequence["RankingRecord"] = proto.RepeatedField( + proto.MESSAGE, + number=5, + message="RankingRecord", + ) + + +__all__ = tuple(sorted(__protobuf__.manifest)) diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_service.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_service.py index 04b7285f6d9a..12a50c6eda1e 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_service.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_service.py @@ -60,9 +60,15 @@ class SearchRequest(proto.Message): page_size (int): Maximum number of [Document][google.cloud.discoveryengine.v1alpha.Document]s - to return. If unspecified, defaults to a reasonable value. - The maximum allowed value is 100. Values above 100 are - coerced to 100. + to return. The maximum allowed value depends on the data + type. Values above the maximum value are coerced to the + maximum value. + + - Websites with basic indexing: Default ``10``, Maximum + ``25``. + - Websites with advanced indexing: Default ``25``, Maximum + ``50``. + - Other: Default ``50``, Maximum ``100``. If this field is negative, an ``INVALID_ARGUMENT`` is returned. @@ -204,7 +210,7 @@ class SearchRequest(proto.Message): If [SearchRequest.EmbeddingSpec.EmbeddingVector.field_path][google.cloud.discoveryengine.v1alpha.SearchRequest.EmbeddingSpec.EmbeddingVector.field_path] is not provided, it will use - [ServingConfig.EmbeddingConfig.field_path][]. + [ServingConfig.EmbeddingConfig.field_path][google.cloud.discoveryengine.v1alpha.ServingConfig.embedding_config]. ranking_expression (str): The ranking expression controls the customized ranking on retrieval documents. This overrides @@ -526,9 +532,131 @@ class ConditionBoostSpec(proto.Message): ranking, but it is not blocked out completely. Setting to 0.0 means no boost applied. The boosting - condition is ignored. + condition is ignored. Only one of the (condition, boost) + combination or the boost_control_spec below are set. If both + are set then the global boost is ignored and the more + fine-grained boost_control_spec is applied. + boost_control_spec (google.cloud.discoveryengine_v1alpha.types.SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec): + Complex specification for custom ranking + based on customer defined attribute value. """ + class BoostControlSpec(proto.Message): + r"""Specification for custom ranking based on customer specified + attribute value. It provides more controls for customized + ranking than the simple (condition, boost) combination above. + + Attributes: + field_name (str): + The name of the field whose value will be + used to determine the boost amount. + attribute_type (google.cloud.discoveryengine_v1alpha.types.SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.AttributeType): + The attribute type to be used to determine the boost amount. + The attribute value can be derived from the field value of + the specified field_name. In the case of numerical it is + straightforward i.e. attribute_value = + numerical_field_value. In the case of freshness however, + attribute_value = (time.now() - datetime_field_value). + interpolation_type (google.cloud.discoveryengine_v1alpha.types.SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.InterpolationType): + The interpolation type to be applied to + connect the control points listed below. + control_points (MutableSequence[google.cloud.discoveryengine_v1alpha.types.SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.ControlPoint]): + The control points used to define the curve. The monotonic + function (defined through the interpolation_type above) + passes through the control points listed here. + """ + + class AttributeType(proto.Enum): + r"""The attribute(or function) for which the custom ranking is to + be applied. + + Values: + ATTRIBUTE_TYPE_UNSPECIFIED (0): + Unspecified AttributeType. + NUMERICAL (1): + The value of the numerical field will be used to dynamically + update the boost amount. In this case, the attribute_value + (the x value) of the control point will be the actual value + of the numerical field for which the boost_amount is + specified. + FRESHNESS (2): + For the freshness use case the attribute value will be the + duration between the current time and the date in the + datetime field specified. The value must be formatted as an + XSD ``dayTimeDuration`` value (a restricted subset of an ISO + 8601 duration value). The pattern for this is: + ``[nD][T[nH][nM][nS]]``. E.g. ``5D``, ``3DT12H30M``, + ``T24H``. + """ + ATTRIBUTE_TYPE_UNSPECIFIED = 0 + NUMERICAL = 1 + FRESHNESS = 2 + + class InterpolationType(proto.Enum): + r"""The interpolation type to be applied. Default will be linear + (Piecewise Linear). + + Values: + INTERPOLATION_TYPE_UNSPECIFIED (0): + Interpolation type is unspecified. In this + case, it defaults to Linear. + LINEAR (1): + Piecewise linear interpolation will be + applied. + """ + INTERPOLATION_TYPE_UNSPECIFIED = 0 + LINEAR = 1 + + class ControlPoint(proto.Message): + r"""The control points used to define the curve. The curve + defined through these control points can only be monotonically + increasing or decreasing(constant values are acceptable). + + Attributes: + attribute_value (str): + Can be one of: + + 1. The numerical field value. + 2. The duration spec for freshness: The value must be + formatted as an XSD ``dayTimeDuration`` value (a + restricted subset of an ISO 8601 duration value). The + pattern for this is: ``[nD][T[nH][nM][nS]]``. + boost_amount (float): + The value between -1 to 1 by which to boost the score if the + attribute_value evaluates to the value specified above. + """ + + attribute_value: str = proto.Field( + proto.STRING, + number=1, + ) + boost_amount: float = proto.Field( + proto.FLOAT, + number=2, + ) + + field_name: str = proto.Field( + proto.STRING, + number=1, + ) + attribute_type: "SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.AttributeType" = proto.Field( + proto.ENUM, + number=2, + enum="SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.AttributeType", + ) + interpolation_type: "SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.InterpolationType" = proto.Field( + proto.ENUM, + number=3, + enum="SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.InterpolationType", + ) + control_points: MutableSequence[ + "SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.ControlPoint" + ] = proto.RepeatedField( + proto.MESSAGE, + number=4, + message="SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec.ControlPoint", + ) + condition: str = proto.Field( proto.STRING, number=1, @@ -537,6 +665,11 @@ class ConditionBoostSpec(proto.Message): proto.FLOAT, number=2, ) + boost_control_spec: "SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec" = proto.Field( + proto.MESSAGE, + number=3, + message="SearchRequest.BoostSpec.ConditionBoostSpec.BoostControlSpec", + ) condition_boost_specs: MutableSequence[ "SearchRequest.BoostSpec.ConditionBoostSpec" @@ -655,6 +788,12 @@ class ContentSearchSpec(proto.Message): - If [DataStore.DocumentProcessingConfig.chunking_config][] is specified, it defaults to ``CHUNKS``. - Otherwise, it defaults to ``DOCUMENTS``. + chunk_spec (google.cloud.discoveryengine_v1alpha.types.SearchRequest.ContentSearchSpec.ChunkSpec): + Specifies the chunk spec to be returned from the search + response. Only available if the + [SearchRequest.ContentSearchSpec.search_result_mode][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.search_result_mode] + is set to + [CHUNKS][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.SearchResultMode.CHUNKS] """ class SearchResultMode(proto.Enum): @@ -781,6 +920,15 @@ class SummarySpec(proto.Message): model_spec (google.cloud.discoveryengine_v1alpha.types.SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec): If specified, the spec will be used to modify the model specification provided to the LLM. + use_semantic_chunks (bool): + If true, answer will be generated from most + relevant chunks from top search results. This + feature will improve summary quality. Please + note that with this feature enabled, not all top + search results will be referenced and included + in the reference list, so the citation source + index only points to the search results listed + in the reference list. """ class ModelPromptSpec(proto.Message): @@ -808,11 +956,14 @@ class ModelSpec(proto.Message): Supported values are: - ``stable``: string. Default value when no value is - specified. Uses a generally available, fine-tuned version - of the text-bison@001 model. - - ``preview``: string. (Public preview) Uses a fine-tuned - version of the text-bison@002 model. This model works - only for summaries in English. + specified. Uses a generally available, fine-tuned model. + For more information, see `Answer generation model + versions and + lifecycle `__. + - ``preview``: string. (Public preview) Uses a preview + model. For more information, see `Answer generation model + versions and + lifecycle `__. """ version: str = proto.Field( @@ -852,6 +1003,10 @@ class ModelSpec(proto.Message): message="SearchRequest.ContentSearchSpec.SummarySpec.ModelSpec", ) ) + use_semantic_chunks: bool = proto.Field( + proto.BOOL, + number=8, + ) class ExtractiveContentSpec(proto.Message): r"""A specification for configuring the extractive content in a @@ -895,12 +1050,10 @@ class ExtractiveContentSpec(proto.Message): ``max_extractive_segment_count``. return_extractive_segment_score (bool): Specifies whether to return the confidence score from the - extractive segments in each search result. The default value - is ``false``. - - Note: this is a priavte preview feature and only works for - allowlisted users, please reach out to Cloud Support team if - you want to use it. + extractive segments in each search result. This feature is + available only for new or allowlisted data stores. To + allowlist your data store, please contact your Customer + Engineer. The default value is ``false``. num_previous_segments (int): Specifies whether to also include the adjacent from each selected segments. Return at most ``num_previous_segments`` @@ -931,6 +1084,35 @@ class ExtractiveContentSpec(proto.Message): number=5, ) + class ChunkSpec(proto.Message): + r"""Specifies the chunk spec to be returned from the search response. + Only available if the + [SearchRequest.ContentSearchSpec.search_result_mode][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.search_result_mode] + is set to + [CHUNKS][google.cloud.discoveryengine.v1alpha.SearchRequest.ContentSearchSpec.SearchResultMode.CHUNKS] + + Attributes: + num_previous_chunks (int): + The number of previous chunks to be returned + of the current chunk. The maximum allowed value + is 3. If not specified, no previous chunks will + be returned. + num_next_chunks (int): + The number of next chunks to be returned of + the current chunk. The maximum allowed value is + 3. If not specified, no next chunks will be + returned. + """ + + num_previous_chunks: int = proto.Field( + proto.INT32, + number=1, + ) + num_next_chunks: int = proto.Field( + proto.INT32, + number=2, + ) + snippet_spec: "SearchRequest.ContentSearchSpec.SnippetSpec" = proto.Field( proto.MESSAGE, number=1, @@ -953,6 +1135,11 @@ class ExtractiveContentSpec(proto.Message): enum="SearchRequest.ContentSearchSpec.SearchResultMode", ) ) + chunk_spec: "SearchRequest.ContentSearchSpec.ChunkSpec" = proto.Field( + proto.MESSAGE, + number=5, + message="SearchRequest.ContentSearchSpec.ChunkSpec", + ) class EmbeddingSpec(proto.Message): r"""The specification that uses customized query embedding vector @@ -1480,8 +1667,30 @@ class Reference(proto.Message): ``projects/*/locations/*/collections/*/dataStores/*/branches/*/documents/*``. uri (str): Cloud Storage or HTTP uri for the document. + chunk_contents (MutableSequence[google.cloud.discoveryengine_v1alpha.types.SearchResponse.Summary.Reference.ChunkContent]): + List of cited chunk contents derived from + document content. """ + class ChunkContent(proto.Message): + r"""Chunk content. + + Attributes: + content (str): + Chunk textual content. + page_identifier (str): + Page identifier. + """ + + content: str = proto.Field( + proto.STRING, + number=1, + ) + page_identifier: str = proto.Field( + proto.STRING, + number=2, + ) + title: str = proto.Field( proto.STRING, number=1, @@ -1494,6 +1703,13 @@ class Reference(proto.Message): proto.STRING, number=3, ) + chunk_contents: MutableSequence[ + "SearchResponse.Summary.Reference.ChunkContent" + ] = proto.RepeatedField( + proto.MESSAGE, + number=4, + message="SearchResponse.Summary.Reference.ChunkContent", + ) class SummaryWithMetadata(proto.Message): r"""Summary with metadata information. diff --git a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_tuning_service.py b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_tuning_service.py index d9ff396cb75a..79af1dbe973d 100644 --- a/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_tuning_service.py +++ b/packages/google-cloud-discoveryengine/google/cloud/discoveryengine_v1alpha/types/search_tuning_service.py @@ -72,8 +72,8 @@ class GcsTrainingInput(proto.Message): delimited jsonl/ndjson file. For search-tuning model, each line should have the \_id, - title and text. Example: {"_id": "doc1", title: "relevant - doc", "text": "relevant text"} + title and text. Example: + ``{"_id": "doc1", title: "relevant doc", "text": "relevant text"}`` query_data_path (str): The gcs query data which could be associated in train data. The data path format is @@ -160,6 +160,8 @@ class TrainCustomModelResponse(proto.Message): Won't deploy. - **in-progress**: Model training is in progress. - **ready**: The model is ready for serving. + metrics (MutableMapping[str, float]): + The metrics of the trained model. """ error_samples: MutableSequence[status_pb2.Status] = proto.RepeatedField( @@ -176,6 +178,11 @@ class TrainCustomModelResponse(proto.Message): proto.STRING, number=3, ) + metrics: MutableMapping[str, float] = proto.MapField( + proto.STRING, + proto.DOUBLE, + number=4, + ) class TrainCustomModelMetadata(proto.Message): diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_async.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_async.py new file mode 100644 index 000000000000..0258cb3c2030 --- /dev/null +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_async.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for GetProcessedDocument +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-discoveryengine + + +# [START discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_async] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import discoveryengine_v1alpha + + +async def sample_get_processed_document(): + # Create a client + client = discoveryengine_v1alpha.DocumentServiceAsyncClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.GetProcessedDocumentRequest( + name="name_value", + processed_document_type="CHUNKED_DOCUMENT", + ) + + # Make the request + response = await client.get_processed_document(request=request) + + # Handle the response + print(response) + +# [END discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_async] diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_sync.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_sync.py new file mode 100644 index 000000000000..7581ccb4e33d --- /dev/null +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_document_service_get_processed_document_sync.py @@ -0,0 +1,53 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for GetProcessedDocument +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-discoveryengine + + +# [START discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_sync] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import discoveryengine_v1alpha + + +def sample_get_processed_document(): + # Create a client + client = discoveryengine_v1alpha.DocumentServiceClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.GetProcessedDocumentRequest( + name="name_value", + processed_document_type="CHUNKED_DOCUMENT", + ) + + # Make the request + response = client.get_processed_document(request=request) + + # Handle the response + print(response) + +# [END discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_sync] diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_async.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_async.py index 6f9f40a45c69..5a694ea7edb7 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_async.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_async.py @@ -41,7 +41,7 @@ async def sample_create_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.CreateEngineRequest( parent="parent_value", diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_sync.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_sync.py index 91cdd1082d64..7e0e69e63750 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_sync.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_create_engine_sync.py @@ -41,7 +41,7 @@ def sample_create_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.CreateEngineRequest( parent="parent_value", diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_async.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_async.py index 36c28cbbc63f..c87725dfe050 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_async.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_async.py @@ -41,7 +41,7 @@ async def sample_update_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateEngineRequest( engine=engine, diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_sync.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_sync.py index 54038b9efbac..68ea0949794f 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_sync.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_engine_service_update_engine_sync.py @@ -41,7 +41,7 @@ def sample_update_engine(): # Initialize request argument(s) engine = discoveryengine_v1alpha.Engine() engine.display_name = "display_name_value" - engine.solution_type = "SOLUTION_TYPE_CHAT" + engine.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateEngineRequest( engine=engine, diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_async.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_async.py new file mode 100644 index 000000000000..e6b87332a454 --- /dev/null +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_async.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for Rank +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-discoveryengine + + +# [START discoveryengine_v1alpha_generated_RankService_Rank_async] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import discoveryengine_v1alpha + + +async def sample_rank(): + # Create a client + client = discoveryengine_v1alpha.RankServiceAsyncClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.RankRequest( + ranking_config="ranking_config_value", + ) + + # Make the request + response = await client.rank(request=request) + + # Handle the response + print(response) + +# [END discoveryengine_v1alpha_generated_RankService_Rank_async] diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_sync.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_sync.py new file mode 100644 index 000000000000..6de6d7dcfd13 --- /dev/null +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_rank_service_rank_sync.py @@ -0,0 +1,52 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Generated code. DO NOT EDIT! +# +# Snippet for Rank +# NOTE: This snippet has been automatically generated for illustrative purposes only. +# It may require modifications to work in your environment. + +# To install the latest published package dependency, execute the following: +# python3 -m pip install google-cloud-discoveryengine + + +# [START discoveryengine_v1alpha_generated_RankService_Rank_sync] +# This snippet has been automatically generated and should be regarded as a +# code template only. +# It will require modifications to work: +# - It may require correct/in-range values for request initialization. +# - It may require specifying regional endpoints when creating the service +# client as shown in: +# https://googleapis.dev/python/google-api-core/latest/client_options.html +from google.cloud import discoveryengine_v1alpha + + +def sample_rank(): + # Create a client + client = discoveryengine_v1alpha.RankServiceClient() + + # Initialize request argument(s) + request = discoveryengine_v1alpha.RankRequest( + ranking_config="ranking_config_value", + ) + + # Make the request + response = client.rank(request=request) + + # Handle the response + print(response) + +# [END discoveryengine_v1alpha_generated_RankService_Rank_sync] diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_async.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_async.py index 091e1cdd0b2e..051ac7c4bce3 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_async.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_async.py @@ -42,7 +42,7 @@ async def sample_update_serving_config(): serving_config = discoveryengine_v1alpha.ServingConfig() serving_config.media_config.content_watched_percentage_threshold = 0.3811 serving_config.display_name = "display_name_value" - serving_config.solution_type = "SOLUTION_TYPE_CHAT" + serving_config.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateServingConfigRequest( serving_config=serving_config, diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_sync.py b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_sync.py index 033587269636..993f234ce4f3 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_sync.py +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/discoveryengine_v1alpha_generated_serving_config_service_update_serving_config_sync.py @@ -42,7 +42,7 @@ def sample_update_serving_config(): serving_config = discoveryengine_v1alpha.ServingConfig() serving_config.media_config.content_watched_percentage_threshold = 0.3811 serving_config.display_name = "display_name_value" - serving_config.solution_type = "SOLUTION_TYPE_CHAT" + serving_config.solution_type = "SOLUTION_TYPE_GENERATIVE_CHAT" request = discoveryengine_v1alpha.UpdateServingConfigRequest( serving_config=serving_config, diff --git a/packages/google-cloud-discoveryengine/samples/generated_samples/snippet_metadata_google.cloud.discoveryengine.v1alpha.json b/packages/google-cloud-discoveryengine/samples/generated_samples/snippet_metadata_google.cloud.discoveryengine.v1alpha.json index eed106675397..48b52a060c5c 100644 --- a/packages/google-cloud-discoveryengine/samples/generated_samples/snippet_metadata_google.cloud.discoveryengine.v1alpha.json +++ b/packages/google-cloud-discoveryengine/samples/generated_samples/snippet_metadata_google.cloud.discoveryengine.v1alpha.json @@ -3742,6 +3742,167 @@ ], "title": "discoveryengine_v1alpha_generated_document_service_get_document_sync.py" }, + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.discoveryengine_v1alpha.DocumentServiceAsyncClient", + "shortName": "DocumentServiceAsyncClient" + }, + "fullName": "google.cloud.discoveryengine_v1alpha.DocumentServiceAsyncClient.get_processed_document", + "method": { + "fullName": "google.cloud.discoveryengine.v1alpha.DocumentService.GetProcessedDocument", + "service": { + "fullName": "google.cloud.discoveryengine.v1alpha.DocumentService", + "shortName": "DocumentService" + }, + "shortName": "GetProcessedDocument" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.cloud.discoveryengine_v1alpha.types.ProcessedDocument", + "shortName": "get_processed_document" + }, + "description": "Sample for GetProcessedDocument", + "file": "discoveryengine_v1alpha_generated_document_service_get_processed_document_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_async", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 46, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 49, + "start": 47, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 50, + "type": "RESPONSE_HANDLING" + } + ], + "title": "discoveryengine_v1alpha_generated_document_service_get_processed_document_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.discoveryengine_v1alpha.DocumentServiceClient", + "shortName": "DocumentServiceClient" + }, + "fullName": "google.cloud.discoveryengine_v1alpha.DocumentServiceClient.get_processed_document", + "method": { + "fullName": "google.cloud.discoveryengine.v1alpha.DocumentService.GetProcessedDocument", + "service": { + "fullName": "google.cloud.discoveryengine.v1alpha.DocumentService", + "shortName": "DocumentService" + }, + "shortName": "GetProcessedDocument" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.discoveryengine_v1alpha.types.GetProcessedDocumentRequest" + }, + { + "name": "name", + "type": "str" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.cloud.discoveryengine_v1alpha.types.ProcessedDocument", + "shortName": "get_processed_document" + }, + "description": "Sample for GetProcessedDocument", + "file": "discoveryengine_v1alpha_generated_document_service_get_processed_document_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "discoveryengine_v1alpha_generated_DocumentService_GetProcessedDocument_sync", + "segments": [ + { + "end": 52, + "start": 27, + "type": "FULL" + }, + { + "end": 52, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 46, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 49, + "start": 47, + "type": "REQUEST_EXECUTION" + }, + { + "end": 53, + "start": 50, + "type": "RESPONSE_HANDLING" + } + ], + "title": "discoveryengine_v1alpha_generated_document_service_get_processed_document_sync.py" + }, { "canonical": true, "clientMethod": { @@ -5843,6 +6004,159 @@ ], "title": "discoveryengine_v1alpha_generated_estimate_billing_service_estimate_data_size_sync.py" }, + { + "canonical": true, + "clientMethod": { + "async": true, + "client": { + "fullName": "google.cloud.discoveryengine_v1alpha.RankServiceAsyncClient", + "shortName": "RankServiceAsyncClient" + }, + "fullName": "google.cloud.discoveryengine_v1alpha.RankServiceAsyncClient.rank", + "method": { + "fullName": "google.cloud.discoveryengine.v1alpha.RankService.Rank", + "service": { + "fullName": "google.cloud.discoveryengine.v1alpha.RankService", + "shortName": "RankService" + }, + "shortName": "Rank" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.discoveryengine_v1alpha.types.RankRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.cloud.discoveryengine_v1alpha.types.RankResponse", + "shortName": "rank" + }, + "description": "Sample for Rank", + "file": "discoveryengine_v1alpha_generated_rank_service_rank_async.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "discoveryengine_v1alpha_generated_RankService_Rank_async", + "segments": [ + { + "end": 51, + "start": 27, + "type": "FULL" + }, + { + "end": 51, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 52, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "discoveryengine_v1alpha_generated_rank_service_rank_async.py" + }, + { + "canonical": true, + "clientMethod": { + "client": { + "fullName": "google.cloud.discoveryengine_v1alpha.RankServiceClient", + "shortName": "RankServiceClient" + }, + "fullName": "google.cloud.discoveryengine_v1alpha.RankServiceClient.rank", + "method": { + "fullName": "google.cloud.discoveryengine.v1alpha.RankService.Rank", + "service": { + "fullName": "google.cloud.discoveryengine.v1alpha.RankService", + "shortName": "RankService" + }, + "shortName": "Rank" + }, + "parameters": [ + { + "name": "request", + "type": "google.cloud.discoveryengine_v1alpha.types.RankRequest" + }, + { + "name": "retry", + "type": "google.api_core.retry.Retry" + }, + { + "name": "timeout", + "type": "float" + }, + { + "name": "metadata", + "type": "Sequence[Tuple[str, str]" + } + ], + "resultType": "google.cloud.discoveryengine_v1alpha.types.RankResponse", + "shortName": "rank" + }, + "description": "Sample for Rank", + "file": "discoveryengine_v1alpha_generated_rank_service_rank_sync.py", + "language": "PYTHON", + "origin": "API_DEFINITION", + "regionTag": "discoveryengine_v1alpha_generated_RankService_Rank_sync", + "segments": [ + { + "end": 51, + "start": 27, + "type": "FULL" + }, + { + "end": 51, + "start": 27, + "type": "SHORT" + }, + { + "end": 40, + "start": 38, + "type": "CLIENT_INITIALIZATION" + }, + { + "end": 45, + "start": 41, + "type": "REQUEST_INITIALIZATION" + }, + { + "end": 48, + "start": 46, + "type": "REQUEST_EXECUTION" + }, + { + "end": 52, + "start": 49, + "type": "RESPONSE_HANDLING" + } + ], + "title": "discoveryengine_v1alpha_generated_rank_service_rank_sync.py" + }, { "canonical": true, "clientMethod": { diff --git a/packages/google-cloud-discoveryengine/scripts/fixup_discoveryengine_v1alpha_keywords.py b/packages/google-cloud-discoveryengine/scripts/fixup_discoveryengine_v1alpha_keywords.py index 2b7f91b13aab..6ab12b1138c7 100644 --- a/packages/google-cloud-discoveryengine/scripts/fixup_discoveryengine_v1alpha_keywords.py +++ b/packages/google-cloud-discoveryengine/scripts/fixup_discoveryengine_v1alpha_keywords.py @@ -67,11 +67,12 @@ class discoveryengineCallTransformer(cst.CSTTransformer): 'get_document': ('name', ), 'get_document_processing_config': ('name', ), 'get_engine': ('name', ), + 'get_processed_document': ('name', 'processed_document_type', 'processed_document_format', ), 'get_schema': ('name', ), 'get_serving_config': ('name', ), 'get_site_search_engine': ('name', ), 'get_target_site': ('name', ), - 'import_documents': ('parent', 'inline_source', 'gcs_source', 'bigquery_source', 'error_config', 'reconciliation_mode', 'auto_generate_ids', 'id_field', ), + 'import_documents': ('parent', 'inline_source', 'gcs_source', 'bigquery_source', 'fhir_store_source', 'spanner_source', 'cloud_sql_source', 'firestore_source', 'bigtable_source', 'error_config', 'reconciliation_mode', 'update_mask', 'auto_generate_ids', 'id_field', ), 'import_suggestion_deny_list_entries': ('parent', 'inline_source', 'gcs_source', ), 'import_user_events': ('parent', 'inline_source', 'gcs_source', 'bigquery_source', 'error_config', ), 'list_chunks': ('parent', 'page_size', 'page_token', ), @@ -86,6 +87,7 @@ class discoveryengineCallTransformer(cst.CSTTransformer): 'purge_documents': ('parent', 'filter', 'gcs_source', 'error_config', 'force', ), 'purge_suggestion_deny_list_entries': ('parent', ), 'purge_user_events': ('parent', 'filter', 'force', ), + 'rank': ('ranking_config', 'records', 'model', 'top_n', 'query', 'ignore_record_details_in_response', ), 'recommend': ('serving_config', 'user_event', 'page_size', 'filter', 'validate_only', 'params', 'user_labels', ), 'recrawl_uris': ('site_search_engine', 'uris', ), 'resume_engine': ('name', ), diff --git a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_conversational_search_service.py b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_conversational_search_service.py index d74d95e7aca4..3d22a0023c63 100644 --- a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_conversational_search_service.py +++ b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_conversational_search_service.py @@ -3616,6 +3616,12 @@ def test_create_conversation_rest(request_type): "title": "title_value", "document": "document_value", "uri": "uri_value", + "chunk_contents": [ + { + "content": "content_value", + "page_identifier": "page_identifier_value", + } + ], } ], }, @@ -4291,6 +4297,12 @@ def test_update_conversation_rest(request_type): "title": "title_value", "document": "document_value", "uri": "uri_value", + "chunk_contents": [ + { + "content": "content_value", + "page_identifier": "page_identifier_value", + } + ], } ], }, diff --git a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_document_service.py b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_document_service.py index 47206f0e969f..13565fe5a3cc 100644 --- a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_document_service.py +++ b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_document_service.py @@ -3304,6 +3304,308 @@ async def test_purge_documents_field_headers_async(): ) in kw["metadata"] +@pytest.mark.parametrize( + "request_type", + [ + document_service.GetProcessedDocumentRequest, + dict, + ], +) +def test_get_processed_document(request_type, transport: str = "grpc"): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = document.ProcessedDocument( + document="document_value", + json_data="json_data_value", + ) + response = client.get_processed_document(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = document_service.GetProcessedDocumentRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, document.ProcessedDocument) + assert response.document == "document_value" + + +def test_get_processed_document_empty_call(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + client.get_processed_document() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == document_service.GetProcessedDocumentRequest() + + +def test_get_processed_document_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are + # automatically populated, according to AIP-4235, with non-empty requests. + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Populate all string fields in the request which are not UUID4 + # since we want to check that UUID4 are populated automatically + # if they meet the requirements of AIP 4235. + request = document_service.GetProcessedDocumentRequest( + name="name_value", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + client.get_processed_document(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == document_service.GetProcessedDocumentRequest( + name="name_value", + ) + + +@pytest.mark.asyncio +async def test_get_processed_document_empty_call_async(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = DocumentServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + document.ProcessedDocument( + document="document_value", + ) + ) + response = await client.get_processed_document() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == document_service.GetProcessedDocumentRequest() + + +@pytest.mark.asyncio +async def test_get_processed_document_async( + transport: str = "grpc_asyncio", + request_type=document_service.GetProcessedDocumentRequest, +): + client = DocumentServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + document.ProcessedDocument( + document="document_value", + ) + ) + response = await client.get_processed_document(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = document_service.GetProcessedDocumentRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, document.ProcessedDocument) + assert response.document == "document_value" + + +@pytest.mark.asyncio +async def test_get_processed_document_async_from_dict(): + await test_get_processed_document_async(request_type=dict) + + +def test_get_processed_document_field_headers(): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = document_service.GetProcessedDocumentRequest() + + request.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + call.return_value = document.ProcessedDocument() + client.get_processed_document(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=name_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_get_processed_document_field_headers_async(): + client = DocumentServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = document_service.GetProcessedDocumentRequest() + + request.name = "name_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + document.ProcessedDocument() + ) + await client.get_processed_document(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=name_value", + ) in kw["metadata"] + + +def test_get_processed_document_flattened(): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = document.ProcessedDocument() + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + client.get_processed_document( + name="name_value", + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + arg = args[0].name + mock_val = "name_value" + assert arg == mock_val + + +def test_get_processed_document_flattened_error(): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.get_processed_document( + document_service.GetProcessedDocumentRequest(), + name="name_value", + ) + + +@pytest.mark.asyncio +async def test_get_processed_document_flattened_async(): + client = DocumentServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object( + type(client.transport.get_processed_document), "__call__" + ) as call: + # Designate an appropriate return value for the call. + call.return_value = document.ProcessedDocument() + + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + document.ProcessedDocument() + ) + # Call the method with a truthy value for each flattened field, + # using the keyword arguments to the method. + response = await client.get_processed_document( + name="name_value", + ) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + arg = args[0].name + mock_val = "name_value" + assert arg == mock_val + + +@pytest.mark.asyncio +async def test_get_processed_document_flattened_error_async(): + client = DocumentServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + await client.get_processed_document( + document_service.GetProcessedDocumentRequest(), + name="name_value", + ) + + @pytest.mark.parametrize( "request_type", [ @@ -5436,6 +5738,300 @@ def test_purge_documents_rest_error(): ) +@pytest.mark.parametrize( + "request_type", + [ + document_service.GetProcessedDocumentRequest, + dict, + ], +) +def test_get_processed_document_rest(request_type): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # send a request that will satisfy transcoding + request_init = { + "name": "projects/sample1/locations/sample2/dataStores/sample3/branches/sample4/documents/sample5" + } + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = document.ProcessedDocument( + document="document_value", + json_data="json_data_value", + ) + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + # Convert return value to protobuf type + return_value = document.ProcessedDocument.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + response = client.get_processed_document(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, document.ProcessedDocument) + assert response.document == "document_value" + + +def test_get_processed_document_rest_required_fields( + request_type=document_service.GetProcessedDocumentRequest, +): + transport_class = transports.DocumentServiceRestTransport + + request_init = {} + request_init["name"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson(pb_request, use_integers_for_enums=False) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).get_processed_document._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["name"] = "name_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).get_processed_document._get_unset_required_fields(jsonified_request) + # Check that path parameters and body parameters are not mixing in. + assert not set(unset_fields) - set( + ( + "processed_document_format", + "processed_document_type", + ) + ) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "name" in jsonified_request + assert jsonified_request["name"] == "name_value" + + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = document.ProcessedDocument() + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "get", + "query_params": pb_request, + } + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = document.ProcessedDocument.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + response = client.get_processed_document(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_get_processed_document_rest_unset_required_fields(): + transport = transports.DocumentServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.get_processed_document._get_unset_required_fields({}) + assert set(unset_fields) == ( + set( + ( + "processedDocumentFormat", + "processedDocumentType", + ) + ) + & set( + ( + "name", + "processedDocumentType", + ) + ) + ) + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_get_processed_document_rest_interceptors(null_interceptor): + transport = transports.DocumentServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.DocumentServiceRestInterceptor(), + ) + client = DocumentServiceClient(transport=transport) + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.DocumentServiceRestInterceptor, "post_get_processed_document" + ) as post, mock.patch.object( + transports.DocumentServiceRestInterceptor, "pre_get_processed_document" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + pb_message = document_service.GetProcessedDocumentRequest.pb( + document_service.GetProcessedDocumentRequest() + ) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = Response() + req.return_value.status_code = 200 + req.return_value.request = PreparedRequest() + req.return_value._content = document.ProcessedDocument.to_json( + document.ProcessedDocument() + ) + + request = document_service.GetProcessedDocumentRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = document.ProcessedDocument() + + client.get_processed_document( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + + +def test_get_processed_document_rest_bad_request( + transport: str = "rest", request_type=document_service.GetProcessedDocumentRequest +): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # send a request that will satisfy transcoding + request_init = { + "name": "projects/sample1/locations/sample2/dataStores/sample3/branches/sample4/documents/sample5" + } + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.get_processed_document(request) + + +def test_get_processed_document_rest_flattened(): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = document.ProcessedDocument() + + # get arguments that satisfy an http rule for this method + sample_request = { + "name": "projects/sample1/locations/sample2/dataStores/sample3/branches/sample4/documents/sample5" + } + + # get truthy value for each flattened field + mock_args = dict( + name="name_value", + ) + mock_args.update(sample_request) + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + # Convert return value to protobuf type + return_value = document.ProcessedDocument.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + client.get_processed_document(**mock_args) + + # Establish that the underlying call was made with the expected + # request object values. + assert len(req.mock_calls) == 1 + _, args, _ = req.mock_calls[0] + assert path_template.validate( + "%s/v1alpha/{name=projects/*/locations/*/dataStores/*/branches/*/documents/*}:getProcessedDocument" + % client.transport._host, + args[1], + ) + + +def test_get_processed_document_rest_flattened_error(transport: str = "rest"): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Attempting to call a method with both a request object and flattened + # fields is an error. + with pytest.raises(ValueError): + client.get_processed_document( + document_service.GetProcessedDocumentRequest(), + name="name_value", + ) + + +def test_get_processed_document_rest_error(): + client = DocumentServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + def test_credentials_transport_error(): # It is an error to provide credentials and a transport instance. transport = transports.DocumentServiceGrpcTransport( @@ -5582,6 +6178,7 @@ def test_document_service_base_transport(): "delete_document", "import_documents", "purge_documents", + "get_processed_document", "get_operation", "list_operations", ) @@ -5885,6 +6482,9 @@ def test_document_service_client_transport_session_collision(transport_name): session1 = client1.transport.purge_documents._session session2 = client2.transport.purge_documents._session assert session1 != session2 + session1 = client1.transport.get_processed_document._session + session2 = client2.transport.get_processed_document._session + assert session1 != session2 def test_document_service_grpc_transport_channel(): @@ -6110,8 +6710,39 @@ def test_parse_document_path(): assert expected == actual +def test_fhir_store_path(): + project = "cuttlefish" + location = "mussel" + dataset = "winkle" + fhir_store = "nautilus" + expected = "projects/{project}/locations/{location}/datasets/{dataset}/fhirStores/{fhir_store}".format( + project=project, + location=location, + dataset=dataset, + fhir_store=fhir_store, + ) + actual = DocumentServiceClient.fhir_store_path( + project, location, dataset, fhir_store + ) + assert expected == actual + + +def test_parse_fhir_store_path(): + expected = { + "project": "scallop", + "location": "abalone", + "dataset": "squid", + "fhir_store": "clam", + } + path = DocumentServiceClient.fhir_store_path(**expected) + + # Check that the path construction is reversible. + actual = DocumentServiceClient.parse_fhir_store_path(path) + assert expected == actual + + def test_common_billing_account_path(): - billing_account = "cuttlefish" + billing_account = "whelk" expected = "billingAccounts/{billing_account}".format( billing_account=billing_account, ) @@ -6121,7 +6752,7 @@ def test_common_billing_account_path(): def test_parse_common_billing_account_path(): expected = { - "billing_account": "mussel", + "billing_account": "octopus", } path = DocumentServiceClient.common_billing_account_path(**expected) @@ -6131,7 +6762,7 @@ def test_parse_common_billing_account_path(): def test_common_folder_path(): - folder = "winkle" + folder = "oyster" expected = "folders/{folder}".format( folder=folder, ) @@ -6141,7 +6772,7 @@ def test_common_folder_path(): def test_parse_common_folder_path(): expected = { - "folder": "nautilus", + "folder": "nudibranch", } path = DocumentServiceClient.common_folder_path(**expected) @@ -6151,7 +6782,7 @@ def test_parse_common_folder_path(): def test_common_organization_path(): - organization = "scallop" + organization = "cuttlefish" expected = "organizations/{organization}".format( organization=organization, ) @@ -6161,7 +6792,7 @@ def test_common_organization_path(): def test_parse_common_organization_path(): expected = { - "organization": "abalone", + "organization": "mussel", } path = DocumentServiceClient.common_organization_path(**expected) @@ -6171,7 +6802,7 @@ def test_parse_common_organization_path(): def test_common_project_path(): - project = "squid" + project = "winkle" expected = "projects/{project}".format( project=project, ) @@ -6181,7 +6812,7 @@ def test_common_project_path(): def test_parse_common_project_path(): expected = { - "project": "clam", + "project": "nautilus", } path = DocumentServiceClient.common_project_path(**expected) @@ -6191,8 +6822,8 @@ def test_parse_common_project_path(): def test_common_location_path(): - project = "whelk" - location = "octopus" + project = "scallop" + location = "abalone" expected = "projects/{project}/locations/{location}".format( project=project, location=location, @@ -6203,8 +6834,8 @@ def test_common_location_path(): def test_parse_common_location_path(): expected = { - "project": "oyster", - "location": "nudibranch", + "project": "squid", + "location": "clam", } path = DocumentServiceClient.common_location_path(**expected) diff --git a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_rank_service.py b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_rank_service.py new file mode 100644 index 000000000000..969430be6605 --- /dev/null +++ b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_rank_service.py @@ -0,0 +1,2662 @@ +# -*- coding: utf-8 -*- +# Copyright 2024 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +import os + +# try/except added for compatibility with python < 3.8 +try: + from unittest import mock + from unittest.mock import AsyncMock # pragma: NO COVER +except ImportError: # pragma: NO COVER + import mock + +from collections.abc import Iterable +import json +import math + +from google.api_core import gapic_v1, grpc_helpers, grpc_helpers_async, path_template +from google.api_core import api_core_version, client_options +from google.api_core import exceptions as core_exceptions +import google.auth +from google.auth import credentials as ga_credentials +from google.auth.exceptions import MutualTLSChannelError +from google.cloud.location import locations_pb2 +from google.longrunning import operations_pb2 # type: ignore +from google.oauth2 import service_account +from google.protobuf import json_format +import grpc +from grpc.experimental import aio +from proto.marshal.rules import wrappers +from proto.marshal.rules.dates import DurationRule, TimestampRule +import pytest +from requests import PreparedRequest, Request, Response +from requests.sessions import Session + +from google.cloud.discoveryengine_v1alpha.services.rank_service import ( + RankServiceAsyncClient, + RankServiceClient, + transports, +) +from google.cloud.discoveryengine_v1alpha.types import rank_service + + +def client_cert_source_callback(): + return b"cert bytes", b"key bytes" + + +# If default endpoint is localhost, then default mtls endpoint will be the same. +# This method modifies the default endpoint so the client can produce a different +# mtls endpoint for endpoint testing purposes. +def modify_default_endpoint(client): + return ( + "foo.googleapis.com" + if ("localhost" in client.DEFAULT_ENDPOINT) + else client.DEFAULT_ENDPOINT + ) + + +# If default endpoint template is localhost, then default mtls endpoint will be the same. +# This method modifies the default endpoint template so the client can produce a different +# mtls endpoint for endpoint testing purposes. +def modify_default_endpoint_template(client): + return ( + "test.{UNIVERSE_DOMAIN}" + if ("localhost" in client._DEFAULT_ENDPOINT_TEMPLATE) + else client._DEFAULT_ENDPOINT_TEMPLATE + ) + + +def test__get_default_mtls_endpoint(): + api_endpoint = "example.googleapis.com" + api_mtls_endpoint = "example.mtls.googleapis.com" + sandbox_endpoint = "example.sandbox.googleapis.com" + sandbox_mtls_endpoint = "example.mtls.sandbox.googleapis.com" + non_googleapi = "api.example.com" + + assert RankServiceClient._get_default_mtls_endpoint(None) is None + assert ( + RankServiceClient._get_default_mtls_endpoint(api_endpoint) == api_mtls_endpoint + ) + assert ( + RankServiceClient._get_default_mtls_endpoint(api_mtls_endpoint) + == api_mtls_endpoint + ) + assert ( + RankServiceClient._get_default_mtls_endpoint(sandbox_endpoint) + == sandbox_mtls_endpoint + ) + assert ( + RankServiceClient._get_default_mtls_endpoint(sandbox_mtls_endpoint) + == sandbox_mtls_endpoint + ) + assert RankServiceClient._get_default_mtls_endpoint(non_googleapi) == non_googleapi + + +def test__read_environment_variables(): + assert RankServiceClient._read_environment_variables() == (False, "auto", None) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}): + assert RankServiceClient._read_environment_variables() == (True, "auto", None) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"}): + assert RankServiceClient._read_environment_variables() == (False, "auto", None) + + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "Unsupported"} + ): + with pytest.raises(ValueError) as excinfo: + RankServiceClient._read_environment_variables() + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`" + ) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + assert RankServiceClient._read_environment_variables() == (False, "never", None) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + assert RankServiceClient._read_environment_variables() == ( + False, + "always", + None, + ) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}): + assert RankServiceClient._read_environment_variables() == (False, "auto", None) + + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "Unsupported"}): + with pytest.raises(MutualTLSChannelError) as excinfo: + RankServiceClient._read_environment_variables() + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`" + ) + + with mock.patch.dict(os.environ, {"GOOGLE_CLOUD_UNIVERSE_DOMAIN": "foo.com"}): + assert RankServiceClient._read_environment_variables() == ( + False, + "auto", + "foo.com", + ) + + +def test__get_client_cert_source(): + mock_provided_cert_source = mock.Mock() + mock_default_cert_source = mock.Mock() + + assert RankServiceClient._get_client_cert_source(None, False) is None + assert ( + RankServiceClient._get_client_cert_source(mock_provided_cert_source, False) + is None + ) + assert ( + RankServiceClient._get_client_cert_source(mock_provided_cert_source, True) + == mock_provided_cert_source + ) + + with mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", return_value=True + ): + with mock.patch( + "google.auth.transport.mtls.default_client_cert_source", + return_value=mock_default_cert_source, + ): + assert ( + RankServiceClient._get_client_cert_source(None, True) + is mock_default_cert_source + ) + assert ( + RankServiceClient._get_client_cert_source( + mock_provided_cert_source, "true" + ) + is mock_provided_cert_source + ) + + +@mock.patch.object( + RankServiceClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceClient), +) +@mock.patch.object( + RankServiceAsyncClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceAsyncClient), +) +def test__get_api_endpoint(): + api_override = "foo.com" + mock_client_cert_source = mock.Mock() + default_universe = RankServiceClient._DEFAULT_UNIVERSE + default_endpoint = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=default_universe + ) + mock_universe = "bar.com" + mock_endpoint = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=mock_universe + ) + + assert ( + RankServiceClient._get_api_endpoint( + api_override, mock_client_cert_source, default_universe, "always" + ) + == api_override + ) + assert ( + RankServiceClient._get_api_endpoint( + None, mock_client_cert_source, default_universe, "auto" + ) + == RankServiceClient.DEFAULT_MTLS_ENDPOINT + ) + assert ( + RankServiceClient._get_api_endpoint(None, None, default_universe, "auto") + == default_endpoint + ) + assert ( + RankServiceClient._get_api_endpoint(None, None, default_universe, "always") + == RankServiceClient.DEFAULT_MTLS_ENDPOINT + ) + assert ( + RankServiceClient._get_api_endpoint( + None, mock_client_cert_source, default_universe, "always" + ) + == RankServiceClient.DEFAULT_MTLS_ENDPOINT + ) + assert ( + RankServiceClient._get_api_endpoint(None, None, mock_universe, "never") + == mock_endpoint + ) + assert ( + RankServiceClient._get_api_endpoint(None, None, default_universe, "never") + == default_endpoint + ) + + with pytest.raises(MutualTLSChannelError) as excinfo: + RankServiceClient._get_api_endpoint( + None, mock_client_cert_source, mock_universe, "auto" + ) + assert ( + str(excinfo.value) + == "mTLS is not supported in any universe other than googleapis.com." + ) + + +def test__get_universe_domain(): + client_universe_domain = "foo.com" + universe_domain_env = "bar.com" + + assert ( + RankServiceClient._get_universe_domain( + client_universe_domain, universe_domain_env + ) + == client_universe_domain + ) + assert ( + RankServiceClient._get_universe_domain(None, universe_domain_env) + == universe_domain_env + ) + assert ( + RankServiceClient._get_universe_domain(None, None) + == RankServiceClient._DEFAULT_UNIVERSE + ) + + with pytest.raises(ValueError) as excinfo: + RankServiceClient._get_universe_domain("", None) + assert str(excinfo.value) == "Universe Domain cannot be an empty string." + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc"), + (RankServiceClient, transports.RankServiceRestTransport, "rest"), + ], +) +def test__validate_universe_domain(client_class, transport_class, transport_name): + client = client_class( + transport=transport_class(credentials=ga_credentials.AnonymousCredentials()) + ) + assert client._validate_universe_domain() == True + + # Test the case when universe is already validated. + assert client._validate_universe_domain() == True + + if transport_name == "grpc": + # Test the case where credentials are provided by the + # `local_channel_credentials`. The default universes in both match. + channel = grpc.secure_channel( + "http://localhost/", grpc.local_channel_credentials() + ) + client = client_class(transport=transport_class(channel=channel)) + assert client._validate_universe_domain() == True + + # Test the case where credentials do not exist: e.g. a transport is provided + # with no credentials. Validation should still succeed because there is no + # mismatch with non-existent credentials. + channel = grpc.secure_channel( + "http://localhost/", grpc.local_channel_credentials() + ) + transport = transport_class(channel=channel) + transport._credentials = None + client = client_class(transport=transport) + assert client._validate_universe_domain() == True + + # TODO: This is needed to cater for older versions of google-auth + # Make this test unconditional once the minimum supported version of + # google-auth becomes 2.23.0 or higher. + google_auth_major, google_auth_minor = [ + int(part) for part in google.auth.__version__.split(".")[0:2] + ] + if google_auth_major > 2 or (google_auth_major == 2 and google_auth_minor >= 23): + credentials = ga_credentials.AnonymousCredentials() + credentials._universe_domain = "foo.com" + # Test the case when there is a universe mismatch from the credentials. + client = client_class(transport=transport_class(credentials=credentials)) + with pytest.raises(ValueError) as excinfo: + client._validate_universe_domain() + assert ( + str(excinfo.value) + == "The configured universe domain (googleapis.com) does not match the universe domain found in the credentials (foo.com). If you haven't configured the universe domain explicitly, `googleapis.com` is the default." + ) + + # Test the case when there is a universe mismatch from the client. + # + # TODO: Make this test unconditional once the minimum supported version of + # google-api-core becomes 2.15.0 or higher. + api_core_major, api_core_minor = [ + int(part) for part in api_core_version.__version__.split(".")[0:2] + ] + if api_core_major > 2 or (api_core_major == 2 and api_core_minor >= 15): + client = client_class( + client_options={"universe_domain": "bar.com"}, + transport=transport_class( + credentials=ga_credentials.AnonymousCredentials(), + ), + ) + with pytest.raises(ValueError) as excinfo: + client._validate_universe_domain() + assert ( + str(excinfo.value) + == "The configured universe domain (bar.com) does not match the universe domain found in the credentials (googleapis.com). If you haven't configured the universe domain explicitly, `googleapis.com` is the default." + ) + + # Test that ValueError is raised if universe_domain is provided via client options and credentials is None + with pytest.raises(ValueError): + client._compare_universes("foo.bar", None) + + +@pytest.mark.parametrize( + "client_class,transport_name", + [ + (RankServiceClient, "grpc"), + (RankServiceAsyncClient, "grpc_asyncio"), + (RankServiceClient, "rest"), + ], +) +def test_rank_service_client_from_service_account_info(client_class, transport_name): + creds = ga_credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_info" + ) as factory: + factory.return_value = creds + info = {"valid": True} + client = client_class.from_service_account_info(info, transport=transport_name) + assert client.transport._credentials == creds + assert isinstance(client, client_class) + + assert client.transport._host == ( + "discoveryengine.googleapis.com:443" + if transport_name in ["grpc", "grpc_asyncio"] + else "https://discoveryengine.googleapis.com" + ) + + +@pytest.mark.parametrize( + "transport_class,transport_name", + [ + (transports.RankServiceGrpcTransport, "grpc"), + (transports.RankServiceGrpcAsyncIOTransport, "grpc_asyncio"), + (transports.RankServiceRestTransport, "rest"), + ], +) +def test_rank_service_client_service_account_always_use_jwt( + transport_class, transport_name +): + with mock.patch.object( + service_account.Credentials, "with_always_use_jwt_access", create=True + ) as use_jwt: + creds = service_account.Credentials(None, None, None) + transport = transport_class(credentials=creds, always_use_jwt_access=True) + use_jwt.assert_called_once_with(True) + + with mock.patch.object( + service_account.Credentials, "with_always_use_jwt_access", create=True + ) as use_jwt: + creds = service_account.Credentials(None, None, None) + transport = transport_class(credentials=creds, always_use_jwt_access=False) + use_jwt.assert_not_called() + + +@pytest.mark.parametrize( + "client_class,transport_name", + [ + (RankServiceClient, "grpc"), + (RankServiceAsyncClient, "grpc_asyncio"), + (RankServiceClient, "rest"), + ], +) +def test_rank_service_client_from_service_account_file(client_class, transport_name): + creds = ga_credentials.AnonymousCredentials() + with mock.patch.object( + service_account.Credentials, "from_service_account_file" + ) as factory: + factory.return_value = creds + client = client_class.from_service_account_file( + "dummy/file/path.json", transport=transport_name + ) + assert client.transport._credentials == creds + assert isinstance(client, client_class) + + client = client_class.from_service_account_json( + "dummy/file/path.json", transport=transport_name + ) + assert client.transport._credentials == creds + assert isinstance(client, client_class) + + assert client.transport._host == ( + "discoveryengine.googleapis.com:443" + if transport_name in ["grpc", "grpc_asyncio"] + else "https://discoveryengine.googleapis.com" + ) + + +def test_rank_service_client_get_transport_class(): + transport = RankServiceClient.get_transport_class() + available_transports = [ + transports.RankServiceGrpcTransport, + transports.RankServiceRestTransport, + ] + assert transport in available_transports + + transport = RankServiceClient.get_transport_class("grpc") + assert transport == transports.RankServiceGrpcTransport + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc"), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + ), + (RankServiceClient, transports.RankServiceRestTransport, "rest"), + ], +) +@mock.patch.object( + RankServiceClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceClient), +) +@mock.patch.object( + RankServiceAsyncClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceAsyncClient), +) +def test_rank_service_client_client_options( + client_class, transport_class, transport_name +): + # Check that if channel is provided we won't create a new one. + with mock.patch.object(RankServiceClient, "get_transport_class") as gtc: + transport = transport_class(credentials=ga_credentials.AnonymousCredentials()) + client = client_class(transport=transport) + gtc.assert_not_called() + + # Check that if channel is provided via str we will create a new one. + with mock.patch.object(RankServiceClient, "get_transport_class") as gtc: + client = client_class(transport=transport_name) + gtc.assert_called() + + # Check the case api_endpoint is provided. + options = client_options.ClientOptions(api_endpoint="squid.clam.whelk") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(transport=transport_name, client_options=options) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "never". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT is + # "always". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client.DEFAULT_MTLS_ENDPOINT, + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT has + # unsupported value. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "Unsupported"}): + with pytest.raises(MutualTLSChannelError) as excinfo: + client = client_class(transport=transport_name) + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`" + ) + + # Check the case GOOGLE_API_USE_CLIENT_CERTIFICATE has unsupported value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "Unsupported"} + ): + with pytest.raises(ValueError) as excinfo: + client = client_class(transport=transport_name) + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`" + ) + + # Check the case quota_project_id is provided + options = client_options.ClientOptions(quota_project_id="octopus") + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id="octopus", + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + # Check the case api_endpoint is provided + options = client_options.ClientOptions( + api_audience="https://language.googleapis.com" + ) + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience="https://language.googleapis.com", + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name,use_client_cert_env", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc", "true"), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + "true", + ), + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc", "false"), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + "false", + ), + (RankServiceClient, transports.RankServiceRestTransport, "rest", "true"), + (RankServiceClient, transports.RankServiceRestTransport, "rest", "false"), + ], +) +@mock.patch.object( + RankServiceClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceClient), +) +@mock.patch.object( + RankServiceAsyncClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceAsyncClient), +) +@mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "auto"}) +def test_rank_service_client_mtls_env_auto( + client_class, transport_class, transport_name, use_client_cert_env +): + # This tests the endpoint autoswitch behavior. Endpoint is autoswitched to the default + # mtls endpoint, if GOOGLE_API_USE_CLIENT_CERTIFICATE is "true" and client cert exists. + + # Check the case client_cert_source is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + options = client_options.ClientOptions( + client_cert_source=client_cert_source_callback + ) + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + + if use_client_cert_env == "false": + expected_client_cert_source = None + expected_host = client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ) + else: + expected_client_cert_source = client_cert_source_callback + expected_host = client.DEFAULT_MTLS_ENDPOINT + + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + client_cert_source_for_mtls=expected_client_cert_source, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # Check the case ADC client cert is provided. Whether client cert is used depends on + # GOOGLE_API_USE_CLIENT_CERTIFICATE value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=True, + ): + with mock.patch( + "google.auth.transport.mtls.default_client_cert_source", + return_value=client_cert_source_callback, + ): + if use_client_cert_env == "false": + expected_host = client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ) + expected_client_cert_source = None + else: + expected_host = client.DEFAULT_MTLS_ENDPOINT + expected_client_cert_source = client_cert_source_callback + + patched.return_value = None + client = client_class(transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=expected_host, + scopes=None, + client_cert_source_for_mtls=expected_client_cert_source, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # Check the case client_cert_source and ADC client cert are not provided. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": use_client_cert_env} + ): + with mock.patch.object(transport_class, "__init__") as patched: + with mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=False, + ): + patched.return_value = None + client = client_class(transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + +@pytest.mark.parametrize("client_class", [RankServiceClient, RankServiceAsyncClient]) +@mock.patch.object( + RankServiceClient, "DEFAULT_ENDPOINT", modify_default_endpoint(RankServiceClient) +) +@mock.patch.object( + RankServiceAsyncClient, + "DEFAULT_ENDPOINT", + modify_default_endpoint(RankServiceAsyncClient), +) +def test_rank_service_client_get_mtls_endpoint_and_cert_source(client_class): + mock_client_cert_source = mock.Mock() + + # Test the case GOOGLE_API_USE_CLIENT_CERTIFICATE is "true". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}): + mock_api_endpoint = "foo" + options = client_options.ClientOptions( + client_cert_source=mock_client_cert_source, api_endpoint=mock_api_endpoint + ) + api_endpoint, cert_source = client_class.get_mtls_endpoint_and_cert_source( + options + ) + assert api_endpoint == mock_api_endpoint + assert cert_source == mock_client_cert_source + + # Test the case GOOGLE_API_USE_CLIENT_CERTIFICATE is "false". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "false"}): + mock_client_cert_source = mock.Mock() + mock_api_endpoint = "foo" + options = client_options.ClientOptions( + client_cert_source=mock_client_cert_source, api_endpoint=mock_api_endpoint + ) + api_endpoint, cert_source = client_class.get_mtls_endpoint_and_cert_source( + options + ) + assert api_endpoint == mock_api_endpoint + assert cert_source is None + + # Test the case GOOGLE_API_USE_MTLS_ENDPOINT is "never". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + api_endpoint, cert_source = client_class.get_mtls_endpoint_and_cert_source() + assert api_endpoint == client_class.DEFAULT_ENDPOINT + assert cert_source is None + + # Test the case GOOGLE_API_USE_MTLS_ENDPOINT is "always". + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + api_endpoint, cert_source = client_class.get_mtls_endpoint_and_cert_source() + assert api_endpoint == client_class.DEFAULT_MTLS_ENDPOINT + assert cert_source is None + + # Test the case GOOGLE_API_USE_MTLS_ENDPOINT is "auto" and default cert doesn't exist. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}): + with mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=False, + ): + api_endpoint, cert_source = client_class.get_mtls_endpoint_and_cert_source() + assert api_endpoint == client_class.DEFAULT_ENDPOINT + assert cert_source is None + + # Test the case GOOGLE_API_USE_MTLS_ENDPOINT is "auto" and default cert exists. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}): + with mock.patch( + "google.auth.transport.mtls.has_default_client_cert_source", + return_value=True, + ): + with mock.patch( + "google.auth.transport.mtls.default_client_cert_source", + return_value=mock_client_cert_source, + ): + ( + api_endpoint, + cert_source, + ) = client_class.get_mtls_endpoint_and_cert_source() + assert api_endpoint == client_class.DEFAULT_MTLS_ENDPOINT + assert cert_source == mock_client_cert_source + + # Check the case api_endpoint is not provided and GOOGLE_API_USE_MTLS_ENDPOINT has + # unsupported value. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "Unsupported"}): + with pytest.raises(MutualTLSChannelError) as excinfo: + client_class.get_mtls_endpoint_and_cert_source() + + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_MTLS_ENDPOINT` must be `never`, `auto` or `always`" + ) + + # Check the case GOOGLE_API_USE_CLIENT_CERTIFICATE has unsupported value. + with mock.patch.dict( + os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "Unsupported"} + ): + with pytest.raises(ValueError) as excinfo: + client_class.get_mtls_endpoint_and_cert_source() + + assert ( + str(excinfo.value) + == "Environment variable `GOOGLE_API_USE_CLIENT_CERTIFICATE` must be either `true` or `false`" + ) + + +@pytest.mark.parametrize("client_class", [RankServiceClient, RankServiceAsyncClient]) +@mock.patch.object( + RankServiceClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceClient), +) +@mock.patch.object( + RankServiceAsyncClient, + "_DEFAULT_ENDPOINT_TEMPLATE", + modify_default_endpoint_template(RankServiceAsyncClient), +) +def test_rank_service_client_client_api_endpoint(client_class): + mock_client_cert_source = client_cert_source_callback + api_override = "foo.com" + default_universe = RankServiceClient._DEFAULT_UNIVERSE + default_endpoint = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=default_universe + ) + mock_universe = "bar.com" + mock_endpoint = RankServiceClient._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=mock_universe + ) + + # If ClientOptions.api_endpoint is set and GOOGLE_API_USE_CLIENT_CERTIFICATE="true", + # use ClientOptions.api_endpoint as the api endpoint regardless. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_CLIENT_CERTIFICATE": "true"}): + with mock.patch( + "google.auth.transport.requests.AuthorizedSession.configure_mtls_channel" + ): + options = client_options.ClientOptions( + client_cert_source=mock_client_cert_source, api_endpoint=api_override + ) + client = client_class( + client_options=options, + credentials=ga_credentials.AnonymousCredentials(), + ) + assert client.api_endpoint == api_override + + # If ClientOptions.api_endpoint is not set and GOOGLE_API_USE_MTLS_ENDPOINT="never", + # use the _DEFAULT_ENDPOINT_TEMPLATE populated with GDU as the api endpoint. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + client = client_class(credentials=ga_credentials.AnonymousCredentials()) + assert client.api_endpoint == default_endpoint + + # If ClientOptions.api_endpoint is not set and GOOGLE_API_USE_MTLS_ENDPOINT="always", + # use the DEFAULT_MTLS_ENDPOINT as the api endpoint. + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "always"}): + client = client_class(credentials=ga_credentials.AnonymousCredentials()) + assert client.api_endpoint == client_class.DEFAULT_MTLS_ENDPOINT + + # If ClientOptions.api_endpoint is not set, GOOGLE_API_USE_MTLS_ENDPOINT="auto" (default), + # GOOGLE_API_USE_CLIENT_CERTIFICATE="false" (default), default cert source doesn't exist, + # and ClientOptions.universe_domain="bar.com", + # use the _DEFAULT_ENDPOINT_TEMPLATE populated with universe domain as the api endpoint. + options = client_options.ClientOptions() + universe_exists = hasattr(options, "universe_domain") + if universe_exists: + options = client_options.ClientOptions(universe_domain=mock_universe) + client = client_class( + client_options=options, credentials=ga_credentials.AnonymousCredentials() + ) + else: + client = client_class( + client_options=options, credentials=ga_credentials.AnonymousCredentials() + ) + assert client.api_endpoint == ( + mock_endpoint if universe_exists else default_endpoint + ) + assert client.universe_domain == ( + mock_universe if universe_exists else default_universe + ) + + # If ClientOptions does not have a universe domain attribute and GOOGLE_API_USE_MTLS_ENDPOINT="never", + # use the _DEFAULT_ENDPOINT_TEMPLATE populated with GDU as the api endpoint. + options = client_options.ClientOptions() + if hasattr(options, "universe_domain"): + delattr(options, "universe_domain") + with mock.patch.dict(os.environ, {"GOOGLE_API_USE_MTLS_ENDPOINT": "never"}): + client = client_class( + client_options=options, credentials=ga_credentials.AnonymousCredentials() + ) + assert client.api_endpoint == default_endpoint + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc"), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + ), + (RankServiceClient, transports.RankServiceRestTransport, "rest"), + ], +) +def test_rank_service_client_client_options_scopes( + client_class, transport_class, transport_name +): + # Check the case scopes are provided. + options = client_options.ClientOptions( + scopes=["1", "2"], + ) + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=["1", "2"], + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name,grpc_helpers", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc", grpc_helpers), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + grpc_helpers_async, + ), + (RankServiceClient, transports.RankServiceRestTransport, "rest", None), + ], +) +def test_rank_service_client_client_options_credentials_file( + client_class, transport_class, transport_name, grpc_helpers +): + # Check the case credentials file is provided. + options = client_options.ClientOptions(credentials_file="credentials.json") + + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file="credentials.json", + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + +def test_rank_service_client_client_options_from_dict(): + with mock.patch( + "google.cloud.discoveryengine_v1alpha.services.rank_service.transports.RankServiceGrpcTransport.__init__" + ) as grpc_transport: + grpc_transport.return_value = None + client = RankServiceClient(client_options={"api_endpoint": "squid.clam.whelk"}) + grpc_transport.assert_called_once_with( + credentials=None, + credentials_file=None, + host="squid.clam.whelk", + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + +@pytest.mark.parametrize( + "client_class,transport_class,transport_name,grpc_helpers", + [ + (RankServiceClient, transports.RankServiceGrpcTransport, "grpc", grpc_helpers), + ( + RankServiceAsyncClient, + transports.RankServiceGrpcAsyncIOTransport, + "grpc_asyncio", + grpc_helpers_async, + ), + ], +) +def test_rank_service_client_create_channel_credentials_file( + client_class, transport_class, transport_name, grpc_helpers +): + # Check the case credentials file is provided. + options = client_options.ClientOptions(credentials_file="credentials.json") + + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options, transport=transport_name) + patched.assert_called_once_with( + credentials=None, + credentials_file="credentials.json", + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) + + # test that the credentials from file are saved and used as the credentials. + with mock.patch.object( + google.auth, "load_credentials_from_file", autospec=True + ) as load_creds, mock.patch.object( + google.auth, "default", autospec=True + ) as adc, mock.patch.object( + grpc_helpers, "create_channel" + ) as create_channel: + creds = ga_credentials.AnonymousCredentials() + file_creds = ga_credentials.AnonymousCredentials() + load_creds.return_value = (file_creds, None) + adc.return_value = (creds, None) + client = client_class(client_options=options, transport=transport_name) + create_channel.assert_called_with( + "discoveryengine.googleapis.com:443", + credentials=file_creds, + credentials_file=None, + quota_project_id=None, + default_scopes=("https://www.googleapis.com/auth/cloud-platform",), + scopes=None, + default_host="discoveryengine.googleapis.com", + ssl_credentials=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + +@pytest.mark.parametrize( + "request_type", + [ + rank_service.RankRequest, + dict, + ], +) +def test_rank(request_type, transport: str = "grpc"): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = rank_service.RankResponse() + response = client.rank(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + request = rank_service.RankRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, rank_service.RankResponse) + + +def test_rank_empty_call(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + client.rank() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == rank_service.RankRequest() + + +def test_rank_non_empty_request_with_auto_populated_field(): + # This test is a coverage failsafe to make sure that UUID4 fields are + # automatically populated, according to AIP-4235, with non-empty requests. + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc", + ) + + # Populate all string fields in the request which are not UUID4 + # since we want to check that UUID4 are populated automatically + # if they meet the requirements of AIP 4235. + request = rank_service.RankRequest( + ranking_config="ranking_config_value", + model="model_value", + query="query_value", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + client.rank(request=request) + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == rank_service.RankRequest( + ranking_config="ranking_config_value", + model="model_value", + query="query_value", + ) + + +@pytest.mark.asyncio +async def test_rank_empty_call_async(): + # This test is a coverage failsafe to make sure that totally empty calls, + # i.e. request == None and no flattened fields passed, work. + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc_asyncio", + ) + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + rank_service.RankResponse() + ) + response = await client.rank() + call.assert_called() + _, args, _ = call.mock_calls[0] + assert args[0] == rank_service.RankRequest() + + +@pytest.mark.asyncio +async def test_rank_async( + transport: str = "grpc_asyncio", request_type=rank_service.RankRequest +): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = request_type() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + rank_service.RankResponse() + ) + response = await client.rank(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + request = rank_service.RankRequest() + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, rank_service.RankResponse) + + +@pytest.mark.asyncio +async def test_rank_async_from_dict(): + await test_rank_async(request_type=dict) + + +def test_rank_field_headers(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = rank_service.RankRequest() + + request.ranking_config = "ranking_config_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + call.return_value = rank_service.RankResponse() + client.rank(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "ranking_config=ranking_config_value", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_rank_field_headers_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = rank_service.RankRequest() + + request.ranking_config = "ranking_config_value" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.rank), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + rank_service.RankResponse() + ) + await client.rank(request) + + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "ranking_config=ranking_config_value", + ) in kw["metadata"] + + +@pytest.mark.parametrize( + "request_type", + [ + rank_service.RankRequest, + dict, + ], +) +def test_rank_rest(request_type): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + + # send a request that will satisfy transcoding + request_init = { + "ranking_config": "projects/sample1/locations/sample2/rankingConfigs/sample3" + } + request = request_type(**request_init) + + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = rank_service.RankResponse() + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + # Convert return value to protobuf type + return_value = rank_service.RankResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + response = client.rank(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, rank_service.RankResponse) + + +def test_rank_rest_required_fields(request_type=rank_service.RankRequest): + transport_class = transports.RankServiceRestTransport + + request_init = {} + request_init["ranking_config"] = "" + request = request_type(**request_init) + pb_request = request_type.pb(request) + jsonified_request = json.loads( + json_format.MessageToJson(pb_request, use_integers_for_enums=False) + ) + + # verify fields with default values are dropped + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).rank._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with default values are now present + + jsonified_request["rankingConfig"] = "ranking_config_value" + + unset_fields = transport_class( + credentials=ga_credentials.AnonymousCredentials() + ).rank._get_unset_required_fields(jsonified_request) + jsonified_request.update(unset_fields) + + # verify required fields with non-default values are left alone + assert "rankingConfig" in jsonified_request + assert jsonified_request["rankingConfig"] == "ranking_config_value" + + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request = request_type(**request_init) + + # Designate an appropriate value for the returned response. + return_value = rank_service.RankResponse() + # Mock the http request call within the method and fake a response. + with mock.patch.object(Session, "request") as req: + # We need to mock transcode() because providing default values + # for required fields will fail the real version if the http_options + # expect actual values for those fields. + with mock.patch.object(path_template, "transcode") as transcode: + # A uri without fields and an empty body will force all the + # request fields to show up in the query_params. + pb_request = request_type.pb(request) + transcode_result = { + "uri": "v1/sample_method", + "method": "post", + "query_params": pb_request, + } + transcode_result["body"] = pb_request + transcode.return_value = transcode_result + + response_value = Response() + response_value.status_code = 200 + + # Convert return value to protobuf type + return_value = rank_service.RankResponse.pb(return_value) + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + response = client.rank(request) + + expected_params = [("$alt", "json;enum-encoding=int")] + actual_params = req.call_args.kwargs["params"] + assert expected_params == actual_params + + +def test_rank_rest_unset_required_fields(): + transport = transports.RankServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials + ) + + unset_fields = transport.rank._get_unset_required_fields({}) + assert set(unset_fields) == ( + set(()) + & set( + ( + "rankingConfig", + "records", + ) + ) + ) + + +@pytest.mark.parametrize("null_interceptor", [True, False]) +def test_rank_rest_interceptors(null_interceptor): + transport = transports.RankServiceRestTransport( + credentials=ga_credentials.AnonymousCredentials(), + interceptor=None + if null_interceptor + else transports.RankServiceRestInterceptor(), + ) + client = RankServiceClient(transport=transport) + with mock.patch.object( + type(client.transport._session), "request" + ) as req, mock.patch.object( + path_template, "transcode" + ) as transcode, mock.patch.object( + transports.RankServiceRestInterceptor, "post_rank" + ) as post, mock.patch.object( + transports.RankServiceRestInterceptor, "pre_rank" + ) as pre: + pre.assert_not_called() + post.assert_not_called() + pb_message = rank_service.RankRequest.pb(rank_service.RankRequest()) + transcode.return_value = { + "method": "post", + "uri": "my_uri", + "body": pb_message, + "query_params": pb_message, + } + + req.return_value = Response() + req.return_value.status_code = 200 + req.return_value.request = PreparedRequest() + req.return_value._content = rank_service.RankResponse.to_json( + rank_service.RankResponse() + ) + + request = rank_service.RankRequest() + metadata = [ + ("key", "val"), + ("cephalopod", "squid"), + ] + pre.return_value = request, metadata + post.return_value = rank_service.RankResponse() + + client.rank( + request, + metadata=[ + ("key", "val"), + ("cephalopod", "squid"), + ], + ) + + pre.assert_called_once() + post.assert_called_once() + + +def test_rank_rest_bad_request( + transport: str = "rest", request_type=rank_service.RankRequest +): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # send a request that will satisfy transcoding + request_init = { + "ranking_config": "projects/sample1/locations/sample2/rankingConfigs/sample3" + } + request = request_type(**request_init) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.rank(request) + + +def test_rank_rest_error(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport="rest" + ) + + +def test_credentials_transport_error(): + # It is an error to provide credentials and a transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # It is an error to provide a credentials file and a transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = RankServiceClient( + client_options={"credentials_file": "credentials.json"}, + transport=transport, + ) + + # It is an error to provide an api_key and a transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + options = client_options.ClientOptions() + options.api_key = "api_key" + with pytest.raises(ValueError): + client = RankServiceClient( + client_options=options, + transport=transport, + ) + + # It is an error to provide an api_key and a credential. + options = client_options.ClientOptions() + options.api_key = "api_key" + with pytest.raises(ValueError): + client = RankServiceClient( + client_options=options, credentials=ga_credentials.AnonymousCredentials() + ) + + # It is an error to provide scopes and a transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + with pytest.raises(ValueError): + client = RankServiceClient( + client_options={"scopes": ["1", "2"]}, + transport=transport, + ) + + +def test_transport_instance(): + # A client may be instantiated with a custom transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + client = RankServiceClient(transport=transport) + assert client.transport is transport + + +def test_transport_get_channel(): + # A client may be instantiated with a custom transport instance. + transport = transports.RankServiceGrpcTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + transport = transports.RankServiceGrpcAsyncIOTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + channel = transport.grpc_channel + assert channel + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.RankServiceGrpcTransport, + transports.RankServiceGrpcAsyncIOTransport, + transports.RankServiceRestTransport, + ], +) +def test_transport_adc(transport_class): + # Test default credentials are used if not provided. + with mock.patch.object(google.auth, "default") as adc: + adc.return_value = (ga_credentials.AnonymousCredentials(), None) + transport_class() + adc.assert_called_once() + + +@pytest.mark.parametrize( + "transport_name", + [ + "grpc", + "rest", + ], +) +def test_transport_kind(transport_name): + transport = RankServiceClient.get_transport_class(transport_name)( + credentials=ga_credentials.AnonymousCredentials(), + ) + assert transport.kind == transport_name + + +def test_transport_grpc_default(): + # A client should use the gRPC transport by default. + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + assert isinstance( + client.transport, + transports.RankServiceGrpcTransport, + ) + + +def test_rank_service_base_transport_error(): + # Passing both a credentials object and credentials_file should raise an error + with pytest.raises(core_exceptions.DuplicateCredentialArgs): + transport = transports.RankServiceTransport( + credentials=ga_credentials.AnonymousCredentials(), + credentials_file="credentials.json", + ) + + +def test_rank_service_base_transport(): + # Instantiate the base transport. + with mock.patch( + "google.cloud.discoveryengine_v1alpha.services.rank_service.transports.RankServiceTransport.__init__" + ) as Transport: + Transport.return_value = None + transport = transports.RankServiceTransport( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Every method on the transport should just blindly + # raise NotImplementedError. + methods = ( + "rank", + "get_operation", + "list_operations", + ) + for method in methods: + with pytest.raises(NotImplementedError): + getattr(transport, method)(request=object()) + + with pytest.raises(NotImplementedError): + transport.close() + + # Catch all for all remaining methods and properties + remainder = [ + "kind", + ] + for r in remainder: + with pytest.raises(NotImplementedError): + getattr(transport, r)() + + +def test_rank_service_base_transport_with_credentials_file(): + # Instantiate the base transport with a credentials file + with mock.patch.object( + google.auth, "load_credentials_from_file", autospec=True + ) as load_creds, mock.patch( + "google.cloud.discoveryengine_v1alpha.services.rank_service.transports.RankServiceTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + load_creds.return_value = (ga_credentials.AnonymousCredentials(), None) + transport = transports.RankServiceTransport( + credentials_file="credentials.json", + quota_project_id="octopus", + ) + load_creds.assert_called_once_with( + "credentials.json", + scopes=None, + default_scopes=("https://www.googleapis.com/auth/cloud-platform",), + quota_project_id="octopus", + ) + + +def test_rank_service_base_transport_with_adc(): + # Test the default credentials are used if credentials and credentials_file are None. + with mock.patch.object(google.auth, "default", autospec=True) as adc, mock.patch( + "google.cloud.discoveryengine_v1alpha.services.rank_service.transports.RankServiceTransport._prep_wrapped_messages" + ) as Transport: + Transport.return_value = None + adc.return_value = (ga_credentials.AnonymousCredentials(), None) + transport = transports.RankServiceTransport() + adc.assert_called_once() + + +def test_rank_service_auth_adc(): + # If no credentials are provided, we should use ADC credentials. + with mock.patch.object(google.auth, "default", autospec=True) as adc: + adc.return_value = (ga_credentials.AnonymousCredentials(), None) + RankServiceClient() + adc.assert_called_once_with( + scopes=None, + default_scopes=("https://www.googleapis.com/auth/cloud-platform",), + quota_project_id=None, + ) + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.RankServiceGrpcTransport, + transports.RankServiceGrpcAsyncIOTransport, + ], +) +def test_rank_service_transport_auth_adc(transport_class): + # If credentials and host are not provided, the transport class should use + # ADC credentials. + with mock.patch.object(google.auth, "default", autospec=True) as adc: + adc.return_value = (ga_credentials.AnonymousCredentials(), None) + transport_class(quota_project_id="octopus", scopes=["1", "2"]) + adc.assert_called_once_with( + scopes=["1", "2"], + default_scopes=("https://www.googleapis.com/auth/cloud-platform",), + quota_project_id="octopus", + ) + + +@pytest.mark.parametrize( + "transport_class", + [ + transports.RankServiceGrpcTransport, + transports.RankServiceGrpcAsyncIOTransport, + transports.RankServiceRestTransport, + ], +) +def test_rank_service_transport_auth_gdch_credentials(transport_class): + host = "https://language.com" + api_audience_tests = [None, "https://language2.com"] + api_audience_expect = [host, "https://language2.com"] + for t, e in zip(api_audience_tests, api_audience_expect): + with mock.patch.object(google.auth, "default", autospec=True) as adc: + gdch_mock = mock.MagicMock() + type(gdch_mock).with_gdch_audience = mock.PropertyMock( + return_value=gdch_mock + ) + adc.return_value = (gdch_mock, None) + transport_class(host=host, api_audience=t) + gdch_mock.with_gdch_audience.assert_called_once_with(e) + + +@pytest.mark.parametrize( + "transport_class,grpc_helpers", + [ + (transports.RankServiceGrpcTransport, grpc_helpers), + (transports.RankServiceGrpcAsyncIOTransport, grpc_helpers_async), + ], +) +def test_rank_service_transport_create_channel(transport_class, grpc_helpers): + # If credentials and host are not provided, the transport class should use + # ADC credentials. + with mock.patch.object( + google.auth, "default", autospec=True + ) as adc, mock.patch.object( + grpc_helpers, "create_channel", autospec=True + ) as create_channel: + creds = ga_credentials.AnonymousCredentials() + adc.return_value = (creds, None) + transport_class(quota_project_id="octopus", scopes=["1", "2"]) + + create_channel.assert_called_with( + "discoveryengine.googleapis.com:443", + credentials=creds, + credentials_file=None, + quota_project_id="octopus", + default_scopes=("https://www.googleapis.com/auth/cloud-platform",), + scopes=["1", "2"], + default_host="discoveryengine.googleapis.com", + ssl_credentials=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + +@pytest.mark.parametrize( + "transport_class", + [transports.RankServiceGrpcTransport, transports.RankServiceGrpcAsyncIOTransport], +) +def test_rank_service_grpc_transport_client_cert_source_for_mtls(transport_class): + cred = ga_credentials.AnonymousCredentials() + + # Check ssl_channel_credentials is used if provided. + with mock.patch.object(transport_class, "create_channel") as mock_create_channel: + mock_ssl_channel_creds = mock.Mock() + transport_class( + host="squid.clam.whelk", + credentials=cred, + ssl_channel_credentials=mock_ssl_channel_creds, + ) + mock_create_channel.assert_called_once_with( + "squid.clam.whelk:443", + credentials=cred, + credentials_file=None, + scopes=None, + ssl_credentials=mock_ssl_channel_creds, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + + # Check if ssl_channel_credentials is not provided, then client_cert_source_for_mtls + # is used. + with mock.patch.object(transport_class, "create_channel", return_value=mock.Mock()): + with mock.patch("grpc.ssl_channel_credentials") as mock_ssl_cred: + transport_class( + credentials=cred, + client_cert_source_for_mtls=client_cert_source_callback, + ) + expected_cert, expected_key = client_cert_source_callback() + mock_ssl_cred.assert_called_once_with( + certificate_chain=expected_cert, private_key=expected_key + ) + + +def test_rank_service_http_transport_client_cert_source_for_mtls(): + cred = ga_credentials.AnonymousCredentials() + with mock.patch( + "google.auth.transport.requests.AuthorizedSession.configure_mtls_channel" + ) as mock_configure_mtls_channel: + transports.RankServiceRestTransport( + credentials=cred, client_cert_source_for_mtls=client_cert_source_callback + ) + mock_configure_mtls_channel.assert_called_once_with(client_cert_source_callback) + + +@pytest.mark.parametrize( + "transport_name", + [ + "grpc", + "grpc_asyncio", + "rest", + ], +) +def test_rank_service_host_no_port(transport_name): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="discoveryengine.googleapis.com" + ), + transport=transport_name, + ) + assert client.transport._host == ( + "discoveryengine.googleapis.com:443" + if transport_name in ["grpc", "grpc_asyncio"] + else "https://discoveryengine.googleapis.com" + ) + + +@pytest.mark.parametrize( + "transport_name", + [ + "grpc", + "grpc_asyncio", + "rest", + ], +) +def test_rank_service_host_with_port(transport_name): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + client_options=client_options.ClientOptions( + api_endpoint="discoveryengine.googleapis.com:8000" + ), + transport=transport_name, + ) + assert client.transport._host == ( + "discoveryengine.googleapis.com:8000" + if transport_name in ["grpc", "grpc_asyncio"] + else "https://discoveryengine.googleapis.com:8000" + ) + + +@pytest.mark.parametrize( + "transport_name", + [ + "rest", + ], +) +def test_rank_service_client_transport_session_collision(transport_name): + creds1 = ga_credentials.AnonymousCredentials() + creds2 = ga_credentials.AnonymousCredentials() + client1 = RankServiceClient( + credentials=creds1, + transport=transport_name, + ) + client2 = RankServiceClient( + credentials=creds2, + transport=transport_name, + ) + session1 = client1.transport.rank._session + session2 = client2.transport.rank._session + assert session1 != session2 + + +def test_rank_service_grpc_transport_channel(): + channel = grpc.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.RankServiceGrpcTransport( + host="squid.clam.whelk", + channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +def test_rank_service_grpc_asyncio_transport_channel(): + channel = aio.secure_channel("http://localhost/", grpc.local_channel_credentials()) + + # Check that channel is used if provided. + transport = transports.RankServiceGrpcAsyncIOTransport( + host="squid.clam.whelk", + channel=channel, + ) + assert transport.grpc_channel == channel + assert transport._host == "squid.clam.whelk:443" + assert transport._ssl_channel_credentials == None + + +# Remove this test when deprecated arguments (api_mtls_endpoint, client_cert_source) are +# removed from grpc/grpc_asyncio transport constructor. +@pytest.mark.parametrize( + "transport_class", + [transports.RankServiceGrpcTransport, transports.RankServiceGrpcAsyncIOTransport], +) +def test_rank_service_transport_channel_mtls_with_client_cert_source(transport_class): + with mock.patch( + "grpc.ssl_channel_credentials", autospec=True + ) as grpc_ssl_channel_cred: + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_ssl_cred = mock.Mock() + grpc_ssl_channel_cred.return_value = mock_ssl_cred + + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + + cred = ga_credentials.AnonymousCredentials() + with pytest.warns(DeprecationWarning): + with mock.patch.object(google.auth, "default") as adc: + adc.return_value = (cred, None) + transport = transport_class( + host="squid.clam.whelk", + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=client_cert_source_callback, + ) + adc.assert_called_once() + + grpc_ssl_channel_cred.assert_called_once_with( + certificate_chain=b"cert bytes", private_key=b"key bytes" + ) + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=cred, + credentials_file=None, + scopes=None, + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + assert transport._ssl_channel_credentials == mock_ssl_cred + + +# Remove this test when deprecated arguments (api_mtls_endpoint, client_cert_source) are +# removed from grpc/grpc_asyncio transport constructor. +@pytest.mark.parametrize( + "transport_class", + [transports.RankServiceGrpcTransport, transports.RankServiceGrpcAsyncIOTransport], +) +def test_rank_service_transport_channel_mtls_with_adc(transport_class): + mock_ssl_cred = mock.Mock() + with mock.patch.multiple( + "google.auth.transport.grpc.SslCredentials", + __init__=mock.Mock(return_value=None), + ssl_credentials=mock.PropertyMock(return_value=mock_ssl_cred), + ): + with mock.patch.object( + transport_class, "create_channel" + ) as grpc_create_channel: + mock_grpc_channel = mock.Mock() + grpc_create_channel.return_value = mock_grpc_channel + mock_cred = mock.Mock() + + with pytest.warns(DeprecationWarning): + transport = transport_class( + host="squid.clam.whelk", + credentials=mock_cred, + api_mtls_endpoint="mtls.squid.clam.whelk", + client_cert_source=None, + ) + + grpc_create_channel.assert_called_once_with( + "mtls.squid.clam.whelk:443", + credentials=mock_cred, + credentials_file=None, + scopes=None, + ssl_credentials=mock_ssl_cred, + quota_project_id=None, + options=[ + ("grpc.max_send_message_length", -1), + ("grpc.max_receive_message_length", -1), + ], + ) + assert transport.grpc_channel == mock_grpc_channel + + +def test_ranking_config_path(): + project = "squid" + location = "clam" + ranking_config = "whelk" + expected = "projects/{project}/locations/{location}/rankingConfigs/{ranking_config}".format( + project=project, + location=location, + ranking_config=ranking_config, + ) + actual = RankServiceClient.ranking_config_path(project, location, ranking_config) + assert expected == actual + + +def test_parse_ranking_config_path(): + expected = { + "project": "octopus", + "location": "oyster", + "ranking_config": "nudibranch", + } + path = RankServiceClient.ranking_config_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_ranking_config_path(path) + assert expected == actual + + +def test_common_billing_account_path(): + billing_account = "cuttlefish" + expected = "billingAccounts/{billing_account}".format( + billing_account=billing_account, + ) + actual = RankServiceClient.common_billing_account_path(billing_account) + assert expected == actual + + +def test_parse_common_billing_account_path(): + expected = { + "billing_account": "mussel", + } + path = RankServiceClient.common_billing_account_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_common_billing_account_path(path) + assert expected == actual + + +def test_common_folder_path(): + folder = "winkle" + expected = "folders/{folder}".format( + folder=folder, + ) + actual = RankServiceClient.common_folder_path(folder) + assert expected == actual + + +def test_parse_common_folder_path(): + expected = { + "folder": "nautilus", + } + path = RankServiceClient.common_folder_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_common_folder_path(path) + assert expected == actual + + +def test_common_organization_path(): + organization = "scallop" + expected = "organizations/{organization}".format( + organization=organization, + ) + actual = RankServiceClient.common_organization_path(organization) + assert expected == actual + + +def test_parse_common_organization_path(): + expected = { + "organization": "abalone", + } + path = RankServiceClient.common_organization_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_common_organization_path(path) + assert expected == actual + + +def test_common_project_path(): + project = "squid" + expected = "projects/{project}".format( + project=project, + ) + actual = RankServiceClient.common_project_path(project) + assert expected == actual + + +def test_parse_common_project_path(): + expected = { + "project": "clam", + } + path = RankServiceClient.common_project_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_common_project_path(path) + assert expected == actual + + +def test_common_location_path(): + project = "whelk" + location = "octopus" + expected = "projects/{project}/locations/{location}".format( + project=project, + location=location, + ) + actual = RankServiceClient.common_location_path(project, location) + assert expected == actual + + +def test_parse_common_location_path(): + expected = { + "project": "oyster", + "location": "nudibranch", + } + path = RankServiceClient.common_location_path(**expected) + + # Check that the path construction is reversible. + actual = RankServiceClient.parse_common_location_path(path) + assert expected == actual + + +def test_client_with_default_client_info(): + client_info = gapic_v1.client_info.ClientInfo() + + with mock.patch.object( + transports.RankServiceTransport, "_prep_wrapped_messages" + ) as prep: + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + client_info=client_info, + ) + prep.assert_called_once_with(client_info) + + with mock.patch.object( + transports.RankServiceTransport, "_prep_wrapped_messages" + ) as prep: + transport_class = RankServiceClient.get_transport_class() + transport = transport_class( + credentials=ga_credentials.AnonymousCredentials(), + client_info=client_info, + ) + prep.assert_called_once_with(client_info) + + +@pytest.mark.asyncio +async def test_transport_close_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="grpc_asyncio", + ) + with mock.patch.object( + type(getattr(client.transport, "grpc_channel")), "close" + ) as close: + async with client: + close.assert_not_called() + close.assert_called_once() + + +def test_get_operation_rest_bad_request( + transport: str = "rest", request_type=operations_pb2.GetOperationRequest +): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + request = request_type() + request = json_format.ParseDict( + { + "name": "projects/sample1/locations/sample2/collections/sample3/dataConnector/operations/sample4" + }, + request, + ) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.get_operation(request) + + +@pytest.mark.parametrize( + "request_type", + [ + operations_pb2.GetOperationRequest, + dict, + ], +) +def test_get_operation_rest(request_type): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request_init = { + "name": "projects/sample1/locations/sample2/collections/sample3/dataConnector/operations/sample4" + } + request = request_type(**request_init) + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.Operation() + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + response = client.get_operation(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.Operation) + + +def test_list_operations_rest_bad_request( + transport: str = "rest", request_type=operations_pb2.ListOperationsRequest +): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + request = request_type() + request = json_format.ParseDict( + { + "name": "projects/sample1/locations/sample2/collections/sample3/dataConnector" + }, + request, + ) + + # Mock the http request call within the method and fake a BadRequest error. + with mock.patch.object(Session, "request") as req, pytest.raises( + core_exceptions.BadRequest + ): + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 400 + response_value.request = Request() + req.return_value = response_value + client.list_operations(request) + + +@pytest.mark.parametrize( + "request_type", + [ + operations_pb2.ListOperationsRequest, + dict, + ], +) +def test_list_operations_rest(request_type): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport="rest", + ) + request_init = { + "name": "projects/sample1/locations/sample2/collections/sample3/dataConnector" + } + request = request_type(**request_init) + # Mock the http request call within the method and fake a response. + with mock.patch.object(type(client.transport._session), "request") as req: + # Designate an appropriate value for the returned response. + return_value = operations_pb2.ListOperationsResponse() + + # Wrap the value into a proper Response obj + response_value = Response() + response_value.status_code = 200 + json_return_value = json_format.MessageToJson(return_value) + + response_value._content = json_return_value.encode("UTF-8") + req.return_value = response_value + + response = client.list_operations(request) + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.ListOperationsResponse) + + +def test_get_operation(transport: str = "grpc"): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = operations_pb2.GetOperationRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation() + response = client.get_operation(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.Operation) + + +@pytest.mark.asyncio +async def test_get_operation_async(transport: str = "grpc_asyncio"): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = operations_pb2.GetOperationRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation() + ) + response = await client.get_operation(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.Operation) + + +def test_get_operation_field_headers(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = operations_pb2.GetOperationRequest() + request.name = "locations" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + call.return_value = operations_pb2.Operation() + + client.get_operation(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=locations", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_get_operation_field_headers_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = operations_pb2.GetOperationRequest() + request.name = "locations" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation() + ) + await client.get_operation(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=locations", + ) in kw["metadata"] + + +def test_get_operation_from_dict(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.Operation() + + response = client.get_operation( + request={ + "name": "locations", + } + ) + call.assert_called() + + +@pytest.mark.asyncio +async def test_get_operation_from_dict_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.get_operation), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.Operation() + ) + response = await client.get_operation( + request={ + "name": "locations", + } + ) + call.assert_called() + + +def test_list_operations(transport: str = "grpc"): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = operations_pb2.ListOperationsRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.ListOperationsResponse() + response = client.list_operations(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.ListOperationsResponse) + + +@pytest.mark.asyncio +async def test_list_operations_async(transport: str = "grpc_asyncio"): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + transport=transport, + ) + + # Everything is optional in proto3 as far as the runtime is concerned, + # and we are mocking out the actual API, so just send an empty request. + request = operations_pb2.ListOperationsRequest() + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.ListOperationsResponse() + ) + response = await client.list_operations(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the response is the type that we expect. + assert isinstance(response, operations_pb2.ListOperationsResponse) + + +def test_list_operations_field_headers(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = operations_pb2.ListOperationsRequest() + request.name = "locations" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + call.return_value = operations_pb2.ListOperationsResponse() + + client.list_operations(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=locations", + ) in kw["metadata"] + + +@pytest.mark.asyncio +async def test_list_operations_field_headers_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + + # Any value that is part of the HTTP/1.1 URI should be sent as + # a field header. Set these to a non-empty value. + request = operations_pb2.ListOperationsRequest() + request.name = "locations" + + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.ListOperationsResponse() + ) + await client.list_operations(request) + # Establish that the underlying gRPC stub method was called. + assert len(call.mock_calls) == 1 + _, args, _ = call.mock_calls[0] + assert args[0] == request + + # Establish that the field header was sent. + _, _, kw = call.mock_calls[0] + assert ( + "x-goog-request-params", + "name=locations", + ) in kw["metadata"] + + +def test_list_operations_from_dict(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = operations_pb2.ListOperationsResponse() + + response = client.list_operations( + request={ + "name": "locations", + } + ) + call.assert_called() + + +@pytest.mark.asyncio +async def test_list_operations_from_dict_async(): + client = RankServiceAsyncClient( + credentials=ga_credentials.AnonymousCredentials(), + ) + # Mock the actual call within the gRPC stub, and fake the request. + with mock.patch.object(type(client.transport.list_operations), "__call__") as call: + # Designate an appropriate return value for the call. + call.return_value = grpc_helpers_async.FakeUnaryUnaryCall( + operations_pb2.ListOperationsResponse() + ) + response = await client.list_operations( + request={ + "name": "locations", + } + ) + call.assert_called() + + +def test_transport_close(): + transports = { + "rest": "_session", + "grpc": "_grpc_channel", + } + + for transport, close_name in transports.items(): + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport=transport + ) + with mock.patch.object( + type(getattr(client.transport, close_name)), "close" + ) as close: + with client: + close.assert_not_called() + close.assert_called_once() + + +def test_client_ctx(): + transports = [ + "rest", + "grpc", + ] + for transport in transports: + client = RankServiceClient( + credentials=ga_credentials.AnonymousCredentials(), transport=transport + ) + # Test client calls underlying transport. + with mock.patch.object(type(client.transport), "close") as close: + close.assert_not_called() + with client: + pass + close.assert_called() + + +@pytest.mark.parametrize( + "client_class,transport_class", + [ + (RankServiceClient, transports.RankServiceGrpcTransport), + (RankServiceAsyncClient, transports.RankServiceGrpcAsyncIOTransport), + ], +) +def test_api_key_credentials(client_class, transport_class): + with mock.patch.object( + google.auth._default, "get_api_key_credentials", create=True + ) as get_api_key_credentials: + mock_cred = mock.Mock() + get_api_key_credentials.return_value = mock_cred + options = client_options.ClientOptions() + options.api_key = "api_key" + with mock.patch.object(transport_class, "__init__") as patched: + patched.return_value = None + client = client_class(client_options=options) + patched.assert_called_once_with( + credentials=mock_cred, + credentials_file=None, + host=client._DEFAULT_ENDPOINT_TEMPLATE.format( + UNIVERSE_DOMAIN=client._DEFAULT_UNIVERSE + ), + scopes=None, + client_cert_source_for_mtls=None, + quota_project_id=None, + client_info=transports.base.DEFAULT_CLIENT_INFO, + always_use_jwt_access=True, + api_audience=None, + ) diff --git a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_serving_config_service.py b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_serving_config_service.py index bfc99d86ebed..70ed08b42530 100644 --- a/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_serving_config_service.py +++ b/packages/google-cloud-discoveryengine/tests/unit/gapic/discoveryengine_v1alpha/test_serving_config_service.py @@ -2497,6 +2497,7 @@ def test_update_serving_config_rest(request_type): "model_prompt_spec": {"preamble": "preamble_value"}, "language_code": "language_code_value", "model_spec": {"version": "version_value"}, + "use_semantic_chunks": True, }, "extractive_content_spec": { "max_extractive_answer_count": 2907, @@ -2506,6 +2507,7 @@ def test_update_serving_config_rest(request_type): "num_next_segments": 1843, }, "search_result_mode": 1, + "chunk_spec": {"num_previous_chunks": 2071, "num_next_chunks": 1625}, } }, "name": "projects/sample1/locations/sample2/dataStores/sample3/servingConfigs/sample4",