fix(langchain): add patch test cases, fix double patch issue for embeddings/vectorstores [backport 1.17] (#6512)

Backport 9ec214f from #6475 to 1.17.

This PR adds patch test cases for the langchain integration, and fixes a
bug where we were double patching
`langchain.embeddings.HuggingFaceEmbeddings.embed_query/documents` and
`langchain.vectorstores.Milvus.similarity_search`. This issue stemmed
from the two classes being reused/inherited by other classes
(`SentenceTransformerEmbeddings`, `Zilliz` respectively), meaning that
when we wrapped the latter two classes' methods, we unintentionally
wrapped the former two classes' methods twice.

The fix involves checking that the function to wrap isn't already a
wrapped method (specifically a `wrapt.ObjectProxy` which we use in our
integrations). This could potentially be problematic if a user wraps
this before running our patching code as we would ignore this, but at
the moment of writing I don't see a huge risk nor a cleaner solution to
avoid double patching.

## Checklist

- [x] Change(s) are motivated and described in the PR description.
- [x] Testing strategy is described if automated tests are not included
in the PR.
- [x] Risk is outlined (performance impact, potential for breakage,
maintainability, etc).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] [Library release note
guidelines](https://ddtrace.readthedocs.io/en/stable/releasenotes.html)
are followed. If no release note is required, add label
`changelog/no-changelog`.
- [x] Documentation is included (in-code, generated user docs, [public
corp docs](https://github.com/DataDog/documentation/)).
- [x] Backport labels are set (if
[applicable](https://ddtrace.readthedocs.io/en/latest/contributing.html#backporting))

## Reviewer Checklist

- [x] Title is accurate.
- [x] No unnecessary changes are introduced.
- [x] Description motivates each change.
- [x] Avoids breaking
[API](https://ddtrace.readthedocs.io/en/stable/versioning.html#interfaces)
changes unless absolutely necessary.
- [x] Testing strategy adequately addresses listed risk(s).
- [x] Change is maintainable (easy to change, telemetry, documentation).
- [x] Release note makes sense to a user of the library.
- [x] Reviewer has explicitly acknowledged and discussed the performance
implications of this PR as reported in the benchmarks PR comment.
- [x] 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]>

Author: github-actions[bot]

ddtrace/contrib/langchain/patch.py

@@ -30,6 +30,7 @@
 from ddtrace.internal.utils.formats import asbool
 from ddtrace.internal.utils.formats import deep_getattr
 from ddtrace.pin import Pin
+from ddtrace.vendor import wrapt
 
 
 if TYPE_CHECKING:
@@ -750,9 +751,8 @@ def patch():
             )
         integration.start_log_writer()
 
-    # TODO: check if we need to version gate LLM/Chat/TextEmbedding
     wrap("langchain", "llms.base.BaseLLM.generate", traced_llm_generate(langchain))
-    wrap("langchain", "llms.BaseLLM.agenerate", traced_llm_agenerate(langchain))
+    wrap("langchain", "llms.base.BaseLLM.agenerate", traced_llm_agenerate(langchain))
     wrap("langchain", "chat_models.base.BaseChatModel.generate", traced_chat_model_generate(langchain))
     wrap("langchain", "chat_models.base.BaseChatModel.agenerate", traced_chat_model_agenerate(langchain))
     wrap("langchain", "chains.base.Chain.__call__", traced_chain_call(langchain))
@@ -761,18 +761,32 @@ def patch():
     #  wrap each langchain-provided text embedding model.
     for text_embedding_model in text_embedding_models:
         if hasattr(langchain.embeddings, text_embedding_model):
-            wrap("langchain", "embeddings.%s.embed_query" % text_embedding_model, traced_embedding(langchain))
-            wrap("langchain", "embeddings.%s.embed_documents" % text_embedding_model, traced_embedding(langchain))
-            # TODO: langchain >= 0.0.209 includes async embedding implementation (only for OpenAI)
+            # Ensure not double patched, as some Embeddings interfaces are pointers to other Embeddings.
+            if not isinstance(
+                deep_getattr(langchain.embeddings, "%s.embed_query" % text_embedding_model), wrapt.ObjectProxy
+            ):
+                wrap("langchain", "embeddings.%s.embed_query" % text_embedding_model, traced_embedding(langchain))
+            if not isinstance(
+                deep_getattr(langchain.embeddings, "%s.embed_documents" % text_embedding_model), wrapt.ObjectProxy
+            ):
+                wrap("langchain", "embeddings.%s.embed_documents" % text_embedding_model, traced_embedding(langchain))
+                # TODO: langchain >= 0.0.209 includes async embedding implementation (only for OpenAI)
     # We need to do the same with Vectorstores.
     for vectorstore in vectorstores:
         if hasattr(langchain.vectorstores, vectorstore):
-            wrap("langchain", "vectorstores.%s.similarity_search" % vectorstore, traced_similarity_search(langchain))
+            # Ensure not double patched, as some Embeddings interfaces are pointers to other Embeddings.
+            if not isinstance(
+                deep_getattr(langchain.vectorstores, "%s.similarity_search" % vectorstore), wrapt.ObjectProxy
+            ):
+                wrap(
+                    "langchain", "vectorstores.%s.similarity_search" % vectorstore, traced_similarity_search(langchain)
+                )
 
 
 def unpatch():
-    if getattr(langchain, "_datadog_patch", False):
-        setattr(langchain, "_datadog_patch", False)
+    if not getattr(langchain, "_datadog_patch", False):
+        return
+    setattr(langchain, "_datadog_patch", False)
 
     unwrap(langchain.llms.base.BaseLLM, "generate")
     unwrap(langchain.llms.base.BaseLLM, "agenerate")
@@ -782,10 +796,19 @@ def unpatch():
     unwrap(langchain.chains.base.Chain, "acall")
     for text_embedding_model in text_embedding_models:
         if hasattr(langchain.embeddings, text_embedding_model):
-            unwrap(getattr(langchain.embeddings, text_embedding_model), "embed_query")
-            unwrap(getattr(langchain.embeddings, text_embedding_model), "embed_documents")
+            if isinstance(
+                deep_getattr(langchain.embeddings, "%s.embed_query" % text_embedding_model), wrapt.ObjectProxy
+            ):
+                unwrap(getattr(langchain.embeddings, text_embedding_model), "embed_query")
+            if isinstance(
+                deep_getattr(langchain.embeddings, "%s.embed_documents" % text_embedding_model), wrapt.ObjectProxy
+            ):
+                unwrap(getattr(langchain.embeddings, text_embedding_model), "embed_documents")
     for vectorstore in vectorstores:
         if hasattr(langchain.vectorstores, vectorstore):
-            unwrap(getattr(langchain.vectorstores, vectorstore), "similarity_search")
+            if isinstance(
+                deep_getattr(langchain.vectorstores, "%s.similarity_search" % vectorstore), wrapt.ObjectProxy
+            ):
+                unwrap(getattr(langchain.vectorstores, vectorstore), "similarity_search")
 
     delattr(langchain, "_datadog_integration")

releasenotes/notes/fix-langchain-patching-800e4d50f4b2769a.yaml

@@ -0,0 +1,6 @@
+---
+fixes:
+  - |
+    langchain: This fix resolves an issue where ``langchain.embeddings.HuggingFaceEmbeddings`` embedding
+     methods, and ``langchain.vectorstores.Milvus.similarity_search`` were patched twice
+     due to a nested class hierarchy in ``langchain``. 

tests/contrib/langchain/test_langchain_patch.py

@@ -0,0 +1,63 @@
+from ddtrace.contrib.langchain import patch
+from ddtrace.contrib.langchain import unpatch
+from ddtrace.contrib.langchain.constants import text_embedding_models
+from ddtrace.contrib.langchain.constants import vectorstores
+from tests.contrib.patch import PatchTestCase
+
+
+class TestLangchainPatch(PatchTestCase.Base):
+    __integration_name__ = "langchain"
+    __module_name__ = "langchain"
+    __patch_func__ = patch
+    __unpatch_func__ = unpatch
+
+    def assert_module_patched(self, langchain):
+        self.assert_wrapped(langchain.llms.base.BaseLLM.generate)
+        self.assert_wrapped(langchain.llms.base.BaseLLM.agenerate)
+        self.assert_wrapped(langchain.chat_models.base.BaseChatModel.generate)
+        self.assert_wrapped(langchain.chat_models.base.BaseChatModel.agenerate)
+        self.assert_wrapped(langchain.chains.base.Chain.__call__)
+        self.assert_wrapped(langchain.chains.base.Chain.acall)
+        for text_embedding_model in text_embedding_models:
+            embedding_model = getattr(langchain.embeddings, text_embedding_model, None)
+            if embedding_model:
+                self.assert_wrapped(embedding_model.embed_query)
+                self.assert_wrapped(embedding_model.embed_documents)
+        for vectorstore in vectorstores:
+            vectorstore_interface = getattr(langchain.vectorstores, vectorstore, None)
+            if vectorstore_interface:
+                self.assert_wrapped(vectorstore_interface.similarity_search)
+
+    def assert_not_module_patched(self, langchain):
+        self.assert_not_wrapped(langchain.llms.base.BaseLLM.generate)
+        self.assert_not_wrapped(langchain.llms.base.BaseLLM.agenerate)
+        self.assert_not_wrapped(langchain.chat_models.base.BaseChatModel.generate)
+        self.assert_not_wrapped(langchain.chat_models.base.BaseChatModel.agenerate)
+        self.assert_not_wrapped(langchain.chains.base.Chain.__call__)
+        self.assert_not_wrapped(langchain.chains.base.Chain.acall)
+        for text_embedding_model in text_embedding_models:
+            embedding_model = getattr(langchain.embeddings, text_embedding_model, None)
+            if embedding_model:
+                self.assert_not_wrapped(embedding_model.embed_query)
+                self.assert_not_wrapped(embedding_model.embed_documents)
+        for vectorstore in vectorstores:
+            vectorstore_interface = getattr(langchain.vectorstores, vectorstore, None)
+            if vectorstore_interface:
+                self.assert_not_wrapped(vectorstore_interface.similarity_search)
+
+    def assert_not_module_double_patched(self, langchain):
+        self.assert_not_double_wrapped(langchain.llms.base.BaseLLM.generate)
+        self.assert_not_double_wrapped(langchain.llms.base.BaseLLM.agenerate)
+        self.assert_not_double_wrapped(langchain.chat_models.base.BaseChatModel.generate)
+        self.assert_not_double_wrapped(langchain.chat_models.base.BaseChatModel.agenerate)
+        self.assert_not_double_wrapped(langchain.chains.base.Chain.__call__)
+        self.assert_not_double_wrapped(langchain.chains.base.Chain.acall)
+        for text_embedding_model in text_embedding_models:
+            embedding_model = getattr(langchain.embeddings, text_embedding_model, None)
+            if embedding_model:
+                self.assert_not_double_wrapped(embedding_model.embed_query)
+                self.assert_not_double_wrapped(embedding_model.embed_documents)
+        for vectorstore in vectorstores:
+            vectorstore_interface = getattr(langchain.vectorstores, vectorstore, None)
+            if vectorstore_interface:
+                self.assert_not_double_wrapped(vectorstore_interface.similarity_search)