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

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly
Appearance settings

Commit 46785e2

Browse files
holtskinnerholtskinner
authored
docs(generative_ai): Update RAG Engine Samples to V1 (Googl 8000 eCloudPlatform#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(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