Merge pull request #1510 from weaviate/google_aistudio

Add google ai studio
weaviate · Jan 13, 2025 · 056814c · 056814c
2 parents c02ed69 + e7af01f
commit 056814c
Show file tree

Hide file tree

Showing 3 changed files with 107 additions and 1 deletion.
diff --git a/test/collection/test_config.py b/test/collection/test_config.py
@@ -318,6 +318,15 @@ def test_basic_config():
             }
         },
     ),
+    (
+        Configure.Vectorizer.text2vec_google_aistudio(),
+        {
+            "text2vec-palm": {
+                "apiEndpoint": "generativelanguage.googleapis.com",
+                "vectorizeClassName": True,
+            }
+        },
+    ),
     (
         Configure.Vectorizer.text2vec_palm(
             project_id="project",
@@ -1509,6 +1518,26 @@ def test_vector_config_flat_pq() -> None:
             }
         },
     ),
+    (
+        [
+            Configure.NamedVectors.text2vec_google_aistudio(
+                name="test",
+                source_properties=["prop"],
+            )
+        ],
+        {
+            "test": {
+                "vectorizer": {
+                    "text2vec-palm": {
+                        "apiEndpoint": "generativelanguage.googleapis.com",
+                        "properties": ["prop"],
+                        "vectorizeClassName": True,
+                    }
+                },
+                "vectorIndexType": "hnsw",
+            }
+        },
+    ),
     (
         [Configure.NamedVectors.text2vec_transformers(name="test", source_properties=["prop"])],
         {

diff --git a/weaviate/collections/classes/config_named_vectors.py b/weaviate/collections/classes/config_named_vectors.py
@@ -1087,6 +1087,53 @@ def text2vec_google(
             vector_index_config=vector_index_config,
         )
 
+    @staticmethod
+    def text2vec_google_aistudio(
+        name: str,
+        *,
+        source_properties: Optional[List[str]] = None,
+        vector_index_config: Optional[_VectorIndexConfigCreate] = None,
+        vectorize_collection_name: bool = True,
+        model_id: Optional[str] = None,
+        title_property: Optional[str] = None,
+    ) -> _NamedVectorConfigCreate:
+        """Create a named vector using the `text2vec_palm` model.
+
+        See the [documentation]https://weaviate.io/developers/weaviate/model-providers/google/embeddings)
+        for detailed usage.
+
+        Arguments:
+            `name`
+                The name of the named vector.
+            `source_properties`
+                Which properties should be included when vectorizing. By default all text properties are included.
+            `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`.
+            `model_id`
+                The model ID 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.
+
+        Raises:
+            `pydantic.ValidationError` if `api_endpoint` is not a valid URL.
+        """
+        return _NamedVectorConfigCreate(
+            name=name,
+            source_properties=source_properties,
+            vectorizer=_Text2VecGoogleConfig(
+                projectId=None,
+                apiEndpoint="generativelanguage.googleapis.com",
+                modelId=model_id,
+                vectorizeClassName=vectorize_collection_name,
+                titleProperty=title_property,
+            ),
+            vector_index_config=vector_index_config,
+        )
+
     @staticmethod
     def text2vec_transformers(
         name: str,

diff --git a/weaviate/collections/classes/config_vectorizers.py b/weaviate/collections/classes/config_vectorizers.py
@@ -306,7 +306,7 @@ class _Text2VecGoogleConfig(_VectorizerConfigCreate):
     vectorizer: Union[Vectorizers, _EnumLikeStr] = Field(
         default=Vectorizers.TEXT2VEC_PALM, frozen=True, exclude=True
     )
-    projectId: str
+    projectId: Optional[str]
     apiEndpoint: Optional[str]
     modelId: Optional[str]
     vectorizeClassName: bool
@@ -1084,6 +1084,36 @@ def text2vec_palm(
             titleProperty=title_property,
         )
 
+    @staticmethod
+    def text2vec_google_aistudio(
+        model_id: Optional[str] = None,
+        title_property: Optional[str] = None,
+        vectorize_collection_name: bool = True,
+    ) -> _VectorizerConfigCreate:
+        """Create a `_Text2VecGoogleConfig` object for use when vectorizing using the `text2vec-google` model.
+
+        See the [documentation](https://weaviate.io/developers/weaviate/model-providers/google/embeddings)
+        for detailed usage.
+
+        Arguments:
+            `model_id`
+                The model ID 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.
+            `vectorize_collection_name`
+                Whether to vectorize the collection name. Defaults to `True`.
+
+        Raises:
+            `pydantic.ValidationError` if `api_endpoint` is not a valid URL.
+        """
+        return _Text2VecGoogleConfig(
+            projectId=None,
+            apiEndpoint="generativelanguage.googleapis.com",
+            modelId=model_id,
+            vectorizeClassName=vectorize_collection_name,
+            titleProperty=title_property,
+        )
+
     @staticmethod
     def text2vec_google(
         project_id: str,