import os import sys import json import streamlit as st import warnings import traceback import logs import chromadb import hashlib import sqlite3 import regex as re from pinecone import Pinecone from typing import Optional, Dict, Any from sentence_transformers import util os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3" warnings.filterwarnings("ignore") sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), 'src'))) from sentence_transformers import SentenceTransformer from configuration import Configuration from rag_scripts.rag_pipeline import RAGPipeline from rag_scripts.documents_processing.chunking import PyMuPDFChunker from rag_scripts.embedding.embedder import SentenceTransformerEmbedder from rag_scripts.embedding.vector_db.chroma_db import chromaDBVectorDB from rag_scripts.embedding.vector_db.faiss_db import FAISSVectorDB from rag_scripts.embedding.vector_db.pinecone_db import PineconeVectorDB from rag_scripts.llm.llmResponse import GROQLLM from rag_scripts.evaluation.evaluator import RAGEvaluator class RAGOperations: VALID_VECTOR_DB = {'chroma', 'faiss', 'pinecone'} @staticmethod def check_db(vector_db_type: str, db_path: str, collection_name: str) -> bool: try: if vector_db_type not in RAGOperations.VALID_VECTOR_DB: logs.logger.info(f"Invalid Vector DB: {vector_db_type}") raise if vector_db_type.lower() == 'pinecone': pc = Pinecone(api_key=Configuration.PINECONE_API_KEY) return collection_name in pc.list_indexes().names() elif vector_db_type.lower() == 'chroma': if not os.path.exists(db_path): return False client = chromadb.PersistentClient(path=db_path) try: client.get_collection(collection_name) return True except: return False elif vector_db_type.lower() == "faiss": faiss_index_file = os.path.join(db_path, f"{collection_name}.faiss") faiss_doc_store_file = os.path.join(db_path, f"{collection_name}_docs.pkl") return os.path.exists(faiss_index_file) and os.path.exists(faiss_doc_store_file) except Exception as ex: traceback.print_exc() logs.logger.info(f"Exception in checking {vector_db_type} existence") return False @staticmethod def get_pipeline_params(chunk_size: Optional[int] =None, chunk_overlap: Optional[int]=None, embedding_model: Optional[str]=None, vector_db_type: Optional[str]=None, llm_model: Optional[str] = None, temperature: Optional[float] = None, top_p: Optional[float] = None, max_tokens: Optional[int] = None, re_ranker_model: Optional[str] = None, use_tuned: bool = False) -> Dict[str, Any]: try: best_param_path = os.path.join(Configuration.DATA_DIR, 'best_params.json') params = { 'document_path': Configuration.FULL_PDF_PATH, 'chunk_size': chunk_size if chunk_size is not None else Configuration.DEFAULT_CHUNK_SIZE, 'chunk_overlap': chunk_overlap if chunk_overlap is not None else Configuration.DEFAULT_CHUNK_OVERLAP, 'embedding_model_name': embedding_model if embedding_model is not None else Configuration.DEFAULT_SENTENCE_TRANSFORMER_MODEL, 'vector_db_type': vector_db_type if vector_db_type is not None else "chroma", 'llm_model_name': llm_model if llm_model is not None else llm_model, 'db_path': None, 'collection_name': Configuration.COLLECTION_NAME, 'vector_db': None, 'temperature': temperature if temperature is not None else 0.1, 'top_p': top_p if top_p is not None else .95, 'max_tokens': max_tokens if max_tokens is not None else 1500, 're_ranker_model': re_ranker_model if re_ranker_model is not None else Configuration.DEFAULT_RERANKER, } if use_tuned and os.path.exists(best_param_path): with open(best_param_path, 'rb') as f: best_params = json.load(f) logs.logger.info(f"Best params: {best_params} from the file {best_param_path}") params.update({ 'vector_db_type': best_params.get('vector_db_type', params['vector_db_type']), 'embedding_model_name': best_params.get('embedding_model', params['embedding_model_name']), 'chunk_overlap': best_params.get('chunk_overlap', params['chunk_overlap']), 'chunk_size': best_params.get('chunk_size', params['chunk_size']), 're_ranker_model': best_params.get('re_ranker_model', params['re_ranker_model'])}) use_tuned = True if use_tuned: tuned_db_type = params['vector_db_type'] params['db_path'] = os.path.join(Configuration.DATA_DIR, 'TunedDB', tuned_db_type) if tuned_db_type != 'pinecone' else "" params['collection_name'] = 'tuned-' + Configuration.COLLECTION_NAME if tuned_db_type in ['chroma', 'faiss']: os.makedirs(params['db_path'], exist_ok=True) logs.logger.info(f"Tuned db path: {params['db_path']}") else: params['db_path'] = (Configuration.CHROMA_DB_PATH if params['vector_db_type'] == 'chroma' else Configuration.FAISS_DB_PATH if params['vector_db_type'] == 'faiss' else "") if params['vector_db_type'] in ['chroma', 'faiss']: os.makedirs(params['db_path'], exist_ok=True) logs.logger.info(f"Created directory for {params['vector_db_type']} at {params['db_path']}") return params except Exception as ex: logs.logger.info(f"Exception in get_pipeline_params: {ex}") traceback.print_exc() @staticmethod def check_embedding_dimension(vector_db_type: str, db_path: str, collection_name: str, embedding_model: str) -> bool: if vector_db_type != 'chroma': return True try: client = chromadb.PersistentClient(path=db_path) collection = client.get_collection(collection_name) model = SentenceTransformer(embedding_model) sample_embedding = model.encode(["test"])[0] try: expected_dim = collection._embedding_function.dim except AttributeError: peek_result = collection.peek(limit=1) if 'embedding' in peek_result and peek_result['embedding']: expected_dim = len(peek_result['embedding'][0]) else: return False actual_dim = len(sample_embedding) logs.logger.info(f"Expected dimension: {expected_dim} Actual dimension: {actual_dim}") return expected_dim == actual_dim except Exception as ex: logs.logger.info(f"Error checking embedding dimension: {ex}") return False @staticmethod def initialize_pipeline(params: dict[str, Any]) -> RAGPipeline: try: embedder = SentenceTransformerEmbedder(model_name=params['embedding_model_name']) chunkerObj = PyMuPDFChunker( pdf_path=params['document_path'], chunk_size=params['chunk_size'], chunk_overlap=params['chunk_overlap']) llm_model = params['llm_model_name'] vector_db = None if params['vector_db_type'] == 'chroma': vector_db = chromaDBVectorDB(embedder=embedder, db_path=params['db_path'], collection_name=params['collection_name']) elif params['vector_db_type'] == 'faiss': vector_db = FAISSVectorDB(embedder=embedder, db_path=params['db_path'], collection_name=params['collection_name']) elif params['vector_db_type'] == 'pinecone': vector_db = PineconeVectorDB(embedder=embedder, db_path=params['db_path'], collection_name=params['collection_name']) else: raise ValueError(f"Unknown vector_db_type: {params['vector_db_type']}") return RAGPipeline(document_path=params['document_path'], chunker=chunkerObj, embedder=embedder, vector_db=vector_db, llm=GROQLLM(model_name=llm_model), re_ranker_model_name=params['re_ranker_model'] if params[ 're_ranker_model'] else Configuration.DEFAULT_RERANKER, ) except Exception as ex: logs.logger.info(f"Exception in pipeline initialize: {ex}") traceback.print_exc() sys.exit(1) @staticmethod def run_build_job(chunk_size: Optional[int] = None, chunk_overlap: Optional[int] = None, embedding_model: Optional[str] = None, vector_db_type: Optional[str]= None, llm_model: Optional[str]= None, temperature: Optional[float]= None, top_p: Optional[float]= None, max_tokens: Optional[int]= None, re_ranker_model: Optional[str] =None, use_tuned: bool = False) -> None: try: params = RAGOperations.get_pipeline_params(chunk_size=chunk_size, chunk_overlap=chunk_overlap, embedding_model=embedding_model, vector_db_type=vector_db_type, llm_model=llm_model, temperature=temperature, top_p=top_p, max_tokens=max_tokens, re_ranker_model=re_ranker_model, use_tuned=use_tuned) pipeline = RAGOperations.initialize_pipeline(params) pipeline.build_index() logs.logger.info(f"RAG Build JOB completed") except Exception as ex: logs.logger.info(f"Exception in run build job: {ex}") traceback.print_exc() raise @staticmethod def run_search_job(query: Optional[str] = None, k: int = 5, raw: bool = False, use_tuned: bool = False, llm_model: Optional[str]= None, user_context: Optional[Dict[str,str]] = None, temperature: Optional[float]= None, top_p: Optional[float]= None, max_tokens: Optional[int]= None, chunk_size: Optional[int]= None, chunk_overlap: Optional[int]= None, embedding_model: Optional[str]= None, vector_db_type: Optional[str]= None, re_ranker_model: Optional[str]= None, use_rag:bool = True) -> Dict[str, Any]: try: params = RAGOperations.get_pipeline_params(chunk_size=chunk_size, chunk_overlap=chunk_overlap, embedding_model=embedding_model, vector_db_type=vector_db_type, llm_model=llm_model, temperature=temperature, top_p=top_p, max_tokens=max_tokens, re_ranker_model=re_ranker_model, use_tuned=use_tuned) vector_db_type = params['vector_db_type'] db_path = params['db_path'] collection_name = params['collection_name'] pipeline = RAGOperations.initialize_pipeline(params) db_exists = RAGOperations.check_db(vector_db_type, db_path, collection_name) if use_rag: if not db_exists: pipeline.build_index() elif pipeline.vector_db.count_documents() == 0: pipeline.build_index() elif not RAGOperations.check_embedding_dimension(vector_db_type, db_path, collection_name, params['embedding_model_name']): logs.logger.info(f"Embedding dimension mismatch. rebuilding the index") pipeline.vector_db.delete_collection(collection_name) pipeline.build_index() else: logs.logger.info(f"Using existing {vector_db_type} database with collection: {collection_name}") if pipeline.vector_db.count_documents() == 0: logs.logger.info(f"No Documents found in vector database after re-build") sys.exit(1) evaluator = RAGEvaluator(eval_data_path=Configuration.EVAL_DATA_PATH, pdf_path=Configuration.FULL_PDF_PATH) user_query = query if query else ( input("Enter your Query: ")) if user_query.lower() == 'exit': return expected_answers = None expected_keywords = [] query_found = False try: with open(Configuration.EVAL_DATA_PATH, 'r') as f: eval_data = json.load(f) for item in eval_data: if item.get('query').strip().lower() == user_query.strip().lower(): expected_keywords = item.get('expected_keywords', []) expected_answers = item.get('expected_answer_snippet', "") query_found = True break if not expected_keywords and not expected_answers: logs.logger.info(f"No evaluation data found for query in json") except Exception as ex: logs.logger.info(f"No json file : {ex}") retrieved_documents = [] if raw: retrieved_documents = pipeline.retrieve_raw_documents( user_query, k=k * 2) logs.logger.info("Raw documents retrieved") logs.logger.info(json.dumps(retrieved_documents, indent=4)) if not retrieved_documents: response = {"summary": "No relevant documents found", "sources": []} else: query_embedding = evaluator.embedder.encode(user_query, convert_to_tensor=True, normalize_embeddings=True) similarities = [(doc, util.cos_sim(query_embedding, evaluator.embedder.encode(doc['content'], convert_to_tensor=True, normalize_embeddings=True)).item()) for doc in retrieved_documents] similarities.sort(key=lambda x: x[1], reverse=True) top_docs = similarities[:min(3, len(similarities))] truncated_content = [] for doc, sim in top_docs: content_paragraphs = re.split(r'\n\s*\n', doc['content'].strip()) para_sims = [(para, util.cos_sim(query_embedding, evaluator.embedder.encode(para.strip(), convert_to_tensor=True, normalize_embeddings=True)).item()) for para in content_paragraphs if para.strip()] para_sims.sort(key=lambda x: x[1], reverse=True) top_paras = [para for para, para_sim in para_sims[:2] if para_sim >= 0.3] if len(top_paras) < 1: # Fallback to at least one paragraph top_paras = [para for para, _ in para_sims[:1]] truncated_content.append('\n\n'.join(top_paras)) response = { "summary": "\n".join(truncated_content), "sources": [{"document_id": f"DOC {idx + 1}", "page": str(doc['metadata'].get("page_number", "NA")), "section": doc['metadata'].get("section", "NA"), "clause": doc['metadata'].get("clause", "NA")} for idx, (doc, _) in enumerate(top_docs)]} else: logs.logger.info("LLM+RAG") response = pipeline.query(user_query, k=k, include_metadata=True, user_context=user_context ) retrieved_documents = pipeline.retrieve_raw_documents( user_query, k=k) final_expected_answer = expected_answers if expected_answers is not None else "" additional_eval_metrices = {} if not query_found: logs.logger.info(f"No query found in eval_Data.json: {user_query}") raw_reference_for_score = evaluator._syntesize_raw_reference(retrieved_documents) if not final_expected_answer.strip(): final_expected_answer = raw_reference_for_score retrieved_documents_content = [doc.get('content', '') for doc in retrieved_documents] llm_as_judge = evaluator._evaluate_with_llm(user_query, response.get('summary', ''), retrieved_documents_content) if llm_as_judge: additional_eval_metrices.update(llm_as_judge) output = {"query": user_query, "response": response, "evaluation": llm_as_judge} logs.logger.info(json.dumps(output, indent=4)) return output else: output = {"query": user_query, "response": response, "evaluation": llm_as_judge} logs.logger.info(json.dumps(output, indent=4)) return output else: eval_result = evaluator.evaluate_response(user_query, response, retrieved_documents, expected_keywords, expected_answers) output = {"query": user_query, "response": response, "evaluation": eval_result} logs.logger.info(json.dumps(output, indent=2, ensure_ascii=False)) return output except Exception as ex: logs.logger.info(f"Exception in run search job {ex}") traceback.print_exc() @staticmethod def run_hypertune_job(llm_model: Optional[str] = None, search_type: str = "random", n_iter: int = 3) -> Dict[str,Any]: try: evaluator = RAGEvaluator(eval_data_path=Configuration.EVAL_DATA_PATH, pdf_path=Configuration.FULL_PDF_PATH) result = evaluator.evaluate_combined_params_grid( chunk_size_to_test=[512, 1024, 2048], chunk_overlap_to_test=[100, 200, 400], embedding_models_to_test=["all-MiniLM-L6-v2", "all-mpnet-base-v2", "paraphrase-MiniLM-L3-v2", "multi-qa-mpnet-base-dot-v1"], vector_db_types_to_test=['pinecone'], llm_model_name=llm_model, re_ranker_model=["cross-encoder/ms-marco-MiniLM-L-6-v2", "cross-encoder/ms-marco-TinyBERT-L-2"], search_type=search_type, n_iter=n_iter) best_parameter = result['best_params'] best_score = result['best_score'] pkl_file = result['pkl_file'] best_metrics = result['best_metrics'] best_param_path = os.path.join(Configuration.DATA_DIR, 'best_params.json') with open(best_param_path, 'w') as f: json.dump(best_parameter, f, indent=4) tuned_db = best_parameter['vector_db_type'] tuned_path = os.path.join(Configuration.DATA_DIR, 'TunedDB', tuned_db) if tuned_db != 'pinecone': os.makedirs(tuned_path, exist_ok=True) tuned_collection_name = "tuned-" + Configuration.COLLECTION_NAME tuned_params = { 'document_path': Configuration.FULL_PDF_PATH, 'chunk_size': best_parameter.get('chunk_size', Configuration.DEFAULT_CHUNK_SIZE), 'chunk_overlap': best_parameter.get('chunk_overlap', Configuration.DEFAULT_CHUNK_OVERLAP), 'embedding_model_name': best_parameter.get('embedding_model', Configuration.DEFAULT_SENTENCE_TRANSFORMER_MODEL), 'vector_db_type': tuned_db, 'llm_model_name': llm_model, 'db_path': tuned_path if tuned_db != 'pinecone' else "", 'collection_name': tuned_collection_name, 'vector_db': None, 're_ranker_model': best_parameter.get('re_ranker', Configuration.DEFAULT_RERANKER) } if 're_ranker_model' in best_parameter: tuned_params['re_ranker_model'] = best_parameter['re_ranker_model'] else: tuned_params['re_ranker_model'] = Configuration.DEFAULT_RERANKER tuned_pipeline = RAGOperations.initialize_pipeline(tuned_params) tuned_pipeline.build_index() return result except Exception as ex: logs.logger.info(f"Exception in hypertune: {ex} ") traceback.print_exc() @staticmethod def run_llm_with_prompt(run_type: str, temperature: float=0.1, top_p: float=0.95, max_tokens=1500) -> None: try: params = RAGOperations.get_pipeline_params() pipeline = RAGOperations.initialize_pipeline(params) evaluator = RAGEvaluator(eval_data_path=Configuration.EVAL_DATA_PATH, pdf_path=Configuration.FULL_PDF_PATH) system_message = ( "You are an expert assistant for Flykite Airlines HR Policy Queries." "Provide concise, accurate and policy-specific answers based solely on the the provided context." "Structured your response clearly, using bullet points, newlines if applicable. " "If the context lacks information, state that clearly and speculation." ) if run_type == 'prompting' else None user_query = input("Enter your query: ") expected_answer = None expected_keywords = [] try: with open(Configuration.EVAL_DATA_PATH, 'r') as f: eval_data = json.load(f) for item in eval_data: expected_answer = item.get('expected_answer_snippet', "") expected_keywords = item.get('expected_keywords', []) break except Exception as ex: logs.logger.info(f"Error loading eval_data.json for query {user_query}: {ex}") if run_type == 'prompting': prompt = ( f"You are an expert assistant for Flykite Airlines HR Policy Queries." f"Answer the following question with a structured response, using bullet points or sections where applicable" f"Base your answer solely on the query and avoid hallucination" f"Question: \n {user_query} \n" f"Answer: ") else: prompt = user_query response = pipeline.llm.generate_response( prompt=prompt, system_message=system_message, temperature=temperature, top_p=top_p, max_tokens=max_tokens ) retreived_documents = [] eval_result = evaluator.evaluate_response(user_query, response, retreived_documents, expected_keywords, expected_answer) output = {"query": user_query, "response": { "summary: ": response.strip(), "source: ": ["LLM Response Not RAG loaded"]}, "evaluation": eval_result} logs.logger.info(json.dumps(output, indent=2)) return output except Exception as ex: logs.logger.info(f"Exception in LLm_prompting response: {ex}") traceback.print_exc() return {"error": str(ex)} @staticmethod def login() -> Dict[str, str]: username = input("Enter your username: ") password = input("Enter your password: ") hashed_password = hashlib.sha256(password.encode()).hexdigest() try: conn = sqlite3.connect('users.db') cursor = conn.cursor() cursor.execute( "SELECT username,jobrole,department,location FROM users WHERE username = ? AND password = ?", (username, hashed_password) ) user = cursor.fetchone() logs.logger.info(f"{user}") conn.close() if user: return {"username": user[0], "role": user[1], "department": user[2], "location": user[3]} else: logs.logger.info("Invalid username or password") sys.exit(1) except sqlite3.Error as ex: return False @staticmethod def authenticate_user(username, password) -> Optional[Dict[str, str]]: hashed_password = hashlib.sha256(password.encode()).hexdigest() conn = sqlite3.connect('users.db') cursor = conn.cursor() cursor.execute( "SELECT username, jobrole, department, location FROM users WHERE username = ? AND password = ?", (username, hashed_password) ) user = cursor.fetchone() conn.close() if user: return {"username": user[0], "role": user[1], "department": user[2], "location": user[3]} return None @staticmethod def home_page(): st.title("Welcome to Flykite RAG System") if 'logged_in' not in st.session_state: st.session_state.logged_in = False if 'user_info' not in st.session_state: st.session_state.user_info = None if not st.session_state.logged_in: st.subheader("Login") with st.form("login_form"): username = st.text_input("Username") password = st.text_input("Password", type="password") login_button = st.form_submit_button("Login") if login_button: user_data = RAGOperations.authenticate_user(username, password) if user_data: st.session_state.logged_in = True st.session_state.user_info = user_data st.session_state.user_context = { "role": user_data['role'], "department": user_data['department'], "location": user_data['location'] } st.success(f"Logged in as {user_data['username']} ({user_data['role']})") # No rerun needed here, the main_app will handle navigation st.session_state.page = "User" if user_data['role'] != 'admin' else "Admin" st.rerun() else: st.error("Invalid username or password.") else: st.write( f"You are logged in as **{st.session_state.user_info['username']}** (Role: **{st.session_state.user_info['role']}**)") if st.button("Logout"): st.session_state.logged_in = False st.session_state.user_info = None st.session_state.user_context = None st.session_state.page = "Home" # Redirect to home on logout st.rerun() @staticmethod @staticmethod def admin_page(): st.title("Admin Dashboard") st.write(f"Logged in as: {st.session_state.user_info['username']} (Role: {st.session_state.user_info['role']})") if st.session_state.user_info and st.session_state.user_info['role'] == 'admin': st.header("RAG Hypertuning") st.info("Run hyperparameter tuning to find the best RAG configuration and build a tuned index.") with st.form("hypertune_form"): st.write("Hypertuning parameters:") llm_model_ht = st.selectbox("LLM Model for Hypertuning Evaluation", options=["llama-3.3-70b-versatile", "llama-3.1-8b-instant"], index=["llama-3.3-70b-versatile", "llama-3.1-8b-instant"].index( Configuration.DEFAULT_GROQ_LLM_MODEL) if Configuration.DEFAULT_GROQ_LLM_MODEL in [ "llama-3.3-70b-versatile", "llama-3.1-8b-instant"] else 0, key="llm_model_ht_select") # New inputs for hyperparameter tuning st.subheader("Hyperparameter Ranges/Options:") chunk_sizes = st.multiselect("Chunk Sizes to Test (e.g., 256, 512, 1024)", options=[512, 1024,2048], default=[512], key="chunk_sizes_ht") chunk_overlaps = st.multiselect("Chunk Overlaps to Test (e.g., 50, 100, 200)", options=[150,200,400], default=[150], key="chunk_overlaps_ht") embedding_models = st.multiselect("Embedding Models to Test", options=["all-MiniLM-L6-v2", "all-mpnet-base-v2", "paraphrase-MiniLM-L3-v2", "multi-qa-mpnet-base-dot-v1"], default=["all-MiniLM-L6-v2", "all-mpnet-base-v2"], key="embedding_models_ht") re_ranker_models = st.multiselect("Re-ranker Models to Test", options=["cross-encoder/ms-marco-MiniLM-L-6-v2", "cross-encoder/ms-marco-TinyBERT-L-2", "None"], default=["cross-encoder/ms-marco-MiniLM-L-6-v2"], key="re_ranker_models_ht") vector_db_types = st.multiselect("Vector DB Types to Test", options=['chroma', 'faiss', 'pinecone'], default=['chroma'], key="vector_db_types_ht") search_type = st.radio("Hypertuning Search Type", options=["random", "grid"], index=0, # Default to random key="search_type_ht") n_iter = st.number_input("Number of Hyper-tuning Iterations (for Random Search)", min_value=1, value=3, step=1, help="Only applicable for 'Random' search type.", key="n_iter_ht") hypertune_button = st.form_submit_button("Run Hypertune Job") if hypertune_button: if not chunk_sizes or not chunk_overlaps or not embedding_models or not re_ranker_models or not vector_db_types: st.error("Please select at least one option for all hyperparameter categories.") else: # Handle 'None' for re-ranker model: remove "None" string and pass None object if needed final_re_ranker_models = [ None if model == "None" else model for model in re_ranker_models ] st.write("Starting RAG Hypertuning. This may take a while...") with st.spinner("Running hypertuning..."): try: result = RAGOperations.run_hypertune_job( llm_model=llm_model_ht, chunk_size_to_test=chunk_sizes, chunk_overlap_to_test=chunk_overlaps, embedding_models_to_test=embedding_models, re_ranker_model=final_re_ranker_models, vector_db_types_to_test=vector_db_types, search_type=search_type, n_iter=n_iter if search_type == "random" else None # n_iter only for random search ) if result and "error" not in result: st.success("Hypertuning completed and tuned index built!") st.subheader("Best Parameters Found:") st.json(result.get('best_params', {})) if 'best_score' in result: st.write(f"Best Score: {result['best_score']:.4f}") if 'best_metrics' in result: st.subheader("Best Metrics:") st.json(result['best_metrics']) else: st.error(f"Hypertuning failed: {result.get('error', 'Unknown error')}") except Exception as e: st.error(f"An unexpected error occurred during hypertuning: {e}") st.exception(e) # Display full traceback in Streamlit st.header("RAG Testing") st.info("Test the RAG pipeline with a specific query, optionally using the tuned database.") with st.form("rag_test_form"): test_query = st.text_area("Enter a test query for the RAG system:", value="What is the policy on annual leave?", key="test_query_input") use_tuned_db = st.checkbox("Use Tuned RAG Database (if hypertuned previously)", value=True, key="use_tuned_db_checkbox") display_raw = st.checkbox("Display Raw Retrieved Documents only (no LLM)", key="display_raw_docs_checkbox") k_value = st.slider("Number of documents to retrieve (k)", min_value=1, max_value=10, value=5, key="k_value_slider") test_rag_button = st.form_submit_button("Run RAG Test Query") if test_rag_button: st.write("Running RAG test query...") with st.spinner("Getting RAG response..."): try: result = RAGOperations.run_search_job( query=test_query, k=k_value, raw=display_raw, use_tuned=use_tuned_db, llm_model=st.session_state.get('llm_model_ht_select', Configuration.DEFAULT_GROQ_LLM_MODEL), user_context=st.session_state.user_context ) if result and "error" not in result: st.success("RAG Test Query Completed!") st.subheader("RAG Response:") if display_raw: st.json(result.get('response', {})) else: response_data = result.get('response', {}) if 'summary' in response_data: st.write(response_data['summary']) if 'sources' in response_data and response_data['sources']: st.subheader("Sources:") for source in response_data['sources']: if isinstance(source, dict): st.markdown( f"- **Document ID:** {source.get('document_id', 'N/A')}, **Page:** {source.get('page', 'N/A')}, **Section:** {source.get('section', 'N/A')}, **Clause:** {source.get('clause', 'N/A')}") else: st.markdown(f"- {source}") else: st.json(response_data) if 'evaluation' in result: st.subheader("Evaluation Results:") st.json(result['evaluation']) else: st.error(f"RAG test query failed: {result.get('error', 'Unknown error')}") except Exception as e: st.error(f"An unexpected error occurred during RAG test: {e}") st.exception(e) else: st.warning("You do not have administrative privileges to view this page.") if st.button("Go to User Page"): st.session_state.page = "User" st.rerun() @staticmethod def run_hypertune_job(llm_model: Optional[str] = None, chunk_size_to_test: Optional[list[int]] = None, # Added parameter chunk_overlap_to_test: Optional[list[int]] = None, # Added parameter embedding_models_to_test: Optional[list[str]] = None, # Added parameter vector_db_types_to_test: Optional[list[str]] = None, # Added parameter re_ranker_model: Optional[list[str]] = None, # Added parameter search_type: str = "random", n_iter: Optional[int] = 3) -> Dict[str, Any]: try: evaluator = RAGEvaluator(eval_data_path=Configuration.EVAL_DATA_PATH, pdf_path=Configuration.FULL_PDF_PATH) result = evaluator.evaluate_combined_params_grid( chunk_size_to_test=chunk_size_to_test if chunk_size_to_test is not None else [512, 1024, 2048], chunk_overlap_to_test=chunk_overlap_to_test if chunk_overlap_to_test is not None else [100, 200, 400], embedding_models_to_test=embedding_models_to_test if embedding_models_to_test is not None else [ "all-MiniLM-L6-v2", "all-mpnet-base-v2", "paraphrase-MiniLM-L3-v2", "multi-qa-mpnet-base-dot-v1"], vector_db_types_to_test=vector_db_types_to_test if vector_db_types_to_test is not None else ['chroma'], llm_model_name=llm_model, re_ranker_model=re_ranker_model if re_ranker_model is not None else [ "cross-encoder/ms-marco-MiniLM-L-6-v2", "cross-encoder/ms-marco-TinyBERT-L-2"], search_type=search_type, n_iter=n_iter) best_parameter = result['best_params'] best_score = result['best_score'] pkl_file = result['pkl_file'] best_metrics = result['best_metrics'] best_param_path = os.path.join(Configuration.DATA_DIR, 'best_params.json') with open(best_param_path, 'w') as f: json.dump(best_parameter, f, indent=4) tuned_db = best_parameter['vector_db_type'] tuned_path = os.path.join(Configuration.DATA_DIR, 'TunedDB', tuned_db) if tuned_db != 'pinecone': os.makedirs(tuned_path, exist_ok=True) tuned_collection_name = "tuned-" + Configuration.COLLECTION_NAME tuned_params = { 'document_path': Configuration.FULL_PDF_PATH, 'chunk_size': best_parameter.get('chunk_size', Configuration.DEFAULT_CHUNK_SIZE), 'chunk_overlap': best_parameter.get('chunk_overlap', Configuration.DEFAULT_CHUNK_OVERLAP), 'embedding_model_name': best_parameter.get('embedding_model', Configuration.DEFAULT_SENTENCE_TRANSFORMER_MODEL), 'vector_db_type': tuned_db, 'llm_model_name': llm_model, 'db_path': tuned_path if tuned_db != 'pinecone' else "", 'collection_name': tuned_collection_name, 'vector_db': None, 're_ranker_model': best_parameter.get('re_ranker', Configuration.DEFAULT_RERANKER) } if 're_ranker_model' in best_parameter: tuned_params['re_ranker_model'] = best_parameter['re_ranker_model'] else: tuned_params['re_ranker_model'] = Configuration.DEFAULT_RERANKER tuned_pipeline = RAGOperations.initialize_pipeline(tuned_params) tuned_pipeline.build_index() return result except Exception as ex: logs.logger.info(f"Exception in hypertune: {ex} ") traceback.print_exc() return {"error": str(ex)} # Return error for Streamlit to display @staticmethod def user_page(): st.title("Flykite HR Policy Query") st.write(f"Logged in as: {st.session_state.user_info['username']} (Role: {st.session_state.user_info['role']})") st.info("Ask any question about the Flykite Airlines HR policy document.") with st.form("user_query_form"): user_query = st.text_area("Your Query:", height=100, key="user_query_input") response_type = st.radio("Choose Response Type:", options=["LLM Tuned Response (RAG + LLM)", "RAG Raw Response (Retrieved Docs Only)"], index=0, key="response_type_radio") k_value_user = st.slider("Number of documents to consider (k)", min_value=1, max_value=10, value=5, key="k_value_user_slider") submit_query_button = st.form_submit_button("Get Answer") if submit_query_button and user_query: st.subheader("Response:") with st.spinner("Fetching answer..."): try: display_raw = (response_type == "RAG Raw Response (Retrieved Docs Only)") # Direct call to RAGOperations.run_search_job result = RAGOperations.run_search_job( query=user_query, raw=display_raw, k=k_value_user, use_tuned=True, # User page always uses tuned if available user_context=st.session_state.user_context # Pass user context ) if result and "error" not in result: response_data = result.get('response', {}) evaluation = result.get('evaluation',{}) if display_raw: st.json(response_data) # Raw output from main.py is already formatted else: if 'summary' in response_data: st.markdown(response_data['summary']) if 'sources' in response_data and response_data['sources']: st.subheader("Sources:") for source in response_data['sources']: if isinstance(source, dict): st.markdown( f"- **Document ID:** {source.get('document_id', 'N/A')}, **Page:** {source.get('page', 'N/A')}, **Section:** {source.get('section', 'N/A')}, **Clause:** {source.get('clause', 'N/A')}") else: # Fallback for raw string sources st.markdown(f"- {source}") else: st.json(response_data) if evaluation: #st.markdown(f"**Evaluation Results:** **Groundedness Score** {evaluation.get('Groundedness score', 'N/A')}, **Relevance Score:** {evaluation.get('Relevance score', 'N/A')}, **Reasoning** {evaluation.get('Reasoning', 'N/A')}") st.json(evaluation) else: st.error( f"Failed to get a response: {result.get('error', 'Unknown error')}. Please try again.") except Exception as e: st.error(f"An unexpected error occurred during user query: {e}") st.error(traceback.format_exc()) elif submit_query_button and not user_query: st.warning("Please enter a query.") def main_app(): st.sidebar.title("Navigation") if 'logged_in' not in st.session_state: st.session_state.logged_in = False if 'page' not in st.session_state: st.session_state.page = "Home" if not st.session_state.logged_in: st.session_state.page = "Home" RAGOperations.home_page() else: st.sidebar.button("Home", on_click=lambda: st.session_state.update(page="Home")) if st.session_state.user_info and st.session_state.user_info['role'] == 'admin': st.sidebar.button("Admin Dashboard", on_click=lambda: st.session_state.update(page="Admin")) st.sidebar.button("User Query", on_click=lambda: st.session_state.update(page="User")) else: st.sidebar.button("User Query", on_click=lambda: st.session_state.update(page="User")) if st.session_state.page == "Home": RAGOperations.home_page() elif st.session_state.page == "Admin": RAGOperations.admin_page() elif st.session_state.page == "User": RAGOperations.user_page() if __name__ == "__main__": main_app()