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
792 lines
25 KiB
Python
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])
|