App_Function_Libraries/RAG/RAG_Library_2.py

# RAG_Library_2.py
# Description: This script contains the main RAG pipeline function and related functions for the RAG pipeline.
#
# Import necessary modules and functions
import configparser
import logging
import os
import time
from typing import Dict, Any, List, Optional

from App_Function_Libraries.DB.Character_Chat_DB import get_character_chats, perform_full_text_search_chat, \
    fetch_keywords_for_chats, search_character_chat, search_character_cards, fetch_character_ids_by_keywords
from App_Function_Libraries.DB.RAG_QA_Chat_DB import search_rag_chat, search_rag_notes
#
# Local Imports
from App_Function_Libraries.RAG.ChromaDB_Library import process_and_store_content, vector_search, chroma_client
from App_Function_Libraries.RAG.RAG_Persona_Chat import perform_vector_search_chat
from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_custom_openai
from App_Function_Libraries.Web_Scraping.Article_Extractor_Lib import scrape_article
from App_Function_Libraries.DB.DB_Manager import fetch_keywords_for_media, search_media_db, get_notes_by_keywords, \
    search_conversations_by_keywords
from App_Function_Libraries.Utils.Utils import load_comprehensive_config
from App_Function_Libraries.Metrics.metrics_logger import log_counter, log_histogram
#
# 3rd-Party Imports
import openai
from flashrank import Ranker, RerankRequest
#
########################################################################################################################
#
# Functions:

# Initialize OpenAI client (adjust this based on your API key management)
openai.api_key = "your-openai-api-key"

# Get the directory of the current script
current_dir = os.path.dirname(os.path.abspath(__file__))
# Construct the path to the config file
config_path = os.path.join(current_dir, 'Config_Files', 'config.txt')
# Read the config file
config = configparser.ConfigParser()
# Read the configuration file
config.read('config.txt')


search_functions = {
    "Media DB": search_media_db,
    "RAG Chat": search_rag_chat,
    "RAG Notes": search_rag_notes,
    "Character Chat": search_character_chat,
    "Character Cards": search_character_cards
}

# RAG pipeline function for web scraping
# def rag_web_scraping_pipeline(url: str, query: str, api_choice=None) -> Dict[str, Any]:
#     try:
#         # Extract content
#         try:
#             article_data = scrape_article(url)
#             content = article_data['content']
#             title = article_data['title']
#         except Exception as e:
#             logging.error(f"Error scraping article: {str(e)}")
#             return {"error": "Failed to scrape article", "details": str(e)}
#
#         # Store the article in the database and get the media_id
#         try:
#             media_id = add_media_to_database(url, title, 'article', content)
#         except Exception as e:
#             logging.error(f"Error adding article to database: {str(e)}")
#             return {"error": "Failed to store article in database", "details": str(e)}
#
#         # Process and store content
#         collection_name = f"article_{media_id}"
#         try:
#             # Assuming you have a database object available, let's call it 'db'
#             db = get_database_connection()
#
#             process_and_store_content(
#                 database=db,
#                 content=content,
#                 collection_name=collection_name,
#                 media_id=media_id,
#                 file_name=title,
#                 create_embeddings=True,
#                 create_contextualized=True,
#                 api_name=api_choice
#             )
#         except Exception as e:
#             logging.error(f"Error processing and storing content: {str(e)}")
#             return {"error": "Failed to process and store content", "details": str(e)}
#
#         # Perform searches
#         try:
#             vector_results = vector_search(collection_name, query, k=5)
#             fts_results = search_db(query, ["content"], "", page=1, results_per_page=5)
#         except Exception as e:
#             logging.error(f"Error performing searches: {str(e)}")
#             return {"error": "Failed to perform searches", "details": str(e)}
#
#         # Combine results with error handling for missing 'content' key
#         all_results = []
#         for result in vector_results + fts_results:
#             if isinstance(result, dict) and 'content' in result:
#                 all_results.append(result['content'])
#             else:
#                 logging.warning(f"Unexpected result format: {result}")
#                 all_results.append(str(result))
#
#         context = "\n".join(all_results)
#
#         # Generate answer using the selected API
#         try:
#             answer = generate_answer(api_choice, context, query)
#         except Exception as e:
#             logging.error(f"Error generating answer: {str(e)}")
#             return {"error": "Failed to generate answer", "details": str(e)}
#
#         return {
#             "answer": answer,
#             "context": context
#         }
#
#     except Exception as e:
#         logging.error(f"Unexpected error in rag_pipeline: {str(e)}")
#         return {"error": "An unexpected error occurred", "details": str(e)}


# RAG Search with keyword filtering
# FIXME - Update each called function to support modifiable top-k results
def enhanced_rag_pipeline(query: str, api_choice: str, keywords: str = None, fts_top_k=10, apply_re_ranking=True, database_types: List[str] = "Media DB") -> Dict[str, Any]:
    """

    Perform full text search across specified database type.



    Args:

        query: Search query string

        api_choice: API to use for generating the response

        fts_top_k: Maximum number of results to return

        keywords: Optional list of media IDs to filter results

        database_types: Type of database to search ("Media DB", "RAG Chat", or "Character Chat")



    Returns:

        Dictionary containing search results with content

    """
    log_counter("enhanced_rag_pipeline_attempt", labels={"api_choice": api_choice})
    start_time = time.time()
    try:
        # Load embedding provider from config, or fallback to 'openai'
        embedding_provider = config.get('Embeddings', 'provider', fallback='openai')

        # Log the provider used
        logging.debug(f"Using embedding provider: {embedding_provider}")

        # Process keywords if provided
        keyword_list = [k.strip().lower() for k in keywords.split(',')] if keywords else []
        logging.debug(f"\n\nenhanced_rag_pipeline - Keywords: {keyword_list}")

        relevant_ids = {}

        # Fetch relevant IDs based on keywords if keywords are provided
        if keyword_list:
            try:
                for db_type in database_types:
                    if db_type == "Media DB":
                        relevant_media_ids = fetch_relevant_media_ids(keyword_list)
                        relevant_ids[db_type] = relevant_media_ids
                        logging.debug(f"enhanced_rag_pipeline - {db_type} relevant media IDs: {relevant_media_ids}")

                    elif db_type == "RAG Chat":
                        conversations, total_pages, total_count = search_conversations_by_keywords(
                            keywords=keyword_list)
                        relevant_conversation_ids = [conv['conversation_id'] for conv in conversations]
                        relevant_ids[db_type] = relevant_conversation_ids
                        logging.debug(
                            f"enhanced_rag_pipeline - {db_type} relevant conversation IDs: {relevant_conversation_ids}")

                    elif db_type == "RAG Notes":
                        notes, total_pages, total_count = get_notes_by_keywords(keyword_list)
                        relevant_note_ids = [note_id for note_id, _, _, _ in notes]  # Unpack note_id from the tuple
                        relevant_ids[db_type] = relevant_note_ids
                        logging.debug(f"enhanced_rag_pipeline - {db_type} relevant note IDs: {relevant_note_ids}")

                    elif db_type == "Character Chat":
                        relevant_chat_ids = fetch_keywords_for_chats(keyword_list)
                        relevant_ids[db_type] = relevant_chat_ids
                        logging.debug(f"enhanced_rag_pipeline - {db_type} relevant chat IDs: {relevant_chat_ids}")

                    elif db_type == "Character Cards":
                        # Assuming we have a function to fetch character IDs by keywords
                        relevant_character_ids = fetch_character_ids_by_keywords(keyword_list)
                        relevant_ids[db_type] = relevant_character_ids
                        logging.debug(
                            f"enhanced_rag_pipeline - {db_type} relevant character IDs: {relevant_character_ids}")

                    else:
                        logging.error(f"Unsupported database type: {db_type}")

            except Exception as e:
                logging.error(f"Error fetching relevant IDs: {str(e)}")
        else:
            relevant_ids = None

        # Extract relevant media IDs for each selected DB
        # Prepare a dict to hold relevant_media_ids per DB
        relevant_media_ids_dict = {}
        if relevant_ids:
            for db_type in database_types:
                relevant_media_ids = relevant_ids.get(db_type, None)
                if relevant_media_ids:
                    # Convert to List[str] if not None
                    relevant_media_ids_dict[db_type] = [str(media_id) for media_id in relevant_media_ids]
                else:
                    relevant_media_ids_dict[db_type] = None
        else:
            relevant_media_ids_dict = {db_type: None for db_type in database_types}

        # Perform vector search for all selected databases
        vector_results = []
        for db_type in database_types:
            try:
                db_relevant_ids = relevant_media_ids_dict.get(db_type)
                results = perform_vector_search(query, db_relevant_ids, top_k=fts_top_k)
                vector_results.extend(results)
                logging.debug(f"\nenhanced_rag_pipeline - Vector search results for {db_type}: {results}")
            except Exception as e:
                logging.error(f"Error performing vector search on {db_type}: {str(e)}")

        # Perform vector search
        # FIXME
        vector_results = perform_vector_search(query, relevant_media_ids)
        logging.debug(f"\n\nenhanced_rag_pipeline - Vector search results: {vector_results}")

        # Perform full-text search
        #v1
        #fts_results = perform_full_text_search(query, database_type, relevant_media_ids, fts_top_k)

        # v2
        # Perform full-text search across specified databases
        fts_results = []
        for db_type in database_types:
            try:
                db_relevant_ids = relevant_ids.get(db_type) if relevant_ids else None
                db_results = perform_full_text_search(query, db_type, db_relevant_ids, fts_top_k)
                fts_results.extend(db_results)
                logging.debug(f"enhanced_rag_pipeline - FTS results for {db_type}: {db_results}")
            except Exception as e:
                logging.error(f"Error performing full-text search on {db_type}: {str(e)}")

        logging.debug("\n\nenhanced_rag_pipeline - Full-text search results:")
        logging.debug(
            "\n\nenhanced_rag_pipeline - Full-text search results:\n" + "\n".join(
                [str(item) for item in fts_results]) + "\n"
        )

        # Combine results
        all_results = vector_results + fts_results

        if apply_re_ranking:
            logging.debug(f"\nenhanced_rag_pipeline - Applying Re-Ranking")
            # FIXME - add option to use re-ranking at call time
            # FIXME - specify model + add param to modify at call time
            # FIXME - add option to set a custom top X results
            # You can specify a model if necessary, e.g., model_name="ms-marco-MiniLM-L-12-v2"
            if all_results:
                ranker = Ranker()

                # Prepare passages for re-ranking
                passages = [{"id": i, "text": result['content']} for i, result in enumerate(all_results)]
                rerank_request = RerankRequest(query=query, passages=passages)

                # Rerank the results
                reranked_results = ranker.rerank(rerank_request)

                # Sort results based on the re-ranking score
                reranked_results = sorted(reranked_results, key=lambda x: x['score'], reverse=True)

                # Log reranked results
                logging.debug(f"\n\nenhanced_rag_pipeline - Reranked results: {reranked_results}")

                # Update all_results based on reranking
                all_results = [all_results[result['id']] for result in reranked_results]

        # Extract content from results (top fts_top_k by default)
        context = "\n".join([result['content'] for result in all_results[:fts_top_k]])
        logging.debug(f"Context length: {len(context)}")
        logging.debug(f"Context: {context[:200]}")

        # Generate answer using the selected API
        answer = generate_answer(api_choice, context, query)

        if not all_results:
            logging.info(f"No results found. Query: {query}, Keywords: {keywords}")
            return {
                "answer": "No relevant information based on your query and keywords were found in the database. Your query has been directly passed to the LLM, and here is its answer: \n\n" + answer,
                "context": "No relevant information based on your query and keywords were found in the database. The only context used was your query: \n\n" + query
            }
        # Metrics
        pipeline_duration = time.time() - start_time
        log_histogram("enhanced_rag_pipeline_duration", pipeline_duration, labels={"api_choice": api_choice})
        log_counter("enhanced_rag_pipeline_success", labels={"api_choice": api_choice})
        return {
            "answer": answer,
            "context": context
        }

    except Exception as e:
        # Metrics
        log_counter("enhanced_rag_pipeline_error", labels={"api_choice": api_choice, "error": str(e)})
        logging.error(f"Error in enhanced_rag_pipeline: {str(e)}")
        logging.error(f"Error in enhanced_rag_pipeline: {str(e)}")
        return {
            "answer": "An error occurred while processing your request.",
            "context": ""
        }



# Need to write a test for this function FIXME
def generate_answer(api_choice: str, context: str, query: str) -> str:
    # Metrics
    log_counter("generate_answer_attempt", labels={"api_choice": api_choice})
    start_time = time.time()
    logging.debug("Entering generate_answer function")
    config = load_comprehensive_config()
    logging.debug(f"Config sections: {config.sections()}")
    prompt = f"Context: {context}\n\nQuestion: {query}"
    try:
        if api_choice == "OpenAI":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_openai
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_openai(config['API']['openai_api_key'], prompt, "")

        elif api_choice == "Anthropic":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_anthropic
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_anthropic(config['API']['anthropic_api_key'], prompt, "")

        elif api_choice == "Cohere":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_cohere
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_cohere(config['API']['cohere_api_key'], prompt, "")

        elif api_choice == "Groq":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_groq
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_groq(config['API']['groq_api_key'], prompt, "")

        elif api_choice == "OpenRouter":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_openrouter
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_openrouter(config['API']['openrouter_api_key'], prompt, "")

        elif api_choice == "HuggingFace":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_huggingface
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_huggingface(config['API']['huggingface_api_key'], prompt, "")

        elif api_choice == "DeepSeek":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_deepseek
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_deepseek(config['API']['deepseek_api_key'], prompt, "")

        elif api_choice == "Mistral":
            from App_Function_Libraries.Summarization.Summarization_General_Lib import summarize_with_mistral
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_mistral(config['API']['mistral_api_key'], prompt, "")

        # Local LLM APIs
        elif api_choice == "Local-LLM":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_local_llm
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            # FIXME
            return summarize_with_local_llm(config['Local-API']['local_llm_path'], prompt, "")

        elif api_choice == "Llama.cpp":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_llama
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_llama(prompt, "", config['Local-API']['llama_api_key'], None, None)
        elif api_choice == "Kobold":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_kobold
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_kobold(prompt, config['Local-API']['kobold_api_key'], "", system_message=None, temp=None)

        elif api_choice == "Ooba":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_oobabooga
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_oobabooga(prompt, config['Local-API']['ooba_api_key'], custom_prompt="", system_message=None, temp=None)

        elif api_choice == "TabbyAPI":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_tabbyapi
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_tabbyapi(prompt, None, None, None, None, )

        elif api_choice == "vLLM":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_vllm
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_vllm(prompt, "", config['Local-API']['vllm_api_key'], None, None)

        elif api_choice.lower() == "ollama":
            from App_Function_Libraries.Summarization.Local_Summarization_Lib import summarize_with_ollama
            answer_generation_duration = time.time() - start_time
            log_histogram("generate_answer_duration", answer_generation_duration, labels={"api_choice": api_choice})
            log_counter("generate_answer_success", labels={"api_choice": api_choice})
            return summarize_with_ollama(prompt, "", config['Local-API']['ollama_api_IP'], config['Local-API']['ollama_api_key'], None, None, None)

        elif api_choice.lower() == "custom_openai_api":
            logging.debug(f"RAG Answer Gen: Trying with Custom_OpenAI API")
            summary = summarize_with_custom_openai(prompt, "", config['API']['custom_openai_api_key'], None,
                                                   None)
        else:
            log_counter("generate_answer_error", labels={"api_choice": api_choice, "error": str()})
            raise ValueError(f"Unsupported API choice: {api_choice}")
    except Exception as e:
        log_counter("generate_answer_error", labels={"api_choice": api_choice, "error": str(e)})
        logging.error(f"Error in generate_answer: {str(e)}")
        return "An error occurred while generating the answer."


def perform_vector_search(query: str, relevant_media_ids: List[str] = None, top_k=10) -> List[Dict[str, Any]]:
    log_counter("perform_vector_search_attempt")
    start_time = time.time()
    all_collections = chroma_client.list_collections()
    vector_results = []
    try:
        for collection in all_collections:
            collection_results = vector_search(collection.name, query, k=top_k)
            if not collection_results:
                continue  # Skip empty results
            filtered_results = [
                result for result in collection_results
                if relevant_media_ids is None or result['metadata'].get('media_id') in relevant_media_ids
            ]
            vector_results.extend(filtered_results)
        search_duration = time.time() - start_time
        log_histogram("perform_vector_search_duration", search_duration)
        log_counter("perform_vector_search_success", labels={"result_count": len(vector_results)})
        return vector_results
    except Exception as e:
        log_counter("perform_vector_search_error", labels={"error": str(e)})
        logging.error(f"Error in perform_vector_search: {str(e)}")
        raise


# V2
def perform_full_text_search(query: str, database_type: str, relevant_ids: List[str] = None, fts_top_k=None) -> List[Dict[str, Any]]:
    """

    Perform full-text search on a specified database type.



    Args:

        query: Search query string

        database_type: Type of database to search ("Media DB", "RAG Chat", "RAG Notes", "Character Chat", "Character Cards")

        relevant_ids: Optional list of media IDs to filter results

        fts_top_k: Maximum number of results to return



    Returns:

        List of search results with content and metadata

    """
    log_counter("perform_full_text_search_attempt", labels={"database_type": database_type})
    start_time = time.time()

    try:
        # Set default for fts_top_k
        if fts_top_k is None:
            fts_top_k = 10

        # Call appropriate search function based on database type
        if database_type not in search_functions:
            raise ValueError(f"Unsupported database type: {database_type}")

        # Call the appropriate search function
        results = search_functions[database_type](query, fts_top_k, relevant_ids)

        search_duration = time.time() - start_time
        log_histogram("perform_full_text_search_duration", search_duration,
                      labels={"database_type": database_type})
        log_counter("perform_full_text_search_success",
                    labels={"database_type": database_type, "result_count": len(results)})

        return results

    except Exception as e:
        log_counter("perform_full_text_search_error",
                    labels={"database_type": database_type, "error": str(e)})
        logging.error(f"Error in perform_full_text_search ({database_type}): {str(e)}")
        raise


# v1
# def perform_full_text_search(query: str, relevant_media_ids: List[str] = None, fts_top_k=None) -> List[Dict[str, Any]]:
#     log_counter("perform_full_text_search_attempt")
#     start_time = time.time()
#     try:
#         fts_results = search_db(query, ["content"], "", page=1, results_per_page=fts_top_k or 10)
#         filtered_fts_results = [
#             {
#                 "content": result['content'],
#                 "metadata": {"media_id": result['id']}
#             }
#             for result in fts_results
#             if relevant_media_ids is None or result['id'] in relevant_media_ids
#         ]
#         search_duration = time.time() - start_time
#         log_histogram("perform_full_text_search_duration", search_duration)
#         log_counter("perform_full_text_search_success", labels={"result_count": len(filtered_fts_results)})
#         return filtered_fts_results
#     except Exception as e:
#         log_counter("perform_full_text_search_error", labels={"error": str(e)})
#         logging.error(f"Error in perform_full_text_search: {str(e)}")
#         raise


def fetch_relevant_media_ids(keywords: List[str], top_k=10) -> List[int]:
    log_counter("fetch_relevant_media_ids_attempt", labels={"keyword_count": len(keywords)})
    start_time = time.time()
    relevant_ids = set()
    for keyword in keywords:
        try:
            media_ids = fetch_keywords_for_media(keyword)
            relevant_ids.update(media_ids)
        except Exception as e:
            log_counter("fetch_relevant_media_ids_error", labels={"error": str(e)})
            logging.error(f"Error fetching relevant media IDs for keyword '{keyword}': {str(e)}")
            # Continue processing other keywords

    fetch_duration = time.time() - start_time
    log_histogram("fetch_relevant_media_ids_duration", fetch_duration)
    log_counter("fetch_relevant_media_ids_success", labels={"result_count": len(relevant_ids)})
    return list(relevant_ids)


def filter_results_by_keywords(results: List[Dict[str, Any]], keywords: List[str]) -> List[Dict[str, Any]]:
    log_counter("filter_results_by_keywords_attempt", labels={"result_count": len(results), "keyword_count": len(keywords)})
    start_time = time.time()
    if not keywords:
        return results

    filtered_results = []
    for result in results:
        try:
            metadata = result.get('metadata', {})
            if metadata is None:
                logging.warning(f"No metadata found for result: {result}")
                continue
            if not isinstance(metadata, dict):
                logging.warning(f"Unexpected metadata type: {type(metadata)}. Expected dict.")
                continue

            media_id = metadata.get('media_id')
            if media_id is None:
                logging.warning(f"No media_id found in metadata: {metadata}")
                continue

            media_keywords = fetch_keywords_for_media(media_id)
            if any(keyword.lower() in [mk.lower() for mk in media_keywords] for keyword in keywords):
                filtered_results.append(result)
        except Exception as e:
            logging.error(f"Error processing result: {result}. Error: {str(e)}")

    filter_duration = time.time() - start_time
    log_histogram("filter_results_by_keywords_duration", filter_duration)
    log_counter("filter_results_by_keywords_success", labels={"filtered_count": len(filtered_results)})
    return filtered_results

# FIXME: to be implememted
def extract_media_id_from_result(result: str) -> Optional[int]:
    # Implement this function based on how you store the media_id in your results
    # For example, if it's stored at the beginning of each result:
    try:
        return int(result.split('_')[0])
    except (IndexError, ValueError):
        logging.error(f"Failed to extract media_id from result: {result}")
        return None

#
#
########################################################################################################################


############################################################################################################
#
# Chat RAG

def enhanced_rag_pipeline_chat(query: str, api_choice: str, character_id: int, keywords: Optional[str] = None) -> Dict[str, Any]:
    """

    Enhanced RAG pipeline tailored for the Character Chat tab.



    Args:

        query (str): The user's input query.

        api_choice (str): The API to use for generating the response.

        character_id (int): The ID of the character being interacted with.

        keywords (Optional[str]): Comma-separated keywords to filter search results.



    Returns:

        Dict[str, Any]: Contains the generated answer and the context used.

    """
    log_counter("enhanced_rag_pipeline_chat_attempt", labels={"api_choice": api_choice, "character_id": character_id})
    start_time = time.time()
    try:
        # Load embedding provider from config, or fallback to 'openai'
        embedding_provider = config.get('Embeddings', 'provider', fallback='openai')
        logging.debug(f"Using embedding provider: {embedding_provider}")

        # Process keywords if provided
        keyword_list = [k.strip().lower() for k in keywords.split(',')] if keywords else []
        logging.debug(f"enhanced_rag_pipeline_chat - Keywords: {keyword_list}")

        # Fetch relevant chat IDs based on character_id and keywords
        if keyword_list:
            relevant_chat_ids = fetch_keywords_for_chats(keyword_list)
        else:
            relevant_chat_ids = fetch_all_chat_ids(character_id)
        logging.debug(f"enhanced_rag_pipeline_chat - Relevant chat IDs: {relevant_chat_ids}")

        if not relevant_chat_ids:
            logging.info(f"No chats found for the given keywords and character ID: {character_id}")
            # Fallback to generating answer without context
            answer = generate_answer(api_choice, "", query)
            # Metrics
            pipeline_duration = time.time() - start_time
            log_histogram("enhanced_rag_pipeline_chat_duration", pipeline_duration, labels={"api_choice": api_choice})
            log_counter("enhanced_rag_pipeline_chat_success",
                        labels={"api_choice": api_choice, "character_id": character_id})
            return {
                "answer": answer,
                "context": ""
            }

        # Perform vector search within the relevant chats
        vector_results = perform_vector_search_chat(query, relevant_chat_ids)
        logging.debug(f"enhanced_rag_pipeline_chat - Vector search results: {vector_results}")

        # Perform full-text search within the relevant chats
        # FIXME - Update for DB Selection
        fts_results = perform_full_text_search_chat(query, relevant_chat_ids)
        logging.debug("enhanced_rag_pipeline_chat - Full-text search results:")
        logging.debug("\n".join([str(item) for item in fts_results]))

        # Combine results
        all_results = vector_results + fts_results

        apply_re_ranking = True
        if apply_re_ranking:
            logging.debug("enhanced_rag_pipeline_chat - Applying Re-Ranking")
            ranker = Ranker()

            # Prepare passages for re-ranking
            passages = [{"id": i, "text": result['content']} for i, result in enumerate(all_results)]
            rerank_request = RerankRequest(query=query, passages=passages)

            # Rerank the results
            reranked_results = ranker.rerank(rerank_request)

            # Sort results based on the re-ranking score
            reranked_results = sorted(reranked_results, key=lambda x: x['score'], reverse=True)

            # Log reranked results
            logging.debug(f"enhanced_rag_pipeline_chat - Reranked results: {reranked_results}")

            # Update all_results based on reranking
            all_results = [all_results[result['id']] for result in reranked_results]

        # Extract context from top results (limit to top 10)
        context = "\n".join([result['content'] for result in all_results[:10]])
        logging.debug(f"Context length: {len(context)}")
        logging.debug(f"Context: {context[:200]}")  # Log only the first 200 characters for brevity

        # Generate answer using the selected API
        answer = generate_answer(api_choice, context, query)

        if not all_results:
            logging.info(f"No results found. Query: {query}, Keywords: {keywords}")
            return {
                "answer": "No relevant information based on your query and keywords were found in the database. Your query has been directly passed to the LLM, and here is its answer: \n\n" + answer,
                "context": "No relevant information based on your query and keywords were found in the database. The only context used was your query: \n\n" + query
            }

        return {
            "answer": answer,
            "context": context
        }

    except Exception as e:
        log_counter("enhanced_rag_pipeline_chat_error", labels={"api_choice": api_choice, "character_id": character_id, "error": str(e)})
        logging.error(f"Error in enhanced_rag_pipeline_chat: {str(e)}")
        return {
            "answer": "An error occurred while processing your request.",
            "context": ""
        }


def fetch_relevant_chat_ids(character_id: int, keywords: List[str]) -> List[int]:
    """

    Fetch chat IDs associated with a character and filtered by keywords.



    Args:

        character_id (int): The ID of the character.

        keywords (List[str]): List of keywords to filter chats.



    Returns:

        List[int]: List of relevant chat IDs.

    """
    log_counter("fetch_relevant_chat_ids_attempt", labels={"character_id": character_id, "keyword_count": len(keywords)})
    start_time = time.time()
    relevant_ids = set()
    try:
        media_ids = fetch_keywords_for_chats(keywords)
        fetch_duration = time.time() - start_time
        log_histogram("fetch_relevant_chat_ids_duration", fetch_duration)
        log_counter("fetch_relevant_chat_ids_success",
                    labels={"character_id": character_id, "result_count": len(relevant_ids)})
        relevant_ids.update(media_ids)
        return list(relevant_ids)
    except Exception as e:
        log_counter("fetch_relevant_chat_ids_error", labels={"character_id": character_id, "error": str(e)})
        logging.error(f"Error fetching relevant chat IDs: {str(e)}")
        return []


def fetch_all_chat_ids(character_id: int) -> List[int]:
    """

    Fetch all chat IDs associated with a specific character.



    Args:

        character_id (int): The ID of the character.



    Returns:

        List[int]: List of all chat IDs for the character.

    """
    log_counter("fetch_all_chat_ids_attempt", labels={"character_id": character_id})
    start_time = time.time()
    try:
        chats = get_character_chats(character_id=character_id)
        chat_ids = [chat['id'] for chat in chats]
        fetch_duration = time.time() - start_time
        log_histogram("fetch_all_chat_ids_duration", fetch_duration)
        log_counter("fetch_all_chat_ids_success", labels={"character_id": character_id, "chat_count": len(chat_ids)})
        return chat_ids
    except Exception as e:
        log_counter("fetch_all_chat_ids_error", labels={"character_id": character_id, "error": str(e)})
        logging.error(f"Error fetching all chat IDs for character {character_id}: {str(e)}")
        return []

#
# End of Chat RAG
############################################################################################################

# Function to preprocess and store all existing content in the database
# def preprocess_all_content(database, create_contextualized=True, api_name="gpt-3.5-turbo"):
#     unprocessed_media = get_unprocessed_media()
#     total_media = len(unprocessed_media)
#
#     for index, row in enumerate(unprocessed_media, 1):
#         media_id, content, media_type, file_name = row
#         collection_name = f"{media_type}_{media_id}"
#
#         logger.info(f"Processing media {index} of {total_media}: ID {media_id}, Type {media_type}")
#
#         try:
#             process_and_store_content(
#                 database=database,
#                 content=content,
#                 collection_name=collection_name,
#                 media_id=media_id,
#                 file_name=file_name or f"{media_type}_{media_id}",
#                 create_embeddings=True,
#                 create_contextualized=create_contextualized,
#                 api_name=api_name
#             )
#
#             # Mark the media as processed in the database
#             mark_media_as_processed(database, media_id)
#
#             logger.info(f"Successfully processed media ID {media_id}")
#         except Exception as e:
#             logger.error(f"Error processing media ID {media_id}: {str(e)}")
#
#     logger.info("Finished preprocessing all unprocessed content")

############################################################################################################
#
# ElasticSearch Retriever

# https://github.com/langchain-ai/langchain/tree/44e3e2391c48bfd0a8e6a20adde0b6567f4f43c3/templates/rag-elasticsearch
#
# https://github.com/langchain-ai/langchain/tree/44e3e2391c48bfd0a8e6a20adde0b6567f4f43c3/templates/rag-self-query

#
# End of RAG_Library_2.py
############################################################################################################