diff --git a/newrelic/config.py b/newrelic/config.py index f7b2de267c..6d46522229 100644 --- a/newrelic/config.py +++ b/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", diff --git a/newrelic/hooks/mlmodel_langchain.py b/newrelic/hooks/mlmodel_langchain.py index e0e4e46b0f..69fb279999 100644 --- a/newrelic/hooks/mlmodel_langchain.py +++ b/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 diff --git a/tests/mlmodel_langchain/test_agent.py b/tests/mlmodel_langchain/test_agent.py index 133e235ed2..3633ba5e8b 100644 --- a/tests/mlmodel_langchain/test_agent.py +++ b/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) diff --git a/tests/mlmodel_langchain/test_chain.py b/tests/mlmodel_langchain/test_chain.py index 6e97f3beca..13b9c43677 100644 --- a/tests/mlmodel_langchain/test_chain.py +++ b/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() diff --git a/tests/mlmodel_langchain/test_tool.py b/tests/mlmodel_langchain/test_tool.py index ff46c195c0..35fc2e0c6d 100644 --- a/tests/mlmodel_langchain/test_tool.py +++ b/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() diff --git a/tests/mlmodel_langchain/test_vectorstore.py b/tests/mlmodel_langchain/test_vectorstore.py index 3e813bf5fc..de6619b70c 100644 --- a/tests/mlmodel_langchain/test_vectorstore.py +++ b/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():