trace_weft_openinference/
lib.rs

1use opentelemetry::KeyValue;
2use trace_weft_core::{SpanRecord, TraceWeftSpanKind};
3
4/// Maps a TraceWeft SpanRecord into OpenInference semantic attributes.
5pub fn map_to_openinference_attributes(record: &SpanRecord) -> Vec<KeyValue> {
6    let mut attributes = Vec::new();
7
8    let openinference_span_kind = match record.span_kind {
9        TraceWeftSpanKind::LlmCall => "LLM",
10        TraceWeftSpanKind::Tool => "TOOL",
11        TraceWeftSpanKind::Agent => "AGENT",
12        TraceWeftSpanKind::Retrieval => "RETRIEVER",
13        TraceWeftSpanKind::Embedding => "EMBEDDING",
14        TraceWeftSpanKind::Rerank => "RERANKER",
15        TraceWeftSpanKind::Evaluator => "EVALUATOR",
16        TraceWeftSpanKind::Guardrail => "GUARDRAIL",
17        _ => "CHAIN", // fallback
18    };
19
20    attributes.push(KeyValue::new(
21        "openinference.span.kind",
22        openinference_span_kind,
23    ));
24
25    if let Some(token_usage) = &record.token_usage {
26        attributes.push(KeyValue::new(
27            "llm.token_count.prompt",
28            token_usage.input as i64,
29        ));
30        attributes.push(KeyValue::new(
31            "llm.token_count.completion",
32            token_usage.output as i64,
33        ));
34        let total = token_usage.input + token_usage.output;
35        attributes.push(KeyValue::new("llm.token_count.total", total as i64));
36    }
37
38    if let Some(model_name) = &record.model_name {
39        attributes.push(KeyValue::new("llm.model_name", model_name.clone()));
40    }
41
42    if let Some(tool_name) = &record.tool_name {
43        attributes.push(KeyValue::new("tool.name", tool_name.clone()));
44    }
45
46    // NOTE: In a complete implementation we would also map input/output strings,
47    // prompts, variables, and retrieved documents into their respective
48    // openinference attributes like 'input.value', 'output.value', 'retrieval.documents' etc.
49
50    attributes
51}
52
53#[cfg(test)]
54mod tests {
55    use super::*;
56    use opentelemetry::Value;
57    use trace_weft_core::test_util::{sample_span_full, sample_span_minimal};
58
59    fn attr_value<'a>(attrs: &'a [KeyValue], key: &str) -> Option<&'a Value> {
60        attrs
61            .iter()
62            .find(|kv| kv.key.as_str() == key)
63            .map(|kv| &kv.value)
64    }
65
66    #[test]
67    fn maps_span_kinds_to_openinference_names() {
68        let cases = [
69            (TraceWeftSpanKind::LlmCall, "LLM"),
70            (TraceWeftSpanKind::Tool, "TOOL"),
71            (TraceWeftSpanKind::Agent, "AGENT"),
72            (TraceWeftSpanKind::Retrieval, "RETRIEVER"),
73            (TraceWeftSpanKind::Embedding, "EMBEDDING"),
74            (TraceWeftSpanKind::Rerank, "RERANKER"),
75            (TraceWeftSpanKind::Evaluator, "EVALUATOR"),
76            (TraceWeftSpanKind::Guardrail, "GUARDRAIL"),
77            // Kinds without a direct OpenInference equivalent fall back to CHAIN.
78            (TraceWeftSpanKind::Workflow, "CHAIN"),
79            (TraceWeftSpanKind::Planner, "CHAIN"),
80        ];
81        for (kind, expected) in cases {
82            let mut span = sample_span_minimal();
83            span.span_kind = kind;
84            let attrs = map_to_openinference_attributes(&span);
85            assert_eq!(
86                attr_value(&attrs, "openinference.span.kind"),
87                Some(&Value::from(expected)),
88                "kind {kind:?} should map to {expected}"
89            );
90        }
91    }
92
93    #[test]
94    fn maps_token_counts_including_total() {
95        let attrs = map_to_openinference_attributes(&sample_span_full());
96        assert_eq!(
97            attr_value(&attrs, "llm.token_count.prompt"),
98            Some(&Value::I64(100))
99        );
100        assert_eq!(
101            attr_value(&attrs, "llm.token_count.completion"),
102            Some(&Value::I64(50))
103        );
104        assert_eq!(
105            attr_value(&attrs, "llm.token_count.total"),
106            Some(&Value::I64(150))
107        );
108    }
109
110    #[test]
111    fn maps_model_and_tool_names() {
112        let attrs = map_to_openinference_attributes(&sample_span_full());
113        assert_eq!(
114            attr_value(&attrs, "llm.model_name"),
115            Some(&Value::from("gpt-4.1"))
116        );
117        assert_eq!(
118            attr_value(&attrs, "tool.name"),
119            Some(&Value::from("kb_search"))
120        );
121    }
122
123    #[test]
124    fn omits_llm_attributes_for_bare_spans() {
125        let attrs = map_to_openinference_attributes(&sample_span_minimal());
126        assert!(attr_value(&attrs, "llm.token_count.prompt").is_none());
127        assert!(attr_value(&attrs, "llm.model_name").is_none());
128        assert!(attr_value(&attrs, "tool.name").is_none());
129        assert_eq!(attrs.len(), 1, "only openinference.span.kind expected");
130    }
131}
trace_weft_openinference/lib.rs

trace_weft_openinference/
lib.rs
