Data dog llm obs improvement

.secrets.baseline

@@ -158,47 +158,47 @@
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "3a07fa5a98e5eab18d39dd37f33148d45f8091dc",
         "is_verified": false,
-        "line_number": 203,
+        "line_number": 205,
         "is_secret": false
       },
       {
         "type": "Base64 High Entropy String",
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "71ecb7a9026acd0b5c75ce1cd543c411eaa78a75",
         "is_verified": false,
-        "line_number": 204,
+        "line_number": 206,
         "is_secret": false
       },
       {
         "type": "Base64 High Entropy String",
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "d2aa34aa9083c20fc7eeb1f2e08dfd7f8f6c0022",
         "is_verified": false,
-        "line_number": 205,
+        "line_number": 207,
         "is_secret": false
       },
       {
         "type": "Base64 High Entropy String",
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "28c0d52cedc09b9e794d0e1033ca0b2a06def89a",
         "is_verified": false,
-        "line_number": 206,
+        "line_number": 208,
         "is_secret": false
       },
       {
         "type": "Base64 High Entropy String",
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "8782f26f9343343d50facf336a0befad925537d1",
         "is_verified": false,
-        "line_number": 207,
+        "line_number": 209,
         "is_secret": false
       },
       {
         "type": "Base64 High Entropy String",
         "filename": "django_app/redbox_app/settings.py",
         "hashed_secret": "1bc89245f6e26d516ddc579b00a0d59e4096a765",
         "is_verified": false,
-        "line_number": 216,
+        "line_number": 218,
         "is_secret": false
       }
     ],
@@ -281,5 +281,5 @@
       }
     ]
   },
-  "generated_at": "2026-04-09T11:20:11Z"
+  "generated_at": "2026-04-10T17:06:15Z"
 }

django_app/redbox_app/settings.py

@@ -10,6 +10,8 @@
 import sentry_sdk
 from dbt_copilot_python.database import database_from_env
 from dbt_copilot_python.error_tracking import DatadogErrorTrackingFilter
+from ddtrace import patch
+from ddtrace.llmobs import LLMObs
 from django.urls import reverse_lazy
 from django_log_formatter_asim import ASIMFormatter
 from dotenv import find_dotenv, load_dotenv
@@ -505,3 +507,8 @@ def filter_transactions(event, _hint):
 )
 
 PRODUCT_NAME = env.str("PRODUCT_NAME", "Redbox at DBT")
+
+# datadog
+# enable llm manual instrument
+LLMObs.enable(integrations_enabled=False, api_key=env.str("DATADOG_API_KEY", "Fake"))
+patch(langchain=True, langgraph=True, mcp=True, botocore=False)

django_app/start.sh

@@ -7,5 +7,5 @@ venv/bin/django-admin collectstatic --noinput
 venv/bin/django-admin create_admin_user
 
 echo "Starting daphne on port $PORT"
-#venv/bin/daphne --websocket_timeout 86400 -b 0.0.0.0 -p $PORT redbox_app.asgi:application
-venv/bin/ddtrace-run venv/bin/daphne --websocket_timeout 86400 -b 0.0.0.0 -p $PORT redbox_app.asgi:application
+venv/bin/daphne --websocket_timeout 86400 -b 0.0.0.0 -p $PORT redbox_app.asgi:application
+# venv/bin/ddtrace-run venv/bin/daphne --websocket_timeout 86400 -b 0.0.0.0 -p $PORT redbox_app.asgi:application

redbox/redbox/chains/runnables.py

@@ -2,6 +2,8 @@
 import re
 from typing import Any, Callable, Iterable, Iterator
 
+from ddtrace.llmobs import LLMObs
+from ddtrace.trace import tracer
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun, dispatch_custom_event
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
@@ -314,6 +316,16 @@ def _llm_type(self) -> str:
         return "canned"
 
 
+def extract_model_name(llm_output) -> str:
+    model_name = (llm_output.response_metadata or {}).get("model_name", "unspecified_model")
+    if model_name:
+        if model_name.startswith("arn"):
+            # grab the last bit
+            model_name = model_name.split("/")[-1]
+        model_name = model_name.split(".")[-1]
+    return model_name
+
+
 def basic_chat_chain(
     system_prompt,
     tools=None,
@@ -344,22 +356,36 @@ def _basic_chat_chain(state: RedboxState):
             "question": state.request.question,
             "chat_history": truncated_history if using_chat_history else "",
         } | _additional_variables
-        if parser:
-            if isinstance(parser, StrOutputParser):
-                prompt = ChatPromptTemplate([(system_prompt)])
-            else:
-                format_instructions = parser.get_format_instructions()
-                prompt = ChatPromptTemplate(
-                    [(system_prompt)], partial_variables={"format_instructions": format_instructions}
-                )
-            if using_only_structure:
-                chain = prompt | llm
-            else:
-                chain = prompt | llm | parser
+        if parser and not isinstance(parser, StrOutputParser):
+            format_instructions = parser.get_format_instructions()
+            prompt = ChatPromptTemplate(
+                [(system_prompt)], partial_variables={"format_instructions": format_instructions}
+            )
         else:
             prompt = ChatPromptTemplate([(system_prompt)])
-            chain = prompt | llm
-        return chain.invoke(context)
+        chain = prompt | llm
+
+        output = chain.invoke(context)
+        bedrock_span = tracer.current_span()
+        # formatted_input_dd = prompt.format_messages(**context)
+        # role_map = {"human": "user", "ai": "assistant", "system": "system"}
+        LLMObs.annotate(
+            span=bedrock_span,
+            metadata={
+                "max_tokens": (llm._default_config or {}).get("max_tokens", None),
+                "stop_reason": (output.response_metadata or {}).get("stop_reason", None),
+            },
+            metrics={
+                "input_tokens": (output.usage_metadata or {}).get("input_tokens", None),
+                "output_tokens": (output.usage_metadata or {}).get("output_tokens", None),
+                "total_tokens": (output.usage_metadata or {}).get("total_tokens", None),
+            },
+            tags={"func": "basic_chat_chain"},
+        )
+
+        if parser and not using_only_structure:
+            output = parser.invoke(output)
+        return output
 
     return _basic_chat_chain
 

redbox/redbox/graph/agents/workers.py

@@ -149,5 +149,5 @@ def execute(self):
         return (
             self.reading_task_info()
             | RunnableParallel(state=self.log_agent_activity(), result=self.core_task() | self.post_processing())
-            | (lambda x: x["result"])  # Return only the result
-        )
+            | (lambda x: x["result"])
+        )  # Return only the result

redbox/redbox/graph/nodes/tools.py

@@ -14,13 +14,18 @@
 import numpy as np
 import pandas as pd
 import requests
+from ddtrace.llmobs import LLMObs
+from ddtrace.trace import tracer
 from elasticsearch import Elasticsearch
 from langchain_community.utilities import WikipediaAPIWrapper
 from langchain_core.documents import Document
 from langchain_core.embeddings.embeddings import Embeddings
 from langchain_core.messages import ToolCall
 from langchain_core.tools import Tool, tool
+from langchain_mcp_adapters.tools import load_mcp_tools
 from langgraph.prebuilt import InjectedState
+from mcp import ClientSession
+from mcp.client.streamable_http import streamablehttp_client
 from mohawk import Sender
 from opensearchpy import OpenSearch
 from sklearn.metrics.pairwise import cosine_similarity
@@ -41,9 +46,6 @@
 )
 from redbox.retriever.retrievers import SchematisedTabularChunkRetriever, query_to_documents
 from redbox.transform import bedrock_tokeniser, merge_documents, sort_documents
-from mcp import ClientSession
-from mcp.client.streamable_http import streamablehttp_client
-from langchain_mcp_adapters.tools import load_mcp_tools
 
 log = logging.getLogger(__name__)
 
@@ -220,10 +222,18 @@ def search_repo(query, selected_files, permitted_files, ai_settings, start_time=
             "[_search_documents] Initial query using %s seconds",
             time.time() - start_time,
         )
+        metrics = {
+            "initial_query_time": None,
+            "boosted_query_time": None,
+            "merged_sort_docuemnt_time": None,
+            "no_returned_documents": 0,
+        }
+
+        metrics["initial_query_time"] = time.time() - start_time
 
         # Handle nothing found (as when no files are permitted)
         if not initial_documents:
-            return "", []
+            return "", [], metrics
 
         # Adjacent documents
         with_adjacent_query = add_document_filter_scores_to_query(
@@ -236,6 +246,7 @@ def search_repo(query, selected_files, permitted_files, ai_settings, start_time=
             "[_search_documents] Adjacent boosted query using %s seconds",
             time.time() - start_time,
         )
+        metrics["boosted_query_time"] = time.time() - metrics["initial_query_time"]
 
         # Merge and sort
         merged_documents = merge_documents(initial=initial_documents, adjacent=adjacent_boosted)
@@ -245,9 +256,11 @@ def search_repo(query, selected_files, permitted_files, ai_settings, start_time=
             time.time() - start_time,
         )
         log.warning("[_search_documents] Returning %s documents", len(sorted_documents))
+        metrics["merged_sort_docuemnt_time"] = time.time() - metrics["boosted_query_time"]
+        metrics["no_returned_documents"] = len(sorted_documents)
 
         # Return as state update
-        return format_documents(sorted_documents), sorted_documents
+        return format_documents(sorted_documents), sorted_documents, metrics
 
     @tool(response_format="content_and_artifact")
     def _search_documents(query: str, state: Annotated[RedboxState, InjectedState]) -> tuple[str, list[Document]]:
@@ -264,13 +277,25 @@ def _search_documents(query: str, state: Annotated[RedboxState, InjectedState])
         Returns:
             dict[str, Any]: Collection of matching document snippets with metadata:
         """
-        return search_repo(
+
+        document, artifact, metrics = search_repo(
             query=query,
             selected_files=state.request.s3_keys,
             permitted_files=state.request.permitted_s3_keys,
             ai_settings=state.request.ai_settings,
         )
 
+        LLMObs.annotate(
+            span=tracer.current_span(),
+            input_data=query,
+            output_data=document,
+            metadata={"custom": metrics, "test": "something"},
+            metrics=metrics,
+            tags={"func": "hello-world"},
+        )
+
+        return document, artifact
+
     @tool(response_format="content_and_artifact")
     def _search_knowledge_base(query: str, state: Annotated[RedboxState, InjectedState]) -> tuple[str, list[Document]]:
         """
@@ -286,13 +311,23 @@ def _search_knowledge_base(query: str, state: Annotated[RedboxState, InjectedSta
         Returns:
             dict[str, Any]: Collection of matching document snippets with metadata:
         """
-        return search_repo(
+        document, artifact, metrics = search_repo(
             query=query,
             selected_files=state.request.knowledge_base_s3_keys,
             permitted_files=state.request.knowledge_base_s3_keys,
             ai_settings=state.request.ai_settings,
         )
 
+        LLMObs.annotate(
+            span=tracer.current_span(),
+            input_data=query,
+            output_data=document,
+            metrics=metrics,
+            tags={"func": "hello-world"},
+        )
+
+        return document, artifact
+
     return _search_documents if repository == "user_uploaded" else _search_knowledge_base
 
 

redbox/redbox/graph/root.py

@@ -22,8 +22,8 @@
 from redbox.graph.nodes.processes import (
     build_activity_log_node,
     build_agent_with_loop,
-    build_datahub_agent_with_loop,
     build_chat_pattern,
+    build_datahub_agent_with_loop,
     build_error_pattern,
     build_merge_pattern,
     build_passthrough_pattern,

redbox/redbox/test/data.py

@@ -12,11 +12,10 @@
 from langchain_core.tools import BaseTool
 from pydantic.v1 import BaseModel, Field, validator
 
-from redbox.models.chain import RedboxQuery
+from redbox.models.chain import MultiAgentPlanBase, RedboxQuery
 from redbox.models.chat import ChatRoute, ErrorRoute
 from redbox.models.file import ChunkResolution, TabularSchema, UploadedFileMetadata
 from redbox.models.graph import RedboxActivityEvent
-from redbox.models.chain import MultiAgentPlanBase
 
 log = logging.getLogger()
 
@@ -243,6 +242,10 @@ class GenericFakeChatModelWithTools(GenericFakeChatModel):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
+    @property
+    def _default_config(self):
+        return {}
+
     def bind_tools(
         self,
         tools: Sequence[dict[str, Any] | type | Callable | BaseTool],

redbox/tests/graph/test_app.py

@@ -17,17 +17,17 @@
 from pytest_mock import MockerFixture
 
 from redbox import Redbox
-from redbox.graph.nodes.processes import create_or_update_db_from_tabulars, check_if_task_requires_user_feedback
+from redbox.graph.nodes.processes import check_if_task_requires_user_feedback, create_or_update_db_from_tabulars
 from redbox.models.chain import (
     AISettings,
     DocumentState,
     RedboxQuery,
     RedboxState,
     RequestMetadata,
     StructuredResponseWithCitations,
+    TaskStatus,
     configure_agent_task_plan,
     metadata_reducer,
-    TaskStatus,
 )
 from redbox.models.chat import ChatRoute, ErrorRoute
 from redbox.models.graph import RedboxActivityEvent
@@ -224,7 +224,6 @@ def _search_govuk(query: str) -> dict[str, Any]:
 
     # Mock the LLM and relevant tools
     llm = GenericFakeChatModelWithTools(messages=iter(test_case.test_data.llm_responses))
-    llm._default_config = {"model": "bedrock"}
     mocker.patch("redbox.graph.nodes.processes.get_chat_llm", return_value=llm)
 
     # Instantiate app