Skip to content

Bedrock Usage Tracker

llmfy.llmfy_core.models.google.googleai_usage

track_googleai_usage(func)

Decorator to wrap __call_googleai calls on GoogleAIModel.

Source code in llmfy/llmfy_core/models/google/googleai_usage.py 
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def track_googleai_usage(func):
    """Decorator to wrap `__call_googleai` calls on `GoogleAIModel`."""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        response = func(*args, **kwargs)
        usage_tracker = LLMFY_USAGE_TRACKER_VAR.get()
        if usage_tracker is None:
            return response
        model = args[0]["model"]  # args is tuple[params] and params contain `model`
        if response.usage_metadata:
            usage = _extract_usage(response.usage_metadata)
            usage_tracker.update(
                provider=ServiceProvider.GOOGLE,
                type=ServiceType.LLM,
                model=model,
                usage=usage,
            )
        return response

    return wrapper

track_googleai_stream_usage(func)

Decorator to wrap __call_stream_googleai calls on GoogleAIModel.

Source code in llmfy/llmfy_core/models/google/googleai_usage.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def track_googleai_stream_usage(func):
    """Decorator to wrap `__call_stream_googleai` calls on `GoogleAIModel`."""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        stream_origin = func(*args, **kwargs)
        usage_tracker = LLMFY_USAGE_TRACKER_VAR.get()
        if usage_tracker is None:
            return stream_origin
        model = args[0]["model"]  # args is tuple[params] and params contain `model`

        stream_usage = None

        if stream_origin:
            stream, stream_copy = itertools.tee(
                stream_origin
            )  # Duplicate the generator
            stream_origin = stream  # Replace original stream

            for chunk in stream_copy:  # Iterate over the copy
                if chunk.usage_metadata and chunk.usage_metadata.prompt_token_count:
                    stream_usage = _extract_usage(chunk.usage_metadata)
                    break  # No need to iterate further

        if stream_usage:
            usage_tracker.update(
                provider=ServiceProvider.GOOGLE,
                type=ServiceType.LLM,
                model=model,
                usage=stream_usage,
            )
        return stream_origin

    return wrapper

track_googleai_embedding_usage(func)

Decorator to wrap __call_googleai_embedding calls on GoogleAIEmbedding.

The wrapped function must accept (model, contents, client) so the decorator can call count_tokens before embedding, since EmbedContentResponse carries no usage metadata.

Source code in llmfy/llmfy_core/models/google/googleai_usage.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def track_googleai_embedding_usage(func):
    """Decorator to wrap `__call_googleai_embedding` calls on `GoogleAIEmbedding`.

    The wrapped function must accept (model, contents, client) so the decorator
    can call count_tokens before embedding, since EmbedContentResponse carries
    no usage metadata.
    """

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        response = func(*args, **kwargs)
        usage_tracker = LLMFY_USAGE_TRACKER_VAR.get()
        if usage_tracker is None:
            return response
        model = args[0]
        contents = args[1]
        client = args[2]

        token_count = 0
        try:
            count_response = client.models.count_tokens(
                model=model,
                contents=contents,
            )
            token_count = count_response.total_tokens or 0
        except Exception:
            pass

        usage_tracker.update(
            provider=ServiceProvider.GOOGLE,
            type=ServiceType.EMBEDDING,
            model=model,
            usage={"prompt_token_count": token_count},
        )
        return response

    return wrapper