braintrust-sdk-python/py/src/braintrust/integrations/google_genai/integration.py at 273e201274ea4096436e07743e449bb4410816a5 · braintrustdata/braintrust-sdk-python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
"""Google GenAI integration — orchestration class and setup entry-point."""

import logging

from braintrust.integrations.base import BaseIntegration

from .patchers import (
    AsyncModelsEmbedContentPatcher,
    AsyncModelsGenerateContentPatcher,
    AsyncModelsGenerateContentStreamPatcher,
    AsyncModelsGenerateImagesPatcher,
    ModelsEmbedContentPatcher,
    ModelsGenerateContentPatcher,
    ModelsGenerateContentStreamPatcher,
    ModelsGenerateImagesPatcher,
)


logger = logging.getLogger(__name__)


class GoogleGenAIIntegration(BaseIntegration):
    """Braintrust instrumentation for the Google GenAI Python SDK."""

    name = "google_genai"
    import_names = ("google.genai",)
    patchers = (
        ModelsGenerateContentPatcher,
        ModelsGenerateContentStreamPatcher,
        ModelsEmbedContentPatcher,
        ModelsGenerateImagesPatcher,
        AsyncModelsGenerateContentPatcher,
        AsyncModelsGenerateContentStreamPatcher,
        AsyncModelsEmbedContentPatcher,
        AsyncModelsGenerateImagesPatcher,
    )