chATLAS_Chains 0.1.5__tar.gz → 0.1.6__tar.gz

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.5
+Version: 0.1.6
 Summary: A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project.
 Author-email: Joe Egan <joseph.caimin.egan@cern.ch>
 License: Apache-2.0
@@ -91,6 +91,12 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.6
+
+Fix bug in `reciprocal_rank_fusion` which caused it to silently return only one document
+
+Add `fallback_models` optional argument to `advanced_rag`
+
 #### 0.1.5
 
 Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/README.md

@@ -65,6 +65,12 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.6
+
+Fix bug in `reciprocal_rank_fusion` which caused it to silently return only one document
+
+Add `fallback_models` optional argument to `advanced_rag`
+
 #### 0.1.5
 
 Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains/chains/advanced.py

@@ -11,6 +11,7 @@ upweighting results that appear in both
 """
 
 import argparse
+import logging
 import os
 import sys
 from typing import TypedDict
@@ -30,6 +31,8 @@ from chATLAS_Chains.search.basic import search_runnable
 from chATLAS_Chains.utils.doc_utils import combine_documents
 from chATLAS_Embed.Base import VectorStore
 
+logger = logging.getLogger(__name__)
+
 
 # Define TypedDict for the simplified state
 class HybridGraphState(TypedDict, total=False):
@@ -57,6 +60,7 @@ def advanced_rag(
     rrf_constant: float = 60.0,
     rrf_weights: dict[str, float] | None = None,
     retry_config: RetryConfig | None = None,
+    fallback_models: list[str] | None = None,
 ) -> CompiledStateGraph:
     """
     Advanced Agentic RAG graph with optional query rewriting, dual-stage reranking and self-evaluation.
@@ -77,6 +81,7 @@ def advanced_rag(
     :param rrf_constant: Constant to use for RRF (Reciprocal Rank Fusion) calculations.
     :param rrf_weights: How to weight the RRF score for each retriever. Default is 1.0 for all.
     :param retry_config: Optional RetryConfig for Groq API calls
+    :param fallback_models: Optional list of model names to fall back to if the Groq API fails.
 
     :return: A compiled LangGraph with the chosen features
 
@@ -91,13 +96,15 @@ def advanced_rag(
     searcher = search_runnable(vectorstore)
 
     # Initialise models
-    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models, retry_config)
+    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models, retry_config, fallback_models)
 
     if enable_query_rewriting:
         query_rewriting_model_instance = get_chat_model(
             model_name=query_rewriting_model,
             temperature=query_rewriting_temperature,
             use_preview_models=use_preview_models,
+            retry_config=retry_config,
+            fallback_models=fallback_models,
         )
 
     # --------------- Define functions for the graph nodes ---------------
@@ -107,7 +114,7 @@ def advanced_rag(
         question = state["question"]
 
         if not enable_query_rewriting:
-            print("[WARNING] Query rewriting disabled, but query_rewriting node was called.")
+            logger.warning("Query rewriting disabled, but query_rewriting node was called.")
             # return original state
             return {**state}
 
@@ -128,7 +135,7 @@ def advanced_rag(
 
         search_kwargs = state.get("search_kwargs", {})
         if not search_kwargs:
-            print("[WARNING] No search_kwargs provided, using defaults.")
+            logger.warning("No search_kwargs provided, using defaults.")
             results = searcher.invoke(question)
         else:
             results = searcher.invoke(question, config={"metadata": {"search_kwargs": search_kwargs}})
@@ -136,7 +143,7 @@ def advanced_rag(
         if "docs" not in results:
             raise Exception("Search results missing 'docs' field")
 
-        print(f"Retrieved {len(results['docs'])} documents")
+        logger.debug(f"Retrieved {len(results['docs'])} documents")
 
         return {
             **state,
@@ -150,7 +157,7 @@ def advanced_rag(
         """
 
         if not enable_reranking:
-            print("[WARNING] Reranking disabled, but rerank node was called.")
+            logger.warning("Reranking disabled, but rerank node was called.")
             # return original state
             return {**state}
 
@@ -166,7 +173,7 @@ def advanced_rag(
             )
 
         except Exception as e:
-            print(f"[WARNING] Reranking failed with exception {e}. Returning original documents.")
+            logger.warning(f"Reranking failed with exception {e}. Returning original documents.")
             reranked_docs = docs
 
         return {
@@ -177,13 +184,13 @@ def advanced_rag(
     def rrf(state: HybridGraphState) -> HybridGraphState:
         """Perform Reciprocal Rank Fusion (RRF) on retrieved documents."""
         if not enable_rrf:
-            print("[WARNING] RRF disabled, but rrf node was called.")
+            logger.warning("RRF disabled, but rrf node was called.")
             # return original state
             return {**state}
 
         docs = state.get("docs", [])
         if not docs:
-            print("[WARNING] No documents retrieved for RRF.")
+            logger.warning("No documents retrieved for RRF.")
             return {**state}
 
         try:
@@ -191,7 +198,7 @@ def advanced_rag(
                 results=split_docs_by_retriever(docs), k=rrf_constant, weights=rrf_weights
             )
         except Exception as e:
-            print(f"[WARNING] RRF failed with exception {e}. Returning original documents.")
+            logger.warning(f"RRF failed with exception {e}. Returning original documents.")
             rrf_docs = docs
 
         return {
@@ -201,7 +208,6 @@ def advanced_rag(
 
     def generate_answer(state: HybridGraphState) -> HybridGraphState:
         """Generate a response from the LLM"""
-
         # Format the prompt using LangChain template
         prompt_input = {"context": combine_documents(state["docs"]), "question": state["question"]}
         final_prompt = prompt_template.format_messages(**prompt_input)
@@ -266,6 +272,7 @@ if __name__ == "__main__":
     from chATLAS_Chains.vectorstore import get_vectorstore
 
     twiki = get_vectorstore("twiki_prod")
+    mkdocs = get_vectorstore("mkdocs_prod_v1")
 
     retry_config = RetryConfig(
         max_retries=1,
@@ -274,34 +281,24 @@ if __name__ == "__main__":
 
     # Create the hybrid graph
     graph = advanced_rag(
-        vectorstore=[twiki],
+        vectorstore=[twiki, mkdocs],
         model_name=GROQ_PRODUCTION_MODELS[0],
         enable_query_rewriting=True,
         enable_rrf=True,
-        enable_reranking=True,
+        enable_reranking=False,
     )
 
-    # Test query
-    try:
-        ans = graph.invoke(
-            {
-                "question": "How can one check for and remove bad or corrupted events in the analysis?",
-                "search_kwargs": {
-                    "k_text": 3,
-                    "k": 15,
-                    "date_filter": "01-01-2010",
-                    # "type": ["twiki"],
-                },
-            }
-        )
-
-        print(f"Number of docs is : {len(ans['docs'])}")
-        print(f"Answer: {ans['answer']}")
-
-    except Exception as e:
-        print(f"❌ Graph execution failed with error: {e}")
-        print(f"Error type: {type(e).__name__}")
-        import traceback
+    ans = graph.invoke(
+        {
+            "question": "What is the crack veto in electron reconstruction?",
+            "search_kwargs": {
+                "k_text": 3,
+                "k": 15,
+                "date_filter": "01-01-2010",
+                # "type": ["twiki"],
+            },
+        }
+    )
 
-        traceback.print_exc()
-        sys.exit(1)
+    print(f"Number of docs is : {len(ans['docs'])}")
+    print(f"Answer: {ans['answer']}")

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains/documents/rrf.py

@@ -5,7 +5,9 @@ Functions for performing Reciprocal Rank Fusion (RRF) on documents retrieved fro
 from langchain_core.documents import Document
 
 
-def reciprocal_rank_fusion(results: dict, k: float = 60, weights: dict[str, float] | None = None) -> list:
+def reciprocal_rank_fusion(
+    results: dict[str, list[Document]], k: float = 60, weights: dict[str, float] | None = None
+) -> list:
     """
     Fuse search results using weighted Reciprocal Rank Fusion (RRF) algorithm.
 
@@ -39,24 +41,21 @@ def reciprocal_rank_fusion(results: dict, k: float = 60, weights: dict[str, floa
         if key not in weights:
             weights[key] = 1.0
 
-    if len(results) != len(weights):
-        raise ValueError("Number of results and weights are different")
-
     # Do this for each retriever in the dict
     for retriever, docs in results.items():
         for rank, doc in enumerate(docs, 1):
-            id = doc.id
+            name = doc.metadata.get("name")
 
             # initialise score to zero
-            if id not in rrf_scores:
-                rrf_scores[id] = 0
-                document_store[id] = doc
+            if name not in rrf_scores:
+                rrf_scores[name] = 0
+                document_store[name] = doc
 
             # weighted score for this retriever
-            rrf_scores[id] += weights[retriever] * (1 / (k + rank))
+            rrf_scores[name] += weights[retriever] * (1 / (k + rank))
 
             # track which retrievers had this document
-            metadata = document_store[id].metadata
+            metadata = document_store[name].metadata
             if "retrievers" not in metadata:
                 metadata["retrievers"] = []
             if retriever not in metadata["retrievers"]:
@@ -64,9 +63,10 @@ def reciprocal_rank_fusion(results: dict, k: float = 60, weights: dict[str, floa
 
     # Sort and return results
     sorted_results = []
-    for id in sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True):
-        doc = document_store[id]
-        doc.metadata["rrf_score"] = rrf_scores[id]
+
+    for name in sorted(rrf_scores.keys(), key=lambda x: rrf_scores[x], reverse=True):
+        doc = document_store[name]
+        doc.metadata["rrf_score"] = rrf_scores[name]
 
         sorted_results.append(doc)
 
@@ -90,7 +90,6 @@ def split_docs_by_retriever(docs: list[Document]):
             retriever_types.append(search_type)
 
     split_docs = {retriever_type: [] for retriever_type in retriever_types}
-
     for doc in docs:
         search_type = doc.metadata["search_type"]
         if search_type not in split_docs:
@@ -102,36 +101,40 @@ def split_docs_by_retriever(docs: list[Document]):
 
 if __name__ == "__main__":
     # Example usage
-    from langchain_core.documents import Document
 
     # Create some example documents
     doc1 = Document(
         page_content="In the Standard Model, the Higgs potential is responsible for electroweak symmetry breaking",
-        id="higgs",
-        metadata={"search_type": "vector"},
+        metadata={"search_type": "vector", "name": "higgs"},
+        source="twiki",
+        url="",
     )
     doc2 = Document(
         page_content="The top quark is the heaviest known elementary particle",
-        id="top",
-        metadata={"search_type": "vector"},
+        metadata={"search_type": "vector", "name": "top"},
+        source="twiki",
+        url="",
     )
 
     vector_results = [doc1, doc2]
 
     doc3 = Document(
         page_content="The gluon is the force carrier of Quantum Chromodynamics (QCD)",
-        id="gluon",
-        metadata={"search_type": "text"},
+        metadata={"search_type": "text", "name": "gluon"},
+        source="twiki",
+        url="",
     )
     doc4 = Document(
         page_content="In the Standard Model, the Higgs potential is responsible for electroweak symmetry breaking",
-        id="higgs",
-        metadata={"search_type": "text"},
+        metadata={"search_type": "text", "name": "higgs"},
+        source="twiki",
+        url="",
     )
     doc5 = Document(
         page_content="The Large Hadron Collider (LHC) is the world's largest and most powerful particle accelerator",
-        id="lhc",
-        metadata={"search_type": "text"},
+        metadata={"search_type": "text", "name": "lhc"},
+        source="twiki",
+        url="",
     )
 
     text_results = [doc3, doc4, doc5]
@@ -143,5 +146,5 @@ if __name__ == "__main__":
     print("RRF Results:")
     for doc in rrf_results:
         print(
-            f"Document ID: {doc.id}, RRF Score: {doc.metadata.get('rrf_score', 0)}, Retrievers: {doc.metadata.get('retrievers', [])}"
+            f"Document name: {doc.metadata['name']}, RRF Score: {doc.metadata['rrf_score']}, Retrievers: {doc.metadata['retrievers']}"
         )

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains/llm/groq.py

@@ -139,6 +139,7 @@ class AccGPTChatGroq(BaseChatModel, BaseModel):
     def __init__(self, **data):
         super().__init__(**data)
         if self.retry_config is None:
+            logger.debug("No retry_config provided, using default")
             self.retry_config = RetryConfig()
 
         if self.fallback_models is None:
@@ -565,7 +566,7 @@ if __name__ == "__main__":
     api_key = api_key.strip()
 
     base_url = os.getenv("CHATLAS_GROQ_BASE_URL")
-    base_url = "http://localhost:3000"
+    # base_url = "http://localhost:3000"
     if not base_url or not base_url.strip():
         raise ValueError("CHATLAS_GROQ_BASE_URL not set in environment")
     base_url = base_url.strip()

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains/llm/model_selection.py

@@ -84,6 +84,7 @@ def get_chat_model(
     temperature: float | None = None,
     use_preview_models: bool = False,
     retry_config: RetryConfig | None = None,
+    fallback_models: list[str] | None = None,
 ):
     """
     Initialize chat model with the provided model name (if supported)
@@ -100,6 +101,12 @@ def get_chat_model(
     :param use_preview_models: If True, allows the use of preview models from Groq. Defaults to False.
     :type use_preview_models: bool
 
+    :param retry_config: Configuration for retrying requests to Groq models. If None, the default retry configuration is used.
+    :type retry_config: RetryConfig | None
+
+    :param fallback_models: A list of model names to fall back to if the Groq API fails.
+    :type fallback_models: list[str] | None
+
     :raises ValueError: If the environment variable `CHATLAS_OPENAI_KEY` is not set when using OpenAI models, or if `CHATLAS_GROQ_KEY` and `CHATLAS_GROQ_BASE_URL` are not set when using Groq models.
 
     :return: An instance of the specified chat model.
@@ -143,17 +150,14 @@ def get_chat_model(
                 f"max_tokens ({max_tokens}) exceeds the model's maximum ({get_max_completion_tokens(model_name)})."
             )
 
-        # default retry config
-        if retry_config is None:
-            retry_config = RetryConfig()
-
         llm = AccGPTChatGroq(
             model_name=model_name,
             api_key=api_key,
             base_url=base_url,
             max_tokens=max_tokens,
             temperature=temperature,
-            retry_config=retry_config,
+            retry_config=retry_config,  # if None, AccGPTChatGroq will use the default
+            fallback_models=fallback_models,
         )
 
     return llm

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.5
+Version: 0.1.6
 Summary: A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project.
 Author-email: Joe Egan <joseph.caimin.egan@cern.ch>
 License: Apache-2.0
@@ -91,6 +91,12 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.6
+
+Fix bug in `reciprocal_rank_fusion` which caused it to silently return only one document
+
+Add `fallback_models` optional argument to `advanced_rag`
+
 #### 0.1.5
 
 Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/chATLAS_Chains.egg-info/SOURCES.txt

@@ -33,5 +33,6 @@ tests/conftest.py
 tests/test_chains.py
 tests/test_groq.py
 tests/test_llm.py
+tests/test_rrf.py
 tests/test_search.py
 tests/test_utils.py

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chATLAS_Chains"
-version = "0.1.5"
+version = "0.1.6"
 description = "A modular Python package for implementing Retrieval Augmented Generation chains for the chATLAS project."
 readme = "README.md"
 requires-python = ">=3.11"

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/tests/test_chains.py

@@ -22,7 +22,7 @@ from chATLAS_Chains.vectorstore import get_vectorstore
 @pytest.fixture
 def retry_config():
     """Retry config to use in tests"""
-    return RetryConfig(max_retries=2, max_delay=120.0)
+    return RetryConfig(max_retries=5, max_delay=120.0)
 
 
 def test_basic_retrieval_chain_returns_runnablesequence(twiki_vectorstore):
@@ -55,10 +55,11 @@ def test_basic_retrieval_chain_returns_docs_and_answer(twiki_vectorstore, retry_
     )
     output = chain.invoke("What is the Higgs boson?")
 
-    assert "docs" in output
-    assert "answer" in output
+    assert "docs" in output, "'docs' key not in output"
+    assert "answer" in output, "'answer' key not in output"
 
-    assert isinstance(output["docs"], list)
+    assert isinstance(output["docs"], list), "'docs' is not a list"
+    assert all([isinstance(obj, Document) for obj in output["docs"]]), "'docs' should be a list of Documents"
 
 
 def test_basic_retrieval_chain_multiple_vectorstores(three_vectorstores, retry_config):
@@ -70,10 +71,11 @@ def test_basic_retrieval_chain_multiple_vectorstores(three_vectorstores, retry_c
     )
     output = chain.invoke("What is the Higgs boson?")
 
-    assert "docs" in output
-    assert "answer" in output
+    assert "docs" in output, "'docs' key not in output"
+    assert len(output["docs"]) > 1, "Didn't return more than one doc with k=15, k_text=3, three vectorstores"
+    assert "answer" in output, "'answer' key not in output"
 
-    assert isinstance(output["docs"], list)
+    assert isinstance(output["docs"], list), "'docs' is not a list"
 
     sources = set([doc.metadata.get("source") for doc in output["docs"]])
     expected_sources = {"twiki", "MkDocs", "CDS"}
@@ -101,10 +103,13 @@ def test_advanced_rag_chain(twiki_vectorstore, retry_config):
         {"question": "What is the Higgs boson?", "search_kwargs": search_kwargs},
     )
 
-    assert "docs" in output
-    assert "answer" in output
-    assert isinstance(output["docs"], list)
-    assert isinstance(output["docs"][0], Document)
+    assert "docs" in output, "'docs' key not in output"
+    assert len(output["docs"]) > 1, "Didn't return more than one doc with k=15, k_text=3"
+    assert "answer" in output, "'answer' key not in output"
+    assert isinstance(output["docs"], list), "'docs' is not a list"
+    assert isinstance(output["docs"][0], Document), (
+        f"'docs' should contain instances of {Document}, but instead it has {type(output['docs'][0])}"
+    )
 
     # test with the additional options turned on (besides rerank)
     chain = advanced_rag(
@@ -119,10 +124,14 @@ def test_advanced_rag_chain(twiki_vectorstore, retry_config):
         {"question": "What is the Higgs boson?", "search_kwargs": search_kwargs},
     )
 
-    assert "docs" in output
-    assert "answer" in output
-    assert isinstance(output["docs"], list)
-    assert isinstance(output["docs"][0], Document)
+    assert "docs" in output, "'docs' key not in output"
+    assert len(output["docs"]) > 1, "Didn't return more than one doc with k=15, k_text=3"
+    assert "answer" in output, "'answer' key not in output"
+    assert isinstance(output["docs"], list), "'docs' is not a list"
+    assert all(isinstance(obj, Document) for obj in output["docs"]), (
+        f"'docs' should only contain instances of {Document}"
+    )
+    assert all("rrf_score" in doc.metadata for doc in output["docs"]), "'rrf_score' not in document metadata"
 
 
 def test_advanced_rag_chain_multiple_vectorstores(three_vectorstores, retry_config):
@@ -143,10 +152,13 @@ def test_advanced_rag_chain_multiple_vectorstores(three_vectorstores, retry_conf
         {"question": "What is the Higgs boson?", "search_kwargs": search_kwargs},
     )
 
-    assert "docs" in output
-    assert "answer" in output
-    assert isinstance(output["docs"], list)
-    assert isinstance(output["docs"][0], Document)
+    assert "docs" in output, "'docs' key not in output"
+    assert len(output["docs"]) > 1, "Didn't return more than one doc with k=15, k_text=3, three vectorstores"
+    assert "answer" in output, "'answer' key not in output"
+    assert isinstance(output["docs"], list), "'docs' is not a list"
+    assert isinstance(output["docs"][0], Document), (
+        f"'docs' should contain instances of {Document}, but instead it has {type(output['docs'][0])}"
+    )
 
     sources = set([doc.metadata.get("source") for doc in output["docs"]])
     expected_sources = {"twiki", "MkDocs", "CDS"}

{chatlas_chains-0.1.5 → chatlas_chains-0.1.6}/tests/test_llm.py

@@ -60,7 +60,9 @@ class TestModelSelection:
         """Test that Groq models are correctly initialized."""
         mock_class, mock_instance = mock_groq
 
-        model = get_chat_model(model_name, use_preview_models=True)
+        retry_config = RetryConfig(max_retries=5, base_delay=5.0, max_delay=300.0)
+
+        model = get_chat_model(model_name, use_preview_models=True, retry_config=retry_config)
 
         mock_class.assert_called_once_with(
             model_name=model_name,
@@ -68,7 +70,8 @@ class TestModelSelection:
             base_url="http://fake-groq-url.com",
             max_tokens=get_max_completion_tokens(model_name),
             temperature=None,
-            retry_config=RetryConfig(),
+            retry_config=retry_config,
+            fallback_models=None,
         )
         assert model is mock_instance
 
@@ -97,7 +100,8 @@ class TestModelSelection:
             base_url="http://fake-groq-url.com",
             max_tokens=get_max_completion_tokens(preview_model),
             temperature=None,
-            retry_config=RetryConfig(),
+            retry_config=None,
+            fallback_models=None,
         )
         assert model is mock_instance
 
@@ -197,9 +201,27 @@ class TestModelSelection:
             base_url="http://whitespace-url.com",
             max_tokens=get_max_completion_tokens(GROQ_PRODUCTION_MODELS[0]),
             temperature=None,
-            retry_config=RetryConfig(),
+            retry_config=None,
+            fallback_models=None,
         )
 
     def test_supported_models_lists(self):
         """Test that SUPPORTED_CHAT_MODELS is correctly composed of OPENAI_MODELS and GROQ_MODELS."""
         assert set(SUPPORTED_CHAT_MODELS) == set(OPENAI_MODELS + GROQ_MODELS)
+
+    def test_retry_config_set_correctly(self):
+        """Test that the retry_config is correctly passed to the Groq model."""
+        retry_config = RetryConfig(max_retries=3, base_delay=2.0, max_delay=100.0)
+
+        model = get_chat_model(GROQ_PRODUCTION_MODELS[0], use_preview_models=True, retry_config=retry_config)
+
+        assert isinstance(model, AccGPTChatGroq), "Returned model is not an instance of AccGPTChatGroq"
+        assert model.retry_config.max_retries == 3, (
+            f"max_retries not set correctly on model, it set to {model.retry_config.max_retries}"
+        )
+        assert model.retry_config.base_delay == 2.0, (
+            f"base_delay not set correctly on model, it set to {model.retry_config.base_delay}"
+        )
+        assert model.retry_config.max_delay == 100.0, (
+            f"max_delay not set correctly on model, it set to {model.retry_config.max_delay}"
+        )

chatlas_chains-0.1.6/tests/test_rrf.py

@@ -0,0 +1,126 @@
+import pytest
+from langchain_core.documents import Document
+
+from chATLAS_Chains.documents.rrf import reciprocal_rank_fusion, split_docs_by_retriever
+
+
+@pytest.fixture
+def vector_results():
+    doc1 = Document(
+        page_content="In the Standard Model, the Higgs potential is responsible for electroweak symmetry breaking",
+        metadata={"search_type": "vector", "name": "higgs"},
+        source="twiki",
+        url="",
+    )
+    doc2 = Document(
+        page_content="The top quark is the heaviest known elementary particle",
+        # name="top",
+        metadata={"search_type": "vector", "name": "top"},
+        source="twiki",
+        url="",
+    )
+
+    return [doc1, doc2]
+
+
+@pytest.fixture
+def text_results():
+    doc3 = Document(
+        page_content="The gluon is the force carrier of Quantum Chromodynamics (QCD)",
+        # name="gluon",
+        metadata={"search_type": "text", "name": "gluon"},
+        source="twiki",
+        url="",
+    )
+    doc4 = Document(
+        page_content="In the Standard Model, the Higgs potential is responsible for electroweak symmetry breaking",
+        # name="higgs",
+        metadata={"search_type": "text", "name": "higgs"},
+        source="twiki",
+        url="",
+    )
+    doc5 = Document(
+        page_content="The Large Hadron Collider (LHC) is the world's largest and most powerful particle accelerator",
+        # name="lhc",
+        metadata={"search_type": "text", "name": "lhc"},
+        source="twiki",
+        url="",
+    )
+
+    return [doc3, doc4, doc5]
+
+
+@pytest.fixture
+def docs_dict(vector_results, text_results):
+    return {
+        "vector": vector_results,
+        "text": text_results,
+    }
+
+
+def test_split_docs_by_retriever(vector_results, text_results):
+    # pass it the combined list, it should re-split
+    combined = vector_results + text_results
+
+    split = split_docs_by_retriever(combined)
+
+    assert isinstance(split, dict), "return type should be dict"
+    assert "text" in split, "missing retriever type in returned dict"
+    assert "vector" in split, "missing retriever type in returned dict"
+    assert len(split["text"]) == len(text_results), "text results length mismatch"
+    assert len(split["vector"]) == len(vector_results), "vector results length mismatch"
+
+
+def test_split_docs_by_retriever_missing_search_type(vector_results):
+    # remove search_type from one doc
+    vector_results[0].metadata.pop("search_type")
+
+    with pytest.raises(ValueError, match="search_type metadata field missing from Document.metadata"):
+        split_docs_by_retriever(vector_results)
+
+
+def test_reciprocal_rank_fusion(docs_dict):
+    fused = reciprocal_rank_fusion(docs_dict)
+
+    assert isinstance(fused, list), "return type should be list"
+    assert len(fused) <= 5, "fused results length exceeds top_k"
+    assert all(isinstance(doc, Document) for doc in fused), f"all items in fused results should be {Document} instances"
+
+    assert all("rrf_score" in doc.metadata for doc in fused), "rrf_score missing from Document.metadata"
+
+    # Check that the documents are ordered by score (highest first)
+    scores = [doc.metadata["rrf_score"] for doc in fused]
+    assert scores == sorted(scores, reverse=True), "documents are not ordered by rrf_score"
+
+
+def test_reciprocal_rank_fusion_with_docs_splitter(vector_results, text_results):
+    """
+    Same as above, but check it still works using the actual splitter function
+    """
+    fused = reciprocal_rank_fusion(split_docs_by_retriever(vector_results + text_results))
+
+    assert isinstance(fused, list), "return type should be list"
+    assert len(fused) <= 5, "fused results length exceeds top_k"
+    assert all(isinstance(doc, Document) for doc in fused), f"all items in fused results should be {Document} instances"
+
+    assert all("rrf_score" in doc.metadata for doc in fused), "rrf_score missing from Document.metadata"
+
+    # Check that the documents are ordered by score (highest first)
+    scores = [doc.metadata["rrf_score"] for doc in fused]
+    assert scores == sorted(scores, reverse=True), "documents are not ordered by rrf_score"
+
+
+def test_reciprocal_rank_fusion_raise_for_negative_k(docs_dict):
+    with pytest.raises(ValueError, match="k must be positive"):
+        reciprocal_rank_fusion(docs_dict, k=-7)
+
+
+def test_reciprocal_rank_fusion_raise_for_mismatched_weights(docs_dict):
+    weights = {
+        "vector": 1.0,
+        "text": 1.0,
+        "extra": 1.0,
+    }
+
+    with pytest.raises(ValueError, match="Weight 'extra' not found in provided results dictionary."):
+        reciprocal_rank_fusion(docs_dict, weights=weights)