feat(llmobs): submit langchain similarity search spans

ddtrace/contrib/langchain/patch.py

@@ -936,6 +936,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

@@ -16,16 +16,17 @@
 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
 
+from ..utils import Document
 from .base import BaseLLMIntegration
 
 
 log = get_logger(__name__)
 
-
 API_KEY = "langchain.request.api_key"
 MODEL = "langchain.request.model"
 PROVIDER = "langchain.request.provider"
@@ -48,7 +49,7 @@ class LangChainIntegration(BaseLLMIntegration):
 
     def llmobs_set_tags(
         self,
-        operation: str,  # oneof "llm","chat","chain"
+        operation: str,  # oneof "llm","chat","chain","retrieval"
         span: Span,
         inputs: Any,
         response: Any = None,
@@ -79,6 +80,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 == "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:
@@ -194,6 +197,53 @@ 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_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 "")
+
+        input_tag_key = INPUT_VALUE if is_workflow else INPUT_MESSAGES
+        output_tag_key = OUTPUT_VALUE if is_workflow else OUTPUT_DOCUMENTS
+
+        if input_query is not None:
+            try:
+                formatted_inputs = self.format_io(input_query)
+                if isinstance(formatted_inputs, str):
+                    span.set_tag_str(input_tag_key, formatted_inputs)
+                else:
+                    span.set_tag_str(input_tag_key, json.dumps(self.format_io(input_query)))
+            except TypeError:
+                log.warning("Failed to serialize similarity input query to JSON")
+        if error:
+            span.set_tag_str(output_tag_key, "")
+        elif isinstance(output_documents, list):
+            if is_workflow:
+                span.set_tag_str(output_tag_key, json.dumps(output_documents))
+            else:
+                try:
+                    documents = [
+                        Document(
+                            id=d.id,
+                            text=d.page_content,
+                            name=d.metadata["source"],
+                        )
+                        for d in output_documents
+                    ]
+                except AttributeError:
+                    log.warning("Failed to extract document information from similarity output")
+                    documents = [Document(text=d.page_content) for d in output_documents]
+                try:
+                    span.set_tag_str(output_tag_key, json.dumps(self.format_io(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 similarity search spans to LLM Observability.

tests/contrib/langchain/conftest.py

@@ -147,3 +147,19 @@ def langchain_anthropic(ddtrace_config_langchain, mock_logs, mock_metrics, langc
         yield langchain_anthropic
     except ImportError:
         yield
+
+
+@pytest.fixture
+def langchain_pinecone(ddtrace_config_langchain, mock_logs, mock_metrics, langchain):
+    with override_env(
+        dict(
+            PINECONE_API_KEY=os.getenv("PINECONE_API_KEY", "<not-a-real-key>"),
+            PINECONE_ENV=os.getenv("PINECONE_ENV", "<not-a-real-env>"),
+        )
+    ):
+        try:
+            import langchain_pinecone
+
+            yield langchain_pinecone
+        except ImportError:
+            yield

tests/contrib/langchain/test_langchain_llmobs.py

@@ -123,6 +123,14 @@ def _invoke_chain(cls, chain, prompt, mock_tracer, cassette_name, batch=False):
         LLMObs.disable()
         return mock_tracer.pop_traces()[0]
 
+    @classmethod
+    def _similarity_search(cls, vector_db, prompt, k, mock_tracer, cassette_name):
+        LLMObs.enable(ml_app=cls.ml_app, integrations_enabled=False, _tracer=mock_tracer)
+        with get_request_vcr(subdirectory_name=cls.cassette_subdirectory_name).use_cassette(cassette_name):
+            vector_db.similarity_search(prompt, k)
+        LLMObs.disable()
+        return mock_tracer.pop_traces()[0][0]
+
 
 @pytest.mark.skipif(not PATCH_LANGCHAIN_V0, reason="These tests are for langchain < 0.1.0")
 class TestLLMObsLangchain(BaseTestLLMObsLangchain):
@@ -324,6 +332,41 @@ def test_llmobs_chain_schema_io(self, langchain, mock_llmobs_span_writer, mock_t
         )
         _assert_expected_llmobs_llm_span(trace[1], mock_llmobs_span_writer, mock_io=True)
 
+    @pytest.mark.skipif(sys.version_info < (3, 10, 0), reason="Requires unnecessary cassette file for Python 3.9")
+    def test_llmobs_similarity_search(self, langchain, langchain_pinecone, mock_llmobs_span_writer, mock_tracer):
+        if langchain_pinecone is None:
+            pytest.skip("langchain_pinecone not installed which is required for this test.")
+        embedding = langchain.embeddings.OpenAIEmbeddings(model="text-embedding-ada-002")
+        index_name = "langchain-retrieval"
+        namespace = "langchain-retrieval"
+        vectorstore = langchain_pinecone.PineconeVectorStore(
+            index_name=index_name,
+            embedding=embedding,
+            namespace=namespace,
+        )
+        trace = self._similarity_search(
+            vectorstore, "Who was Alan Turing?", 1, mock_tracer, "openai_pinecone_similarity_search.yaml"
+        )
+        assert mock_llmobs_span_writer.enqueue.call_count == 1
+        span = trace[0] if isinstance(trace, list) else trace
+        mock_llmobs_span_writer.enqueue.assert_called_with(
+            _expected_llmobs_non_llm_span_event(
+                span,
+                span_kind="similarity_search",
+                input_value="Who was Alan Turing?",
+                output_value=[
+                    {
+                        "id": 13,
+                        "title": "Alan Turing",
+                        "text": "A brilliant mathematician and cryptographer Alan was to become the founder of "
+                        "modern-day computer science and artificial intelli...",
+                    }
+                ],
+                tags={"ml_app": ""},
+                integration="langchain",
+            )
+        )
+
 
 @flaky(1735812000, reason="Community cassette tests are flaky")
 @pytest.mark.skipif(PATCH_LANGCHAIN_V0, reason="These tests are for langchain >= 0.1.0")