Skip to content

Commit dbbfe4a

Browse files
Thoughtseize1Thoughtseize1
authored
refactor(generative_ai): (set-4) Refactored GenAI Samples (GoogleCloudPlatform#12447)
* Refactored GenAI Samples * Fix lint
1 parent ffa16d8 commit dbbfe4a

20 files changed

+117
-168
lines changed

‎generative_ai/multimodal_embedding_image.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ def get_image_embeddings() -> MultiModalEmbeddingResponse:
2929
import vertexai
3030
from vertexai.vision_models import Image, MultiModalEmbeddingModel
3131

32-
# TODO(developer): Update project
32+
# TODO(developer): Update project_id and location
3333
vertexai.init(project=PROJECT_ID, location="us-central1")
3434

3535
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")

‎generative_ai/multimodal_embedding_image_video_text.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def get_image_video_text_embeddings() -> MultiModalEmbeddingResponse:
3030
from vertexai.vision_models import Image, MultiModalEmbeddingModel, Video
3131
from vertexai.vision_models import VideoSegmentConfig
3232

33-
# TODO(developer): Update project
33+
# TODO(developer): Update project_id and location
3434
vertexai.init(project=PROJECT_ID, location="us-central1")
3535

3636
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")

‎generative_ai/multimodal_embedding_video.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ def get_video_embeddings() -> MultiModalEmbeddingResponse:
3030
from vertexai.vision_models import MultiModalEmbeddingModel, Video
3131
from vertexai.vision_models import VideoSegmentConfig
3232

33-
# TODO(developer): Update project
33+
# TODO(developer): Update project_id and location
3434
vertexai.init(project=PROJECT_ID, location="us-central1")
3535

3636
model = MultiModalEmbeddingModel.from_pretrained("multimodalembedding")

‎generative_ai/rag.py‎

Lines changed: 43 additions & 47 deletions
Original file line numberDiff line numberDiff line change
@@ -13,12 +13,14 @@
1313
# limitations under the License.
1414

1515
# flake8: noqa ANN001, ANN201
16+
import os
1617

1718
from typing import List, Optional
1819

20+
PROJECT_ID = os.getenv("GOOGLE_CLOUD_PROJECT")
21+
1922

2023
def create_corpus(
21-
project_id: str,
2224
display_name: Optional[str] = None,
2325
description: Optional[str] = None,
2426
):
@@ -28,12 +30,12 @@ def create_corpus(
2830
import vertexai
2931

3032
# TODO(developer): Update and un-comment below lines
31-
# project_id = "PROJECT_ID"
33+
# PROJECT_ID = "your-project-id"
3234
# display_name = "test_corpus"
3335
# description = "Corpus Description"
3436

3537
# Initialize Vertex AI API once per session
36-
vertexai.init(project=project_id, location="us-central1")
38+
vertexai.init(project=PROJECT_ID, location="us-central1")
3739

3840
# Configure embedding model
3941
embedding_model_config = rag.EmbeddingModelConfig(
@@ -50,36 +52,36 @@ def create_corpus(
5052
return corpus
5153

5254

53-
def get_corpus(project_id: str, corpus_name: str):
55+
def get_corpus(corpus_name: str):
5456
# [START generativeaionvertexai_rag_get_corpus]
5557

5658
from vertexai.preview import rag
5759
import vertexai
5860

5961
# TODO(developer): Update and un-comment below lines
60-
# project_id = "PROJECT_ID"
61-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
62+
# PROJECT_ID = "your-project-id"
63+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
6264

6365
# Initialize Vertex AI API once per session
64-
vertexai.init(project=project_id, location="us-central1")
66+
vertexai.init(project=PROJECT_ID, location="us-central1")
6567

6668
corpus = rag.get_corpus(name=corpus_name)
6769
print(corpus)
6870
# [END generativeaionvertexai_rag_get_corpus]
6971
return corpus
7072

7173

72-
def list_corpora(project_id: str):
74+
def list_corpora():
7375
# [START generativeaionvertexai_rag_list_corpora]
7476

7577
from vertexai.preview import rag
7678
import vertexai
7779

7880
# TODO(developer): Update and un-comment below lines
79-
# project_id = "PROJECT_ID"
81+
# PROJECT_ID = "your-project-id"
8082

8183
# Initialize Vertex AI API once per session
82-
vertexai.init(project=project_id, location="us-central1")
84+
vertexai.init(project=PROJECT_ID, location="us-central1")
8385

8486
corpora = rag.list_corpora()
8587
print(corpora)
@@ -88,7 +90,6 @@ def list_corpora(project_id: str):
8890

8991

9092
def upload_file(
91-
project_id: str,
9293
corpus_name: str,
9394
path: str,
9495
display_name: Optional[str] = None,
@@ -100,14 +101,14 @@ def upload_file(
100101
import vertexai
101102

102103
# TODO(developer): Update and un-comment below lines
103-
# project_id = "PROJECT_ID"
104-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
104+
# PROJECT_ID = "your-project-id"
105+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
105106
# path = "path/to/local/file.txt"
106107
# display_name = "file_display_name"
107108
# description = "file description"
108109

109110
# Initialize Vertex AI API once per session
110-
vertexai.init(project=project_id, location="us-central1")
111+
vertexai.init(project=PROJECT_ID, location="us-central1")
111112

112113
rag_file = rag.upload_file(
113114
corpus_name=corpus_name,
@@ -121,7 +122,6 @@ def upload_file(
121122

122123

123124
def import_files(
124-
project_id: str,
125125
corpus_name: str,
126126
paths: List[str],
127127
):
@@ -131,12 +131,12 @@ def import_files(
131131
import vertexai
132132

133133
# TODO(developer): Update and un-comment below lines
134-
# project_id = "PROJECT_ID"
135-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
134+
# PROJECT_ID = "your-project-id"
135+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
136136
# paths = ["https://drive.google.com/file/123", "gs://my_bucket/my_files_dir"] # Supports Google Cloud Storage and Google Drive Links
137137

138138
# Initialize Vertex AI API once per session
139-
vertexai.init(project=project_id, location="us-central1")
139+
vertexai.init(project=PROJECT_ID, location="us-central1")
140140

141141
response = rag.import_files(
142142
corpus_name=corpus_name,
@@ -151,7 +151,6 @@ def import_files(
151151

152152

153153
async def import_files_async(
154-
project_id: str,
155154
corpus_name: str,
156155
paths: List[str],
157156
):
@@ -161,14 +160,14 @@ async def import_files_async(
161160
import vertexai
162161

163162
# TODO(developer): Update and un-comment below lines
164-
# project_id = "PROJECT_ID"
165-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
163+
# PROJECT_ID = "your-project-id"
164+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
166165

167166
# Supports Google Cloud Storage and Google Drive Links
168167
# paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir"]
169168

170169
# Initialize Vertex AI API once per session
171-
vertexai.init(project=project_id, location="us-central1")
170+
vertexai.init(project=PROJECT_ID, location="us-central1")
172171

173172
response = await rag.import_files_async(
174173
corpus_name=corpus_name,
@@ -184,18 +183,18 @@ async def import_files_async(
184183
return result
185184

186185

187-
def get_file(project_id: str, file_name: str):
186+
def get_file(file_name: str):
188187
# [START generativeaionvertexai_rag_get_file]
189188

190189
from vertexai.preview import rag
191190
import vertexai
192191

193192
# TODO(developer): Update and un-comment below lines
194-
# project_id = "PROJECT_ID"
195-
# file_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"
193+
# PROJECT_ID = "your-project-id"
194+
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"
196195

197196
# Initialize Vertex AI API once per session
198-
vertexai.init(project=project_id, location="us-central1")
197+
vertexai.init(project=PROJECT_ID, location="us-central1")
199198

200199
rag_file = rag.get_file(name=file_name)
201200
print(rag_file)
@@ -204,18 +203,18 @@ def get_file(project_id: str, file_name: str):
204203
return rag_file
205204

206205

207-
def list_files(project_id: str, corpus_name: str):
206+
def list_files(corpus_name: str):
208207
# [START generativeaionvertexai_rag_list_files]
209208

210209
from vertexai.preview import rag
211210
import vertexai
212211

213212
# TODO(developer): Update and un-comment below lines
214-
# project_id = "PROJECT_ID"
215-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
213+
# PROJECT_ID = "your-project-id"
214+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
216215

217216
# Initialize Vertex AI API once per session
218-
vertexai.init(project=project_id, location="us-central1")
217+
vertexai.init(project=PROJECT_ID, location="us-central1")
219218

220219
files = rag.list_files(corpus_name=corpus_name)
221220
for file in files:
@@ -225,44 +224,43 @@ def list_files(project_id: str, corpus_name: str):
225224
return files
226225

227226

228-
def delete_file(project_id: str, file_name: str) -> None:
227+
def delete_file(file_name: str) -> None:
229228
# [START generativeaionvertexai_rag_delete_file]
230229

231230
from vertexai.preview import rag
232231
import vertexai
233232

234233
# TODO(developer): Update and un-comment below lines
235-
# project_id = "PROJECT_ID"
236-
# file_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"
234+
# PROJECT_ID = "your-project-id"
235+
# file_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}/ragFiles/{rag_file_id}"
237236

238237
# Initialize Vertex AI API once per session
239-
vertexai.init(project=project_id, location="us-central1")
238+
vertexai.init(project=PROJECT_ID, location="us-central1")
240239

241240
rag.delete_file(name=file_name)
242241
print(f"File {file_name} deleted.")
243242
# [END generativeaionvertexai_rag_delete_file]
244243

245244

246-
def delete_corpus(project_id: str, corpus_name: str) -> None:
245+
def delete_corpus(corpus_name: str) -> None:
247246
# [START generativeaionvertexai_rag_delete_corpus]
248247

249248
from vertexai.preview import rag
250249
import vertexai
251250

252251
# TODO(developer): Update and un-comment below lines
253-
# project_id = "PROJECT_ID"
254-
# corpus_name = "projects/{project_id}/locations/us-central1/ragCorpora/{rag_corpus_id}"
252+
# PROJECT_ID = "your-project-id"
253+
# corpus_name = "projects/{PROJECT_ID}/locations/us-central1/ragCorpora/{rag_corpus_id}"
255254

256255
# Initialize Vertex AI API once per session
257-
vertexai.init(project=project_id, location="us-central1")
256+
vertexai.init(project=PROJECT_ID, location="us-central1")
258257

259258
rag.delete_corpus(name=corpus_name)
260259
print(f"Corpus {corpus_name} deleted.")
261260
# [END generativeaionvertexai_rag_delete_corpus]
262261

263262

264263
def retrieval_query(
265-
project_id: str,
266264
rag_corpus_id: str,
267265
):
268266
# [START generativeaionvertexai_rag_retrieval_query]
@@ -271,11 +269,11 @@ def retrieval_query(
271269
import vertexai
272270

273271
# TODO(developer): Update and un-comment below lines
274-
# project_id = "PROJECT_ID"
272+
# PROJECT_ID = "your-project-id"
275273
# rag_corpus_id = "9183965540115283968" # Only one corpus is supported at this time
276274

277275
# Initialize Vertex AI API once per session
278-
vertexai.init(project=project_id, location="us-central1")
276+
vertexai.init(project=PROJECT_ID, location="us-central1")
279277

280278
response = rag.retrieval_query(
281279
rag_resources=[
@@ -296,7 +294,6 @@ def retrieval_query(
296294

297295

298296
def generate_content_with_rag(
299-
project_id: str,
300297
rag_corpus_id: str,
301298
):
302299
# [START generativeaionvertexai_rag_generate_content]
@@ -306,11 +303,11 @@ def generate_content_with_rag(
306303
import vertexai
307304

308305
# TODO(developer): Update and un-comment below lines
309-
# project_id = "PROJECT_ID"
306+
# PROJECT_ID = "your-project-id"
310307
# rag_corpus_id = "9183965540115283968" # Only one corpus is supported at this time
311308

312309
# Initialize Vertex AI API once per session
313-
vertexai.init(project=project_id, location="us-central1")
310+
vertexai.init(project=PROJECT_ID, location="us-central1")
314311

315312
rag_retrieval_tool = Tool.from_retrieval(
316313
retrieval=rag.Retrieval(
@@ -339,7 +336,6 @@ def generate_content_with_rag(
339336

340337

341338
def quickstart(
342-
project_id: str,
343339
display_name: str,
344340
paths: List[str],
345341
):
@@ -351,12 +347,12 @@ def quickstart(
351347
# Create a RAG Corpus, Import Files, and Generate a response
352348

353349
# TODO(developer): Update and un-comment below lines
354-
# project_id = "PROJECT_ID"
350+
# PROJECT_ID = "your-project-id"
355351
# display_name = "test_corpus"
356352
# paths = ["https://drive.google.com/file/d/123", "gs://my_bucket/my_files_dir"] # Supports Google Cloud Storage and Google Drive Links
357353

358354
# Initialize Vertex AI API once per session
359-
vertexai.init(project=project_id, location="us-central1")
355+
vertexai.init(project=PROJECT_ID, location="us-central1")
360356

361357
# Create RagCorpus
362358
# Configure embedding model, for example "text-embedding-004".

0 commit comments

Comments
 (0)