Skip to content

llmfy_usage_tracker

llmfy.llmfy_core.usage.usage_tracker

llmfy_usage_tracker(openai_pricing=None, bedrock_pricing=None)

LLMfy usage tracker.

Use this to track token usage all provider.

Example: 

with llmfy_usage_tracker() as usage:
        result = llm.generate(messages)
        ...
        print(usage)

Parameters:

Name Type Description Default
openai_pricing Optional[Dict[str, Any]]

OpenAI Pricing dictionary source. Defaults to None. If None then use default pricing from this dependency.

Example pricing structure: 

{
    "gpt-4o": {
        "input": 2.50,
        "output": 10.00
    },
    "gpt-4o-mini": {
        "input": 0.15,
        "output": 0.60
    },
    "gpt-3.5-turbo": {
        "input": 0.05,
        "output": 1.50
    }
}

None
bedrock_pricing Optional[Dict[str, Any]]

Bedrock Pricing dictionary source. Defaults to None. If None then use default pricing from this dependency.

Example pricing structure: { "anthropic.claude-3-5-sonnet-20240620-v1:0": { "us-east-1": { "region": "US East (N. Virginia)", "input": 0.003, "output": 0.015, }, "us-west-2": { "region": "US West (Oregon)", "input": 0.003, "output": 0.015, }, }, "anthropic.claude-3-5-sonnet-20241022-v2:0": { "us-east-1": { "region": "US East (N. Virginia)", "input": 0.003, "output": 0.015, }, "us-west-2": { "region": "US West (Oregon)", "input": 0.003, "output": 0.015, }, } }

 None

Yields:

Name Type Description
OpenAIUsage

OpenAI usage accumulation.
Source code in llmfy/llmfy_core/usage/usage_tracker.py 
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
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
@contextmanager
def llmfy_usage_tracker(
    openai_pricing: Optional[Dict[str, Any]] = None,
    bedrock_pricing: Optional[Dict[str, Any]] = None,
):
    """LLMfy usage tracker.

    Use this to track token usage all provider.

    Example:
    ```python
    with llmfy_usage_tracker() as usage:
            result = llm.generate(messages)
            ...
            print(usage)
    ```

    Args:
        openai_pricing (Optional[Dict[str, Any]], optional): OpenAI Pricing dictionary source. Defaults to None.
            If None then use default pricing from this dependency.

            Example pricing structure:
            ```
            {
                "gpt-4o": {
                    "input": 2.50,
                    "output": 10.00
                },
                "gpt-4o-mini": {
                    "input": 0.15,
                    "output": 0.60
                },
                "gpt-3.5-turbo": {
                    "input": 0.05,
                    "output": 1.50
                }
            }
            ```

        bedrock_pricing (Optional[Dict[str, Any]], optional): Bedrock Pricing dictionary source. Defaults to None.
            If None then use default pricing from this dependency.

            Example pricing structure:
             ```
            {
                "anthropic.claude-3-5-sonnet-20240620-v1:0": {
                    "us-east-1": {
                        "region": "US East (N. Virginia)",
                        "input": 0.003,
                        "output": 0.015,
                    },
                    "us-west-2": {
                        "region": "US West (Oregon)",
                        "input": 0.003,
                        "output": 0.015,
                    },
                },
                "anthropic.claude-3-5-sonnet-20241022-v2:0": {
                    "us-east-1": {
                        "region": "US East (N. Virginia)",
                        "input": 0.003,
                        "output": 0.015,
                    },
                    "us-west-2": {
                        "region": "US West (Oregon)",
                        "input": 0.003,
                        "output": 0.015,
                    },
                }
            }
            ```

    Yields:
            OpenAIUsage: OpenAI usage accumulation.
    """

    usage_tracker = LLMfyUsage(
        openai_pricing=openai_pricing,
        bedrock_pricing=bedrock_pricing,
    )
    LLMFY_USAGE_TRACKER_VAR.set(usage_tracker)  # Store usage_tracker it in the context
    try:
        yield usage_tracker  # Expose tracker to the context
    finally:
        pass

track_bedrock_embedding_usage(func)

Decorator to wrap __call_bedrock_embedding calls on BedrockEmbedding.

Source code in llmfy/llmfy_core/usage/usage_tracker.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
def track_bedrock_embedding_usage(func):
    """Decorator to wrap `__call_bedrock_embedding` calls on `BedrockEmbedding`."""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        response = func(*args, **kwargs)
        model = args[0]
        # Extract token usage from headers
        headers = response.get("ResponseMetadata", {}).get("HTTPHeaders", {})
        input_tokens = int(headers.get("x-amzn-bedrock-input-token-count", 0))
        usage = {"x-amzn-bedrock-input-token-count": input_tokens}
        usage_tracker = LLMFY_USAGE_TRACKER_VAR.get()
        usage_tracker.update(
            provider=ServiceProvider.BEDROCK,
            type=ServiceType.EMBEDDING,
            model=model,
            usage=usage,
        )

        return response

    return wrapper

track_bedrock_usage(func)

Decorator to wrap __call_bedrock calls on BedrockModel.

Source code in llmfy/llmfy_core/usage/usage_tracker.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
def track_bedrock_usage(func):
    """Decorator to wrap `__call_bedrock` calls on `BedrockModel`."""

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

    return wrapper

track_openai_embedding_usage(func)

Decorator to wrap __call_openai_embedding calls on OpenAIEmbedding.

Source code in llmfy/llmfy_core/usage/usage_tracker.py 
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def track_openai_embedding_usage(func):
    """Decorator to wrap `__call_openai_embedding` calls on `OpenAIEmbedding`."""

    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        response = func(*args, **kwargs)
        model = args[0]
        # Extract token usage
        # input_tokens = response.usage.prompt_tokens if response.usage else 0
        # Usage(prompt_tokens=8, total_tokens=8)
        usage = {
            "prompt_tokens": response.usage.prompt_tokens or 0,
            "total_tokens": response.usage.total_tokens or 0,
        }
        print(usage)
        usage_tracker = LLMFY_USAGE_TRACKER_VAR.get()
        usage_tracker.update(
            provider=ServiceProvider.OPENAI,
            type=ServiceType.EMBEDDING,
            model=model,
            usage=usage,
        )

        return response

    return wrapper

track_openai_usage(func)

Decorator to wrap __call_openai calls on OpenAIModel.

Source code in llmfy/llmfy_core/usage/usage_tracker.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
def track_openai_usage(func):
    """Decorator to wrap `__call_openai` calls on `OpenAIModel`."""

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

    return wrapper