import streamlit as st import os import nest_asyncio import re from pathlib import Path import typing as t import base64 from mimetypes import guess_type from llama_parse import LlamaParse from llama_index.core.schema import TextNode from llama_index.core import VectorStoreIndex, StorageContext, load_index_from_storage, Settings from llama_index.embeddings.openai import OpenAIEmbedding from llama_index.llms.openai import OpenAI from llama_index.core.query_engine import CustomQueryEngine from llama_index.multi_modal_llms.openai import OpenAIMultiModal from llama_index.core.prompts import PromptTemplate from llama_index.core.schema import ImageNode from llama_index.core.base.response.schema import Response from typing import Any, List, Optional, Tuple from llama_index.core.postprocessor.types import BaseNodePostprocessor from llama_index.core.query_engine import CustomQueryEngine from llama_index.core.retrievers import BaseRetriever from llama_index.multi_modal_llms.openai import OpenAIMultiModal from llama_index.core.schema import NodeWithScore, MetadataMode, QueryBundle from llama_index.core.base.response.schema import Response from llama_index.core.prompts import PromptTemplate from llama_index.core.schema import ImageNode from langchain_core.prompts import ChatPromptTemplate from langchain_core.prompts import HumanMessagePromptTemplate from langchain_core.messages import SystemMessage from llama_index.embeddings.huggingface import HuggingFaceEmbedding nest_asyncio.apply() # Setting API keys os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY') os.environ["LLAMA_CLOUD_API_KEY"] = os.getenv('LLAMA_CLOUD_API_KEY') # Initialize Streamlit app st.title("Medical Knowledge Base & Query System") st.sidebar.title("Settings") # User input for file upload st.sidebar.subheader("Upload Knowledge Base") uploaded_file = st.sidebar.file_uploader("Upload a medical text book (pdf)", type=["jpg", "png", "pdf"]) # # Ensure the 'files' directory exists # if not os.path.exists("files"): # os.makedirs("files") # Initialize the parser parser = LlamaParse( result_type="markdown", parsing_instruction="You are given a medical textbook on medicine", use_vendor_multimodal_model=True, vendor_multimodal_model_name="gpt-4o-mini-2024-07-18", show_progress=True, verbose=True, invalidate_cache=True, do_not_cache=True, num_workers=8, language="en" ) # Initialize md_json_objs as an empty list md_json_objs = [] # Upload and process file if uploaded_file: st.sidebar.write("Processing file...") file_path = f"{uploaded_file.name}" with open(file_path, "wb") as f: f.write(uploaded_file.read()) # Parse the uploaded image md_json_objs = parser.get_json_result([file_path]) image_dicts = parser.get_images(md_json_objs, download_path="data_images") # Extract and display parsed information st.write("File successfully processed!") st.write(f"Processed file: {uploaded_file.name}") # Function to encode image to data URL def local_image_to_data_url(image_path): mime_type, _ = guess_type(image_path) if mime_type is None: mime_type = 'image/png' with open(image_path, "rb") as image_file: base64_encoded_data = base64.b64encode(image_file.read()).decode('utf-8') return f"data:{mime_type};base64,{base64_encoded_data}" # Function to get sorted image files def get_page_number(file_name): match = re.search(r"-page-(\d+)\.jpg$", str(file_name)) if match: return int(match.group(1)) return 0 def _get_sorted_image_files(image_dir): raw_files = [f for f in list(Path(image_dir).iterdir()) if f.is_file()] sorted_files = sorted(raw_files, key=get_page_number) return sorted_files def get_text_nodes(md_json_objs, image_dir) -> t.List[TextNode]: nodes = [] for result in md_json_objs: json_dicts = result["pages"] document_name = result["file_path"].split('/')[-1] docs = [doc["md"] for doc in json_dicts] image_files = _get_sorted_image_files(image_dir) for idx, doc in enumerate(docs): node = TextNode( text=doc, metadata={"image_path": str(image_files[idx]), "page_num": idx + 1, "document_name": document_name}, ) nodes.append(node) return nodes # Load text nodes if md_json_objs is not empty if md_json_objs: text_nodes = get_text_nodes(md_json_objs, "data_images") else: text_nodes = [] # Setup index and LLM embed_model = HuggingFaceEmbedding(model_name="neuml/pubmedbert-base-embeddings") llm = OpenAI(model="gpt-4o-mini-2024-07-18", temperature=0.1) Settings.llm = llm Settings.embed_model = embed_model if not os.path.exists("storage_manuals"): index = VectorStoreIndex(text_nodes, embed_model=embed_model) index.storage_context.persist(persist_dir="./storage_manuals") else: ctx = StorageContext.from_defaults(persist_dir="./storage_manuals") index = load_index_from_storage(ctx) retriever = index.as_retriever() # Query input st.subheader("Ask a Question") query_text = st.text_input("Enter your query:") uploaded_query_image = st.file_uploader("Upload a query image (if any):", type=["jpg", "png"]) # Encode query image if provided encoded_image_url = None if uploaded_query_image: query_image_path = f"{uploaded_query_image.name}" with open(query_image_path, "wb") as img_file: img_file.write(uploaded_query_image.read()) encoded_image_url = local_image_to_data_url(query_image_path) # Setup query engine # QA_PROMPT_TMPL = """ # You are a friendly medical chatbot designed to assist users by providing accurate and detailed responses to medical questions based on information from medical books. # ### Context: # --------------------- # {context_str} # --------------------- # ### Query Text: # {query_str} # ### Query Image: # --------------------- # {encoded_image_url} # --------------------- # ### Answer: # """ QA_PROMPT_TMPL="""You are a friendly medical chatbot designed to assist users by providing accurate and detailed responses to medical questions based on information from medical books. In this task, you will receive parsed text from books in two formats: **Markdown mode** and **Raw text mode**. Markdown mode converts relevant diagrams into tables for clarity, while raw text mode preserves the original layout of the content. ### Key Guidelines: - **Prioritize Image Information**: Always analyze the image provided first for relevant details. Use the text or markdown information only if the image does not contain the necessary information. - **No Image Links**: Your responses should contain only text explanations. Do not include links to images or other resources. - **Contextual Answers**: Your answers should strictly rely on the provided context information. If the information to answer the query is not present, respond with "I don't know," and provide the page number and document name where similar information can be found. ### Context: --------------------- {context_str} --------------------- ### Query Text: {query_str} ### Query Image: --------------------- {encoded_image_url} --------------------- ### Answer: """ QA_PROMPT = PromptTemplate(QA_PROMPT_TMPL) gpt_4o_mm = OpenAIMultiModal(model="gpt-4o-mini-2024-07-18", temperature=0.1) # 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 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) 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_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": nodes, "image_nodes": image_nodes}, ) query_engine = MultimodalQueryEngine(QA_PROMPT, retriever, gpt_4o_mm) # Handle query if query_text: st.write("Querying...") response = query_engine.custom_query(query_text) st.markdown(response.response)