feat(agent): Add trace for agent (#1407)

dbgpt/util/tracer/base.py

@@ -15,6 +15,7 @@ class SpanType(str, Enum):
     BASE = "base"
     RUN = "run"
     CHAT = "chat"
+    AGENT = "agent"
 
 
 class SpanTypeRunName(str, Enum):
@@ -99,6 +100,21 @@ class Span:
             "metadata": _clean_for_json(self.metadata),
         }
 
+    def copy(self) -> Span:
+        """Create a copy of this span."""
+        metadata = self.metadata.copy() if self.metadata else None
+        span = Span(
+            self.trace_id,
+            self.span_id,
+            self.span_type,
+            self.parent_span_id,
+            self.operation_name,
+            metadata=metadata,
+        )
+        span.start_time = self.start_time
+        span.end_time = self.end_time
+        return span
+
 
 class SpanStorageType(str, Enum):
     ON_CREATE = "on_create"
@@ -191,7 +207,7 @@ class TracerContext:
 
 
 def _clean_for_json(data: Optional[str, Any] = None):
-    if not data:
+    if data is None:
         return None
     if isinstance(data, dict):
         cleaned_dict = {}

dbgpt/util/tracer/span_storage.py

@@ -49,7 +49,9 @@ class SpanStorageContainer(SpanStorage):
         self.flush_thread = threading.Thread(
             target=self._flush_to_storages, daemon=True
         )
+        self._stop_event = threading.Event()
         self.flush_thread.start()
+        self._stop_event.clear()
 
     def append_storage(self, storage: SpanStorage):
         """Append sotrage to container
@@ -68,7 +70,7 @@ class SpanStorageContainer(SpanStorage):
                 pass  # If the signal queue is full, it's okay. The flush thread will handle it.
 
     def _flush_to_storages(self):
-        while True:
+        while not self._stop_event.is_set():
             interval = time.time() - self.last_flush_time
             if interval < self.flush_interval:
                 try:
@@ -90,13 +92,24 @@ class SpanStorageContainer(SpanStorage):
                     try:
                         storage.append_span_batch(spans_to_write)
                     except Exception as e:
-                        logger.warn(
+                        logger.warning(
                             f"Append spans to storage {str(storage)} failed: {str(e)}, span_data: {spans_to_write}"
                         )
 
-                self.executor.submit(append_and_ignore_error, s, spans_to_write)
+                try:
+                    self.executor.submit(append_and_ignore_error, s, spans_to_write)
+                except RuntimeError:
+                    append_and_ignore_error(s, spans_to_write)
             self.last_flush_time = time.time()
 
+    def before_stop(self):
+        try:
+            self.flush_signal_queue.put(True)
+            self._stop_event.set()
+            self.flush_thread.join()
+        except Exception:
+            pass
+
 
 class FileSpanStorage(SpanStorage):
     def __init__(self, filename: str):

dbgpt/util/tracer/tests/test_tracer_impl.py

@@ -52,7 +52,7 @@ def test_start_and_end_span(tracer: Tracer):
     assert span.end_time is not None
 
     stored_span = tracer._get_current_storage().spans[0]
-    assert stored_span == span
+    assert stored_span.span_id == span.span_id
 
 
 def test_start_and_end_span_with_tracer_manager(tracer_manager: TracerManager):
@@ -76,8 +76,12 @@ def test_parent_child_span_relation(tracer: Tracer):
     tracer.end_span(child_span)
     tracer.end_span(parent_span)
 
-    assert parent_span in tracer._get_current_storage().spans
-    assert child_span in tracer._get_current_storage().spans
+    assert parent_span.operation_name in [
+        s.operation_name for s in tracer._get_current_storage().spans
+    ]
+    assert child_span.operation_name in [
+        s.operation_name for s in tracer._get_current_storage().spans
+    ]
 
 
 @pytest.mark.parametrize(

dbgpt/util/tracer/tracer_impl.py

@@ -36,7 +36,7 @@ class DefaultTracer(Tracer):
         self._span_storage_type = span_storage_type
 
     def append_span(self, span: Span):
-        self._get_current_storage().append_span(span)
+        self._get_current_storage().append_span(span.copy())
 
     def start_span(
         self,
@@ -130,9 +130,13 @@ class TracerManager:
         """
         tracer = self._get_tracer()
         if not tracer:
-            return Span("empty_span", "empty_span")
+            return Span(
+                "empty_span", "empty_span", span_type=span_type, metadata=metadata
+            )
         if not parent_span_id:
             parent_span_id = self.get_current_span_id()
+        if not span_type and parent_span_id:
+            span_type = self._get_current_span_type()
         return tracer.start_span(
             operation_name, parent_span_id, span_type=span_type, metadata=metadata
         )
@@ -156,6 +160,10 @@ class TracerManager:
         ctx = self._trace_context_var.get()
         return ctx.span_id if ctx else None
 
+    def _get_current_span_type(self) -> Optional[SpanType]:
+        current_span = self.get_current_span()
+        return current_span.span_type if current_span else None
+
 
 root_tracer: TracerManager = TracerManager()
 
@@ -197,14 +205,19 @@ def _parse_operation_name(func, *args):
 
 
 def initialize_tracer(
-    system_app: SystemApp,
     tracer_filename: str,
     root_operation_name: str = "DB-GPT-Web-Entry",
-    tracer_storage_cls: str = None,
+    system_app: Optional[SystemApp] = None,
+    tracer_storage_cls: Optional[str] = None,
+    create_system_app: bool = False,
 ):
+    """Initialize the tracer with the given filename and system app."""
+    from dbgpt.util.tracer.span_storage import FileSpanStorage, SpanStorageContainer
+
+    if not system_app and create_system_app:
+        system_app = SystemApp()
     if not system_app:
         return
-    from dbgpt.util.tracer.span_storage import FileSpanStorage, SpanStorageContainer
 
     trace_context_var = ContextVar(
         "trace_context",