From 8547164a9e925954ea17f19c343be9b2cd594545 Mon Sep 17 00:00:00 2001 From: Hassieb Pakzad <68423100+hassiebp@users.noreply.github.com> Date: Fri, 31 Jan 2025 11:18:40 +0100 Subject: [PATCH] fix(llama-index): workflow integration (#1095) (#1097) Co-authored-by: Petros Mitseas --- langfuse/llama_index/_span_handler.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/langfuse/llama_index/_span_handler.py b/langfuse/llama_index/_span_handler.py index 33ee8282..14387d5a 100644 --- a/langfuse/llama_index/_span_handler.py +++ b/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