feat(llmobs): submit langchain embedding spans

ddtrace/contrib/langchain/patch.py

@@ -567,6 +567,7 @@ def traced_embedding(langchain, pin, func, instance, args, kwargs):
     """
     try:
         input_texts = get_argument_value(args, kwargs, 0, "texts")
+        chunk_size = get_argument_value(args, kwargs, 1, "chunk_Size")
     except ArgumentError:
         input_texts = get_argument_value(args, kwargs, 0, "text")
 
@@ -575,11 +576,13 @@ def traced_embedding(langchain, pin, func, instance, args, kwargs):
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__module__, instance.__class__.__name__),
+        submit_to_llmobs=True,
         interface_type="embedding",
         provider=provider,
         model=_extract_model_name(instance),
         api_key=_extract_api_key(instance),
     )
+    embeddings = None
     try:
         if isinstance(input_texts, str):
             if integration.is_pc_sampled_span(span):
@@ -590,6 +593,7 @@ def traced_embedding(langchain, pin, func, instance, args, kwargs):
                 for idx, text in enumerate(input_texts):
                     span.set_tag_str("langchain.request.inputs.%d.text" % idx, integration.trunc(text))
             span.set_metric("langchain.request.input_count", len(input_texts))
+            span.set_tag_str("langchain.request.inputs.chunk_size", chunk_size)
         # langchain currently does not support token tracking for OpenAI embeddings:
         #  https://github.com/hwchase17/langchain/issues/945
         embeddings = func(*args, **kwargs)
@@ -603,6 +607,14 @@ def traced_embedding(langchain, pin, func, instance, args, kwargs):
         integration.metric(span, "incr", "request.error", 1)
         raise
     finally:
+        if integration.is_pc_sampled_llmobs(span):
+            integration.llmobs_set_tags(
+                "embedding",
+                span,
+                input_texts=input_texts,
+                output_embedding=embeddings,
+                error=bool(span.error),
+            )
         span.finish()
         integration.metric(span, "dist", "request.duration", span.duration_ns)
         if integration.is_pc_sampled_log(span):

ddtrace/llmobs/_integrations/langchain.py

@@ -10,6 +10,7 @@
 from ddtrace.constants import ERROR_TYPE
 from ddtrace.internal.logger import get_logger
 from ddtrace.llmobs import LLMObs
+from ddtrace.llmobs._constants import INPUT_DOCUMENTS
 from ddtrace.llmobs._constants import INPUT_MESSAGES
 from ddtrace.llmobs._constants import INPUT_VALUE
 from ddtrace.llmobs._constants import METADATA
@@ -42,21 +43,31 @@
     "system": "system",
 }
 
+SUPPORTED_OPERATIONS = ["llm",
+                        "chat",
+                        "chain",
+                        "embedding"
+                        ]
+
 
 class LangChainIntegration(BaseLLMIntegration):
     _integration_name = "langchain"
 
     def llmobs_set_tags(
-        self,
-        operation: str,  # oneof "llm","chat","chain"
-        span: Span,
-        inputs: Any,
-        response: Any = None,
-        error: bool = False,
+            self,
+            operation: str,  # oneof "llm","chat","chain","embedding"
+            span: Span,
+            inputs: Any,
+            response: Any = None,
+            error: bool = False,
     ) -> None:
         """Sets meta tags and metrics for span events to be sent to LLMObs."""
         if not self.llmobs_enabled:
             return
+        if operation not in SUPPORTED_OPERATIONS:
+            log.warning("Unsupported operation : %s" % operation)
+            return
+
         model_provider = span.get_tag(PROVIDER)
         self._llmobs_set_metadata(span, model_provider)
 
@@ -79,6 +90,8 @@ def llmobs_set_tags(
             self._llmobs_set_meta_tags_from_chat_model(span, inputs, response, error, is_workflow=is_workflow)
         elif operation == "chain":
             self._llmobs_set_meta_tags_from_chain(span, inputs, response, error)
+        elif operation == "embedding":
+            self._llmobs_set_meta_tags_from_embedding(span, inputs, response, error, is_workflow=is_workflow)
         span.set_tag_str(METRICS, json.dumps({}))
 
     def _llmobs_set_metadata(self, span: Span, model_provider: Optional[str] = None) -> None:
@@ -194,6 +207,50 @@ def _llmobs_set_meta_tags_from_chain(
             except TypeError:
                 log.warning("Failed to serialize chain output data to JSON")
 
+    def _llmobs_set_meta_tags_from_embedding(
+            self,
+            span: Span,
+            input_texts: Union[str, List[str]],
+            output_embedding: Union[List[float], List[List[float]], None],
+            error: bool = False,
+            is_workflow: bool = False,
+    ) -> None:
+        span.set_tag_str(SPAN_KIND, "workflow" if is_workflow else "embedding")
+        span.set_tag_str(MODEL_NAME, span.get_tag(MODEL) or "")
+        span.set_tag_str(MODEL_PROVIDER, span.get_tag(PROVIDER) or "")
+
+        input_tag_key = INPUT_VALUE if is_workflow else INPUT_DOCUMENTS
+        output_tag_key = OUTPUT_VALUE
+
+        if (input_texts is not None
+                and (isinstance(input_texts, str) or all(isinstance(text, str) for text in input_texts))):
+            if is_workflow:
+                try:
+                    formatted_inputs = self.format_io(input_texts)
+                    if isinstance(input_texts, str):
+                        span.set_tag_str(input_tag_key, formatted_inputs)
+                    else:
+                        span.set_tag_str(input_tag_key, json.dumps(self.format_io(input_texts)))
+                except TypeError:
+                    log.warning("Failed to serialize embedding input data to JSON")
+            else:
+                if isinstance(input_texts, str):
+                    span.set_tag_str(input_tag_key, input_texts)
+                elif (isinstance(input_texts, list)
+                      and all(isinstance(text, str) for text in input_texts) and not is_workflow):
+                    json_documents_list = json.dumps(input_texts)
+                    span.set_tag_str(input_tag_key, '{"documents":%s}' % json_documents_list)
+
+        if error:
+            span.set_tag_str(output_tag_key, "")
+        elif output_embedding is not None:
+            if not isinstance(output_embedding[0], list):
+                output_embedding = [output_embedding]
+            embedding_dim = len(output_embedding[0])
+            span.set_tag_str(
+                output_tag_key, "[{} embedding(s) returned with size {}]".format(len(output_embedding), embedding_dim)
+            )
+
     def _set_base_span_tags(  # type: ignore[override]
         self,
         span: Span,