Upload 3 files

app.py

@@ -1,352 +1,348 @@
-import logging
-import json
-import pandas as pd
 import streamlit as st
-from pinecone import Pinecone
-from llama_index.vector_stores.pinecone import PineconeVectorStore
-from llama_index.core import (
-    StorageContext, VectorStoreIndex, SimpleDirectoryReader, 
-    get_response_synthesizer, Settings
-)
-from llama_index.core.node_parser import SentenceSplitter
-from llama_index.core.retrievers import (
-    VectorIndexRetriever, RouterRetriever
-)
-from llama_index.retrievers.bm25 import BM25Retriever
-from llama_index.core.tools import RetrieverTool
-from llama_index.core.query_engine import (
-    RetrieverQueryEngine, FLAREInstructQueryEngine, MultiStepQueryEngine
-)
-from llama_index.core.indices.query.query_transform import (
-    StepDecomposeQueryTransform
-)
-from llama_index.llms.groq import Groq
+from llama_index.core.memory.chat_memory_buffer import ChatMemoryBuffer
+from retrievers import PARetriever
+from utils_code import  create_chat_engine
 from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.llms.azure_openai import AzureOpenAI
-from llama_index.embeddings.openai import OpenAIEmbedding
-from llama_index.readers.file import PyMuPDFReader
-import traceback
-from oauth2client.service_account import ServiceAccountCredentials
-import gspread
-import uuid
-from dotenv import load_dotenv
+from llama_index.core import Settings
 import os
-from datetime import datetime
-
-# Load environment variables
-load_dotenv()
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-
-# Google Sheets setup
-scope = ["https://spreadsheets.google.com/feeds", "https://www.googleapis.com/auth/drive"]
-creds_dict = {
-    "type": os.getenv("type"),
-    "project_id": os.getenv("project_id"),
-    "private_key_id": os.getenv("private_key_id"),
-    "private_key": os.getenv("private_key"),
-    "client_email": os.getenv("client_email"),
-    "client_id": os.getenv("client_id"),
-    "auth_uri": os.getenv("auth_uri"),
-    "token_uri": os.getenv("token_uri"),
-    "auth_provider_x509_cert_url": os.getenv("auth_provider_x509_cert_url"),
-    "client_x509_cert_url": os.getenv("client_x509_cert_url")
-}
-creds_dict['private_key'] = creds_dict['private_key'].replace('\\n', '\n')
-creds = ServiceAccountCredentials.from_json_keyfile_dict(creds_dict, scope)
-client = gspread.authorize(creds)
-sheet = client.open("RAG").sheet1
-
-# Fixed variables
-AZURE_DEPLOYMENT_NAME = os.getenv("AZURE_DEPLOYMENT_NAME")
-AZURE_API_VERSION = os.getenv("AZURE_API_VERSION")
-AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
-
-# Global variables for lazy loading
-llm = None
-pinecone_index = None
-
-def log_and_exit(message):
-    logging.error(message)
-    raise SystemExit(message)
-
-def initialize_apis(api, model, pinecone_api_key, groq_api_key, azure_api_key):
-    global llm, pinecone_index
-    try:
-        if llm is None:
-            llm = initialize_llm(api, model, groq_api_key, azure_api_key)
-        if pinecone_index is None:
-            pinecone_client = Pinecone(pinecone_api_key)
-            pinecone_index = pinecone_client.Index("ll144")
-        logging.info("Initialized LLM and Pinecone.")
-    except Exception as e:
-        log_and_exit(f"Error initializing APIs: {e}")
-
-def initialize_llm(api, model, groq_api_key, azure_api_key):
-    if api == 'groq':
-        model_mappings = {
-            'mixtral-8x7b': "mixtral-8x7b-32768",
-            'llama3-8b': "llama3-8b-8192",
-            'llama3-70b': "llama3-70b-8192",
-            'gemma-7b': "gemma-7b-it"
-        }
-        return Groq(model=model_mappings[model], api_key=groq_api_key)
-    elif api == 'azure':
-        if model == 'gpt35':
-            return AzureOpenAI(
-                deployment_name=AZURE_DEPLOYMENT_NAME,
-                temperature=0,
-                api_key=azure_api_key,
-                azure_endpoint=AZURE_OPENAI_ENDPOINT,
-                api_version=AZURE_API_VERSION
-            )
+from llama_index.llms.azure_openai import AzureOpenAI
+from dotenv import load_dotenv, find_dotenv
+from retrievers import HyPARetriever, PARetriever
+from llama_index.vector_stores.pinecone import PineconeVectorStore
+from llama_index.core import VectorStoreIndex
+from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
+from llama_index.core import PropertyGraphIndex
+from llama_index.core.vector_stores import MetadataFilter, MetadataFilters, FilterOperator
+from llama_index.retrievers.bm25 import BM25Retriever
 
-def load_pdf_data(chunk_size):
-    reader = PyMuPDFReader()
-    file_extractor = {".pdf": reader}
-    documents = SimpleDirectoryReader(input_files=['LL144.pdf', 'LL144_Definitions.pdf'], file_extractor=file_extractor).load_data()
-    return documents
-
-def create_index(documents, embedding_model_type="HF", embedding_model="BAAI/bge-large-en-v1.5", retriever_method="BM25", chunk_size=512):
-    global llm, pinecone_index
-    try:
-        embed_model = select_embedding_model(embedding_model_type, embedding_model)
-
-        Settings.llm = llm
-        Settings.embed_model = embed_model
-        Settings.chunk_size = chunk_size
-
-        if retriever_method in ["BM25", "BM25+Vector"]:
-            nodes = create_bm25_nodes(documents, chunk_size)
-            logging.info("Created BM25 nodes from documents.")
-            if retriever_method == "BM25+Vector":
-                vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
-                storage_context = StorageContext.from_defaults(vector_store=vector_store)
-                index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
-                logging.info("Created index for BM25+Vector from documents.")
-                return index, nodes
-            return None, nodes
-        else:
-            vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
-            storage_context = StorageContext.from_defaults(vector_store=vector_store)
-            index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
-            logging.info("Created index from documents.")
-            return index, None
-    except Exception as e:
-        log_and_exit(f"Error creating index: {e}")
-
-def select_embedding_model(embedding_model_type, embedding_model):
-    if embedding_model_type == "HF":
-        return HuggingFaceEmbedding(model_name=embedding_model)
-    elif embedding_model_type == "OAI":
-        return OpenAIEmbedding()  # Implement OAI Embedding if needed
-
-def create_bm25_nodes(documents, chunk_size):
-    splitter = SentenceSplitter(chunk_size=chunk_size)
-    nodes = splitter.get_nodes_from_documents(documents)
-    return nodes
-
-def select_retriever(index, nodes, retriever_method, top_k):
-    logging.info(f"Selecting retriever with method: {retriever_method}")
-    if nodes is not None:
-        logging.info(f"Available document IDs: {list(range(len(nodes)))}")
-    else:
-        logging.warning("Nodes are None")
+# Load environment variables from the .env file
+dotenv_path = find_dotenv()
+#print(f"Dotenv Path: {dotenv_path}")
+load_dotenv(dotenv_path)
 
-    if retriever_method == 'BM25':
-        return BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=top_k)
-    elif retriever_method == "BM25+Vector":
-        if index is None:
-            log_and_exit("Index must be initialized when using BM25+Vector retriever method.")
-        
-        bm25_retriever = RetrieverTool.from_defaults(
-            retriever=BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=top_k),
-            description="BM25 Retriever"
-        )
-        
-        vector_retriever = RetrieverTool.from_defaults(
-            retriever=VectorIndexRetriever(index=index),
-            description="Vector Retriever"
-        )
+
+embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5")
+Settings.embed_model = embed_model
+
+# Set Azure OpenAI keys for Giskard if needed
+#os.environ["AZURE_OPENAI_API_KEY"] = os.getenv("GSK_AZURE_OPENAI_API_KEY")
+#os.environ["AZURE_OPENAI_ENDPOINT"] = os.getenv("GSK_AZURE_OPENAI_ENDPOINT")
+os.environ["GSK_LLM_MODEL"] = "gpt-4o-mini"
+
+# Pinecone and Neo4j credentials
+pinecone_api_key = os.getenv("PINECONE_API_KEY")
+ll144_index_name = 'll144'
+euaiact_index_name = 'euaiact'
+
+# Initialize Pinecone
+from pinecone import Pinecone
+pc = Pinecone(api_key=pinecone_api_key)
+
+
+def metadata_filter(corpus_name):
+
+    if corpus_name == "EUAIACT":
+
+        # Filter for 'EUAIACT.pdf'
+        filter = MetadataFilters(filters=[MetadataFilter(key="filepath", value="'EUAIACT.pdf'", operator=FilterOperator.CONTAINS)])
+
+    elif corpus_name == "LL144":
+    # Filter for 'LLL144.pdf' or 'LL144_Definitions.pdf'
+        filter = MetadataFilters(filters=[
+            MetadataFilter(key="filepath", value="'LL144.pdf'", operator=FilterOperator.CONTAINS),
+            MetadataFilter(key="filepath", value="'LL144_Definitions.pdf'", operator=FilterOperator.CONTAINS)
+        ])
+
+    return filter
         
-        router_retriever = RouterRetriever.from_defaults(
-            retriever_tools=[bm25_retriever, vector_retriever],
+
+# Load vector index
+#@st.cache_data(ttl=None, persist=None)
+def load_vector_index(corpus_name):
+    if corpus_name == "LL144":
+        pinecone_index = pc.Index(ll144_index_name)
+    elif corpus_name == "EUAIACT":
+        pinecone_index = pc.Index(euaiact_index_name)
+
+    vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
+    vector_index = VectorStoreIndex.from_vector_store(vector_store)
+    
+    return vector_index
+
+# Load property graph index
+#@st.cache_data(ttl=None, persist=None)
+def load_pg_index():
+    neo4j_username = os.getenv("NEO4J_USERNAME")
+    neo4j_password = os.getenv("NEO4J_PASSWORD")
+    neo4j_url = os.getenv("NEO4J_URI")
+
+    graph_store = Neo4jPropertyGraphStore(username=neo4j_username, password=neo4j_password, url=neo4j_url)
+    pg_index = PropertyGraphIndex.from_existing(property_graph_store=graph_store)
+    return pg_index
+
+# Initialize the retriever (HyPA or PA)
+def init_retriever(retriever_type, corpus_name, use_reranker, use_rewriter, classifier_model):
+    # Check if vector index is cached, if not, load it
+    if "vector_index" not in st.session_state:
+        st.session_state.vector_index = load_vector_index(corpus_name)
+
+    # Check if property graph index is cached, if not, load it
+    if "pg_index" not in st.session_state:
+        st.session_state.pg_index = load_pg_index()
+
+    vector_index = st.session_state.vector_index
+    graph_index = st.session_state.pg_index
+    llm = st.session_state.llm
+
+    filter = metadata_filter(corpus_name=corpus_name)
+    # Set the reranker model if selected
+    reranker_model_name = "BAAI/bge-reranker-large" if use_reranker else None
+
+    # Choose the appropriate retriever based on user selection
+    if retriever_type == "HyPA":
+        retriever = HyPARetriever(
             llm=llm,
-            select_multi=True
+            vector_retriever=vector_index.as_retriever(similarity_top_k=10),
+            bm25_retriever=None,#BM25Retriever.from_defaults(index=vector_index, similarity_top_k=10),
+            kg_index=graph_index,  # Include KG for HyPA
+            rewriter=use_rewriter,  # Set rewriter option
+            classifier_model=classifier_model,  # Use the selected classifier model
+            verbose=False,
+            property_index=True,  # Use property graph index
+            reranker_model_name=reranker_model_name,  # Use reranker if selected
+            pg_filters=filter
         )
-        return router_retriever
-    elif retriever_method == "Vector Search":
-        if index is None:
-            log_and_exit("Index must be initialized when using Vector Search retriever method.")
-        return VectorIndexRetriever(index=index, similarity_top_k=top_k)
     else:
-        log_and_exit(f"Unsupported retriever method: {retriever_method}")
+        retriever = PARetriever(
+            llm=llm,
+            vector_retriever=vector_index.as_retriever(similarity_top_k=10),
+            bm25_retriever=None,#BM25Retriever.from_defaults(index=vector_index, similarity_top_k=10),
+            rewriter=use_rewriter,  # Set rewriter option
+            classifier_model=classifier_model,  # Use the selected classifier model
+            verbose=False,
+            reranker_model_name=reranker_model_name  # Use reranker if selected
+        )
 
-def setup_query_engine(index, response_mode, nodes=None, query_engine_method=None, retriever_method=None, top_k=2):
-    global llm
-    try:
-        logging.info(f"Setting up query engine with retriever_method: {retriever_method} and query_engine_method: {query_engine_method}")
-        retriever = select_retriever(index, nodes, retriever_method, top_k)
-        
-        if retriever is None:
-            log_and_exit("Failed to create retriever. Index or nodes might be None.")
+    memory = ChatMemoryBuffer.from_defaults(token_limit=8192)
+    chat_engine = create_chat_engine(retriever=retriever, memory=memory, llm=llm)
+    st.session_state.chat_engine = chat_engine
+    #return chat_engine
+
+
+def process_query(query):
+    """Processes the input query and displays it along with the response in the main chat area."""
+    # Append the user query to the message history and display it
+    st.session_state.messages.append({"role": "user", "content": query})
+    with st.chat_message("user"):
+        st.write(query)
+
+    # Ensure the chat engine is initialized
+    chat_engine = st.session_state.get('chat_engine', None)
+    if chat_engine:
+        # Process the query through the chat engine
+        with st.chat_message("assistant"):
+            with st.spinner("Retrieving Knowledge..."):
+                response = chat_engine.stream_chat(query)
+                response_str = ""
+                response_container = st.empty()
+                for token in response.response_gen:
+                    response_str += token
+                    response_container.write(response_str)
+                # Append the assistant's response to the message history
+                st.session_state.messages.append({"role": "assistant", "content": response_str})
+                
+            # Expander for additional info
+            with st.expander("Source Nodes"):
+                # Display source nodes
+                if hasattr(response, 'source_nodes') and response.source_nodes:
+                    
+                    for idx, node in enumerate(response.source_nodes):
+                        st.markdown(f"#### Source Node {idx + 1}")
+                        st.write(f"**Node ID:** {node.node_id}")
+                        st.write(f"**Node Score:** {node.score}")
+                        
+                        st.write("**Metadata:**")
+                        for key, value in node.metadata.items():
+                            st.write(f"- **{key}:** {value}")
+                        
+                        st.write("**Content:**")
+                        st.write(node.node.get_content())
+
+                        # Add a horizontal line to separate nodes
+                        st.markdown("---")
+                else:
+                    st.write("No additional source nodes available.")
+
+        st.session_state.messages.append({"role": "assistant", "content": str(response)})
+
+
+
+
+
+# Streamlit App
+def main():
+
+    
+    # Sidebar for retriever options
+    with st.sidebar:
+        st.image('holisticai.svg', use_column_width=True)
+        st.title("Retriever Settings")
+
+        # Azure OpenAI credentials input fields (start with blank fields)
+        azure_api_key = st.text_input("Azure OpenAI API Key", value="", type="password")
+        azure_endpoint = st.text_input("Azure OpenAI Endpoint", value="", type="password")
+
+        llm_model_choice = st.selectbox("Select LLM Model", ["gpt-4o-mini", "gpt35"])
+
+        # Let the user make selections without updating session state yet
+        retriever_type = st.selectbox("Select Retriever Method", ["PA", "HyPA"])
+        corpus_name = st.selectbox("Select Corpus", ["LL144", "EUAIACT"])
+        temperature = st.slider("Set LLM Temperature", min_value=0.0, max_value=1.0, value=0.0, step=0.1)
+
+        # Display a red warning about non-zero temperature
+        if temperature > 0:
+            st.markdown(
+                "<p style='color:red;'>Warning: A non-zero temperature may lead to hallucinations in the generated responses.</p>",
+                unsafe_allow_html=True
+            )
 
-        response_synthesizer = get_response_synthesizer(response_mode=response_mode)
-        index_query_engine = index.as_query_engine(similarity_top_k=top_k) if index else None
+        # Checkboxes for reranker and rewriter options
+        use_reranker = st.checkbox("Use Reranker")
+        use_rewriter = st.checkbox("Use Rewriter")
 
-        if query_engine_method == "FLARE":
-            query_engine = FLAREInstructQueryEngine(
-                query_engine=index_query_engine,
-                max_iterations=4,
-                verbose=False
-            )
-        elif query_engine_method == "MS":
-            query_engine = MultiStepQueryEngine(
-                query_engine=index_query_engine,
-                query_transform=StepDecomposeQueryTransform(llm=llm, verbose=False),
-                index_summary="Used to answer questions about the regulation"
+        # Radio buttons for classifier model
+        classifier_type = st.radio("Select Classifier Type", ["2-Class", "3-Class"])
+        classifier_model = "rk68/distilbert-q-classifier-2" if classifier_type == "2-Class" else "rk68/distilbert-q-classifier-3"
+
+
+
+        # When the user clicks "Initialize", store everything in session state
+        if st.button("Initialize"):
+            st.session_state.retriever_type = retriever_type
+            st.session_state.corpus_name = corpus_name
+            st.session_state.temperature = temperature
+            st.session_state.use_reranker = use_reranker
+            st.session_state.use_rewriter = use_rewriter
+            st.session_state.classifier_type = classifier_type
+            st.session_state.classifier_model = classifier_model
+
+            # Store the user inputs in session state
+            st.session_state.azure_api_key = azure_api_key
+            st.session_state.azure_endpoint = azure_endpoint
+
+            # Set the environment variables from user inputs
+            os.environ["AZURE_OPENAI_API_KEY"] = azure_api_key
+            os.environ["AZURE_OPENAI_ENDPOINT"] = azure_endpoint
+
+            llm = AzureOpenAI(
+                deployment_name=llm_model_choice, temperature=temperature,
+                api_key=azure_api_key, azure_endpoint=azure_endpoint,
+                api_version=os.getenv("AZURE_API_VERSION")
             )
-        else:
-            query_engine = RetrieverQueryEngine(retriever=retriever, response_synthesizer=response_synthesizer)
+            Settings.llm = llm
+            st.session_state.llm = llm
+
+            # Initialize retriever after storing the settings
+            init_retriever(retriever_type, corpus_name, use_reranker, use_rewriter, classifier_model)
+            st.success("Retriever Initialized")
+
+        # Example questions based on selected corpus
+        st.markdown("### Example Queries")
+        # Example questions with unique button handling
+        example_questions = {
+            "LL144": [
+                "What is a bias audit?",
+                "When does it come into effect?",
+                "Summarise Local Law 144"
+            ],
+            "EUAIACT": [
+                "What is an AI system?",
+                "What are the key takeaways?",
+                "Explain the key provisions of EUAIACT."
+            ]
+        }
+
+
+            # Display buttons for the example queries
+        for idx, question in enumerate(example_questions.get(corpus_name, [])):
+            if st.button(f"{question} [{idx}]"):
+                process_query(question)
+
+
+
+
         
-        if query_engine is None:
-            log_and_exit("Failed to create query engine.")
+        # Add a disclaimer at the bottom
+        st.markdown("---")  # Horizontal line for separation
         
-        return query_engine
-    except Exception as e:
-        logging.error(f"Error setting up query engine: {e}")
-        traceback.print_exc()
-        log_and_exit(f"Error setting up query engine: {e}")
-
-def log_to_google_sheets(data):
-    try:
-        sheet.append_row(data)
-        logging.info("Logged data to Google Sheets.")
-    except Exception as e:
-        logging.error(f"Error logging data to Google Sheets: {e}")
-
-def update_google_sheets(question_id, feedback=None, detailed_feedback=None, annotated_answer=None):
-    try:
-        existing_data = sheet.get_all_values()
-        headers = existing_data[0]
-        for i, row in enumerate(existing_data):
-            if row[0] == question_id:
-                if feedback is not None:
-                    sheet.update_cell(i+1, headers.index("Feedback") + 1, feedback)
-                if detailed_feedback is not None:
-                    sheet.update_cell(i+1, headers.index("Detailed Feedback") + 1, detailed_feedback)
-                if annotated_answer is not None:
-                    sheet.update_cell(i+1, headers.index("annotated_answer") + 1, annotated_answer)
-                logging.info("Updated data in Google Sheets.")
-                return
-    except Exception as e:
-        logging.error(f"Error updating data in Google Sheets: {e}")
-
-def run_streamlit_app():
-    if 'query_engine' not in st.session_state:
-        st.session_state.query_engine = None
-
-    st.title("RAG Chat Application")
-
-    col1, col2 = st.columns(2)
-
-    with col1:
-        pinecone_api_key = st.text_input("Pinecone API Key")
-        azure_api_key = st.text_input("Azure API Key")
-        groq_api_key = st.text_input("Groq API Key")
-
-    def update_api_based_on_model():
-        selected_model = st.session_state['selected_model']
-        if selected_model == 'gpt35':
-            st.session_state['selected_api'] = 'azure'
-        else:
-            st.session_state['selected_api'] = 'groq'
-
-    with col2:
-        selected_model = st.selectbox("Select Model", ["llama3-8b", "llama3-70b", "mixtral-8x7b", "gemma-7b", "gpt35"], index=4, key='selected_model', on_change=update_api_based_on_model)
-        selected_api = st.selectbox("Select API", ["azure", "groq"], index=0, key='selected_api', disabled=True)
-        embedding_model_type = "HF"
-        embedding_model = st.selectbox("Select Embedding Model", ["BAAI/bge-large-en-v1.5", "other_model"])
-        retriever_method = st.selectbox("Select Retriever Method", ["Vector Search", "BM25", "BM25+Vector"])
-
-    col3, col4 = st.columns(2)
-    with col3:
-        chunk_size = st.selectbox("Select Chunk Size", [128, 256, 512, 1024], index=2)
-    with col4:
-        top_k = st.selectbox("Select Top K", [1, 2, 3, 5, 6], index=1)
-
-    if st.button("Initialize"):
-        initialize_apis(st.session_state['selected_api'], selected_model, pinecone_api_key, groq_api_key, azure_api_key)
-        documents = load_pdf_data(chunk_size)
-        index, nodes = create_index(documents, embedding_model_type=embedding_model_type, embedding_model=embedding_model, retriever_method=retriever_method, chunk_size=chunk_size)
-        st.session_state.query_engine = setup_query_engine(index, response_mode="compact", nodes=nodes, query_engine_method=None, retriever_method=retriever_method, top_k=top_k)
-        st.success("Initialization complete.")
-
-    if 'chat_history' not in st.session_state:
-        st.session_state.chat_history = []
-
-    for chat_index, chat in enumerate(st.session_state.chat_history):
-        with st.chat_message("user"):
-            st.markdown(chat['user'])
-        with st.chat_message("bot"):
-            st.markdown("### Retrieved Contexts")
-            for node in chat.get('contexts', []):
-                st.markdown(
-                    f"<div style='border:1px solid #ccc; padding:10px; margin:10px 0; font-size:small;'>{node.text}</div>",
-                    unsafe_allow_html=True
-                )
-            st.markdown("### Answer")
-            st.markdown(chat['response'])
-
-            col1, col2 = st.columns([1, 1])
-            with col1:
-                if st.button("Annotate 👎", key=f"annotate_{chat_index}"):
-                    chat['annotate'] = True
-                    chat['feedback'] = -1
-                    st.session_state.chat_history[chat_index] = chat
-                    update_google_sheets(chat['id'], feedback=-1)
-                    st.rerun()
-            with col2:
-                if st.button("Approve 👍", key=f"approve_{chat_index}"):
-                    chat['approved'] = True
-                    chat['feedback'] = 1
-                    st.session_state.chat_history[chat_index] = chat
-                    update_google_sheets(chat['id'], feedback=1, annotated_answer=chat['response'])
-
-            if chat.get('annotate', False):
-                annotated_answer = st.text_area("Annotate Answer", value=chat['response'], key=f"annotate_text_{chat_index}")
-                if st.button("Submit Annotated Answer", key=f"submit_annotate_{chat_index}"):
-                    chat['annotated_answer'] = annotated_answer
-                    chat['annotate'] = False
-                    st.session_state.chat_history[chat_index] = chat
-                    update_google_sheets(chat['id'], annotated_answer=annotated_answer)
-            
-            feedback = st.text_area("How was the response? Does it match the context? Does it answer the question fully?", key=f"textarea_{chat_index}")
-            if st.button("Submit Feedback", key=f"submit_{chat_index}"):
-                chat['detailed_feedback'] = feedback
-                st.session_state.chat_history[chat_index] = chat
-                update_google_sheets(chat['id'], detailed_feedback=feedback)
-
-    if question := st.chat_input("Enter your question"):
-        if st.session_state.query_engine:
-            with st.spinner('Generating response...'):
-                # Compile chat history for context
-                history = "\n".join([f"Q: {chat['user']}\nA: {chat['response']}" for chat in st.session_state.chat_history])
-                full_query = f"{history}\nQ: {question}"
-                response = st.session_state.query_engine.query(full_query)
-                logging.info(f"Generated response: {response.response}")
-                logging.info(f"Retrieved contexts: {[node.text for node in response.source_nodes]}")
-            question_id = str(uuid.uuid4())
-            timestamp = datetime.now().isoformat()
-            st.session_state.chat_history.append({'id': question_id, 'user': question, 'response': response.response, 'contexts': response.source_nodes, 'feedback': 0, 'detailed_feedback': '', 'annotated_answer': '', 'timestamp': timestamp})
-            
-            # Log initial query and response to Google Sheets without feedback
-            log_to_google_sheets([question_id, question, response.response, st.session_state['selected_api'], selected_model, embedding_model, retriever_method, chunk_size, top_k, 0, "", "", timestamp])
+        st.markdown(
+            """
+            <p style="color:grey; font-size:12px;">
+            <strong>Disclaimer:</strong> This system is an academic prototype demonstration of our hybrid parameter-adaptive retrieval-augmented generation system. It is <strong>NOT</strong> a production-ready application. All outputs should be considered experimental and may not be fully accurate. This system should not be used for making important legal decisions. For complete, specific, and tailored legal advice, please consult a licensed legal professional.<br><br>
+            </p>
+            """, 
+            unsafe_allow_html=True
+        )
+
+
+
+    # Check if the retriever is initialized
+    if "chat_engine" in st.session_state:
+        chat_engine = st.session_state.chat_engine
+    else:
+        st.warning("Please initialize the retriever from the sidebar.")
+
+
+    # Initialize session state for chat messages
+    if "messages" not in st.session_state:
+        st.session_state.messages = [{"role": "assistant", "content": "How may I assist you?"}]
+
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.write(message["content"])
+
             
-            st.rerun()
-        else:
-            st.error("Query engine is not initialized. Please initialize it first.")
+
+    # User-provided prompt
+    if prompt := st.chat_input():
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.write(prompt)
+
+        # Generate a response if the last message is from the user
+        if st.session_state.messages[-1]["role"] == "user":
+            with st.chat_message("assistant"):
+                with st.spinner("Retrieving Knowledge..."):
+                    response = chat_engine.stream_chat(prompt)
+                    response_str = ""
+                    response_container = st.empty()
+                    for token in response.response_gen:
+                        response_str += token
+                        response_container.write(response_str)
+                # Expander for additional info
+                with st.expander("Source Nodes"):
+                    # Display source nodes
+                    if hasattr(response, 'source_nodes') and response.source_nodes:
+                        
+                        for idx, node in enumerate(response.source_nodes):
+                            st.markdown(f"#### Source Node {idx + 1}")
+                            st.write(f"**Node ID:** {node.node_id}")
+                            st.write(f"**Node Score:** {node.score}")
+                            
+                            st.write("**Metadata:**")
+                            for key, value in node.metadata.items():
+                                st.write(f"- **{key}:** {value}")
+                            
+                            st.write("**Content:**")
+                            st.write(node.node.get_content())
+
+                            # Add a horizontal line to separate nodes
+                            st.markdown("---")
+                    else:
+                        st.write("No additional source nodes available.")
+
+            st.session_state.messages.append({"role": "assistant", "content": str(response)})
 
 if __name__ == "__main__":
-    run_streamlit_app()
+    main()

retrievers.py

@@ -0,0 +1,465 @@
+from prompts import get_classification_prompt, get_query_generation_prompt
+from utils_code import initialize_openai_creds, create_llm
+from llama_index.core.schema import QueryBundle, NodeWithScore
+from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever
+from transformers import pipeline
+from typing import List, Optional
+import asyncio
+from llama_index.core.postprocessor import SentenceTransformerRerank
+from llama_index.core.indices.property_graph import LLMSynonymRetriever
+from llama_index.core.indices.property_graph import VectorContextRetriever, PGRetriever
+from llama_index.core.retrievers import BaseRetriever, VectorIndexRetriever, KGTableRetriever
+import os
+
+
+class PARetriever(BaseRetriever):
+    """Custom retriever that performs query rewriting, Vector search, and BM25 search without Knowledge Graph search."""
+
+    def __init__(
+        self,
+        llm,  # LLM for query generation
+        vector_retriever: Optional[VectorIndexRetriever] = None,
+        bm25_retriever: Optional[BaseRetriever] = None,
+        mode: str = "OR",
+        rewriter: bool = True,
+        classifier_model: Optional[str] = None,  # Optional classifier model
+        device: str = 'mps',  # Set to 'mps' as the default device
+        reranker_model_name: Optional[str] = None,  # Model name for SentenceTransformerRerank
+        verbose: bool = False,  # Verbose flag
+        fixed_params: Optional[dict] = None,  # New parameter to pass in fixed parameters
+        categories_list: Optional[List[str]] = None,  # List of categories for query classification
+        param_mappings: Optional[dict] = None  # Custom parameter mappings based on classifier labels
+    ) -> None:
+        """Initialize PARetriever parameters."""
+        self._vector_retriever = vector_retriever
+        self._bm25_retriever = bm25_retriever
+        self._llm = llm
+        self._rewriter = rewriter
+        self._mode = mode
+        self._reranker_model_name = reranker_model_name
+        self._reranker = None  # Initialize reranker as None
+        self.verbose = verbose
+        self.fixed_params = fixed_params
+        self.categories_list = categories_list
+        self.param_mappings = param_mappings or {  
+            "label_0": {"top_k": 5, "max_keywords_per_query": 3, "max_knowledge_sequence": 1},
+            "label_1": {"top_k": 7, "max_keywords_per_query": 4, "max_knowledge_sequence": 2},
+            "label_2": {"top_k": 10, "max_keywords_per_query": 5, "max_knowledge_sequence": 3}
+        }
+
+        # Initialize the classifier if provided
+        self.classifier = None
+        if classifier_model:
+            self.classifier = pipeline("text-classification", model=classifier_model, device=device)
+
+        if mode not in ("AND", "OR"):
+            raise ValueError("Invalid mode.")
+
+    def classify_query_and_get_params(self, query: str) -> (str, dict):
+        """Classify the query and determine adaptive parameters or use fixed parameters."""
+        if self.fixed_params:
+            # Use fixed parameters from the dictionary if provided
+            params = self.fixed_params
+            classification_result = "Fixed"
+            if self.verbose:
+                print(f"Using fixed parameters: {params}")
+        else:
+            params = {
+                "top_k": 5,  # Default top-k
+                "max_keywords_per_query": 4,  # Default max keywords
+                "max_knowledge_sequence": 2  # Default max knowledge sequence
+            }
+            classification_result = None
+
+            if self.classifier:
+                classification = self.classifier(query)[0]
+                label = classification['label']  # Get the classification label directly
+                classification_result = label  # Store the classification result
+                if self.verbose:
+                    print(f"Query Classification: {classification['label']} with score {classification['score']}")
+
+                # Use custom mappings or default mappings
+                if label in self.param_mappings:
+                    params = self.param_mappings[label]
+                else:
+                    if self.verbose:
+                        print(f"Warning: No mapping found for label {label}, using default parameters.")
+
+        self._classification_result = classification_result
+        return classification_result, params
+
+    def classify_query(self, query_str: str) -> Optional[str]:
+        """Classify the query into one of the predefined categories using LLM, or skip if no categories are provided."""
+        if not self.categories_list:
+            if self.verbose:
+                print("No categories provided, skipping query classification.")
+            return None
+
+        # Generate the classification prompt using external function
+        classification_prompt = get_classification_prompt(self.categories_list) + f" Query: '{query_str}'"
+
+        response = self._llm.complete(classification_prompt)
+        category = response.text.strip()
+
+        # Return the category only if it's in the categories list
+        return category if category in self.categories_list else None
+
+    def generate_queries(self, query_str: str, category: Optional[str], num_queries: int = 3) -> List[str]:
+        """Generate query variations using the LLM, taking into account the category if applicable."""
+
+        # Generate query generation prompt using external function
+        query_gen_prompt = get_query_generation_prompt(query_str, num_queries)
+
+        response = self._llm.complete(query_gen_prompt)
+        queries = response.text.split("\n")
+
+        queries = [query.strip() for query in queries if query.strip()]
+
+        if category:
+            category_query = f"{category}"
+            queries.append(category_query)
+
+        return queries
+
+    async def run_queries(self, queries: List[str], retrievers: List[BaseRetriever]) -> dict:
+        """Run queries against retrievers."""
+        tasks = []
+        for query in queries:
+            for retriever in retrievers:
+                tasks.append(retriever.aretrieve(query))
+
+        task_results = await asyncio.gather(*tasks)
+
+        results_dict = {}
+        for i, (query, query_result) in enumerate(zip(queries, task_results)):
+            results_dict[(query, i)] = query_result
+        return results_dict
+
+    def fuse_vector_and_bm25_results(self, results_dict, similarity_top_k: int) -> List[NodeWithScore]:
+        """Fuse results from Vector and BM25 retrievers."""
+        k = 60.0  # `k` is a parameter used to control the impact of outlier rankings.
+        fused_scores = {}
+        text_to_node = {}
+
+        for nodes_with_scores in results_dict.values():
+            for rank, node_with_score in enumerate(
+                sorted(nodes_with_scores, key=lambda x: x.score or 0.0, reverse=True)
+            ):
+                text = node_with_score.node.get_content()
+                text_to_node[text] = node_with_score
+                if text not in fused_scores:
+                    fused_scores[text] = 0.0
+                fused_scores[text] += 1.0 / (rank + k)
+
+        reranked_results = dict(sorted(fused_scores.items(), key=lambda x: x[1], reverse=True))
+
+        reranked_nodes: List[NodeWithScore] = []
+        for text, score in reranked_results.items():
+            if text in text_to_node:
+                node = text_to_node[text]
+                node.score = score
+                reranked_nodes.append(node)
+            else:
+                if self.verbose:
+                    print(f"Warning: Text not found in `text_to_node`: {text}")
+
+        return reranked_nodes[:similarity_top_k]
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Retrieve nodes given query."""
+        if self._rewriter:
+            category = self.classify_query(query_bundle.query_str)
+            if self.verbose and category:
+                print(f"Classified Category: {category}")
+
+        classification_result, params = self.classify_query_and_get_params(query_bundle.query_str)
+        self._classification_result = classification_result
+
+        top_k = params["top_k"]
+
+        if self._reranker_model_name:
+            self._reranker = SentenceTransformerRerank(model=self._reranker_model_name, top_n=top_k)
+            if self.verbose:
+                print(f"Initialized reranker with top_n: {top_k}")
+
+        num_queries = 3 if top_k == 5 else 5 if top_k == 7 else 7
+        if self.verbose:
+            print(f"Number of Query Rewrites: {num_queries}")
+
+        if self._rewriter:
+            queries = self.generate_queries(query_bundle.query_str, category, num_queries=num_queries)
+            if self.verbose:
+                print(f"Generated Queries: {queries}")
+        else:
+            queries = [query_bundle.query_str]
+
+        active_retrievers = []
+        if self._vector_retriever:
+            active_retrievers.append(self._vector_retriever)
+        if self._bm25_retriever:
+            active_retrievers.append(self._bm25_retriever)
+
+        if not active_retrievers:
+            raise ValueError("No active retriever provided!")
+
+        results = {}
+        if active_retrievers:
+            results = asyncio.run(self.run_queries(queries, active_retrievers))
+            if self.verbose:
+                print(f"Fusion Results: {results}")
+
+        final_results = self.fuse_vector_and_bm25_results(results, similarity_top_k=top_k)
+
+        if self._reranker:
+            final_results = self._reranker.postprocess_nodes(final_results, query_bundle)
+            if self.verbose:
+                print(f"Reranked Results: {final_results}")
+        else:
+            final_results = final_results[:top_k]
+
+        if self._rewriter:
+            unique_nodes = {}
+            for node in final_results:
+                content = node.node.get_content()
+                if content not in unique_nodes:
+                    unique_nodes[content] = node
+            final_results = list(unique_nodes.values())
+
+        if self.verbose:
+            print(f"Final Results: {final_results}")
+
+        return final_results
+
+    def get_classification_result(self) -> str:
+        return getattr(self, "_classification_result", None)
+
+
+class HyPARetriever(PARetriever):
+    """Custom retriever that extends PARetriever with knowledge graph (KG) search."""
+    
+    def __init__(
+        self,
+        llm,  # LLM for query generation
+        vector_retriever: Optional[VectorIndexRetriever] = None,
+        bm25_retriever: Optional[BaseRetriever] = None,
+        kg_index=None,  # Pass the knowledge graph index
+        property_index: bool = True,  # Whether to use the property graph for retrieval
+        pg_filters=None,
+        **kwargs,  # Pass any additional arguments to PARetriever
+    ):
+        # Initialize PARetriever to reuse all its functionality
+        super().__init__(
+            llm=llm,
+            vector_retriever=vector_retriever,
+            bm25_retriever=bm25_retriever,
+            **kwargs
+        )
+
+        # Initialize knowledge graph (KG) specific components
+        self._pg_filters = pg_filters
+        self._kg_index = kg_index
+        self.property_index = property_index
+
+    def _initialize_kg_retriever(self, params):
+        """Initialize the KG retriever based on retrieval mode."""
+        graph_index = self._kg_index
+        filters = self._pg_filters
+
+        if self._kg_index and not self.property_index:
+            # If not using property index, use KGTableRetriever
+            return KGTableRetriever(
+                index=self._kg_index,
+                retriever_mode='hybrid',
+                max_keywords_per_query=params["max_keywords_per_query"],
+                max_knowledge_sequence=params["max_knowledge_sequence"]
+            )
+        
+        elif self._kg_index and self.property_index:
+            # If using property index, use the simpler graph index retriever
+            # Use this for the DEMO 
+
+            vector_retriever = VectorContextRetriever(
+                graph_store=graph_index.property_graph_store,
+                similarity_top_k=params["max_keywords_per_query"],
+                path_depth=params["max_knowledge_sequence"],
+                include_text=True,
+                filters=filters
+            )
+            synonym_retriever = LLMSynonymRetriever(
+                graph_store=graph_index.property_graph_store,
+                llm=self._llm,
+                include_text=True,
+                filters=filters
+            )
+            return graph_index.as_retriever(sub_retrievers=[vector_retriever, synonym_retriever])
+            #return graph_index.as_retriever(similarity_top_k=params["top_k"])
+        
+        return None
+
+    def _combine_with_kg_results(self, vector_bm25_results, kg_results):
+        """Combine KG results with vector and BM25 results."""
+        vector_ids = {n.node.id_ for n in vector_bm25_results}
+        kg_ids = {n.node.id_ for n in kg_results}
+        
+        combined_results = {n.node.id_: n for n in vector_bm25_results}
+        combined_results.update({n.node.id_: n for n in kg_results})
+
+        if self._mode == "AND":
+            result_ids = vector_ids.intersection(kg_ids)
+        else:
+            result_ids = vector_ids.union(kg_ids)
+
+        return [combined_results[rid] for rid in result_ids]
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Retrieve nodes with KG integration."""
+        # Call PARetriever's _retrieve to get the vector and BM25 results
+        final_results = super()._retrieve(query_bundle)
+
+        # If we have a KG index, initialize the retriever
+        if self._kg_index:
+            kg_retriever = self._initialize_kg_retriever(self.classify_query_and_get_params(query_bundle.query_str)[1])
+            
+            if kg_retriever:
+                kg_nodes = kg_retriever.retrieve(query_bundle)
+                
+                # Only combine KG and vector/BM25 results if property_index is True
+                if self.property_index:
+                    final_results = self._combine_with_kg_results(final_results, kg_nodes)
+        
+        return final_results
+
+
+
+import os
+from dotenv import load_dotenv
+from llama_index.llms.azure_openai import AzureOpenAI
+from llama_index.core import VectorStoreIndex, Settings
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.retrievers import KGTableRetriever, VectorIndexRetriever
+from llama_index.retrievers.bm25 import BM25Retriever
+from llama_index.readers.file import PyMuPDFReader
+from llama_index.core.chat_engine import ContextChatEngine
+from llama_index.core.memory.chat_memory_buffer import ChatMemoryBuffer
+from llama_index.core import KnowledgeGraphIndex
+from retrievers import PARetriever, HyPARetriever
+
+
+def load_documents():
+    """Load and return documents from specified file paths."""
+    loader = PyMuPDFReader()
+    documents1 = loader.load(file_path="../../legal_data/LL144/LL144.pdf")
+    documents2 = loader.load(file_path="../../legal_data/LL144/LL144_Definitions.pdf")
+    return documents1 + documents2
+
+def create_indices(documents, llm, embed_model):
+    """Create and return VectorStoreIndex and KnowledgeGraphIndex from documents."""
+    splitter = SentenceSplitter(chunk_size=512)
+    
+    vector_index = VectorStoreIndex.from_documents(
+        documents,
+        embed_model=embed_model,
+        transformations=[splitter]
+    )
+    
+    """graph_index = KnowledgeGraphIndex.from_documents(
+        documents,
+        max_triplets_per_chunk=5,
+        llm=llm,
+        embed_model=embed_model,
+        include_embeddings=True,
+        transformations=[splitter]
+    )"""
+
+    return vector_index#, graph_index
+
+def create_retrievers(vector_index, graph_index, llm, category_list):
+    """Create and return the PA and HyPA retrievers."""
+    vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=10)
+    bm25_retriever = BM25Retriever.from_defaults(index=vector_index, similarity_top_k=10)
+
+    PA_retriever = PARetriever(
+        llm=llm,
+        categories_list=category_list,
+        rewriter=True,
+        vector_retriever=vector_retriever,
+        bm25_retriever=bm25_retriever,
+        classifier_model="rk68/distilbert-q-classifier-3",
+        verbose=False
+    )
+
+    HyPA_retriever = HyPARetriever(
+        llm=llm,
+        categories_list=category_list,
+        rewriter=True,
+        kg_index=graph_index,
+        vector_retriever=vector_retriever,
+        bm25_retriever=bm25_retriever,
+        classifier_model="rk68/distilbert-q-classifier-3",
+        verbose=False,
+        property_index=False
+    )
+    
+    return PA_retriever, HyPA_retriever
+
+def create_chat_engine(retriever, memory):
+    """Create and return the ContextChatEngine using the provided retriever and memory."""
+    return ContextChatEngine.from_defaults(
+        retriever=retriever,
+        verbose=False,
+        chat_mode="context",
+        memory_cls=memory,
+        memory=memory
+    )
+
+def main():
+    # Initialize environment and LLM
+    gpt35_creds, gpt4o_mini_creds, gpt4o_creds = initialize_openai_creds()
+    llm_gpt35 = create_llm(gpt35_creds=gpt35_creds, gpt4o_mini_creds=gpt4o_mini_creds, gpt4o_creds=gpt4o_creds)
+    
+    # Set global settings for embedding model and LLM
+    embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5")
+    Settings.embed_model = embed_model
+    Settings.llm = llm_gpt35
+    
+    category_list = [
+        '5-301 Bias Audit',
+        '5-302 Data Requirements',
+        '§ 5-303 Published Results',
+        '§ 5-304 Notice to Candidates and Employees'
+    ]
+    
+    # Load documents and create indices
+    documents = load_documents()
+    vector_index, graph_index = create_indices(documents, llm_gpt35, embed_model)
+    
+    # Create retrievers
+    PA_retriever, HyPA_retriever = create_retrievers(vector_index, graph_index, llm_gpt35, category_list)
+    
+    # Initialize chat memory
+    memory = ChatMemoryBuffer.from_defaults(token_limit=8192)
+    
+    # Create chat engines
+    PA_chat_engine = create_chat_engine(PA_retriever, memory)
+    HyPA_chat_engine = create_chat_engine(HyPA_retriever, memory)
+    
+    # Sample question and response
+    question = "What is a bias audit?"
+    PA_response = PA_chat_engine.chat(question)
+    HyPA_response = HyPA_chat_engine.chat(question)
+    
+    # Output responses in a nicely formatted manner
+    print("\n" + "="*50)
+    print(f"Question: {question}")
+    print("="*50)
+    
+    print("\n------- PA Retriever Response -------")
+    print(PA_response)
+    
+    print("\n------- HyPA Retriever Response -------")
+    print(HyPA_response)
+    print("="*50 + "\n")
+
+if __name__ == '__main__':
+    main()

utils_code.py

@@ -0,0 +1,212 @@
+
+import os
+from dotenv import load_dotenv, find_dotenv
+from llama_index.llms.azure_openai import AzureOpenAI
+from llama_index.readers.file import PyMuPDFReader
+from llama_index.core.chat_engine import ContextChatEngine
+from llama_index.core import KnowledgeGraphIndex
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+
+def initialize_openai_creds():
+    """Load environment variables and set API keys."""
+    dotenv_path = find_dotenv()
+    if dotenv_path == "":
+        print("No .env file found. Make sure the .env file is in the correct directory.")
+    else:
+        print(f".env file found at: {dotenv_path}")
+
+    load_dotenv(dotenv_path)
+
+    # General Azure OpenAI settings for gpt35 and gpt-4o-mini
+    general_creds = {
+        "api_key": os.getenv('AZURE_OPENAI_API_KEY'),
+        "api_version": os.getenv("AZURE_API_VERSION"),
+        "endpoint": os.getenv("AZURE_OPENAI_ENDPOINT"),
+        "temperature": 0,  # Default temperature for models
+        "gpt35_deployment_name": os.getenv("AZURE_DEPLOYMENT_NAME"),
+        "gpt4o_mini_deployment_name": os.getenv("GPT4O_MINI_DEPLOYMENT_NAME")
+    }
+
+    # GPT-4o specific settings
+    gpt4o_creds = {
+        "api_key": os.getenv('GPT4O_API_KEY'),
+        "api_version": os.getenv("GPT4O_API_VERSION"),
+        "endpoint": os.getenv("GPT4O_AZURE_ENDPOINT"),
+        "deployment_name": os.getenv("GPT4O_DEPLOYMENT_NAME"),
+        "temperature": os.getenv("GPT4O_TEMPERATURE", 0)  # Default temperature for GPT-4o
+    }
+
+    return general_creds, gpt4o_creds
+
+
+
+def initialize_openai_creds():
+    """Load environment variables and set API keys."""
+    dotenv_path = find_dotenv()
+    if dotenv_path == "":
+        print("No .env file found. Make sure the .env file is in the correct directory.")
+    else:
+        print(f".env file found at: {dotenv_path}")
+
+    load_dotenv(dotenv_path)
+
+    # GPT-3.5 Credentials
+    gpt35_creds = {
+        "api_key": os.getenv('AZURE_OPENAI_API_KEY_GPT35'),
+        "api_version": os.getenv("AZURE_API_VERSION"),
+        "endpoint": os.getenv("AZURE_OPENAI_ENDPOINT_GPT35"),
+        "temperature": 0,  # Default temperature for models
+        "deployment_name": os.getenv("AZURE_DEPLOYMENT_NAME_GPT35")
+    }
+
+    # GPT-4o-mini Credentials (shares the same API key as GPT-3.5 but different deployment name and endpoint)
+    gpt4o_mini_creds = {
+        "api_key": os.getenv('AZURE_OPENAI_API_KEY_GPT4O_MINI'),
+        "api_version": os.getenv("AZURE_API_VERSION"),
+        "endpoint": os.getenv("AZURE_OPENAI_ENDPOINT_GPT4O_MINI"),
+        "temperature": 0,  # Default temperature for models
+        "deployment_name": os.getenv("GPT4O_MINI_DEPLOYMENT_NAME")
+    }
+
+    # GPT-4o specific credentials
+    gpt4o_creds = {
+        "api_key": os.getenv('GPT4O_API_KEY'),
+        "api_version": os.getenv("GPT4O_API_VERSION"),
+        "endpoint": os.getenv("GPT4O_AZURE_ENDPOINT"),
+        "deployment_name": os.getenv("GPT4O_DEPLOYMENT_NAME"),
+        "temperature": os.getenv("GPT4O_TEMPERATURE", 0)  # Default temperature for GPT-4o
+    }
+
+    return gpt35_creds, gpt4o_mini_creds, gpt4o_creds
+
+
+
+def create_llm(model: str, gpt35_creds: dict, gpt4o_mini_creds: dict, gpt4o_creds: dict):
+    """
+    Initialize and return the Azure OpenAI LLM based on the selected model.
+    
+    :param model: The model to initialize ("gpt35", "gpt4o", or "gpt-4o-mini").
+    :param gpt35_creds: Credentials for gpt35.
+    :param gpt4o_mini_creds: Credentials for gpt-4o-mini.
+    :param gpt4o_creds: Credentials for gpt4o.
+    """
+    if model == "gpt35":
+        return AzureOpenAI(
+            deployment_name=gpt35_creds["deployment_name"],
+            temperature=gpt35_creds["temperature"],
+            api_key=gpt35_creds["api_key"],
+            azure_endpoint=gpt35_creds["endpoint"],
+            api_version=gpt35_creds["api_version"]
+        )
+    elif model == "gpt-4o-mini":
+        return AzureOpenAI(
+            deployment_name=gpt4o_mini_creds["deployment_name"],
+            temperature=gpt4o_mini_creds["temperature"],
+            api_key=gpt4o_mini_creds["api_key"],
+            azure_endpoint=gpt4o_mini_creds["endpoint"],
+            api_version=gpt4o_mini_creds["api_version"]
+        )
+    elif model == "gpt4o":
+        return AzureOpenAI(
+            deployment_name=gpt4o_creds["deployment_name"],
+            temperature=gpt4o_creds["temperature"],
+            api_key=gpt4o_creds["api_key"],
+            azure_endpoint=gpt4o_creds["endpoint"],
+            api_version=gpt4o_creds["api_version"]
+        )
+    else:
+        raise ValueError(f"Invalid model: {model}. Choose from 'gpt35', 'gpt4o', or 'gpt-4o-mini'.")
+
+    
+    
+def create_chat_engine(retriever, memory, llm):
+    """Create and return the ContextChatEngine using the provided retriever and memory."""
+    chat_engine = ContextChatEngine.from_defaults(
+        retriever=retriever,
+        memory=memory,
+        llm=llm
+    )
+    return chat_engine
+
+
+def load_documents(filepaths):
+    """
+    Load and return documents from specified file paths.
+    
+    :param filepaths: A string (single file path) or a list of strings (multiple file paths).
+    :return: A list of loaded documents.
+    """
+    loader = PyMuPDFReader()
+
+    # If a single string is passed, convert it to a list for consistent handling
+    if isinstance(filepaths, str):
+        filepaths = [filepaths]
+
+    # Load and accumulate documents
+    all_documents = []
+    for filepath in filepaths:
+        documents = loader.load(file_path=filepath)
+        all_documents += documents
+
+    return all_documents
+
+
+def create_kg_index(
+    documents,
+    storage_context,
+    llm,
+    max_triplets_per_chunk=10,
+    embed_model=HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5"),
+    include_embeddings=True,
+    chunk_size=512
+):
+    splitter = SentenceSplitter(chunk_size=chunk_size)
+    graph_index = KnowledgeGraphIndex.from_documents(
+        documents,
+        storage_context=storage_context,
+        max_triplets_per_chunk=max_triplets_per_chunk,
+        llm=llm,
+        embed_model=embed_model,
+        include_embeddings=include_embeddings,
+        transformations=[splitter]
+    )
+    return graph_index
+
+
+from llama_index.core.indices.property_graph import SimpleLLMPathExtractor
+from llama_index.core.indices.property_graph import DynamicLLMPathExtractor
+from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore
+from llama_index.core import PropertyGraphIndex
+
+
+def create_pg_index(
+    llm,
+    documents,
+    graph_store,
+    max_triplets_per_chunk=10,
+    num_workers=4,
+    embed_kg_nodes=True,
+    embed_model=HuggingFaceEmbedding(model_name="BAAI/bge-large-en-v1.5")
+):
+    
+    splitter = SentenceSplitter(chunk_size=512)
+    # Initialize the LLM path extractor
+    kg_extractor = DynamicLLMPathExtractor(
+        llm=llm,
+        max_triplets_per_chunk=max_triplets_per_chunk,
+        num_workers=num_workers
+    )
+
+
+    # Create the Property Graph Index
+    graph_index = PropertyGraphIndex.from_documents(
+        documents,
+        property_graph_store=graph_store,
+        embed_model=embed_model,
+        embed_kg_nodes=embed_kg_nodes,
+        kg_extractors=[kg_extractor],
+        transformations=[splitter]
+    )
+
+    return graph_index