8000 docs(generative_ai): Update RAG Engine Samples to V1 (#13287) · yan283/python-docs-samples@46785e2 · GitHub
[go: up one dir, main page]
Skip to content

Commit 46785e2

Browse files
holtskinnerholtskinner
authored
docs(generative_ai): Update RAG Engine Samples to V1 (GoogleCloudPlatform#13287)
* docs(generative_ai): Update RAG Engine Samples to V1 * Formatting * Undo test changes
1 parent ac3e8be commit 46785e2

17 files changed

+94
-65
lines changed

‎generative_ai/rag/create_corpus_example.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def create_corpus(
2626
) -> RagCorpus:
2727
# [START generativeaionvertexai_rag_create_corpus]
2828

29-
from vertexai.preview import rag
29+
from vertexai import rag
3030
import vertexai
3131

3232
# TODO(developer): Update and un-comment below lines
@@ -37,15 +37,19 @@ def create_corpus(
3737
# Initialize Vertex AI API once per session
3838
vertexai.init(project=PROJECT_ID, location="us-central1")
3939

40-
# Configure embedding model (Optional)
41-
embedding_model_config = rag.EmbeddingModelConfig(
42-
publisher_model="publishers/google/models/text-embedding-004"
40+
# Configure backend_config
41+
backend_config = rag.RagVectorDbConfig(
42+
rag_embedding_model_config=rag.RagEmbeddingModelConfig(
43+
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
44+
publisher_model="publishers/google/models/text-embedding-005"
45+
)
46+
)
4347
)
4448

4549
corpus = rag.create_corpus(
4650
display_name=display_name,
4751
description=description,
48-
embedding_model_config=embedding_model_config,
52+
backend_config=backend_config,
4953
)
5054
print(corpus)
5155
# Example response:

‎generative_ai/rag/create_corpus_pinecone_example.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def create_corpus_pinecone(
2828
) -> RagCorpus:
2929
# [START generativeaionvertexai_rag_create_corpus_pinecone]
3030

31-
from vertexai.preview import rag
31+
from vertexai import rag
3232
import vertexai
3333

3434
# TODO(developer): Update and un-comment below lines
@@ -42,8 +42,10 @@ def create_corpus_pinecone(
4242
vertexai.init(project=PROJECT_ID, location="us-central1")
4343

4444
# Configure embedding model (Optional)
45-
embedding_model_config = rag.EmbeddingModelConfig(
46-
publisher_model="publishers/google/models/text-embedding-004"
45+
embedding_model_config = rag.RagEmbeddingModelConfig(
46+
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
47+
publisher_model="publishers/google/models/text-embedding-005"
48+
)
4749
)
4850

4951
# Configure Vector DB
@@ -55,8 +57,10 @@ def create_corpus_pinecone(
5557
corpus = rag.create_corpus(
5658
display_name=display_name,
5759
description=description,
58-
embedding_model_config=embedding_model_config,
59-
vector_db=vector_db,
60+
backend_config=rag.RagVectorDbConfig(
61+
rag_embedding_model_config=embedding_model_config,
62+
vector_db=vector_db,
63+
),
6064
)
6165
print F438 (corpus)
6266
# Example response:

‎generative_ai/rag/create_corpus_vector_search_example.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def create_corpus_vector_search(
2828
) -> RagCorpus:
2929
# [START generativeaionvertexai_rag_create_corpus_vector_search]
3030

31-
from vertexai.preview import rag
31+
from vertexai import rag
3232
import vertexai
3333

3434
# TODO(developer): Update and un-comment below lines
@@ -42,8 +42,10 @@ def create_corpus_vector_search(
4242
vertexai.init(project=PROJECT_ID, location="us-central1")
4343

4444
# Configure embedding model (Optional)
45-
embedding_model_config = rag.EmbeddingModelConfig(
46-
publisher_model="publishers/google/models/text-embedding-004"
45+
embedding_model_config = rag.RagEmbeddingModelConfig(
46+
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
47+
publisher_model="publishers/google/models/text-embedding-005"
48+
)
4749
)
4850

4951
# Configure Vector DB
@@ -54,8 +56,10 @@ def create_corpus_vector_search(
5456
corpus = rag.create_corpus(
5557
display_name=display_name,
5658
description=description,
57-
embedding_model_config=embedding_model_config,
58-
vector_db=vector_db,
59+
backend_config=rag.RagVectorDbConfig(
60+
rag_embedding_model_config=embedding_model_config,
61+
vector_db=vector_db,
62+
),
5963
)
6064
print(corpus)
6165
# Example response:

‎generative_ai/rag/create_corpus_vertex_ai_search_example.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
from typing import Optional
1717

18-
from vertexai.preview.rag import RagCorpus
18+
from vertexai import rag
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

@@ -24,10 +24,10 @@ def create_corpus_vertex_ai_search(
2424
vertex_ai_search_engine_name: str,
2525
display_name: Optional[str] = None,
2626
description: Optional[str] = None,
27-
) -> RagCorpus:
27+
) -> rag.RagCorpus:
2828
# [START generativeaionvertexai_rag_create_corpus_vertex_ai_search]
2929

30-
from vertexai.preview import rag
30+
from vertexai import rag
3131
import vertexai
3232

3333
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/delete_corpus_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
def delete_corpus(corpus_name: str) -> None:
2121
# [START generativeaionvertexai_rag_delete_corpus]
2222

23-
from vertexai.preview import rag
23+
from vertexai import rag
2424
import vertexai
2525

2626
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/delete_file_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@
2020
def delete_file(file_name: str) -> None:
2121
# [START generativeaionvertexai_rag_delete_file]
2222

23-
from vertexai.preview import rag
23+
from vertexai import rag
2424
import vertexai
2525

2626
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/generate_content_example.py

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -24,8 +24,8 @@ def generate_content_with_rag(
2424
) -> GenerationResponse:
2525
# [START generativeaionvertexai_rag_generate_content]
2626

27-
from vertexai.preview import rag
28-
from vertexai.preview.generative_models import GenerativeModel, Tool
27+
from vertexai import rag
28+
from vertexai.generative_models import GenerativeModel, Tool
2929
import vertexai
3030

3131
# TODO(developer): Update and un-comment below lines
@@ -45,14 +45,16 @@ def generate_content_with_rag(
4545
# rag_file_ids=["rag-file-1", "rag-file-2", ...],
4646
)
4747
],
48-
similarity_top_k=3, # Optional
49-
vector_distance_threshold=0.5, # Optional
48+
rag_retrieval_config=rag.RagRetrievalConfig(
49+
top_k=10,
50+
filter=rag.utils.resources.Filter(vector_distance_threshold=0.5),
51+
),
5052
),
5153
)
5254
)
5355

5456
rag_model = GenerativeModel(
55-
model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
57+
model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
5658
)
5759
response = rag_model.generate_content("Why is the sky blue?")
5860
print(response.text)

‎generative_ai/rag/get_corpus_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@
2222
def get_corpus(corpus_name: str) -> RagCorpus:
2323
# [START generativeaionvertexai_rag_get_corpus]
2424

25-
from vertexai.preview import rag
25+
from vertexai import rag
2626
import vertexai
2727

2828
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/get_file_example.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,15 @@
1414

1515
import os
1616

17-
from google.cloud.aiplatform_v1beta1 import RagFile
17+
from google.cloud.aiplatform_v1 import RagFile
1818

1919
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2020

2121

2222
def get_file(file_name: str) -> RagFile:
2323
# [START generativeaionvertexai_rag_get_file]
2424

25-
from vertexai.preview import rag
25+
from vertexai import rag
2626
import vertexai
2727

2828
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/import_files_async_example.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616

1717
from typing import List
1818

19-
from google.cloud.aiplatform_v1beta1 import ImportRagFilesResponse
19+
from google.cloud.aiplatform_v1 import ImportRagFilesResponse
2020

2121
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2222

@@ -27,7 +27,7 @@ async def import_files_async(
2727
) -> ImportRagFilesResponse:
2828
# [START generativeaionvertexai_rag_import_files_async]
2929

30-
from vertexai.preview import rag
30+
from vertexai import rag
3131
import vertexai
3232

3333
# TODO(developer): Update and un-comment below lines
@@ -40,11 +40,12 @@ async def import_files_async(
4040
# Initialize Vertex AI API once per session
4141
vertexai.init(project=PROJECT_ID, location="us-central1")
4242

43-
response = await rag.import_files_async(
43+
response = await rag.import_files(
4444
corpus_name=corpus_name,
4545
paths=paths,
46-
chunk_size=512, # Optional
47-
chunk_overlap=100, # Optional
46+
transformation_config=rag.TransformationConfig(
47+
rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
48+
),
4849
max_embedding_requests_per_min=900, # Optional
4950
)
5051

‎generative_ai/rag/import_files_example.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
import os
1616
from typing import List
1717

18-
from google.cloud.aiplatform_v1beta1 import ImportRagFilesResponse
18+
from google.cloud.aiplatform_v1 import ImportRagFilesResponse
1919

2020
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2121

@@ -26,7 +26,7 @@ def import_files(
2626
) -> ImportRagFilesResponse:
2727
# [START generativeaionvertexai_rag_import_files]
2828

29-
from vertexai.preview import rag
29+
from vertexai import rag
3030
import vertexai
3131

3232
# TODO(developer): Update and un-comment below lines
@@ -40,8 +40,9 @@ def import_files(
4040
response = rag.import_files(
4141
corpus_name=corpus_name,
4242
paths=paths,
43-
chunk_size=512, # Optional
44-
chunk_overlap=100, # Optional
43+
transformation_config=rag.TransformationConfig(
44+
rag.ChunkingConfig(chunk_size=512, chunk_overlap=100)
45+
),
4546
max_embedding_requests_per_min=900, # Optional
4647
)
4748
print(f"Imported {response.imported_rag_files_count} files.")

‎generative_ai/rag/list_corpora_example.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424
def list_corpora() -> ListRagCorporaPager:
2525
# [START generativeaionvertexai_rag_list_corpora]
2626

27-
from vertexai.preview import rag
27+
from vertexai import rag
2828
import vertexai
2929

3030
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/list_files_example.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414

1515
import os
1616

17-
from google.cloud.aiplatform_v1beta1.services.vertex_rag_data_service.pagers import (
17+
from google.cloud.aiplatform_v1.services.vertex_rag_data_service.pagers import (
1818
ListRagFilesPager,
1919
)
2020

@@ -24,7 +24,7 @@
2424
def list_files(corpus_name: str) -> ListRagFilesPager:
2525
# [START generativeaionvertexai_rag_list_files]
2626

27-
from vertexai.preview import rag
27+
from vertexai import rag
2828
import vertexai
2929

3030
# TODO(developer): Update and un-comment below lines

‎generative_ai/rag/quickstart_example.py

Lines changed: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,19 +16,19 @@
1616

1717
from typing import List, Tuple
1818

19+
from vertexai import rag
1920
from vertexai.generative_models import GenerationResponse
20-
from vertexai.preview.rag.utils.resources import RagCorpus
2121

2222
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
2323

2424

2525
def quickstart(
2626
display_name: str,
2727
paths: List[str],
28-
) -> Tuple[RagCorpus, GenerationResponse]:
28+
) -> Tuple[rag.RagCorpus, GenerationResponse]:
2929
# [START generativeaionvertexai_rag_quickstart]
30-
from vertexai.preview import rag
31-
from vertexai.preview.generative_models import GenerativeModel, Tool
30+
from vertexai import rag
31+
from vertexai.generative_models import GenerativeModel, Tool
3232
import vertexai
3333

3434
# Create a RAG Corpus, Import Files, and Generate a response
@@ -42,26 +42,39 @@ def quickstart(
4242
vertexai.init(project=PROJECT_ID, location="us-central1")
4343

4444
# Create RagCorpus
45-
# Configure embedding model, for example "text-embedding-004".
46-
embedding_model_config = rag.EmbeddingModelConfig(
47-
publisher_model="publishers/google/models/text-embedding-004"
45+
# Configure embedding model, for example "text-embedding-005".
46+
embedding_model_config = rag.RagEmbeddingModelConfig(
47+
vertex_prediction_endpoint=rag.VertexPredictionEndpoint(
48+
publisher_model="publishers/google/models/text-embedding-005"
49+
)
4850
)
4951

5052
rag_corpus = rag.create_corpus(
5153
display_name=display_name,
52-
embedding_model_config=embedding_model_config,
54+
backend_config=rag.RagVectorDbConfig(
55+
rag_embedding_model_config=embedding_model_config
56+
),
5357
)
5458

5559
# Import Files to the RagCorpus
5660
rag.import_files(
5761
rag_corpus.name,
5862
paths,
59-
chunk_size=512, # Optional
60-
chunk_overlap=100, # Optional
61-
max_embedding_requests_per_min=900, # Optional
63+
# Optional
64+
transformation_config=rag.TransformationConfig(
65+
chunking_config=rag.ChunkingConfig(
66+
chunk_size=512,
67+
chunk_overlap=100,
68+
),
69+
),
70+
max_embedding_requests_per_min=1000, # Optional
6271
)
6372

6473
# Direct context retrieval
74+
rag_retrieval_config = rag.RagRetrievalConfig(
75+
top_k=3, # Optional
76+
filter=rag.Filter(vector_distance_threshold=0.5), # Optional
77+
)
6578
response = rag.retrieval_query(
6679
rag_resources=[
6780
rag.RagResource(
@@ -71,8 +84,7 @@ def quickstart(
7184
)
7285
],
7386
text="What is RAG and why it is helpful?",
74-
similarity_top_k=10, # Optional
75-
vector_distance_threshold=0.5, # Optional
87+
rag_retrieval_config=rag_retrieval_config,
7688
)
7789
print(response)
7890

@@ -88,14 +100,14 @@ def quickstart(
88100
# rag_file_ids=["rag-file-1", "rag-file-2", ...],
89101
)
90102
],
91-
similarity_top_k=3, # Optional
92-
vector_distance_threshold=0.5, # Optional
103+
rag_retrieval_config=rag_retrieval_config,
93104
),
94105
)
95106
)
96-
# Create a gemini-pro model instance
107+
108+
# Create a Gemini model instance
97109
rag_model = GenerativeModel(
98-
model_name="gemini-1.5-flash-001", tools=[rag_retrieval_tool]
110+
model_name="gemini-2.0-flash-001", tools=[rag_retrieval_tool]
99111
)
100112

101113
# Generate response

‎generative_ai/rag/requirements.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +1 @@
1-
google-cloud-aiplatform[all]==1.74.0
1+
google-cloud-aiplatform==1.87.0

0 commit comments

Comments
 (0)
0