import random
import pytest
from ai.analysis.model.assistant_chat_model import AssistantChatModel
from ai.analysis.money.money import Money
from ai.analysis.run.analysis_result import AnalysisResult
from ai.assistant.assistant import Assistant, AssistantName
from ai.assistant.model.model_provider import ModelProvider
default = AnalysisResult.empty()
[docs]
@pytest.fixture
def some_assistant():
model_provider = ModelProvider(
name="OpenAI", base_url="https://api.openai.com", api_key="sk_test_1234567890"
)
return Assistant(
name=AssistantName.Email_Answer,
base_prompt="This is a test prompt",
ai_model=AssistantChatModel.GPT_4O_MINI,
ai_model_provider=model_provider
)
[docs]
def test_total_cost():
analysis_result = default.update(prompts_cost=Money(5), completions_cost=Money(3))
assert analysis_result.total_cost.amount == 8
[docs]
def test_get_share():
analysis_result = default.update(prompts_cost=Money(5), completions_cost=Money(5))
global_total_cost = Money(20)
share = analysis_result.get_share(global_total_cost)
assert share == 0.5
[docs]
def test_get_cost_per_thousand_tickets():
analysis_result = default.update(prompts_cost=Money(5), completions_cost=Money(3))
cost = analysis_result.get_cost_per_thousand_tickets(number_tickets=10_000)
assert cost.amount == 0.8
[docs]
def test_addition(some_assistant):
analysis_result_1 = default.update(
assistant=some_assistant,
prompt_tokens=100,
prompts_cost=Money(5),
completion_tokens=50,
completions_cost=Money(3),
)
analysis_result_2 = default.update(
prompt_tokens=200,
prompts_cost=Money(10),
completion_tokens=100,
completions_cost=Money(6),
)
result = analysis_result_1 + analysis_result_2
assert result == default.update(
assistant=some_assistant,
prompt_tokens=300,
prompts_cost=Money(15),
completion_tokens=150,
completions_cost=Money(9),
)
[docs]
def test_sum_assistant_analysis(some_assistant):
# Arrange
analysis_result_1 = AnalysisResult(
assistant=some_assistant,
model="gpt-4o",
prompt_tokens=100,
prompts_cost=Money(5),
completion_tokens=50,
completions_cost=Money(3),
)
analysis_result_2 = default.update(
prompt_tokens=200,
prompts_cost=Money(10),
completion_tokens=100,
completions_cost=Money(6),
)
# Act
result = sum([analysis_result_1, analysis_result_2], AnalysisResult.empty())
# Assert
assert result == default.update(
assistant=some_assistant,
model="gpt-4o",
prompt_tokens=300,
prompts_cost=Money(15),
completion_tokens=150,
completions_cost=Money(9),
)