lcc add root_span_id

Author: JqRrt

cozeloop/integration/langchain/trace_callback.py

@@ -76,6 +76,7 @@ def __init__(
         self.tags_fn = tags_fn
         self._tags = tags if tags else {}
         self.trace_id: Optional[str] = None
+        self.root_span_id: Optional[str] = None
         self._id_lock = threading.Lock()
         self._child_of = child_of
         self._state_span_ctx_key = state_span_ctx_key
@@ -357,11 +358,14 @@ def _new_flow_span(self, node_name: str, span_type: str, **kwargs: Any) -> Span:
         span_name = node_name
         # set parent span
         parent_span: Span = None
+        is_root_span = False
         if 'parent_run_id' in kwargs and kwargs['parent_run_id'] is not None and str(kwargs['parent_run_id']) in self.run_map:
             parent_span = self.run_map[str(kwargs['parent_run_id'])].span
         # only root span use child_of
-        if parent_span is None and self._child_of:
-            parent_span = self._child_of
+        if parent_span is None:
+            is_root_span = True
+            if self._child_of:
+                parent_span = self._child_of
         # modify name
         error_tag = {}
         try:
@@ -373,6 +377,11 @@ def _new_flow_span(self, node_name: str, span_type: str, **kwargs: Any) -> Span:
             error_tag = {'error_info': f'name_fn error {repr(e)}', 'error_trace': traceback.format_exc()}
         # new span
         flow_span = self._client.start_span(span_name, span_type, child_of=parent_span)
+        if is_root_span:
+            if hasattr(flow_span, 'context'):
+                self.root_span_id = flow_span.context.span_id
+            else:
+                self.trace_id = flow_span.span_id
         run_id = str(kwargs['run_id'])
         self.run_map[run_id] = Run(run_id, flow_span, span_type)
         # set runtime