Update app.py

app.py

@@ -161,17 +161,74 @@ QA_PROMPT = PromptTemplate(QA_PROMPT_TMPL)
 gpt_4o_mm = OpenAIMultiModal(model="gpt-4o-mini-2024-07-18")
 
 class MultimodalQueryEngine(CustomQueryEngine):
-    def __init__(self, qa_prompt, retriever, multi_modal_llm, node_postprocessors=[]):
-        super().__init__(qa_prompt=qa_prompt, retriever=retriever, multi_modal_llm=multi_modal_llm, node_postprocessors=node_postprocessors)
+    # def __init__(self, qa_prompt, retriever, multi_modal_llm, node_postprocessors=[]):
+    #     super().__init__(qa_prompt=qa_prompt, retriever=retriever, multi_modal_llm=multi_modal_llm, node_postprocessors=node_postprocessors)
 
-    def custom_query(self, query_str):
+    # def custom_query(self, query_str):
+    #     nodes = self.retriever.retrieve(query_str)
+    #     image_nodes = [NodeWithScore(node=ImageNode(image_path=n.node.metadata["image_path"])) for n in nodes]
+    #     ctx_str = "\n\n".join([r.node.get_content().strip() for r in nodes])
+    #     fmt_prompt = self.qa_prompt.format(context_str=ctx_str, query_str=query_str, encoded_image_url=encoded_image_url)
+    #     llm_response = self.multi_modal_llm.complete(prompt=fmt_prompt, image_documents=[image_node.node for image_node in image_nodes])
+    #     return Response(response=str(llm_response), source_nodes=nodes, metadata={"text_nodes": text_nodes, "image_nodes": image_nodes})
+
+
+
+class MultimodalQueryEngine(CustomQueryEngine):
+    qa_prompt: PromptTemplate
+    retriever: BaseRetriever
+    multi_modal_llm: OpenAIMultiModal
+    node_postprocessors: Optional[List[BaseNodePostprocessor]]
+
+    def __init__(
+        self,
+        qa_prompt: PromptTemplate,
+        retriever: BaseRetriever,
+        multi_modal_llm: OpenAIMultiModal,
+        node_postprocessors: Optional[List[BaseNodePostprocessor]] = [],
+    ):
+        super().__init__(
+            qa_prompt=qa_prompt,
+            retriever=retriever,
+            multi_modal_llm=multi_modal_llm,
+            node_postprocessors=node_postprocessors
+        )
+
+    def custom_query(self, query_str: str):
+        # retrieve most relevant nodes
         nodes = self.retriever.retrieve(query_str)
-        image_nodes = [NodeWithScore(node=ImageNode(image_path=n.node.metadata["image_path"])) for n in nodes]
-        ctx_str = "\n\n".join([r.node.get_content().strip() for r in nodes])
-        fmt_prompt = self.qa_prompt.format(context_str=ctx_str, query_str=query_str, encoded_image_url=encoded_image_url)
-        llm_response = self.multi_modal_llm.complete(prompt=fmt_prompt, image_documents=[image_node.node for image_node in image_nodes])
-        return Response(response=str(llm_response), source_nodes=nodes, metadata={"text_nodes": text_nodes, "image_nodes": image_nodes})
 
+        for postprocessor in self.node_postprocessors:
+            nodes = postprocessor.postprocess_nodes(
+                nodes, query_bundle=QueryBundle(query_str)
+            )
+
+
+        # create image nodes from the image associated with those nodes
+        image_nodes = [
+            NodeWithScore(node=ImageNode(image_path=n.node.metadata["image_path"]))
+            for n in nodes
+        ]
+
+        # create context string from parsed markdown text
+        ctx_str = "\n\n".join(
+            [r.node.get_content(metadata_mode=MetadataMode.LLM).strip() for r in nodes]
+        )
+
+        # prompt for the LLM
+        fmt_prompt = self.qa_prompt.format(context_str=ctx_str, query_str=query_str,encoded_image_url=encoded_image_url)
+
+        # use the multimodal LLM to interpret images and generate a response to the prompt
+        llm_repsonse = self.multi_modal_llm.complete(
+            prompt=fmt_prompt,
+            image_documents=[image_node.node for image_node in image_nodes],
+        )
+        return Response(
+            response=str(llm_repsonse),
+            source_nodes=nodes,
+            metadata={"text_nodes": text_nodes, "image_nodes": image_nodes},
+        )
+        
 query_engine = MultimodalQueryEngine(QA_PROMPT, retriever, gpt_4o_mm)
 
 # Handle query