From a95614c38e77c33146b02494d19d40fa95d09f82 Mon Sep 17 00:00:00 2001
From: lievan <evan.li@datadoqhq.com>
Date: Fri, 19 Jul 2024 15:12:04 -0400
Subject: [PATCH] dev mode poc

---
 ddtrace/llmobs/_llmobs.py          |  7 +++++++
 ddtrace/llmobs/_trace_processor.py | 23 +++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/ddtrace/llmobs/_llmobs.py b/ddtrace/llmobs/_llmobs.py
index 1367ae231c4..012c2df4833 100644
--- a/ddtrace/llmobs/_llmobs.py
+++ b/ddtrace/llmobs/_llmobs.py
@@ -115,6 +115,7 @@ def enable(
         api_key: Optional[str] = None,
         env: Optional[str] = None,
         service: Optional[str] = None,
+        dev_mode: bool = True,
         _tracer: Optional[ddtrace.Tracer] = None,
     ) -> None:
         """
@@ -288,6 +289,7 @@ def llm(
         session_id: Optional[str] = None,
         ml_app: Optional[str] = None,
     ) -> Span:
+        print("[✧ LLM Observability] LLM ✨: {} running ...".format(name))
         """
         Trace an invocation call to an LLM where inputs and outputs are represented as text.
 
@@ -325,6 +327,7 @@ def tool(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_a
 
         :returns: The Span object representing the traced operation.
         """
+        print("[✧ LLM Observability] Tool 🔧: {} running ...".format(name))
         if cls.enabled is False:
             log.warning(SPAN_START_WHILE_DISABLED_WARNING)
         return cls._instance._start_span("tool", name=name, session_id=session_id, ml_app=ml_app)
@@ -341,12 +344,14 @@ def task(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_a
 
         :returns: The Span object representing the traced operation.
         """
+        print("[✧ LLM Observability] Task 📌: {} running...".format(name))
         if cls.enabled is False:
             log.warning(SPAN_START_WHILE_DISABLED_WARNING)
         return cls._instance._start_span("task", name=name, session_id=session_id, ml_app=ml_app)
 
     @classmethod
     def agent(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None) -> Span:
+        print("[✧ LLM Observability] Agent 🤖: {} running ...".format(name))
         """
         Trace a dynamic workflow in which an embedded language model (agent) decides what sequence of actions to take.
 
@@ -365,6 +370,7 @@ def agent(cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_
     def workflow(
         cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None
     ) -> Span:
+        print("[✧ LLM Observability] Workflow 🔗: {} running ...".format(name))
         """
         Trace a predefined or static sequence of operations.
 
@@ -422,6 +428,7 @@ def embedding(
     def retrieval(
         cls, name: Optional[str] = None, session_id: Optional[str] = None, ml_app: Optional[str] = None
     ) -> Span:
+        print("[✧ LLM Observability] Retrieval 🔎: {} running ...".format(name))
         """
         Trace a vector search operation involving a list of documents being returned from an external knowledge base.
 
diff --git a/ddtrace/llmobs/_trace_processor.py b/ddtrace/llmobs/_trace_processor.py
index 7c2e4608567..5b3426938d4 100644
--- a/ddtrace/llmobs/_trace_processor.py
+++ b/ddtrace/llmobs/_trace_processor.py
@@ -68,6 +68,7 @@ def submit_llmobs_span(self, span: Span) -> None:
     def _llmobs_span_event(self, span: Span) -> Dict[str, Any]:
         """Span event object structure."""
         span_kind = span._meta.pop(SPAN_KIND)
+
         meta: Dict[str, Any] = {"span.kind": span_kind, "input": {}, "output": {}}
         if span_kind in ("llm", "embedding") and span.get_tag(MODEL_NAME) is not None:
             meta["model_name"] = span._meta.pop(MODEL_NAME)
@@ -103,6 +104,28 @@ def _llmobs_span_event(self, span: Span) -> Dict[str, Any]:
         span.set_tag_str(SESSION_ID, session_id)
         parent_id = str(_get_llmobs_parent_id(span) or "undefined")
         span._meta.pop(PARENT_ID_KEY, None)
+
+        name = _get_span_name(span)
+        if span_kind == "llm":
+            print("[✧ LLM Observability] LLM ✨: {} finished in {} seconds!".format(name, span.duration))
+        elif span_kind == "workflow":
+            print("[✧ LLM Observability] Workflow 🔗: {} finished in {} seconds!".format(name, span.duration))
+        elif span_kind == "agent":
+            print("[✧ LLM Observability] Agent 🤖: {} finished in {} seconds!".format(name, span.duration))
+            url = """
+            View your agent run:
+            https://app.datadoghq.com/llm/traces?query=%40event_type%3Aspan%20%40parent_id%3Aundefined%20%40trace_id%3A{}%20&agg_m=count&agg_m_source=base&agg_t=count&fromUser=false&llmPanels=%5B%7B%22t%22%3A%22sampleDetailPanel%22%2C%22rEID%22%3A%22AgAAAZDMT2fSc-LOggAAAAAAAAAYAAAAAEFaRE1UMS1vQUFBMl9fZXBadnc3QUFBQQAAACQAAAAAMDE5MGNjNGYtODc3MC00YmY0LTg5NGItZmFiNTY1NDk1ZjE0%22%7D%5D&sidepanelTab=trace&viz=stream
+            """.format(
+                span.trace_id
+            )
+            print(url)
+        elif span_kind == "tool":
+            print("[✧ LLM Observability] Tool 🔧: {} finished in {} seconds!".format(name, span.duration))
+        elif span_kind == "task":
+            print("[✧ LLM Observability] Task 📌: {} finished in {} seconds!".format(name, span.duration))
+        elif span_kind == "retrieval":
+            print("[✧ LLM Observability] Retrieval 🔎: {} finished in {} seconds!".format(name, span.duration))
+
         return {
             "trace_id": "{:x}".format(span.trace_id),
             "span_id": str(span.span_id),