Propagate context vars in all classes/methods (#15329)

- Any direct usage of ThreadPoolExecutor or asyncio.run_in_executor
needs manual handling of context vars

<!-- Thank you for contributing to LangChain!

Please title your PR "<package>: <description>", where <package> is
whichever of langchain, community, core, experimental, etc. is being
modified.

Replace this entire comment with:
  - **Description:** a description of the change, 
  - **Issue:** the issue # it fixes if applicable,
  - **Dependencies:** any dependencies required for this change,
- **Twitter handle:** we announce bigger features on Twitter. If your PR
gets announced, and you'd like a mention, we'll gladly shout you out!

Please make sure your PR is passing linting and testing before
submitting. Run `make format`, `make lint` and `make test` from the root
of the package you've modified to check this locally.

See contribution guidelines for more information on how to write/run
tests, lint, etc: https://python.langchain.com/docs/contributing/

If you're adding a new integration, please include:
1. a test for the integration, preferably unit tests that do not rely on
network access,
2. an example notebook showing its use. It lives in
`docs/docs/integrations` directory.

If no one reviews your PR within a few days, please @-mention one of
@baskaryan, @eyurtsev, @hwchase17.
 -->

libs/community/langchain_community/chat_models/human.py

@@ -1,12 +1,9 @@
 """ChatModel wrapper which returns user input as the response.."""
-import asyncio
-from functools import partial
 from io import StringIO
 from typing import Any, Callable, Dict, List, Mapping, Optional
 
 import yaml
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models.chat_models import BaseChatModel
@@ -111,15 +108,3 @@ def _generate(
         self.message_func(messages, **self.message_kwargs)
         user_input = self.input_func(messages, stop=stop, **self.input_kwargs)
         return ChatResult(generations=[ChatGeneration(message=user_input)])
-
-    async def _agenerate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        func = partial(
-            self._generate, messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return await asyncio.get_event_loop().run_in_executor(None, func)

libs/community/langchain_community/chat_models/mlflow.py

@@ -1,11 +1,8 @@
-import asyncio
 import logging
-from functools import partial
 from typing import Any, Dict, List, Mapping, Optional
 from urllib.parse import urlparse
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models import BaseChatModel
@@ -125,18 +122,6 @@ def _generate(
         resp = self._client.predict(endpoint=self.endpoint, inputs=data)
         return ChatMlflow._create_chat_result(resp)
 
-    async def _agenerate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        func = partial(
-            self._generate, messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return await asyncio.get_event_loop().run_in_executor(None, func)
-
     @property
     def _identifying_params(self) -> Dict[str, Any]:
         return self._default_params

libs/community/langchain_community/chat_models/mlflow_ai_gateway.py

@@ -1,11 +1,8 @@
-import asyncio
 import logging
 import warnings
-from functools import partial
 from typing import Any, Dict, List, Mapping, Optional
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForLLMRun,
     CallbackManagerForLLMRun,
 )
 from langchain_core.language_models.chat_models import BaseChatModel
@@ -116,18 +113,6 @@ def _generate(
         resp = mlflow.gateway.query(self.route, data=data)
         return ChatMLflowAIGateway._create_chat_result(resp)
 
-    async def _agenerate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        func = partial(
-            self._generate, messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return await asyncio.get_event_loop().run_in_executor(None, func)
-
     @property
     def _identifying_params(self) -> Dict[str, Any]:
         return self._default_params

libs/community/langchain_community/chat_models/pai_eas_endpoint.py

@@ -1,7 +1,5 @@
-import asyncio
 import json
 import logging
-from functools import partial
 from typing import Any, AsyncIterator, Dict, List, Optional, cast
 
 import requests
@@ -300,25 +298,3 @@ async def _astream(
                 # break if stop sequence found
                 if stop_seq_found:
                     break
-
-    async def _agenerate(
-        self,
-        messages: List[BaseMessage],
-        stop: Optional[List[str]] = None,
-        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
-        stream: Optional[bool] = None,
-        **kwargs: Any,
-    ) -> ChatResult:
-        if stream if stream is not None else self.streaming:
-            generation: Optional[ChatGenerationChunk] = None
-            async for chunk in self._astream(
-                messages=messages, stop=stop, run_manager=run_manager, **kwargs
-            ):
-                generation = chunk
-            assert generation is not None
-            return ChatResult(generations=[generation])
-
-        func = partial(
-            self._generate, messages, stop=stop, run_manager=run_manager, **kwargs
-        )
-        return await asyncio.get_event_loop().run_in_executor(None, func)

libs/community/langchain_community/embeddings/bedrock.py

@@ -1,11 +1,11 @@
 import asyncio
 import json
 import os
-from functools import partial
 from typing import Any, Dict, List, Optional
 
 from langchain_core.embeddings import Embeddings
 from langchain_core.pydantic_v1 import BaseModel, Extra, root_validator
+from langchain_core.runnables.config import run_in_executor
 
 
 class BedrockEmbeddings(BaseModel, Embeddings):
@@ -181,9 +181,7 @@ async def aembed_query(self, text: str) -> List[float]:
             Embeddings for the text.
         """
 
-        return await asyncio.get_running_loop().run_in_executor(
-            None, partial(self.embed_query, text)
-        )
+        return await run_in_executor(None, self.embed_query, text)
 
     async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
         """Asynchronous compute doc embeddings using a Bedrock model.

libs/community/langchain_community/embeddings/ernie.py

@@ -1,12 +1,12 @@
 import asyncio
 import logging
 import threading
-from functools import partial
 from typing import Dict, List, Optional
 
 import requests
 from langchain_core.embeddings import Embeddings
 from langchain_core.pydantic_v1 import BaseModel, root_validator
+from langchain_core.runnables.config import run_in_executor
 from langchain_core.utils import get_from_dict_or_env
 
 logger = logging.getLogger(__name__)
@@ -134,9 +134,7 @@ async def aembed_query(self, text: str) -> List[float]:
             List[float]: Embeddings for the text.
         """
 
-        return await asyncio.get_running_loop().run_in_executor(
-            None, partial(self.embed_query, text)
-        )
+        return await run_in_executor(None, self.embed_query, text)
 
     async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
         """Asynchronous Embed search docs.

libs/community/langchain_community/tools/multion/close_session.py

@@ -1,8 +1,6 @@
-import asyncio
 from typing import TYPE_CHECKING, Optional, Type
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
 from langchain_core.pydantic_v1 import BaseModel, Field
@@ -57,11 +55,3 @@ def _run(
                 print(f"{e}, retrying...")
         except Exception as e:
             raise Exception(f"An error occurred: {e}")
-
-    async def _arun(
-        self,
-        sessionId: str,
-        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
-    ) -> None:
-        loop = asyncio.get_running_loop()
-        await loop.run_in_executor(None, self._run, sessionId)

libs/community/langchain_community/tools/multion/create_session.py

@@ -1,8 +1,6 @@
-import asyncio
 from typing import TYPE_CHECKING, Optional, Type
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
 from langchain_core.pydantic_v1 import BaseModel, Field
@@ -67,14 +65,3 @@ def _run(
             }
         except Exception as e:
             raise Exception(f"An error occurred: {e}")
-
-    async def _arun(
-        self,
-        query: str,
-        url: Optional[str] = "https://www.google.com/",
-        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
-    ) -> dict:
-        loop = asyncio.get_running_loop()
-        result = await loop.run_in_executor(None, self._run, query, url)
-
-        return result

libs/community/langchain_community/tools/multion/update_session.py

@@ -1,8 +1,6 @@
-import asyncio
 from typing import TYPE_CHECKING, Optional, Type
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
 from langchain_core.pydantic_v1 import BaseModel, Field
@@ -74,15 +72,3 @@ def _run(
                 return {"error": f"{e}", "Response": "retrying..."}
         except Exception as e:
             raise Exception(f"An error occurred: {e}")
-
-    async def _arun(
-        self,
-        sessionId: str,
-        query: str,
-        url: Optional[str] = "https://www.google.com/",
-        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
-    ) -> dict:
-        loop = asyncio.get_running_loop()
-        result = await loop.run_in_executor(None, self._run, sessionId, query, url)
-
-        return result

libs/community/langchain_community/tools/shell/tool.py

@@ -1,10 +1,8 @@
-import asyncio
 import platform
 import warnings
 from typing import Any, List, Optional, Type, Union
 
 from langchain_core.callbacks import (
-    AsyncCallbackManagerForToolRun,
     CallbackManagerForToolRun,
 )
 from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
@@ -77,13 +75,3 @@ def _run(
     ) -> str:
         """Run commands and return final output."""
         return self.process.run(commands)
-
-    async def _arun(
-        self,
-        commands: Union[str, List[str]],
-        run_manager: Optional[AsyncCallbackManagerForToolRun] = None,
-    ) -> str:
-        """Run commands asynchronously and return final output."""
-        return await asyncio.get_event_loop().run_in_executor(
-            None, self.process.run, commands
-        )

libs/community/langchain_community/vectorstores/faiss.py

@@ -1,13 +1,11 @@
 from __future__ import annotations
 
-import asyncio
 import logging
 import operator
 import os
 import pickle
 import uuid
 import warnings
-from functools import partial
 from pathlib import Path
 from typing import (
     Any,
@@ -24,6 +22,7 @@
 import numpy as np
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.runnables.config import run_in_executor
 from langchain_core.vectorstores import VectorStore
 
 from langchain_community.docstore.base import AddableMixin, Docstore
@@ -359,15 +358,15 @@ async def asimilarity_search_with_score_by_vector(
         """
 
         # This is a temporary workaround to make the similarity search asynchronous.
-        func = partial(
+        return await run_in_executor(
+            None,
             self.similarity_search_with_score_by_vector,
             embedding,
             k=k,
             filter=filter,
             fetch_k=fetch_k,
             **kwargs,
         )
-        return await asyncio.get_event_loop().run_in_executor(None, func)
 
     def similarity_search_with_score(
         self,
@@ -640,15 +639,15 @@ async def amax_marginal_relevance_search_with_score_by_vector(
                 relevance and score for each.
         """
         # This is a temporary workaround to make the similarity search asynchronous.
-        func = partial(
+        return await run_in_executor(
+            None,
             self.max_marginal_relevance_search_with_score_by_vector,
             embedding,
             k=k,
             fetch_k=fetch_k,
             lambda_mult=lambda_mult,
             filter=filter,
         )
-        return await asyncio.get_event_loop().run_in_executor(None, func)
 
     def max_marginal_relevance_search_by_vector(
         self,

libs/community/langchain_community/vectorstores/pgvector.py

@@ -1,11 +1,9 @@
 from __future__ import annotations
 
-import asyncio
 import contextlib
 import enum
 import logging
 import uuid
-from functools import partial
 from typing import (
     Any,
     Callable,
@@ -31,6 +29,7 @@
 
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.runnables.config import run_in_executor
 from langchain_core.utils import get_from_dict_or_env
 from langchain_core.vectorstores import VectorStore
 
@@ -941,7 +940,8 @@ async def amax_marginal_relevance_search_by_vector(
         # This is a temporary workaround to make the similarity search
         # asynchronous. The proper solution is to make the similarity search
         # asynchronous in the vector store implementations.
-        func = partial(
+        return await run_in_executor(
+            None,
             self.max_marginal_relevance_search_by_vector,
             embedding,
             k=k,
@@ -950,4 +950,3 @@ async def amax_marginal_relevance_search_by_vector(
             filter=filter,
             **kwargs,
         )
-        return await asyncio.get_event_loop().run_in_executor(None, func)

libs/community/langchain_community/vectorstores/qdrant.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import functools
 import uuid
 import warnings
@@ -25,6 +24,7 @@
 import numpy as np
 from langchain_core.documents import Document
 from langchain_core.embeddings import Embeddings
+from langchain_core.runnables.config import run_in_executor
 from langchain_core.vectorstores import VectorStore
 
 from langchain_community.vectorstores.utils import maximal_marginal_relevance
@@ -58,10 +58,9 @@ async def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
             # by removing the first letter from the method name. For example,
             # if the async method is called ``aaad_texts``, the synchronous method
             # will be called ``aad_texts``.
-            sync_method = functools.partial(
-                getattr(self, method.__name__[1:]), *args, **kwargs
+            return await run_in_executor(
+                None, getattr(self, method.__name__[1:]), *args, **kwargs
             )
-            return await asyncio.get_event_loop().run_in_executor(None, sync_method)
 
     return wrapper