Skip to content

Commit 76a2b80

Browse files
holtskinnergcf-owl-bot[bot]
holtskinner
and
gcf-owl-bot[bot]
authored
samples(discoveryengine): Add Grounded Generation API Samples (#12064)
* samples(discoveryengine): Add Grounded Generation API Samples NOTE: This feature is allowlist only, but is visible in the public documentation. --------- Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
1 parent c3081f1 commit 76a2b80

File tree

6 files changed

+216
-8
lines changed

6 files changed

+216
-8
lines changed

‎discoveryengine/requirements-test.txt

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,3 @@
11
pytest==8.2.0
2-
google-api-core==2.17.1
2+
google-api-core==2.21.0
3+
google-cloud-resource-manager==1.12.5

‎discoveryengine/requirements.txt

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,2 +1,2 @@
1-
google-cloud-discoveryengine==0.12.1
2-
google-api-core==2.19.0
1+
google-cloud-discoveryengine==0.12.3
2+
google-api-core==2.21.0

‎discoveryengine/site_search_engine_sample.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,8 @@ def create_target_site(
2929
# project_id = "YOUR_PROJECT_ID"
3030
# location = "YOUR_LOCATION" # Values: "global"
3131
# data_store_id = "YOUR_DATA_STORE_ID"
32-
# uri_pattern = "https://cloud.google.com/generative-ai-app-builder/docs/*"
32+
# NOTE: Do not include http or https protocol in the URI pattern
33+
# uri_pattern = "cloud.google.com/generative-ai-app-builder/docs/*"
3334

3435
# For more information, refer to:
3536
# https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store

‎discoveryengine/site_search_engine_sample_test.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ def test_create_target_site():
2828
project_id,
2929
location,
3030
data_store_id,
31-
uri_pattern="https://cloud.google.com/generative-ai-app-builder/docs/*",
31+
uri_pattern="cloud.google.com/generative-ai-app-builder/docs/*",
3232
)
3333
assert response, response
3434
match = re.search(r"\/targetSites\/([^\/]+)", response.name)

‎discoveryengine/standalone_apis_sample.py

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

16-
from google.cloud import discoveryengine_v1alpha as discoveryengine
16+
from google.cloud import discoveryengine_v1 as discoveryengine
1717

1818

1919
def check_grounding_sample(
2020
project_id: str,
2121
) -> discoveryengine.CheckGroundingResponse:
2222
# [START genappbuilder_check_grounding]
23-
from google.cloud import discoveryengine_v1alpha as discoveryengine
23+
from google.cloud import discoveryengine_v1 as discoveryengine
2424

2525
# TODO(developer): Uncomment these variables before running the sample.
2626
# project_id = "YOUR_PROJECT_ID"
@@ -78,7 +78,7 @@ def rank_sample(
7878
project_id: str,
7979
) -> discoveryengine.RankResponse:
8080
# [START genappbuilder_rank]
81-
from google.cloud import discoveryengine_v1alpha as discoveryengine
81+
from google.cloud import discoveryengine_v1 as discoveryengine
8282

8383
# TODO(developer): Uncomment these variables before running the sample.
8484
# project_id = "YOUR_PROJECT_ID"
@@ -123,3 +123,183 @@ def rank_sample(
123123
# [END genappbuilder_rank]
124124

125125
return response
126+
127+
128+
def grounded_generation_inline_vais_sample(
129+
project_number: str,
130+
engine_id: str,
131+
) -> discoveryengine.GenerateGroundedContentResponse:
132+
# [START genappbuilder_grounded_generation_inline_vais]
133+
from google.cloud import discoveryengine_v1 as discoveryengine
134+
135+
# TODO(developer): Uncomment these variables before running the sample.
136+
# project_number = "YOUR_PROJECT_NUMBER"
137+
# engine_id = "YOUR_ENGINE_ID"
138+
139+
client = discoveryengine.GroundedGenerationServiceClient()
140+
141+
request = discoveryengine.GenerateGroundedContentRequest(
142+
# The full resource name of the location.
143+
# Format: projects/{project_number}/locations/{location}
144+
location=client.common_location_path(project=project_number, location="global"),
145+
generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
146+
model_id="gemini-1.5-flash",
147+
),
148+
# Conversation between user and model
149+
contents=[
150+
discoveryengine.GroundedGenerationContent(
151+
role="user",
152+
parts=[
153+
discoveryengine.GroundedGenerationContent.Part(
154+
text="How did Google do in 2020? Where can I find BigQuery docs?"
155+
)
156+
],
157+
)
158+
],
159+
system_instruction=discoveryengine.GroundedGenerationContent(
160+
parts=[
161+
discoveryengine.GroundedGenerationContent.Part(
162+
text="Add a smiley emoji after the answer."
163+
)
164+
],
165+
),
166+
# What to ground on.
167+
grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
168+
grounding_sources=[
169+
discoveryengine.GenerateGroundedContentRequest.GroundingSource(
170+
inline_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.InlineSource(
171+
grounding_facts=[
172+
discoveryengine.GroundingFact(
173+
fact_text=(
174+
"The BigQuery documentation can be found at https://cloud.google.com/bigquery/docs/introduction"
175+
),
176+
attributes={
177+
"title": "BigQuery Overview",
178+
"uri": "https://cloud.google.com/bigquery/docs/introduction",
179+
},
180+
),
181+
]
182+
),
183+
),
184+
discoveryengine.GenerateGroundedContentRequest.GroundingSource(
185+
search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.SearchSource(
186+
# The full resource name of the serving config for a Vertex AI Search App
187+
serving_config=f"projects/{project_number}/locations/global/collections/default_collection/engines/{engine_id}/servingConfigs/default_search",
188+
),
189+
),
190+
]
191+
),
192+
)
193+
response = client.generate_grounded_content(request)
194+
195+
# Handle the response
196+
print(response)
197+
# [END genappbuilder_grounded_generation_inline_vais]
198+
199+
return response
200+
201+
202+
def grounded_generation_google_search_sample(
203+
project_number: str,
204+
) -> discoveryengine.GenerateGroundedContentResponse:
205+
# [START genappbuilder_grounded_generation_google_search]
206+
from google.cloud import discoveryengine_v1 as discoveryengine
207+
208+
# TODO(developer): Uncomment these variables before running the sample.
209+
# project_number = "YOUR_PROJECT_NUMBER"
210+
211+
client = discoveryengine.GroundedGenerationServiceClient()
212+
213+
request = discoveryengine.GenerateGroundedContentRequest(
214+
# The full resource name of the location.
215+
# Format: projects/{project_number}/locations/{location}
216+
location=client.common_location_path(project=project_number, location="global"),
217+
generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
218+
model_id="gemini-1.5-flash",
219+
),
220+
# Conversation between user and model
221+
contents=[
222+
discoveryengine.GroundedGenerationContent(
223+
role="user",
224+
parts=[
225+
discoveryengine.GroundedGenerationContent.Part(
226+
text="How much is Google stock?"
227+
)
228+
],
229+
)
230+
],
231+
system_instruction=discoveryengine.GroundedGenerationContent(
232+
parts=[
233+
discoveryengine.GroundedGenerationContent.Part(text="Be comprehensive.")
234+
],
235+
),
236+
# What to ground on.
237+
grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
238+
grounding_sources=[
239+
discoveryengine.GenerateGroundedContentRequest.GroundingSource(
240+
google_search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.GoogleSearchSource(
241+
# Optional: For Dynamic Retrieval
242+
dynamic_retrieval_config=discoveryengine.GenerateGroundedContentRequest.DynamicRetrievalConfiguration(
243+
predictor=discoveryengine.GenerateGroundedContentRequest.DynamicRetrievalConfiguration.DynamicRetrievalPredictor(
244+
threshold=0.7
245+
)
246+
)
247+
)
248+
),
249+
]
250+
),
251+
)
252+
response = client.generate_grounded_content(request)
253+
254+
# Handle the response
255+
print(response)
256+
# [END genappbuilder_grounded_generation_google_search]
257+
258+
return response
259+
260+
261+
def grounded_generation_streaming_sample(
262+
project_number: str,
263+
) -> discoveryengine.GenerateGroundedContentResponse:
264+
# [START genappbuilder_grounded_generation_streaming]
265+
from google.cloud import discoveryengine_v1 as discoveryengine
266+
267+
# TODO(developer): Uncomment these variables before running the sample.
268+
# project_id = "YOUR_PROJECT_ID"
269+
270+
client = discoveryengine.GroundedGenerationServiceClient()
271+
272+
request = discoveryengine.GenerateGroundedContentRequest(
273+
# The full resource name of the location.
274+
# Format: projects/{project_number}/locations/{location}
275+
location=client.common_location_path(project=project_number, location="global"),
276+
generation_spec=discoveryengine.GenerateGroundedContentRequest.GenerationSpec(
277+
model_id="gemini-1.5-flash",
278+
),
279+
# Conversation between user and model
280+
contents=[
281+
discoveryengine.GroundedGenerationContent(
282+
role="user",
283+
parts=[
284+
discoveryengine.GroundedGenerationContent.Part(
285+
text="Summarize how to delete a data store in Vertex AI Agent Builder?"
286+
)
287+
],
288+
)
289+
],
290+
grounding_spec=discoveryengine.GenerateGroundedContentRequest.GroundingSpec(
291+
grounding_sources=[
292+
discoveryengine.GenerateGroundedContentRequest.GroundingSource(
293+
google_search_source=discoveryengine.GenerateGroundedContentRequest.GroundingSource.GoogleSearchSource()
294+
),
295+
]
296+
),
297+
)
298+
responses = client.stream_generate_grounded_content(iter([request]))
299+
300+
for response in responses:
301+
# Handle the response
302+
print(response)
303+
# [END genappbuilder_grounded_generation_streaming]
304+
305+
return response

‎discoveryengine/standalone_apis_sample_test.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717

1818
from discoveryengine import standalone_apis_sample
1919

20+
from google.cloud import resourcemanager_v3
21+
2022
project_id = os.environ["GOOGLE_CLOUD_PROJECT"]
2123

2224

@@ -32,3 +34,27 @@ def test_rank():
3234
response = standalone_apis_sample.rank_sample(project_id)
3335
assert response
3436
assert response.records
37+
38+
39+
def test_grounded_generation_inline_vais_sample():
40+
# Grounded Generation requires Project Number
41+
client = resourcemanager_v3.ProjectsClient()
42+
project = client.get_project(name=client.project_path(project_id))
43+
project_number = client.parse_project_path(project.name)["project"]
44+
45+
response = standalone_apis_sample.grounded_generation_inline_vais_sample(
46+
project_number, engine_id="test-search-engine_1689960780551"
47+
)
48+
assert response
49+
50+
51+
def test_grounded_generation_google_search_sample():
52+
# Grounded Generation requires Project Number
53+
client = resourcemanager_v3.ProjectsClient()
54+
project = client.get_project(name=client.project_path(project_id))
55+
project_number = client.parse_project_path(project.name)["project"]
56+
57+
response = standalone_apis_sample.grounded_generation_google_search_sample(
58+
project_number
59+
)
60+
assert response

0 commit comments

Comments
 (0)