app.py

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)