From 03467aeaf29b216ed8ae72c84f4f945fbc4aa202 Mon Sep 17 00:00:00 2001
From: BaoVu2k4 <vuquangbao2004@gmail.com>
Date: Tue, 25 Aug 2026 16:44:55 +0700
Subject: [PATCH] Attach real input/output/token usage to provider spans for
 Langfuse cost tracking

Langfuse showed $0.00 cost on every trace even after tracing was wired up:
provider.bedrock_converse.* spans carried only latency, no content or usage.
Converse's own response already includes token usage -- it just wasn't being
read. Captured on the adapter (last_usage, no interface change) and attached
to the span using Langfuse's own OTel attribute convention
(langfuse.observation.*), which its docs say takes precedence over generic
GenAI attributes. Verified live: a fresh trace now shows real prompt/completion
tokens and Langfuse computes real cost once given the model's actual per-token
price.
---
 apps/ai-service/adapters/bedrock_converse.py |  6 +++
 apps/ai-service/rag/instrumentation.py       | 49 +++++++++++++++++++-
 2 files changed, 53 insertions(+), 2 deletions(-)

diff --git a/apps/ai-service/adapters/bedrock_converse.py b/apps/ai-service/adapters/bedrock_converse.py
index 8c078f5..68d4d0c 100644
--- a/apps/ai-service/adapters/bedrock_converse.py
+++ b/apps/ai-service/adapters/bedrock_converse.py
@@ -78,6 +78,10 @@ class BedrockConverseAnswerGenerator:
         self._client = client
         self._model_id = model_id
         self._max_tokens = max_tokens
+        # Converse's own response usage, captured for the observability layer
+        # (rag/instrumentation.py) to attach to its span. Not part of the
+        # `generate()` return contract -- every existing caller is unaffected.
+        self.last_usage: dict[str, int] | None = None
 
     @property
     def model_id(self) -> str:
@@ -126,6 +130,8 @@ class BedrockConverseAnswerGenerator:
                 f"{self._model_id} could not be invoked: {type(error).__name__}"
             ) from error
 
+        self.last_usage = dict(response.get("usage") or {}) or None
+
         if response.get("stopReason") in _EMPTY_STOP_REASONS:
             raise AnswerGenerationUnavailable(
                 f"{self._model_id} produced no usable content "
diff --git a/apps/ai-service/rag/instrumentation.py b/apps/ai-service/rag/instrumentation.py
index 5868202..1aed483 100644
--- a/apps/ai-service/rag/instrumentation.py
+++ b/apps/ai-service/rag/instrumentation.py
@@ -6,6 +6,7 @@ those files or duplicating their domain decisions.
 """
 from __future__ import annotations
 
+import json
 from typing import Any
 
 from .agent import RagAgent
@@ -42,6 +43,40 @@ def _failure_reason(exc: BaseException) -> str:
     return "error"
 
 
+_USAGE_KEY_MAP = {
+    "inputTokens": "input",
+    "outputTokens": "output",
+    "totalTokens": "total",
+}
+
+
+def _annotate_generation(
+    span: Any, model_id: str, system: str, user: str, output: str, usage: dict | None
+) -> None:
+    """Attach real request/response content to a provider span.
+
+    Attribute names are Langfuse's own OTel convention (`langfuse.observation.*`),
+    which its docs say take precedence over generic OpenTelemetry GenAI
+    attributes -- so this is read reliably regardless of what else instruments
+    the same span. Only real values from the actual call are recorded; usage is
+    omitted entirely (not estimated) when the provider response carried none.
+    """
+    span.set_attribute("langfuse.observation.model.name", model_id)
+    span.set_attribute(
+        "langfuse.observation.input",
+        json.dumps({"system": system, "user": user}, ensure_ascii=False),
+    )
+    span.set_attribute("langfuse.observation.output", output)
+    if usage:
+        span.set_attribute(
+            "langfuse.observation.usage_details",
+            json.dumps(
+                {_USAGE_KEY_MAP.get(k, k): v for k, v in usage.items()},
+                ensure_ascii=False,
+            ),
+        )
+
+
 class InstrumentedGenerator:
     def __init__(self, delegate: Any, metrics: Metrics) -> None:
         self._delegate = delegate
@@ -58,9 +93,9 @@ class InstrumentedGenerator:
             "generation": "generate",
             "entailment": "entailment",
         }.get(current_stage(), "generate")
-        with dependency_span(self._provider, operation):
+        with dependency_span(self._provider, operation) as span:
             try:
-                return self._delegate.generate(system, user, schema)
+                result = self._delegate.generate(system, user, schema)
             except Exception as exc:
                 self._metrics.increment(
                     PROVIDER_FAILURE,
@@ -69,6 +104,16 @@ class InstrumentedGenerator:
                     reason=_failure_reason(exc),
                 )
                 raise
+            if span is not None:
+                _annotate_generation(
+                    span,
+                    self.model_id,
+                    system,
+                    user,
+                    result,
+                    getattr(self._delegate, "last_usage", None),
+                )
+            return result
 
 
 class InstrumentedEmbedder: