[pull] main from GoogleCloudPlatform:main

generative_ai/rag/create_corpus_example.py

@@ -15,41 +15,44 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
 
 def create_corpus(
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus]
 
-    from vertexai import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
-    # Configure backend_config
-    backend_config = rag.RagVectorDbConfig(
-        rag_embedding_model_config=rag.RagEmbeddingModelConfig(
-            vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
-                publisher_model="publishers/google/models/text-embedding-005"
+    # Configure project-level config
+    backend_config = types.RagVectorDbConfig(
+        rag_embedding_model_config=types.RagEmbeddingModelConfig(
+            vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
+                endpoint="publishers/google/models/text-embedding-005"
             )
         )
     )
 
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        backend_config=backend_config,
+    # Create a corpus
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            rag_vector_db_config=backend_config,
+        )
     )
     print(corpus)
     # Example response:

generative_ai/rag/create_corpus_feature_store_example.py

@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -24,34 +24,39 @@ def create_corpus_feature_store(
     feature_view_name: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_feature_store]
 
-    from vertexai.preview import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
     # feature_view_name = "projects/{PROJECT_ID}/locations/{LOCATION}/featureOnlineStores/{FEATURE_ONLINE_STORE_ID}/featureViews/{FEATURE_VIEW_ID}"
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
     # Configure embedding model (Optional)
-    embedding_model_config = rag.EmbeddingModelConfig(
-        publisher_model="publishers/google/models/text-embedding-004"
+    backend_config = types.RagVectorDbConfig(
+        rag_embedding_model_config=types.RagEmbeddingModelConfig(
+            vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
+                endpoint="publishers/google/models/text-embedding-005"
+            ),
+        ),
+        vertex_feature_store=types.RagVectorDbConfigVertexFeatureStore(
+            feature_view_resource_name=feature_view_name
+        )
     )
 
-    # Configure Vector DB
-    vector_db = rag.VertexFeatureStore(resource_name=feature_view_name)
-
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        embedding_model_config=embedding_model_config,
-        vector_db=vector_db,
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            rag_vector_db_config=backend_config,
+        )
     )
     print(corpus)
     # Example response:

generative_ai/rag/create_corpus_pinecone_example.py

@@ -15,52 +15,53 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
 
+
 def create_corpus_pinecone(
     pinecone_index_name: str,
     pinecone_api_key_secret_manager_version: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_pinecone]
 
-    from vertexai import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
     # pinecone_index_name = "pinecone-index-name"
-    # pinecone_api_key_secret_manager_version = "projects/{PROJECT_ID}/secrets/{SECRET_NAME}/versions/latest"
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
     # Configure embedding model (Optional)
-    embedding_model_config = rag.RagEmbeddingModelConfig(
-        vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
-            publisher_model="publishers/google/models/text-embedding-005"
+    embedding_model_config = types.RagEmbeddingModelConfig(
+        vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
+            endpoint="publishers/google/models/text-embedding-005"
         )
     )
 
     # Configure Vector DB
-    vector_db = rag.Pinecone(
-        index_name=pinecone_index_name,
-        api_key=pinecone_api_key_secret_manager_version,
+    vector_db = types.RagVectorDbConfig(
+        pinecone=types.RagVectorDbConfigPinecone(
+         index_name=pinecone_index_name,
+        ),
+        rag_embedding_model_config=embedding_model_config,
     )
 
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        backend_config=rag.RagVectorDbConfig(
-            rag_embedding_model_config=embedding_model_config,
-            vector_db=vector_db,
-        ),
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            rag_vector_db_config=vector_db,
+        )
     )
     print(corpus)
     # Example response:

generative_ai/rag/create_corpus_vector_search_example.py

@@ -15,21 +15,20 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
-
 def create_corpus_vector_search(
     vector_search_index_name: str,
     vector_search_index_endpoint_name: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_vector_search]
 
-    from vertexai import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
@@ -38,28 +37,31 @@ def create_corpus_vector_search(
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
     # Configure embedding model (Optional)
-    embedding_model_config = rag.RagEmbeddingModelConfig(
-        vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
-            publisher_model="publishers/google/models/text-embedding-005"
+    embedding_model_config = types.RagEmbeddingModelConfig(
+        vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
+            endpoint="publishers/google/models/text-embedding-005"
         )
     )
 
     # Configure Vector DB
-    vector_db = rag.VertexVectorSearch(
-        index=vector_search_index_name, index_endpoint=vector_search_index_endpoint_name
+    vector_db = types.RagVectorDbConfigVertexVectorSearch(
+        index=vector_search_index_name,
+        index_endpoint=vector_search_index_endpoint_name
     )
 
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        backend_config=rag.RagVectorDbConfig(
-            rag_embedding_model_config=embedding_model_config,
-            vector_db=vector_db,
-        ),
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            rag_vector_db_config=types.RagVectorDbConfig(
+                rag_embedding_model_config=embedding_model_config,
+                vertex_vector_search=vector_db,
+            ),
+        )
     )
     print(corpus)
     # Example response:

generative_ai/rag/create_corpus_vertex_ai_search_example.py

@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from vertexai import rag
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -24,30 +24,32 @@ def create_corpus_vertex_ai_search(
     vertex_ai_search_engine_name: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> rag.RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_vertex_ai_search]
 
-    from vertexai import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
     # vertex_ai_search_engine_name = "projects/{PROJECT_ID}/locations/{LOCATION}/collections/default_collection/engines/{ENGINE_ID}"
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
     # Configure Search
-    vertex_ai_search_config = rag.VertexAiSearchConfig(
+    vertex_ai_search_config = types.VertexAiSearchConfig(
         serving_config=f"{vertex_ai_search_engine_name}/servingConfigs/default_search",
     )
 
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        vertex_ai_search_config=vertex_ai_search_config,
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            vertex_ai_search_config=vertex_ai_search_config,
+        ),
     )
     print(corpus)
     # Example response:

generative_ai/rag/create_corpus_weaviate_example.py

@@ -15,7 +15,7 @@
 
 from typing import Optional
 
-from vertexai.preview.rag import RagCorpus
+from agentplatform import types
 
 PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
 
@@ -26,11 +26,11 @@ def create_corpus_weaviate(
     weaviate_api_key_secret_manager_version: str,
     display_name: Optional[str] = None,
     description: Optional[str] = None,
-) -> RagCorpus:
+) -> types.RagCorpus:
     # [START generativeaionvertexai_rag_create_corpus_weaviate]
 
-    from vertexai.preview import rag
-    import vertexai
+    import agentplatform
+    from agentplatform import types
 
     # TODO(developer): Update and un-comment below lines
     # PROJECT_ID = "your-project-id"
@@ -40,26 +40,31 @@ def create_corpus_weaviate(
     # display_name = "test_corpus"
     # description = "Corpus Description"
 
-    # Initialize Vertex AI API once per session
-    vertexai.init(project=PROJECT_ID, location="us-central1")
+    # Initialize Agent Platform client once per session
+    client = agentplatform.Client(project=PROJECT_ID, location="us-central1")
 
     # Configure embedding model (Optional)
-    embedding_model_config = rag.EmbeddingModelConfig(
-        publisher_model="publishers/google/models/text-embedding-004"
+    embedding_model_config = types.RagEmbeddingModelConfig(
+        vertex_prediction_endpoint=types.RagEmbeddingModelConfigVertexPredictionEndpoint(
+            endpoint="publishers/google/models/text-embedding-004"
+        )
     )
 
     # Configure Vector DB
-    vector_db = rag.Weaviate(
-        weaviate_http_endpoint=weaviate_http_endpoint,
+    vector_db = types.RagVectorDbConfigWeaviate(
+        http_endpoint=weaviate_http_endpoint,
         collection_name=weaviate_collection_name,
-        api_key=weaviate_api_key_secret_manager_version,
     )
 
-    corpus = rag.create_corpus(
-        display_name=display_name,
-        description=description,
-        embedding_model_config=embedding_model_config,
-        vector_db=vector_db,
+    corpus = client.rag.create_corpus(
+        rag_corpus=types.RagCorpus(
+            display_name=display_name,
+            description=description,
+            rag_embedding_model_config=embedding_model_config,
+            rag_vector_db_config=types.RagVectorDbConfig(
+                weaviate=vector_db
+            ),
+        )
     )
     print(corpus)
     # Example response: