add itemgetter to fetch context from the vectorstore

app.py

@@ -12,6 +12,7 @@ from langchain.schema.output_parser import StrOutputParser
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.schema.runnable.config import RunnableConfig
 
+
 # GLOBAL SCOPE - ENTIRE APPLICATION HAS ACCESS TO VALUES SET IN THIS SCOPE #
 # ---- ENV VARIABLES ---- # 
 """
@@ -136,7 +137,7 @@ async def start_chat():
     """
 
     ### BUILD LCEL RAG CHAIN THAT ONLY RETURNS TEXT
-    lcel_rag_chain = rag_prompt | hf_llm
+    lcel_rag_chain = {"context": itemgetter("query") | hf_retriever, "query": itemgetter("query")}| rag_prompt | hf_llm
 
     cl.user_session.set("lcel_rag_chain", lcel_rag_chain)
 
@@ -150,23 +151,11 @@ async def main(message: cl.Message):
     The LCEL RAG chain is stored in the user session, and is unique to each user session - this is why we can access it here.
     """
     lcel_rag_chain = cl.user_session.get("lcel_rag_chain")
-    
-    # Retrieve context for the user's query
-    context = hf_retriever.retrieve(query=message.content)
-
-    # Combine context into a single string
-    context_text = "\n".join([doc.page_content for doc in context])
-
-    # Prepare input for the prompt
-    input_dict = {
-        "query": message.content,
-        "context": context_text
-    }
 
     msg = cl.Message(content="")
 
     async for chunk in lcel_rag_chain.astream(
-        input_dict,
+        {"query": message.content},
         config=RunnableConfig(callbacks=[cl.LangchainCallbackHandler()]),
     ):
         await msg.stream_token(chunk)