Update rag_llamaindex.py

rag_llamaindex.py

@@ -5,6 +5,7 @@ from llama_index import download_loader, PromptTemplate, ServiceContext
 from llama_index.embeddings import OpenAIEmbedding
 from llama_index.indices.vector_store.base import VectorStoreIndex
 from llama_index.llms import OpenAI
+from llama_index.prompts import LangchainPromptTemplate
 from llama_index.storage.storage_context import StorageContext
 from llama_index.vector_stores.mongodb import MongoDBAtlasVectorSearch
 
@@ -14,7 +15,7 @@ from rag_base import BaseRAG
 
 class LlamaIndexRAG(BaseRAG):
     MONGODB_DB_NAME = "llamaindex_db"
-    
+
     def load_documents(self):
         docs = []
     
@@ -102,5 +103,17 @@ class LlamaIndexRAG(BaseRAG):
             service_context = self.get_service_context(config),
             similarity_top_k = config["k"]
         )
- 
+
+        ###
+        prompt_template = LangchainPromptTemplate(
+            template = os.environ["RAG_TEMPLATE"],
+            template_var_mappings = {"context": "context", 
+                                     "question": "question"},
+        )
+        
+        query_engine.update_prompts(
+            {"response_synthesizer:text_qa_template": prompt_template}
+        )
+        ###
+        
         return query_engine.query(prompt)