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

@@ -355,6 +355,9 @@ def filter_transactions(event, _hint):
 
 LOG_LEVEL = env.str("DJANGO_LOG_LEVEL", "WARNING")
 LOG_FORMAT = env.str("DJANGO_LOG_FORMAT", "asim_formatter")
+DD_TRACE_DEBUG_ENABLED = env.bool("DD_TRACE_DEBUG", default=False)
+DD_TRACE_LOG_LEVEL = env.str("DD_TRACE_LOG_LEVEL", "DEBUG" if DD_TRACE_DEBUG_ENABLED else "ERROR")
+DD_TRACE_LOG_HANDLER = "console" if DD_TRACE_DEBUG_ENABLED else "asim"
 LOGGING = {
     "version": 1,
     "disable_existing_loggers": False,
@@ -412,8 +415,8 @@ def filter_transactions(event, _hint):
             "level": "WARNING",
         },
         "ddtrace": {
-            "handlers": ["asim"],
-            "level": "ERROR",
+            "handlers": [DD_TRACE_LOG_HANDLER],
+            "level": DD_TRACE_LOG_LEVEL,
             "propagate": False,
         },
     },
@@ -493,7 +496,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)
 

django_app/start.sh

@@ -2,6 +2,13 @@
 
 PORT=8080
 
+if [ "${REDBOX_ENABLE_DD_DIAGNOSTICS:-false}" = "1" ] || [ "${REDBOX_ENABLE_DD_DIAGNOSTICS:-false}" = "true" ] || [ "${REDBOX_ENABLE_DD_DIAGNOSTICS:-false}" = "TRUE" ]; then
+	export DD_TRACE_DEBUG=true
+	export DD_TRACE_LOG_LEVEL=DEBUG
+	export REDBOX_BEDROCK_DIAGNOSTICS=true
+	echo "Datadog diagnostics mode enabled"
+fi
+
 venv/bin/django-admin migrate
 venv/bin/django-admin collectstatic --noinput
 venv/bin/django-admin create_admin_user

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", {}