Merge branch 'master' into master

ZenGuard-AI · Jun 19, 2024 · 24a29df · 24a29df
2 parents 44bcc46 + db6f46c
commit 24a29df
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/libs/langchain/langchain/retrievers/multi_query.py b/libs/langchain/langchain/retrievers/multi_query.py
@@ -9,6 +9,7 @@
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.output_parsers import BaseOutputParser
+from langchain_core.prompts import BasePromptTemplate
 from langchain_core.prompts.prompt import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
 from langchain_core.runnables import Runnable
@@ -62,7 +63,7 @@ def from_llm(
         cls,
         retriever: BaseRetriever,
         llm: BaseLanguageModel,
-        prompt: PromptTemplate = DEFAULT_QUERY_PROMPT,
+        prompt: BasePromptTemplate = DEFAULT_QUERY_PROMPT,
         parser_key: Optional[str] = None,
         include_original: bool = False,
     ) -> "MultiQueryRetriever":

diff --git a/libs/langchain/langchain/retrievers/re_phraser.py b/libs/langchain/langchain/retrievers/re_phraser.py
@@ -7,6 +7,7 @@
 )
 from langchain_core.documents import Document
 from langchain_core.language_models import BaseLLM
+from langchain_core.prompts import BasePromptTemplate
 from langchain_core.prompts.prompt import PromptTemplate
 from langchain_core.retrievers import BaseRetriever
 
@@ -36,7 +37,7 @@ def from_llm(
         cls,
         retriever: BaseRetriever,
         llm: BaseLLM,
-        prompt: PromptTemplate = DEFAULT_QUERY_PROMPT,
+        prompt: BasePromptTemplate = DEFAULT_QUERY_PROMPT,
     ) -> "RePhraseQueryRetriever":
         """Initialize from llm using default template.