feat(llama-index): add LlamaIndexInstrumentor (#946)

Author: hassiebp

langfuse/llama_index/__init__.py

@@ -1,4 +1,7 @@
 from .llama_index import LlamaIndexCallbackHandler
-from .span_handler import LlamaIndexSpanHandler
+from ._instrumentor import LlamaIndexInstrumentor
 
-__all__ = ["LlamaIndexCallbackHandler", "LlamaIndexSpanHandler"]
+__all__ = [
+    "LlamaIndexCallbackHandler",
+    "LlamaIndexInstrumentor",
+]

langfuse/llama_index/_context.py

@@ -0,0 +1,132 @@
+from contextvars import ContextVar
+from typing import Optional, Any, List, Dict
+from ._types import InstrumentorContextData
+
+
+default_context: InstrumentorContextData = {
+    "trace_id": None,
+    "parent_observation_id": None,
+    "update_parent": True,
+    "trace_name": None,
+    "root_llama_index_span_id": None,
+    "is_user_managed_trace": None,
+    "user_id": None,
+    "session_id": None,
+    "version": None,
+    "release": None,
+    "metadata": None,
+    "tags": None,
+    "public": None,
+}
+
+langfuse_instrumentor_context: ContextVar[InstrumentorContextData] = ContextVar(
+    "langfuse_instrumentor_context",
+    default={**default_context},
+    # The spread operator (**) is used here to create a new dictionary
+    # that is a shallow copy of default_trace_attributes.
+    # This ensures that each ContextVar instance gets its own copy of the default attributes,
+    # preventing accidental shared state between different contexts.
+    # If we didn't use the spread operator, all contexts would reference the same dictionary,
+    # which could lead to unexpected behavior if the dictionary is modified.
+)
+
+
+class InstrumentorContext:
+    @staticmethod
+    def _get_context():
+        return langfuse_instrumentor_context.get()
+
+    @property
+    def trace_id(self) -> Optional[str]:
+        return self._get_context()["trace_id"]
+
+    @property
+    def parent_observation_id(self) -> Optional[str]:
+        return self._get_context()["parent_observation_id"]
+
+    @property
+    def root_llama_index_span_id(self) -> Optional[str]:
+        return self._get_context()["root_llama_index_span_id"]
+
+    @property
+    def is_user_managed_trace(self) -> Optional[bool]:
+        return self._get_context()["is_user_managed_trace"]
+
+    @property
+    def update_parent(self) -> Optional[bool]:
+        return self._get_context()["update_parent"]
+
+    @property
+    def trace_name(self) -> Optional[str]:
+        return self._get_context()["trace_name"]
+
+    @property
+    def trace_data(self):
+        return {
+            "user_id": self._get_context()["user_id"],
+            "session_id": self._get_context()["session_id"],
+            "version": self._get_context()["version"],
+            "release": self._get_context()["release"],
+            "metadata": self._get_context()["metadata"],
+            "tags": self._get_context()["tags"],
+            "public": self._get_context()["public"],
+        }
+
+    @staticmethod
+    def reset():
+        langfuse_instrumentor_context.set({**default_context})
+
+    def reset_trace_id(self):
+        previous_context = self._get_context()
+
+        langfuse_instrumentor_context.set(
+            {**previous_context, "trace_id": None, "root_llama_index_span_id": None}
+        )
+
+    @staticmethod
+    def update(
+        *,
+        trace_id: Optional[str] = None,
+        parent_observation_id: Optional[str] = None,
+        update_parent: Optional[bool] = None,
+        root_llama_index_span_id: Optional[str] = None,
+        is_user_managed_trace: Optional[bool] = None,
+        trace_name: Optional[str] = None,
+        user_id: Optional[str] = None,
+        session_id: Optional[str] = None,
+        version: Optional[str] = None,
+        release: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        tags: Optional[List[str]] = None,
+        public: Optional[bool] = None,
+    ):
+        updates = {}
+
+        if trace_id is not None:
+            updates["trace_id"] = trace_id
+        if parent_observation_id is not None:
+            updates["parent_observation_id"] = parent_observation_id
+        if update_parent is not None:
+            updates["update_parent"] = update_parent
+        if trace_name is not None:
+            updates["trace_name"] = trace_name
+        if root_llama_index_span_id is not None:
+            updates["root_llama_index_span_id"] = root_llama_index_span_id
+        if is_user_managed_trace is not None:
+            updates["is_user_managed_trace"] = is_user_managed_trace
+        if user_id is not None:
+            updates["user_id"] = user_id
+        if session_id is not None:
+            updates["session_id"] = session_id
+        if version is not None:
+            updates["version"] = version
+        if release is not None:
+            updates["release"] = release
+        if metadata is not None:
+            updates["metadata"] = metadata
+        if tags is not None:
+            updates["tags"] = tags
+        if public is not None:
+            updates["public"] = public
+
+        langfuse_instrumentor_context.get().update(updates)

langfuse/llama_index/_event_handler.py

@@ -0,0 +1,193 @@
+from typing import Optional, Any, Union, Dict, Mapping
+
+from langfuse.client import (
+    Langfuse,
+)
+from langfuse.model import ModelUsage
+
+
+try:
+    from llama_index.core.base.llms.types import (
+        ChatResponse,
+        CompletionResponse,
+    )
+    from llama_index.core.instrumentation.events import BaseEvent
+    from llama_index.core.instrumentation.events.embedding import (
+        EmbeddingStartEvent,
+        EmbeddingEndEvent,
+    )
+    from llama_index.core.instrumentation.event_handlers import BaseEventHandler
+    from llama_index.core.instrumentation.events.llm import (
+        LLMCompletionEndEvent,
+        LLMCompletionStartEvent,
+        LLMChatEndEvent,
+        LLMChatStartEvent,
+    )
+    from llama_index.core.utilities.token_counting import TokenCounter
+
+except ImportError:
+    raise ModuleNotFoundError(
+        "Please install llama-index to use the Langfuse llama-index integration: 'pip install llama-index'"
+    )
+
+from logging import getLogger
+
+logger = getLogger(__name__)
+
+
+class LlamaIndexEventHandler(BaseEventHandler, extra="allow"):
+    def __init__(
+        self,
+        *,
+        langfuse_client: Langfuse,
+        observation_updates: Dict[str, Dict[str, Any]],
+    ):
+        super().__init__()
+
+        self._langfuse = langfuse_client
+        self._observation_updates = observation_updates
+        self._token_counter = TokenCounter()
+
+    @classmethod
+    def class_name(cls) -> str:
+        """Class name."""
+        return "LlamaIndexEventHandler"
+
+    def handle(self, event: BaseEvent) -> None:
+        logger.debug(f"Event {type(event).__name__} received: {event}")
+
+        if isinstance(
+            event, (LLMCompletionStartEvent, LLMChatStartEvent, EmbeddingStartEvent)
+        ):
+            self.update_generation_from_start_event(event)
+        elif isinstance(
+            event, (LLMCompletionEndEvent, LLMChatEndEvent, EmbeddingEndEvent)
+        ):
+            self.update_generation_from_end_event(event)
+
+    def update_generation_from_start_event(
+        self,
+        event: Union[LLMCompletionStartEvent, LLMChatStartEvent, EmbeddingStartEvent],
+    ) -> None:
+        if event.span_id is None:
+            logger.warning("Span ID is not set")
+            return
+
+        model_data = event.model_dict
+        model = model_data.pop("model", None) or model_data.pop("model_name", None)
+        traced_model_data = {
+            k: str(v)
+            for k, v in model_data.items()
+            if v is not None
+            and k
+            in [
+                "max_tokens",
+                "max_retries",
+                "temperature",
+                "timeout",
+                "strict",
+                "top_logprobs",
+                "logprobs",
+                "embed_batch_size",
+            ]
+        }
+
+        self._update_observation_updates(
+            event.span_id, model=model, model_parameters=traced_model_data
+        )
+
+    def update_generation_from_end_event(
+        self, event: Union[LLMCompletionEndEvent, LLMChatEndEvent, EmbeddingEndEvent]
+    ) -> None:
+        if event.span_id is None:
+            logger.warning("Span ID is not set")
+            return
+
+        usage = None
+
+        if isinstance(event, (LLMCompletionEndEvent, LLMChatEndEvent)):
+            usage = self._parse_token_usage(event.response) if event.response else None
+
+        if isinstance(event, EmbeddingEndEvent):
+            token_count = sum(
+                self._token_counter.get_string_tokens(chunk) for chunk in event.chunks
+            )
+
+            usage = {
+                "input": 0,
+                "output": 0,
+                "total": token_count or None,
+            }
+
+        self._update_observation_updates(event.span_id, usage=usage)
+
+    def _update_observation_updates(self, id_: str, **kwargs) -> None:
+        if id_ not in self._observation_updates:
+            return
+
+        self._observation_updates[id_].update(kwargs)
+
+    def _parse_token_usage(
+        self, response: Union[ChatResponse, CompletionResponse]
+    ) -> Optional[ModelUsage]:
+        if (
+            (raw := getattr(response, "raw", None))
+            and hasattr(raw, "get")
+            and (usage := raw.get("usage"))
+        ):
+            return _parse_usage_from_mapping(usage)
+
+        if additional_kwargs := getattr(response, "additional_kwargs", None):
+            return _parse_usage_from_mapping(additional_kwargs)
+
+
+def _parse_usage_from_mapping(
+    usage: Union[object, Mapping[str, Any]],
+) -> ModelUsage:
+    if isinstance(usage, Mapping):
+        return _get_token_counts_from_mapping(usage)
+
+    return _parse_usage_from_object(usage)
+
+
+def _parse_usage_from_object(usage: object) -> ModelUsage:
+    model_usage: ModelUsage = {
+        "unit": None,
+        "input": None,
+        "output": None,
+        "total": None,
+        "input_cost": None,
+        "output_cost": None,
+        "total_cost": None,
+    }
+
+    if (prompt_tokens := getattr(usage, "prompt_tokens", None)) is not None:
+        model_usage["input"] = prompt_tokens
+    if (completion_tokens := getattr(usage, "completion_tokens", None)) is not None:
+        model_usage["output"] = completion_tokens
+    if (total_tokens := getattr(usage, "total_tokens", None)) is not None:
+        model_usage["total"] = total_tokens
+
+    return model_usage
+
+
+def _get_token_counts_from_mapping(
+    usage_mapping: Mapping[str, Any],
+) -> ModelUsage:
+    model_usage: ModelUsage = {
+        "unit": None,
+        "input": None,
+        "output": None,
+        "total": None,
+        "input_cost": None,
+        "output_cost": None,
+        "total_cost": None,
+    }
+    if (prompt_tokens := usage_mapping.get("prompt_tokens")) is not None:
+        model_usage["input"] = prompt_tokens
+    if (completion_tokens := usage_mapping.get("completion_tokens")) is not None:
+        model_usage["output"] = completion_tokens
+    if (total_tokens := usage_mapping.get("total_tokens")) is not None:
+        model_usage["total"] = total_tokens
+
+    return model_usage