Merge pull request #1199 from newrelic/fix-ai-bug

Fixes for langchain & openai

Author: hmstepanek

newrelic/config.py

@@ -2086,6 +2086,12 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.mlmodel_openai",
         "instrument_openai_resources_chat_completions",
     )
+
+    _process_module_definition(
+        "openai.resources.completions",
+        "newrelic.hooks.mlmodel_openai",
+        "instrument_openai_resources_chat_completions",
+    )
     _process_module_definition(
         "openai._base_client",
         "newrelic.hooks.mlmodel_openai",
@@ -2103,6 +2109,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.mlmodel_langchain",
         "instrument_langchain_runnables_chains_base",
     )
+    _process_module_definition(
+        "langchain_core.runnables.config",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_core_runnables_config",
+    )
     _process_module_definition(
         "langchain.chains.base",
         "newrelic.hooks.mlmodel_langchain",
@@ -2133,6 +2144,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.mlmodel_langchain",
         "instrument_langchain_vectorstore_similarity_search",
     )
+    _process_module_definition(
+        "langchain_community.vectorstores.aerospike",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
     _process_module_definition(
         "langchain_community.vectorstores.analyticdb",
         "newrelic.hooks.mlmodel_langchain",
@@ -2148,6 +2164,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.mlmodel_langchain",
         "instrument_langchain_vectorstore_similarity_search",
     )
+    _process_module_definition(
+        "langchain_community.vectorstores.aperturedb",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
     _process_module_definition(
         "langchain_community.vectorstores.astradb",
         "newrelic.hooks.mlmodel_langchain",
@@ -2163,6 +2184,11 @@ def _process_module_builtin_defaults():
         "newrelic.hooks.mlmodel_langchain",
         "instrument_langchain_vectorstore_similarity_search",
     )
+    _process_module_definition(
+        "langchain_community.vectorstores.azure_cosmos_db_no_sql",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
     _process_module_definition(
         "langchain_community.vectorstores.azure_cosmos_db",
         "newrelic.hooks.mlmodel_langchain",
@@ -2349,6 +2375,12 @@ def _process_module_builtin_defaults():
         "instrument_langchain_vectorstore_similarity_search",
     )
 
+    _process_module_definition(
+        "langchain_community.vectorstores.manticore_search",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
+
     _process_module_definition(
         "langchain_community.vectorstores.marqo",
         "newrelic.hooks.mlmodel_langchain",
@@ -2397,6 +2429,12 @@ 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.nucliadb",
         "newrelic.hooks.mlmodel_langchain",
@@ -2625,6 +2663,12 @@ def _process_module_builtin_defaults():
         "instrument_langchain_vectorstore_similarity_search",
     )
 
+    _process_module_definition(
+        "langchain_community.vectorstores.zep_cloud",
+        "newrelic.hooks.mlmodel_langchain",
+        "instrument_langchain_vectorstore_similarity_search",
+    )
+
     _process_module_definition(
         "langchain_community.vectorstores.zep",
         "newrelic.hooks.mlmodel_langchain",

newrelic/hooks/mlmodel_langchain.py

@@ -18,26 +18,29 @@
 import uuid
 
 from newrelic.api.function_trace import FunctionTrace
-from newrelic.api.time_trace import get_trace_linking_metadata
+from newrelic.api.time_trace import current_trace, get_trace_linking_metadata
 from newrelic.api.transaction import current_transaction
 from newrelic.common.object_wrapper import wrap_function_wrapper
 from newrelic.common.package_version_utils import get_package_version
 from newrelic.common.signature import bind_args
 from newrelic.core.config import global_settings
+from newrelic.core.context import context_wrapper
 
 _logger = logging.getLogger(__name__)
 LANGCHAIN_VERSION = get_package_version("langchain")
 EXCEPTION_HANDLING_FAILURE_LOG_MESSAGE = "Exception occurred in langchain instrumentation: While reporting an exception in langchain, another exception occurred. Report this issue to New Relic Support.\n%s"
 RECORD_EVENTS_FAILURE_LOG_MESSAGE = "Exception occurred in langchain instrumentation: Failed to record LLM events. Report this issue to New Relic Support.\n%s"
-
 VECTORSTORE_CLASSES = {
+    "langchain_community.vectorstores.aerospike": "Aerospike",
     "langchain_community.vectorstores.alibabacloud_opensearch": "AlibabaCloudOpenSearch",
     "langchain_community.vectorstores.analyticdb": "AnalyticDB",
     "langchain_community.vectorstores.annoy": "Annoy",
     "langchain_community.vectorstores.apache_doris": "ApacheDoris",
+    "langchain_community.vectorstores.aperturedb": "ApertureDB",
     "langchain_community.vectorstores.astradb": "AstraDB",
     "langchain_community.vectorstores.atlas": "AtlasDB",
     "langchain_community.vectorstores.awadb": "AwaDB",
+    "langchain_community.vectorstores.azure_cosmos_db_no_sql": "AzureCosmosDBNoSqlVectorSearch",
     "langchain_community.vectorstores.azure_cosmos_db": "AzureCosmosDBVectorSearch",
     "langchain_community.vectorstores.azuresearch": "AzureSearch",
     "langchain_community.vectorstores.baiduvectordb": "BaiduVectorDB",
@@ -71,6 +74,7 @@
     "langchain_community.vectorstores.lancedb": "LanceDB",
     "langchain_community.vectorstores.lantern": "Lantern",
     "langchain_community.vectorstores.llm_rails": "LLMRails",
+    "langchain_community.vectorstores.manticore_search": "ManticoreSearch",
     "langchain_community.vectorstores.marqo": "Marqo",
     "langchain_community.vectorstores.matching_engine": "MatchingEngine",
     "langchain_community.vectorstores.meilisearch": "Meilisearch",
@@ -79,7 +83,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.thirdai_neuraldb": ["NeuralDBClientVectorStore", "NeuralDBVectorStore"],
     "langchain_community.vectorstores.nucliadb": "NucliaDB",
     "langchain_community.vectorstores.oraclevs": "OracleVS",
     "langchain_community.vectorstores.opensearch_vector_search": "OpenSearchVectorSearch",
@@ -118,17 +122,35 @@
     "langchain_community.vectorstores.weaviate": "Weaviate",
     "langchain_community.vectorstores.xata": "XataVectorStore",
     "langchain_community.vectorstores.yellowbrick": "Yellowbrick",
+    "langchain_community.vectorstores.zep_cloud": "ZepCloudVectorStore",
     "langchain_community.vectorstores.zep": "ZepVectorStore",
     "langchain_community.vectorstores.docarray.hnsw": "DocArrayHnswSearch",
     "langchain_community.vectorstores.docarray.in_memory": "DocArrayInMemorySearch",
 }
 
 
-def _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata):
+def bind_submit(func, *args, **kwargs):
+    return {"func": func, "args": args, "kwargs": kwargs}
+
+
+def wrap_ContextThreadPoolExecutor_submit(wrapped, instance, args, kwargs):
+    trace = current_trace()
+    if not trace:
+        return wrapped(*args, **kwargs)
+
+    # Use hardened function signature bind so we have safety net catchall of args and kwargs.
+    bound_args = bind_submit(*args, **kwargs)
+    bound_args["func"] = context_wrapper(bound_args["func"], trace=trace, strict=True)
+    return wrapped(bound_args["func"], *bound_args["args"], **bound_args["kwargs"])
+
+
+def _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata, wrapped):
     settings = transaction.settings if transaction.settings is not None else global_settings()
     span_id = linking_metadata.get("span.id")
     trace_id = linking_metadata.get("trace.id")
-    request_query, request_k = bind_similarity_search(*args, **kwargs)
+    bound_args = bind_args(wrapped, args, kwargs)
+    request_query = bound_args["query"]
+    request_k = bound_args["k"]
     llm_metadata_dict = _get_llm_metadata(transaction)
     vectorstore_error_dict = {
         "request.k": request_k,
@@ -169,21 +191,17 @@ async def wrap_asimilarity_search(wrapped, instance, args, kwargs):
     except Exception as exc:
         ft.notice_error(attributes={"vector_store_id": search_id})
         ft.__exit__(*sys.exc_info())
-        _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata)
+        _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata, wrapped)
         raise
     ft.__exit__(None, None, None)
 
     if not response:
         return response
 
-    _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response)
+    _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response, wrapped)
     return response
 
 
-def bind_similarity_search(query, k, *args, **kwargs):
-    return query, k
-
-
 def wrap_similarity_search(wrapped, instance, args, kwargs):
     transaction = current_transaction()
     if not transaction:
@@ -206,20 +224,22 @@ def wrap_similarity_search(wrapped, instance, args, kwargs):
     except Exception as exc:
         ft.notice_error(attributes={"vector_store_id": search_id})
         ft.__exit__(*sys.exc_info())
-        _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata)
+        _create_error_vectorstore_events(transaction, search_id, args, kwargs, linking_metadata, wrapped)
         raise
     ft.__exit__(None, None, None)
 
     if not response:
         return response
 
-    _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response)
+    _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response, wrapped)
     return response
 
 
-def _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response):
+def _record_vector_search_success(transaction, linking_metadata, ft, search_id, args, kwargs, response, wrapped):
     settings = transaction.settings if transaction.settings is not None else global_settings()
-    request_query, request_k = bind_similarity_search(*args, **kwargs)
+    bound_args = bind_args(wrapped, args, kwargs)
+    request_query = bound_args["query"]
+    request_k = bound_args["k"]
     duration = ft.duration * 1000
     response_number_of_documents = len(response)
     llm_metadata_dict = _get_llm_metadata(transaction)
@@ -855,12 +875,20 @@ def instrument_langchain_chains_base(module):
 
 
 def instrument_langchain_vectorstore_similarity_search(module):
-    vector_class = VECTORSTORE_CLASSES.get(module.__name__)
-
-    if vector_class and hasattr(getattr(module, vector_class, ""), "similarity_search"):
-        wrap_function_wrapper(module, "%s.similarity_search" % vector_class, wrap_similarity_search)
-    if vector_class and hasattr(getattr(module, vector_class, ""), "asimilarity_search"):
-        wrap_function_wrapper(module, "%s.asimilarity_search" % vector_class, wrap_asimilarity_search)
+    def _instrument_class(module, vector_class):
+        if hasattr(getattr(module, vector_class, ""), "similarity_search"):
+            wrap_function_wrapper(module, "%s.similarity_search" % vector_class, wrap_similarity_search)
+        if hasattr(getattr(module, vector_class, ""), "asimilarity_search"):
+            wrap_function_wrapper(module, "%s.asimilarity_search" % vector_class, wrap_asimilarity_search)
+
+    vector_classes = VECTORSTORE_CLASSES.get(module.__name__)
+    if vector_classes is None:
+        return
+    if isinstance(vector_classes, list):
+        for vector_class in vector_classes:
+            _instrument_class(module, vector_class)
+    else:
+        _instrument_class(module, vector_classes)
 
 
 def instrument_langchain_core_tools(module):
@@ -879,3 +907,8 @@ def instrument_langchain_callbacks_manager(module):
         wrap_function_wrapper(module, "CallbackManager.on_chain_start", wrap_on_chain_start)
     if hasattr(getattr(module, "AsyncCallbackManager"), "on_chain_start"):
         wrap_function_wrapper(module, "AsyncCallbackManager.on_chain_start", wrap_async_on_chain_start)
+
+
+def instrument_langchain_core_runnables_config(module):
+    if hasattr(module, "ContextThreadPoolExecutor"):
+        wrap_function_wrapper(module, "ContextThreadPoolExecutor.submit", wrap_ContextThreadPoolExecutor_submit)

newrelic/hooks/mlmodel_openai.py

@@ -492,7 +492,9 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa
             finish_reason = None
             choices = response.get("choices") or []
             if choices:
-                output_message_list = [choices[0].get("message")]
+                output_message_list = [
+                    choices[0].get("message") or {"content": choices[0].get("text"), "role": "assistant"}
+                ]
                 finish_reason = choices[0].get("finish_reason")
         else:
             response_model = kwargs.get("response.model")
@@ -507,7 +509,7 @@ def _record_completion_success(transaction, linking_metadata, completion_id, kwa
 
         request_id = response_headers.get("x-request-id")
         organization = response_headers.get("openai-organization") or getattr(response, "organization", None)
-        messages = kwargs.get("messages", None) or []
+        messages = kwargs.get("messages") or [{"content": kwargs.get("prompt"), "role": "user"}]
         input_message_list = list(messages)
         full_chat_completion_summary_dict = {
             "id": completion_id,