arize-phoenix 11.23.1__py3-none-any.whl → 12.28.1__py3-none-any.whl

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 (221) hide show
  1. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/METADATA +61 -36
  2. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/RECORD +212 -162
  3. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/WHEEL +1 -1
  4. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/licenses/IP_NOTICE +1 -1
  5. phoenix/__generated__/__init__.py +0 -0
  6. phoenix/__generated__/classification_evaluator_configs/__init__.py +20 -0
  7. phoenix/__generated__/classification_evaluator_configs/_document_relevance_classification_evaluator_config.py +17 -0
  8. phoenix/__generated__/classification_evaluator_configs/_hallucination_classification_evaluator_config.py +17 -0
  9. phoenix/__generated__/classification_evaluator_configs/_models.py +18 -0
  10. phoenix/__generated__/classification_evaluator_configs/_tool_selection_classification_evaluator_config.py +17 -0
  11. phoenix/__init__.py +2 -1
  12. phoenix/auth.py +27 -2
  13. phoenix/config.py +1594 -81
  14. phoenix/db/README.md +546 -28
  15. phoenix/db/bulk_inserter.py +119 -116
  16. phoenix/db/engines.py +140 -33
  17. phoenix/db/facilitator.py +22 -1
  18. phoenix/db/helpers.py +818 -65
  19. phoenix/db/iam_auth.py +64 -0
  20. phoenix/db/insertion/dataset.py +133 -1
  21. phoenix/db/insertion/document_annotation.py +9 -6
  22. phoenix/db/insertion/evaluation.py +2 -3
  23. phoenix/db/insertion/helpers.py +2 -2
  24. phoenix/db/insertion/session_annotation.py +176 -0
  25. phoenix/db/insertion/span_annotation.py +3 -4
  26. phoenix/db/insertion/trace_annotation.py +3 -4
  27. phoenix/db/insertion/types.py +41 -18
  28. phoenix/db/migrations/versions/01a8342c9cdf_add_user_id_on_datasets.py +40 -0
  29. phoenix/db/migrations/versions/0df286449799_add_session_annotations_table.py +105 -0
  30. phoenix/db/migrations/versions/272b66ff50f8_drop_single_indices.py +119 -0
  31. phoenix/db/migrations/versions/58228d933c91_dataset_labels.py +67 -0
  32. phoenix/db/migrations/versions/699f655af132_experiment_tags.py +57 -0
  33. phoenix/db/migrations/versions/735d3d93c33e_add_composite_indices.py +41 -0
  34. phoenix/db/migrations/versions/ab513d89518b_add_user_id_on_dataset_versions.py +40 -0
  35. phoenix/db/migrations/versions/d0690a79ea51_users_on_experiments.py +40 -0
  36. phoenix/db/migrations/versions/deb2c81c0bb2_dataset_splits.py +139 -0
  37. phoenix/db/migrations/versions/e76cbd66ffc3_add_experiments_dataset_examples.py +87 -0
  38. phoenix/db/models.py +364 -56
  39. phoenix/db/pg_config.py +10 -0
  40. phoenix/db/types/trace_retention.py +7 -6
  41. phoenix/experiments/functions.py +69 -19
  42. phoenix/inferences/inferences.py +1 -2
  43. phoenix/server/api/auth.py +9 -0
  44. phoenix/server/api/auth_messages.py +46 -0
  45. phoenix/server/api/context.py +60 -0
  46. phoenix/server/api/dataloaders/__init__.py +36 -0
  47. phoenix/server/api/dataloaders/annotation_summaries.py +60 -8
  48. phoenix/server/api/dataloaders/average_experiment_repeated_run_group_latency.py +50 -0
  49. phoenix/server/api/dataloaders/average_experiment_run_latency.py +17 -24
  50. phoenix/server/api/dataloaders/cache/two_tier_cache.py +1 -2
  51. phoenix/server/api/dataloaders/dataset_dataset_splits.py +52 -0
  52. phoenix/server/api/dataloaders/dataset_example_revisions.py +0 -1
  53. phoenix/server/api/dataloaders/dataset_example_splits.py +40 -0
  54. phoenix/server/api/dataloaders/dataset_examples_and_versions_by_experiment_run.py +47 -0
  55. phoenix/server/api/dataloaders/dataset_labels.py +36 -0
  56. phoenix/server/api/dataloaders/document_evaluation_summaries.py +2 -2
  57. phoenix/server/api/dataloaders/document_evaluations.py +6 -9
  58. phoenix/server/api/dataloaders/experiment_annotation_summaries.py +88 -34
  59. phoenix/server/api/dataloaders/experiment_dataset_splits.py +43 -0
  60. phoenix/server/api/dataloaders/experiment_error_rates.py +21 -28
  61. phoenix/server/api/dataloaders/experiment_repeated_run_group_annotation_summaries.py +77 -0
  62. phoenix/server/api/dataloaders/experiment_repeated_run_groups.py +57 -0
  63. phoenix/server/api/dataloaders/experiment_runs_by_experiment_and_example.py +44 -0
  64. phoenix/server/api/dataloaders/latency_ms_quantile.py +40 -8
  65. phoenix/server/api/dataloaders/record_counts.py +37 -10
  66. phoenix/server/api/dataloaders/session_annotations_by_session.py +29 -0
  67. phoenix/server/api/dataloaders/span_cost_summary_by_experiment_repeated_run_group.py +64 -0
  68. phoenix/server/api/dataloaders/span_cost_summary_by_project.py +28 -14
  69. phoenix/server/api/dataloaders/span_costs.py +3 -9
  70. phoenix/server/api/dataloaders/table_fields.py +2 -2
  71. phoenix/server/api/dataloaders/token_prices_by_model.py +30 -0
  72. phoenix/server/api/dataloaders/trace_annotations_by_trace.py +27 -0
  73. phoenix/server/api/exceptions.py +5 -1
  74. phoenix/server/api/helpers/playground_clients.py +263 -83
  75. phoenix/server/api/helpers/playground_spans.py +2 -1
  76. phoenix/server/api/helpers/playground_users.py +26 -0
  77. phoenix/server/api/helpers/prompts/conversions/google.py +103 -0
  78. phoenix/server/api/helpers/prompts/models.py +61 -19
  79. phoenix/server/api/input_types/{SpanAnnotationFilter.py → AnnotationFilter.py} +22 -14
  80. phoenix/server/api/input_types/ChatCompletionInput.py +3 -0
  81. phoenix/server/api/input_types/CreateProjectSessionAnnotationInput.py +37 -0
  82. phoenix/server/api/input_types/DatasetFilter.py +5 -2
  83. phoenix/server/api/input_types/ExperimentRunSort.py +237 -0
  84. phoenix/server/api/input_types/GenerativeModelInput.py +3 -0
  85. phoenix/server/api/input_types/ProjectSessionSort.py +158 -1
  86. phoenix/server/api/input_types/PromptVersionInput.py +47 -1
  87. phoenix/server/api/input_types/SpanSort.py +3 -2
  88. phoenix/server/api/input_types/UpdateAnnotationInput.py +34 -0
  89. phoenix/server/api/input_types/UserRoleInput.py +1 -0
  90. phoenix/server/api/mutations/__init__.py +8 -0
  91. phoenix/server/api/mutations/annotation_config_mutations.py +8 -8
  92. phoenix/server/api/mutations/api_key_mutations.py +15 -20
  93. phoenix/server/api/mutations/chat_mutations.py +106 -37
  94. phoenix/server/api/mutations/dataset_label_mutations.py +243 -0
  95. phoenix/server/api/mutations/dataset_mutations.py +21 -16
  96. phoenix/server/api/mutations/dataset_split_mutations.py +351 -0
  97. phoenix/server/api/mutations/experiment_mutations.py +2 -2
  98. phoenix/server/api/mutations/export_events_mutations.py +3 -3
  99. phoenix/server/api/mutations/model_mutations.py +11 -9
  100. phoenix/server/api/mutations/project_mutations.py +4 -4
  101. phoenix/server/api/mutations/project_session_annotations_mutations.py +158 -0
  102. phoenix/server/api/mutations/project_trace_retention_policy_mutations.py +8 -4
  103. phoenix/server/api/mutations/prompt_label_mutations.py +74 -65
  104. phoenix/server/api/mutations/prompt_mutations.py +65 -129
  105. phoenix/server/api/mutations/prompt_version_tag_mutations.py +11 -8
  106. phoenix/server/api/mutations/span_annotations_mutations.py +15 -10
  107. phoenix/server/api/mutations/trace_annotations_mutations.py +13 -8
  108. phoenix/server/api/mutations/trace_mutations.py +3 -3
  109. phoenix/server/api/mutations/user_mutations.py +55 -26
  110. phoenix/server/api/queries.py +501 -617
  111. phoenix/server/api/routers/__init__.py +2 -2
  112. phoenix/server/api/routers/auth.py +141 -87
  113. phoenix/server/api/routers/ldap.py +229 -0
  114. phoenix/server/api/routers/oauth2.py +349 -101
  115. phoenix/server/api/routers/v1/__init__.py +22 -4
  116. phoenix/server/api/routers/v1/annotation_configs.py +19 -30
  117. phoenix/server/api/routers/v1/annotations.py +455 -13
  118. phoenix/server/api/routers/v1/datasets.py +355 -68
  119. phoenix/server/api/routers/v1/documents.py +142 -0
  120. phoenix/server/api/routers/v1/evaluations.py +20 -28
  121. phoenix/server/api/routers/v1/experiment_evaluations.py +16 -6
  122. phoenix/server/api/routers/v1/experiment_runs.py +335 -59
  123. phoenix/server/api/routers/v1/experiments.py +475 -47
  124. phoenix/server/api/routers/v1/projects.py +16 -50
  125. phoenix/server/api/routers/v1/prompts.py +50 -39
  126. phoenix/server/api/routers/v1/sessions.py +108 -0
  127. phoenix/server/api/routers/v1/spans.py +156 -96
  128. phoenix/server/api/routers/v1/traces.py +51 -77
  129. phoenix/server/api/routers/v1/users.py +64 -24
  130. phoenix/server/api/routers/v1/utils.py +3 -7
  131. phoenix/server/api/subscriptions.py +257 -93
  132. phoenix/server/api/types/Annotation.py +90 -23
  133. phoenix/server/api/types/ApiKey.py +13 -17
  134. phoenix/server/api/types/AuthMethod.py +1 -0
  135. phoenix/server/api/types/ChatCompletionSubscriptionPayload.py +1 -0
  136. phoenix/server/api/types/Dataset.py +199 -72
  137. phoenix/server/api/types/DatasetExample.py +88 -18
  138. phoenix/server/api/types/DatasetExperimentAnnotationSummary.py +10 -0
  139. phoenix/server/api/types/DatasetLabel.py +57 -0
  140. phoenix/server/api/types/DatasetSplit.py +98 -0
  141. phoenix/server/api/types/DatasetVersion.py +49 -4
  142. phoenix/server/api/types/DocumentAnnotation.py +212 -0
  143. phoenix/server/api/types/Experiment.py +215 -68
  144. phoenix/server/api/types/ExperimentComparison.py +3 -9
  145. phoenix/server/api/types/ExperimentRepeatedRunGroup.py +155 -0
  146. phoenix/server/api/types/ExperimentRepeatedRunGroupAnnotationSummary.py +9 -0
  147. phoenix/server/api/types/ExperimentRun.py +120 -70
  148. phoenix/server/api/types/ExperimentRunAnnotation.py +158 -39
  149. phoenix/server/api/types/GenerativeModel.py +95 -42
  150. phoenix/server/api/types/GenerativeProvider.py +1 -1
  151. phoenix/server/api/types/ModelInterface.py +7 -2
  152. phoenix/server/api/types/PlaygroundModel.py +12 -2
  153. phoenix/server/api/types/Project.py +218 -185
  154. phoenix/server/api/types/ProjectSession.py +146 -29
  155. phoenix/server/api/types/ProjectSessionAnnotation.py +187 -0
  156. phoenix/server/api/types/ProjectTraceRetentionPolicy.py +1 -1
  157. phoenix/server/api/types/Prompt.py +119 -39
  158. phoenix/server/api/types/PromptLabel.py +42 -25
  159. phoenix/server/api/types/PromptVersion.py +11 -8
  160. phoenix/server/api/types/PromptVersionTag.py +65 -25
  161. phoenix/server/api/types/Span.py +130 -123
  162. phoenix/server/api/types/SpanAnnotation.py +189 -42
  163. phoenix/server/api/types/SystemApiKey.py +65 -1
  164. phoenix/server/api/types/Trace.py +184 -53
  165. phoenix/server/api/types/TraceAnnotation.py +149 -50
  166. phoenix/server/api/types/User.py +128 -33
  167. phoenix/server/api/types/UserApiKey.py +73 -26
  168. phoenix/server/api/types/node.py +10 -0
  169. phoenix/server/api/types/pagination.py +11 -2
  170. phoenix/server/app.py +154 -36
  171. phoenix/server/authorization.py +5 -4
  172. phoenix/server/bearer_auth.py +13 -5
  173. phoenix/server/cost_tracking/cost_model_lookup.py +42 -14
  174. phoenix/server/cost_tracking/model_cost_manifest.json +1085 -194
  175. phoenix/server/daemons/generative_model_store.py +61 -9
  176. phoenix/server/daemons/span_cost_calculator.py +10 -8
  177. phoenix/server/dml_event.py +13 -0
  178. phoenix/server/email/sender.py +29 -2
  179. phoenix/server/grpc_server.py +9 -9
  180. phoenix/server/jwt_store.py +8 -6
  181. phoenix/server/ldap.py +1449 -0
  182. phoenix/server/main.py +9 -3
  183. phoenix/server/oauth2.py +330 -12
  184. phoenix/server/prometheus.py +43 -6
  185. phoenix/server/rate_limiters.py +4 -9
  186. phoenix/server/retention.py +33 -20
  187. phoenix/server/session_filters.py +49 -0
  188. phoenix/server/static/.vite/manifest.json +51 -53
  189. phoenix/server/static/assets/components-BreFUQQa.js +6702 -0
  190. phoenix/server/static/assets/{index-BPCwGQr8.js → index-CTQoemZv.js} +42 -35
  191. phoenix/server/static/assets/pages-DBE5iYM3.js +9524 -0
  192. phoenix/server/static/assets/vendor-BGzfc4EU.css +1 -0
  193. phoenix/server/static/assets/vendor-DCE4v-Ot.js +920 -0
  194. phoenix/server/static/assets/vendor-codemirror-D5f205eT.js +25 -0
  195. phoenix/server/static/assets/{vendor-recharts-Bw30oz1A.js → vendor-recharts-V9cwpXsm.js} +7 -7
  196. phoenix/server/static/assets/{vendor-shiki-DZajAPeq.js → vendor-shiki-Do--csgv.js} +1 -1
  197. phoenix/server/static/assets/vendor-three-CmB8bl_y.js +3840 -0
  198. phoenix/server/templates/index.html +7 -1
  199. phoenix/server/thread_server.py +1 -2
  200. phoenix/server/utils.py +74 -0
  201. phoenix/session/client.py +55 -1
  202. phoenix/session/data_extractor.py +5 -0
  203. phoenix/session/evaluation.py +8 -4
  204. phoenix/session/session.py +44 -8
  205. phoenix/settings.py +2 -0
  206. phoenix/trace/attributes.py +80 -13
  207. phoenix/trace/dsl/query.py +2 -0
  208. phoenix/trace/projects.py +5 -0
  209. phoenix/utilities/template_formatters.py +1 -1
  210. phoenix/version.py +1 -1
  211. phoenix/server/api/types/Evaluation.py +0 -39
  212. phoenix/server/static/assets/components-D0DWAf0l.js +0 -5650
  213. phoenix/server/static/assets/pages-Creyamao.js +0 -8612
  214. phoenix/server/static/assets/vendor-CU36oj8y.js +0 -905
  215. phoenix/server/static/assets/vendor-CqDb5u4o.css +0 -1
  216. phoenix/server/static/assets/vendor-arizeai-Ctgw0e1G.js +0 -168
  217. phoenix/server/static/assets/vendor-codemirror-Cojjzqb9.js +0 -25
  218. phoenix/server/static/assets/vendor-three-BLWp5bic.js +0 -2998
  219. phoenix/utilities/deprecation.py +0 -31
  220. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/entry_points.txt +0 -0
  221. {arize_phoenix-11.23.1.dist-info → arize_phoenix-12.28.1.dist-info}/licenses/LICENSE +0 -0
@@ -174,40 +174,41 @@ class QueueInserter(ABC, Generic[_PrecursorT, _InsertableT, _RowT, _DmlEventT]):
174
174
  class Precursors(ABC):
175
175
  @dataclass(frozen=True)
176
176
  class SpanAnnotation:
177
+ updated_at: datetime
177
178
  span_id: str
178
179
  obj: models.SpanAnnotation
179
180
 
180
181
  def as_insertable(
181
182
  self,
182
183
  span_rowid: int,
183
- id_: Optional[int] = None,
184
184
  ) -> Insertables.SpanAnnotation:
185
185
  return Insertables.SpanAnnotation(
186
+ updated_at=self.updated_at,
186
187
  span_id=self.span_id,
187
188
  obj=self.obj,
188
189
  span_rowid=span_rowid,
189
- id_=id_,
190
190
  )
191
191
 
192
192
  @dataclass(frozen=True)
193
193
  class TraceAnnotation:
194
+ updated_at: datetime
194
195
  trace_id: str
195
196
  obj: models.TraceAnnotation
196
197
 
197
198
  def as_insertable(
198
199
  self,
199
200
  trace_rowid: int,
200
- id_: Optional[int] = None,
201
201
  ) -> Insertables.TraceAnnotation:
202
202
  return Insertables.TraceAnnotation(
203
+ updated_at=self.updated_at,
203
204
  trace_id=self.trace_id,
204
205
  obj=self.obj,
205
206
  trace_rowid=trace_rowid,
206
- id_=id_,
207
207
  )
208
208
 
209
209
  @dataclass(frozen=True)
210
210
  class DocumentAnnotation:
211
+ updated_at: datetime
211
212
  span_id: str
212
213
  document_position: int
213
214
  obj: models.DocumentAnnotation
@@ -215,56 +216,78 @@ class Precursors(ABC):
215
216
  def as_insertable(
216
217
  self,
217
218
  span_rowid: int,
218
- id_: Optional[int] = None,
219
219
  ) -> Insertables.DocumentAnnotation:
220
220
  return Insertables.DocumentAnnotation(
221
+ updated_at=self.updated_at,
221
222
  span_id=self.span_id,
222
223
  document_position=self.document_position,
223
224
  obj=self.obj,
224
225
  span_rowid=span_rowid,
225
- id_=id_,
226
+ )
227
+
228
+ @dataclass(frozen=True)
229
+ class SessionAnnotation:
230
+ updated_at: datetime
231
+ session_id: str
232
+ obj: models.ProjectSessionAnnotation
233
+
234
+ def as_insertable(
235
+ self,
236
+ project_session_rowid: int,
237
+ ) -> Insertables.SessionAnnotation:
238
+ return Insertables.SessionAnnotation(
239
+ updated_at=self.updated_at,
240
+ session_id=self.session_id,
241
+ obj=self.obj,
242
+ project_session_rowid=project_session_rowid,
226
243
  )
227
244
 
228
245
 
229
246
  class Insertables(ABC):
230
247
  @dataclass(frozen=True)
231
248
  class SpanAnnotation(Precursors.SpanAnnotation):
249
+ updated_at: datetime
232
250
  span_rowid: int
233
- identifier: str = ""
234
- id_: Optional[int] = None
235
251
 
236
252
  @property
237
253
  def row(self) -> models.SpanAnnotation:
238
254
  obj = copy(self.obj)
239
255
  obj.span_rowid = self.span_rowid
240
- if self.id_ is not None:
241
- obj.id = self.id_
256
+ obj.updated_at = self.updated_at
242
257
  return obj
243
258
 
244
259
  @dataclass(frozen=True)
245
260
  class TraceAnnotation(Precursors.TraceAnnotation):
261
+ updated_at: datetime
246
262
  trace_rowid: int
247
- identifier: str = ""
248
- id_: Optional[int] = None
249
263
 
250
264
  @property
251
265
  def row(self) -> models.TraceAnnotation:
252
266
  obj = copy(self.obj)
253
267
  obj.trace_rowid = self.trace_rowid
254
- if self.id_ is not None:
255
- obj.id = self.id_
268
+ obj.updated_at = self.updated_at
256
269
  return obj
257
270
 
258
271
  @dataclass(frozen=True)
259
272
  class DocumentAnnotation(Precursors.DocumentAnnotation):
273
+ updated_at: datetime
260
274
  span_rowid: int
261
- identifier: str = ""
262
- id_: Optional[int] = None
263
275
 
264
276
  @property
265
277
  def row(self) -> models.DocumentAnnotation:
266
278
  obj = copy(self.obj)
267
279
  obj.span_rowid = self.span_rowid
268
- if self.id_ is not None:
269
- obj.id = self.id_
280
+ obj.updated_at = self.updated_at
281
+ return obj
282
+
283
+ @dataclass(frozen=True)
284
+ class SessionAnnotation(Precursors.SessionAnnotation):
285
+ updated_at: datetime
286
+ project_session_rowid: int
287
+
288
+ @property
289
+ def row(self) -> models.ProjectSessionAnnotation:
290
+ obj = copy(self.obj)
291
+ obj.project_session_id = self.project_session_rowid
292
+ obj.updated_at = self.updated_at
270
293
  return obj
@@ -0,0 +1,40 @@
1
+ """add user_id on datasets
2
+
3
+ Revision ID: 01a8342c9cdf
4
+ Revises: 0df286449799
5
+ Create Date: 2025-09-25 16:08:51.254947
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ import sqlalchemy as sa
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision: str = "01a8342c9cdf"
16
+ down_revision: Union[str, None] = "0df286449799"
17
+ branch_labels: Union[str, Sequence[str], None] = None
18
+ depends_on: Union[str, Sequence[str], None] = None
19
+
20
+ _Integer = sa.Integer().with_variant(
21
+ sa.BigInteger(),
22
+ "postgresql",
23
+ )
24
+
25
+
26
+ def upgrade() -> None:
27
+ with op.batch_alter_table("datasets") as batch_op:
28
+ batch_op.add_column(
29
+ sa.Column(
30
+ "user_id",
31
+ _Integer,
32
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
33
+ nullable=True,
34
+ ),
35
+ )
36
+
37
+
38
+ def downgrade() -> None:
39
+ with op.batch_alter_table("datasets") as batch_op:
40
+ batch_op.drop_column("user_id")
@@ -0,0 +1,105 @@
1
+ """add session annotations table
2
+
3
+ Revision ID: 0df286449799
4
+ Revises: 735d3d93c33e
5
+ Create Date: 2025-08-06 11:27:01.479664
6
+ """
7
+
8
+ from typing import Any, Sequence, Union
9
+
10
+ import sqlalchemy as sa
11
+ from alembic import op
12
+ from sqlalchemy import JSON
13
+ from sqlalchemy.dialects import postgresql
14
+ from sqlalchemy.ext.compiler import compiles
15
+
16
+ # revision identifiers, used by Alembic.
17
+ revision: str = "0df286449799"
18
+ down_revision: Union[str, None] = "735d3d93c33e"
19
+ branch_labels: Union[str, Sequence[str], None] = None
20
+ depends_on: Union[str, Sequence[str], None] = None
21
+
22
+
23
+ class JSONB(JSON):
24
+ # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
25
+ __visit_name__ = "JSONB"
26
+
27
+
28
+ @compiles(JSONB, "sqlite")
29
+ def _(*args: Any, **kwargs: Any) -> str:
30
+ # See https://docs.sqlalchemy.org/en/20/core/custom_types.html
31
+ return "JSONB"
32
+
33
+
34
+ JSON_ = (
35
+ JSON()
36
+ .with_variant(
37
+ postgresql.JSONB(),
38
+ "postgresql",
39
+ )
40
+ .with_variant(
41
+ JSONB(),
42
+ "sqlite",
43
+ )
44
+ )
45
+
46
+ _Integer = sa.Integer().with_variant(
47
+ sa.BigInteger(),
48
+ "postgresql",
49
+ )
50
+
51
+
52
+ def upgrade() -> None:
53
+ op.create_table(
54
+ "project_session_annotations",
55
+ sa.Column("id", _Integer, primary_key=True),
56
+ sa.Column(
57
+ "project_session_id",
58
+ _Integer,
59
+ sa.ForeignKey("project_sessions.id", ondelete="CASCADE"),
60
+ nullable=False,
61
+ index=True,
62
+ ),
63
+ sa.Column("name", sa.String, nullable=False),
64
+ sa.Column("label", sa.String),
65
+ sa.Column("score", sa.Float),
66
+ sa.Column("explanation", sa.String),
67
+ sa.Column("metadata", JSON_, nullable=False),
68
+ sa.Column(
69
+ "annotator_kind",
70
+ sa.String,
71
+ sa.CheckConstraint(
72
+ "annotator_kind IN ('LLM', 'CODE', 'HUMAN')",
73
+ name="valid_annotator_kind",
74
+ ),
75
+ nullable=False,
76
+ ),
77
+ sa.Column(
78
+ "user_id",
79
+ _Integer,
80
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
81
+ nullable=True,
82
+ ),
83
+ sa.Column("identifier", sa.String, server_default="", nullable=False),
84
+ sa.Column(
85
+ "source",
86
+ sa.String,
87
+ sa.CheckConstraint("source IN ('API', 'APP')", name="valid_source"),
88
+ nullable=False,
89
+ ),
90
+ sa.Column(
91
+ "created_at", sa.TIMESTAMP(timezone=True), server_default=sa.func.now(), nullable=False
92
+ ),
93
+ sa.Column(
94
+ "updated_at",
95
+ sa.TIMESTAMP(timezone=True),
96
+ server_default=sa.func.now(),
97
+ onupdate=sa.func.now(),
98
+ nullable=False,
99
+ ),
100
+ sa.UniqueConstraint("name", "project_session_id", "identifier"),
101
+ )
102
+
103
+
104
+ def downgrade() -> None:
105
+ op.drop_table("project_session_annotations")
@@ -0,0 +1,119 @@
1
+ """drop single indices from traces, project_sessions, and experiment_runs
2
+
3
+ Revision ID: 272b66ff50f8
4
+ Revises: a20694b15f82
5
+ Create Date: 2025-08-11 20:37:46.941940
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ from alembic import op
12
+
13
+ # revision identifiers, used by Alembic.
14
+ revision: str = "272b66ff50f8"
15
+ down_revision: Union[str, None] = "a20694b15f82"
16
+ branch_labels: Union[str, Sequence[str], None] = None
17
+ depends_on: Union[str, Sequence[str], None] = None
18
+
19
+
20
+ def upgrade() -> None:
21
+ op.drop_index(
22
+ "ix_traces_project_rowid",
23
+ table_name="traces",
24
+ if_exists=True,
25
+ )
26
+ op.drop_index(
27
+ "ix_traces_start_time",
28
+ table_name="traces",
29
+ if_exists=True,
30
+ )
31
+
32
+ op.drop_index(
33
+ "ix_project_sessions_project_id",
34
+ table_name="project_sessions",
35
+ if_exists=True,
36
+ )
37
+ op.drop_index(
38
+ "ix_project_sessions_start_time",
39
+ table_name="project_sessions",
40
+ if_exists=True,
41
+ )
42
+
43
+ op.drop_index(
44
+ "ix_experiment_runs_experiment_id",
45
+ table_name="experiment_runs",
46
+ if_exists=True,
47
+ )
48
+ op.drop_index(
49
+ "ix_experiment_run_annotations_experiment_run_id",
50
+ table_name="experiment_run_annotations",
51
+ if_exists=True,
52
+ )
53
+
54
+ op.drop_index(
55
+ "ix_dataset_example_revisions_dataset_example_id",
56
+ table_name="dataset_example_revisions",
57
+ if_exists=True,
58
+ )
59
+
60
+ op.drop_index(
61
+ "ix_span_cost_details_span_cost_id",
62
+ table_name="span_cost_details",
63
+ if_exists=True,
64
+ )
65
+
66
+
67
+ def downgrade() -> None:
68
+ op.create_index(
69
+ "ix_traces_project_rowid",
70
+ "traces",
71
+ ["project_rowid"],
72
+ if_not_exists=True,
73
+ )
74
+ op.create_index(
75
+ "ix_traces_start_time",
76
+ "traces",
77
+ ["start_time"],
78
+ if_not_exists=True,
79
+ )
80
+
81
+ op.create_index(
82
+ "ix_project_sessions_project_id",
83
+ "project_sessions",
84
+ ["project_id"],
85
+ if_not_exists=True,
86
+ )
87
+ op.create_index(
88
+ "ix_project_sessions_start_time",
89
+ "project_sessions",
90
+ ["start_time"],
91
+ if_not_exists=True,
92
+ )
93
+
94
+ op.create_index(
95
+ "ix_experiment_runs_experiment_id",
96
+ "experiment_runs",
97
+ ["experiment_id"],
98
+ if_not_exists=True,
99
+ )
100
+ op.create_index(
101
+ "ix_experiment_run_annotations_experiment_run_id",
102
+ "experiment_run_annotations",
103
+ ["experiment_run_id"],
104
+ if_not_exists=True,
105
+ )
106
+
107
+ op.create_index(
108
+ "ix_dataset_example_revisions_dataset_example_id",
109
+ "dataset_example_revisions",
110
+ ["dataset_example_id"],
111
+ if_not_exists=True,
112
+ )
113
+
114
+ op.create_index(
115
+ "ix_span_cost_details_span_cost_id",
116
+ "span_cost_details",
117
+ ["span_cost_id"],
118
+ if_not_exists=True,
119
+ )
@@ -0,0 +1,67 @@
1
+ """dataset_labels
2
+
3
+ Revision ID: 58228d933c91
4
+ Revises: 699f655af132
5
+ Create Date: 2025-09-05 17:47:34.637329
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ import sqlalchemy as sa
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision: str = "58228d933c91"
16
+ down_revision: Union[str, None] = "699f655af132"
17
+ branch_labels: Union[str, Sequence[str], None] = None
18
+ depends_on: Union[str, Sequence[str], None] = None
19
+
20
+ _Integer = sa.Integer().with_variant(
21
+ sa.BigInteger(),
22
+ "postgresql",
23
+ )
24
+
25
+
26
+ def upgrade() -> None:
27
+ op.create_table(
28
+ "dataset_labels",
29
+ sa.Column("id", _Integer, primary_key=True),
30
+ sa.Column("name", sa.String, nullable=False, unique=True),
31
+ sa.Column("description", sa.String, nullable=True),
32
+ sa.Column("color", sa.String, nullable=False),
33
+ sa.Column(
34
+ "user_id",
35
+ _Integer,
36
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
37
+ nullable=True,
38
+ index=True,
39
+ ),
40
+ )
41
+
42
+ op.create_table(
43
+ "datasets_dataset_labels",
44
+ sa.Column(
45
+ "dataset_id",
46
+ _Integer,
47
+ sa.ForeignKey("datasets.id", ondelete="CASCADE"),
48
+ nullable=False,
49
+ ),
50
+ sa.Column(
51
+ "dataset_label_id",
52
+ _Integer,
53
+ sa.ForeignKey("dataset_labels.id", ondelete="CASCADE"),
54
+ nullable=False,
55
+ # index on the second element of the composite primary key
56
+ index=True,
57
+ ),
58
+ sa.PrimaryKeyConstraint(
59
+ "dataset_id",
60
+ "dataset_label_id",
61
+ ),
62
+ )
63
+
64
+
65
+ def downgrade() -> None:
66
+ op.drop_table("datasets_dataset_labels")
67
+ op.drop_table("dataset_labels")
@@ -0,0 +1,57 @@
1
+ """experiment_tags
2
+
3
+ Revision ID: 699f655af132
4
+ Revises: d0690a79ea51
5
+ Create Date: 2025-09-05 13:14:22.676233
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ import sqlalchemy as sa
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision: str = "699f655af132"
16
+ down_revision: Union[str, None] = "d0690a79ea51"
17
+ branch_labels: Union[str, Sequence[str], None] = None
18
+ depends_on: Union[str, Sequence[str], None] = None
19
+
20
+ _Integer = sa.Integer().with_variant(
21
+ sa.BigInteger(),
22
+ "postgresql",
23
+ )
24
+
25
+
26
+ def upgrade() -> None:
27
+ op.create_table(
28
+ "experiment_tags",
29
+ sa.Column("id", _Integer, primary_key=True),
30
+ sa.Column(
31
+ "experiment_id",
32
+ _Integer,
33
+ sa.ForeignKey("experiments.id", ondelete="CASCADE"),
34
+ nullable=False,
35
+ index=True,
36
+ ),
37
+ sa.Column(
38
+ "dataset_id",
39
+ _Integer,
40
+ sa.ForeignKey("datasets.id", ondelete="CASCADE"),
41
+ nullable=False,
42
+ ),
43
+ sa.Column(
44
+ "user_id",
45
+ _Integer,
46
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
47
+ nullable=True,
48
+ index=True,
49
+ ),
50
+ sa.Column("name", sa.String, nullable=False),
51
+ sa.Column("description", sa.String, nullable=True),
52
+ sa.UniqueConstraint("dataset_id", "name"),
53
+ )
54
+
55
+
56
+ def downgrade() -> None:
57
+ op.drop_table("experiment_tags")
@@ -0,0 +1,41 @@
1
+ """add composite indices to traces and project_sessions
2
+
3
+ Revision ID: 735d3d93c33e
4
+ Revises: 272b66ff50f8
5
+ Create Date: 2025-08-11 20:52:47.477712
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ from alembic import op
12
+
13
+ # revision identifiers, used by Alembic.
14
+ revision: str = "735d3d93c33e"
15
+ down_revision: Union[str, None] = "272b66ff50f8"
16
+ branch_labels: Union[str, Sequence[str], None] = None
17
+ depends_on: Union[str, Sequence[str], None] = None
18
+
19
+
20
+ def upgrade() -> None:
21
+ op.execute(
22
+ "CREATE INDEX IF NOT EXISTS ix_traces_project_rowid_start_time "
23
+ "ON traces (project_rowid, start_time DESC)"
24
+ )
25
+ op.execute(
26
+ "CREATE INDEX IF NOT EXISTS ix_project_sessions_project_id_start_time "
27
+ "ON project_sessions (project_id, start_time DESC)"
28
+ )
29
+
30
+
31
+ def downgrade() -> None:
32
+ op.drop_index(
33
+ "ix_traces_project_rowid_start_time",
34
+ table_name="traces",
35
+ if_exists=True,
36
+ )
37
+ op.drop_index(
38
+ "ix_project_sessions_project_id_start_time",
39
+ table_name="project_sessions",
40
+ if_exists=True,
41
+ )
@@ -0,0 +1,40 @@
1
+ """add user_id on dataset_versions
2
+
3
+ Revision ID: ab513d89518b
4
+ Revises: 01a8342c9cdf
5
+ Create Date: 2025-09-26 11:00:06.961920
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ import sqlalchemy as sa
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision: str = "ab513d89518b"
16
+ down_revision: Union[str, None] = "01a8342c9cdf"
17
+ branch_labels: Union[str, Sequence[str], None] = None
18
+ depends_on: Union[str, Sequence[str], None] = None
19
+
20
+ _Integer = sa.Integer().with_variant(
21
+ sa.BigInteger(),
22
+ "postgresql",
23
+ )
24
+
25
+
26
+ def upgrade() -> None:
27
+ with op.batch_alter_table("dataset_versions") as batch_op:
28
+ batch_op.add_column(
29
+ sa.Column(
30
+ "user_id",
31
+ _Integer,
32
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
33
+ nullable=True,
34
+ ),
35
+ )
36
+
37
+
38
+ def downgrade() -> None:
39
+ with op.batch_alter_table("dataset_versions") as batch_op:
40
+ batch_op.drop_column("user_id")
@@ -0,0 +1,40 @@
1
+ """add user_id on experiments
2
+
3
+ Revision ID: d0690a79ea51
4
+ Revises: ab513d89518b
5
+ Create Date: 2025-08-26 19:12:47.849806
6
+
7
+ """
8
+
9
+ from typing import Sequence, Union
10
+
11
+ import sqlalchemy as sa
12
+ from alembic import op
13
+
14
+ # revision identifiers, used by Alembic.
15
+ revision: str = "d0690a79ea51"
16
+ down_revision: Union[str, None] = "ab513d89518b"
17
+ branch_labels: Union[str, Sequence[str], None] = None
18
+ depends_on: Union[str, Sequence[str], None] = None
19
+
20
+ _Integer = sa.Integer().with_variant(
21
+ sa.BigInteger(),
22
+ "postgresql",
23
+ )
24
+
25
+
26
+ def upgrade() -> None:
27
+ with op.batch_alter_table("experiments") as batch_op:
28
+ batch_op.add_column(
29
+ sa.Column(
30
+ "user_id",
31
+ _Integer,
32
+ sa.ForeignKey("users.id", ondelete="SET NULL"),
33
+ nullable=True,
34
+ ),
35
+ )
36
+
37
+
38
+ def downgrade() -> None:
39
+ with op.batch_alter_table("experiments") as batch_op:
40
+ batch_op.drop_column("user_id")