main.py

import streamlit
import os
import sys
import json
import argparse
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 SentenceTransformer, 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':
                return os.path.exists(db_path) and os.listdir(db_path)
            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:
            logs.logger.info(f"Exception in checking {vector_db_type} existence")
            logs.logger.info(traceback.print_exc())
            return False

    @staticmethod
    def get_pipeline_params(args: argparse.Namespace, 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':args.chunk_size,
                'chunk_overlap':args.chunk_overlap,
                'embedding_model_name':args.embedding_model,
                'vector_db_type':args.vector_db_type,
                'llm_model_name':args.llm_model,
                'db_path': None,
                'collection_name': Configuration.COLLECTION_NAME,
                'vector_db': None,
                'temperature': args.temperature,
                'top_p':args.top_p,
                'max_tokens':args.max_tokens,
                're_ranker_model':args.re_ranker_model
            }

            if 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}")
            logs.logger.info(traceback.print_exc())
            sys.exit(1)


    @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(args: argparse.Namespace) -> None:
        try:
            params = RAGOperations.get_pipeline_params(args)
            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()
            sys.exit(1)


    @staticmethod
    def run_search_job(args: argparse.Namespace,user_info: Dict[str,str]) -> None:
        try:
            params = RAGOperations.get_pipeline_params(args, use_tuned=args.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 args.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 = args.query if args.query else  (
                            input("Enter your Query: "))
            if user_query.lower() == 'exit':
                return
            user_context = {"role": user_info['role'],
                            "location": user_info['location'],
                            "department":user_info['department'] }

            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 args.raw:
                retrieved_documents = pipeline.retrieve_raw_documents(
                                        user_query, k=args.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=args.k,
                                          include_metadata=True,
                                          user_context=user_context
                                          )
                retrieved_documents = pipeline.retrieve_raw_documents(
                    user_query, k=args.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": additional_eval_metrices}
                    logs.logger.info(json.dumps(output, indent=4))
                    return json.dumps(output)
                else:
                    output = { "query": user_query, "response":response, "evaluation":llm_as_judge }
                    logs.logger.info(json.dumps(output, indent=4))
                    return json.dumps(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 json.dumps(output)


        except Exception as ex:
            logs.logger.info(f"Exception in run search job {ex}")
            traceback.print_exc()

    @staticmethod
    def run_hypertune_job(args: argparse.Namespace) -> None:
        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=args.llm_model,
                re_ranker_model = [ "cross-encoder/ms-marco-MiniLM-L-6-v2",
                    "cross-encoder/ms-marco-TinyBERT-L-2"],
                search_type='random', n_iter=1 )
            # embedding_models_to_test = ["all-MiniLM-L6-v2",
            #                             "all-mpnet-base-v2",
            #                             "paraphrase-MiniLM-L3-v2",
            #                             "multi-qa-mpnet-base-dot-v1"]
            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':args.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()

        except Exception as ex:
            logs.logger.info(f"Exception in hypertune: {ex} ")
            traceback.print_exc()


    @staticmethod
    def run_llm_with_prompt(args: argparse.Namespace,run_type: str) -> None:
        try:
            params = RAGOperations.get_pipeline_params(args,
                                        use_tuned=args.use_tuned)
            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 = args.temperature,
                    top_p = args.top_p,
                    max_tokens = args.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))

        except Exception as ex:
            logs.logger.info(f"Exception in LLm_prompting response: {ex}")
            traceback.print_exc()
            sys.exit(1)

    @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