app.py

import os
import streamlit as st
from datetime import datetime
import json
import requests
import uuid
from datetime import date, datetime
import requests
from pydantic import BaseModel, Field
from typing import Optional
from retriver import retriever
import pandas as pd
import os

df_chunks = pd.read_pickle('Chunks_Complete.pkl')

placeHolderPersona1 = """##Mission
Please create a highly targeted query for a semantic search engine. The query must represent the conversation to date. 
** You will be given the converstaion to date in the user prompt.
** If no converstaion provided then this is the first converstaion

##Rules
Ensure the query is concise
Do not respond with anything other than the query for the Semantic Search Engine.
Respond with just a plain string """

class ChatRequestClient(BaseModel):

    user_id: str
    user_input: str
    numberOfQuestions: int
    welcomeMessage: str
    llm1: str
    tokens1: int
    temperature1: float
    persona1SystemMessage: str
    persona2SystemMessage: str
    userMessage2: str
    llm2: str
    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 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, k):
    """
    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 = retriever(query, k)
    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'],
            "score": result['score'],
            "PageNumber": result['PageNumber'],
            "ReconstructedText": reconstructed_text
        })

    return reconstructed_results

def call_chat_api(data: ChatRequestClient, k):
    url = "https://agent-builder-api.greensea-b20be511.northeurope.azurecontainerapps.io/chat/"
    # Validate and convert the data to a dictionary
    validated_data = data.dict()
    
    # Make the POST request to the FastAPI server
    response = requests.post(url, json=validated_data)
    
    if response.status_code == 200:
        body = response.json()
        query = body.get("content")
        final_results = search_and_reconstruct(query, df_chunks, k)
        return body, final_results # Return the JSON response if successful
    else:
        return "An error occured"  # Return the raw response text if not successful



# Title of the application
# st.image('agentBuilderLogo.png')
st.title('RAG Design and Evaluator')

# Sidebar for inputting personas
st.sidebar.image('cognizant_logo.jpg')
st.sidebar.header("Query Designer")
# st.sidebar.subheader("Welcome Message")
# welcomeMessage = st.sidebar.text_area("Define Intake Persona", value=welcomeMessage, height=300)
st.sidebar.subheader("Query Designer Config")
# numberOfQuestions = st.sidebar.slider("Number of Questions", min_value=0, max_value=10, step=1, value=5, key='persona1_questions')
persona1SystemMessage = st.sidebar.text_area("Query Designer System Message", value=placeHolderPersona1, height=300)

llm1 = st.sidebar.selectbox("Model Selection", ['GPT-4', 'GPT3.5'], key='persona1_size')
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')

k = st.sidebar.slider("Returned Docs", min_value=1, max_value=10, step=1, value=3, key='k')

st.sidebar.caption(f"Session ID: {genuuid()}")

# Main chat interface
st.markdown("""#### Query Translation in RAG Architecture

Query translation in a Retrieval-Augmented Generation (RAG) architecture is the process where an LLM acts as a translator between the user's natural language input and the retrieval system. 

##### Key Functions of Query Translation:
1. **Adds Context**  
   The LLM enriches the user's input with relevant context (e.g., expanding vague questions or specifying details) to make it more precise.
   
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
user_id = st.text_input("Experiment ID:", key="user_id")

# Ensure user_id is defined or fallback to a default value
if not user_id:
    st.warning("Please provide an experiment ID to start the chat.")
else:
    # Initialize chat history in session state
    if "messages" not in st.session_state:
        st.session_state.messages = []

    retrival = []
    response = {}

    if user_input := st.chat_input("Start chat:"):
        st.session_state.messages.append({"role": "user", "content": user_input})
        data = ChatRequestClient(
            user_id=user_id,
            user_input=user_input,
            numberOfQuestions=1000,
            welcomeMessage="",
            llm1=llm1,
            tokens1=tokens1,
            temperature1=temp1,
            persona1SystemMessage=persona1SystemMessage,
            persona2SystemMessage="",
            userMessage2="",
            llm2="GPT3.5",
            tokens2=1000,
            temperature2=0.2
        
        )

        response, retrival = call_chat_api(data, k)
        agent_message = response.get("content", "No response received from the agent.")
        elapsed_time = response.get("elapsed_time", 0)
        st.session_state.messages.append({"role": "assistant", "content": agent_message})

    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"**Score** {entry['score']}")
                st.write(f"**Page Number:** {entry['PageNumber']}")
                st.write("Grounding Text", entry['ReconstructedText'], height=150)