Merge pull request #1630 from newrelic/feat-langgraph

LangChain / LangGraph Agent Instrumentation

Author: umaannamalai

newrelic/common/llm_utils.py

@@ -12,13 +12,97 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import itertools
+import logging
+
+from newrelic.api.transaction import current_transaction
+from newrelic.common.object_wrapper import ObjectProxy
+
+_logger = logging.getLogger(__name__)
+
 
 def _get_llm_metadata(transaction):
-    # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events
-    custom_attrs_dict = transaction._custom_params
-    llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")}
-    llm_context_attrs = getattr(transaction, "_llm_context_attrs", None)
-    if llm_context_attrs:
-        llm_metadata_dict.update(llm_context_attrs)
+    if not transaction:
+        return {}
+    try:
+        # Grab LLM-related custom attributes off of the transaction to store as metadata on LLM events
+        custom_attrs_dict = getattr(transaction, "_custom_params", {})
+        llm_metadata_dict = {key: value for key, value in custom_attrs_dict.items() if key.startswith("llm.")}
+        llm_context_attrs = getattr(transaction, "_llm_context_attrs", None)
+        if llm_context_attrs:
+            llm_metadata_dict.update(llm_context_attrs)
+    except Exception:
+        _logger.warning("Unable to capture custom metadata attributes to record on LLM events.")
+        return {}
 
     return llm_metadata_dict
+
+
+class GeneratorProxy(ObjectProxy):
+    def __init__(self, wrapped, on_stop_iteration, on_error):
+        super().__init__(wrapped)
+        self._nr_on_stop_iteration = on_stop_iteration
+        self._nr_on_error = on_error
+
+    def __iter__(self):
+        self._nr_wrapped_iter = self.__wrapped__.__iter__()
+        return self
+
+    def __next__(self):
+        transaction = current_transaction()
+        if not transaction:
+            return self._nr_wrapped_iter.__next__()
+
+        return_val = None
+        try:
+            return_val = self._nr_wrapped_iter.__next__()
+        except StopIteration:
+            self._nr_on_stop_iteration(self, transaction)
+            raise
+        except Exception:
+            self._nr_on_error(self, transaction)
+            raise
+        return return_val
+
+    def close(self):
+        return self.__wrapped__.close()
+
+    def __copy__(self):
+        # Required to properly interface with itertool.tee, which can be called by LangChain on generators
+        self.__wrapped__, copy = itertools.tee(self.__wrapped__, 2)
+        return GeneratorProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)
+
+
+class AsyncGeneratorProxy(ObjectProxy):
+    def __init__(self, wrapped, on_stop_iteration, on_error):
+        super().__init__(wrapped)
+        self._nr_on_stop_iteration = on_stop_iteration
+        self._nr_on_error = on_error
+
+    def __aiter__(self):
+        self._nr_wrapped_iter = self.__wrapped__.__aiter__()
+        return self
+
+    async def __anext__(self):
+        transaction = current_transaction()
+        if not transaction:
+            return await self._nr_wrapped_iter.__anext__()
+
+        return_val = None
+        try:
+            return_val = await self._nr_wrapped_iter.__anext__()
+        except StopAsyncIteration:
+            self._nr_on_stop_iteration(self, transaction)
+            raise
+        except Exception:
+            self._nr_on_error(self, transaction)
+            raise
+        return return_val
+
+    async def aclose(self):
+        return await self.__wrapped__.aclose()
+
+    def __copy__(self):
+        # Required to properly interface with itertool.tee, which can be called by LangChain on generators
+        self.__wrapped__, copy = itertools.tee(self.__wrapped__, n=2)
+        return AsyncGeneratorProxy(copy, self._nr_on_stop_iteration, self._nr_on_error)

newrelic/config.py

@@ -2088,6 +2088,10 @@ def _process_module_builtin_defaults():
 
     _process_module_definition("asyncio.runners", "newrelic.hooks.coroutines_asyncio", "instrument_asyncio_runners")
 
+    _process_module_definition(
+        "langgraph.prebuilt.tool_node", "newrelic.hooks.mlmodel_langgraph", "instrument_langgraph_prebuilt_tool_node"
+    )
+
     _process_module_definition(
         "langchain_core.runnables.base",
         "newrelic.hooks.mlmodel_langchain",
@@ -2099,13 +2103,19 @@ def _process_module_builtin_defaults():
         "instrument_langchain_core_runnables_config",
     )
     _process_module_definition(
-        "langchain.chains.base", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_chains_base"
+        "langchain_core.tools.structured",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_core_tools_structured",
     )
+
     _process_module_definition(
-        "langchain_classic.chains.base", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_chains_base"
+        "langchain.agents.factory", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_agents_factory"
     )
     _process_module_definition(
-        "langchain_core.callbacks.manager", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_callbacks_manager"
+        "langchain.chains.base", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_chains_base"
+    )
+    _process_module_definition(
+        "langchain_classic.chains.base", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_chains_base"
     )
 
     # VectorStores with similarity_search method
@@ -2671,10 +2681,6 @@ def _process_module_builtin_defaults():
         "langchain_core.tools", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_core_tools"
     )
 
-    _process_module_definition(
-        "langchain_core.callbacks.manager", "newrelic.hooks.mlmodel_langchain", "instrument_langchain_callbacks_manager"
-    )
-
     _process_module_definition("asgiref.sync", "newrelic.hooks.adapter_asgiref", "instrument_asgiref_sync")
 
     _process_module_definition(

newrelic/hooks/external_botocore.py

@@ -1070,7 +1070,7 @@ def __next__(self):
         return return_val
 
     def close(self):
-        return super().close()
+        return self.__wrapped__.close()
 
 
 class AsyncEventStreamWrapper(ObjectProxy):
@@ -1108,7 +1108,7 @@ async def __anext__(self):
         return return_val
 
     async def aclose(self):
-        return await super().aclose()
+        return await self.__wrapped__.aclose()
 
 
 def handle_embedding_event(transaction, bedrock_attrs):