refactor(llama-index): send generation updates directly from event handler (#981)

Author: hassiebp

langfuse/llama_index/_event_handler.py

@@ -1,10 +1,14 @@
-from typing import Optional, Any, Union, Dict, Mapping
+from typing import Optional, Any, Union, Mapping
 
 from langfuse.client import (
     Langfuse,
+    StatefulGenerationClient,
+    StateType,
 )
+from langfuse.utils import _get_timestamp
 from langfuse.model import ModelUsage
-
+from ._context import InstrumentorContext
+from uuid import uuid4 as create_uuid
 
 try:
     from llama_index.core.base.llms.types import (
@@ -36,17 +40,12 @@
 
 
 class LlamaIndexEventHandler(BaseEventHandler, extra="allow"):
-    def __init__(
-        self,
-        *,
-        langfuse_client: Langfuse,
-        observation_updates: Dict[str, Dict[str, Any]],
-    ):
+    def __init__(self, *, langfuse_client: Langfuse):
         super().__init__()
 
         self._langfuse = langfuse_client
-        self._observation_updates = observation_updates
         self._token_counter = TokenCounter()
+        self._context = InstrumentorContext()
 
     @classmethod
     def class_name(cls) -> str:
@@ -92,8 +91,8 @@ def update_generation_from_start_event(
             ]
         }
 
-        self._update_observation_updates(
-            event.span_id, model=model, model_parameters=traced_model_data
+        self._get_generation_client(event.span_id).update(
+            model=model, model_parameters=traced_model_data
         )
 
     def update_generation_from_end_event(
@@ -119,13 +118,9 @@ def update_generation_from_end_event(
                 "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)
+        self._get_generation_client(event.span_id).update(
+            usage=usage, end_time=_get_timestamp()
+        )
 
     def _parse_token_usage(
         self, response: Union[ChatResponse, CompletionResponse]
@@ -140,6 +135,22 @@ def _parse_token_usage(
         if additional_kwargs := getattr(response, "additional_kwargs", None):
             return _parse_usage_from_mapping(additional_kwargs)
 
+    def _get_generation_client(self, id: str) -> StatefulGenerationClient:
+        trace_id = self._context.trace_id
+        if trace_id is None:
+            logger.warning(
+                "Trace ID is not set. Creating generation client with new trace id."
+            )
+            trace_id = str(create_uuid())
+
+        return StatefulGenerationClient(
+            client=self._langfuse.client,
+            id=id,
+            trace_id=trace_id,
+            task_manager=self._langfuse.task_manager,
+            state_type=StateType.OBSERVATION,
+        )
+
 
 def _parse_usage_from_mapping(
     usage: Union[object, Mapping[str, Any]],

langfuse/llama_index/_instrumentor.py

@@ -98,15 +98,8 @@ def __init__(
             mask=mask,
             sdk_integration="llama-index_instrumentation",
         )
-        self._observation_updates = {}
-        self._span_handler = LlamaIndexSpanHandler(
-            langfuse_client=self._langfuse,
-            observation_updates=self._observation_updates,
-        )
-        self._event_handler = LlamaIndexEventHandler(
-            langfuse_client=self._langfuse,
-            observation_updates=self._observation_updates,
-        )
+        self._span_handler = LlamaIndexSpanHandler(langfuse_client=self._langfuse)
+        self._event_handler = LlamaIndexEventHandler(langfuse_client=self._langfuse)
         self._context = InstrumentorContext()
 
     def start(self):

langfuse/llama_index/_span_handler.py

@@ -1,5 +1,5 @@
 import inspect
-from typing import Optional, Any, Tuple, Dict, Generator, AsyncGenerator
+from typing import Optional, Any, Tuple, Generator, AsyncGenerator
 import uuid
 
 from langfuse.client import (
@@ -39,16 +39,10 @@ class LangfuseSpan(BaseSpan):
 
 
 class LlamaIndexSpanHandler(BaseSpanHandler[LangfuseSpan], extra="allow"):
-    def __init__(
-        self,
-        *,
-        langfuse_client: Langfuse,
-        observation_updates: Dict[str, Dict[str, Any]],
-    ):
+    def __init__(self, *, langfuse_client: Langfuse):
         super().__init__()
 
         self._langfuse_client = langfuse_client
-        self._observation_updates = observation_updates
         self._context = InstrumentorContext()
 
     def new_span(
@@ -109,9 +103,6 @@ def new_span(
                 metadata=kwargs,
             )
 
-        # Initialize observation update for the span to be populated by event handler
-        self._observation_updates[id_] = {}
-
     def prepare_to_exit_span(
         self,
         id_: str,
@@ -122,7 +113,6 @@ def prepare_to_exit_span(
     ) -> Optional[LangfuseSpan]:
         logger.debug(f"Exiting span {instance.__class__.__name__} with ID {id_}")
 
-        observation_updates = self._observation_updates.pop(id_, {})
         output, metadata = self._parse_output_metadata(instance, result)
 
         # Reset the context root if the span is the root span
@@ -138,15 +128,13 @@ def prepare_to_exit_span(
         if self._is_generation(id_, instance):
             generationClient = self._get_generation_client(id_)
             generationClient.end(
-                **observation_updates,
                 output=output,
                 metadata=metadata,
             )
 
         else:
             spanClient = self._get_span_client(id_)
             spanClient.end(
-                **observation_updates,
                 output=output,
                 metadata=metadata,
             )
@@ -161,8 +149,6 @@ def prepare_to_drop_span(
     ) -> Optional[LangfuseSpan]:
         logger.debug(f"Dropping span {instance.__class__.__name__} with ID {id_}")
 
-        observation_updates = self._observation_updates.pop(id_, {})
-
         # Reset the context root if the span is the root span
         if id_ == self._context.root_llama_index_span_id:
             if self._context.update_parent:
@@ -177,15 +163,13 @@ def prepare_to_drop_span(
         if self._is_generation(id_, instance):
             generationClient = self._get_generation_client(id_)
             generationClient.end(
-                **observation_updates,
                 level="ERROR",
                 status_message=str(err),
             )
 
         else:
             spanClient = self._get_span_client(id_)
             spanClient.end(
-                **observation_updates,
                 level="ERROR",
                 status_message=str(err),
             )
@@ -217,7 +201,10 @@ def _is_generation(self, id_: str, instance: Optional[Any] = None) -> bool:
     def _get_generation_client(self, id: str) -> StatefulGenerationClient:
         trace_id = self._context.trace_id
         if trace_id is None:
-            raise ValueError("Trace ID is not set")
+            logger.warning(
+                "Trace ID is not set. Creating generation client with new trace id."
+            )
+            trace_id = str(uuid.uuid4())
 
         return StatefulGenerationClient(
             client=self._langfuse_client.client,
@@ -230,7 +217,10 @@ def _get_generation_client(self, id: str) -> StatefulGenerationClient:
     def _get_span_client(self, id: str) -> StatefulSpanClient:
         trace_id = self._context.trace_id
         if trace_id is None:
-            raise ValueError("Trace ID is not set")
+            logger.warning(
+                "Trace ID is not set. Creating generation client with new trace id."
+            )
+            trace_id = str(uuid.uuid4())
 
         return StatefulSpanClient(
             client=self._langfuse_client.client,