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

ragmint/qa_generator.py

@@ -0,0 +1,190 @@
+"""
+Batched Validation QA Generator for Ragmint (Functional Version)
+
+Generates a JSON QA dataset from a large corpus using an LLM.
+Processes documents in batches to avoid token limits and API errors.
+Uses topic-aware dynamic question count estimation.
+"""
+
+import os
+import re
+import json
+import math
+import time
+import argparse
+from pathlib import Path
+from dotenv import load_dotenv
+import numpy as np
+from sklearn.cluster import KMeans
+from sentence_transformers import SentenceTransformer
+
+# --- Load .env if available ---
+load_dotenv()
+
+
+# ---------- Utility functions ----------
+
+def extract_json_from_markdown(text: str):
+    """Extract JSON from a markdown-style code block."""
+    match = re.search(r"```(?:json)?\s*(\[\s*[\s\S]*?\s*\])\s*```", text, re.MULTILINE)
+    if match:
+        json_str = match.group(1)
+        return json.loads(json_str)
+    else:
+        return json.loads(text.strip())
+
+
+def read_corpus(docs_path: str):
+    """Load all text documents from a folder."""
+    docs = []
+    for file in Path(docs_path).glob("**/*.txt"):
+        with open(file, "r", encoding="utf-8") as f:
+            text = f.read().strip()
+            if text:
+                docs.append({"filename": file.name, "text": text})
+    return docs
+
+
+def determine_question_count(text: str, embedder, min_q=3, max_q=25):
+    """Estimate number of questions dynamically based on text length and topic diversity."""
+    sentences = [s.strip() for s in text.split('.') if len(s.strip().split()) > 3]
+    word_count = len(text.split())
+
+    if word_count == 0:
+        return min_q
+
+    base_q = math.log1p(word_count / 150)
+
+    # Topic diversity via clustering
+    n_sent = len(sentences)
+    if n_sent < 5:
+        topic_factor = 1.0
+    else:
+        try:
+            emb = embedder.encode(sentences, normalize_embeddings=True)
+            n_clusters = min(max(2, n_sent // 10), 8)
+            km = KMeans(n_clusters=n_clusters, n_init=10, random_state=42)
+            labels = km.fit_predict(emb)
+            topic_factor = len(set(labels)) / n_clusters
+        except Exception as e:
+            print(f"[WARN] Clustering failed ({type(e).__name__}): {e}")
+            topic_factor = 1.0
+
+    score = base_q * (1 + 0.8 * topic_factor)
+    question_count = round(min_q + score)
+    return int(max(min_q, min(question_count, max_q)))
+
+
+def setup_llm(llm_model="gemini-2.5-flash-lite"):
+    """Configure Gemini or Claude based on available environment keys."""
+    google_key = os.getenv("GOOGLE_API_KEY")
+    anthropic_key = os.getenv("ANTHROPIC_API_KEY")
+
+    if google_key:
+        import google.generativeai as genai
+        genai.configure(api_key=google_key)
+        llm = genai.GenerativeModel(llm_model)
+        return llm, "gemini"
+
+    elif anthropic_key:
+        from anthropic import Anthropic
+        llm = Anthropic(api_key=anthropic_key)
+        return llm, "claude"
+
+    else:
+        raise ValueError("Set ANTHROPIC_API_KEY or GOOGLE_API_KEY in your environment.")
+
+
+def generate_qa_for_batch(batch, llm, backend, embedder, min_q=3, max_q=25):
+    """Send one LLM call for a batch of documents."""
+    prompt_texts = []
+    for doc in batch:
+        n_questions = determine_question_count(doc["text"], embedder, min_q, max_q)
+        prompt_texts.append(
+            f"Document: {doc['text'][:1000]}\n"
+            f"Generate {n_questions} factual question-answer pairs in JSON format."
+        )
+
+    prompt = "\n\n".join(prompt_texts)
+    prompt += "\n\nReturn a single JSON array of objects like:\n" \
+              '[{"query": "string", "expected_answer": "string"}]'
+
+    try:
+        if backend == "gemini":
+            response = llm.generate_content(prompt)
+            text_out = getattr(response, "text", None)
+            if not text_out and hasattr(response, "candidates"):
+                text_out = response.candidates[0].content.parts[0].text
+            return extract_json_from_markdown(text_out)
+
+        elif backend == "claude":
+            response = llm.messages.create(
+                model="claude-3-opus-20240229",
+                messages=[{"role": "user", "content": prompt}],
+                max_tokens=2000,
+            )
+            return json.loads(response.content[0].text)
+
+    except Exception as e:
+        print(f"[WARN] Failed to parse batch: {e}")
+        return []
+
+
+def save_json(output_path, data):
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    with open(output_path, "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2, ensure_ascii=False)
+    print(f"[INFO] Saved {len(data)} QAs → {output_path}")
+
+
+def generate_validation_qa(
+    docs_path="data/docs",
+    output_path="experiments/validation_qa.json",
+    llm_model="gemini-2.5-flash-lite",
+    batch_size=5,
+    sleep_between_batches=2,
+    min_q=3,
+    max_q=25,
+):
+    """Main pipeline to generate QAs."""
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    llm, backend = setup_llm(llm_model)
+    all_qa = []
+
+    corpus = read_corpus(docs_path)
+    print(f"[INFO] Loaded {len(corpus)} documents from {docs_path}")
+
+    for i in range(0, len(corpus), batch_size):
+        batch = corpus[i: i + batch_size]
+        batch_qa = generate_qa_for_batch(batch, llm, backend, embedder, min_q, max_q)
+        all_qa.extend(batch_qa)
+        print(f"[INFO] Batch {i // batch_size + 1}: {len(batch_qa)} QAs (Total: {len(all_qa)})")
+        time.sleep(sleep_between_batches)
+
+    save_json(output_path, all_qa)
+
+
+# ---------- CLI entry point ----------
+
+def main():
+    parser = argparse.ArgumentParser(description="Generate validation QA dataset for Ragmint.")
+    parser.add_argument("--docs_path", type=str, default="data/docs")
+    parser.add_argument("--output", type=str, default="experiments/validation_qa.json")
+    parser.add_argument("--batch_size", type=int, default=5)
+    parser.add_argument("--sleep", type=int, default=2)
+    parser.add_argument("--min_q", type=int, default=3)
+    parser.add_argument("--max_q", type=int, default=25)
+    args = parser.parse_args()
+
+    generate_validation_qa(
+        docs_path=args.docs_path,
+        output_path=args.output,
+        batch_size=args.batch_size,
+        sleep_between_batches=args.sleep,
+        min_q=args.min_q,
+        max_q=args.max_q,
+    )
+
+
+if __name__ == "__main__":
+    main()

ragmint/tests/test_autotuner.py

@@ -1,42 +1,64 @@
+import os
+import json
 import pytest
 from ragmint.autotuner import AutoRAGTuner
 
 
-def test_autorag_recommend_small():
-    """Small corpus should trigger BM25 + OpenAI."""
-    tuner = AutoRAGTuner({"size": 500, "avg_len": 150})
-    rec = tuner.recommend()
-    assert rec["retriever"] == "BM25"
-    assert rec["embedding_model"] == "OpenAI"
+def setup_docs(tmp_path):
+    """Create a temporary corpus with multiple text files for testing."""
+    corpus = tmp_path / "corpus"
+    corpus.mkdir()
+    (corpus / "short_doc.txt").write_text("AI is changing the world.")
+    (corpus / "long_doc.txt").write_text("Machine learning enables RAG pipelines to optimize retrievals. " * 50)
+    return str(corpus)
 
 
-def test_autorag_recommend_medium():
-    """Medium corpus should trigger Chroma + SentenceTransformers."""
-    tuner = AutoRAGTuner({"size": 5000, "avg_len": 200})
-    rec = tuner.recommend()
-    assert rec["retriever"] == "Chroma"
-    assert rec["embedding_model"] == "SentenceTransformers"
+def test_analyze_corpus(tmp_path):
+    """Ensure AutoRAGTuner analyzes corpus correctly."""
+    docs_path = setup_docs(tmp_path)
+    tuner = AutoRAGTuner(docs_path)
+    stats = tuner.corpus_stats
 
+    assert stats["num_docs"] == 2, "Should detect all documents"
+    assert stats["size"] > 0, "Corpus size should be positive"
+    assert stats["avg_len"] > 0, "Average document length should be computed"
 
-def test_autorag_recommend_large():
-    """Large corpus should trigger FAISS + InstructorXL."""
-    tuner = AutoRAGTuner({"size": 50000, "avg_len": 300})
-    rec = tuner.recommend()
-    assert rec["retriever"] == "FAISS"
-    assert rec["embedding_model"] == "InstructorXL"
 
+@pytest.mark.parametrize("size,expected_retriever", [
+    (10_000, "Chroma"),
+    (500_000, "FAISS"),
+    (1_000, "BM25"),
+])
+def test_recommendation_logic(tmp_path, monkeypatch, size, expected_retriever):
+    """Validate retriever recommendation based on corpus size and chunk suggestion."""
+    docs_path = setup_docs(tmp_path)
+    tuner = AutoRAGTuner(docs_path)
 
-def test_autorag_auto_tune(monkeypatch):
-    """Test auto_tune with a mock validation dataset."""
-    tuner = AutoRAGTuner({"size": 12000, "avg_len": 250})
+    # Mock corpus stats manually
+    tuner.corpus_stats = {"size": size, "avg_len": 300, "num_docs": 10}
 
-    # Monkeypatch evaluate_config inside autotuner
-    import ragmint.autotuner as autotuner
-    def mock_eval(config, data):
-        return {"faithfulness": 0.9, "latency": 0.01}
-    monkeypatch.setattr(autotuner, "evaluate_config", mock_eval)
+    # Provide mandatory num_chunk_pairs
+    rec = tuner.recommend(num_chunk_pairs=3)
 
-    result = tuner.auto_tune([{"question": "What is AI?", "answer": "Artificial Intelligence"}])
-    assert "recommended" in result
-    assert "results" in result
-    assert isinstance(result["results"], dict)
+    assert "retriever" in rec and "embedding_model" in rec
+    assert rec["retriever"] == expected_retriever, f"Expected {expected_retriever}"
+    assert rec["chunk_size"] > 0 and rec["overlap"] >= 0
+    assert "chunk_candidates" in rec, "Should include suggested chunk pairs"
+    assert len(rec["chunk_candidates"]) == 3, "Should generate correct number of chunk pairs"
+
+
+def test_invalid_corpus_path(tmp_path):
+    """Should handle missing directories gracefully."""
+    missing_path = tmp_path / "nonexistent"
+    tuner = AutoRAGTuner(str(missing_path))
+    assert tuner.corpus_stats["size"] == 0
+    assert tuner.corpus_stats["num_docs"] == 0
+
+
+def test_suggest_chunk_sizes_requires_num_pairs(tmp_path):
+    """Ensure suggest_chunk_sizes raises error if num_pairs is not provided."""
+    docs_path = setup_docs(tmp_path)
+    tuner = AutoRAGTuner(docs_path)
+
+    with pytest.raises(ValueError):
+        tuner.suggest_chunk_sizes()

ragmint/tests/test_config_adapter.py

@@ -0,0 +1,39 @@
+import pytest
+from ragmint.integrations.config_adapter import LangchainConfigAdapter
+
+def test_default_conversion():
+    """Test that default config values are applied correctly."""
+    cfg = {
+        "retriever": "FAISS",
+        "embedding_model": "all-MiniLM-L6-v2",
+        "chunk_size": 500,
+        "overlap": 100
+    }
+
+    adapter = LangchainConfigAdapter(cfg)
+    result = adapter.to_standard_config()
+
+    assert result["retriever"].lower() == "faiss"
+    assert result["embedding_model"] == "sentence-transformers/all-MiniLM-L6-v2"
+    assert result["chunk_size"] == 500
+    assert result["overlap"] == 100
+
+
+def test_missing_fields_are_defaulted():
+    """Ensure missing optional fields (e.g. chunk params) are filled in."""
+    cfg = {"retriever": "BM25", "embedding_model": "all-MiniLM-L6-v2"}
+    adapter = LangchainConfigAdapter(cfg)
+    result = adapter.to_standard_config()
+
+    assert "chunk_size" in result
+    assert "overlap" in result
+    assert result["chunk_size"] > 0
+    assert result["overlap"] >= 0
+
+
+def test_validation_of_invalid_retriever():
+    """Ensure invalid retriever names raise an informative error."""
+    cfg = {"retriever": "InvalidBackend", "embedding_model": "all-MiniLM-L6-v2"}
+
+    with pytest.raises(ValueError, match="Unsupported retriever backend"):
+        LangchainConfigAdapter(cfg).to_standard_config()

ragmint/tests/test_embeddings.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pytest
+from ragmint.core.embeddings import Embeddings
+
+
+def test_dummy_backend_output_shape():
+    model = Embeddings(backend="dummy")
+    texts = ["hello", "world"]
+    embeddings = model.encode(texts)
+
+    # Expect 2x768 array
+    assert isinstance(embeddings, np.ndarray)
+    assert embeddings.shape == (2, 768)
+    assert embeddings.dtype == np.float32
+
+
+def test_dummy_backend_single_string():
+    model = Embeddings(backend="dummy")
+    text = "test"
+    embeddings = model.encode(text)
+
+    assert embeddings.shape == (1, 768)
+    assert isinstance(embeddings, np.ndarray)
+
+
+'''@pytest.mark.skipif(
+    not hasattr(__import__('importlib').util.find_spec("sentence_transformers"), "loader"),
+    reason="sentence-transformers not installed"
+)
+def test_huggingface_backend_output_shape():
+    model = Embeddings(backend="huggingface", model_name="all-MiniLM-L6-v2")
+    texts = ["This is a test.", "Another sentence."]
+    embeddings = model.encode(texts)
+
+    # Expect 2x384 for MiniLM-L6-v2
+    assert isinstance(embeddings, np.ndarray)
+    assert embeddings.ndim == 2
+    assert embeddings.shape[0] == len(texts)
+    assert embeddings.dtype == np.float32
+'''
+
+def test_invalid_backend():
+    try:
+        Embeddings(backend="unknown")
+    except ValueError as e:
+        assert "Unsupported embedding backend" in str(e)

ragmint/tests/test_explainer.py

@@ -1,20 +1,36 @@
 import pytest
+import sys
+import types
 from ragmint.explainer import explain_results
 
 
-def test_explain_results_gemini():
-    """Gemini explanation should contain model-specific phrasing."""
-    config_a = {"retriever": "FAISS", "embedding_model": "OpenAI"}
-    config_b = {"retriever": "Chroma", "embedding_model": "SentenceTransformers"}
-    result = explain_results(config_a, config_b, model="gemini")
-    assert isinstance(result, str)
-    assert "Gemini" in result or "gemini" in result
+def test_explain_results_with_claude(monkeypatch):
+    """Claude explanation should use Anthropic API path when ANTHROPIC_API_KEY is set."""
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "fake-key")
+    monkeypatch.delenv("GOOGLE_API_KEY", raising=False)
+
+    # Create a fake anthropic module with the required interface
+    mock_anthropic = types.ModuleType("anthropic")
+
+    class MockContent:
+        text = "Claude: The best configuration performs well due to optimized chunk size."
+
+    class MockMessages:
+        def create(self, *args, **kwargs):
+            return type("MockResponse", (), {"content": [MockContent()]})()
+
+    class MockClient:
+        def __init__(self, api_key):
+            self.messages = MockMessages()
+
+    mock_anthropic.Anthropic = MockClient
+    sys.modules["anthropic"] = mock_anthropic  # Inject fake module
+
+    best = {"retriever": "Chroma", "metric": 0.9}
+    all_results = [{"retriever": "FAISS", "metric": 0.85}]
+    corpus_stats = {"size": 10000, "avg_len": 400, "num_docs": 20}
 
+    result = explain_results(best, all_results, corpus_stats, model="claude-3-opus-20240229")
 
-def test_explain_results_claude():
-    """Claude explanation should contain model-specific phrasing."""
-    config_a = {"retriever": "FAISS"}
-    config_b = {"retriever": "Chroma"}
-    result = explain_results(config_a, config_b, model="claude")
     assert isinstance(result, str)
     assert "Claude" in result or "claude" in result

ragmint/tests/test_integration_autotuner_ragmint.py

@@ -1,60 +1,47 @@
+import os
+import json
 import pytest
-from ragmint.tuner import RAGMint
 from ragmint.autotuner import AutoRAGTuner
+from ragmint.tuner import RAGMint
 
 
-def test_integration_ragmint_autotune(monkeypatch, tmp_path):
-    """
-    Smoke test for integration between AutoRAGTuner and RAGMint.
-    Ensures end-to-end flow runs without real retrievers or embeddings.
-    """
-
-    # --- Mock corpus and validation data ---
+def setup_docs(tmp_path):
+    """Create a temporary corpus for integration testing."""
     corpus = tmp_path / "docs"
     corpus.mkdir()
-    (corpus / "doc1.txt").write_text("This is an AI document.")
-    validation_data = [{"question": "What is AI?", "answer": "Artificial Intelligence"}]
-
-    # --- Mock RAGMint.optimize() to avoid real model work ---
-    def mock_optimize(self, validation_set=None, metric="faithfulness", trials=2):
-        return (
-            {"retriever": "FAISS", "embedding_model": "OpenAI", "score": 0.88},
-            [{"trial": 1, "score": 0.88}],
-        )
-
-    monkeypatch.setattr(RAGMint, "optimize", mock_optimize)
-
-    # --- Mock evaluation used by AutoRAGTuner ---
-    def mock_evaluate_config(config, data):
-        return {"faithfulness": 0.9, "latency": 0.01}
-
-    import ragmint.autotuner as autotuner
-    monkeypatch.setattr(autotuner, "evaluate_config", mock_evaluate_config)
-
-    # --- Create AutoRAGTuner and RAGMint instances ---
-    ragmint = RAGMint(
-        docs_path=str(corpus),
-        retrievers=["faiss", "chroma"],
-        embeddings=["text-embedding-3-small"],
-        rerankers=["mmr"],
+    (corpus / "doc1.txt").write_text("This document discusses Artificial Intelligence and Machine Learning.")
+    (corpus / "doc2.txt").write_text("Retrieval-Augmented Generation combines retrievers and LLMs effectively.")
+    return str(corpus)
+
+
+def setup_validation_file(tmp_path):
+    """Create a temporary validation QA dataset."""
+    data = [
+        {"question": "What is AI?", "answer": "Artificial Intelligence"},
+        {"question": "Define RAG", "answer": "Retrieval-Augmented Generation"},
+    ]
+    file = tmp_path / "validation_qa.json"
+    with open(file, "w", encoding="utf-8") as f:
+        json.dump(data, f)
+    return str(file)
+
+
+def test_autotune_integration(tmp_path):
+    """Test that AutoRAGTuner can fully run a RAGMint optimization."""
+    docs_path = setup_docs(tmp_path)
+    val_file = setup_validation_file(tmp_path)
+
+    tuner = AutoRAGTuner(docs_path)
+    best, results = tuner.auto_tune(
+        validation_set=val_file,
+        metric="faithfulness",
+        trials=2,
+        search_type="random",
     )
 
-    tuner = AutoRAGTuner({"size": 2000, "avg_len": 150})
-
-    # --- Run Auto-Tune and RAG Optimization ---
-    recommendation = tuner.recommend()
-    assert "retriever" in recommendation
-    assert "embedding_model" in recommendation
-
-    tuning_results = tuner.auto_tune(validation_data)
-    assert "results" in tuning_results
-    assert isinstance(tuning_results["results"], dict)
-
-    # --- Run RAGMint optimization flow (mocked) ---
-    best_config, results = ragmint.optimize(validation_set=validation_data, trials=2)
-    assert isinstance(best_config, dict)
-    assert "score" in best_config
-    assert isinstance(results, list)
-
-    # --- Integration Success ---
-    print(f"Integration OK: AutoRAG recommended {recommendation}, RAGMint best {best_config}")
+    # Assertions on the results
+    assert isinstance(best, dict), "Best configuration should be a dict"
+    assert isinstance(results, list), "Results should be a list"
+    assert len(results) > 0, "Optimization should produce results"
+    assert "retriever" in best and "embedding_model" in best
+    assert best.get("faithfulness", 0.0) >= 0.0, "Metric value should be non-negative"

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())