model-library 0.1.4__tar.gz → 0.1.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (126) hide show
  1. {model_library-0.1.4 → model_library-0.1.5}/PKG-INFO +1 -1
  2. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/output.py +24 -10
  3. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/utils.py +27 -5
  4. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/all_models.json +51 -1
  5. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/google_models.yaml +15 -0
  6. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/perplexity_models.yaml +2 -0
  7. {model_library-0.1.4 → model_library-0.1.5}/model_library/exceptions.py +1 -1
  8. {model_library-0.1.4 → model_library-0.1.5}/model_library/registry_utils.py +60 -0
  9. {model_library-0.1.4 → model_library-0.1.5}/model_library.egg-info/PKG-INFO +1 -1
  10. {model_library-0.1.4 → model_library-0.1.5}/model_library.egg-info/SOURCES.txt +1 -0
  11. model_library-0.1.5/tests/unit/test_result_metadata.py +206 -0
  12. {model_library-0.1.4 → model_library-0.1.5}/.gitattributes +0 -0
  13. {model_library-0.1.4 → model_library-0.1.5}/.github/workflows/publish.yml +0 -0
  14. {model_library-0.1.4 → model_library-0.1.5}/.github/workflows/style.yaml +0 -0
  15. {model_library-0.1.4 → model_library-0.1.5}/.github/workflows/test.yaml +0 -0
  16. {model_library-0.1.4 → model_library-0.1.5}/.github/workflows/typecheck.yml +0 -0
  17. {model_library-0.1.4 → model_library-0.1.5}/.gitignore +0 -0
  18. {model_library-0.1.4 → model_library-0.1.5}/LICENSE +0 -0
  19. {model_library-0.1.4 → model_library-0.1.5}/Makefile +0 -0
  20. {model_library-0.1.4 → model_library-0.1.5}/README.md +0 -0
  21. {model_library-0.1.4 → model_library-0.1.5}/examples/README.md +0 -0
  22. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/batch.py +0 -0
  23. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/custom_retrier.py +0 -0
  24. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/deep_research.py +0 -0
  25. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/stress.py +0 -0
  26. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/structured_output.py +0 -0
  27. {model_library-0.1.4 → model_library-0.1.5}/examples/advanced/web_search.py +0 -0
  28. {model_library-0.1.4 → model_library-0.1.5}/examples/basics.py +0 -0
  29. {model_library-0.1.4 → model_library-0.1.5}/examples/data/files.py +0 -0
  30. {model_library-0.1.4 → model_library-0.1.5}/examples/data/images.py +0 -0
  31. {model_library-0.1.4 → model_library-0.1.5}/examples/embeddings.py +0 -0
  32. {model_library-0.1.4 → model_library-0.1.5}/examples/files.py +0 -0
  33. {model_library-0.1.4 → model_library-0.1.5}/examples/images.py +0 -0
  34. {model_library-0.1.4 → model_library-0.1.5}/examples/prompt_caching.py +0 -0
  35. {model_library-0.1.4 → model_library-0.1.5}/examples/setup.py +0 -0
  36. {model_library-0.1.4 → model_library-0.1.5}/examples/tool_calls.py +0 -0
  37. {model_library-0.1.4 → model_library-0.1.5}/model_library/__init__.py +0 -0
  38. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/__init__.py +0 -0
  39. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/base.py +0 -0
  40. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/batch.py +0 -0
  41. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/delegate_only.py +0 -0
  42. {model_library-0.1.4 → model_library-0.1.5}/model_library/base/input.py +0 -0
  43. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/README.md +0 -0
  44. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/ai21labs_models.yaml +0 -0
  45. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/alibaba_models.yaml +0 -0
  46. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/amazon_models.yaml +0 -0
  47. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/anthropic_models.yaml +0 -0
  48. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/cohere_models.yaml +0 -0
  49. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/deepseek_models.yaml +0 -0
  50. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/dummy_model.yaml +0 -0
  51. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/fireworks_models.yaml +0 -0
  52. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/inception_models.yaml +0 -0
  53. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/kimi_models.yaml +0 -0
  54. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/minimax_models.yaml +0 -0
  55. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/mistral_models.yaml +0 -0
  56. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/openai_models.yaml +0 -0
  57. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/together_models.yaml +0 -0
  58. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/xai_models.yaml +0 -0
  59. {model_library-0.1.4 → model_library-0.1.5}/model_library/config/zai_models.yaml +0 -0
  60. {model_library-0.1.4 → model_library-0.1.5}/model_library/file_utils.py +0 -0
  61. {model_library-0.1.4 → model_library-0.1.5}/model_library/logging.py +0 -0
  62. {model_library-0.1.4 → model_library-0.1.5}/model_library/model_utils.py +0 -0
  63. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/__init__.py +0 -0
  64. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/ai21labs.py +0 -0
  65. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/alibaba.py +0 -0
  66. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/amazon.py +0 -0
  67. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/anthropic.py +0 -0
  68. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/azure.py +0 -0
  69. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/cohere.py +0 -0
  70. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/deepseek.py +0 -0
  71. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/fireworks.py +0 -0
  72. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/google/__init__.py +0 -0
  73. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/google/batch.py +0 -0
  74. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/google/google.py +0 -0
  75. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/inception.py +0 -0
  76. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/kimi.py +0 -0
  77. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/minimax.py +0 -0
  78. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/mistral.py +0 -0
  79. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/openai.py +0 -0
  80. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/perplexity.py +0 -0
  81. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/together.py +0 -0
  82. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/vals.py +0 -0
  83. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/xai.py +0 -0
  84. {model_library-0.1.4 → model_library-0.1.5}/model_library/providers/zai.py +0 -0
  85. {model_library-0.1.4 → model_library-0.1.5}/model_library/py.typed +0 -0
  86. {model_library-0.1.4 → model_library-0.1.5}/model_library/register_models.py +0 -0
  87. {model_library-0.1.4 → model_library-0.1.5}/model_library/settings.py +0 -0
  88. {model_library-0.1.4 → model_library-0.1.5}/model_library/utils.py +0 -0
  89. {model_library-0.1.4 → model_library-0.1.5}/model_library.egg-info/dependency_links.txt +0 -0
  90. {model_library-0.1.4 → model_library-0.1.5}/model_library.egg-info/requires.txt +0 -0
  91. {model_library-0.1.4 → model_library-0.1.5}/model_library.egg-info/top_level.txt +0 -0
  92. {model_library-0.1.4 → model_library-0.1.5}/pyproject.toml +0 -0
  93. {model_library-0.1.4 → model_library-0.1.5}/scripts/browse_models.py +0 -0
  94. {model_library-0.1.4 → model_library-0.1.5}/scripts/config.py +0 -0
  95. {model_library-0.1.4 → model_library-0.1.5}/scripts/publish.py +0 -0
  96. {model_library-0.1.4 → model_library-0.1.5}/scripts/run_models.py +0 -0
  97. {model_library-0.1.4 → model_library-0.1.5}/setup.cfg +0 -0
  98. {model_library-0.1.4 → model_library-0.1.5}/tests/README.md +0 -0
  99. {model_library-0.1.4 → model_library-0.1.5}/tests/__init__.py +0 -0
  100. {model_library-0.1.4 → model_library-0.1.5}/tests/conftest.py +0 -0
  101. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/__init__.py +0 -0
  102. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/conftest.py +0 -0
  103. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_batch.py +0 -0
  104. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_completion.py +0 -0
  105. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_files.py +0 -0
  106. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_reasoning.py +0 -0
  107. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_retry.py +0 -0
  108. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_streaming.py +0 -0
  109. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_structured_output.py +0 -0
  110. {model_library-0.1.4 → model_library-0.1.5}/tests/integration/test_tools.py +0 -0
  111. {model_library-0.1.4 → model_library-0.1.5}/tests/test_helpers.py +0 -0
  112. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/__init__.py +0 -0
  113. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/conftest.py +0 -0
  114. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/providers/__init__.py +0 -0
  115. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/providers/test_fireworks_provider.py +0 -0
  116. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/providers/test_google_provider.py +0 -0
  117. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_batch.py +0 -0
  118. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_context_window.py +0 -0
  119. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_deep_research.py +0 -0
  120. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_perplexity_provider.py +0 -0
  121. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_prompt_caching.py +0 -0
  122. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_registry.py +0 -0
  123. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_retry.py +0 -0
  124. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_streaming.py +0 -0
  125. {model_library-0.1.4 → model_library-0.1.5}/tests/unit/test_tools.py +0 -0
  126. {model_library-0.1.4 → model_library-0.1.5}/uv.lock +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: model-library
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Summary: Model Library for vals.ai
5
5
  Author-email: "Vals AI, Inc." <contact@vals.ai>
6
6
  License: MIT
@@ -9,9 +9,7 @@ from pydantic import BaseModel, Field, computed_field, field_validator
9
9
  from typing_extensions import override
10
10
 
11
11
  from model_library.base.input import InputItem, ToolCall
12
- from model_library.base.utils import (
13
- sum_optional,
14
- )
12
+ from model_library.base.utils import add_optional
15
13
  from model_library.utils import truncate_str
16
14
 
17
15
 
@@ -42,10 +40,14 @@ class QueryResultCost(BaseModel):
42
40
  reasoning: float | None = None
43
41
  cache_read: float | None = None
44
42
  cache_write: float | None = None
43
+ total_override: float | None = None
45
44
 
46
45
  @computed_field
47
46
  @property
48
47
  def total(self) -> float:
48
+ if self.total_override is not None:
49
+ return self.total_override
50
+
49
51
  return sum(
50
52
  filter(
51
53
  None,
@@ -86,6 +88,16 @@ class QueryResultCost(BaseModel):
86
88
  )
87
89
  )
88
90
 
91
+ def __add__(self, other: "QueryResultCost") -> "QueryResultCost":
92
+ return QueryResultCost(
93
+ input=self.input + other.input,
94
+ output=self.output + other.output,
95
+ reasoning=add_optional(self.reasoning, other.reasoning),
96
+ cache_read=add_optional(self.cache_read, other.cache_read),
97
+ cache_write=add_optional(self.cache_write, other.cache_write),
98
+ total_override=add_optional(self.total_override, other.total_override),
99
+ )
100
+
89
101
  @override
90
102
  def __repr__(self):
91
103
  use_cents = self.total < 1
@@ -150,18 +162,20 @@ class QueryResultMetadata(BaseModel):
150
162
  return QueryResultMetadata(
151
163
  in_tokens=self.in_tokens + other.in_tokens,
152
164
  out_tokens=self.out_tokens + other.out_tokens,
153
- reasoning_tokens=sum_optional(
154
- self.reasoning_tokens, other.reasoning_tokens
165
+ reasoning_tokens=cast(
166
+ int | None, add_optional(self.reasoning_tokens, other.reasoning_tokens)
155
167
  ),
156
- cache_read_tokens=sum_optional(
157
- self.cache_read_tokens, other.cache_read_tokens
168
+ cache_read_tokens=cast(
169
+ int | None,
170
+ add_optional(self.cache_read_tokens, other.cache_read_tokens),
158
171
  ),
159
- cache_write_tokens=sum_optional(
160
- self.cache_write_tokens, other.cache_write_tokens
172
+ cache_write_tokens=cast(
173
+ int | None,
174
+ add_optional(self.cache_write_tokens, other.cache_write_tokens),
161
175
  ),
162
176
  duration_seconds=self.default_duration_seconds
163
177
  + other.default_duration_seconds,
164
- cost=self.cost,
178
+ cost=cast(QueryResultCost | None, add_optional(self.cost, other.cost)),
165
179
  )
166
180
 
167
181
  @override
@@ -1,4 +1,4 @@
1
- from typing import Sequence, cast
1
+ from typing import Sequence, TypeVar, cast
2
2
 
3
3
  from model_library.base.input import (
4
4
  FileBase,
@@ -8,17 +8,39 @@ from model_library.base.input import (
8
8
  ToolResult,
9
9
  )
10
10
  from model_library.utils import truncate_str
11
+ from pydantic import BaseModel
11
12
 
13
+ T = TypeVar("T", bound=BaseModel)
12
14
 
13
- def sum_optional(a: int | None, b: int | None) -> int | None:
14
- """Sum two optional integers, returning None if both are None.
15
+
16
+ def add_optional(
17
+ a: int | float | T | None, b: int | float | T | None
18
+ ) -> int | float | T | None:
19
+ """Add two optional objects, returning None if both are None.
15
20
 
16
21
  Preserves None to indicate "unknown/not provided" when both inputs are None,
17
- otherwise treats None as 0 for summation.
22
+ otherwise returns the non-None value or their sum.
18
23
  """
19
24
  if a is None and b is None:
20
25
  return None
21
- return (a or 0) + (b or 0)
26
+
27
+ if a is None or b is None:
28
+ return a or b
29
+
30
+ if isinstance(a, (int, float)) and isinstance(b, (int, float)):
31
+ return a + b
32
+
33
+ # NOTE: Ensure that the subtypes are the same so we can use the __add__ method just from one
34
+ if type(a) is type(b):
35
+ add_method = getattr(a, "__add__", None)
36
+ if add_method is not None:
37
+ return add_method(b)
38
+ else:
39
+ raise ValueError(
40
+ f"Cannot add {type(a)} and {type(b)} because they are not the same subclass"
41
+ )
42
+
43
+ return None
22
44
 
23
45
 
24
46
  def get_pretty_input_types(input: Sequence["InputItem"], verbose: bool = False) -> str:
@@ -1,4 +1,54 @@
1
1
  {
2
+ "google/gemini-3-flash-preview": {
3
+ "company": "Google",
4
+ "label": "Gemini 3 Flash (12/25)",
5
+ "description": "Google's newest budget workhorse model",
6
+ "release_date": "2025-12-17",
7
+ "open_source": false,
8
+ "documentation_url": "https://ai.google.dev/gemini-api/docs/models",
9
+ "properties": {
10
+ "context_window": 1048576,
11
+ "max_tokens": 65536,
12
+ "training_cutoff": null,
13
+ "reasoning_model": true
14
+ },
15
+ "supports": {
16
+ "images": true,
17
+ "videos": true,
18
+ "files": true,
19
+ "batch": true,
20
+ "temperature": true,
21
+ "tools": true
22
+ },
23
+ "metadata": {
24
+ "deprecated": false,
25
+ "available_for_everyone": true,
26
+ "available_as_evaluator": false,
27
+ "ignored_for_cost": false
28
+ },
29
+ "provider_properties": {},
30
+ "costs_per_million_token": {
31
+ "input": 0.5,
32
+ "output": 3.0,
33
+ "cache": {
34
+ "read_discount": 0.1,
35
+ "write_markup": 1.0
36
+ },
37
+ "batch": {
38
+ "input_discount": 0.5,
39
+ "output_discount": 0.5
40
+ }
41
+ },
42
+ "alternative_keys": [],
43
+ "default_parameters": {
44
+ "temperature": 1.0,
45
+ "reasoning_effort": "high"
46
+ },
47
+ "provider_endpoint": "gemini-3-flash-preview",
48
+ "provider_name": "google",
49
+ "full_key": "google/gemini-3-flash-preview",
50
+ "slug": "google_gemini-3-flash-preview"
51
+ },
2
52
  "openai/gpt-5.2-pro-2025-12-11": {
3
53
  "company": "OpenAI",
4
54
  "label": "GPT 5.2 Pro",
@@ -15428,7 +15478,7 @@
15428
15478
  "tools": false
15429
15479
  },
15430
15480
  "metadata": {
15431
- "deprecated": false,
15481
+ "deprecated": true,
15432
15482
  "available_for_everyone": true,
15433
15483
  "available_as_evaluator": false,
15434
15484
  "ignored_for_cost": false
@@ -54,6 +54,21 @@ gemini-3-models:
54
54
  temperature: 1
55
55
  reasoning_effort: "high"
56
56
 
57
+ google/gemini-3-flash-preview:
58
+ label: Gemini 3 Flash (12/25)
59
+ description: Google's newest budget workhorse model
60
+ release_date: 2025-12-17
61
+ properties:
62
+ context_window: 1048576
63
+ max_tokens: 65536
64
+ reasoning_model: true
65
+ costs_per_million_token:
66
+ input: 0.50
67
+ output: 3.00
68
+ default_parameters:
69
+ temperature: 1
70
+ reasoning_effort: "high"
71
+
57
72
  google/gemini-3-pro-preview:
58
73
  label: Gemini 3 Pro (11/25)
59
74
  description: Gemini 3 Pro, Google's most powerful model.
@@ -46,6 +46,8 @@ perplexity-models:
46
46
  label: Sonar Reasoning
47
47
  description: Reasoning-focused search model that exposes intermediate thinking for step-by-step answers.
48
48
  documentation_url: https://docs.perplexity.ai/models/models/sonar-reasoning
49
+ metadata:
50
+ deprecated: true
49
51
  properties:
50
52
  context_window: 128000
51
53
  reasoning_model: true
@@ -183,8 +183,8 @@ RETRIABLE_EXCEPTION_CODES = [
183
183
  "server_error",
184
184
  "overloaded",
185
185
  "throttling", # AWS throttling errors
186
- "throttlingexception", # AWS throttling errors
187
186
  "internal server error",
187
+ "InternalServerError",
188
188
  ]
189
189
 
190
190
 
@@ -1,9 +1,11 @@
1
1
  from functools import cache
2
2
  from pathlib import Path
3
+ from typing import TypedDict
3
4
 
4
5
  import tiktoken
5
6
 
6
7
  from model_library.base import LLM, LLMConfig, ProviderConfig
8
+ from model_library.base.output import QueryResultCost, QueryResultMetadata
7
9
  from model_library.register_models import (
8
10
  CostProperties,
9
11
  ModelConfig,
@@ -129,6 +131,64 @@ def get_model_cost(model_str: str) -> CostProperties | None:
129
131
  return model_config.costs_per_million_token
130
132
 
131
133
 
134
+ class TokenDict(TypedDict, total=False):
135
+ """Token counts for cost calculation."""
136
+
137
+ in_tokens: int
138
+ out_tokens: int
139
+ reasoning_tokens: int | None
140
+ cache_read_tokens: int | None
141
+ cache_write_tokens: int | None
142
+
143
+
144
+ async def recompute_cost(
145
+ model_str: str,
146
+ tokens: TokenDict,
147
+ ) -> QueryResultCost:
148
+ """
149
+ Recompute the cost for a model based on token information.
150
+
151
+ Uses the model provider's existing _calculate_cost method to ensure
152
+ provider-specific cost calculations are applied.
153
+
154
+ Args:
155
+ model_str: The model identifier (e.g., "openai/gpt-4o")
156
+ tokens: Dictionary containing token counts with keys:
157
+ - in_tokens (required): Number of input tokens
158
+ - out_tokens (required): Number of output tokens
159
+ - reasoning_tokens (optional): Number of reasoning tokens
160
+ - cache_read_tokens (optional): Number of cache read tokens
161
+ - cache_write_tokens (optional): Number of cache write tokens
162
+
163
+ Returns:
164
+ QueryResultCost with computed costs
165
+
166
+ Raises:
167
+ ValueError: If required token parameters are missing
168
+ Exception: If model not found in registry or costs not configured
169
+ """
170
+ if "in_tokens" not in tokens:
171
+ raise ValueError("Token dict must contain 'in_tokens'")
172
+ if "out_tokens" not in tokens:
173
+ raise ValueError("Token dict must contain 'out_tokens'")
174
+
175
+ model = get_registry_model(model_str)
176
+
177
+ metadata = QueryResultMetadata(
178
+ in_tokens=tokens["in_tokens"],
179
+ out_tokens=tokens["out_tokens"],
180
+ reasoning_tokens=tokens.get("reasoning_tokens"),
181
+ cache_read_tokens=tokens.get("cache_read_tokens"),
182
+ cache_write_tokens=tokens.get("cache_write_tokens"),
183
+ )
184
+
185
+ cost = await model._calculate_cost(metadata) # type: ignore[arg-type]
186
+ if cost is None:
187
+ raise Exception(f"No cost information available for model {model_str}")
188
+
189
+ return cost
190
+
191
+
132
192
  @cache
133
193
  def get_provider_names() -> list[str]:
134
194
  """Return all provider names in the registry"""
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: model-library
3
- Version: 0.1.4
3
+ Version: 0.1.5
4
4
  Summary: Model Library for vals.ai
5
5
  Author-email: "Vals AI, Inc." <contact@vals.ai>
6
6
  License: MIT
@@ -115,6 +115,7 @@ tests/unit/test_deep_research.py
115
115
  tests/unit/test_perplexity_provider.py
116
116
  tests/unit/test_prompt_caching.py
117
117
  tests/unit/test_registry.py
118
+ tests/unit/test_result_metadata.py
118
119
  tests/unit/test_retry.py
119
120
  tests/unit/test_streaming.py
120
121
  tests/unit/test_tools.py
@@ -0,0 +1,206 @@
1
+ import pytest
2
+
3
+ from model_library.base.output import QueryResultCost, QueryResultMetadata
4
+
5
+
6
+ @pytest.mark.unit
7
+ class TestQueryResultCostAddition:
8
+ def test_add_full_costs(self):
9
+ cost1 = QueryResultCost(
10
+ input=0.01,
11
+ output=0.02,
12
+ reasoning=0.005,
13
+ cache_read=0.001,
14
+ cache_write=0.002,
15
+ )
16
+ cost2 = QueryResultCost(
17
+ input=0.02,
18
+ output=0.03,
19
+ reasoning=0.01,
20
+ cache_read=0.002,
21
+ cache_write=0.003,
22
+ )
23
+
24
+ result = cost1 + cost2
25
+
26
+ assert result.input == 0.03
27
+ assert result.output == 0.05
28
+ assert result.reasoning == 0.015
29
+ assert result.cache_read == 0.003
30
+ assert result.cache_write == 0.005
31
+
32
+ def test_add_costs_with_none_fields(self):
33
+ cost1 = QueryResultCost(input=0.01, output=0.02, reasoning=0.005)
34
+ cost2 = QueryResultCost(input=0.02, output=0.03, cache_read=0.002)
35
+
36
+ result = cost1 + cost2
37
+
38
+ assert result.input == 0.03
39
+ assert result.output == 0.05
40
+ assert result.reasoning == 0.005
41
+ assert result.cache_read == 0.002
42
+ assert result.cache_write is None
43
+
44
+ def test_add_costs_both_none_fields(self):
45
+ cost1 = QueryResultCost(input=0.01, output=0.02)
46
+ cost2 = QueryResultCost(input=0.02, output=0.03)
47
+
48
+ result = cost1 + cost2
49
+
50
+ assert result.input == 0.03
51
+ assert result.output == 0.05
52
+ assert result.reasoning is None
53
+ assert result.cache_read is None
54
+ assert result.cache_write is None
55
+
56
+ def test_cost_total_computed(self):
57
+ cost1 = QueryResultCost(input=0.01, output=0.02, reasoning=0.005)
58
+ cost2 = QueryResultCost(input=0.02, output=0.03, cache_read=0.001)
59
+
60
+ result = cost1 + cost2
61
+
62
+ expected_total = 0.03 + 0.05 + 0.005 + 0.001
63
+ assert abs(result.total - expected_total) < 1e-10
64
+
65
+
66
+ @pytest.mark.unit
67
+ class TestQueryResultMetadataAddition:
68
+ def test_add_full_metadata(self):
69
+ meta1 = QueryResultMetadata(
70
+ in_tokens=100,
71
+ out_tokens=50,
72
+ reasoning_tokens=20,
73
+ cache_read_tokens=10,
74
+ cache_write_tokens=5,
75
+ duration_seconds=1.5,
76
+ cost=QueryResultCost(input=0.01, output=0.02),
77
+ )
78
+ meta2 = QueryResultMetadata(
79
+ in_tokens=200,
80
+ out_tokens=100,
81
+ reasoning_tokens=30,
82
+ cache_read_tokens=15,
83
+ cache_write_tokens=10,
84
+ duration_seconds=2.0,
85
+ cost=QueryResultCost(input=0.02, output=0.03),
86
+ )
87
+
88
+ result = meta1 + meta2
89
+
90
+ assert result.in_tokens == 300
91
+ assert result.out_tokens == 150
92
+ assert result.reasoning_tokens == 50
93
+ assert result.cache_read_tokens == 25
94
+ assert result.cache_write_tokens == 15
95
+ assert result.duration_seconds == 3.5
96
+ assert result.cost is not None
97
+ assert result.cost.input == 0.03
98
+ assert result.cost.output == 0.05
99
+
100
+ def test_add_metadata_missing_optional_tokens(self):
101
+ meta1 = QueryResultMetadata(
102
+ in_tokens=100,
103
+ out_tokens=50,
104
+ reasoning_tokens=20,
105
+ duration_seconds=1.0,
106
+ )
107
+ meta2 = QueryResultMetadata(
108
+ in_tokens=200,
109
+ out_tokens=100,
110
+ cache_read_tokens=15,
111
+ duration_seconds=2.0,
112
+ )
113
+
114
+ result = meta1 + meta2
115
+
116
+ assert result.in_tokens == 300
117
+ assert result.out_tokens == 150
118
+ assert result.reasoning_tokens == 20
119
+ assert result.cache_read_tokens == 15
120
+ assert result.cache_write_tokens is None
121
+
122
+ def test_add_metadata_both_missing_optional_tokens(self):
123
+ meta1 = QueryResultMetadata(in_tokens=100, out_tokens=50)
124
+ meta2 = QueryResultMetadata(in_tokens=200, out_tokens=100)
125
+
126
+ result = meta1 + meta2
127
+
128
+ assert result.in_tokens == 300
129
+ assert result.out_tokens == 150
130
+ assert result.reasoning_tokens is None
131
+ assert result.cache_read_tokens is None
132
+ assert result.cache_write_tokens is None
133
+
134
+ def test_add_metadata_one_has_cost(self):
135
+ meta1 = QueryResultMetadata(
136
+ in_tokens=100,
137
+ out_tokens=50,
138
+ cost=QueryResultCost(input=0.01, output=0.02),
139
+ )
140
+ meta2 = QueryResultMetadata(in_tokens=200, out_tokens=100, cost=None)
141
+
142
+ result = meta1 + meta2
143
+
144
+ assert result.cost is not None
145
+ assert result.cost.input == 0.01
146
+ assert result.cost.output == 0.02
147
+
148
+ def test_add_metadata_other_has_cost(self):
149
+ meta1 = QueryResultMetadata(in_tokens=100, out_tokens=50, cost=None)
150
+ meta2 = QueryResultMetadata(
151
+ in_tokens=200,
152
+ out_tokens=100,
153
+ cost=QueryResultCost(input=0.02, output=0.03),
154
+ )
155
+
156
+ result = meta1 + meta2
157
+
158
+ assert result.cost is not None
159
+ assert result.cost.input == 0.02
160
+ assert result.cost.output == 0.03
161
+
162
+ def test_add_metadata_neither_has_cost(self):
163
+ meta1 = QueryResultMetadata(in_tokens=100, out_tokens=50)
164
+ meta2 = QueryResultMetadata(in_tokens=200, out_tokens=100)
165
+
166
+ result = meta1 + meta2
167
+
168
+ assert result.cost is None
169
+
170
+ def test_add_metadata_default_duration(self):
171
+ meta1 = QueryResultMetadata(in_tokens=100, out_tokens=50, duration_seconds=None)
172
+ meta2 = QueryResultMetadata(in_tokens=200, out_tokens=100, duration_seconds=2.0)
173
+
174
+ result = meta1 + meta2
175
+
176
+ assert result.duration_seconds == 2.0
177
+
178
+ def test_add_metadata_both_none_duration(self):
179
+ meta1 = QueryResultMetadata(in_tokens=100, out_tokens=50, duration_seconds=None)
180
+ meta2 = QueryResultMetadata(
181
+ in_tokens=200, out_tokens=100, duration_seconds=None
182
+ )
183
+
184
+ result = meta1 + meta2
185
+
186
+ assert result.duration_seconds == 0.0
187
+
188
+ def test_computed_total_tokens(self):
189
+ meta1 = QueryResultMetadata(
190
+ in_tokens=100,
191
+ out_tokens=50,
192
+ reasoning_tokens=20,
193
+ cache_read_tokens=10,
194
+ cache_write_tokens=5,
195
+ )
196
+ meta2 = QueryResultMetadata(
197
+ in_tokens=200,
198
+ out_tokens=100,
199
+ reasoning_tokens=30,
200
+ cache_read_tokens=15,
201
+ )
202
+
203
+ result = meta1 + meta2
204
+
205
+ assert result.total_input_tokens == 300 + 25 + 5
206
+ assert result.total_output_tokens == 150 + 50
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes