Update app.py

app.py

@@ -8,7 +8,11 @@ from datetime import date, datetime
 import requests
 from pydantic import BaseModel, Field
 from typing import Optional
-from 
+from retriver import retriever
+import pandas as pd
+import os
+
+df_chunks = pd.read_pickle('Chunks_Complete.pkl')
 
 placeHolderPersona1 = """
 ##Mission
@@ -37,6 +41,102 @@ class ChatRequestClient(BaseModel):
     tokens2: int
     temperature2: float
 
+
+def genuuid ():
+    return uuid.uuid4()
+
+def format_elapsed_time(time):
+    # Format the elapsed time to two decimal places
+    return "{:.2f}".format(time)
+
+def search_knowledgebase(query):
+    results = retriever(query)
+    return results
+
+def process_search_results(search_results):
+    """
+    Processes search results to extract and organize metadata and other details.
+
+    :param search_results: List of search result matches from Pinecone.
+    :return: A list of dictionaries containing relevant metadata and scores.
+    """
+    processed_results = []
+
+    for result in search_results:
+        processed_results.append({
+            "id": result['id'],
+            "score": result['score'],
+            "Title": result['metadata'].get('Title', ''),
+            "ChunkText": result['metadata'].get('ChunkText', ''),
+            "PageNumber": result['metadata'].get('PageNumber', ''),
+            "Chunk": result['metadata'].get('Chunk', '')
+        })
+
+    return processed_results
+
+def reconstruct_text_from_chunks(df_chunks):
+    """
+    Reconstructs a single string of text from the chunks in the DataFrame.
+
+    :param df_chunks: DataFrame with columns ['Title', 'Chunk', 'ChunkText', 'TokenCount', 'PageNumber', 'ChunkID']
+    :return: A string combining all chunk texts in order.
+    """
+    return " ".join(df_chunks.sort_values(by=['Chunk'])['ChunkText'].tolist())
+
+def lookup_related_chunks(df_chunks, chunk_id):
+    """
+    Returns all chunks matching the title and page number of the specified chunk ID,
+    including chunks from the previous and next pages, handling edge cases where
+    there is no preceding or succeeding page.
+
+    :param df_chunks: DataFrame with columns ['Title', 'Chunk', 'ChunkText', 'TokenCount', 'PageNumber', 'ChunkID']
+    :param chunk_id: The unique ID of the chunk to look up.
+    :return: DataFrame with all chunks matching the title and page range of the specified chunk ID.
+    """
+    target_chunk = df_chunks[df_chunks['ChunkID'] == chunk_id]
+    if target_chunk.empty:
+        raise ValueError("Chunk ID not found")
+
+    title = target_chunk.iloc[0]['Title']
+    page_number = target_chunk.iloc[0]['PageNumber']
+
+    # Determine the valid page range
+    min_page = df_chunks[df_chunks['Title'] == title]['PageNumber'].min()
+    max_page = df_chunks[df_chunks['Title'] == title]['PageNumber'].max()
+
+    page_range = [page for page in [page_number - 1, page_number, page_number + 1] if min_page <= page <= max_page]
+
+    return df_chunks[(df_chunks['Title'] == title) & (df_chunks['PageNumber'].isin(page_range))]
+
+
+def search_and_reconstruct(query, df_chunks):
+    """
+    Combines search, lookup of related chunks, and text reconstruction.
+
+    :param query: The query string to search for.
+    :param df_chunks: DataFrame with chunk data.
+    :param namespace: Pinecone namespace to search within.
+    :param top_k: Number of top search results to retrieve.
+    :return: A list of dictionaries with document title, page number, and reconstructed text.
+    """
+    search_results = search_knowledgebase(query)
+    processed_results = process_search_results(search_results)
+
+    reconstructed_results = []
+
+    for result in processed_results:
+        chunk_id = result['id']
+        related_chunks = lookup_related_chunks(df_chunks, chunk_id)
+        reconstructed_text = reconstruct_text_from_chunks(related_chunks)
+
+        reconstructed_results.append({
+            "Title": result['Title'],
+            "PageNumber": result['PageNumber'],
+            "ReconstructedText": reconstructed_text
+        })
+
+    return reconstructed_results
+
 def call_chat_api(data: ChatRequestClient):
     url = "https://agent-builder-api.greensea-b20be511.northeurope.azurecontainerapps.io/chat/"
     # Validate and convert the data to a dictionary
@@ -46,24 +146,18 @@ def call_chat_api(data: ChatRequestClient):
     response = requests.post(url, json=validated_data)
     
     if response.status_code == 200:
-        return response.json()  # Return the JSON response if successful
+        body = response.json()
+        query = body.get("content")
+        final_results = search_and_reconstruct(query, df_chunks)
+        return body, final_results # Return the JSON response if successful
     else:
         return "An error occured"  # Return the raw response text if not successful
 
-def genuuid ():
-    return uuid.uuid4()
-
-def format_elapsed_time(time):
-    # Format the elapsed time to two decimal places
-    return "{:.2f}".format(time)
 
-def search_knowledgebase(query)
-    return results
-    
 
 # Title of the application
 # st.image('agentBuilderLogo.png')
-st.title('RAG Query Designer')
+st.title('RAG Design and Evaluator')
 
 # Sidebar for inputting personas
 st.sidebar.image('cognizant_logo.jpg')
@@ -78,19 +172,8 @@ llm1 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona
 temp1 = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.6, key='persona1_temp')
 tokens1 = st.sidebar.slider("Tokens", min_value=0, max_value=4000, step=100, value=500, key='persona1_tokens')
 
-# # Persona 2
-# st.sidebar.subheader("Recommendation and Next Best Action AI")
-# persona2SystemMessage = st.sidebar.text_area("Define Recommendation Persona", value=placeHolderPersona2, height=300)
-# with st.sidebar.expander("See explanation"):
-#     st.write("This AI persona uses the output of the symptom intake AI as its input. This AI’s job is to augment a health professional by assisting with a diagnosis and possible next best action. The teams will need to determine if this should be a tool used directly by the patient, as an assistant to the health professional or a hybrid of the two. ")
-#     st.image("agentPersona2.png")
-# llm2 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona2_size')
-# temp2 = st.sidebar.slider("Temperature", min_value=0.0, max_value=1.0, step=0.1, value=0.5, key='persona2_temp')
-# tokens2 = st.sidebar.slider("Tokens", min_value=0, max_value=4000, step=100, value=500, key='persona2_tokens')
-# userMessage2 = st.sidebar.text_area("Define User Message", value="This is the conversation todate, ", height=150)
 st.sidebar.caption(f"Session ID: {genuuid()}")
 
-
 # Main chat interface
 st.markdown("""#### Query Translation in RAG Architecture
 
@@ -103,6 +186,9 @@ Query translation in a Retrieval-Augmented Generation (RAG) architecture is the
 2. **Converts to Concise Query**  
    The LLM reformulates the input into a succinct and effective query optimized for the retrieval system's semantic search capabilities.
 
+3. **Uses Concise Query to serach Vector DB**  
+   The query is used to search the vector DB for suitable grounding information. 
+            
 ##### Purpose
 This ensures that the retrieval system receives a clear and focused query, increasing the relevance of the information it retrieves. The query translator acts as a bridge between human conversational language and the technical requirements of a semantic retrieval system.""")
 # User ID Input
@@ -116,20 +202,13 @@ else:
     if "messages" not in st.session_state:
         st.session_state.messages = []
 
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
+    retrival = []
+    response = {}
 
-    # Collect user input
     if user_input := st.chat_input("Start chat:"):
-        # Add user message to the chat history
         st.session_state.messages.append({"role": "user", "content": user_input})
-        st.chat_message("user").markdown(user_input)
-
-        # Prepare data for API call
         data = ChatRequestClient(
-            user_id=user_id,  # Ensure user_id is passed correctly
+            user_id=user_id,
             user_input=user_input,
             numberOfQuestions=1000,
             welcomeMessage="",
@@ -144,21 +223,27 @@ else:
             temperature2=0.2
         )
 
-        # Call the API
-        response = call_chat_api(data)
-
-        # Process the API response
+        response, retrival = call_chat_api(data)
         agent_message = response.get("content", "No response received from the agent.")
         elapsed_time = response.get("elapsed_time", 0)
-        count = response.get("count", 0)
-
-        # Add agent response to the chat history
         st.session_state.messages.append({"role": "assistant", "content": agent_message})
-        with st.chat_message("assistant"):
-            st.markdown(agent_message)
 
-        # Display additional metadata
-        st.caption(f"##### Time taken: {format_elapsed_time(elapsed_time)} seconds")
-        # st.caption(f"##### Question Count: {count} of {numberOfQuestions}")
+    col1, col2 = st.columns(2)
+
+    with col1:
+        for message in st.session_state.messages:
+            with st.chat_message(message["role"]):
+                st.markdown(message["content"])
+
+        if response:
+            st.chat_message("assistant").markdown(response.get("content", "No response"))
+            st.caption(f"##### Time taken: {format_elapsed_time(response.get('elapsed_time', 0))} seconds")
+
+    with col2:
+        for entry in retrival:
+            with st.container():
+                st.write(f"**Title:** {entry['Title']}")
+                st.write(f"**Page Number:** {entry['PageNumber']}")
+                st.text_area("Grounding Text", entry['ReconstructedText'], height=150)