Upload 11 files

README.md

@@ -1,14 +1,12 @@
 ---
-title: HyPA RAG
-emoji: 🌍
-colorFrom: indigo
-colorTo: blue
+title: Ragdemo
+emoji: 🦀
+colorFrom: pink
+colorTo: green
 sdk: streamlit
-sdk_version: 1.39.0
+sdk_version: 1.36.0
 app_file: app.py
 pinned: false
-license: mit
-short_description: RAG demo to test queries against the NYC Local Law 144
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,348 @@
+import streamlit as st
+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.core import Settings
+import os
+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
+
+# Load environment variables from the .env file
+dotenv_path = find_dotenv()
+#print(f"Dotenv Path: {dotenv_path}")
+load_dotenv(dotenv_path)
+
+
+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
+        
+
+# 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,
+            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
+        )
+    else:
+        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
+        )
+
+    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
+            )
+
+        # Checkboxes for reranker and rewriter options
+        use_reranker = st.checkbox("Use Reranker")
+        use_rewriter = st.checkbox("Use Rewriter")
+
+        # 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")
+            )
+            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)
+
+
+
+
+        
+        # Add a disclaimer at the bottom
+        st.markdown("---")  # Horizontal line for separation
+        
+        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"])
+
+            
+
+    # 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__":
+    main()

prompts.py

@@ -0,0 +1,27 @@
+from typing import List
+
+def get_classification_prompt(categories_list: List[str]) -> str:
+    """Generate classification prompt based on the categories list."""
+    categories_str = ', '.join([f"'{category}'" for category in categories_list])
+    return (
+        f"Classify the following query into one of the following categories: {categories_str}. "
+        f"If it doesn't fit into any category, respond with 'None'. "
+        f"Return the classification, do not output absolutely anything else."
+    )
+
+
+def get_query_generation_prompt(query_str: str, num_queries: int) -> str:
+    """Generate query generation prompt based on query string and number of sub-queries."""
+    return (
+        f"You are an expert at distilling a user question into sub-questions that can be used to fully answer the original query. "
+        f"First, identify the key words from the original question below: \n"
+        f"{query_str}"
+        f"Generate {num_queries} sub-queries that cover the different aspects needed to fully address the user's query.\n\n"
+        f"Here is an example: \n"
+        f"Original Question: What does test data mean and what do I need to know about it?\n"
+        f"Output:\n"
+        f"definition of 'test data'\n"
+        f"test data requirements and conditions for a bias audit\n"
+        f"examples of the use of test data in a bias audit\n\n"
+        f"Output the rewritten sub-queries, one on each line, do not output absolutely anything else."
+    )

requirements.txt

@@ -0,0 +1,18 @@
+pinecone-client
+llama-index
+llama-index-core
+llama-index-llms-openai
+llama-index-llms-replicate
+llama-index-embeddings-huggingface
+llama-index-vector-stores-pinecone
+llama-index-readers-file
+llama-index-retrievers-bm25
+llama-index-llms-groq
+llama-index-llms-azure-openai
+llama-index-llms-openai
+llama-index-readers-file
+llama-index-graph-stores-neo4j
+oauth2client
+gspread
+python-dotenv
+PyMuPDF==1.24.0

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 = 'cpu',  # Device to CPU for huggingface demo
+        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