From ce3ebe903fb3d7e6beed2f5d1352554c39ab3ccf Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Sun, 2 Jun 2024 19:56:03 +0200 Subject: [PATCH 1/6] propagate cost and tokens --- agenta-cli/agenta/sdk/tracing/llm_tracing.py | 23 +++++++++++++++++--- 1 file changed, 20 insertions(+), 3 deletions(-) diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py index 263334adac..c7bc8e2c43 100644 --- a/agenta-cli/agenta/sdk/tracing/llm_tracing.py +++ b/agenta-cli/agenta/sdk/tracing/llm_tracing.py @@ -164,6 +164,20 @@ def start_span( def update_span_status(self, span: CreateSpan, value: str): span.status = value + def _update_span_cost(self, span: CreateSpan, cost: Optional[float]): + if cost is not None and isinstance(cost, float): + if span.cost is None: + span.cost = cost + else: + span.cost += cost + + def _update_span_tokens(self, span: CreateSpan, tokens: Optional[int]): + if tokens is not None and isinstance(tokens, int): + if span.tokens is None: + span.tokens = tokens + else: + span.tokens += tokens + def end_span(self, outputs: Dict[str, Any]): """ Ends the active span, if it is a parent span, ends the trace too. @@ -172,8 +186,8 @@ def end_span(self, outputs: Dict[str, Any]): raise ValueError("There is no active span to end.") self.active_span.end_time = datetime.now(timezone.utc) self.active_span.outputs = [outputs.get("message", "")] - self.active_span.cost = outputs.get("cost", None) - self.active_span.tokens = outputs.get("usage", None) + self._update_span_cost(self.active_span, outputs.get("cost", None)) + self._update_span_tokens(self.active_span, outputs.get("usage", None)) # Push span to list of recorded spans self.pending_spans.append(self.active_span) @@ -183,7 +197,10 @@ def end_span(self, outputs: Dict[str, Any]): if self.active_span.parent_span_id is None: self.end_trace(parent_span=self.active_span) else: - self.active_span = self.span_dict[self.active_span.parent_span_id] + parent_span = self.span_dict[self.active_span.parent_span_id] + self._update_span_cost(parent_span, self.active_span.cost) + self._update_span_tokens(parent_span, self.active_span.tokens) + self.active_span = parent_span def end_trace(self, parent_span: CreateSpan): if self.api_key == "": From 915c5abb874caaa970f4d5a5604aaec07afff412 Mon Sep 17 00:00:00 2001 From: Abram Date: Wed, 5 Jun 2024 08:49:33 +0100 Subject: [PATCH 2/6] refactor (tool): set LlmTokens frozen config to False --- agenta-cli/agenta/client/backend/types/llm_tokens.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/agenta-cli/agenta/client/backend/types/llm_tokens.py b/agenta-cli/agenta/client/backend/types/llm_tokens.py index 0c5a71755b..7e9fd8358b 100644 --- a/agenta-cli/agenta/client/backend/types/llm_tokens.py +++ b/agenta-cli/agenta/client/backend/types/llm_tokens.py @@ -33,6 +33,6 @@ def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]: return super().dict(**kwargs_with_defaults) class Config: - frozen = True + frozen = False smart_union = True json_encoders = {dt.datetime: serialize_datetime} From 8970b80c70c8757cb3598793e5dbc546fc09e7d2 Mon Sep 17 00:00:00 2001 From: Abram Date: Wed, 5 Jun 2024 08:53:50 +0100 Subject: [PATCH 3/6] refactor (sdk): improve tokens argument at _update_span_tokens method --- agenta-cli/agenta/sdk/tracing/llm_tracing.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py index c7bc8e2c43..fb5b9719f9 100644 --- a/agenta-cli/agenta/sdk/tracing/llm_tracing.py +++ b/agenta-cli/agenta/sdk/tracing/llm_tracing.py @@ -7,7 +7,11 @@ from agenta.sdk.tracing.tasks_manager import TaskQueue from agenta.client.backend.client import AsyncAgentaApi from agenta.client.backend.client import AsyncObservabilityClient -from agenta.client.backend.types.create_span import CreateSpan, SpanKind, SpanStatusCode +from agenta.client.backend.types.create_span import ( + CreateSpan, + LlmTokens, + SpanStatusCode, +) from bson.objectid import ObjectId @@ -171,12 +175,14 @@ def _update_span_cost(self, span: CreateSpan, cost: Optional[float]): else: span.cost += cost - def _update_span_tokens(self, span: CreateSpan, tokens: Optional[int]): - if tokens is not None and isinstance(tokens, int): + def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]): + if tokens is not None and isinstance(tokens, dict): if span.tokens is None: - span.tokens = tokens + span.tokens = LlmTokens(**tokens) else: - span.tokens += tokens + span.tokens.prompt_tokens += tokens["prompt_tokens"] + span.tokens.completion_tokens += tokens["completion_tokens"] + span.tokens.total_tokens += tokens["total_tokens"] def end_span(self, outputs: Dict[str, Any]): """ @@ -184,6 +190,7 @@ def end_span(self, outputs: Dict[str, Any]): """ if self.active_span is None: raise ValueError("There is no active span to end.") + self.active_span.end_time = datetime.now(timezone.utc) self.active_span.outputs = [outputs.get("message", "")] self._update_span_cost(self.active_span, outputs.get("cost", None)) From a35a98133c4621f693e18e2fce1674f36ad90b04 Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Thu, 6 Jun 2024 20:32:16 +0200 Subject: [PATCH 4/6] fix(sdk): Fixed double containing in cost --- agenta-cli/agenta/sdk/tracing/llm_tracing.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py index fb5b9719f9..8928ecffc3 100644 --- a/agenta-cli/agenta/sdk/tracing/llm_tracing.py +++ b/agenta-cli/agenta/sdk/tracing/llm_tracing.py @@ -169,11 +169,13 @@ def update_span_status(self, span: CreateSpan, value: str): span.status = value def _update_span_cost(self, span: CreateSpan, cost: Optional[float]): + print(f"Updating {span.name} span with cost: {cost}") if cost is not None and isinstance(cost, float): if span.cost is None: span.cost = cost else: span.cost += cost + print(f"New cost {span.name} {span.cost}") def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]): if tokens is not None and isinstance(tokens, dict): @@ -188,13 +190,18 @@ def end_span(self, outputs: Dict[str, Any]): """ Ends the active span, if it is a parent span, ends the trace too. """ + print(f"Ending span {self.active_span.name}") if self.active_span is None: raise ValueError("There is no active span to end.") self.active_span.end_time = datetime.now(timezone.utc) self.active_span.outputs = [outputs.get("message", "")] - self._update_span_cost(self.active_span, outputs.get("cost", None)) - self._update_span_tokens(self.active_span, outputs.get("usage", None)) + if self.active_span.spankind in [ + "LLM", + "RETRIEVER", + ]: # TODO: Remove this whole part. Setting the cost should be done through set_span_attribute + self._update_span_cost(self.active_span, outputs.get("cost", None)) + self._update_span_tokens(self.active_span, outputs.get("usage", None)) # Push span to list of recorded spans self.pending_spans.append(self.active_span) From a0cf1e8e0db76d24f62b8bd1ea18063c1160bf45 Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Thu, 6 Jun 2024 20:34:12 +0200 Subject: [PATCH 5/6] chore(sdk): Removed prints --- agenta-cli/agenta/sdk/tracing/llm_tracing.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py index 8928ecffc3..3ba58e27b0 100644 --- a/agenta-cli/agenta/sdk/tracing/llm_tracing.py +++ b/agenta-cli/agenta/sdk/tracing/llm_tracing.py @@ -169,13 +169,11 @@ def update_span_status(self, span: CreateSpan, value: str): span.status = value def _update_span_cost(self, span: CreateSpan, cost: Optional[float]): - print(f"Updating {span.name} span with cost: {cost}") if cost is not None and isinstance(cost, float): if span.cost is None: span.cost = cost else: span.cost += cost - print(f"New cost {span.name} {span.cost}") def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]): if tokens is not None and isinstance(tokens, dict): @@ -190,7 +188,6 @@ def end_span(self, outputs: Dict[str, Any]): """ Ends the active span, if it is a parent span, ends the trace too. """ - print(f"Ending span {self.active_span.name}") if self.active_span is None: raise ValueError("There is no active span to end.") From 12bde7c8a4f28ae1dd748757667fda0429ed5a11 Mon Sep 17 00:00:00 2001 From: Mahmoud Mabrouk Date: Thu, 6 Jun 2024 20:41:14 +0200 Subject: [PATCH 6/6] fix(sdk): Fixed token propagation in spans --- agenta-cli/agenta/sdk/tracing/llm_tracing.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py index 3ba58e27b0..628baffaee 100644 --- a/agenta-cli/agenta/sdk/tracing/llm_tracing.py +++ b/agenta-cli/agenta/sdk/tracing/llm_tracing.py @@ -176,6 +176,8 @@ def _update_span_cost(self, span: CreateSpan, cost: Optional[float]): span.cost += cost def _update_span_tokens(self, span: CreateSpan, tokens: Optional[dict]): + if isinstance(tokens, LlmTokens): + tokens = tokens.dict() if tokens is not None and isinstance(tokens, dict): if span.tokens is None: span.tokens = LlmTokens(**tokens)