agenta 0.26.0a0__py3-none-any.whl → 0.27.0__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.

Potentially problematic release.


This version of agenta might be problematic. Click here for more details.

Files changed (85) hide show
  1. agenta/__init__.py +29 -10
  2. agenta/cli/helper.py +5 -1
  3. agenta/client/backend/__init__.py +14 -0
  4. agenta/client/backend/apps/client.py +28 -20
  5. agenta/client/backend/client.py +47 -16
  6. agenta/client/backend/containers/client.py +5 -1
  7. agenta/client/backend/core/__init__.py +2 -1
  8. agenta/client/backend/core/client_wrapper.py +6 -6
  9. agenta/client/backend/core/file.py +33 -11
  10. agenta/client/backend/core/http_client.py +45 -31
  11. agenta/client/backend/core/pydantic_utilities.py +144 -29
  12. agenta/client/backend/core/request_options.py +3 -0
  13. agenta/client/backend/core/serialization.py +139 -42
  14. agenta/client/backend/evaluations/client.py +7 -2
  15. agenta/client/backend/evaluators/client.py +349 -1
  16. agenta/client/backend/observability/client.py +11 -2
  17. agenta/client/backend/testsets/client.py +10 -10
  18. agenta/client/backend/types/__init__.py +14 -0
  19. agenta/client/backend/types/app.py +1 -0
  20. agenta/client/backend/types/app_variant_response.py +3 -1
  21. agenta/client/backend/types/config_dto.py +32 -0
  22. agenta/client/backend/types/config_response_model.py +32 -0
  23. agenta/client/backend/types/create_span.py +3 -2
  24. agenta/client/backend/types/environment_output.py +1 -0
  25. agenta/client/backend/types/environment_output_extended.py +1 -0
  26. agenta/client/backend/types/evaluation.py +1 -2
  27. agenta/client/backend/types/evaluator.py +2 -0
  28. agenta/client/backend/types/evaluator_config.py +1 -0
  29. agenta/client/backend/types/evaluator_mapping_output_interface.py +21 -0
  30. agenta/client/backend/types/evaluator_output_interface.py +21 -0
  31. agenta/client/backend/types/human_evaluation.py +1 -2
  32. agenta/client/backend/types/lifecycle_dto.py +24 -0
  33. agenta/client/backend/types/llm_tokens.py +2 -2
  34. agenta/client/backend/types/reference_dto.py +23 -0
  35. agenta/client/backend/types/reference_request_model.py +23 -0
  36. agenta/client/backend/types/span.py +1 -0
  37. agenta/client/backend/types/span_detail.py +7 -1
  38. agenta/client/backend/types/test_set_output_response.py +5 -2
  39. agenta/client/backend/types/trace_detail.py +7 -1
  40. agenta/client/backend/types/with_pagination.py +4 -2
  41. agenta/client/backend/variants/client.py +1565 -272
  42. agenta/docker/docker-assets/Dockerfile.cloud.template +1 -1
  43. agenta/sdk/__init__.py +44 -7
  44. agenta/sdk/agenta_init.py +85 -33
  45. agenta/sdk/context/__init__.py +0 -0
  46. agenta/sdk/context/routing.py +26 -0
  47. agenta/sdk/context/tracing.py +3 -0
  48. agenta/sdk/decorators/__init__.py +0 -0
  49. agenta/sdk/decorators/{llm_entrypoint.py → routing.py} +216 -191
  50. agenta/sdk/decorators/tracing.py +218 -99
  51. agenta/sdk/litellm/__init__.py +1 -0
  52. agenta/sdk/litellm/litellm.py +288 -0
  53. agenta/sdk/managers/__init__.py +6 -0
  54. agenta/sdk/managers/config.py +318 -0
  55. agenta/sdk/managers/deployment.py +45 -0
  56. agenta/sdk/managers/shared.py +639 -0
  57. agenta/sdk/managers/variant.py +182 -0
  58. agenta/sdk/router.py +0 -7
  59. agenta/sdk/tracing/__init__.py +1 -0
  60. agenta/sdk/tracing/attributes.py +141 -0
  61. agenta/sdk/tracing/context.py +24 -0
  62. agenta/sdk/tracing/conventions.py +49 -0
  63. agenta/sdk/tracing/exporters.py +65 -0
  64. agenta/sdk/tracing/inline.py +1252 -0
  65. agenta/sdk/tracing/processors.py +117 -0
  66. agenta/sdk/tracing/spans.py +136 -0
  67. agenta/sdk/tracing/tracing.py +233 -0
  68. agenta/sdk/types.py +49 -2
  69. agenta/sdk/utils/{helper/openai_cost.py → costs.py} +3 -0
  70. agenta/sdk/utils/debug.py +5 -5
  71. agenta/sdk/utils/exceptions.py +52 -0
  72. agenta/sdk/utils/globals.py +3 -5
  73. agenta/sdk/{tracing/logger.py → utils/logging.py} +3 -5
  74. agenta/sdk/utils/singleton.py +13 -0
  75. {agenta-0.26.0a0.dist-info → agenta-0.27.0.dist-info}/METADATA +5 -1
  76. {agenta-0.26.0a0.dist-info → agenta-0.27.0.dist-info}/RECORD +78 -57
  77. agenta/sdk/config_manager.py +0 -205
  78. agenta/sdk/context.py +0 -41
  79. agenta/sdk/decorators/base.py +0 -10
  80. agenta/sdk/tracing/callbacks.py +0 -187
  81. agenta/sdk/tracing/llm_tracing.py +0 -617
  82. agenta/sdk/tracing/tasks_manager.py +0 -129
  83. agenta/sdk/tracing/tracing_context.py +0 -27
  84. {agenta-0.26.0a0.dist-info → agenta-0.27.0.dist-info}/WHEEL +0 -0
  85. {agenta-0.26.0a0.dist-info → agenta-0.27.0.dist-info}/entry_points.txt +0 -0
@@ -0,0 +1,1252 @@
1
+ ############################
2
+ ### services.shared.dtos ###
3
+ ### -------------------- ###
4
+
5
+ from typing import Optional
6
+
7
+ from pydantic import BaseModel
8
+ from uuid import UUID
9
+ from datetime import datetime
10
+ from enum import Enum
11
+ from collections import OrderedDict
12
+
13
+
14
+ class ProjectScopeDTO(BaseModel):
15
+ project_id: UUID
16
+
17
+
18
+ class LifecycleDTO(BaseModel):
19
+ created_at: datetime
20
+ updated_at: Optional[datetime] = None
21
+
22
+ updated_by_id: Optional[UUID] = None
23
+
24
+
25
+ ### -------------------- ###
26
+ ### services.shared.dtos ###
27
+ ############################
28
+
29
+
30
+ ###################################
31
+ ### services.observability.dtos ###
32
+ ### --------------------------- ###
33
+
34
+ from typing import List, Dict, Any, Union, Optional
35
+
36
+ from enum import Enum
37
+ from datetime import datetime
38
+ from uuid import UUID
39
+
40
+
41
+ class TimeDTO(BaseModel):
42
+ start: datetime
43
+ end: datetime
44
+ span: int
45
+
46
+
47
+ class StatusCode(Enum):
48
+ UNSET = "UNSET"
49
+ OK = "OK"
50
+ ERROR = "ERROR"
51
+
52
+
53
+ class StatusDTO(BaseModel):
54
+ code: StatusCode
55
+ message: Optional[str] = None
56
+ stacktrace: Optional[str] = None
57
+
58
+
59
+ AttributeValueType = Any
60
+ Attributes = Dict[str, AttributeValueType]
61
+
62
+
63
+ class TreeType(Enum):
64
+ # --- VARIANTS --- #
65
+ INVOCATION = "invocation"
66
+ # --- VARIANTS --- #
67
+
68
+
69
+ class NodeType(Enum):
70
+ # --- VARIANTS --- #
71
+ ## SPAN_KIND_SERVER
72
+ AGENT = "agent"
73
+ WORKFLOW = "workflow"
74
+ CHAIN = "chain"
75
+ ## SPAN_KIND_INTERNAL
76
+ TASK = "task"
77
+ ## SPAN_KIND_CLIENT
78
+ TOOL = "tool"
79
+ EMBEDDING = "embedding"
80
+ QUERY = "query"
81
+ COMPLETION = "completion"
82
+ CHAT = "chat"
83
+ RERANK = "rerank"
84
+ # --- VARIANTS --- #
85
+
86
+
87
+ class RootDTO(BaseModel):
88
+ id: UUID
89
+
90
+
91
+ class TreeDTO(BaseModel):
92
+ id: UUID
93
+ type: Optional[TreeType] = None
94
+
95
+
96
+ class NodeDTO(BaseModel):
97
+ id: UUID
98
+ type: Optional[NodeType] = None
99
+ name: str
100
+
101
+
102
+ Data = Dict[str, Any]
103
+ Metrics = Dict[str, Any]
104
+ Metadata = Dict[str, Any]
105
+ Tags = Dict[str, str]
106
+ Refs = Dict[str, str]
107
+
108
+
109
+ class LinkDTO(BaseModel):
110
+ type: str
111
+ id: UUID
112
+ tree_id: Optional[UUID] = None
113
+
114
+
115
+ class ParentDTO(BaseModel):
116
+ id: UUID
117
+
118
+
119
+ class OTelSpanKind(Enum):
120
+ SPAN_KIND_UNSPECIFIED = "SPAN_KIND_UNSPECIFIED"
121
+ # INTERNAL
122
+ SPAN_KIND_INTERNAL = "SPAN_KIND_INTERNAL"
123
+ # SYNCHRONOUS
124
+ SPAN_KIND_SERVER = "SPAN_KIND_SERVER"
125
+ SPAN_KIND_CLIENT = "SPAN_KIND_CLIENT"
126
+ # ASYNCHRONOUS
127
+ SPAN_KIND_PRODUCER = "SPAN_KIND_PRODUCER"
128
+ SPAN_KIND_CONSUMER = "SPAN_KIND_CONSUMER"
129
+
130
+
131
+ class OTelStatusCode(Enum):
132
+ STATUS_CODE_OK = "STATUS_CODE_OK"
133
+ STATUS_CODE_ERROR = "STATUS_CODE_ERROR"
134
+ STATUS_CODE_UNSET = "STATUS_CODE_UNSET"
135
+
136
+
137
+ class OTelContextDTO(BaseModel):
138
+ trace_id: str
139
+ span_id: str
140
+
141
+
142
+ class OTelEventDTO(BaseModel):
143
+ name: str
144
+ timestamp: datetime
145
+
146
+ attributes: Optional[Attributes] = None
147
+
148
+
149
+ class OTelLinkDTO(BaseModel):
150
+ context: OTelContextDTO
151
+
152
+ attributes: Optional[Attributes] = None
153
+
154
+
155
+ class OTelExtraDTO(BaseModel):
156
+ kind: Optional[str] = None
157
+
158
+ attributes: Optional[Attributes] = None
159
+ events: Optional[List[OTelEventDTO]] = None
160
+ links: Optional[List[OTelLinkDTO]] = None
161
+
162
+
163
+ class SpanDTO(BaseModel):
164
+ scope: Optional[ProjectScopeDTO] = None
165
+
166
+ lifecycle: Optional[LifecycleDTO] = None
167
+
168
+ root: RootDTO
169
+ tree: TreeDTO
170
+ node: NodeDTO
171
+
172
+ parent: Optional[ParentDTO] = None
173
+
174
+ time: TimeDTO
175
+ status: StatusDTO
176
+
177
+ data: Optional[Data] = None
178
+ metrics: Optional[Metrics] = None
179
+ meta: Optional[Metadata] = None
180
+ tags: Optional[Tags] = None
181
+ refs: Optional[Refs] = None
182
+
183
+ links: Optional[List[LinkDTO]] = None
184
+
185
+ otel: Optional[OTelExtraDTO] = None
186
+
187
+ nodes: Optional[Dict[str, Union["SpanDTO", List["SpanDTO"]]]] = None
188
+
189
+
190
+ class OTelSpanDTO(BaseModel):
191
+ context: OTelContextDTO
192
+
193
+ name: str
194
+ kind: OTelSpanKind = OTelSpanKind.SPAN_KIND_UNSPECIFIED
195
+
196
+ start_time: datetime
197
+ end_time: datetime
198
+
199
+ status_code: OTelStatusCode = OTelStatusCode.STATUS_CODE_UNSET
200
+ status_message: Optional[str] = None
201
+
202
+ attributes: Optional[Attributes] = None
203
+ events: Optional[List[OTelEventDTO]] = None
204
+
205
+ parent: Optional[OTelContextDTO] = None
206
+ links: Optional[List[OTelLinkDTO]] = None
207
+
208
+
209
+ ### --------------------------- ###
210
+ ### services.observability.dtos ###
211
+ ###################################
212
+
213
+
214
+ ####################################
215
+ ### services.observability.utils ###
216
+ ### ---------------------------- ###
217
+
218
+ from typing import List, Dict, OrderedDict
219
+
220
+
221
+ def parse_span_dtos_to_span_idx(
222
+ span_dtos: List[SpanDTO],
223
+ ) -> Dict[str, SpanDTO]:
224
+ span_idx = {span_dto.node.id: span_dto for span_dto in span_dtos}
225
+
226
+ return span_idx
227
+
228
+
229
+ def parse_span_idx_to_span_id_tree(
230
+ span_idx: Dict[str, SpanDTO],
231
+ ) -> OrderedDict:
232
+ span_id_tree = OrderedDict()
233
+ index = {}
234
+
235
+ def push(span_dto: SpanDTO) -> None:
236
+ if span_dto.parent is None:
237
+ span_id_tree[span_dto.node.id] = OrderedDict()
238
+ index[span_dto.node.id] = span_id_tree[span_dto.node.id]
239
+ elif span_dto.parent.id in index:
240
+ index[span_dto.parent.id][span_dto.node.id] = OrderedDict()
241
+ index[span_dto.node.id] = index[span_dto.parent.id][span_dto.node.id]
242
+
243
+ for span_dto in sorted(span_idx.values(), key=lambda span_dto: span_dto.time.start):
244
+ push(span_dto)
245
+
246
+ return span_id_tree
247
+
248
+
249
+ def cumulate_costs(
250
+ spans_id_tree: OrderedDict,
251
+ spans_idx: Dict[str, SpanDTO],
252
+ ) -> None:
253
+ def _get_unit(span: SpanDTO):
254
+ if span.metrics is not None:
255
+ return span.metrics.get("unit.costs.total", 0.0)
256
+
257
+ return 0.0
258
+
259
+ def _get_acc(span: SpanDTO):
260
+ if span.metrics is not None:
261
+ return span.metrics.get("acc.costs.total", 0.0)
262
+
263
+ return 0.0
264
+
265
+ def _acc(a: float, b: float):
266
+ return a + b
267
+
268
+ def _set(span: SpanDTO, cost: float):
269
+ if span.metrics is None:
270
+ span.metrics = {}
271
+
272
+ if cost != 0.0:
273
+ span.metrics["acc.costs.total"] = cost
274
+
275
+ _cumulate_tree_dfs(spans_id_tree, spans_idx, _get_unit, _get_acc, _acc, _set)
276
+
277
+
278
+ def cumulate_tokens(
279
+ spans_id_tree: OrderedDict,
280
+ spans_idx: Dict[str, dict],
281
+ ) -> None:
282
+ def _get_unit(span: SpanDTO):
283
+ _tokens = {
284
+ "prompt": 0.0,
285
+ "completion": 0.0,
286
+ "total": 0.0,
287
+ }
288
+
289
+ if span.metrics is not None:
290
+ return {
291
+ "prompt": span.metrics.get("unit.tokens.prompt", 0.0),
292
+ "completion": span.metrics.get("unit.tokens.completion", 0.0),
293
+ "total": span.metrics.get("unit.tokens.total", 0.0),
294
+ }
295
+
296
+ return _tokens
297
+
298
+ def _get_acc(span: SpanDTO):
299
+ _tokens = {
300
+ "prompt": 0.0,
301
+ "completion": 0.0,
302
+ "total": 0.0,
303
+ }
304
+
305
+ if span.metrics is not None:
306
+ return {
307
+ "prompt": span.metrics.get("acc.tokens.prompt", 0.0),
308
+ "completion": span.metrics.get("acc.tokens.completion", 0.0),
309
+ "total": span.metrics.get("acc.tokens.total", 0.0),
310
+ }
311
+
312
+ return _tokens
313
+
314
+ def _acc(a: dict, b: dict):
315
+ return {
316
+ "prompt": a.get("prompt", 0.0) + b.get("prompt", 0.0),
317
+ "completion": a.get("completion", 0.0) + b.get("completion", 0.0),
318
+ "total": a.get("total", 0.0) + b.get("total", 0.0),
319
+ }
320
+
321
+ def _set(span: SpanDTO, tokens: dict):
322
+ if span.metrics is None:
323
+ span.metrics = {}
324
+
325
+ if tokens.get("prompt", 0.0) != 0.0:
326
+ span.metrics["acc.tokens.prompt"] = tokens.get("prompt", 0.0)
327
+ if tokens.get("completion", 0.0) != 0.0:
328
+ span.metrics["acc.tokens.completion"] = (
329
+ tokens.get("completion", 0.0)
330
+ if tokens.get("completion", 0.0) != 0.0
331
+ else None
332
+ )
333
+ if tokens.get("total", 0.0) != 0.0:
334
+ span.metrics["acc.tokens.total"] = (
335
+ tokens.get("total", 0.0) if tokens.get("total", 0.0) != 0.0 else None
336
+ )
337
+
338
+ _cumulate_tree_dfs(spans_id_tree, spans_idx, _get_unit, _get_acc, _acc, _set)
339
+
340
+
341
+ def _cumulate_tree_dfs(
342
+ spans_id_tree: OrderedDict,
343
+ spans_idx: Dict[str, SpanDTO],
344
+ get_unit_metric,
345
+ get_acc_metric,
346
+ accumulate_metric,
347
+ set_metric,
348
+ ):
349
+ for span_id, children_spans_id_tree in spans_id_tree.items():
350
+ children_spans_id_tree: OrderedDict
351
+
352
+ cumulated_metric = get_unit_metric(spans_idx[span_id])
353
+
354
+ _cumulate_tree_dfs(
355
+ children_spans_id_tree,
356
+ spans_idx,
357
+ get_unit_metric,
358
+ get_acc_metric,
359
+ accumulate_metric,
360
+ set_metric,
361
+ )
362
+
363
+ for child_span_id in children_spans_id_tree.keys():
364
+ marginal_metric = get_acc_metric(spans_idx[child_span_id])
365
+ cumulated_metric = accumulate_metric(cumulated_metric, marginal_metric)
366
+
367
+ set_metric(spans_idx[span_id], cumulated_metric)
368
+
369
+
370
+ def connect_children(
371
+ spans_id_tree: OrderedDict,
372
+ spans_idx: Dict[str, dict],
373
+ ) -> None:
374
+ _connect_tree_dfs(spans_id_tree, spans_idx)
375
+
376
+
377
+ def _connect_tree_dfs(
378
+ spans_id_tree: OrderedDict,
379
+ spans_idx: Dict[str, SpanDTO],
380
+ ):
381
+ for span_id, children_spans_id_tree in spans_id_tree.items():
382
+ children_spans_id_tree: OrderedDict
383
+
384
+ parent_span = spans_idx[span_id]
385
+
386
+ parent_span.nodes = dict()
387
+
388
+ _connect_tree_dfs(children_spans_id_tree, spans_idx)
389
+
390
+ for child_span_id in children_spans_id_tree.keys():
391
+ child_span_name = spans_idx[child_span_id].node.name
392
+ if child_span_name not in parent_span.nodes:
393
+ parent_span.nodes[child_span_name] = spans_idx[child_span_id]
394
+ else:
395
+ if not isinstance(parent_span.nodes[child_span_name], list):
396
+ parent_span.nodes[child_span_name] = [
397
+ parent_span.nodes[child_span_name]
398
+ ]
399
+
400
+ parent_span.nodes[child_span_name].append(spans_idx[child_span_id])
401
+
402
+ if len(parent_span.nodes) == 0:
403
+ parent_span.nodes = None
404
+
405
+
406
+ ### ---------------------------- ###
407
+ ### services.observability.utils ###
408
+ ####################################
409
+
410
+
411
+ ########################################################
412
+ ### apis.fastapi.observability.opentelemetry.semconv ###
413
+ ### ------------------------------------------------ ###
414
+
415
+ from json import loads
416
+
417
+ VERSION = "0.4.1"
418
+
419
+ V_0_4_1_ATTRIBUTES_EXACT = [
420
+ # OPENLLMETRY
421
+ ("gen_ai.system", "ag.meta.system"),
422
+ ("gen_ai.request.base_url", "ag.meta.request.base_url"),
423
+ ("gen_ai.request.endpoint", "ag.meta.request.endpoint"),
424
+ ("gen_ai.request.headers", "ag.meta.request.headers"),
425
+ ("gen_ai.request.type", "ag.type.node"),
426
+ ("gen_ai.request.streaming", "ag.meta.request.streaming"),
427
+ ("gen_ai.request.model", "ag.meta.request.model"),
428
+ ("gen_ai.request.max_tokens", "ag.meta.request.max_tokens"),
429
+ ("gen_ai.request.temperature", "ag.meta.request.temperature"),
430
+ ("gen_ai.request.top_p", "ag.meta.request.top_p"),
431
+ ("gen_ai.response.model", "ag.meta.response.model"),
432
+ ("gen_ai.usage.prompt_tokens", "ag.metrics.unit.tokens.prompt"),
433
+ ("gen_ai.usage.completion_tokens", "ag.metrics.unit.tokens.completion"),
434
+ ("gen_ai.usage.total_tokens", "ag.metrics.unit.tokens.total"),
435
+ ("llm.headers", "ag.meta.request.headers"),
436
+ ("llm.request.type", "ag.type.node"),
437
+ ("llm.top_k", "ag.meta.request.top_k"),
438
+ ("llm.is_streaming", "ag.meta.request.streaming"),
439
+ ("llm.usage.total_tokens", "ag.metrics.unit.tokens.total"),
440
+ ("gen_ai.openai.api_base", "ag.meta.request.base_url"),
441
+ ("db.system", "ag.meta.system"),
442
+ ("db.vector.query.top_k", "ag.meta.request.top_k"),
443
+ ("pinecone.query.top_k", "ag.meta.request.top_k"),
444
+ ("traceloop.span.kind", "ag.type.node"),
445
+ ("traceloop.entity.name", "ag.node.name"),
446
+ # OPENINFERENCE
447
+ ("output.value", "ag.data.outputs"),
448
+ ("input.value", "ag.data.inputs"),
449
+ ("embedding.model_name", "ag.meta.request.model"),
450
+ ("llm.invocation_parameters", "ag.meta.request"),
451
+ ("llm.model_name", "ag.meta.request.model"),
452
+ ("llm.provider", "ag.meta.provider"),
453
+ ("llm.system", "ag.meta.system"),
454
+ ]
455
+ V_0_4_1_ATTRIBUTES_PREFIX = [
456
+ # OPENLLMETRY
457
+ ("gen_ai.prompt", "ag.data.inputs.prompt"),
458
+ ("gen_ai.completion", "ag.data.outputs.completion"),
459
+ ("llm.request.functions", "ag.data.inputs.functions"),
460
+ ("llm.request.tools", "ag.data.inputs.tools"),
461
+ # OPENINFERENCE
462
+ ("llm.token_count", "ag.metrics.unit.tokens"),
463
+ ("llm.input_messages", "ag.data.inputs.prompt"),
464
+ ("llm.output_messages", "ag.data.outputs.completion"),
465
+ ]
466
+
467
+ V_0_4_1_ATTRIBUTES_DYNAMIC = [
468
+ # OPENLLMETRY
469
+ ("traceloop.entity.input", lambda x: ("ag.data.inputs", loads(x).get("inputs"))),
470
+ ("traceloop.entity.output", lambda x: ("ag.data.outputs", loads(x).get("outputs"))),
471
+ ]
472
+
473
+
474
+ V_0_4_1_MAPS = {
475
+ "attributes": {
476
+ "exact": {
477
+ "from": {otel: agenta for otel, agenta in V_0_4_1_ATTRIBUTES_EXACT[::-1]},
478
+ "to": {agenta: otel for otel, agenta in V_0_4_1_ATTRIBUTES_EXACT[::-1]},
479
+ },
480
+ "prefix": {
481
+ "from": {otel: agenta for otel, agenta in V_0_4_1_ATTRIBUTES_PREFIX[::-1]},
482
+ "to": {agenta: otel for otel, agenta in V_0_4_1_ATTRIBUTES_PREFIX[::-1]},
483
+ },
484
+ "dynamic": {
485
+ "from": {otel: agenta for otel, agenta in V_0_4_1_ATTRIBUTES_DYNAMIC[::-1]}
486
+ },
487
+ },
488
+ }
489
+ V_0_4_1_KEYS = {
490
+ "attributes": {
491
+ "exact": {
492
+ "from": list(V_0_4_1_MAPS["attributes"]["exact"]["from"].keys()),
493
+ "to": list(V_0_4_1_MAPS["attributes"]["exact"]["to"].keys()),
494
+ },
495
+ "prefix": {
496
+ "from": list(V_0_4_1_MAPS["attributes"]["prefix"]["from"].keys()),
497
+ "to": list(V_0_4_1_MAPS["attributes"]["prefix"]["to"].keys()),
498
+ },
499
+ "dynamic": {
500
+ "from": list(V_0_4_1_MAPS["attributes"]["dynamic"]["from"].keys()),
501
+ },
502
+ },
503
+ }
504
+
505
+
506
+ MAPS = {
507
+ "0.4.1": V_0_4_1_MAPS, # LATEST
508
+ }
509
+ KEYS = {
510
+ "0.4.1": V_0_4_1_KEYS, # LATEST
511
+ }
512
+
513
+ CODEX = {"maps": MAPS[VERSION], "keys": KEYS[VERSION]}
514
+
515
+
516
+ ### ------------------------------------------------ ###
517
+ ### apis.fastapi.observability.opentelemetry.semconv ###
518
+ ########################################################
519
+
520
+
521
+ ########################################
522
+ ### apis.fastapi.observability.utils ###
523
+ ### -------------------------------- ###
524
+
525
+ from typing import Optional, Union, Tuple, Any, List, Dict
526
+ from uuid import UUID
527
+ from collections import OrderedDict
528
+ from json import loads, JSONDecodeError, dumps
529
+ from copy import copy
530
+
531
+
532
+ def _unmarshal_attributes(
533
+ marshalled: Dict[str, Any],
534
+ ) -> Dict[str, Any]:
535
+ """
536
+ Unmarshals a dictionary of marshalled attributes into a nested dictionary
537
+
538
+ Example:
539
+ marshalled = {
540
+ "ag.type": "tree",
541
+ "ag.node.name": "root",
542
+ "ag.node.children.0.name": "child1",
543
+ "ag.node.children.1.name": "child2"
544
+ }
545
+ unmarshalled = {
546
+ "ag": {
547
+ "type": "tree",
548
+ "node": {
549
+ "name": "root",
550
+ "children": [
551
+ {
552
+ "name": "child1",
553
+ },
554
+ {
555
+ "name": "child2",
556
+ }
557
+ ]
558
+ }
559
+ }
560
+ }
561
+ """
562
+ unmarshalled = {}
563
+
564
+ for key, value in marshalled.items():
565
+ keys = key.split(".")
566
+
567
+ level = unmarshalled
568
+
569
+ for i, part in enumerate(keys[:-1]):
570
+ if part.isdigit():
571
+ part = int(part)
572
+
573
+ if not isinstance(level, list):
574
+ level = []
575
+
576
+ while len(level) <= part:
577
+ level.append({})
578
+
579
+ level = level[part]
580
+
581
+ else:
582
+ if part not in level:
583
+ level[part] = {} if not keys[i + 1].isdigit() else []
584
+
585
+ level = level[part]
586
+
587
+ last_key = keys[-1]
588
+
589
+ if last_key.isdigit():
590
+ last_key = int(last_key)
591
+
592
+ if not isinstance(level, list):
593
+ level = []
594
+
595
+ while len(level) <= last_key:
596
+ level.append(None)
597
+
598
+ level[last_key] = value
599
+
600
+ else:
601
+ level[last_key] = value
602
+
603
+ return unmarshalled
604
+
605
+
606
+ def _encode_key(
607
+ namespace,
608
+ key: str,
609
+ ) -> str:
610
+ return f"ag.{namespace}.{key}"
611
+
612
+
613
+ def _decode_key(
614
+ namespace,
615
+ key: str,
616
+ ) -> str:
617
+ return key.replace(f"ag.{namespace}.", "")
618
+
619
+
620
+ def _decode_value(
621
+ value: Any,
622
+ ) -> Any:
623
+ if isinstance(value, (int, float, bool, bytes)):
624
+ return value
625
+
626
+ if isinstance(value, str):
627
+ if value == "@ag.type=none:":
628
+ return None
629
+
630
+ if value.startswith("@ag.type=json:"):
631
+ encoded = value[len("@ag.type=json:") :]
632
+ value = loads(encoded)
633
+ return value
634
+
635
+ return value
636
+
637
+ return value
638
+
639
+
640
+ def _get_attributes(
641
+ attributes: Attributes,
642
+ namespace: str,
643
+ ):
644
+ return {
645
+ _decode_key(namespace, key): _decode_value(value)
646
+ for key, value in attributes.items()
647
+ if key != _decode_key(namespace, key)
648
+ }
649
+
650
+
651
+ def _parse_from_types(
652
+ otel_span_dto: OTelSpanDTO,
653
+ ) -> dict:
654
+ types = _get_attributes(otel_span_dto.attributes, "type")
655
+
656
+ if types.get("tree"):
657
+ del otel_span_dto.attributes[_encode_key("type", "tree")]
658
+
659
+ if types.get("node"):
660
+ del otel_span_dto.attributes[_encode_key("type", "node")]
661
+
662
+ return types
663
+
664
+
665
+ def _parse_from_semconv(
666
+ attributes: Attributes,
667
+ ) -> None:
668
+ _attributes = copy(attributes)
669
+
670
+ for old_key, value in _attributes.items():
671
+ if old_key in CODEX["keys"]["attributes"]["exact"]["from"]:
672
+ new_key = CODEX["maps"]["attributes"]["exact"]["from"][old_key]
673
+
674
+ attributes[new_key] = value
675
+
676
+ del attributes[old_key]
677
+
678
+ else:
679
+ for prefix_key in CODEX["keys"]["attributes"]["prefix"]["from"]:
680
+ if old_key.startswith(prefix_key):
681
+ prefix = CODEX["maps"]["attributes"]["prefix"]["from"][prefix_key]
682
+
683
+ new_key = old_key.replace(prefix_key, prefix)
684
+
685
+ attributes[new_key] = value
686
+
687
+ del attributes[old_key]
688
+
689
+ for dynamic_key in CODEX["keys"]["attributes"]["dynamic"]["from"]:
690
+ if old_key == dynamic_key:
691
+ try:
692
+ new_key, new_value = CODEX["maps"]["attributes"]["dynamic"][
693
+ "from"
694
+ ][dynamic_key](value)
695
+
696
+ attributes[new_key] = new_value
697
+
698
+ except: # pylint: disable=bare-except
699
+ pass
700
+
701
+
702
+ def _parse_from_links(
703
+ otel_span_dto: OTelSpanDTO,
704
+ ) -> dict:
705
+ # TESTING
706
+ otel_span_dto.links = [
707
+ OTelLinkDTO(
708
+ context=otel_span_dto.context,
709
+ attributes={"ag.type.link": "testcase"},
710
+ )
711
+ ]
712
+ # -------
713
+
714
+ # LINKS
715
+ links = None
716
+ otel_links = None
717
+
718
+ if otel_span_dto.links:
719
+ links = list()
720
+ otel_links = list()
721
+
722
+ for link in otel_span_dto.links:
723
+ _links = _get_attributes(link.attributes, "type")
724
+
725
+ if _links:
726
+ link_type = _links.get("link")
727
+ link_tree_id = str(UUID(link.context.trace_id[2:]))
728
+ link_node_id = str(
729
+ UUID(link.context.trace_id[2 + 16 :] + link.context.span_id[2:])
730
+ )
731
+
732
+ links.append(
733
+ LinkDTO(
734
+ type=link_type,
735
+ tree_id=link_tree_id,
736
+ id=link_node_id,
737
+ )
738
+ )
739
+ else:
740
+ otel_links.append(link)
741
+
742
+ links = links if links else None
743
+ otel_links = otel_links if otel_links else None
744
+
745
+ otel_span_dto.links = otel_links
746
+
747
+ return links
748
+
749
+
750
+ def _parse_from_attributes(
751
+ otel_span_dto: OTelSpanDTO,
752
+ ) -> Tuple[dict, dict, dict, dict, dict]:
753
+ # DATA
754
+ _data = _get_attributes(otel_span_dto.attributes, "data")
755
+
756
+ for key in _data.keys():
757
+ del otel_span_dto.attributes[_encode_key("data", key)]
758
+
759
+ # _data = _unmarshal_attributes(_data)
760
+ _data = _data if _data else None
761
+
762
+ # METRICS
763
+ _metrics = _get_attributes(otel_span_dto.attributes, "metrics")
764
+
765
+ for key in _metrics.keys():
766
+ del otel_span_dto.attributes[_encode_key("metrics", key)]
767
+
768
+ # _metrics = _unmarshal_attributes(_metrics)
769
+ _metrics = _metrics if _metrics else None
770
+
771
+ # META
772
+ _meta = _get_attributes(otel_span_dto.attributes, "meta")
773
+
774
+ for key in _meta.keys():
775
+ del otel_span_dto.attributes[_encode_key("meta", key)]
776
+
777
+ # _meta = _unmarshal_attributes(_meta)
778
+ _meta = _meta if _meta else None
779
+
780
+ # TAGS
781
+ _tags = _get_attributes(otel_span_dto.attributes, "tags")
782
+
783
+ for key in _tags.keys():
784
+ del otel_span_dto.attributes[_encode_key("tags", key)]
785
+
786
+ _tags = _tags if _tags else None
787
+
788
+ # REFS
789
+ _refs = _get_attributes(otel_span_dto.attributes, "refs")
790
+
791
+ for key in _refs.keys():
792
+ del otel_span_dto.attributes[_encode_key("refs", key)]
793
+
794
+ _refs = _refs if _refs else None
795
+
796
+ if len(otel_span_dto.attributes.keys()) < 1:
797
+ otel_span_dto.attributes = None
798
+
799
+ return _data, _metrics, _meta, _tags, _refs
800
+
801
+
802
+ def parse_from_otel_span_dto(
803
+ otel_span_dto: OTelSpanDTO,
804
+ ) -> SpanDTO:
805
+ lifecyle = LifecycleDTO(
806
+ created_at=datetime.now(),
807
+ )
808
+
809
+ _parse_from_semconv(otel_span_dto.attributes)
810
+
811
+ types = _parse_from_types(otel_span_dto)
812
+
813
+ tree_id = UUID(otel_span_dto.context.trace_id[2:])
814
+
815
+ tree_type: str = types.get("tree")
816
+
817
+ tree = TreeDTO(
818
+ id=tree_id,
819
+ type=tree_type.lower() if tree_type else None,
820
+ )
821
+
822
+ node_id = UUID(tree_id.hex[16:] + otel_span_dto.context.span_id[2:])
823
+
824
+ node_type = NodeType.TASK
825
+ try:
826
+ node_type = NodeType(types.get("node", "").lower())
827
+ except: # pylint: disable=bare-except
828
+ pass
829
+
830
+ node = NodeDTO(
831
+ id=node_id,
832
+ type=node_type,
833
+ name=otel_span_dto.name,
834
+ )
835
+
836
+ parent = (
837
+ ParentDTO(
838
+ id=(
839
+ UUID(
840
+ otel_span_dto.parent.trace_id[2 + 16 :]
841
+ + otel_span_dto.parent.span_id[2:]
842
+ )
843
+ )
844
+ )
845
+ if otel_span_dto.parent
846
+ else None
847
+ )
848
+
849
+ duration = (otel_span_dto.end_time - otel_span_dto.start_time).total_seconds()
850
+
851
+ time = TimeDTO(
852
+ start=otel_span_dto.start_time,
853
+ end=otel_span_dto.end_time,
854
+ span=round(duration * 1_000_000), # microseconds
855
+ )
856
+
857
+ status = StatusDTO(
858
+ code=otel_span_dto.status_code.value.replace("STATUS_CODE_", ""),
859
+ message=otel_span_dto.status_message,
860
+ )
861
+
862
+ links = _parse_from_links(otel_span_dto)
863
+
864
+ data, metrics, meta, tags, refs = _parse_from_attributes(otel_span_dto)
865
+
866
+ root_id = str(tree_id)
867
+ if refs is not None:
868
+ root_id = refs.get("scenario.id", root_id)
869
+
870
+ root = RootDTO(id=UUID(root_id))
871
+
872
+ otel = OTelExtraDTO(
873
+ kind=otel_span_dto.kind.value,
874
+ attributes=otel_span_dto.attributes,
875
+ events=otel_span_dto.events,
876
+ links=otel_span_dto.links,
877
+ )
878
+
879
+ span_dto = SpanDTO(
880
+ lifecycle=lifecyle,
881
+ root=root,
882
+ tree=tree,
883
+ node=node,
884
+ parent=parent,
885
+ time=time,
886
+ status=status,
887
+ data=data,
888
+ metrics=metrics,
889
+ meta=meta,
890
+ tags=tags,
891
+ refs=refs,
892
+ links=links,
893
+ otel=otel,
894
+ )
895
+
896
+ return span_dto
897
+
898
+
899
+ def parse_to_agenta_span_dto(
900
+ span_dto: SpanDTO,
901
+ ) -> SpanDTO:
902
+ # DATA
903
+ if span_dto.data:
904
+ span_dto.data = _unmarshal_attributes(span_dto.data)
905
+
906
+ # if "outputs" in span_dto.data:
907
+ # if "__default__" in span_dto.data["outputs"]:
908
+ # span_dto.data["outputs"] = span_dto.data["outputs"]["__default__"]
909
+
910
+ # METRICS
911
+ if span_dto.metrics:
912
+ span_dto.metrics = _unmarshal_attributes(span_dto.metrics)
913
+
914
+ # META
915
+ if span_dto.meta:
916
+ span_dto.meta = _unmarshal_attributes(span_dto.meta)
917
+
918
+ # TAGS
919
+ if span_dto.tags:
920
+ span_dto.tags = _unmarshal_attributes(span_dto.tags)
921
+
922
+ # REFS
923
+ if span_dto.refs:
924
+ span_dto.refs = _unmarshal_attributes(span_dto.refs)
925
+
926
+ for link in span_dto.links:
927
+ link.tree_id = None
928
+
929
+ if span_dto.nodes:
930
+ for v in span_dto.nodes.values():
931
+ if isinstance(v, list):
932
+ for n in v:
933
+ parse_to_agenta_span_dto(n)
934
+ else:
935
+ parse_to_agenta_span_dto(v)
936
+
937
+ return span_dto
938
+
939
+
940
+ ### -------------------------------- ###
941
+ ### apis.fastapi.observability.utils ###
942
+ ########################################
943
+
944
+
945
+ from litellm import cost_calculator
946
+ from opentelemetry.sdk.trace import ReadableSpan
947
+
948
+
949
+ def parse_inline_trace(
950
+ spans: Dict[str, ReadableSpan],
951
+ ):
952
+ otel_span_dtos = _parse_readable_spans(spans)
953
+
954
+ ############################################################
955
+ ### apis.fastapi.observability.api.otlp_collect_traces() ###
956
+ ### ---------------------------------------------------- ###
957
+ span_dtos = [
958
+ parse_from_otel_span_dto(otel_span_dto) for otel_span_dto in otel_span_dtos
959
+ ]
960
+ ### ---------------------------------------------------- ###
961
+ ### apis.fastapi.observability.api.otlp_collect_traces() ###
962
+ ############################################################
963
+
964
+ #####################################################
965
+ ### services.observability.service.ingest/query() ###
966
+ ### --------------------------------------------- ###
967
+ span_idx = parse_span_dtos_to_span_idx(span_dtos)
968
+ span_id_tree = parse_span_idx_to_span_id_tree(span_idx)
969
+ ### --------------------------------------------- ###
970
+ ### services.observability.service.ingest/query() ###
971
+ #####################################################
972
+
973
+ ###############################################
974
+ ### services.observability.service.ingest() ###
975
+ ### --------------------------------------- ###
976
+ calculate_costs(span_idx)
977
+ cumulate_costs(span_id_tree, span_idx)
978
+ cumulate_tokens(span_id_tree, span_idx)
979
+ ### --------------------------------------- ###
980
+ ### services.observability.service.ingest() ###
981
+ ###############################################
982
+
983
+ ##############################################
984
+ ### services.observability.service.query() ###
985
+ ### -------------------------------------- ###
986
+ connect_children(span_id_tree, span_idx)
987
+ root_span_dtos = [span_idx[span_id] for span_id in span_id_tree.keys()]
988
+ agenta_span_dtos = [
989
+ parse_to_agenta_span_dto(span_dto) for span_dto in root_span_dtos
990
+ ]
991
+ ### -------------------------------------- ###
992
+ ### services.observability.service.query() ###
993
+ ##############################################
994
+
995
+ LEGACY = True
996
+ inline_trace = None
997
+
998
+ if LEGACY:
999
+ legacy_spans = [
1000
+ _parse_to_legacy_span(span_dto) for span_dto in span_idx.values()
1001
+ ]
1002
+
1003
+ root_span = agenta_span_dtos[0]
1004
+
1005
+ trace_id = root_span.root.id.hex
1006
+ latency = root_span.time.span / 1_000_000
1007
+ cost = root_span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
1008
+ tokens = {
1009
+ "prompt_tokens": root_span.metrics.get("acc", {})
1010
+ .get("tokens", {})
1011
+ .get("prompt", 0),
1012
+ "completion_tokens": root_span.metrics.get("acc", {})
1013
+ .get("tokens", {})
1014
+ .get("completion", 0),
1015
+ "total_tokens": root_span.metrics.get("acc", {})
1016
+ .get("tokens", {})
1017
+ .get("total", 0),
1018
+ }
1019
+
1020
+ spans = [
1021
+ loads(span.model_dump_json(exclude_none=True)) for span in legacy_spans
1022
+ ]
1023
+
1024
+ inline_trace = {
1025
+ "trace_id": trace_id,
1026
+ "latency": latency,
1027
+ "cost": cost,
1028
+ "usage": tokens,
1029
+ "spans": spans,
1030
+ }
1031
+
1032
+ else:
1033
+ spans = [
1034
+ loads(span_dto.model_dump_json(exclude_none=True))
1035
+ for span_dto in agenta_span_dtos
1036
+ ]
1037
+
1038
+ inline_trace = spans # turn into Agenta Model ?
1039
+
1040
+ return inline_trace
1041
+
1042
+
1043
+ def _parse_readable_spans(
1044
+ spans: List[ReadableSpan],
1045
+ ) -> List[OTelSpanDTO]:
1046
+ otel_span_dtos = list()
1047
+
1048
+ for span in spans:
1049
+ otel_span_dto = OTelSpanDTO(
1050
+ context=OTelContextDTO(
1051
+ trace_id=_int_to_hex(span.get_span_context().trace_id, 128),
1052
+ span_id=_int_to_hex(span.get_span_context().span_id, 64),
1053
+ ),
1054
+ name=span.name,
1055
+ kind=OTelSpanKind(
1056
+ "SPAN_KIND_"
1057
+ + (span.kind if isinstance(span.kind, str) else span.kind.name)
1058
+ ),
1059
+ start_time=_timestamp_ns_to_datetime(span.start_time),
1060
+ end_time=_timestamp_ns_to_datetime(span.end_time),
1061
+ status_code=OTelStatusCode("STATUS_CODE_" + span.status.status_code.name),
1062
+ status_message=span.status.description,
1063
+ attributes=span.attributes,
1064
+ events=[
1065
+ OTelEventDTO(
1066
+ name=event.name,
1067
+ timestamp=_timestamp_ns_to_datetime(event.timestamp),
1068
+ attributes=event.attributes,
1069
+ )
1070
+ for event in span.events
1071
+ ],
1072
+ parent=(
1073
+ OTelContextDTO(
1074
+ trace_id=_int_to_hex(span.parent.trace_id, 128),
1075
+ span_id=_int_to_hex(span.parent.span_id, 64),
1076
+ )
1077
+ if span.parent
1078
+ else None
1079
+ ),
1080
+ links=[
1081
+ OTelLinkDTO(
1082
+ context=OTelContextDTO(
1083
+ trace_id=_int_to_hex(link.context.trace_id, 128),
1084
+ span_id=_int_to_hex(link.context.span_id, 64),
1085
+ ),
1086
+ attributes=link.attributes,
1087
+ )
1088
+ for link in span.links
1089
+ ],
1090
+ )
1091
+
1092
+ otel_span_dtos.append(otel_span_dto)
1093
+
1094
+ return otel_span_dtos
1095
+
1096
+
1097
+ def _int_to_hex(integer, bits):
1098
+ _hex = hex(integer)[2:]
1099
+
1100
+ _hex = _hex.zfill(bits // 4)
1101
+
1102
+ _hex = "0x" + _hex
1103
+
1104
+ return _hex
1105
+
1106
+
1107
+ def _timestamp_ns_to_datetime(timestamp_ns):
1108
+ _datetime = datetime.fromtimestamp(
1109
+ timestamp_ns / 1_000_000_000,
1110
+ ).isoformat(
1111
+ timespec="microseconds",
1112
+ )
1113
+
1114
+ return _datetime
1115
+
1116
+
1117
+ class LlmTokens(BaseModel):
1118
+ prompt_tokens: Optional[int] = 0
1119
+ completion_tokens: Optional[int] = 0
1120
+ total_tokens: Optional[int] = 0
1121
+
1122
+
1123
+ class CreateSpan(BaseModel):
1124
+ id: str
1125
+ app_id: str
1126
+ variant_id: Optional[str] = None
1127
+ variant_name: Optional[str] = None
1128
+ inputs: Optional[Dict[str, Optional[Any]]] = None
1129
+ internals: Optional[Dict[str, Optional[Any]]] = None
1130
+ outputs: Optional[Union[str, Dict[str, Optional[Any]], List[Any]]] = None
1131
+ config: Optional[Dict[str, Optional[Any]]] = None
1132
+ environment: Optional[str] = None
1133
+ tags: Optional[List[str]] = None
1134
+ token_consumption: Optional[int] = None
1135
+ name: str
1136
+ parent_span_id: Optional[str] = None
1137
+ attributes: Optional[Dict[str, Optional[Any]]] = None
1138
+ spankind: str
1139
+ status: str
1140
+ user: Optional[str] = None
1141
+ start_time: datetime
1142
+ end_time: datetime
1143
+ tokens: Optional[LlmTokens] = None
1144
+ cost: Optional[float] = None
1145
+
1146
+
1147
+ def _parse_to_legacy_span(span: SpanDTO) -> CreateSpan:
1148
+ attributes = None
1149
+ if span.otel:
1150
+ attributes = span.otel.attributes or {}
1151
+
1152
+ if span.otel.events:
1153
+ for event in span.otel.events:
1154
+ if event.name == "exception":
1155
+ attributes.update(**event.attributes)
1156
+
1157
+ legacy_span = CreateSpan(
1158
+ id=span.node.id.hex[:24],
1159
+ spankind=span.node.type,
1160
+ name=span.node.name,
1161
+ #
1162
+ status=span.status.code.name,
1163
+ #
1164
+ start_time=span.time.start,
1165
+ end_time=span.time.end,
1166
+ #
1167
+ parent_span_id=span.parent.id.hex[:24] if span.parent else None,
1168
+ #
1169
+ inputs=span.data.get("inputs") if span.data else {},
1170
+ internals=span.data.get("internals") if span.data else {},
1171
+ outputs=span.data.get("outputs") if span.data else {},
1172
+ #
1173
+ environment=span.meta.get("environment") if span.meta else None,
1174
+ config=span.meta.get("configuration") if span.meta else None,
1175
+ #
1176
+ tokens=(
1177
+ LlmTokens(
1178
+ prompt_tokens=span.metrics.get("acc", {})
1179
+ .get("tokens", {})
1180
+ .get("prompt", 0.0),
1181
+ completion_tokens=span.metrics.get("acc", {})
1182
+ .get("tokens", {})
1183
+ .get("completion", 0.0),
1184
+ total_tokens=span.metrics.get("acc", {})
1185
+ .get("tokens", {})
1186
+ .get("total", 0.0),
1187
+ )
1188
+ if span.metrics
1189
+ else None
1190
+ ),
1191
+ cost=(
1192
+ span.metrics.get("acc", {}).get("costs", {}).get("total", 0.0)
1193
+ if span.metrics
1194
+ else None
1195
+ ),
1196
+ #
1197
+ app_id=(
1198
+ span.refs.get("application", {}).get("id", "missing-app-id")
1199
+ if span.refs
1200
+ else "missing-app-id"
1201
+ ),
1202
+ #
1203
+ attributes=attributes,
1204
+ #
1205
+ variant_id=None,
1206
+ variant_name=None,
1207
+ tags=None,
1208
+ token_consumption=None,
1209
+ user=None,
1210
+ )
1211
+
1212
+ return legacy_span
1213
+
1214
+
1215
+ TYPES_WITH_COSTS = [
1216
+ "embedding",
1217
+ "query",
1218
+ "completion",
1219
+ "chat",
1220
+ "rerank",
1221
+ ]
1222
+
1223
+
1224
+ def calculate_costs(span_idx: Dict[str, SpanDTO]):
1225
+ for span in span_idx.values():
1226
+ if (
1227
+ span.node.type
1228
+ and span.node.type.name.lower() in TYPES_WITH_COSTS
1229
+ and span.meta
1230
+ and span.metrics
1231
+ ):
1232
+ try:
1233
+ costs = cost_calculator.cost_per_token(
1234
+ model=span.meta.get("response.model"),
1235
+ prompt_tokens=span.metrics.get("unit.tokens.prompt", 0.0),
1236
+ completion_tokens=span.metrics.get("unit.tokens.completion", 0.0),
1237
+ call_type=span.node.type.name.lower(),
1238
+ response_time_ms=span.time.span // 1_000,
1239
+ )
1240
+
1241
+ if not costs:
1242
+ continue
1243
+
1244
+ prompt_cost, completion_cost = costs
1245
+ total_cost = prompt_cost + completion_cost
1246
+
1247
+ span.metrics["unit.costs.prompt"] = prompt_cost
1248
+ span.metrics["unit.costs.completion"] = completion_cost
1249
+ span.metrics["unit.costs.total"] = total_cost
1250
+
1251
+ except:
1252
+ pass