braintrust-sdk-python/py/src/braintrust/integrations/google_genai/integration.py at main · 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
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
"""Google GenAI integration — orchestration class and setup entry-point."""

import logging

from braintrust.integrations.base import BaseIntegration

from .patchers import (
    AsyncGAOSInteractionsCancelPatcher,
    AsyncGAOSInteractionsCreatePatcher,
    AsyncGAOSInteractionsDeletePatcher,
    AsyncGAOSInteractionsGetPatcher,
    AsyncInteractionsCancelPatcher,
    AsyncInteractionsCreatePatcher,
    AsyncInteractionsDeletePatcher,
    AsyncInteractionsGetPatcher,
    AsyncModelsEmbedContentPatcher,
    AsyncModelsGenerateContentPatcher,
    AsyncModelsGenerateContentStreamPatcher,
    AsyncModelsGenerateImagesPatcher,
    GAOSInteractionsCancelPatcher,
    GAOSInteractionsCreatePatcher,
    GAOSInteractionsDeletePatcher,
    GAOSInteractionsGetPatcher,
    InteractionsCancelPatcher,
    InteractionsCreatePatcher,
    InteractionsDeletePatcher,
    InteractionsGetPatcher,
    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",)
    min_version = "1.30.0"
    patchers = (
        ModelsGenerateContentPatcher,
        ModelsGenerateContentStreamPatcher,
        ModelsEmbedContentPatcher,
        ModelsGenerateImagesPatcher,
        InteractionsCreatePatcher,
        InteractionsGetPatcher,
        InteractionsCancelPatcher,
        InteractionsDeletePatcher,
        GAOSInteractionsCreatePatcher,
        GAOSInteractionsGetPatcher,
        GAOSInteractionsCancelPatcher,
        GAOSInteractionsDeletePatcher,
        AsyncModelsGenerateContentPatcher,
        AsyncModelsGenerateContentStreamPatcher,
        AsyncModelsEmbedContentPatcher,
        AsyncModelsGenerateImagesPatcher,
        AsyncInteractionsCreatePatcher,
        AsyncInteractionsGetPatcher,
        AsyncInteractionsCancelPatcher,
        AsyncInteractionsDeletePatcher,
        AsyncGAOSInteractionsCreatePatcher,
        AsyncGAOSInteractionsGetPatcher,
        AsyncGAOSInteractionsCancelPatcher,
        AsyncGAOSInteractionsDeletePatcher,
    )