feat(datadog_agent source): accept LLMObs telemetry at /api/v2/llmobs

changelog.d/25441_llmobs_endpoint.feature.md

@@ -0,0 +1,3 @@
+The `datadog_agent` source now accepts LLM Observability (LLMObs) telemetry at `/api/v2/llmobs`. When `multiple_outputs` is enabled, LLMObs span events are available as log events on the `llmobs` output port.
+
+authors: ronitanilkumar

src/sources/datadog_agent/llmobs.rs

@@ -0,0 +1,157 @@
+use bytes::Bytes;
+use http::StatusCode;
+use serde::Deserialize;
+use serde_json::Value;
+use std::sync::Arc;
+use warp::{Filter, filters::BoxedFilter, path, path::FullPath, reply::Response};
+
+use super::{ApiKeyQueryParams, DatadogAgentSource, RequestHandler};
+use crate::{
+    common::http::ErrorMessage,
+    event::{Event, LogEvent},
+    internal_events::DatadogAgentJsonParseError,
+};
+
+pub(super) fn build_warp_filter(
+    handler: RequestHandler,
+    source: DatadogAgentSource,
+) -> BoxedFilter<(Response,)> {
+    warp::post()
+        .and(path!("api" / "v2" / "llmobs" / ..))
+        .and(warp::path::full())
+        .and(warp::header::optional::<String>("content-encoding"))
+        .and(warp::header::optional::<String>("dd-api-key"))
+        .and(warp::query::<ApiKeyQueryParams>())
+        .and(warp::body::bytes())
+        .and_then(
+            move |path: FullPath,
+                  encoding_header: Option<String>,
+                  api_token: Option<String>,
+                  query_params: ApiKeyQueryParams,
+                  body: Bytes| {
+                let events = source
+                    .decode(&encoding_header, body, path.as_str())
+                    .and_then(|body| {
+                        decode_llmobs_body(
+                            body,
+                            source.api_key_extractor.extract(
+                                path.as_str(),
+                                api_token,
+                                query_params.dd_api_key,
+                            ),
+                        )
+                    });
+                handler.clone().handle_request(events, super::LLMOBS)
+            },
+        )
+        .boxed()
+}
+
+#[derive(Deserialize)]
+struct LLMObsEnvelopeItem {
+    #[serde(rename = "event_type")]
+    _event_type: Option<String>,
+    spans: Vec<LLMObsSpan>,
+    #[serde(rename = "_dd.tracer_version")]
+    dd_tracer_version: Option<String>,
+    #[serde(rename = "_dd.scope")]
+    _dd_scope: Option<String>,
+}
+
+#[derive(Deserialize)]
+struct LLMObsSpan {
+    span_id: String,
+    trace_id: String,
+    parent_id: Option<String>,
+    name: Option<String>,
+    session_id: Option<String>,
+    service: Option<String>,
+    start_ns: Option<i64>,
+    duration: Option<i64>,
+    status: Option<String>,
+    status_message: Option<String>,
+    meta: Option<Value>,
+    metrics: Option<Value>,
+    #[serde(default)]
+    tags: Vec<String>,
+    #[serde(rename = "_dd")]
+    dd: Option<Value>,
+}
+
+pub(crate) fn decode_llmobs_body(
+    body: Bytes,
+    api_key: Option<Arc<str>>,
+) -> Result<Vec<Event>, ErrorMessage> {
+    let envelope: Vec<LLMObsEnvelopeItem> = serde_json::from_slice(&body).map_err(|error| {
+        emit!(DatadogAgentJsonParseError { error: &error });
+        ErrorMessage::new(
+            StatusCode::BAD_REQUEST,
+            format!("Error parsing JSON: {error:?}"),
+        )
+    })?;
+
+    let events = envelope
+        .into_iter()
+        .flat_map(|item| {
+            let tracer_version = item.dd_tracer_version.clone();
+            item.spans.into_iter().map(move |span| {
+                let mut log = LogEvent::default();
+                log.insert("span_id", span.span_id);
+                log.insert("trace_id", span.trace_id);
+                if let Some(v) = span.parent_id {
+                    log.insert("parent_id", v);
+                }
+                if let Some(v) = span.name {
+                    log.insert("name", v);
+                }
+                if let Some(v) = span.session_id {
+                    log.insert("session_id", v);
+                }
+                if let Some(v) = span.service {
+                    log.insert("service", v);
+                }
+                if let Some(v) = span.start_ns {
+                    log.insert("start_ns", v);
+                }
+                if let Some(v) = span.duration {
+                    log.insert("duration", v);
+                }
+                if let Some(v) = span.status {
+                    log.insert("status", v);
+                }
+                if let Some(v) = span.status_message {
+                    log.insert("status_message", v);
+                }
+                if let Some(v) = span.meta {
+                    log.insert("meta", v);
+                }
+                if let Some(v) = span.metrics {
+                    log.insert("metrics", v);
+                }
+                if !span.tags.is_empty() {
+                    log.insert("tags", span.tags);
+                }
+                if let Some(ml_app) = span
+                    .dd
+                    .as_ref()
+                    .and_then(|dd| dd.get("ml_app"))
+                    .and_then(|v| v.as_str())
+                {
+                    log.insert("ml_app", ml_app.to_owned());
+                }
+                if let Some(v) = tracer_version.clone() {
+                    log.insert("_dd.tracer_version", v);
+                }
+                Event::Log(log)
+            })
+        })
+        .map(|mut event| {
+            if let Some(k) = &api_key {
+                event.metadata_mut().set_datadog_api_key(Arc::clone(k));
+            }
+            event
+        })
+        .collect();
+
+    Ok(events)
+}

src/sources/datadog_agent/mod.rs

@@ -6,6 +6,7 @@ mod tests;
 pub mod logs;
 pub mod metrics;
 pub mod traces;
+pub mod llmobs;
 
 #[allow(warnings, clippy::pedantic, clippy::nursery)]
 pub(crate) mod ddmetric_proto {
@@ -69,6 +70,7 @@ use crate::{
 pub const LOGS: &str = "logs";
 pub const METRICS: &str = "metrics";
 pub const TRACES: &str = "traces";
+pub const LLMOBS: &str = "llmobs";
 
 /// Configuration for the `datadog_agent` source.
 #[configurable_component(source(
@@ -106,6 +108,11 @@ pub struct DatadogAgentConfig {
     #[serde(default = "crate::serde::default_false")]
     disable_traces: bool,
 
+    /// If this is set to `true`, LLM Observability events are not accepted by the component.
+    #[configurable(metadata(docs::advanced))]
+    #[serde(default = "crate::serde::default_false")]
+    disable_llmobs: bool,
+
     /// If this is set to `true`, logs, metrics (beta), and traces (alpha) are sent to different outputs.
     ///
     ///
@@ -179,6 +186,7 @@ impl GenerateConfig for DatadogAgentConfig {
             disable_logs: false,
             disable_metrics: false,
             disable_traces: false,
+            disable_llmobs: false,
             multiple_outputs: false,
             parse_ddtags: false,
             split_metric_namespace: true,
@@ -322,6 +330,7 @@ impl SourceConfig for DatadogAgentConfig {
             .with_standard_vector_source_metadata();
 
         let mut output = Vec::with_capacity(1);
+        let llmobs_definition = definition.clone();
 
         if self.multiple_outputs {
             if !self.disable_logs {
@@ -333,6 +342,9 @@ impl SourceConfig for DatadogAgentConfig {
             if !self.disable_traces {
                 output.push(SourceOutput::new_traces().with_port(TRACES))
             }
+            if !self.disable_llmobs {
+                output.push(SourceOutput::new_maybe_logs(DataType::Log, llmobs_definition).with_port(LLMOBS))
+            }
         } else {
             output.push(SourceOutput::new_maybe_logs(
                 DataType::all_bits(),
@@ -459,12 +471,19 @@ impl DatadogAgentSource {
         }
 
         if !config.disable_metrics {
-            let metrics_filter = metrics::build_warp_filter(handler, self.clone());
+            let metrics_filter = metrics::build_warp_filter(handler.clone(), self.clone());
             filters = filters
                 .map(|f| f.or(metrics_filter.clone()).unify().boxed())
                 .or(Some(metrics_filter));
         }
 
+        if !config.disable_llmobs {
+            let llmobs_filter = llmobs::build_warp_filter(handler.clone(), self.clone());
+            filters = filters
+                .map(|f| f.or(llmobs_filter.clone()).unify().boxed())
+                .or(Some(llmobs_filter));
+        }
+
         filters.ok_or_else(|| "At least one of the supported data type shall be enabled".into())
     }