Files
wehub-resource-sync 5a558eb09e
TypeScript SDK Compatibility V1.x E2E Tests / Select Node version matrix (push) Has been cancelled
TypeScript SDK Compatibility V1.x E2E Tests / TypeScript SDK Compatibility V1.x E2E Tests Node ${{matrix.node_version}} (push) Has been cancelled
TypeScript SDK E2E Tests / TypeScript SDK E2E Tests Node ${{matrix.node_version}} (push) Has been cancelled
Opik Optimizer - E2E Tests / build-opik (push) Has been cancelled
TypeScript SDK Compatibility V1.x E2E Tests / build-opik (push) Has been cancelled
Python SDK E2E Tests / Select Python version matrix (push) Has been cancelled
Python SDK E2E Tests / Python SDK E2E Tests ${{matrix.python_version}} (push) Has been cancelled
Python SDK E2E Tests / build-opik (push) Has been cancelled
Python SDK Compatibility V1.x E2E Tests / Select Python version matrix (push) Has been cancelled
Python SDK Compatibility V1.x E2E Tests / Python SDK Compatibility V1.x E2E Tests ${{matrix.python_version}} (push) Has been cancelled
Python SDK Compatibility V1.x E2E Tests / build-opik (push) Has been cancelled
TypeScript SDK E2E Tests / Select Node version matrix (push) Has been cancelled
TypeScript SDK E2E Tests / build-opik (push) Has been cancelled
Opik Optimizer - E2E Tests / Opik Optimizer E2E Tests Python ${{matrix.python_version}} (push) Has been cancelled
Opik Optimizer - E2E Tests / Opik Optimizer Integration Smoke Tests (push) Has been cancelled
🐙 Code Quality / detect (push) Has been cancelled
🐙 Code Quality / lint (${{ matrix.leg.name }}) (push) Has been cancelled
🐙 Code Quality / summary (push) Has been cancelled
TypeScript SDK Library Integration Tests / Check Secrets (push) Has been cancelled
TypeScript SDK Library Integration Tests / opik-vercel (Vercel AI SDK / eve) (push) Has been cancelled
SDK Library Integration Tests Runner / Check Secrets (push) Has been cancelled
SDK Library Integration Tests Runner / Missed OpenAI API Key Warning (push) Has been cancelled
SDK Library Integration Tests Runner / Build (push) Has been cancelled
SDK Library Integration Tests Runner / openai_tests (push) Has been cancelled
SDK Library Integration Tests Runner / langchain_tests (push) Has been cancelled
SDK Library Integration Tests Runner / langchain_legacy_tests (push) Has been cancelled
SDK Library Integration Tests Runner / llama_index_tests (push) Has been cancelled
SDK Library Integration Tests Runner / anthropic_tests (push) Has been cancelled
SDK Library Integration Tests Runner / mistral_tests (push) Has been cancelled
SDK Library Integration Tests Runner / groq_tests (push) Has been cancelled
SDK Library Integration Tests Runner / aisuite_tests (push) Has been cancelled
SDK Library Integration Tests Runner / haystack_tests (push) Has been cancelled
SDK Library Integration Tests Runner / dspy_tests (push) Has been cancelled
SDK Library Integration Tests Runner / crewai_v0_tests (push) Has been cancelled
SDK Library Integration Tests Runner / crewai_v1_tests (push) Has been cancelled
SDK Library Integration Tests Runner / genai_tests (push) Has been cancelled
SDK Library Integration Tests Runner / adk_tests (push) Has been cancelled
SDK Library Integration Tests Runner / adk_legacy_1_3_0_tests (push) Has been cancelled
SDK Library Integration Tests Runner / evaluation_metrics_tests (push) Has been cancelled
SDK Library Integration Tests Runner / bedrock_tests (push) Has been cancelled
SDK Library Integration Tests Runner / litellm_tests (push) Has been cancelled
SDK Library Integration Tests Runner / harbor_tests (push) Has been cancelled
SDK Library Integration Tests Runner / Slack Notification (push) Has been cancelled
Lint Opik Helm Chart / render-equality (push) Has been cancelled
Opik Optimizer - Unit Tests / Opik Optimizer Unit Tests Python ${{matrix.python_version}} (push) Has been cancelled
Python BE E2E Tests / Python BE E2E (push) Has been cancelled
Python Backend Tests / run-python-backend-tests (push) Has been cancelled
Python SDK Unit Tests / Python SDK Unit Tests ${{matrix.python_version}} (push) Has been cancelled
Release Drafter / update_release_draft (push) Has been cancelled
SDK E2E Libraries Integration Tests / Check Secrets (push) Has been cancelled
SDK E2E Libraries Integration Tests / Missed OpenAI API Key Warning (push) Has been cancelled
SDK E2E Libraries Integration Tests / build-opik (push) Has been cancelled
SDK E2E Libraries Integration Tests / E2E Lib Integration Python ${{matrix.python_version}} (push) Has been cancelled
TypeScript SDK Integration Build & Publish / build-and-publish (opik-gemini) (push) Has been cancelled
TypeScript SDK Integration Build & Publish / build-and-publish (opik-langchain) (push) Has been cancelled
TypeScript SDK Integration Build & Publish / build-and-publish (opik-openai) (push) Has been cancelled
TypeScript SDK Integration Build & Publish / build-and-publish (opik-otel) (push) Has been cancelled
TypeScript SDK Integration Build & Publish / build-and-publish (opik-vercel) (push) Has been cancelled
TypeScript SDK Build & Publish / build-and-publish (push) Has been cancelled
TypeScript SDK Unit Tests / Test on Node ${{ matrix.node-version }} (push) Has been cancelled
Backend Tests / discover-tests (push) Has been cancelled
Backend Tests / ${{ matrix.name }} (push) Has been cancelled
Build and Publish SDK / build-and-publish (push) Has been cancelled
Build Opik Docker Images / set-version (push) Has been cancelled
Build Opik Docker Images / build-backend (push) Has been cancelled
Build Opik Docker Images / build-sandbox-executor-python (push) Has been cancelled
Build Opik Docker Images / build-python-backend (push) Has been cancelled
Build Opik Docker Images / build-frontend (push) Has been cancelled
Build Opik Docker Images / create-git-tag (push) Has been cancelled
ClickHouse Migration Cluster Check / validate-clickhouse-migrations (push) Has been cancelled
Docs - Publish / run (push) Has been cancelled
E2E Tests - Post Merge (v2) / 🧪 E2E v2 Tests (${{ github.event.inputs.tier || 't1' }}) (push) Has been cancelled
E2E Tests - Post Merge (v2) / 📢 Slack Notification (push) Has been cancelled
Frontend Unit Tests / Test on Node 20 (push) Has been cancelled
Guardrails E2E Tests / Select Python version matrix (push) Has been cancelled
Guardrails E2E Tests / Guardrails E2E Tests ${{matrix.python_version}} (push) Has been cancelled
Guardrails E2E Tests / 📢 Slack Notification (push) Has been cancelled
Guardrails Backend Unit Tests / Guardrails Backend Unit Tests (push) Has been cancelled
Guardrails Backend Unit Tests / 📢 Slack Notification (push) Has been cancelled
Lint Opik Helm Chart / lint-helm-chart (Helm v3.21.0) (push) Has been cancelled
Lint Opik Helm Chart / lint-helm-chart (Helm v4.2.0) (push) Has been cancelled
Lint Opik Helm Chart / unittest-helm-chart (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 13:25:44 +08:00

792 lines
25 KiB
Python

import asyncio
from typing import Any, Dict
import pytest
from google import genai
from google.genai.types import HttpOptions, GenerateContentConfig
import opik
from opik.config import OPIK_PROJECT_DEFAULT_NAME
from opik.integrations.genai import track_genai
from ... import llm_constants
from ...testlib import (
ANY_BUT_NONE,
ANY_DICT,
ANY_LIST,
ANY_STRING,
SpanModel,
TraceModel,
assert_dict_has_keys,
assert_equal,
)
pytestmark = pytest.mark.usefixtures("ensure_vertexai_configured")
MODEL = llm_constants.GEMINI_FLASH
EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT = ANY_DICT.containing(
{
"prompt_tokens": ANY_BUT_NONE,
"completion_tokens": ANY_BUT_NONE,
"total_tokens": ANY_BUT_NONE,
"original_usage.total_token_count": ANY_BUT_NONE,
"original_usage.prompt_token_count": ANY_BUT_NONE,
}
)
def _assert_metadata_contains_required_keys(metadata: Dict[str, Any]):
REQUIRED_METADATA_KEYS = [
"model",
"created_from",
"model_version",
"usage_metadata",
]
assert_dict_has_keys(metadata, REQUIRED_METADATA_KEYS)
@pytest.mark.parametrize(
"project_name, expected_project_name",
[
(None, OPIK_PROJECT_DEFAULT_NAME),
("genai-integration-test", "genai-integration-test"),
],
)
def test_genai_client__generate_content__happyflow(
fake_backend, project_name, expected_project_name
):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client, project_name=project_name)
client.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
config=GenerateContentConfig(max_output_tokens=10),
)
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?", "config": ANY_BUT_NONE},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
project_name=expected_project_name,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={
"contents": "What is the capital of Belarus?",
"config": ANY_BUT_NONE,
},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
project_name=expected_project_name,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__async_generate_content__happyflow(fake_backend):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
response = client.aio.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
)
asyncio.run(response)
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"async_generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"async_generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
@pytest.mark.asyncio
async def test_genai_client__async_generate_content__opik_args__happyflow(fake_backend):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
args_dict = {
"span": {"tags": ["span_tag"], "metadata": {"span_key": "span_value"}},
"trace": {
"thread_id": "conversation-2",
"tags": ["trace_tag"],
"metadata": {"trace_key": "trace_value"},
},
}
_ = await client.aio.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
opik_args=args_dict,
)
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"async_generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai", "span_tag", "trace_tag"],
metadata=ANY_DICT.containing({"trace_key": "trace_value"}),
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
thread_id="conversation-2",
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"async_generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai", "span_tag"],
metadata=ANY_DICT.containing({"span_key": "span_value"}),
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
@pytest.mark.parametrize(
"project_name, expected_project_name",
[
(None, OPIK_PROJECT_DEFAULT_NAME),
("genai-integration-test", "genai-integration-test"),
],
)
def test_genai_client__generate_content_called_inside_another_tracked_function__happyflow(
fake_backend, project_name, expected_project_name
):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
@opik.track(project_name=project_name)
def f():
client.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
)
f()
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name="f",
input={},
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
project_name=expected_project_name,
spans=[
SpanModel(
id=ANY_BUT_NONE,
name="f",
input={},
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
project_name=expected_project_name,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
project_name=expected_project_name,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__async_generate_content_called_inside_another_tracked_function__happyflow(
fake_backend,
):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
@opik.track
async def f():
_ = await client.aio.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
)
asyncio.run(f())
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name="f",
input={},
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
name="f",
input={},
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(
f"async_generate_content: {MODEL}"
),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__generate_content_stream__happyflow(fake_backend):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client, project_name="genai-integration-test")
stream = client.models.generate_content_stream(
model=MODEL,
contents="What is the capital of Belarus?",
)
for _ in stream:
pass
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"generate_content_stream: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
project_name="genai-integration-test",
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"generate_content_stream: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
project_name="genai-integration-test",
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__async_generate_content_stream__happyflow(fake_backend):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
async def stream_example():
stream = await client.aio.models.generate_content_stream(
model=MODEL,
contents="What is the capital of Belarus?",
)
async for _ in stream:
pass
asyncio.run(stream_example())
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"async_generate_content_stream: {MODEL}"),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(
f"async_generate_content_stream: {MODEL}"
),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__generate_content_stream_called_inside_another_tracked_function__generations_started_after_the_parent_span_closed__llm_span_attached_to_a_parent_function_span(
fake_backend,
):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
@opik.track
def f():
stream = client.models.generate_content_stream(
model=MODEL,
contents="What is the capital of Belarus?",
)
return stream
stream = f()
for _ in stream:
pass
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name="f",
input={},
output=ANY_BUT_NONE, # tracked generator
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
name="f",
input={},
output=ANY_BUT_NONE, # tracked generator
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(
f"generate_content_stream: {MODEL}"
),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__async_generate_content_stream_called_inside_another_tracked_function__generations_started_after_the_parent_span_closed__llm_span_has_a_separate_trace(
fake_backend,
):
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client)
@opik.track
async def f():
stream = await client.aio.models.generate_content_stream(
model=MODEL,
contents="What is the capital of Belarus?",
)
return stream
async def stream_outside_of_parent_function_example():
stream = await f()
async for _ in stream:
pass
asyncio.run(stream_outside_of_parent_function_example())
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name="f",
input={},
output=ANY_BUT_NONE, # tracked generator
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
name="f",
input={},
output=ANY_BUT_NONE, # tracked generator
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(
f"async_generate_content_stream: {MODEL}"
),
input={"contents": "What is the capital of Belarus?"},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
@pytest.mark.parametrize(
"project_name, expected_project_name",
[
(None, OPIK_PROJECT_DEFAULT_NAME),
("genai-integration-test", "genai-integration-test"),
],
)
def test_genai_client__generate_content__opik_args__happyflow(
fake_backend, project_name, expected_project_name
):
# test that opik_args are passed to the logged traces and spans
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client, project_name=project_name)
args_dict = {
"span": {"tags": ["span_tag"], "metadata": {"span_key": "span_value"}},
"trace": {
"thread_id": "conversation-2",
"tags": ["trace_tag"],
"metadata": {"trace_key": "trace_value"},
},
}
client.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
config=GenerateContentConfig(max_output_tokens=10),
opik_args=args_dict,
)
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?", "config": ANY_BUT_NONE},
output={"candidates": ANY_LIST},
tags=["genai", "span_tag", "trace_tag"],
metadata=ANY_DICT.containing({"trace_key": "trace_value"}),
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
project_name=expected_project_name,
thread_id="conversation-2",
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={
"contents": "What is the capital of Belarus?",
"config": ANY_BUT_NONE,
},
output={"candidates": ANY_LIST},
tags=["genai", "span_tag"],
metadata=ANY_DICT.containing({"span_key": "span_value"}),
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
project_name=expected_project_name,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
trace_tree = fake_backend.trace_trees[0]
assert_equal(EXPECTED_TRACE_TREE, trace_tree)
llm_span_metadata = trace_tree.spans[0].metadata
_assert_metadata_contains_required_keys(llm_span_metadata)
def test_genai_client__generate_content__cost_callback__sets_span_total_cost(
fake_backend,
):
CUSTOM_COST = 0.042
def cost_callback(output):
return CUSTOM_COST
client = genai.Client(
vertexai=True,
http_options=HttpOptions(api_version="v1"),
)
client = track_genai(client, cost_callback=cost_callback)
client.models.generate_content(
model=MODEL,
contents="What is the capital of Belarus?",
config=GenerateContentConfig(max_output_tokens=10),
)
opik.flush_tracker()
EXPECTED_TRACE_TREE = TraceModel(
id=ANY_BUT_NONE,
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={"contents": "What is the capital of Belarus?", "config": ANY_BUT_NONE},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
last_updated_at=ANY_BUT_NONE,
spans=[
SpanModel(
id=ANY_BUT_NONE,
type="llm",
name=ANY_STRING.starting_with(f"generate_content: {MODEL}"),
input={
"contents": "What is the capital of Belarus?",
"config": ANY_BUT_NONE,
},
output={"candidates": ANY_LIST},
tags=["genai"],
metadata=ANY_DICT,
start_time=ANY_BUT_NONE,
end_time=ANY_BUT_NONE,
usage=EXPECTED_GOOGLE_USAGE_LOGGED_FORMAT,
spans=[],
model=ANY_STRING.starting_with(MODEL),
provider="google_vertexai",
total_cost=CUSTOM_COST,
source="sdk",
)
],
source="sdk",
)
assert len(fake_backend.trace_trees) == 1
assert_equal(EXPECTED_TRACE_TREE, fake_backend.trace_trees[0])