added async functionality to chain execution

app.py

@@ -23,6 +23,7 @@ import langchain
 from langchain.cache import InMemoryCache
 from langchain_core.messages.human import HumanMessage
 from langchain.memory import ConversationBufferMemory
+from chainlit import make_async
 
 load_dotenv()
 YOUR_API_KEY = os.environ["PINECONE_API_KEY"]
@@ -119,7 +120,7 @@ async def main(message: cl.Message):
         memory: ConversationBufferMemory = tools['memory']
 
         # using query search for ArXiv documents and index files(on message) 
-        await cl.make_async(search_and_index)(message=message, embedder=embedder, index=index)
+        await cl.make_async(search_and_index)(message=message, quantity=1, embedder=embedder, index=index)
 
         text_field = "source_document"
         index = pinecone.Index(INDEX_NAME)
@@ -139,7 +140,7 @@ async def main(message: cl.Message):
         await sys_message.update()
     
     # run
-    for chunk in retrieval_augmented_qa_chain.stream({"question": f"{message.content}", "chat_history": memory.buffer_as_messages}):
+    async for chunk in retrieval_augmented_qa_chain.astream({"question": f"{message.content}", "chat_history": memory.buffer_as_messages}):
         if res:= chunk.get('response'):
             await sys_message.stream_token(res.content)
     await sys_message.send()

utils/store.py

@@ -48,13 +48,13 @@ def index_documents(
         index.upsert(vectors=zip(ids, embeds, metadatas))
 
 
-def search_and_index(message: cl.Message, embedder: CacheBackedEmbeddings, index: GRPCIndex) -> None:
+def search_and_index(message: cl.Message, quantity: int, embedder: CacheBackedEmbeddings, index: GRPCIndex) -> None:
             
     arxiv_client = arxiv.Client()
     
     search = arxiv.Search(
         query = message.content,
-        max_results = 10,
+        max_results = quantity,
         sort_by = arxiv.SortCriterion.Relevance
     )
     paper_urls = []