ensure initial flow state

Author: Yun-Kim

ddtrace/contrib/internal/crewai/patch.py

@@ -158,10 +158,15 @@ async def traced_flow_method(crewai, pin, func, instance, args, kwargs):
         submit_to_llmobs=True,
         flow_instance=instance,
     ) as span:
-        initial_flow_state = getattr(instance, "state", {})
+        flow_state = getattr(instance, "state", {})
+        initial_flow_state = {}
+        if isinstance(flow_state, dict):
+            initial_flow_state = {**flow_state}
+        elif hasattr(flow_state, "model_dump"):
+            initial_flow_state = flow_state.model_dump()
         result = await func(*args, **kwargs)
         kwargs["_dd.instance"] = instance
-        kwargs["_dd.flow_state"] = initial_flow_state
+        kwargs["_dd.initial_flow_state"] = initial_flow_state
         integration.llmobs_set_tags(span, args=args, kwargs=kwargs, response=result, operation="flow_method")
         return result
 

ddtrace/llmobs/_integrations/crewai.py

@@ -280,13 +280,7 @@ def _llmobs_set_tags_flow(self, span, args, kwargs, response):
 
     def _llmobs_set_tags_flow_method(self, span, args, kwargs, response):
         flow_instance = kwargs.pop("_dd.instance", None)
-        state_value = kwargs.pop("_dd.flow_state", {})
-        initial_flow_state = {}
-        if isinstance(state_value, dict):
-            initial_flow_state = state_value
-        elif hasattr(state_value, "model_dump"):
-            initial_flow_state = state_value.model_dump()
-
+        initial_flow_state = kwargs.pop("_dd.initial_flow_state", {})
         input_dict = {
             "args": [safe_json(arg) for arg in args[2:]],
             "kwargs": {k: safe_json(v) for k, v in kwargs.items()},