|
| 1 | +from __future__ import annotations |
| 2 | + |
1 | 3 | import pytest
|
2 | 4 | from google.api_core.exceptions import BadRequest, NotFound
|
3 | 5 | from vertexai.generative_models import (
|
|
6 | 8 | GenerativeModel,
|
7 | 9 | Part,
|
8 | 10 | )
|
| 11 | +from vertexai.preview.generative_models import ( |
| 12 | + GenerativeModel as PreviewGenerativeModel, |
| 13 | +) |
9 | 14 |
|
10 | 15 | from opentelemetry.instrumentation.vertexai import VertexAIInstrumentor
|
11 | 16 | from opentelemetry.sdk._logs._internal.export.in_memory_log_exporter import (
|
@@ -278,10 +283,37 @@ def test_generate_content_all_input_events(
|
278 | 283 | log_exporter: InMemoryLogExporter,
|
279 | 284 | instrument_with_content: VertexAIInstrumentor,
|
280 | 285 | ):
|
281 |
| - model = GenerativeModel( |
282 |
| - "gemini-1.5-flash-002", |
283 |
| - system_instruction=Part.from_text("You are a clever language model"), |
| 286 | + generate_content_all_input_events( |
| 287 | + GenerativeModel( |
| 288 | + "gemini-1.5-flash-002", |
| 289 | + system_instruction=Part.from_text( |
| 290 | + "You are a clever language model" |
| 291 | + ), |
| 292 | + ), |
| 293 | + log_exporter, |
284 | 294 | )
|
| 295 | + |
| 296 | + |
| 297 | +@pytest.mark.vcr |
| 298 | +def test_preview_generate_content_all_input_events( |
| 299 | + log_exporter: InMemoryLogExporter, |
| 300 | + instrument_with_content: VertexAIInstrumentor, |
| 301 | +): |
| 302 | + generate_content_all_input_events( |
| 303 | + PreviewGenerativeModel( |
| 304 | + "gemini-1.5-flash-002", |
| 305 | + system_instruction=Part.from_text( |
| 306 | + "You are a clever language model" |
| 307 | + ), |
| 308 | + ), |
| 309 | + log_exporter, |
| 310 | + ) |
| 311 | + |
| 312 | + |
| 313 | +def generate_content_all_input_events( |
| 314 | + model: GenerativeModel | PreviewGenerativeModel, |
| 315 | + log_exporter: InMemoryLogExporter, |
| 316 | +): |
285 | 317 | model.generate_content(
|
286 | 318 | [
|
287 | 319 | Content(
|
|
0 commit comments