feat(llmobs): submit langchain similarity search spans (#9971)

This PR adds instrumentation to submit langchain vector similarity
search spans to LLM Observability. The similarity search spans sent to
LLM Observability will contain the following I/O data:

input.value: a text field annotating the search query
output.value: The output document objects for the vector similarity
search

Ticket : https://datadoghq.atlassian.net/browse/MLOB-938

## Checklist
- [x] PR author has checked that all the criteria below are met
- The PR description includes an overview of the change
- The PR description articulates the motivation for the change
- The change includes tests OR the PR description describes a testing
strategy
- The PR description notes risks associated with the change, if any
- Newly-added code is easy to change
- The change follows the [library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
- The change includes or references documentation updates if necessary
- Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist
- [x] Reviewer has checked that all the criteria below are met 
- Title is accurate
- All changes are related to the pull request's stated goal
- Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes
- Testing strategy adequately addresses listed risks
- Newly-added code is easy to change
- Release note makes sense to a user of the library
- If necessary, author has acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment
- Backport labels are set in a manner that is consistent with the
[release branch maintenance
policy](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting)

---------

Co-authored-by: Yun Kim <[email protected]>

ddtrace/contrib/internal/langchain/patch.py

@@ -905,6 +905,7 @@ def traced_similarity_search(langchain, pin, func, instance, args, kwargs):
     span = integration.trace(
         pin,
         "%s.%s" % (instance.__module__, instance.__class__.__name__),
+        submit_to_llmobs=True,
         interface_type="similarity_search",
         provider=provider,
         api_key=_extract_api_key(instance),
@@ -947,6 +948,14 @@ def traced_similarity_search(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(
+                "retrieval",
+                span,
+                query,
+                documents,
+                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

@@ -17,6 +17,7 @@
 from ddtrace.llmobs._constants import METRICS
 from ddtrace.llmobs._constants import MODEL_NAME
 from ddtrace.llmobs._constants import MODEL_PROVIDER
+from ddtrace.llmobs._constants import OUTPUT_DOCUMENTS
 from ddtrace.llmobs._constants import OUTPUT_MESSAGES
 from ddtrace.llmobs._constants import OUTPUT_VALUE
 from ddtrace.llmobs._constants import SPAN_KIND
@@ -44,15 +45,15 @@
     "system": "system",
 }
 
-SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding"]
+SUPPORTED_OPERATIONS = ["llm", "chat", "chain", "embedding", "retrieval"]
 
 
 class LangChainIntegration(BaseLLMIntegration):
     _integration_name = "langchain"
 
     def llmobs_set_tags(
         self,
-        operation: str,  # oneof "llm","chat","chain","embedding"
+        operation: str,  # oneof "llm","chat","chain","embedding","retrieval"
         span: Span,
         inputs: Any,
         response: Any = None,
@@ -89,6 +90,8 @@ def llmobs_set_tags(
             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)
+        elif operation == "retrieval":
+            self._llmobs_set_meta_tags_from_similarity_search(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:
@@ -260,6 +263,47 @@ def _llmobs_set_meta_tags_from_embedding(
             except (TypeError, IndexError):
                 log.warning("Failed to write output vectors", output_embedding)
 
+    def _llmobs_set_meta_tags_from_similarity_search(
+        self,
+        span: Span,
+        input_query: str,
+        output_documents: Union[List[Any], None],
+        error: bool = False,
+        is_workflow: bool = False,
+    ) -> None:
+        span.set_tag_str(SPAN_KIND, "workflow" if is_workflow else "retrieval")
+        span.set_tag_str(MODEL_NAME, span.get_tag(MODEL) or "")
+        span.set_tag_str(MODEL_PROVIDER, span.get_tag(PROVIDER) or "")
+
+        if input_query is not None:
+            try:
+                formatted_inputs = self.format_io(input_query)
+                if isinstance(formatted_inputs, str):
+                    span.set_tag_str(INPUT_VALUE, formatted_inputs)
+                else:
+                    span.set_tag_str(INPUT_VALUE, json.dumps(formatted_inputs))
+            except TypeError:
+                log.warning("Failed to serialize similarity search input to JSON")
+        if error:
+            span.set_tag_str(OUTPUT_VALUE, "")
+        elif isinstance(output_documents, list):
+            if is_workflow:
+                span.set_tag_str(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(output_documents)))
+            else:
+                documents = []
+                for d in output_documents:
+                    doc = Document(text=d.page_content)
+                    doc["id"] = getattr(d, "id", "")
+                    metadata = getattr(d, "metadata", {})
+                    doc["name"] = metadata.get("name", doc["id"])
+                    documents.append(doc)
+                try:
+                    span.set_tag_str(OUTPUT_DOCUMENTS, json.dumps(self.format_io(documents)))
+                    # we set the value as well to ensure that the UI would display it in case the span was the root
+                    span.set_tag_str(OUTPUT_VALUE, "[{} document(s) retrieved]".format(len(documents)))
+                except TypeError:
+                    log.warning("Failed to serialize similarity output documents to JSON")
+
     def _set_base_span_tags(  # type: ignore[override]
         self,
         span: Span,

releasenotes/notes/feat-llmobs-submit-langchain-similarity-search-spans-c5581deda8b6878e.yaml

@@ -0,0 +1,4 @@
+---
+features:
+  - |
+    LLM Observability: The LangChain integration now submits vectorstore ``similarity_search`` spans to LLM Observability as retrieval spans.