Add LLM attribute to LangChain (#1057)

* Add HanaDB and KDBAI instrumentation.

* Add LLM attr to transactions.

* Add testing to vectorstores.

* [Mega-Linter] Apply linters fixes

---------

Co-authored-by: umaannamalai <[email protected]>

newrelic/config.py

@@ -2423,6 +2423,11 @@ def _process_module_builtin_defaults():
         "instrument_langchain_vectorstore_similarity_search",
     )
 
+    _process_module_definition(
+        "langchain_community.vectorstores.thirdai_neuraldb",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
     _process_module_definition(
         "langchain_community.vectorstores.tigris",
         "newrelic.hooks.mlmodel_langchain",

newrelic/hooks/mlmodel_langchain.py

@@ -65,6 +65,7 @@
     "langchain_community.vectorstores.mongodb_atlas": "MongoDBAtlasVectorSearch",
     "langchain_community.vectorstores.myscale": "MyScale",
     "langchain_community.vectorstores.neo4j_vector": "Neo4jVector",
+    "langchain_community.vectorstores.thirdai_neuraldb": "NeuralDBVectorStore",
     "langchain_community.vectorstores.nucliadb": "NucliaDB",
     "langchain_community.vectorstores.opensearch_vector_search": "OpenSearchVectorSearch",
     "langchain_community.vectorstores.pgembedding": "PGEmbedding",
@@ -128,6 +129,7 @@ async def wrap_asimilarity_search(wrapped, instance, args, kwargs):
         return await wrapped(*args, **kwargs)
 
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
 
     function_name = callable_name(wrapped)
 
@@ -209,6 +211,8 @@ def wrap_similarity_search(wrapped, instance, args, kwargs):
         return wrapped(*args, **kwargs)
 
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
+
     function_name = callable_name(wrapped)
 
     # LLMVectorSearch and Error data
@@ -290,6 +294,8 @@ def wrap_tool_sync_run(wrapped, instance, args, kwargs):
 
     # Framework metric also used for entity tagging in the UI
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
+
     tool_id = str(uuid.uuid4())
 
     run_args = bind_args(wrapped, args, kwargs)
@@ -401,6 +407,7 @@ async def wrap_tool_async_run(wrapped, instance, args, kwargs):
         return await wrapped(*args, **kwargs)
     # Framework metric also used for entity tagging in the UI
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
 
     run_args = bind_args(wrapped, args, kwargs)
 
@@ -543,6 +550,8 @@ async def wrap_chain_async_run(wrapped, instance, args, kwargs):
 
     # Framework metric also used for entity tagging in the UI
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
+
     run_args = bind_args(wrapped, args, kwargs)
     span_id = None
     trace_id = None
@@ -588,6 +597,8 @@ def wrap_chain_sync_run(wrapped, instance, args, kwargs):
 
     # Framework metric also used for entity tagging in the UI
     transaction.add_ml_model_info("Langchain", LANGCHAIN_VERSION)
+    transaction._add_agent_attribute("llm", True)
+
     run_args = bind_args(wrapped, args, kwargs)
     span_id = None
     trace_id = None

tests/mlmodel_langchain/test_agent.py

@@ -18,7 +18,7 @@
 from langchain.prompts import ChatPromptTemplate
 from langchain.tools import tool
 from langchain_core.prompts import MessagesPlaceholder
-from testing_support.fixtures import reset_core_stats_engine
+from testing_support.fixtures import reset_core_stats_engine, validate_attributes
 from testing_support.validators.validate_transaction_metrics import (
     validate_transaction_metrics,
 )
@@ -70,6 +70,7 @@ def prompt():
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_sync_agent(chat_openai_client, tools, prompt):
     agent = create_openai_functions_agent(chat_openai_client, tools, prompt)
@@ -89,6 +90,7 @@ def test_sync_agent(chat_openai_client, tools, prompt):
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_async_agent(loop, chat_openai_client, tools, prompt):
     agent = create_openai_functions_agent(chat_openai_client, tools, prompt)

tests/mlmodel_langchain/test_chain.py

@@ -26,6 +26,7 @@
 from mock import patch
 from testing_support.fixtures import (
     reset_core_stats_engine,
+    validate_attributes,
     validate_custom_event_count,
     validate_transaction_error_event_count,
 )
@@ -1136,6 +1137,7 @@ def test_langchain_chain(
         ],
         background_task=True,
     )
+    @validate_attributes("agent", ["llm"])
     @background_task()
     def _test():
         set_trace_info()
@@ -1597,6 +1599,7 @@ def test_async_langchain_chain(
         ],
         background_task=True,
     )
+    @validate_attributes("agent", ["llm"])
     @background_task()
     def _test():
         set_trace_info()

tests/mlmodel_langchain/test_tool.py

@@ -23,6 +23,7 @@
 from testing_support.fixtures import (  # override_application_settings,
     override_application_settings,
     reset_core_stats_engine,
+    validate_attributes,
     validate_custom_event_count,
     validate_transaction_error_event_count,
 )
@@ -95,6 +96,7 @@ def _multi_arg_tool(first_num: int, second_num: int):
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_langchain_single_arg_tool(set_trace_info, single_arg_tool):
     set_trace_info()
@@ -113,6 +115,7 @@ def test_langchain_single_arg_tool(set_trace_info, single_arg_tool):
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_langchain_single_arg_tool_async(set_trace_info, single_arg_tool, loop):
     set_trace_info()

tests/mlmodel_langchain/test_vectorstore.py

@@ -19,6 +19,7 @@
 from langchain_community.vectorstores.faiss import FAISS
 from testing_support.fixtures import (
     reset_core_stats_engine,
+    validate_attributes,
     validate_custom_event_count,
 )
 from testing_support.validators.validate_custom_events import validate_custom_events
@@ -121,6 +122,7 @@ def test_vectorstore_modules_instrumented():
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_pdf_pagesplitter_vectorstore_in_txn(set_trace_info, embedding_openai_client):
     set_trace_info()
@@ -159,6 +161,7 @@ def test_pdf_pagesplitter_vectorstore_outside_txn(set_trace_info, embedding_open
     ],
     background_task=True,
 )
+@validate_attributes("agent", ["llm"])
 @background_task()
 def test_async_pdf_pagesplitter_vectorstore_in_txn(loop, set_trace_info, embedding_openai_client):
     async def _test():