Skip to content

Commit de2d6b6

Browse files
feat(community): Add support for Blended Search in VertexAISearchRetriever
1 parent 9c3a0ee commit de2d6b6

File tree

1 file changed

+18
-7
lines changed

1 file changed

+18
-7
lines changed

libs/community/langchain_google_community/vertex_ai_search.py

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)