langchain[patch]: Add async methods to VectorstoreIndexCreator (#19582)

pull/19612/head^2
Christophe Bornet 3 months ago committed by GitHub
parent 241774012a
commit a7274f006e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

@ -43,7 +43,22 @@ class VectorStoreIndexWrapper(BaseModel):
chain = RetrievalQA.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return chain.run(question)
return chain.invoke({chain.input_key: question})[chain.output_key]
async def aquery(
self,
question: str,
llm: Optional[BaseLanguageModel] = None,
retriever_kwargs: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> str:
"""Query the vectorstore."""
llm = llm or OpenAI(temperature=0)
retriever_kwargs = retriever_kwargs or {}
chain = RetrievalQA.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return (await chain.ainvoke({chain.input_key: question}))[chain.output_key]
def query_with_sources(
self,
@ -58,7 +73,22 @@ class VectorStoreIndexWrapper(BaseModel):
chain = RetrievalQAWithSourcesChain.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return chain({chain.question_key: question})
return chain.invoke({chain.question_key: question})
async def aquery_with_sources(
self,
question: str,
llm: Optional[BaseLanguageModel] = None,
retriever_kwargs: Optional[Dict[str, Any]] = None,
**kwargs: Any,
) -> dict:
"""Query the vectorstore and get back sources."""
llm = llm or OpenAI(temperature=0)
retriever_kwargs = retriever_kwargs or {}
chain = RetrievalQAWithSourcesChain.from_chain_type(
llm, retriever=self.vectorstore.as_retriever(**retriever_kwargs), **kwargs
)
return await chain.ainvoke({chain.question_key: question})
class VectorstoreIndexCreator(BaseModel):
@ -82,6 +112,14 @@ class VectorstoreIndexCreator(BaseModel):
docs.extend(loader.load())
return self.from_documents(docs)
async def afrom_loaders(self, loaders: List[BaseLoader]) -> VectorStoreIndexWrapper:
"""Create a vectorstore index from loaders."""
docs = []
for loader in loaders:
async for doc in loader.alazy_load():
docs.append(doc)
return await self.afrom_documents(docs)
def from_documents(self, documents: List[Document]) -> VectorStoreIndexWrapper:
"""Create a vectorstore index from documents."""
sub_docs = self.text_splitter.split_documents(documents)
@ -89,3 +127,13 @@ class VectorstoreIndexCreator(BaseModel):
sub_docs, self.embedding, **self.vectorstore_kwargs
)
return VectorStoreIndexWrapper(vectorstore=vectorstore)
async def afrom_documents(
self, documents: List[Document]
) -> VectorStoreIndexWrapper:
"""Create a vectorstore index from documents."""
sub_docs = self.text_splitter.split_documents(documents)
vectorstore = await self.vectorstore_cls.afrom_documents(
sub_docs, self.embedding, **self.vectorstore_kwargs
)
return VectorStoreIndexWrapper(vectorstore=vectorstore)

Loading…
Cancel
Save