From 9cb2c9f03ade379c5ccef9a94f1564549c339258 Mon Sep 17 00:00:00 2001 From: Patrick Marlow Date: Fri, 27 Sep 2024 14:19:15 -0500 Subject: [PATCH] chore: update to use latest GA library for discoveryengine --- src/dfcx_scrapi/core/data_stores.py | 28 +++++++++++-------- src/dfcx_scrapi/core/engines.py | 28 +++++++++++-------- src/dfcx_scrapi/core/sites.py | 43 +++++++++++++++++------------ 3 files changed, 60 insertions(+), 39 deletions(-) diff --git a/src/dfcx_scrapi/core/data_stores.py b/src/dfcx_scrapi/core/data_stores.py index 4fd944f1..b1561d8b 100644 --- a/src/dfcx_scrapi/core/data_stores.py +++ b/src/dfcx_scrapi/core/data_stores.py @@ -17,8 +17,14 @@ import logging from typing import Dict, List from google.longrunning.operations_pb2 import Operation -from google.cloud import discoveryengine_v1alpha -from google.cloud.discoveryengine_v1alpha.types import DataStore +from google.cloud.discoveryengine import ( + DataStore, + DataStoreServiceClient, + ListDataStoresRequest, + GetDataStoreRequest, + CreateDataStoreRequest, + DeleteDataStoreRequest + ) from dfcx_scrapi.core import scrapi_base @@ -85,10 +91,10 @@ def list_data_stores( client_options = self._client_options_discovery_engine( f"projects/{self.project_id}/locations/{location}" ) - client = discoveryengine_v1alpha.DataStoreServiceClient( + client = DataStoreServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.ListDataStoresRequest(parent=parent) + request = ListDataStoresRequest(parent=parent) page_result = client.list_data_stores(request=request) datastores = [] @@ -100,10 +106,10 @@ def list_data_stores( def get_data_store(self, data_store_id: str) -> DataStore: """Get a single Data Store by specified ID.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.DataStoreServiceClient( + client = DataStoreServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.GetDataStoreRequest( + request = GetDataStoreRequest( name=data_store_id ) @@ -133,16 +139,16 @@ def create_datastore( client_options = self._client_options_discovery_engine( f"projects/{self.project_id}/locations/{location}" ) - client = discoveryengine_v1alpha.DataStoreServiceClient( + client = DataStoreServiceClient( credentials=self.creds, client_options=client_options ) - data_store = discoveryengine_v1alpha.DataStore() + data_store = DataStore() data_store.display_name = display_name data_store.industry_vertical = 1 data_store.solution_types = [self._get_solution_type(solution_type)] data_store.content_config = self._get_content_config(datastore_type) - request = discoveryengine_v1alpha.CreateDataStoreRequest( + request = CreateDataStoreRequest( parent=parent, data_store=data_store, data_store_id=data_store.display_name, @@ -156,11 +162,11 @@ def create_datastore( def delete_datastore(self, data_store_id: str): """Delete the specified Data Store by ID.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.DataStoreServiceClient( + client = DataStoreServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.DeleteDataStoreRequest( + request = DeleteDataStoreRequest( name=data_store_id ) operation = client.delete_data_store(request=request) diff --git a/src/dfcx_scrapi/core/engines.py b/src/dfcx_scrapi/core/engines.py index 3b5fa0a1..6439a480 100644 --- a/src/dfcx_scrapi/core/engines.py +++ b/src/dfcx_scrapi/core/engines.py @@ -17,8 +17,14 @@ import logging from typing import Dict, List from google.longrunning.operations_pb2 import Operation -from google.cloud import discoveryengine_v1alpha -from google.cloud.discoveryengine_v1alpha.types import Engine +from google.cloud.discoveryengine import ( + Engine, + EngineServiceClient, + ListEnginesRequest, + GetEngineRequest, + CreateEngineRequest, + DeleteEngineRequest + ) from dfcx_scrapi.core import scrapi_base @@ -151,8 +157,8 @@ def list_engines( self, location: str = "global") -> List[Engine]: """List all Engines for a given project and location.""" parent = self._build_data_store_parent(location) - client = discoveryengine_v1alpha.EngineServiceClient() - request = discoveryengine_v1alpha.ListEnginesRequest(parent=parent) + client = EngineServiceClient() + request = ListEnginesRequest(parent=parent) page_result = client.list_engines(request=request) engines = [] @@ -163,8 +169,8 @@ def list_engines( def get_engine(self, engine_id: str) -> Engine: """Get a single Engine by specified ID.""" - client = discoveryengine_v1alpha.EngineServiceClient() - request = discoveryengine_v1alpha.GetEngineRequest(name=engine_id) + client = EngineServiceClient() + request = GetEngineRequest(name=engine_id) response = client.get_engine(request=request) return response @@ -178,17 +184,17 @@ def create_engine(self, engine: Engine, Args: location, the GCP region to create the Engine in - engine, a proto object of type discoveryengine_v1alpha.types.Engine + engine, a proto object of type discoveryengine.Engine Note that at this time only "Chat" engines are supported solution_type, "chat" is the only value supported at this time. """ parent = self._build_data_store_parent(location) client_options = self._client_options_discovery_engine(parent) - client = discoveryengine_v1alpha.EngineServiceClient( + client = EngineServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.CreateEngineRequest( + request = CreateEngineRequest( parent=parent, engine=engine, engine_id=engine.display_name @@ -204,11 +210,11 @@ def create_engine(self, engine: Engine, def delete_engine(self, engine_id: str) -> Operation: """Deletes the specified Engine.""" client_options = self._client_options_discovery_engine(engine_id) - client = discoveryengine_v1alpha.EngineServiceClient( + client = EngineServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.DeleteEngineRequest( + request = DeleteEngineRequest( name=engine_id ) operation = client.delete_engine(request=request) diff --git a/src/dfcx_scrapi/core/sites.py b/src/dfcx_scrapi/core/sites.py index 50836565..bcc60cb4 100644 --- a/src/dfcx_scrapi/core/sites.py +++ b/src/dfcx_scrapi/core/sites.py @@ -17,8 +17,17 @@ import logging from typing import Dict, List from google.longrunning.operations_pb2 import Operation -from google.cloud import discoveryengine_v1alpha -from google.cloud.discoveryengine_v1alpha.types import TargetSite +from google.cloud.discoveryengine import ( + TargetSite, + SiteSearchEngineServiceClient, + ListTargetSitesRequest, + GetTargetSiteRequest, + CreateTargetSiteRequest, + DeleteTargetSiteRequest, + EnableAdvancedSiteSearchRequest, + DisableAdvancedSiteSearchRequest, + RecrawlUrisRequest +) from dfcx_scrapi.core import scrapi_base @@ -85,12 +94,12 @@ def get_sites_map( def list_sites(self, data_store_id: str) -> List[TargetSite]: """List all URL patterns for a given Data Store ID.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) parent = self.__build_site_search_parent(data_store_id) - request = discoveryengine_v1alpha.ListTargetSitesRequest( + request = ListTargetSitesRequest( parent=parent, page_size=1000 ) @@ -106,10 +115,10 @@ def list_sites(self, data_store_id: str) -> List[TargetSite]: def get_site(self, site_id: str) -> TargetSite: """Get a single Site by specified ID.""" client_options = self._client_options_discovery_engine(site_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.GetTargetSiteRequest(name=site_id) + request = GetTargetSiteRequest(name=site_id) response = client.get_target_site(request=request) return response @@ -134,17 +143,17 @@ def create_site(self, data_store_id: str, uri_pattern: str, search engine. """ client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) - target_site = discoveryengine_v1alpha.TargetSite() + target_site = TargetSite() target_site.provided_uri_pattern = uri_pattern target_site.type_ = self.__get_site_type(include_site) target_site.exact_match = exact_match parent = self.__build_site_search_parent(data_store_id) - request = discoveryengine_v1alpha.CreateTargetSiteRequest( + request = CreateTargetSiteRequest( parent=parent, target_site=target_site ) @@ -156,11 +165,11 @@ def create_site(self, data_store_id: str, uri_pattern: str, def delete_site(self, site_id: str) -> Operation: """Deletes a TargetSite in a Data Store by the specified ID.""" client_options = self._client_options_discovery_engine(site_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) - request = discoveryengine_v1alpha.DeleteTargetSiteRequest(name=site_id) + request = DeleteTargetSiteRequest(name=site_id) operation = client.delete_target_site(request=request) return operation @@ -168,12 +177,12 @@ def delete_site(self, site_id: str) -> Operation: def enable_advanced_site_search(self, data_store_id: str) -> Operation: """Enables Advanced Site Search for the provided Data Store ID.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) parent = self.__build_site_search_parent(data_store_id) - request = discoveryengine_v1alpha.EnableAdvancedSiteSearchRequest( + request = EnableAdvancedSiteSearchRequest( site_search_engine=parent ) @@ -184,12 +193,12 @@ def enable_advanced_site_search(self, data_store_id: str) -> Operation: def disable_advanced_site_search(self, data_store_id: str) -> Operation: """Disable Advanced Site Search for the provided Data Store ID.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) parent = self.__build_site_search_parent(data_store_id) - request = discoveryengine_v1alpha.DisableAdvancedSiteSearchRequest( + request = DisableAdvancedSiteSearchRequest( site_search_engine=parent ) @@ -200,12 +209,12 @@ def disable_advanced_site_search(self, data_store_id: str) -> Operation: def recrawl_uris(self, data_store_id: str, uris: List[str]) -> Operation: """Recrawl the specified set of URIs for the Data Store.""" client_options = self._client_options_discovery_engine(data_store_id) - client = discoveryengine_v1alpha.SiteSearchEngineServiceClient( + client = SiteSearchEngineServiceClient( credentials=self.creds, client_options=client_options ) parent = self.__build_site_search_parent(data_store_id) - request = discoveryengine_v1alpha.RecrawlUrisRequest( + request = RecrawlUrisRequest( site_search_engine=parent, uris=uris )