chore: update to use latest GA library for discoveryengine

GoogleCloudPlatform · Sep 27, 2024 · 9cb2c9f · 9cb2c9f
1 parent 0926656
commit 9cb2c9f
Show file tree

Hide file tree

Showing 3 changed files with 60 additions and 39 deletions.
diff --git 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
@@ -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
@@ -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,24 +165,24 @@ 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
 
     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
             )