Merge pull request #291 from lmnr-ai/dev

Fix association properties bug and add prompt from ai sdk
lmnr-ai · Dec 25, 2024 · 50defb1 · 50defb1
2 parents 5cf0c46 + 9898691
commit 50defb1
Showing 1 changed file with 15 additions and 5 deletions.
diff --git a/app-server/src/traces/spans.rs b/app-server/src/traces/spans.rs
@@ -59,7 +59,7 @@ impl SpanAttributes {
     pub fn session_id(&self) -> Option<String> {
         match self
             .attributes
-            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}session_id").as_str())
+            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}.session_id").as_str())
         {
             Some(Value::String(s)) => Some(s.clone()),
             _ => None,
@@ -69,7 +69,7 @@ impl SpanAttributes {
     pub fn user_id(&self) -> Option<String> {
         match self
             .attributes
-            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}user_id").as_str())
+            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}.user_id").as_str())
         {
             Some(Value::String(s)) => Some(s.clone()),
             _ => None,
@@ -78,7 +78,7 @@ impl SpanAttributes {
 
     pub fn trace_type(&self) -> Option<TraceType> {
         self.attributes
-            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}trace_type").as_str())
+            .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}.trace_type").as_str())
             .and_then(|s| serde_json::from_value(s.clone()).ok())
     }
 
@@ -131,7 +131,7 @@ impl SpanAttributes {
             if provider == "Langchain" {
                 let ls_provider = self
                     .attributes
-                    .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}ls_provider").as_str())
+                    .get(format!("{ASSOCIATION_PROPERTIES_PREFIX}.ls_provider").as_str())
                     .and_then(|s| serde_json::from_value(s.clone()).ok());
                 if let Some(ls_provider) = ls_provider {
                     ls_provider
@@ -375,6 +375,16 @@ impl Span {
                 );
             }
         }
+
+        // Vercel AI SDK wraps "raw" LLM spans in an additional `ai.generateText` span.
+        // Which is not really an LLM span, but it has the prompt in its attributes.
+        // Set the input to the prompt.
+        if let Some(serde_json::Value::String(s)) = attributes.get("ai.prompt") {
+            span.input = Some(
+                serde_json::from_str::<Value>(s).unwrap_or(serde_json::Value::String(s.clone())),
+            );
+        }
+
         // If an LLM span is sent manually, we prefer `lmnr.span.input` and `lmnr.span.output`
         // attributes over gen_ai/vercel/LiteLLM attributes.
         // Therefore this block is outside and after the LLM span type check.
@@ -453,7 +463,7 @@ impl Span {
         };
         let mut attributes = HashMap::new();
         attributes.insert(
-            format!("{ASSOCIATION_PROPERTIES_PREFIX}trace_type",),
+            format!("{ASSOCIATION_PROPERTIES_PREFIX}.trace_type",),
             json!(trace_type),
         );
         attributes.insert(SPAN_PATH.to_string(), json!(path));