import logging
from openai.types.chat import ChatCompletion
from ai.analysis.model.token_price_calculator import TokenPriceCalculator
from ai.analysis.model.token_type import TokenType
from ai.analysis.money.money import Money
from ai.analysis.run.abstract_analyzer import AbstractAnalyzer
from ai.analysis.run.analysis_result import AnalysisResult
from ai.assistant.assistant import Assistant
[docs]
class AssistantsDontMatchError(Exception):
"""Raised when two assistants in combined analyses do not match."""
pass
[docs]
class AssistantRun(AbstractAnalyzer):
"""Analyze a single assistant run to compute token usage costs.
Attributes:
_completion_result (ChatCompletion): The raw OpenAI chat completion.
assistant (Assistant): The assistant used for this run.
_model_price_calculator (TokenPriceCalculator): Computes per-token costs.
_logger (logging.Logger): Logger for warnings and debug messages.
"""
[docs]
def __init__(self, completion_result: ChatCompletion, assistant: Assistant):
"""Initialize the analysis for an assistant run.
Args:
completion_result (ChatCompletion): The result of the chat completion.
assistant (Assistant): The assistant instance used.
Raises:
ValueError: If `completion_result` is not a ChatCompletion or has no model.
"""
if not isinstance(completion_result, ChatCompletion):
raise ValueError(f"Invalid run object: {completion_result}")
if completion_result.model is None:
raise ValueError("Model not found in completion result")
self._completion_result = completion_result
self.assistant = assistant
self._model_price_calculator = TokenPriceCalculator()
self._logger = logging.getLogger(__name__)
[docs]
def get_cost_analysis(self) -> AnalysisResult:
"""Generate cost analysis from the completion result.
Returns:
AnalysisResult: Includes token counts and cost breakdown.
"""
return AnalysisResult(
assistant=self.assistant,
model=self._model,
prompt_tokens=self._prompt_tokens,
prompts_cost=self._input_cost,
completion_tokens=self._completion_tokens,
completions_cost=self._output_cost,
)
@property
def _prompt_tokens(self) -> int:
"""Count of input tokens used.
Returns:
int: Number of prompt tokens, or 0 if unavailable.
"""
usage = self._completion_result.usage
if usage is None:
self._logger.warning("No usage data found for completion")
return 0
return usage.prompt_tokens
@property
def _completion_tokens(self) -> int:
"""Count of output tokens generated.
Returns:
int: Number of completion tokens, or 0 if unavailable.
"""
usage = self._completion_result.usage
if usage is None:
self._logger.warning("No usage data found for completion")
return 0
return usage.completion_tokens
@property
def _model(self) -> str:
"""Model identifier used for the completion.
Returns:
str: Model name.
"""
return self._completion_result.model
@property
def _input_cost(self) -> Money:
"""Cost incurred by input tokens.
Returns:
Money: Cost for all prompt tokens.
"""
return self._get_cost_per_token(TokenType.INPUT) * self._prompt_tokens
@property
def _output_cost(self) -> Money:
"""Cost incurred by output tokens.
Returns:
Money: Cost for all completion tokens.
"""
return self._get_cost_per_token(TokenType.OUTPUT) * self._completion_tokens
@property
def _total_cost(self) -> Money:
"""Combined cost of input and output tokens.
Returns:
Money: Sum of input and output costs.
"""
return self._input_cost + self._output_cost
def _get_cost_per_token(self, cost_type: TokenType) -> Money:
"""Retrieve per-token cost for a given token type.
Args:
cost_type (TokenType): INPUT or OUTPUT token type.
Returns:
Money: Cost per single token of the specified type.
"""
return self._model_price_calculator.get_cost(self.assistant.ai_model, cost_type)
[docs]
def get_assistant_updated(self, new_assistant: Assistant) -> "AssistantRun":
"""Return a copy of this run using a different assistant.
Args:
new_assistant (Assistant): The assistant to replace in the new run.
Returns:
AssistantRun: New instance with the same completion result but updated assistant.
"""
return AssistantRun(completion_result=self._completion_result, assistant=new_assistant)