import dataclasses
import logging
from typing import Self
from openai.types import ChatModel
from ai.analysis.money.currency import Currency
from ai.analysis.money.money import Money
from ai.assistant.assistant import Assistant
[docs]
class AssistantsDontMatchError(Exception):
"""Raised when two AnalysisResult instances have conflicting assistant values."""
pass
[docs]
@dataclasses.dataclass
class AnalysisResult:
"""Contains cost analysis data for a single assistant run.
Attributes:
model (str | None): The name of the model used.
assistant (Assistant | None): The assistant instance used.
prompt_tokens (int): Number of tokens in the prompt.
prompts_cost (Money): Cost incurred by the prompt tokens.
completion_tokens (int): Number of tokens in the completion.
completions_cost (Money): Cost incurred by the completion tokens.
"""
model: str | None
assistant: Assistant | None
prompt_tokens: int
prompts_cost: Money
completion_tokens: int
completions_cost: Money
def __post_init__(self):
"""Initialize internal logger."""
self._logger = logging.getLogger(__name__)
[docs]
@classmethod
def empty(cls) -> Self:
"""Create an empty AnalysisResult with zeroed costs and token counts.
Returns:
AnalysisResult: An instance with no assistant, no model, and zero costs/tokens.
"""
return AnalysisResult(
assistant=None,
model=None,
prompt_tokens=0,
prompts_cost=Money(0),
completion_tokens=0,
completions_cost=Money(0),
)
[docs]
def update(
self,
assistant: Assistant = None,
prompt_tokens: int = None,
model: ChatModel = None,
prompts_cost: Money = None,
completion_tokens: int = None,
completions_cost: Money = None,
) -> Self:
"""Return a new AnalysisResult with updated fields.
Any parameter left as None will retain its current value.
Args:
assistant (Assistant, optional): New assistant instance.
prompt_tokens (int, optional): New prompt token count.
model (ChatModel, optional): New model name.
prompts_cost (Money, optional): New prompts cost.
completion_tokens (int, optional): New completion token count.
completions_cost (Money, optional): New completions cost.
Returns:
AnalysisResult: A new instance reflecting the updates.
"""
return AnalysisResult(
assistant=assistant if assistant is not None else self.assistant,
model=model if model is not None else self.model,
prompt_tokens=prompt_tokens if prompt_tokens is not None else self.prompt_tokens,
prompts_cost=prompts_cost if prompts_cost is not None else self.prompts_cost,
completion_tokens=completion_tokens if completion_tokens is not None else self.completion_tokens,
completions_cost=completions_cost if completions_cost is not None else self.completions_cost,
)
@property
def total_cost(self) -> Money:
"""Compute the total cost (prompts + completions).
Returns:
Money: Sum of prompts_cost and completions_cost.
"""
return self.prompts_cost + self.completions_cost
[docs]
def convert_to(self, currency: Currency) -> Self:
"""Convert both prompt and completion costs to a target currency.
Args:
currency (Currency): The currency to convert all costs into.
Returns:
AnalysisResult: New instance with costs converted.
"""
return AnalysisResult(
assistant=self.assistant,
model=self.model,
prompt_tokens=self.prompt_tokens,
prompts_cost=self.prompts_cost.convert_to(currency),
completion_tokens=self.completion_tokens,
completions_cost=self.completions_cost.convert_to(currency),
)
[docs]
def get_share(self, total_cost: Money) -> float:
"""Calculate this result's share of a total cost.
Args:
total_cost (Money): The total cost against which to compare.
Returns:
float: Fraction of total_cost represented by this result.
"""
return self.total_cost.amount / total_cost.amount
[docs]
def get_cost_per_thousand_tickets(self, number_tickets: int) -> Money:
"""Scale total cost to a per-1000-tickets basis.
Args:
number_tickets (int): Number of tickets over which the total_cost was incurred.
Returns:
Money: Total cost scaled to 1000 tickets.
"""
return self.total_cost * 1_000 / number_tickets
def _validate_and_get_single_value(
self, attribute_name: str, *values
):
"""Ensure provided values are either identical or None and return the single non-None.
Args:
attribute_name (str): Name used for logging in case of mismatch.
*values: Values to validate.
Returns:
Any: The single non-None value, or None if all are None.
Logs:
Debug message if conflicting non-None values are found.
"""
unique_values = {value for value in values if value is not None}
if len(unique_values) > 1:
self._logger.debug(
f"{attribute_name.capitalize()}s should match for addition. "
f"The values are: {list(unique_values)}"
)
return unique_values.pop() if unique_values else None
def __add__(self, other: Self) -> Self:
"""Combine two AnalysisResult instances by summing token counts and costs.
Ensures model and assistant match or logs a debug message if they conflict.
Args:
other (AnalysisResult): Another result to add.
Returns:
AnalysisResult: A new instance with aggregated values.
"""
new_assistant = self._validate_and_get_single_value(
"assistant", self.assistant, other.assistant
)
new_model = self._validate_and_get_single_value(
"model", self.model, other.model
)
return AnalysisResult(
assistant=new_assistant,
model=new_model,
prompt_tokens=self.prompt_tokens + other.prompt_tokens,
prompts_cost=self.prompts_cost + other.prompts_cost,
completion_tokens=self.completion_tokens + other.completion_tokens,
completions_cost=self.completions_cost + other.completions_cost,
)
def __radd__(self, other: Self) -> Self:
"""Support sum() by handling reversed addition."""
return self + other