From 82b5b77940e97f65179efa0268031c47d0584a1c Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Tue, 24 Sep 2024 20:18:20 -0700
Subject: [PATCH] [Core] Add more interops tests (#26841)

To test that the client propagates both ways
---
 libs/core/langchain_core/callbacks/manager.py |  6 ++-
 libs/core/langchain_core/tracers/core.py      |  2 +-
 libs/core/langchain_core/tracers/langchain.py |  2 +-
 .../runnables/test_tracing_interops.py        | 42 ++++++++++++-------
 .../unit_tests/tracers/test_langchain.py      |  2 +-
 5 files changed, 36 insertions(+), 18 deletions(-)

diff --git a/libs/core/langchain_core/callbacks/manager.py b/libs/core/langchain_core/callbacks/manager.py
index 1f1c2cd8f44a3..c3214ea1123bd 100644
--- a/libs/core/langchain_core/callbacks/manager.py
+++ b/libs/core/langchain_core/callbacks/manager.py
@@ -2333,7 +2333,11 @@ def _configure(
                 try:
                     handler = LangChainTracer(
                         project_name=tracer_project,
-                        client=run_tree.client if run_tree is not None else None,
+                        client=(
+                            run_tree.client
+                            if run_tree is not None
+                            else tracing_context["client"]
+                        ),
                     )
                     callback_manager.add_handler(handler, True)
                 except Exception as e:
diff --git a/libs/core/langchain_core/tracers/core.py b/libs/core/langchain_core/tracers/core.py
index f4d7bbdccbf55..0cadbf44babcf 100644
--- a/libs/core/langchain_core/tracers/core.py
+++ b/libs/core/langchain_core/tracers/core.py
@@ -118,7 +118,7 @@ def _start_trace(self, run: Run) -> Union[None, Coroutine[Any, Any, None]]:  # t
                     self._add_child_run(parent_run, run)
             else:
                 if self.log_missing_parent:
-                    logger.warning(
+                    logger.debug(
                         f"Parent run {run.parent_run_id} not found for run {run.id}."
                         " Treating as a root run."
                     )
diff --git a/libs/core/langchain_core/tracers/langchain.py b/libs/core/langchain_core/tracers/langchain.py
index 09c053f9bce35..e04fccaa103d3 100644
--- a/libs/core/langchain_core/tracers/langchain.py
+++ b/libs/core/langchain_core/tracers/langchain.py
@@ -123,7 +123,7 @@ def _start_trace(self, run: Run) -> None:
 
         super()._start_trace(run)
         if run._client is None:
-            run._client = self.client
+            run._client = self.client  # type: ignore
 
     def on_chat_model_start(
         self,
diff --git a/libs/core/tests/unit_tests/runnables/test_tracing_interops.py b/libs/core/tests/unit_tests/runnables/test_tracing_interops.py
index b6318ef5b3bbb..3054b04f89b5c 100644
--- a/libs/core/tests/unit_tests/runnables/test_tracing_interops.py
+++ b/libs/core/tests/unit_tests/runnables/test_tracing_interops.py
@@ -10,6 +10,7 @@
 from langsmith.run_helpers import tracing_context
 from langsmith.run_trees import RunTree
 from langsmith.utils import get_env_var
+from typing_extensions import Literal
 
 from langchain_core.runnables.base import RunnableLambda, RunnableParallel
 from langchain_core.tracers.langchain import LangChainTracer
@@ -361,7 +362,8 @@ def parent(a: int) -> int:
             assert dotted_order.split(".")[0] == dotted_order
 
 
-def test_tree_is_constructed() -> None:
+@pytest.mark.parametrize("parent_type", ("ls", "lc"))
+def test_tree_is_constructed(parent_type: Literal["ls", "lc"]) -> None:
     mock_session = MagicMock()
     mock_client_ = Client(
         session=mock_session, api_key="test", auto_batch_tracing=False
@@ -379,27 +381,39 @@ def grandchild(x: str) -> str:
     def child(x: str) -> str:
         return grandchild.invoke(x)
 
-    @traceable
-    def parent() -> str:
-        return child.invoke("foo")
-
-    collected: dict[str, RunTree] = {}  # noqa
-
-    def collect_run(run: RunTree) -> None:
-        collected[str(run.id)] = run
-
     rid = uuid.uuid4()
-
     with tracing_context(
         client=mock_client_,
         enabled=True,
         metadata={"some_foo": "some_bar"},
         tags=["afoo"],
     ):
-        assert parent(langsmith_extra={"on_end": collect_run, "run_id": rid}) == "foo"
+        if parent_type == "ls":
+            collected: dict[str, RunTree] = {}  # noqa
+
+            def collect_run(run: RunTree) -> None:
+                collected[str(run.id)] = run
+
+            @traceable
+            def parent() -> str:
+                return child.invoke("foo")
+
+            assert (
+                parent(langsmith_extra={"on_end": collect_run, "run_id": rid}) == "foo"
+            )
+            assert collected
+            run = collected.get(str(rid))
+
+        else:
+
+            @RunnableLambda
+            def parent(_) -> str:  # type: ignore
+                return child.invoke("foo")
+
+            tracer = LangChainTracer()
+            assert parent.invoke(..., {"run_id": rid, "callbacks": [tracer]}) == "foo"  # type: ignore
+            run = tracer.latest_run
 
-    assert collected
-    run = collected.get(str(rid))
     assert run is not None
     assert run.name == "parent"
     assert run.child_runs
diff --git a/libs/core/tests/unit_tests/tracers/test_langchain.py b/libs/core/tests/unit_tests/tracers/test_langchain.py
index a8c897615ac3b..7e03717bd08fd 100644
--- a/libs/core/tests/unit_tests/tracers/test_langchain.py
+++ b/libs/core/tests/unit_tests/tracers/test_langchain.py
@@ -65,7 +65,7 @@ def new_persist_run_single(run: Run) -> None:
 def test_tracer_with_run_tree_parent() -> None:
     mock_session = unittest.mock.MagicMock()
     client = Client(session=mock_session, api_key="test")
-    parent = RunTree(name="parent", inputs={"input": "foo"}, _client=client)
+    parent = RunTree(name="parent", inputs={"input": "foo"}, _client=client)  # type: ignore
     run_id = uuid.uuid4()
     tracer = LangChainTracer(client=client)
     tracer.order_map[parent.id] = (parent.trace_id, parent.dotted_order)