From bb2b5d6ae820c748348ef8041eb837802fbf0231 Mon Sep 17 00:00:00 2001 From: luijait Date: Mon, 19 May 2025 13:39:03 +0200 Subject: [PATCH] Pricing in streaming --- src/cai/sdk/agents/models/openai_chatcompletions.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 2d39e536..7b6c0900 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -14,7 +14,7 @@ import asyncio from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Literal, cast, overload -from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost +from cai.util import get_ollama_api_base, fix_message_list, cli_print_agent_messages, create_agent_streaming_context, update_agent_streaming_content, finish_agent_streaming, calculate_model_cost, COST_TRACKER from cai.util import start_idle_timer, stop_idle_timer, start_active_timer, stop_active_timer from wasabi import color from cai.sdk.agents.run_to_jsonl import get_session_recorder @@ -1291,6 +1291,10 @@ class OpenAIChatCompletionsModel(Model): interaction_cost = float(interaction_cost if interaction_cost is not None else 0.0) total_cost = float(total_cost if total_cost is not None else 0.0) + # Update the global COST_TRACKER with the cost of this specific interaction + if hasattr(COST_TRACKER, "add_interaction_cost") and interaction_cost > 0.0: + COST_TRACKER.add_interaction_cost(interaction_cost) + # Store the total cost for future recording self.total_cost = total_cost