Skip to content

Commit 7c804a3

Browse files
committed
feat: add span processor for tool input.value and output.value
1 parent bcd7948 commit 7c804a3

2 files changed

Lines changed: 93 additions & 0 deletions

File tree

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
"""OpenTelemetry SpanProcessor for normalizing LlamaIndex tool call attributes.
2+
3+
LlamaIndex wraps tool arguments in {"kwargs": {...}} which differs from other
4+
frameworks like LangChain that use flat {"arg": value} format. This processor
5+
normalizes the format at the span level before exporters or dev terminal read it.
6+
"""
7+
8+
import json
9+
import logging
10+
from typing import Any, Optional
11+
12+
from opentelemetry.context import Context
13+
from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
14+
15+
logger = logging.getLogger(__name__)
16+
17+
18+
class AttributeNormalizingSpanProcessor(SpanProcessor):
19+
"""Normalizes LlamaIndex tool call attributes to match other frameworks.
20+
21+
Unwraps {"kwargs": {...}} to flat {...} format for consistency with LangChain.
22+
"""
23+
24+
def on_start(self, span: Span, parent_context: Optional[Context] = None) -> None:
25+
"""Called when span starts - no action needed."""
26+
pass
27+
28+
def on_end(self, span: ReadableSpan) -> None:
29+
"""Normalize tool call attributes before span is consumed by exporters/terminal."""
30+
# Access internal mutable attributes (BoundedAttributes at runtime)
31+
# ReadableSpan._attributes is typed as Mapping but mutable in practice
32+
attrs = getattr(span, "_attributes", None)
33+
if not attrs:
34+
return
35+
36+
try:
37+
if attrs.get("openinference.span.kind", None) == "TOOL":
38+
for key in ("input.value", "output.value"):
39+
if key in attrs:
40+
original = attrs[key]
41+
normalized = self._normalize_attribute(key, original)
42+
43+
if normalized != original:
44+
attrs[key] = normalized
45+
if logger.isEnabledFor(logging.DEBUG):
46+
logger.debug(
47+
f"Normalized {key} in span '{span.name}': "
48+
f"{str(original)[:50]}... → {str(normalized)[:50]}..."
49+
)
50+
51+
except Exception as e:
52+
logger.debug(
53+
f"Failed to normalize span '{getattr(span, 'name', 'unknown')}': {e}"
54+
)
55+
56+
def _normalize_attribute(self, key: str, value: Any) -> str:
57+
"""Unwrap LlamaIndex's kwargs wrapper if present."""
58+
if isinstance(value, str):
59+
try:
60+
value = json.loads(value)
61+
except Exception:
62+
pass
63+
if isinstance(value, dict):
64+
if key == "input.value":
65+
if "kwargs" in value:
66+
value = json.dumps(value["kwargs"])
67+
elif key == "output.value":
68+
value = json.dumps(
69+
{
70+
"content": value.get("raw_output"),
71+
"status": "success"
72+
if not value.get("is_error", False)
73+
else "error",
74+
"tool_call_id": value.get("tool_call_id"),
75+
}
76+
)
77+
return str(value)
78+
79+
def shutdown(self) -> None:
80+
"""Called on processor shutdown - no cleanup needed."""
81+
pass
82+
83+
def force_flush(self, timeout_millis: int = 30000) -> bool:
84+
"""Force flush - always succeeds (nothing to flush)."""
85+
return True

src/uipath_llamaindex/runtime/factory.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,9 @@
1717
from uipath.runtime.errors import UiPathErrorCategory
1818
from workflows import Workflow
1919

20+
from uipath_llamaindex.runtime._attribute_normalizer import (
21+
AttributeNormalizingSpanProcessor,
22+
)
2023
from uipath_llamaindex.runtime.config import LlamaIndexConfig
2124
from uipath_llamaindex.runtime.errors import (
2225
UiPathLlamaIndexErrorCode,
@@ -55,6 +58,11 @@ def _setup_instrumentation(self, trace_manager: UiPathTraceManager | None) -> No
5558
LlamaIndexInstrumentor().instrument()
5659
UiPathSpanUtils.register_current_span_provider(get_current_span)
5760

61+
if trace_manager:
62+
trace_manager.tracer_provider.add_span_processor(
63+
AttributeNormalizingSpanProcessor()
64+
)
65+
5866
def _get_storage_path(self) -> str:
5967
"""Get the storage path for workflow state."""
6068
if self._storage_path is None:

0 commit comments

Comments
 (0)