chore: add structured span elements for AI observability

django_app/pyproject.toml

@@ -58,7 +58,7 @@ mohawk = "^1.1.0"
 uwotm8 = "^0.1.6"
 pymupdf = "^1.27.2"
 openpyxl = "^3.1.5"
-ddtrace = "^4.10.3"
+ddtrace = "^4.11.0"
 django-requestlogs = "^0.8"
 unoserver = "^3.6"
 coverage = "^7.13.5"

django_app/redbox_app/settings.py

@@ -493,7 +493,7 @@ def filter_transactions(event, _hint):
 
 EMBEDDING_BACKEND = env.str("EMBEDDING_BACKEND", "amazon.titan-embed-text-v2:0")
 
-DEFAULT_MODEL_ID = env.str("DEFAULT_MODEL_ID", "anthropic.claude-3-sonnet-20240229-v1:0")
+DEFAULT_MODEL_ID = env.str("DEFAULT_MODEL_ID", "claude-sonnet-4-6")
 
 WEB_SEARCH_API_LIMIT = env.int("WEB_SEARCH_API_LIMIT", 100)
 

pyproject.toml

@@ -20,7 +20,7 @@ django = "5.2.13"
 django-staff-sso-client = "^5.2.0"
 dbt-copilot-python = "^2.1.0"
 uwotm8 = "^0.1.6"
-ddtrace = "^4.10.3"
+ddtrace = "^4.11.0"
 django-log-formatter-asim = "^1.3.0"
 django-requestlogs = "^0.8"
 openpyxl = "^3.1.5"

redbox/pyproject.toml

@@ -33,7 +33,7 @@ mohawk = "^1.1.0"
 django-waffle = "^4.2.0"
 langchain-mcp-adapters = "^0.0.11"
 pymupdf = "^1.26.0"
-ddtrace = "^4.10.3"
+ddtrace = "^4.11.0"
 django-requestlogs = "^0.8"
 django-log-formatter-asim = "^1.3.0"
 pytest = "^8.4.2"

redbox/redbox/chains/components.py

@@ -30,7 +30,7 @@
     KnowledgeBaseTabularMetadataRetriever,
     TabularMetadataRetriever,
 )
-from redbox.transform import bedrock_tokeniser
+from redbox.transform import bedrock_tokeniser, ensure_bedrock_client_token_metrics
 
 logger = logging.getLogger(__name__)
 load_dotenv()
@@ -46,6 +46,8 @@ def get_chat_llm(
     ai_settings: AISettings = AISettings(),
     tools: list[StructuredTool] | None = None,
 ):
+    ensure_bedrock_client_token_metrics()
+
     fallback_backend = ChatLLMBackend(
         name="anthropic.claude-3-7-sonnet-20250219-v1:0",
         provider="bedrock",

redbox/redbox/graph/nodes/tools.py

@@ -43,7 +43,13 @@
     get_knowledge_base,
 )
 from redbox.retriever.retrievers import SchematisedTabularChunkRetriever, query_to_documents
-from redbox.transform import bedrock_tokeniser, merge_documents, sort_documents
+from redbox.transform import (
+    annotate_span_with_token_metrics,
+    bedrock_tokeniser,
+    ensure_bedrock_client_token_metrics,
+    merge_documents,
+    sort_documents,
+)
 
 log = logging.getLogger(__name__)
 
@@ -680,11 +686,12 @@ def _search_wikipedia(query: str) -> tuple[str, list[Document]]:
 
 @waffle_flag("DATA_HUB_API_ROUTE_ON")
 def parse_filters_bedrock(prompt: str):
+    ensure_bedrock_client_token_metrics()
     client = boto3.client("bedrock-runtime", region_name="eu-west-2")
 
     settings = get_settings()
 
-    model_id = settings.default_model_id if settings.default_model_id else "anthropic.claude-3-sonnet-20240229-v1:0"
+    model_id = settings.default_model_id if settings.default_model_id else "claude-sonnet-4-6"
 
     response = client.invoke_model(
         modelId=model_id,
@@ -719,9 +726,25 @@ def parse_filters_bedrock(prompt: str):
 
     body = json.loads(response["body"].read())
     try:
-        response_json = json.loads(body["content"][0]["text"].strip())
+        response_text = body["content"][0]["text"].strip()
+        response_json = json.loads(response_text)
+        usage = body.get("usage", {})
+        input_tokens = usage.get("input_tokens", bedrock_tokeniser(prompt))
+        output_tokens = usage.get("output_tokens", bedrock_tokeniser(response_text))
+        annotate_span_with_token_metrics(
+            model=model_id,
+            input_tokens=input_tokens,
+            output_tokens=output_tokens,
+            provider="bedrock",
+        )
         return response_json.get("dataset", "companies-dataset"), response_json.get("filters", {})
     except Exception:
+        annotate_span_with_token_metrics(
+            model=model_id,
+            input_tokens=bedrock_tokeniser(prompt),
+            output_tokens=bedrock_tokeniser(json.dumps(body)),
+            provider="bedrock",
+        )
         return "companies-dataset", {}
 
 

redbox/redbox/transform.py

@@ -1,10 +1,14 @@
 import itertools
+import json
 import logging
-import re
 import math
+import re
+from functools import wraps
 from typing import Dict, Iterable
 from uuid import NAMESPACE_DNS, UUID, uuid5
 
+import boto3
+from ddtrace import tracer
 from langchain_core.callbacks.manager import dispatch_custom_event
 from langchain_core.documents import Document
 from langchain_core.messages import AIMessage, AnyMessage
@@ -15,6 +19,238 @@
 
 log = logging.getLogger(__name__)
 
+_BEDROCK_CONTEXT_KEY = "_redbox_bedrock_api_params"
+_BEDROCK_CLIENT_PATCHED = False
+
+
+def _serialize_text(text: bytes | bytearray | str | None) -> str:
+    if text is None:
+        return ""
+    if isinstance(text, (bytes, bytearray)):
+        try:
+            return text.decode("utf-8")
+        except Exception:
+            return text.decode("utf-8", errors="ignore")
+    return str(text)
+
+
+def _bedrock_request_text(body: str) -> str:
+    if not body:
+        return ""
+
+    try:
+        payload = json.loads(body)
+    except Exception:
+        return body
+
+    if isinstance(payload, dict):
+        if "messages" in payload and isinstance(payload["messages"], list):
+            return " ".join(
+                str(message.get("content", "")) for message in payload["messages"] if isinstance(message, dict)
+            )
+        if "input" in payload:
+            return str(payload["input"])
+        if "content" in payload:
+            return str(payload["content"])
+    return body
+
+
+def _bedrock_response_text(parsed: dict | None, http_response: object | None) -> str:
+    if isinstance(parsed, dict):
+        content = parsed.get("content")
+        if isinstance(content, list) and content:
+            first = content[0]
+            if isinstance(first, dict):
+                return str(first.get("text", ""))
+            if isinstance(first, str):
+                return first
+        if isinstance(content, str):
+            return content
+
+    if http_response is not None:
+        content_attr = getattr(http_response, "content", None)
+        if isinstance(content_attr, (bytes, bytearray)):
+            return _serialize_text(content_attr)
+
+    return ""
+
+
+def annotate_span_with_token_metrics(model: str, input_tokens: int, output_tokens: int, provider: str = "bedrock"):
+    span = tracer.current_span()
+    if span is None:
+        return
+
+    total_tokens = input_tokens + output_tokens
+
+    span.set_tag("input_tokens", input_tokens)
+    span.set_tag("output_tokens", output_tokens)
+    span.set_tag("total_tokens", total_tokens)
+    span.set_tag("model", model)
+    span.set_tag("llm.model", model)
+    span.set_tag("provider", provider)
+    span.set_tag("llm.provider", provider)
+
+    span.set_metric("input_tokens", input_tokens)
+    span.set_metric("output_tokens", output_tokens)
+    span.set_metric("total_tokens", total_tokens)
+
+    span.set_metric("gen_ai.usage.input_tokens", input_tokens)
+    span.set_metric("gen_ai.usage.output_tokens", output_tokens)
+    span.set_metric("gen_ai.usage.total_tokens", total_tokens)
+
+
+def _extract_text_from_content_blocks(content: list | str | None) -> str:
+    if isinstance(content, str):
+        return content
+    if not isinstance(content, list):
+        return ""
+
+    parts = []
+    for item in content:
+        if isinstance(item, str):
+            parts.append(item)
+        elif isinstance(item, dict):
+            text = item.get("text")
+            if text:
+                parts.append(str(text))
+    return " ".join(parts)
+
+
+def _bedrock_request_text_from_params(api_params: dict | None) -> str:
+    if not isinstance(api_params, dict):
+        return ""
+
+    if body := api_params.get("body"):
+        return _bedrock_request_text(_serialize_text(body))
+
+    if messages := api_params.get("messages"):
+        return " ".join(
+            _extract_text_from_content_blocks(message.get("content"))
+            for message in messages
+            if isinstance(message, dict)
+        )
+
+    if system := api_params.get("system"):
+        return _extract_text_from_content_blocks(system)
+
+    if input_text := api_params.get("inputText"):
+        return str(input_text)
+
+    if input_body := api_params.get("input"):
+        return str(input_body)
+
+    try:
+        return json.dumps(api_params)
+    except TypeError:
+        return str(api_params)
+
+
+def _bedrock_response_text_from_parsed(parsed: dict | None) -> str:
+    if not isinstance(parsed, dict):
+        return ""
+
+    if output := parsed.get("output"):
+        if isinstance(output, dict):
+            if message := output.get("message"):
+                if isinstance(message, dict):
+                    return _extract_text_from_content_blocks(message.get("content"))
+            if text := output.get("text"):
+                return str(text)
+
+    return _bedrock_response_text(parsed=parsed, http_response=None)
+
+
+def _usage_value(usage: dict | None, *keys: str) -> int | None:
+    if not isinstance(usage, dict):
+        return None
+
+    for key in keys:
+        value = usage.get(key)
+        if isinstance(value, (int, float)):
+            return int(value)
+    return None
+
+
+def _bedrock_model_from_params(api_params: dict | None) -> str:
+    if not isinstance(api_params, dict):
+        return "unknown"
+
+    return str(api_params.get("modelId") or api_params.get("model_id") or "unknown")
+
+
+def _capture_bedrock_request_for_metrics(model=None, params=None, context=None, **kwargs):
+    if isinstance(context, dict):
+        context[_BEDROCK_CONTEXT_KEY] = params if isinstance(params, dict) else {}
+
+
+def _annotate_bedrock_response_metrics(parsed=None, context=None, **kwargs):
+    api_params = context.get(_BEDROCK_CONTEXT_KEY, {}) if isinstance(context, dict) else {}
+    usage = parsed.get("usage", {}) if isinstance(parsed, dict) else {}
+
+    input_tokens = _usage_value(usage, "input_tokens", "inputTokens")
+    output_tokens = _usage_value(usage, "output_tokens", "outputTokens")
+
+    if input_tokens is None:
+        request_text = _bedrock_request_text_from_params(api_params)
+        if request_text:
+            input_tokens = bedrock_tokeniser(request_text)
+
+    if output_tokens is None:
+        response_text = _bedrock_response_text_from_parsed(parsed)
+        if response_text:
+            output_tokens = bedrock_tokeniser(response_text)
+
+    if input_tokens is None or output_tokens is None:
+        return
+
+    annotate_span_with_token_metrics(
+        model=_bedrock_model_from_params(api_params),
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        provider="bedrock",
+    )
+
+
+def _register_bedrock_client_token_handlers(client):
+    meta = getattr(client, "meta", None)
+    if meta is None or getattr(meta, "service_model", None) is None:
+        return client
+
+    if meta.service_model.service_name != "bedrock-runtime":
+        return client
+
+    if getattr(meta, "_redbox_token_metrics_registered", False):
+        return client
+
+    meta.events.register("before-call.bedrock-runtime", _capture_bedrock_request_for_metrics)
+    meta.events.register("after-call.bedrock-runtime", _annotate_bedrock_response_metrics)
+    meta._redbox_token_metrics_registered = True
+    return client
+
+
+def ensure_bedrock_client_token_metrics():
+    global _BEDROCK_CLIENT_PATCHED
+
+    if _BEDROCK_CLIENT_PATCHED:
+        return
+
+    original_boto3_client = boto3.client
+    original_session_client = boto3.session.Session.client
+
+    @wraps(original_boto3_client)
+    def _patched_boto3_client(*args, **kwargs):
+        client = original_boto3_client(*args, **kwargs)
+        return _register_bedrock_client_token_handlers(client)
+
+    @wraps(original_session_client)
+    def _patched_session_client(self, *args, **kwargs):
+        client = original_session_client(self, *args, **kwargs)
+        return _register_bedrock_client_token_handlers(client)
+
+    boto3.client = _patched_boto3_client
+    boto3.session.Session.client = _patched_session_client
+    _BEDROCK_CLIENT_PATCHED = True
+
 
 def bedrock_tokeniser_tokens(text: str) -> list[str]:
     # Simple tokeniser that counts the number of words in the text
@@ -250,6 +486,13 @@ def to_request_metadata(obj: dict) -> RequestMetadata:
     except Exception:
         output_tokens = len(response[0].get("text", []))
 
+    annotate_span_with_token_metrics(
+        model=model,
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        provider="bedrock",
+    )
+
     metadata_event = RequestMetadata(
         llm_calls=[
             LLMCallMetadata(