@@ -70,11 +70,12 @@ class _BaseVertexAISearchRetriever(Serializable):
7070 """The default custom credentials (google.auth.credentials.Credentials) to use
7171 when making API calls. If not provided, credentials will be ascertained from
7272 the environment."""
73- engine_data_type : int = Field (default = 0 , ge = 0 , le = 2 )
73+ engine_data_type : int = Field (default = 0 , ge = 0 , le = 3 )
7474 """ Defines the Vertex AI Search data type
7575 0 - Unstructured data
7676 1 - Structured data
7777 2 - Website data
78+ 3 - Blended data
7879 """
7980 _beta : bool = PrivateAttr (default = False )
8081 """Whether to use beta version of Vertex AI Search."""
@@ -493,9 +494,9 @@ def _get_content_spec_kwargs(self) -> Optional[Dict[str, Any]]:
493494 )
494495 else :
495496 raise NotImplementedError (
496- "Only data store type 0 (Unstructured), 1 (Structured),"
497- "or 2 (Website) are supported currently."
498- + f" Got { self .engine_data_type } "
497+ "Only data store types 0 (Unstructured), 1 (Structured), "
498+ "2 (Website), or 3 (Blended) are supported currently. "
499+ + f"Got { self .engine_data_type } "
499500 )
500501 return content_search_spec
501502
@@ -574,11 +575,21 @@ def _get_relevant_documents(
574575 documents = self ._convert_website_search_response (
575576 response .results , chunk_type
576577 )
578+ elif self .engine_data_type == 3 :
579+ # Blended search might return varied results; treat similarly to unstructured/website
580+ chunk_type = (
581+ "extractive_answers"
582+ if self .get_extractive_answers
583+ else "extractive_segments" # Or perhaps 'snippets', depending on API behavior
584+ )
585+ documents = self ._convert_unstructured_search_response (
586+ response .results , chunk_type
587+ )
577588 else :
578589 raise NotImplementedError (
579- "Only data store type 0 (Unstructured), 1 (Structured),"
580- "or 2 (Website) are supported currently."
581- + f" Got { self .engine_data_type } "
590+ "Only data store types 0 (Unstructured), 1 (Structured), "
591+ "2 (Website), or 3 (Blended) are supported currently. "
592+ + f"Got { self .engine_data_type } "
582593 )
583594
584595 return documents
0 commit comments