rag_func.py

"""
## Sources
1.  Connecting to a Pinecone index: https://docs.pinecone.io/guides/inference/generate-embeddings
2.  Creating a text loader from a directory: https://github.com/langchain-ai/langchain/discussions/18559
3.  Using an LLM: https://python.langchain.com/docs/tutorials/rag/
4.  Evaluation with RAGAS: https://docs.ragas.io/en/stable/getstarted/evals/#analyzing-results
5.  Dataset creation: https://www.mongodb.com/developer/products/atlas/evaluate-llm-applications-rag/

## Preparation
First, the packages required for handling connections to the LLM and the vector database must be installed.

Subsequently, the corresponding packages are imported along with the `os` module. This module is utilized to handle the assignments to the environment variables for the API keys."""

from pinecone import Pinecone
import os

from langchain_mistralai import ChatMistralAI
from langchain_openai import ChatOpenAI

from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.messages import HumanMessage, SystemMessage
import pathlib
from langchain.schema import Document
from langchain_community.document_loaders import (
    CSVLoader, PyPDFLoader, UnstructuredWordDocumentLoader,
    UnstructuredPowerPointLoader, UnstructuredMarkdownLoader,
    UnstructuredHTMLLoader, NotebookLoader
)

# -------------------------
# UTF-8 safe Text Loader
# -------------------------
class SafeTextLoader:
    """Loads a text file as a single Document, safely handling UTF-8 decoding errors."""
    def __init__(self, file_path):
        self.file_path = file_path

    def load(self):
        try:
            with open(self.file_path, "rb") as f:  # open in binary mode
                raw_bytes = f.read()
            text = raw_bytes.decode("utf-8", errors="ignore")  # decode safely
            return [Document(page_content=text, metadata={"source": str(self.file_path)})]
        except Exception as e:
            print(f"[Error] Failed to read {self.file_path}: {e}")
            return []

# -------------------------
# Loader mapping
# -------------------------
LOADER_MAPPING = {
    # Text
    ".txt": SafeTextLoader,
    ".json": SafeTextLoader,
    ".md": UnstructuredMarkdownLoader,
    ".csv": CSVLoader,
    ".yaml": SafeTextLoader,
    ".yml": SafeTextLoader,

    # Documents
    ".pdf": PyPDFLoader,
    ".docx": UnstructuredWordDocumentLoader,
    ".pptx": UnstructuredPowerPointLoader,
    ".html": UnstructuredHTMLLoader,
    ".htm": UnstructuredHTMLLoader,

    # Code / Notebook
    ".ipynb": NotebookLoader,
    ".py": SafeTextLoader,
    ".js": SafeTextLoader,
    ".sql": SafeTextLoader,
}

import pathlib

CONTEXT_ROOT = pathlib.Path(__file__).parent / "context"

# -------------------------
# Directory loader (recursive)
# -------------------------
def create_text_dir_loader(directory_path: str = ""):
    """
    Loads all supported files in the context directory (and subfolders) using the appropriate loader.
    - If directory_path is empty -> scans entire 'context' folder recursively.
    - If directory_path is given -> scans only that subfolder inside 'context'.
    """
    # Resolve the target directory
    target_dir = CONTEXT_ROOT / directory_path if directory_path else CONTEXT_ROOT

    if not target_dir.exists() or not target_dir.is_dir():
        print(f"[Error] Target directory does not exist: {target_dir}")
        return []

    documents = []

    for file_path in target_dir.rglob("*"):  # recursive
        if not file_path.is_file():
            continue

        ext = file_path.suffix.lower()
        loader_cls = LOADER_MAPPING.get(ext)

        if loader_cls is None:
            print(f"[Skip] Unsupported file type: {file_path}")
            continue

        try:
            loader = loader_cls(str(file_path))
            docs = loader.load()
            documents.extend(docs)
            print(f"[Loaded] {file_path} ({len(docs)} docs)")
        except Exception as e:
            print(f"[Error] Failed to load {file_path}: {e}")

    print(f"[Done] Finished scanning {target_dir}")
    return documents

# -------------------------
# Dataset creation
# -------------------------
def create_dataset(directory_path: str = "context"):
    """
    Loads all supported files from the given directory (recursively).
    Defaults to 'context' if no path is given.
    """
    target_dir = pathlib.Path(directory_path).resolve()

    if not target_dir.exists() or not target_dir.is_dir():
        print(f"[Error] Target directory does not exist: {target_dir}")
        return []

    documents = []

    for file_path in target_dir.rglob("*"):  # recursive
        if not file_path.is_file():
            continue

        ext = file_path.suffix.lower()
        loader_cls = LOADER_MAPPING.get(ext)

        if loader_cls is None:
            print(f"[Skip] Unsupported file type: {file_path}")
            continue

        try:
            loader = loader_cls(str(file_path))
            docs = loader.load()
            documents.extend(docs)
            print(f"[Loaded] {file_path} ({len(docs)} docs)")
        except Exception as e:
            print(f"[Error] Failed to load {file_path}: {e}")

    print(f"[Done] Finished scanning {target_dir}")
    print(f"Total documents loaded: {len(documents)}")
    return documents

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

import time
import os
import math


def prepare_RAG(
    pinecone_API,
    index_name,
    chunk_size=1800,
    chunk_overlap=200,
    llm_model="gpt-5-nano",
    dir_name="context",
    info=True
):
    import time
    from langchain_openai import ChatOpenAI
    from langchain_mistralai import ChatMistralAI
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    from pinecone import Pinecone

    if info:
        print(f"Prepare RAG with model={llm_model}, dir={dir_name}")
    
    if "gpt" in llm_model:
        llm = ChatOpenAI(model=llm_model, streaming=True)   # streaming enabled
    else:
        llm = ChatMistralAI(model=llm_model, streaming=True) 


    documents = create_dataset(dir_name)
    if not documents:
        print(f"[Warning] No documents found in directory '{dir_name}'. Using existing Pinecone index without upserting new vectors.")
        pc = Pinecone(api_key=pinecone_API)
        index = pc.Index(index_name)
        return index, pc, llm

    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    all_splits = text_splitter.split_documents(documents)
    if not all_splits:
        print(f"[Warning] No text chunks were created. Using existing Pinecone index.")
        pc = Pinecone(api_key=pinecone_API)
        index = pc.Index(index_name)
        return index, pc, llm

    if info:
        print("Number of chunks:", len(all_splits))

    pc = Pinecone(api_key=pinecone_API)
    index = pc.Index(index_name)

    initial_batch_size = 94
    min_batch_size = 1
    total_chunks = len(all_splits)
    records = []

    def retry_forever(func, *args, **kwargs):
        attempt = 1
        while True:
            try:
                return func(*args, **kwargs)
            except Exception as e:
                wait_time = min(60, 2 ** min(attempt, 6))
                print(f"[Retry] {func.__name__} failed (attempt {attempt}): {e}")
                print(f"Sleeping {wait_time} seconds before retry...")
                time.sleep(wait_time)
                attempt += 1

    start_idx = 0
    batch_size = initial_batch_size

    while start_idx < total_chunks:
        end_idx = min(start_idx + batch_size, total_chunks)
        batch_splits = all_splits[start_idx:end_idx]

        if info:
            print(f"Embedding batch {start_idx} → {end_idx - 1} (batch_size={batch_size})")

        data = [
            {"id": f"vec{idx}", "text": chunk.page_content, "metadata": chunk.metadata or {}}
            for idx, chunk in enumerate(batch_splits, start=start_idx)
        ]

        # Attempt embedding
        embeddings = retry_forever(
            pc.inference.embed,
            model="llama-text-embed-v2",
            inputs=[d['text'] for d in data],
            parameters={"input_type": "passage", "truncate": "END"}
        )

        # Prepare records for upsert
        batch_records = [
            {"id": d['id'], "values": e['values'], "metadata": {"text": d['text'], **d.get('metadata', {})}}
            for d, e in zip(data, embeddings)
        ]

        # Dynamic upsert with batch size adjustment
        while batch_records:
            try:
                retry_forever(index.upsert, vectors=batch_records, namespace="example-namespace")
                records.extend(batch_records)
                break  # success → exit loop
            except Exception as e:
                if "Request size" in str(e) or "exceeds the maximum" in str(e):
                    if len(batch_records) == 1:
                        print(f"[Error] Single vector too large to upsert: {batch_records[0]['id']}")
                        break
                    # Reduce batch size by half
                    batch_records = batch_records[:len(batch_records) // 2]
                    print(f"[Warning] Batch too large. Reducing batch size to {len(batch_records)} and retrying upsert...")
                    time.sleep(1)  # small delay before retry
                else:
                    raise e

        start_idx += batch_size  # move to next batch
        # Optionally increase batch size gradually if previous upsert succeeded
        batch_size = min(initial_batch_size, batch_size * 2)

    if info:
        print(f"Completed upsert of {len(records)} vectors.")

    return index, pc, llm


def retrieve_RAG(prompt_message, pc, index, top_k=5, info=True):
    if info:
        print("Retrieve RAG with", prompt_message, pc, index, top_k)
    """## Retrieval
    The user query is embedded using the same model (specified in the `model` parameter of the `pc.inference.embed()` function) as for embedding the chunks originally upserted into the vector database.
    """

    query_embedding = pc.inference.embed(
        model="llama-text-embed-v2",
        inputs=[prompt_message],
        parameters={
            "input_type": "query"
        }
    )

    """The relevant chunks are retrieved using semantic search with the cosine distance similarity measure. The number of chunks to be retrieved is passed in the `top_k` variable."""

    retrieved_chunks_raw = index.query(
        namespace="example-namespace",
        vector=query_embedding[0].values,
        top_k=top_k,
        include_values=False,
        include_metadata=True
    )

    """The result of the retrieval is processed, so that the variable `retrieved_chunks` contains a list of the corresponding text splits.
    """

    retrieved_chunks = []
    for match in retrieved_chunks_raw.matches:
        retrieved_chunks.append({
            "text": match.metadata.get("text", ""),
            "source": match.metadata.get("source", "")
        })

    return retrieved_chunks

def generate_RAG(prompt_message, llm, retrieved_chunks, info=True):
    if info:
        print("Generate RAG with", prompt_message, llm)
    """## Generation
    The prompt is sent to the LLM together with the context consisting of the data considered most relevant according to the semantic search. These data are stored in the variable `retrieved_chunks`. The part of the prompt containing the explaination of the task to be performed on the provided context is in the variable `prompt_message`.
    """

    context_message = "You are an expert in the field of the documents that are loaded. Use only the these documents and source files as the context. If you don't know the answer or the answer is not included in the context, then do not answer."
    prompt = [HumanMessage(prompt_message),
            SystemMessage(context_message),
                HumanMessage(str(retrieved_chunks))]

    """The prompt is sent to the selected LLM using the `invoke()` function. As a return value, this function passes the respective answer of the LLM."""

    response = llm.invoke(prompt)
    return response