RuntimeError: Tool context error detected. This can occur due to parallelization.

This is an issue opened in langchain but I think the problem originates here, so I am opening it here as I am obtaining the same error: https://github.com/langchain-ai/langchain/issues/20929

I copy the problem as my setup is more complex (VertexAIEmbeddings inside a langchain chain inside an FastAPI). And would be more dificult to explain.

In my setup when there are concurrent calls i get the paralellization error but it gives no other option, I cannot parallelize requests and the error persists. The error is raised in the context manager.

Setup:

OS type and version:
Python version: 3.10
pip version: 24.0
google-cloud-aiplatform version: 1.52

Steps to reproduce

Code example

retrieve_chain = self.cat_query_prompt | self.retriever_llm | StrOutputParser()

json_s = retrieve_chain.invoke({"question": query,"schema":self.cat_schema_json, "brand_list":list(df['brand_name'].unique()), "few_shot":self.cat_schema_fewshot})

Stack trace

File "/workspace/selfquery_utils.py", line 151, in get_categorical_filters
json_s = retrieve_chain.invoke({"question": query,"schema":self.cat_schema_json, "brand_list":list(df['brand_name'].unique()), "few_shot":self.cat_schema_fewshot})
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/runnables/base.py", line 2499, in invoke
input = step.invoke(
^^^^^^^^^^^^
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/language_models/llms.py", line 276, in invoke
self.generate_prompt(
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/language_models/llms.py", line 633, in generate_prompt
return self.generate(prompt_strings, stop=stop, callbacks=callbacks, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/language_models/llms.py", line 803, in generate
output = self._generate_helper(
^^^^^^^^^^^^^^^^^^^^^^
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/language_models/llms.py", line 670, in _generate_helper
raise e
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_core/language_models/llms.py", line 657, in _generate_helper
self._generate(
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_google_vertexai/llms.py", line 223, in _generate
res = _completion_with_retry(
^^^^^^^^^^^^^^^^^^^^^^^
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/langchain_google_vertexai/llms.py", line 72, in _completion_with_retry
with telemetry.tool_context_manager(llm._user_agent):
File "/layers/google.python.runtime/python/lib/python3.11/contextlib.py", line 144, in exit
next(self.gen)
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/google/cloud/aiplatform/telemetry.py", line 48, in tool_context_manager
_pop_tool_name(tool_name)
File "/layers/google.python.pip/pip/lib/python3.11/site-packages/google/cloud/aiplatform/telemetry.py", line 57, in _pop_tool_name
raise RuntimeError(
RuntimeError: Tool context error detected. This can occur due to parallelization.

In my setup the error is similar but I add the trace in case it is valuable:

  File "/usr/local/lib/python3.10/site-packages/langchain_core/retrievers.py", line 316, in get_relevant_documents
    result = self._get_relevant_documents(
  File "/usr/local/lib/python3.10/site-packages/langchain/retrievers/contextual_compression.py", line 48, in _get_relevant_documents
    compressed_docs = self.base_compressor.compress_documents(
  File "/usr/local/lib/python3.10/site-packages/langchain/retrievers/document_compressors/base.py", line 46, in compress_documents
    documents = _transformer.transform_documents(documents)
  File "/usr/local/lib/python3.10/site-packages/langchain_community/document_transformers/embeddings_redundant_filter.py", line 151, in transform_documents
    embedded_documents = _get_embeddings_from_stateful_docs(
  File "/usr/local/lib/python3.10/site-packages/langchain_community/document_transformers/embeddings_redundant_filter.py", line 70, in _get_embeddings_from_stateful_docs
    embedded_documents = embeddings.embed_documents(
  File "/usr/local/lib/python3.10/site-packages/langchain_google_vertexai/embeddings.py", line 380, in embed_documents
    return self.embed(texts, batch_size, "RETRIEVAL_DOCUMENT")
  File "/usr/local/lib/python3.10/site-packages/langchain_google_vertexai/embeddings.py", line 363, in embed
    embeddings.extend(t.result())
  File "/usr/local/lib/python3.10/concurrent/futures/_base.py", line 451, in result
    return self.__get_result()
  File "/usr/local/lib/python3.10/concurrent/futures/_base.py", line 403, in __get_result
    raise self._exception
  File "/usr/local/lib/python3.10/concurrent/futures/thread.py", line 58, in run
    result = self.fn(*self.args, **self.kwargs)
  File "/usr/local/lib/python3.10/site-packages/langchain_google_vertexai/embeddings.py", line 194, in _get_embeddings_with_retry
    with telemetry.tool_context_manager(self._user_agent):
  File "/usr/local/lib/python3.10/contextlib.py", line 142, in __exit__
    next(self.gen)
  File "/usr/local/lib/python3.10/site-packages/google/cloud/aiplatform/telemetry.py", line 48, in tool_context_manager
    _pop_tool_name(tool_name)
  File "/usr/local/lib/python3.10/site-packages/google/cloud/aiplatform/telemetry.py", line 57, in _pop_tool_name
    raise RuntimeError(
RuntimeError: Tool context error detected. This can occur due to parallelization.

If you need more information in order to resolve the issue please let me know. Thanks!

googleapis / python-aiplatform