diff --git a/generative_ai/rag/create_corpus_example.py b/generative_ai/rag/create_corpus_example.py index 90b1aa60401..f5bec091527 100644 --- a/generative_ai/rag/create_corpus_example.py +++ b/generative_ai/rag/create_corpus_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") @@ -23,33 +23,36 @@ 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: diff --git a/generative_ai/rag/create_corpus_feature_store_example.py b/generative_ai/rag/create_corpus_feature_store_example.py index 8674887c1fe..39f61898361 100644 --- a/generative_ai/rag/create_corpus_feature_store_example.py +++ b/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,11 +24,11 @@ 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" @@ -36,22 +36,27 @@ def create_corpus_feature_store( # 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: diff --git a/generative_ai/rag/create_corpus_pinecone_example.py b/generative_ai/rag/create_corpus_pinecone_example.py index ebca30385e8..74a2c6081e5 100644 --- a/generative_ai/rag/create_corpus_pinecone_example.py +++ b/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: diff --git a/generative_ai/rag/create_corpus_vector_search_example.py b/generative_ai/rag/create_corpus_vector_search_example.py index 5db30008046..8f3905adbd1 100644 --- a/generative_ai/rag/create_corpus_vector_search_example.py +++ b/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: diff --git a/generative_ai/rag/create_corpus_vertex_ai_search_example.py b/generative_ai/rag/create_corpus_vertex_ai_search_example.py index 6d3fca5ab9c..01231e73a16 100644 --- a/generative_ai/rag/create_corpus_vertex_ai_search_example.py +++ b/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,11 +24,11 @@ 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" @@ -36,18 +36,20 @@ def create_corpus_vertex_ai_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 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: diff --git a/generative_ai/rag/create_corpus_weaviate_example.py b/generative_ai/rag/create_corpus_weaviate_example.py index 9823b8332f8..a12452af2e6 100644 --- a/generative_ai/rag/create_corpus_weaviate_example.py +++ b/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: diff --git a/generative_ai/rag/delete_corpus_example.py b/generative_ai/rag/delete_corpus_example.py index 4255110fe14..527126a92a7 100644 --- a/generative_ai/rag/delete_corpus_example.py +++ b/generative_ai/rag/delete_corpus_example.py @@ -20,17 +20,16 @@ def delete_corpus(corpus_name: str) -> None: # [START generativeaionvertexai_rag_delete_corpus] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # 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") - rag.delete_corpus(name=corpus_name) + client.rag.delete_corpus(name=corpus_name) print(f"Corpus {corpus_name} deleted.") # Example response: # Successfully deleted the RagCorpus. diff --git a/generative_ai/rag/delete_file_example.py b/generative_ai/rag/delete_file_example.py index e11afc71d96..477f93a97d0 100644 --- a/generative_ai/rag/delete_file_example.py +++ b/generative_ai/rag/delete_file_example.py @@ -20,17 +20,16 @@ def delete_file(file_name: str) -> None: # [START generativeaionvertexai_rag_delete_file] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}" - # 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") - rag.delete_file(name=file_name) + client.rag.delete_file(name=file_name) print(f"File {file_name} deleted.") # Example response: # Successfully deleted the RagFile. diff --git a/generative_ai/rag/generate_content_example.py b/generative_ai/rag/generate_content_example.py index a02b8bfb7f1..6bff7cb1910 100644 --- a/generative_ai/rag/generate_content_example.py +++ b/generative_ai/rag/generate_content_example.py @@ -14,49 +14,51 @@ import os -from vertexai.generative_models import GenerationResponse +from genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") def generate_content_with_rag( corpus_name: str, -) -> GenerationResponse: +) -> genai_types.GenerateContentResponse: # [START generativeaionvertexai_rag_generate_content] - from vertexai import rag - from vertexai.generative_models import GenerativeModel, Tool - import vertexai + from google import genai + from google.genai import types as genai_types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") - - rag_retrieval_tool = Tool.from_retrieval( - retrieval=rag.Retrieval( - source=rag.VertexRagStore( + rag_retrieval_tool = genai_types.Tool( + retrieval=genai_types.Retrieval( + vertex_rag_store=genai_types.VertexRagStore( rag_resources=[ - rag.RagResource( - rag_corpus=corpus_name, - # Optional: supply IDs from `rag.list_files()`. - # rag_file_ids=["rag-file-1", "rag-file-2", ...], + genai_types.VertexRagStoreRagResource( + rag_corpus=corpus_name ) ], - rag_retrieval_config=rag.RagRetrievalConfig( + rag_retrieval_config=genai_types.RagRetrievalConfig( top_k=10, - filter=rag.utils.resources.Filter(vector_distance_threshold=0.5), + filter=genai_types.RagRetrievalConfigFilter( + vector_distance_threshold=0.5 + ), ), ), ) ) - rag_model = GenerativeModel( - model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool] + # Create a GenAI SDK client to make a generate_content request + genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-central1") + + response = genai_client.models.generate_content( + model="gemini-2.5-pro", + contents="Why is the sky blue?", + config=genai_types.GenerateContentConfig( + tools=[rag_retrieval_tool] + ) ) - response = rag_model.generate_content("Why is the sky blue?") print(response.text) # Example response: # The sky appears blue due to a phenomenon called Rayleigh scattering. diff --git a/generative_ai/rag/get_corpus_example.py b/generative_ai/rag/get_corpus_example.py index 849995156d0..2e0a1ae868c 100644 --- a/generative_ai/rag/get_corpus_example.py +++ b/generative_ai/rag/get_corpus_example.py @@ -14,25 +14,24 @@ import os -from google.cloud.aiplatform_v1beta1 import RagCorpus +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def get_corpus(corpus_name: str) -> RagCorpus: +def get_corpus(corpus_name: str) -> types.RagCorpus: # [START generativeaionvertexai_rag_get_corpus] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # 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") - corpus = rag.get_corpus(name=corpus_name) + corpus = client.rag.get_corpus(name=corpus_name) print(corpus) # Example response: # RagCorpus(name='projects/[PROJECT_ID]/locations/us-central1/ragCorpora/1234567890', diff --git a/generative_ai/rag/get_file_example.py b/generative_ai/rag/get_file_example.py index 90c461ae4d9..af64b386f83 100644 --- a/generative_ai/rag/get_file_example.py +++ b/generative_ai/rag/get_file_example.py @@ -14,25 +14,24 @@ import os -from google.cloud.aiplatform_v1 import RagFile +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def get_file(file_name: str) -> RagFile: +def get_file(file_name: str) -> types.RagFile: # [START generativeaionvertexai_rag_get_file] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}" - # 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") - rag_file = rag.get_file(name=file_name) + rag_file = client.rag.get_file(name=file_name) print(rag_file) # Example response: # RagFile(name='projects/1234567890/locations/us-central1/ragCorpora/11111111111/ragFiles/22222222222', diff --git a/generative_ai/rag/import_files_async_example.py b/generative_ai/rag/import_files_async_example.py index 7485b951ff0..8ca344ce9ee 100644 --- a/generative_ai/rag/import_files_async_example.py +++ b/generative_ai/rag/import_files_async_example.py @@ -16,7 +16,7 @@ from typing import List -from google.cloud.aiplatform_v1 import ImportRagFilesResponse +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -24,38 +24,52 @@ async def import_files_async( corpus_name: str, paths: List[str], -) -> ImportRagFilesResponse: +) -> types.ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files_async] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types + + from google.genai import types as genai_types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" # Supports Google Cloud Storage and Google Drive Links - # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir"] - - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") - - response = await rag.import_files( - corpus_name=corpus_name, - paths=paths, - transformation_config=rag.TransformationConfig( - rag.ChunkingConfig(chunk_size=512, chunk_overlap=100) - ), - max_embedding_requests_per_min=900, # Optional + # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir/*"] + + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") + + response = await client.aio.rag.import_files( + name=corpus_name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[paths[1]]), + google_drive_source=types.GoogleDriveSource( + resource_ids=[ + types.GoogleDriveSourceResourceId( + resource_id=paths[0], + resource_type=types.ResourceType.RESOURCE_TYPE_FILE + ) + ] + ), # optional + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + chunk_overlap=100, + ) + ), # optional + max_embedding_requests_per_min=900, # optional + ) ) - result = await response.result() - print(f"Imported {result.imported_rag_files_count} files.") + print(f"Imported {response.imported_rag_files_count} files.") # Example response: # Imported 2 files. # [END generativeaionvertexai_rag_import_files_async] - return result + return response if __name__ == "__main__": diff --git a/generative_ai/rag/import_files_example.py b/generative_ai/rag/import_files_example.py index c21f68c28d2..d2d92192c8a 100644 --- a/generative_ai/rag/import_files_example.py +++ b/generative_ai/rag/import_files_example.py @@ -15,7 +15,7 @@ import os from typing import List -from google.cloud.aiplatform_v1 import ImportRagFilesResponse +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -23,29 +23,46 @@ def import_files( corpus_name: str, paths: List[str], -) -> ImportRagFilesResponse: +) -> types.ImportRagFilesResponse: # [START generativeaionvertexai_rag_import_files] - from vertexai import rag - import vertexai + import agentplatform + from agentplatform import types + + from google.genai import types as genai_types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # paths = ["https://drive.google.com/file/123", "gs://my_bucket/my_files_dir"] # Supports Google Cloud Storage and Google Drive Links - - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-central1") - - response = rag.import_files( - corpus_name=corpus_name, - paths=paths, - transformation_config=rag.TransformationConfig( - rag.ChunkingConfig(chunk_size=512, chunk_overlap=100) - ), - import_result_sink="gs://sample-existing-folder/sample_import_result_unique.ndjson", # Optional, this has to be an existing storage bucket folder, and file name has to be unique (non-existent). - max_embedding_requests_per_min=900, # Optional + + # Supports Google Cloud Storage and Google Drive Links + # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir/*"] + + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-central1") + + response = client.rag.import_files( + name=corpus_name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[paths[1]]), + google_drive_source=types.GoogleDriveSource( + resource_ids=[ + types.GoogleDriveSourceResourceId( + resource_id=paths[0], + resource_type=types.ResourceType.RESOURCE_TYPE_FILE + ) + ] + ), # optional + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + chunk_overlap=100, + ) + ), # optional + max_embedding_requests_per_min=900, # optional + ) ) + print(f"Imported {response.imported_rag_files_count} files.") # Example response: # Imported 2 files. diff --git a/generative_ai/rag/list_corpora_example.py b/generative_ai/rag/list_corpora_example.py index 138a47f4330..1e9ba774384 100644 --- a/generative_ai/rag/list_corpora_example.py +++ b/generative_ai/rag/list_corpora_example.py @@ -14,26 +14,23 @@ import os -from google.cloud.aiplatform_v1beta1.services.vertex_rag_data_service.pagers import ( - ListRagCorporaPager, -) +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") -def list_corpora() -> ListRagCorporaPager: +def list_corpora() -> types.ListRagCorporaResponse: # [START generativeaionvertexai_rag_list_corpora] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" - # 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") - corpora = rag.list_corpora() + corpora = client.rag.list_corpora() print(corpora) # Example response: # ListRagCorporaPager ListRagFilesPager: +def list_files(corpus_name: str) -> types.ListRagFilesResponse: # [START generativeaionvertexai_rag_list_files] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" - # 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") - files = rag.list_files(corpus_name=corpus_name) - for file in files: + files_response = client.rag.list_files(name=corpus_name) + for file in files_response.rag_files: print(file.display_name) print(file.name) # Example response: @@ -45,7 +42,7 @@ def list_files(corpus_name: str) -> ListRagFilesPager: # projects/1234567890/locations/us-central1/ragCorpora/111111111111/ragFiles/333333333333 # [END generativeaionvertexai_rag_list_files] - return files + return files_response if __name__ == "__main__": diff --git a/generative_ai/rag/quickstart_example.py b/generative_ai/rag/quickstart_example.py index 32649f64aeb..b57ae049e9e 100644 --- a/generative_ai/rag/quickstart_example.py +++ b/generative_ai/rag/quickstart_example.py @@ -14,88 +14,101 @@ import os -from typing import List, Tuple +from typing import Tuple -from vertexai import rag -from vertexai.generative_models import GenerationResponse +from agentplatform import types +from google.genai import types as genai_types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") +MODEL_ID = os.getenv("MODEL_ID") def quickstart( display_name: str, - paths: List[str], -) -> Tuple[rag.RagCorpus, GenerationResponse]: + gcs_path: str, +) -> Tuple[types.RagCorpus, genai_types.GenerateContentResponse]: # [START generativeaionvertexai_rag_quickstart] - from vertexai import rag - from vertexai.generative_models import GenerativeModel, Tool - import vertexai + import agentplatform + + from agentplatform import types + from google import genai + from google.genai import types as genai_types + # Create a RAG Corpus, Import Files, and Generate a response # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" + # MODEL_ID = "gemini-3.5-flash" # display_name = "test_corpus" - # paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir"] # Supports Google Cloud Storage and Google Drive Links + # gcs_path = "gs://my_bucket/my_files_dir/*" + # google_drive_path ="https://drive.google.com/file/d/123" - # Initialize Vertex AI API once per session - vertexai.init(project=PROJECT_ID, location="us-east4") + # Initialize Agent Platform client once per session + client = agentplatform.Client(project=PROJECT_ID, location="us-east4") - # Create RagCorpus # Configure embedding model, for example "text-embedding-005". - 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" + ), ) - rag_corpus = rag.create_corpus( - display_name=display_name, - backend_config=rag.RagVectorDbConfig( - rag_embedding_model_config=embedding_model_config - ), + # Create RagCorpus + rag_corpus = client.rag.create_corpus( + rag_corpus=types.RagCorpus( + display_name=display_name, + rag_vector_db_config=types.RagVectorDbConfig( + rag_embedding_model_config=embedding_model_config + ) + ) ) # Import Files to the RagCorpus - rag.import_files( - rag_corpus.name, - paths, - # Optional - transformation_config=rag.TransformationConfig( - chunking_config=rag.ChunkingConfig( - chunk_size=512, - chunk_overlap=100, - ), - ), - max_embedding_requests_per_min=1000, # Optional + client.rag.import_files( + name=rag_corpus.name, + import_config=types.ImportRagFilesConfig( + gcs_source=genai_types.GcsSource(uris=[gcs_path]), + rag_file_transformation_config=types.RagFileTransformationConfig( + rag_file_chunking_config=types.RagFileChunkingConfig( + chunk_size=512, + chunk_overlap=100, + ) + ), # optional + max_embedding_requests_per_min=1000, # optional + ) ) # Direct context retrieval - rag_retrieval_config = rag.RagRetrievalConfig( + rag_retrieval_config = genai_types.RagRetrievalConfig( top_k=3, # Optional - filter=rag.Filter(vector_distance_threshold=0.5), # Optional + filter=genai_types.RagRetrievalConfigFilter( + vector_distance_threshold=0.5 + ), # Optional ) - response = rag.retrieval_query( - rag_resources=[ - rag.RagResource( - rag_corpus=rag_corpus.name, - # Optional: supply IDs from `rag.list_files()`. - # rag_file_ids=["rag-file-1", "rag-file-2", ...], - ) - ], - text="What is RAG and why it is helpful?", - rag_retrieval_config=rag_retrieval_config, + response = client.rag.retrieve_contexts( + vertex_rag_store=genai_types.VertexRagStore( + rag_resources=[ + genai_types.VertexRagStoreRagResource( + rag_corpus=rag_corpus.name, + ) + ], + ), + query=types.RagQuery( + text="What is RAG and why it is helpful?", + rag_retrieval_config=rag_retrieval_config, + ) ) print(response) # Enhance generation # Create a RAG retrieval tool - rag_retrieval_tool = Tool.from_retrieval( - retrieval=rag.Retrieval( - source=rag.VertexRagStore( + rag_retrieval_tool = genai_types.Tool( + retrieval=genai_types.Retrieval( + vertex_rag_store=genai_types.VertexRagStore( rag_resources=[ - rag.RagResource( - rag_corpus=rag_corpus.name, # Currently only 1 corpus is allowed. + genai_types.VertexRagStoreRagResource( + rag_corpus=rag_corpus.name, # Optional: supply IDs from `rag.list_files()`. # rag_file_ids=["rag-file-1", "rag-file-2", ...], ) @@ -105,13 +118,19 @@ def quickstart( ) ) - # Create a Gemini model instance - rag_model = GenerativeModel( - model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool] - ) + # Call generate_content with the tool using the GenAI SDK + + # Create a GenAI SDK client + genai_client = genai.Client(enterprise=True, project=PROJECT_ID, location="us-east4") + - # Generate response - response = rag_model.generate_content("What is RAG and why it is helpful?") + response = genai_client.models.generate_content( + model=MODEL_ID, + contents="What is RAG and why it is helpful?", + config=genai_types.GenerateContentConfig( + tools=[rag_retrieval_tool] + ) + ) print(response.text) # Example response: # RAG stands for Retrieval-Augmented Generation. diff --git a/generative_ai/rag/retrieval_query_example.py b/generative_ai/rag/retrieval_query_example.py index 6d949b8268b..c887298550b 100644 --- a/generative_ai/rag/retrieval_query_example.py +++ b/generative_ai/rag/retrieval_query_example.py @@ -14,39 +14,47 @@ import os -from google.cloud.aiplatform_v1beta1 import RetrieveContextsResponse +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") def retrieval_query( corpus_name: str, -) -> RetrieveContextsResponse: +) -> types.RetrieveContextsResponse: # [START generativeaionvertexai_rag_retrieval_query] - from vertexai import rag - import vertexai + import agentplatform + + from agentplatform import types + from google.genai import types as genai_types # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/[PROJECT_ID]/locations/us-central1/ragCorpora/[rag_corpus_id]" - # 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-east4") - response = rag.retrieval_query( - rag_resources=[ - rag.RagResource( - rag_corpus=corpus_name, - # Optional: supply IDs from `rag.list_files()`. - # rag_file_ids=["rag-file-1", "rag-file-2", ...], - ) - ], - text="Hello World!", - rag_retrieval_config=rag.RagRetrievalConfig( - top_k=10, - filter=rag.utils.resources.Filter(vector_distance_threshold=0.5), + response = client.rag.retrieve_contexts( + vertex_rag_store=genai_types.VertexRagStore( + rag_resources=[ + genai_types.VertexRagStoreRagResource( + rag_corpus=corpus_name, + # Optional: supply IDs from `rag.list_files()`. + # rag_file_ids=["rag-file-1", "rag-file-2", ...], + ) + ], ), + query=types.RagQuery( + text="Hello World!", + rag_retrieval_config=genai_types.RagRetrievalConfig( + top_k=10, + filter=genai_types.RagRetrievalConfigFilter( + vector_distance_threshold=0.5 + ), + ), + ) ) print(response) # Example response: diff --git a/generative_ai/rag/test_rag_examples.py b/generative_ai/rag/test_rag_examples.py index 3d562f5463c..c35b3c63bba 100644 --- a/generative_ai/rag/test_rag_examples.py +++ b/generative_ai/rag/test_rag_examples.py @@ -16,9 +16,6 @@ import os from pathlib import Path -import pytest -import vertexai - import create_corpus_example import create_corpus_feature_store_example import create_corpus_pinecone_example @@ -34,6 +31,7 @@ import import_files_example import list_corpora_example import list_files_example +import pytest import quickstart_example import retrieval_query_example import upload_file_example @@ -47,9 +45,6 @@ GCS_FILE = "gs://cloud-samples-data/generative-ai/pdf/earnings_statement.pdf" -vertexai.init(project=PROJECT_ID, location=LOCATION) - - @pytest.fixture(scope="module", name="test_file") def test_file_fixture() -> None: file_path = Path("./hello.txt") @@ -156,14 +151,14 @@ def test_get_corpus(test_corpus: pytest.fixture) -> None: def test_list_corpora(test_corpus: pytest.fixture) -> None: corpora = list_corpora_example.list_corpora() - assert any(c.display_name == test_corpus.display_name for c in corpora) + assert any(c.display_name == test_corpus.display_name for c in corpora.rag_corpora) def test_upload_file(test_corpus: pytest.fixture, test_file: pytest.fixture) -> None: rag_file = upload_file_example.upload_file(test_corpus.name, test_file) assert rag_file files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -171,7 +166,7 @@ def test_import_files(test_corpus: pytest.fixture) -> None: response = import_files_example.import_files(test_corpus.name, [GCS_FILE]) assert response.imported_rag_files_count > 0 files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -182,7 +177,7 @@ async def test_import_files_async(test_corpus: pytest.fixture) -> None: ) assert result.imported_rag_files_count > 0 files = list_files_example.list_files(test_corpus.name) - imported_file = next(iter(files)) + imported_file = next(iter(files.rag_files)) delete_file_example.delete_file(imported_file.name) @@ -193,7 +188,7 @@ def test_get_file(uploaded_file: pytest.fixture) -> None: def test_list_files(test_corpus: pytest.fixture, uploaded_file: pytest.fixture) -> None: files = list_files_example.list_files(test_corpus.name) - assert any(f.name == uploaded_file.name for f in files) + assert any(f.name == uploaded_file.name for f in files.rag_files) def test_retrieval_query(test_corpus: pytest.fixture) -> None: @@ -210,8 +205,8 @@ def test_generate_content_with_rag(test_corpus: pytest.fixture) -> None: def test_quickstart() -> None: corpus, response = quickstart_example.quickstart( - "test_corpus_quickstart", [GCS_FILE] + "test_corpus_quickstart", GCS_FILE ) assert response assert response.text - delete_corpus_example.delete_corpus(corpus.name) + delete_corpus_example.delete_corpus(corpus.name) \ No newline at end of file diff --git a/generative_ai/rag/upload_file_example.py b/generative_ai/rag/upload_file_example.py index f56cf23f2dc..d3d964d700a 100644 --- a/generative_ai/rag/upload_file_example.py +++ b/generative_ai/rag/upload_file_example.py @@ -16,7 +16,7 @@ from typing import Optional -from vertexai import rag +from agentplatform import types PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT") @@ -26,31 +26,28 @@ def upload_file( path: str, display_name: Optional[str] = None, description: Optional[str] = None, -) -> rag.RagFile: +) -> types.RagFile: # [START generativeaionvertexai_rag_upload_file] - from vertexai import rag - import vertexai + import agentplatform # TODO(developer): Update and un-comment below lines # PROJECT_ID = "your-project-id" # corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}" # path = "path/to/local/file.txt" # display_name = "file_display_name" - # description = "file 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-east4") - rag_file = rag.upload_file( + rag_file = client.rag.upload_file( corpus_name=corpus_name, path=path, display_name=display_name, - description=description, ) print(rag_file) # RagFile(name='projects/[PROJECT_ID]/locations/us-central1/ragCorpora/1234567890/ragFiles/09876543', - # display_name='file_display_name', description='file description') + # display_name='file_display_name') # [END generativeaionvertexai_rag_upload_file] return rag_file