ragmint 0.2.1__py3-none-any.whl → 0.4.6__py3-none-any.whl

ragmint/tests/test_leaderboard.py

@@ -1,39 +1,92 @@
+import os
 import json
 import tempfile
-from pathlib import Path
+import pytest
+from datetime import datetime
 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))
+@pytest.fixture
+def temp_leaderboard():
+    """Create a temporary leaderboard file for testing."""
+    with tempfile.TemporaryDirectory() as tmpdir:
+        path = os.path.join(tmpdir, "leaderboard.jsonl")
+        lb = Leaderboard(storage_path=path)
+        yield lb, 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
+def test_upload_and_persistence(temp_leaderboard):
+    lb, path = temp_leaderboard
+
+    # --- Mock experiment data ---
+    run_id = "run_001"
+    best_config = {"retriever": "FAISS", "embedding_model": "all-MiniLM"}
+    best_score = 0.92
+    all_results = [
+        {"retriever": "FAISS", "score": 0.92},
+        {"retriever": "BM25", "score": 0.85},
+    ]
+    documents = ["docs/a.txt", "docs/b.txt"]
+    model = "gemini"
+    corpus_stats = {"size": 20000, "avg_len": 400, "num_docs": 10}
+
+    # --- Upload ---
+    record = lb.upload(
+        run_id=run_id,
+        best_config=best_config,
+        best_score=best_score,
+        all_results=all_results,
+        documents=documents,
+        model=model,
+        corpus_stats=corpus_stats,
+    )
 
-    # Get top results
-    top = lb.top_results(limit=1)
-    assert isinstance(top, list)
-    assert len(top) == 1
-    assert "score" in top[0]
+    # --- Validate returned record ---
+    assert record["run_id"] == run_id
+    assert record["model"] == "gemini"
+    assert "timestamp" in record
+    assert record["best_score"] == 0.92
+    assert all(doc in record["documents"] for doc in ["a.txt", "b.txt"])
 
+    # --- File should contain JSON line ---
+    with open(path, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    assert len(lines) == 1
+    parsed = json.loads(lines[0])
+    assert parsed["run_id"] == run_id
 
-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)
+def test_top_results_ordering(temp_leaderboard):
+    lb, _ = temp_leaderboard
 
+    # Upload multiple runs with varying scores
+    for i, score in enumerate([0.8, 0.95, 0.7]):
+        lb.upload(
+            run_id=f"run_{i}",
+            best_config={"retriever": "FAISS"},
+            best_score=score,
+            all_results=[],
+            documents=["file.txt"],
+            model="claude",
+        )
+
+    # --- Get top results ---
     top = lb.top_results(limit=2)
     assert len(top) == 2
-    assert top[0]["score"] >= top[1]["score"]
+
+    # --- Ensure ordering descending by score ---
+    assert top[0]["best_score"] >= top[1]["best_score"]
+    assert top[0]["best_score"] == 0.95
+
+
+def test_all_results_reads_all_entries(temp_leaderboard):
+    lb, _ = temp_leaderboard
+
+    # Add two runs
+    lb.upload("run_a", {}, 0.5, [], ["doc1.txt"], "gemini")
+    lb.upload("run_b", {}, 0.7, [], ["doc2.txt"], "claude")
+
+    results = lb.all_results()
+    assert len(results) == 2
+    run_ids = {r["run_id"] for r in results}
+    assert {"run_a", "run_b"} <= run_ids

ragmint/tests/test_pipeline.py

@@ -1,14 +1,15 @@
 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"]
-    embeddings = [np.random.rand(4) for _ in range(2)]
-    retriever = Retriever(embeddings, docs)
+    embedder = Embeddings(backend="dummy")
+    retriever = Retriever(embedder=embedder, documents=docs)
     reranker = Reranker("mmr")
     evaluator = Evaluator()
     pipeline = RAGPipeline(retriever, reranker, evaluator)

ragmint/tests/test_qa_generator.py

@@ -0,0 +1,66 @@
+import os
+import json
+import tempfile
+from pathlib import Path
+import pytest
+
+from ragmint.qa_generator import generate_validation_qa
+
+
+class DummyLLM:
+    """Mock LLM that returns predictable JSON output."""
+    def generate_content(self, prompt):
+        class DummyResponse:
+            text = '[{"query": "What is X?", "expected_answer": "Y"}]'
+        return DummyResponse()
+
+
+@pytest.fixture
+def dummy_docs(tmp_path):
+    docs_dir = tmp_path / "docs"
+    docs_dir.mkdir()
+    for i in range(3):
+        (docs_dir / f"doc_{i}.txt").write_text(f"This is test document number {i}. It contains some content.")
+    return docs_dir
+
+
+@pytest.fixture
+def output_path(tmp_path):
+    return tmp_path / "validation_qa.json"
+
+
+def test_generate_validation_qa(monkeypatch, dummy_docs, output_path):
+    """Ensure QA generator runs end-to-end with mocked LLM."""
+    # --- Mock LLM setup ---
+    from sentence_transformers import SentenceTransformer
+    monkeypatch.setattr("ragmint.qa_generator.setup_llm", lambda *_: (DummyLLM(), "gemini"))
+    monkeypatch.setattr(SentenceTransformer, "encode", lambda self, x, normalize_embeddings=True: [[0.1] * 3] * len(x))
+
+    # --- Run function ---
+    generate_validation_qa(
+        docs_path=dummy_docs,
+        output_path=output_path,
+        llm_model="gemini-2.5-flash-lite",
+        batch_size=2,
+        sleep_between_batches=0,
+    )
+
+    # --- Validate output ---
+    assert output_path.exists(), "Output JSON file should be created"
+    data = json.loads(output_path.read_text())
+    assert isinstance(data, list), "Output must be a list"
+    assert all("query" in d and "expected_answer" in d for d in data), "Each entry must have query and answer"
+    assert len(data) > 0, "At least one QA pair should be generated"
+
+
+def test_handles_empty_folder(monkeypatch, tmp_path):
+    """Ensure no crash when docs folder is empty."""
+    empty_dir = tmp_path / "empty"
+    empty_dir.mkdir()
+    output_file = tmp_path / "qa.json"
+
+    monkeypatch.setattr("ragmint.qa_generator.setup_llm", lambda *_: (DummyLLM(), "gemini"))
+
+    generate_validation_qa(docs_path=empty_dir, output_path=output_file, sleep_between_batches=0)
+    data = json.loads(output_file.read_text())
+    assert data == [], "Empty folder should produce empty QA list"

ragmint/tests/test_retriever.py

@@ -1,11 +1,12 @@
 import numpy as np
 from ragmint.core.retriever import Retriever
+from ragmint.core.embeddings import Embeddings
 
 
 def test_retrieve_basic():
-    embeddings = [np.random.rand(5) for _ in range(3)]
     docs = ["doc A", "doc B", "doc C"]
-    retriever = Retriever(embeddings, docs)
+    embedder = Embeddings(backend="dummy")
+    retriever = Retriever(embedder=embedder, documents=docs)
 
     results = retriever.retrieve("sample query", top_k=2)
     assert isinstance(results, list)

ragmint/tests/test_tuner.py

@@ -46,7 +46,7 @@ def test_optimize_ragmint(tmp_path, validation_mode, monkeypatch):
     rag = RAGMint(
         docs_path=docs_path,
         retrievers=["faiss"],
-        embeddings=["text-embedding-3-small"],
+        embeddings=["all-MiniLM-L6-v2"],
         rerankers=["mmr"]
     )
 

ragmint/tuner.py

@@ -1,17 +1,17 @@
 import os
-import json
 import logging
-from typing import Any, Dict, List, Tuple, Optional
+from typing import Any, Dict, List, Tuple
 from time import perf_counter
 
 from .core.pipeline import RAGPipeline
-from .core.embeddings import EmbeddingModel
+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
+from .leaderboard import Leaderboard
+from uuid import uuid4
 
 logging.basicConfig(level=logging.INFO, format="[%(levelname)s] %(message)s")
 
@@ -19,6 +19,8 @@ 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__(
@@ -27,62 +29,119 @@ class RAGMint:
         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") or file.endswith(".md") or file.endswith(".rst"):
+            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}")
+
+        logging.info(f"📚 Loaded {len(docs)} documents from {self.docs_path}")
         return docs
 
-    def _embed_docs(self, model_name: str):
+    # -------------------------
+    # 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 = EmbeddingModel(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:
-        emb_model = EmbeddingModel(config["embedding_model"])
-        embeddings = self._embed_docs(config["embedding_model"])
-        retriever = Retriever(embeddings, self.documents)
-        reranker = Reranker(config["reranker"])
+        """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()
-        return RAGPipeline(retriever, reranker, 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")
-            reference = sample.get("answer")
+            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
 
+        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))}
 
+        return {
+            metric: avg_score,
+            "latency": elapsed / max(1, len(validation)),
+        }
+
+    # -------------------------
+    # Optimize
+    # -------------------------
     def optimize(
         self,
         validation_set: str,
@@ -90,16 +149,22 @@ class RAGMint:
         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")
 
+        # 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")
+        logging.info(f"🚀 Starting {search_type} optimization with {trials} trials")
 
+        # Select search strategy
         try:
             if search_type == "grid":
                 searcher = GridSearch(search_space)
@@ -108,16 +173,38 @@ class RAGMint:
             else:
                 searcher = RandomSearch(search_space, n_trials=trials)
         except Exception as e:
-            logging.warning(f"Falling back to RandomSearch due to missing deps: {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}")
+            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 found: {best}")
+        logging.info(f"🏆 Best configuration: {best}")
+
+        # Save to leaderboard
+        run_id = f"run_{uuid4().hex[:8]}"
+        leaderboard = Leaderboard()
+
+        corpus_stats = {
+            "num_docs": len(self.documents),
+            "avg_len": sum(len(d.split()) for d in self.documents) / max(1, len(self.documents)),
+            "corpus_size": sum(len(d) for d in self.documents),
+        }
+
+        leaderboard.upload(
+            run_id=run_id,
+            best_config=best,
+            best_score=best.get(metric, 0.0),
+            all_results=results,
+            documents=os.listdir(self.docs_path),
+            model=best.get("embedding_model", "unknown"),
+            corpus_stats=corpus_stats,
+        )
+
         return best, results