Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.


TASK: user input : may one pdf or list(of pdf) in a folder , along  with query   , we should able retrieve most relevant data, with user number k with similarity searching with relevant like k: int, each Simulink for chuck click, we developing rag system for one pdf or list(of pdf) 


import os
import sys
import logging
import argparse
import shutil
import multiprocessing
from multiprocessing import Pool
from pathlib import Path
from typing import List, Tuple, Optional

import fitz  # PyMuPDF for PDF processing
from tqdm import tqdm
from transformers import AutoTokenizer, AutoModel
import torch
import faiss
import numpy as np

# Set up logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout)
    ]
)

logger = logging.getLogger(__name__)

# Constants
EMBEDDING_MODEL_NAME = 'sentence-transformers/all-MiniLM-L6-v2'  # Fast and lightweight model
CHUNK_SIZE = 512  # Adjust based on your needs
CPU_COUNT = multiprocessing.cpu_count()


def extract_text_from_pdf(pdf_path: Path) -> str:
    """
    Extract text from a PDF file using PyMuPDF.
    """
    try:
        with fitz.open(pdf_path) as doc:
            text = ""
            for page in doc:
                text += page.get_text()
        return text
    except Exception as e:
        logger.error(f"Error reading PDF file {pdf_path}: {e}")
        return ""


def extract_text_from_txt(txt_path: Path) -> str:
    """
    Extract text from a txt file.
    """
    try:
        with open(txt_path, 'r', encoding='utf-8') as f:
            return f.read()
    except Exception as e:
        logger.error(f"Error reading TXT file {txt_path}: {e}")
        return ""


def process_file(file_path: Path) -> Tuple[str, str]:
    """
    Reads a file and returns its text content along with the file name.
    """
    ext = file_path.suffix.lower()
    if ext == '.pdf':
        text = extract_text_from_pdf(file_path)
    elif ext == '.txt':
        text = extract_text_from_txt(file_path)
    else:
        logger.warning(f"Unsupported file type: {file_path}")
        return "", ""
    return text, str(file_path)


def chunk_text(text: str, chunk_size: int) -> List[str]:
    """
    Splits the text into chunks of specified size.
    """
    tokenizer = AutoTokenizer.from_pretrained(EMBEDDING_MODEL_NAME)
    tokens = tokenizer.tokenize(text)
    chunks = []
    for i in range(0, len(tokens), chunk_size):
        chunk_tokens = tokens[i:i + chunk_size]
        chunk_text = tokenizer.convert_tokens_to_string(chunk_tokens)
        chunks.append(chunk_text)
    return chunks


def embed_texts(texts: List[str], model_name: str) -> np.ndarray:
    """
    Generates embeddings for a list of texts using a specified model.
    """
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModel.from_pretrained(model_name).to(device)

    embeddings = []
    with torch.no_grad():
        for text in texts:
            inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True).to(device)
            outputs = model(**inputs)
            # Mean pooling
            embedding = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
            embeddings.append(embedding)
    embeddings = np.vstack(embeddings)
    return embeddings


def create_symlink(target: Path, link_name: Path) -> None:
    """
    Creates a symlink pointing to target named link_name.
    """
    try:
        link_name.parent.mkdir(parents=True, exist_ok=True)
        if link_name.exists() or link_name.is_symlink():
            link_name.unlink()
        os.symlink(target, link_name)
    except OSError as e:
        logger.error(f"Failed to create symlink {link_name} -> {target}: {e}")


def process_and_chunk_file(args: Tuple[Path, int]) -> List[Tuple[str, np.ndarray, str]]:
    """
    Processes a single file: extracts text, chunks it, and generates embeddings.
    Returns a list of tuples with (chunk_text, embedding, source_info).
    """
    file_path, chunk_size = args
    text, source = process_file(file_path)
    if not text:
        return []
    chunks = chunk_text(text, chunk_size)
    if not chunks:
        return []
    embeddings = embed_texts(chunks, EMBEDDING_MODEL_NAME)
    results = []
    for chunk_text, embedding in zip(chunks, embeddings):
        results.append((chunk_text, embedding, source))
    return results


def build_vector_database(data: List[Tuple[str, np.ndarray, str]]) -> faiss.IndexFlatL2:
    """
    Builds a FAISS vector database from the embeddings.
    Returns the FAISS index.
    """
    if not data:
        logger.error("No data provided to build vector database.")
        raise ValueError("Data list is empty.")
    embeddings = np.vstack([item[1] for item in data])
    index = faiss.IndexFlatL2(embeddings.shape[1])
    index.add(embeddings)
    return index


def save_vector_database(index: faiss.IndexFlatL2, path: Path) -> None:
    """
    Saves the FAISS index to the specified path.
    """
    try:
        faiss.write_index(index, str(path))
        logger.info(f"Vector database saved to {path}")
    except Exception as e:
        logger.error(f"Failed to save vector database: {e}")


def main(input_paths: List[Path], output_dir: Path) -> None:
    """
    Main processing function.
    """
    try:
        pool = Pool(processes=CPU_COUNT)
        args = [(path, CHUNK_SIZE) for path in input_paths]
        results = []

        logger.info("Starting file processing and embedding generation...")
        for res in tqdm(pool.imap_unordered(process_and_chunk_file, args), total=len(args), desc="Processing files"):
            if res:
                results.extend(res)
        pool.close()
        pool.join()

        if not results:
            logger.error("No results to process.")
            return

        # Build vector database
        logger.info("Building vector database...")
        index = build_vector_database(results)

        # Save vector database
        vector_db_path = output_dir / 'vector_database.index'
        vector_db_path.parent.mkdir(parents=True, exist_ok=True)
        save_vector_database(index, vector_db_path)

        # Save metadata and create symlinks
        logger.info("Saving metadata and creating symlinks...")
        metadata_path = output_dir / 'metadata.txt'
        symlink_dir = output_dir / 'symlinks'
        symlink_dir.mkdir(parents=True, exist_ok=True)

        with open(metadata_path, 'w', encoding='utf-8') as f:
            for i, (chunk_text, _, source) in enumerate(results):
                f.write(f"{i}\t{source}\n")
                # Save chunk text to file
                chunk_file = output_dir / f"chunk_{i}.txt"
                with open(chunk_file, 'w', encoding='utf-8') as cf:
                    cf.write(chunk_text)
                # Create symlink in symlink directory
                symlink_path = symlink_dir / f"chunk_{i}.txt"
                create_symlink(chunk_file, symlink_path)

        logger.info(f"Processing completed successfully. Outputs are saved in {output_dir}")

    except Exception as e:
        logger.error(f"An error occurred during processing: {e}")