fix(llama-index): workflow integration (#1095) (#1097)

Co-authored-by: Petros Mitseas <petros_94@icloud.com>

Author: hassiebp

langfuse/llama_index/_span_handler.py

@@ -11,6 +11,8 @@
 )
 
 from logging import getLogger
+from pydantic import BaseModel
+
 from ._context import InstrumentorContext
 
 logger = getLogger(__name__)
@@ -25,6 +27,7 @@
         StreamingResponse,
         AsyncStreamingResponse,
     )
+    from llama_index.core.workflow import Context
 
 except ImportError:
     raise ModuleNotFoundError(
@@ -243,6 +246,9 @@ def _parse_generation_input(
     def _parse_output_metadata(
         self, instance: Optional[Any], result: Optional[Any]
     ) -> Tuple[Optional[Any], Optional[Any]]:
+        if isinstance(result, BaseModel):
+            return result.__dict__, None
+
         if not result or isinstance(
             result,
             (Generator, AsyncGenerator, StreamingResponse, AsyncStreamingResponse),
@@ -289,4 +295,12 @@ def _parse_input(self, *, bound_args: inspect.BoundArguments):
         if "nodes" in arguments:
             return {"num_nodes": len(arguments["nodes"])}
 
+        # Remove Context since it is in not properly serialized
+        ctx_key = None
+        for arg, val in arguments.items():
+            if isinstance(val, Context):
+                ctx_key = arg
+        if ctx_key in arguments:
+            return {arg: val for arg, val in arguments.items() if arg != ctx_key}
+
         return arguments