Merge pull request #1900 from weaviate/bugfix/model-integration-parameters

feat: add taskType parameter to google vectorizer

Author: databyjp

weaviate/collections/classes/config_named_vectors.py

@@ -981,6 +981,7 @@ def text2vec_palm(
                 modelId=model_id,
                 vectorizeClassName=vectorize_collection_name,
                 titleProperty=title_property,
+                taskType=None,
             ),
             vector_index_config=vector_index_config,
         )
@@ -1026,6 +1027,7 @@ def text2vec_google(
                 modelId=model_id,
                 vectorizeClassName=vectorize_collection_name,
                 titleProperty=title_property,
+                taskType=None,
             ),
             vector_index_config=vector_index_config,
         )
@@ -1067,6 +1069,7 @@ def text2vec_google_aistudio(
                 modelId=model_id,
                 vectorizeClassName=vectorize_collection_name,
                 titleProperty=title_property,
+                taskType=None,
             ),
             vector_index_config=vector_index_config,
         )

weaviate/collections/classes/config_vectorizers.py

@@ -351,6 +351,7 @@ class _Text2VecGoogleConfig(_VectorizerConfigCreate):
     modelId: Optional[str]
     vectorizeClassName: bool
     titleProperty: Optional[str]
+    taskType: Optional[str]
 
 
 class _Text2VecTransformersConfig(_VectorizerConfigCreate):
@@ -1149,6 +1150,7 @@ def text2vec_palm(
             modelId=model_id,
             vectorizeClassName=vectorize_collection_name,
             titleProperty=title_property,
+            taskType=None,
         )
 
     @staticmethod
@@ -1177,6 +1179,7 @@ def text2vec_google_aistudio(
             modelId=model_id,
             vectorizeClassName=vectorize_collection_name,
             titleProperty=title_property,
+            taskType=None,
         )
 
     @staticmethod
@@ -1210,6 +1213,7 @@ def text2vec_google(
             modelId=model_id,
             vectorizeClassName=vectorize_collection_name,
             titleProperty=title_property,
+            taskType=None,
         )
 
     @staticmethod

weaviate/collections/classes/config_vectors.py

@@ -175,6 +175,7 @@ def multi(
         return _IndexWrappers.single(vector_index_config, quantizer)
 
 
+# TODO: Consider refactoring to move the default values to the pydantic classes themselves (e.g. _VectorConfigCreate, _Text2VecCohereConfig, etc.)
 class _MultiVectors:
     @staticmethod
     def self_provided(
@@ -1188,6 +1189,7 @@ def text2vec_google(
         model: Optional[str] = None,
         project_id: str,
         title_property: Optional[str] = None,
+        task_type: Optional[str] = None,
         source_properties: Optional[List[str]] = None,
         vector_index_config: Optional[_VectorIndexConfigCreate] = None,
         vectorize_collection_name: bool = True,
@@ -1205,6 +1207,7 @@ def text2vec_google(
             model: The model to use. Defaults to `None`, which uses the server-defined default.
             project_id: The project ID to use, REQUIRED.
             title_property: The Weaviate property name for the `gecko-002` or `gecko-003` model to use as the title.
+            task_type: The task type to use (e.g. `RETRIEVAL_QUERY`, `RETRIEVAL_DOCUMENT`). Defaults to `None`, which uses the server-defined default.
             source_properties: Which properties should be included when vectorizing. By default all text properties are included.
             vector_index_config: The configuration for Weaviate's vector index. Use `wvc.config.Configure.VectorIndex` to create a vector index configuration. None by default.
             vectorize_collection_name: Whether to vectorize the collection name. Defaults to `True`.
@@ -1222,6 +1225,7 @@ def text2vec_google(
                 modelId=model,
                 vectorizeClassName=vectorize_collection_name,
                 titleProperty=title_property,
+                taskType=task_type,
             ),
             vector_index_config=_IndexWrappers.single(vector_index_config, quantizer),
         )
@@ -1234,6 +1238,7 @@ def text2vec_google_aistudio(
         dimensions: Optional[int] = None,
         model: Optional[str] = None,
         title_property: Optional[str] = None,
+        task_type: Optional[str] = None,
         source_properties: Optional[List[str]] = None,
         vector_index_config: Optional[_VectorIndexConfigCreate] = None,
         vectorize_collection_name: bool = True,
@@ -1249,6 +1254,7 @@ def text2vec_google_aistudio(
             dimenions: The dimensionality of the vectors. Defaults to `None`, which uses the server-defined default.
             model: The model to use. Defaults to `None`, which uses the server-defined default.
             title_property: The Weaviate property name for the `gecko-002` or `gecko-003` model to use as the title.
+            task_type: The task type to use (e.g. `RETRIEVAL_QUERY`, `RETRIEVAL_DOCUMENT`). Defaults to `None`, which uses the server-defined default.
             source_properties: Which properties should be included when vectorizing. By default all text properties are included.
             vector_index_config: The configuration for Weaviate's vector index. Use `wvc.config.Configure.VectorIndex` to create a vector index configuration. None by default
             vectorize_collection_name: Whether to vectorize the collection name. Defaults to `True`.
@@ -1266,6 +1272,7 @@ def text2vec_google_aistudio(
                 modelId=model,
                 vectorizeClassName=vectorize_collection_name,
                 titleProperty=title_property,
+                taskType=task_type,
             ),
             vector_index_config=_IndexWrappers.single(vector_index_config, quantizer),
         )