From 5876431b7225e466c4b788e36c73dc7c6c1840e3 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Tue, 21 Jul 2026 00:03:16 +0000 Subject: [PATCH 1/2] fix(llamaindex): align spans with framework-delegation rule, add provider metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Emit task (not llm) spans for LlamaIndex LLM classes; they always delegate transport to a separately-instrumentable provider client, so typing them llm would produce token-less doubly-nested llm spans under auto_instrument. - Add metadata.provider via module-namespace allowlist (llama_index.llms.*, llama_index.embeddings.*) with a getattr(instance, "provider") fallback. - Pass the Exception instance to span.log(error=...) instead of pre-formatting. - Drop str(result) fallback in output extraction — Braintrust serializes at send time; unknown result types return None. - Update existing VCR-backed tests to match new span shape. --- .../llamaindex/test_llamaindex.py | 34 +++++++------ .../integrations/llamaindex/tracing.py | 48 +++++++++++-------- 2 files changed, 48 insertions(+), 34 deletions(-) diff --git a/py/src/braintrust/integrations/llamaindex/test_llamaindex.py b/py/src/braintrust/integrations/llamaindex/test_llamaindex.py index 2c27cfc6..36fd2099 100644 --- a/py/src/braintrust/integrations/llamaindex/test_llamaindex.py +++ b/py/src/braintrust/integrations/llamaindex/test_llamaindex.py @@ -110,16 +110,17 @@ def test_llm_complete(logger_memory_logger): spans = memory_logger.pop() assert len(spans) >= 2 - llm_spans = _find_spans_by_attributes(spans, type="llm") - assert len(llm_spans) >= 1 + openai_spans = [s for s in spans if s.get("span_attributes", {}).get("name") == "OpenAI"] + assert len(openai_spans) >= 1 - llm_span = llm_spans[0] + llm_span = openai_spans[0] + assert llm_span["span_attributes"]["type"] == "task" assert llm_span["context"]["span_origin"]["instrumentation"]["name"] == "llamaindex-auto" - assert llm_span["span_attributes"]["name"] == "OpenAI" assert llm_span["input"] is not None assert llm_span["output"] is not None assert llm_span["metadata"]["class"] == "OpenAI" assert llm_span["metadata"]["model"] == "gpt-4o-mini" + assert llm_span["metadata"]["provider"] == "openai" @pytest.mark.vcr @@ -142,14 +143,16 @@ def test_llm_chat(logger_memory_logger): spans = memory_logger.pop() assert len(spans) >= 2 - llm_spans = _find_spans_by_attributes(spans, type="llm") - assert len(llm_spans) >= 1 + openai_spans = [s for s in spans if s.get("span_attributes", {}).get("name") == "OpenAI"] + assert len(openai_spans) >= 1 - llm_span = llm_spans[0] + llm_span = openai_spans[0] + assert llm_span["span_attributes"]["type"] == "task" assert llm_span["input"] is not None assert llm_span["output"] is not None assert isinstance(llm_span["output"], dict) assert "content" in llm_span["output"] or "role" in llm_span["output"] + assert llm_span["metadata"]["provider"] == "openai" def test_document_processing(logger_memory_logger): @@ -224,7 +227,6 @@ def test_query_engine(logger_memory_logger): span_types = {s.get("span_attributes", {}).get("type") for s in spans} assert "task" in span_types - assert "llm" in span_types or "function" in span_types def test_span_hierarchy(logger_memory_logger): @@ -266,9 +268,9 @@ def test_llm_error_handling(logger_memory_logger): spans = memory_logger.pop() assert len(spans) >= 2 - llm_spans = _find_spans_by_attributes(spans, type="llm") - assert len(llm_spans) >= 1 - assert llm_spans[0].get("error") is not None + openai_spans = [s for s in spans if s.get("span_attributes", {}).get("name") == "OpenAI"] + assert len(openai_spans) >= 1 + assert openai_spans[0].get("error") is not None @pytest.mark.vcr @@ -287,8 +289,9 @@ async def test_async_llm_complete(logger_memory_logger): spans = memory_logger.pop() assert len(spans) >= 2 - llm_spans = _find_spans_by_attributes(spans, type="llm") - assert len(llm_spans) >= 1 + openai_spans = [s for s in spans if s.get("span_attributes", {}).get("name") == "OpenAI"] + assert len(openai_spans) >= 1 + assert openai_spans[0]["span_attributes"]["type"] == "task" @pytest.mark.vcr @@ -311,5 +314,6 @@ async def test_async_llm_chat(logger_memory_logger): spans = memory_logger.pop() assert len(spans) >= 2 - llm_spans = _find_spans_by_attributes(spans, type="llm") - assert len(llm_spans) >= 1 + openai_spans = [s for s in spans if s.get("span_attributes", {}).get("name") == "OpenAI"] + assert len(openai_spans) >= 1 + assert openai_spans[0]["span_attributes"]["type"] == "task" diff --git a/py/src/braintrust/integrations/llamaindex/tracing.py b/py/src/braintrust/integrations/llamaindex/tracing.py index f63a60bb..1be2a4d2 100644 --- a/py/src/braintrust/integrations/llamaindex/tracing.py +++ b/py/src/braintrust/integrations/llamaindex/tracing.py @@ -6,10 +6,15 @@ from braintrust.logger import NOOP_SPAN, Span, current_span from braintrust.logger import start_span as _bt_start_span +from braintrust.span_types import SpanTypeAttribute _INSTRUMENTATION = "llamaindex-auto" +_LLM_METADATA_ALLOWLIST = ("model", "model_name", "temperature", "max_tokens") +_LLM_MODULE_PROVIDER_PREFIX = "llama_index.llms." +_EMBEDDING_MODULE_PROVIDER_PREFIX = "llama_index.embeddings." + def start_span(*args, **kwargs): internal = dict(kwargs.get("internal") or {}) @@ -18,7 +23,15 @@ def start_span(*args, **kwargs): return _bt_start_span(*args, **kwargs) -from braintrust.span_types import SpanTypeAttribute +def _extract_provider(instance: Any) -> str | None: + provider = getattr(instance, "provider", None) + if isinstance(provider, str) and provider: + return provider + module = getattr(type(instance), "__module__", "") or "" + for prefix in (_LLM_MODULE_PROVIDER_PREFIX, _EMBEDDING_MODULE_PROVIDER_PREFIX): + if module.startswith(prefix): + return module[len(prefix):].split(".", 1)[0] or None + return None def _extract_block_content(message: Any) -> str | None: @@ -55,11 +68,8 @@ def _extract_messages(messages: Any) -> list[dict[str, Any]] | None: def _extract_response_output(result: Any) -> Any: if result is None: return None - # Streaming/coroutine responses are consumed outside this span handler. - # Do not log unstable object reprs such as "". if inspect.isgenerator(result) or inspect.isasyncgen(result) or inspect.iscoroutine(result): return None - # ChatResponse if hasattr(result, "message") and hasattr(result, "raw"): msg = result.message if not msg: @@ -70,21 +80,18 @@ def _extract_response_output(result: Any) -> Any: if content is not None: output["content"] = content return output - # CompletionResponse if hasattr(result, "text") and hasattr(result, "raw"): return {"text": result.text} - # Query response if hasattr(result, "response") and hasattr(result, "source_nodes"): output = {"response": result.response} if result.source_nodes: output["source_nodes"] = _extract_nodes(result.source_nodes) return output - # List of NodeWithScore if isinstance(result, list) and result and hasattr(result[0], "node"): return _extract_nodes(result) if isinstance(result, str): return result - return str(result) + return None def _extract_nodes(nodes: list[Any]) -> list[dict[str, Any]]: @@ -111,8 +118,12 @@ def _classify_instance(instance: Any) -> tuple[SpanTypeAttribute, str]: cls_name = type(instance).__name__ mro_names = {c.__name__ for c in type(instance).__mro__} + # LLM classes always delegate to a separately-instrumentable provider client + # (openai, anthropic, ...). Two nested llm spans for one API call is + # confusing and would produce token-less framework spans — so type this as + # task and let the underlying provider integration own the llm leaf. if "BaseLLM" in mro_names or "LLM" in mro_names: - return SpanTypeAttribute.LLM, cls_name + return SpanTypeAttribute.TASK, cls_name if "BaseTool" in mro_names or "FunctionTool" in mro_names: return SpanTypeAttribute.TOOL, getattr(instance, "name", None) or cls_name @@ -200,10 +211,13 @@ def new_span( metadata: dict[str, Any] = {} if instance is not None: metadata["class"] = type(instance).__name__ - for attr in ("model", "model_name", "temperature", "max_tokens"): + for attr in _LLM_METADATA_ALLOWLIST: val = getattr(instance, attr, None) if val is not None: metadata[attr] = val + provider = _extract_provider(instance) + if provider is not None: + metadata["provider"] = provider parent_bt_span = self._find_parent_bt_span(parent_span_id) @@ -244,15 +258,10 @@ def prepare_to_exit_span( bt_span = record.bt_span output = _extract_response_output(result) - # Token usage is intentionally not logged on LlamaIndex spans. - # LlamaIndex is an orchestration layer; provider integrations own - # token accounting. Emitting usage here would double-count when - # provider spans are also present. - log_kwargs: dict[str, Any] = {} + # Token accounting belongs to the underlying provider integration + # (openai, anthropic, ...) — do not log usage here. if output is not None: - log_kwargs["output"] = output - if log_kwargs: - bt_span.log(**log_kwargs) + bt_span.log(output=output) bt_span.unset_current() bt_span.end() @@ -271,7 +280,8 @@ def prepare_to_drop_span( return None bt_span = record.bt_span - bt_span.log(error=f"{type(err).__name__}: {err}" if err else "Unknown error") + if err is not None: + bt_span.log(error=err) bt_span.unset_current() bt_span.end() From bdd4c0220994816cd840787319fcf948c941d2d9 Mon Sep 17 00:00:00 2001 From: Starfolk Date: Tue, 21 Jul 2026 00:08:16 +0000 Subject: [PATCH 2/2] style: ruff-format slice whitespace --- py/src/braintrust/integrations/llamaindex/tracing.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/py/src/braintrust/integrations/llamaindex/tracing.py b/py/src/braintrust/integrations/llamaindex/tracing.py index 1be2a4d2..12491d26 100644 --- a/py/src/braintrust/integrations/llamaindex/tracing.py +++ b/py/src/braintrust/integrations/llamaindex/tracing.py @@ -30,7 +30,7 @@ def _extract_provider(instance: Any) -> str | None: module = getattr(type(instance), "__module__", "") or "" for prefix in (_LLM_MODULE_PROVIDER_PREFIX, _EMBEDDING_MODULE_PROVIDER_PREFIX): if module.startswith(prefix): - return module[len(prefix):].split(".", 1)[0] or None + return module[len(prefix) :].split(".", 1)[0] or None return None