ragmint 0.3.1__py3-none-any.whl

ragmint/tests/test_langchain_prebuilder.py

@@ -0,0 +1,82 @@
+import pytest
+from unittest.mock import MagicMock, patch
+from ragmint.integrations.langchain_prebuilder import LangchainPrebuilder
+
+
+@pytest.fixture
+def sample_docs():
+    """Small sample corpus for testing."""
+    return ["AI is transforming the world.", "RAG pipelines improve retrieval."]
+
+
+@pytest.fixture
+def sample_config():
+    """Default configuration for tests."""
+    return {
+        "retriever": "faiss",
+        "embedding_model": "sentence-transformers/all-MiniLM-L6-v2",
+        "chunk_size": 200,
+        "overlap": 50,
+    }
+
+
+@patch("ragmint.integrations.langchain_prebuilder.HuggingFaceEmbeddings", autospec=True)
+@patch("ragmint.integrations.langchain_prebuilder.RecursiveCharacterTextSplitter", autospec=True)
+def test_prepare_creates_components(mock_splitter, mock_embedder, sample_config, sample_docs):
+    """Ensure prepare() builds retriever and embedding components properly."""
+    mock_splitter.return_value.create_documents.return_value = ["doc1", "doc2"]
+    mock_embedder.return_value = MagicMock()
+
+    # Patch FAISS to avoid building a real index
+    with patch("ragmint.integrations.langchain_prebuilder.FAISS", autospec=True) as mock_faiss:
+        mock_db = MagicMock()
+        mock_faiss.from_documents.return_value = mock_db
+        mock_db.as_retriever.return_value = "mock_retriever"
+
+        builder = LangchainPrebuilder(sample_config)
+        retriever, embeddings = builder.prepare(sample_docs)
+
+        assert retriever == "mock_retriever"
+        assert embeddings == mock_embedder.return_value
+
+        mock_splitter.assert_called_once()
+        mock_embedder.assert_called_once_with(model_name="sentence-transformers/all-MiniLM-L6-v2")
+        mock_faiss.from_documents.assert_called_once()
+
+
+@pytest.mark.parametrize("backend", ["faiss", "chroma", "bm25"])
+def test_build_retriever_backends(sample_config, sample_docs, backend):
+    """Check retriever creation for each backend."""
+    cfg = dict(sample_config)
+    cfg["retriever"] = backend
+
+    builder = LangchainPrebuilder(cfg)
+
+    # Mock embeddings + docs
+    fake_embeddings = MagicMock()
+    fake_docs = ["d1", "d2"]
+
+    with patch("ragmint.integrations.langchain_prebuilder.FAISS.from_documents", return_value=MagicMock()) as mock_faiss, \
+         patch("ragmint.integrations.langchain_prebuilder.Chroma.from_documents", return_value=MagicMock()) as mock_chroma, \
+         patch("ragmint.integrations.langchain_prebuilder.BM25Retriever.from_texts", return_value=MagicMock()) as mock_bm25:
+        retriever = builder._build_retriever(fake_docs, fake_embeddings)
+
+        # Validate retriever creation per backend
+        if backend == "faiss":
+            mock_faiss.assert_called_once()
+        elif backend == "chroma":
+            mock_chroma.assert_called_once()
+        elif backend == "bm25":
+            mock_bm25.assert_called_once()
+
+        assert retriever is not None
+
+
+def test_invalid_backend_raises(sample_config):
+    """Ensure ValueError is raised for unsupported retriever."""
+    cfg = dict(sample_config)
+    cfg["retriever"] = "invalid"
+
+    builder = LangchainPrebuilder(cfg)
+    with pytest.raises(ValueError, match="Unsupported retriever backend"):
+        builder._build_retriever(["doc"], MagicMock())

ragmint/tests/test_leaderboard.py

@@ -0,0 +1,39 @@
+import json
+import tempfile
+from pathlib import Path
+from ragmint.leaderboard import Leaderboard
+
+
+def test_leaderboard_add_and_top(tmp_path):
+    """Ensure local leaderboard persistence works without Supabase."""
+    file_path = tmp_path / "leaderboard.jsonl"
+    lb = Leaderboard(storage_path=str(file_path))
+
+    # Add two runs
+    lb.upload("run1", {"retriever": "FAISS"}, 0.91)
+    lb.upload("run2", {"retriever": "Chroma"}, 0.85)
+
+    # Verify file content
+    assert file_path.exists()
+    with open(file_path, "r", encoding="utf-8") as f:
+        lines = [json.loads(line) for line in f]
+    assert len(lines) == 2
+
+    # Get top results
+    top = lb.top_results(limit=1)
+    assert isinstance(top, list)
+    assert len(top) == 1
+    assert "score" in top[0]
+
+
+def test_leaderboard_append_existing(tmp_path):
+    """Ensure multiple uploads append properly."""
+    file_path = tmp_path / "leaderboard.jsonl"
+    lb = Leaderboard(storage_path=str(file_path))
+
+    for i in range(3):
+        lb.upload(f"run{i}", {"retriever": "BM25"}, 0.8 + i * 0.05)
+
+    top = lb.top_results(limit=2)
+    assert len(top) == 2
+    assert top[0]["score"] >= top[1]["score"]

ragmint/tests/test_pipeline.py

@@ -0,0 +1,20 @@
+import numpy as np
+from ragmint.core.pipeline import RAGPipeline
+from ragmint.core.retriever import Retriever
+from ragmint.core.embeddings import Embeddings
+from ragmint.core.reranker import Reranker
+from ragmint.core.evaluation import Evaluator
+
+
+def test_pipeline_run():
+    docs = ["doc1 text", "doc2 text"]
+    embedder = Embeddings(backend="dummy")
+    retriever = Retriever(embedder=embedder, documents=docs)
+    reranker = Reranker("mmr")
+    evaluator = Evaluator()
+    pipeline = RAGPipeline(retriever, reranker, evaluator)
+
+    result = pipeline.run("what is doc1?")
+    assert "query" in result
+    assert "answer" in result
+    assert "metrics" in result

ragmint/tests/test_retriever.py

@@ -0,0 +1,15 @@
+import numpy as np
+from ragmint.core.retriever import Retriever
+from ragmint.core.embeddings import Embeddings
+
+
+def test_retrieve_basic():
+    docs = ["doc A", "doc B", "doc C"]
+    embedder = Embeddings(backend="dummy")
+    retriever = Retriever(embedder=embedder, documents=docs)
+
+    results = retriever.retrieve("sample query", top_k=2)
+    assert isinstance(results, list)
+    assert len(results) == 2
+    assert "text" in results[0]
+    assert "score" in results[0]

ragmint/tests/test_search.py

@@ -0,0 +1,17 @@
+from ragmint.optimization.search import GridSearch, RandomSearch
+
+
+def test_grid_search_iterates():
+    space = {"retriever": ["faiss"], "embedding_model": ["openai"], "reranker": ["mmr"]}
+    search = GridSearch(space)
+    combos = list(search)
+    assert len(combos) == 1
+    assert "retriever" in combos[0]
+
+
+def test_random_search_n_trials():
+    space = {"retriever": ["faiss", "bm25"], "embedding_model": ["openai", "st"], "reranker": ["mmr"]}
+    search = RandomSearch(space, n_trials=5)
+    combos = list(search)
+    assert len(combos) == 5
+    assert all("retriever" in c for c in combos)

ragmint/tests/test_tuner.py

@@ -0,0 +1,71 @@
+import os
+import json
+import pytest
+from ragmint.tuner import RAGMint
+
+
+def setup_validation_file(tmp_path):
+    """Create a temporary validation QA dataset."""
+    data = [
+        {"question": "What is AI?", "answer": "Artificial Intelligence"},
+        {"question": "Define ML", "answer": "Machine Learning"}
+    ]
+    file = tmp_path / "validation_qa.json"
+    with open(file, "w", encoding="utf-8") as f:
+        json.dump(data, f)
+    return str(file)
+
+
+def setup_docs(tmp_path):
+    """Create a small document corpus for testing."""
+    corpus = tmp_path / "corpus"
+    corpus.mkdir()
+    (corpus / "doc1.txt").write_text("This is about Artificial Intelligence.")
+    (corpus / "doc2.txt").write_text("This text explains Machine Learning.")
+    return str(corpus)
+
+
+@pytest.mark.parametrize("validation_mode", [
+    None,  # Built-in dataset
+    "data/custom_eval.json",  # Custom dataset path (mocked below)
+])
+def test_optimize_ragmint(tmp_path, validation_mode, monkeypatch):
+    """Test RAGMint.optimize() with different dataset modes."""
+    docs_path = setup_docs(tmp_path)
+    val_file = setup_validation_file(tmp_path)
+
+    # If using custom dataset, mock the path
+    if validation_mode and "custom_eval" in validation_mode:
+        custom_path = tmp_path / "custom_eval.json"
+        os.rename(val_file, custom_path)
+        validation_mode = str(custom_path)
+
+    metric = "faithfulness"
+
+    # Initialize RAGMint
+    rag = RAGMint(
+        docs_path=docs_path,
+        retrievers=["faiss"],
+        embeddings=["all-MiniLM-L6-v2"],
+        rerankers=["mmr"]
+    )
+
+    # Run optimization
+    best, results = rag.optimize(
+        validation_set=validation_mode,
+        metric=metric,
+        trials=2
+    )
+
+    # Validate results
+    assert isinstance(best, dict), "Best config should be a dict"
+    assert isinstance(results, list), "Results should be a list of trials"
+    assert len(results) > 0, "Optimization should produce results"
+
+    # The best result can expose either 'score' or the metric name (e.g. 'faithfulness')
+    assert any(k in best for k in ("score", metric)), \
+        f"Best config should include either 'score' or '{metric}'"
+
+    # Ensure the metric value is valid
+    assert best.get(metric, best.get("score")) >= 0, \
+        f"{metric} score should be non-negative"

ragmint/tuner.py

@@ -0,0 +1,189 @@
+import os
+import json
+import logging
+from typing import Any, Dict, List, Tuple
+from time import perf_counter
+
+from .core.pipeline import RAGPipeline
+from .core.embeddings import Embeddings
+from .core.retriever import Retriever
+from .core.reranker import Reranker
+from .core.evaluation import Evaluator
+from .optimization.search import GridSearch, RandomSearch, BayesianSearch
+from .utils.data_loader import load_validation_set
+
+logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
+
+
+class RAGMint:
+    """
+    Main RAG pipeline optimizer and evaluator.
+    Runs combinations of retrievers, embeddings, rerankers, and chunking parameters
+    to find the best performing RAG configuration.
+    """
+
+    def __init__(
+        self,
+        docs_path: str,
+        retrievers: List[str],
+        embeddings: List[str],
+        rerankers: List[str],
+        chunk_sizes: List[int] = [400, 600],
+        overlaps: List[int] = [50, 100],
+        strategies: List[str] = ["fixed"],
+    ):
+        self.docs_path = docs_path
+        self.retrievers = retrievers
+        self.embeddings = embeddings
+        self.rerankers = rerankers
+        self.chunk_sizes = chunk_sizes
+        self.overlaps = overlaps
+        self.strategies = strategies
+
+        self.documents: List[str] = self._load_docs()
+        self.embeddings_cache: Dict[str, Any] = {}
+
+    # -------------------------
+    # Document Loading
+    # -------------------------
+    def _load_docs(self) -> List[str]:
+        if not os.path.exists(self.docs_path):
+            logging.warning(f"Corpus path not found: {self.docs_path}")
+            return []
+
+        docs = []
+        for file in os.listdir(self.docs_path):
+            if file.endswith((".txt", ".md", ".rst")):
+                with open(os.path.join(self.docs_path, file), "r", encoding="utf-8") as f:
+                    docs.append(f.read())
+
+        logging.info(f"📚 Loaded {len(docs)} documents from {self.docs_path}")
+        return docs
+
+    # -------------------------
+    # Embedding Cache
+    # -------------------------
+    def _embed_docs(self, model_name: str) -> Any:
+        """Compute and cache document embeddings."""
+        if model_name in self.embeddings_cache:
+            return self.embeddings_cache[model_name]
+
+        model = Embeddings(backend="huggingface", model_name=model_name)
+        embeddings = model.encode(self.documents)
+        self.embeddings_cache[model_name] = embeddings
+        return embeddings
+
+    # -------------------------
+    # Build Pipeline
+    # -------------------------
+    def _build_pipeline(self, config: Dict[str, str]) -> RAGPipeline:
+        """Builds a pipeline from one configuration."""
+        retriever_backend = config["retriever"]
+        model_name = config["embedding_model"]
+        reranker_name = config["reranker"]
+
+        # Chunking params (use defaults if missing)
+        chunk_size = int(config.get("chunk_size", 500))
+        overlap = int(config.get("overlap", 100))
+        strategy = config.get("strategy", "fixed")
+
+        # Load embeddings (cached)
+        embeddings = self._embed_docs(model_name)
+        embedder = Embeddings(backend="huggingface", model_name=model_name)
+
+        # Initialize retriever with backend
+        logging.info(f"⚙️ Initializing retriever backend: {retriever_backend}")
+        retriever = Retriever(
+            embedder=embedder,
+            documents=self.documents,
+            embeddings=embeddings,
+            backend=retriever_backend,
+        )
+
+        reranker = Reranker(reranker_name)
+        evaluator = Evaluator()
+
+        # ✅ Pass chunking parameters into RAGPipeline
+        return RAGPipeline(
+            retriever,
+            reranker,
+            evaluator,
+            chunk_size=chunk_size,
+            overlap=overlap,
+            chunking_strategy=strategy,
+        )
+
+    # -------------------------
+    # Evaluate Configuration
+    # -------------------------
+    def _evaluate_config(
+        self, config: Dict[str, Any], validation: List[Dict[str, str]], metric: str
+    ) -> Dict[str, float]:
+        """Evaluates a single configuration."""
+        pipeline = self._build_pipeline(config)
+        scores = []
+        start = perf_counter()
+
+        for sample in validation:
+            query = sample.get("question") or sample.get("query") or ""
+            result = pipeline.run(query)
+            score = result["metrics"].get(metric, 0.0)
+            scores.append(score)
+
+        elapsed = perf_counter() - start
+        avg_score = sum(scores) / len(scores) if scores else 0.0
+
+        return {
+            metric: avg_score,
+            "latency": elapsed / max(1, len(validation)),
+        }
+
+    # -------------------------
+    # Optimize
+    # -------------------------
+    def optimize(
+        self,
+        validation_set: str,
+        metric: str = "faithfulness",
+        search_type: str = "random",
+        trials: int = 10,
+    ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]:
+        """Run optimization search over retrievers, embeddings, rerankers, and chunking."""
+        validation = load_validation_set(validation_set or "default")
+
+        # ✅ Add chunking parameters to the search space
+        search_space = {
+            "retriever": self.retrievers,
+            "embedding_model": self.embeddings,
+            "reranker": self.rerankers,
+            "chunk_size": self.chunk_sizes,
+            "overlap": self.overlaps,
+            "strategy": self.strategies,
+        }
+
+        logging.info(f"🚀 Starting {search_type} optimization with {trials} trials")
+
+        # Select search strategy
+        try:
+            if search_type == "grid":
+                searcher = GridSearch(search_space)
+            elif search_type == "bayesian":
+                searcher = BayesianSearch(search_space)
+            else:
+                searcher = RandomSearch(search_space, n_trials=trials)
+        except Exception as e:
+            logging.warning(f"⚠️ Fallback to RandomSearch due to missing deps: {e}")
+            searcher = RandomSearch(search_space, n_trials=trials)
+
+        # Run trials
+        results = []
+        for config in searcher:
+            metrics = self._evaluate_config(config, validation, metric)
+            result = {**config, **metrics}
+            results.append(result)
+            logging.info(f"🔹 Tested config: {config} -> {metrics}")
+
+        best = max(results, key=lambda r: r.get(metric, 0.0)) if results else {}
+        logging.info(f"🏆 Best configuration: {best}")
+
+        return best, results

ragmint/utils/caching.py

@@ -0,0 +1,37 @@
+import os
+import json
+import hashlib
+import pickle
+from typing import Any
+
+
+class Cache:
+    """
+    Simple file-based cache for embeddings or retrievals.
+    """
+
+    def __init__(self, cache_dir: str = ".ragmint_cache"):
+        self.cache_dir = cache_dir
+        os.makedirs(cache_dir, exist_ok=True)
+
+    def _hash_key(self, key: str) -> str:
+        return hashlib.md5(key.encode()).hexdigest()
+
+    def exists(self, key: str) -> bool:
+        return os.path.exists(os.path.join(self.cache_dir, self._hash_key(key)))
+
+    def get(self, key: str) -> Any:
+        path = os.path.join(self.cache_dir, self._hash_key(key))
+        if not os.path.exists(path):
+            return None
+        with open(path, "rb") as f:
+            return pickle.load(f)
+
+    def set(self, key: str, value: Any):
+        path = os.path.join(self.cache_dir, self._hash_key(key))
+        with open(path, "wb") as f:
+            pickle.dump(value, f)
+
+    def clear(self):
+        for file in os.listdir(self.cache_dir):
+            os.remove(os.path.join(self.cache_dir, file))

ragmint/utils/data_loader.py

@@ -0,0 +1,65 @@
+import json
+import csv
+from typing import List, Dict
+from pathlib import Path
+import os
+
+try:
+    from datasets import load_dataset
+except ImportError:
+    load_dataset = None  # optional dependency
+
+DEFAULT_VALIDATION_PATH = Path(__file__).parent.parent / "experiments" / "validation_qa.json"
+
+
+def load_json(path: str) -> List[Dict]:
+    with open(path, "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def load_csv(path: str) -> List[Dict]:
+    with open(path, newline="", encoding="utf-8") as csvfile:
+        reader = csv.DictReader(csvfile)
+        return list(reader)
+
+
+def save_json(path: str, data: Dict):
+    with open(path, "w", encoding="utf-8") as f:
+        json.dump(data, f, ensure_ascii=False, indent=2)
+
+def load_validation_set(path: str | None = None) -> List[Dict]:
+    """
+    Loads a validation dataset (QA pairs) from:
+    - Built-in default JSON file
+    - User-provided JSON or CSV
+    - Hugging Face dataset by name
+    """
+    # Default behavior
+    if path is None or path == "default":
+        if not DEFAULT_VALIDATION_PATH.exists():
+            raise FileNotFoundError(f"Default validation set not found at {DEFAULT_VALIDATION_PATH}")
+        return load_json(DEFAULT_VALIDATION_PATH)
+
+    # Hugging Face dataset
+    if not os.path.exists(path) and load_dataset:
+        try:
+            dataset = load_dataset(path, split="validation")
+            data = [
+                {"question": q, "answer": a}
+                for q, a in zip(dataset["question"], dataset["answers"])
+            ]
+            return data
+        except Exception:
+            pass  # fall through to file loading
+
+    # Local file
+    p = Path(path)
+    if not p.exists():
+        raise FileNotFoundError(f"Validation file not found: {path}")
+
+    if p.suffix.lower() == ".json":
+        return load_json(path)
+    elif p.suffix.lower() in [".csv", ".tsv"]:
+        return load_csv(path)
+    else:
+        raise ValueError("Unsupported validation set format. Use JSON, CSV, or a Hugging Face dataset name.")

ragmint/utils/logger.py

@@ -0,0 +1,36 @@
+import logging
+from tqdm import tqdm
+
+
+class Logger:
+    """
+    Centralized logger with optional tqdm integration and color formatting.
+    """
+
+    def __init__(self, name: str = "ragmint", level: int = logging.INFO):
+        self.logger = logging.getLogger(name)
+        self.logger.setLevel(level)
+
+        if not self.logger.handlers:
+            handler = logging.StreamHandler()
+            formatter = logging.Formatter(
+                "\033[96m[%(asctime)s]\033[0m \033[93m%(levelname)s\033[0m: %(message)s",
+                "%H:%M:%S",
+            )
+            handler.setFormatter(formatter)
+            self.logger.addHandler(handler)
+
+    def info(self, msg: str):
+        self.logger.info(msg)
+
+    def warning(self, msg: str):
+        self.logger.warning(msg)
+
+    def error(self, msg: str):
+        self.logger.error(msg)
+
+    def progress(self, iterable, desc="Processing", total=None):
+        return tqdm(iterable, desc=desc, total=total)
+
+def get_logger(name: str = "ragmint") -> Logger:
+    return Logger(name)

ragmint/utils/metrics.py

@@ -0,0 +1,27 @@
+from typing import List
+import numpy as np
+from difflib import SequenceMatcher
+
+
+def bleu_score(reference: str, candidate: str) -> float:
+    """
+    Simple BLEU-like precision approximation.
+    """
+    ref_tokens = reference.split()
+    cand_tokens = candidate.split()
+    if not cand_tokens:
+        return 0.0
+
+    matches = sum(1 for token in cand_tokens if token in ref_tokens)
+    return matches / len(cand_tokens)
+
+
+def rouge_l(reference: str, candidate: str) -> float:
+    """
+    Approximation of ROUGE-L using sequence matcher ratio.
+    """
+    return SequenceMatcher(None, reference, candidate).ratio()
+
+
+def mean_score(scores: List[float]) -> float:
+    return float(np.mean(scores)) if scores else 0.0

ragmint-0.3.1.data/data/LICENSE

@@ -0,0 +1,19 @@
+Apache License
+Version 2.0, January 2004
+http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+Copyright 2025 André Oliveira
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.