chATLAS_Chains 0.1.4__tar.gz → 0.1.5__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.4
+Version: 0.1.5
 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,10 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.5
+
+Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload
+
 #### 0.1.4
 
 Support for Groq-hosted models

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

@@ -65,6 +65,10 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.5
+
+Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload
+
 #### 0.1.4
 
 Support for Groq-hosted models

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

@@ -22,6 +22,7 @@ from langgraph.graph.state import CompiledStateGraph
 
 from chATLAS_Chains.documents.rerank import rerank_documents
 from chATLAS_Chains.documents.rrf import reciprocal_rank_fusion, split_docs_by_retriever
+from chATLAS_Chains.llm.groq import RetryConfig
 from chATLAS_Chains.llm.model_selection import GROQ_PRODUCTION_MODELS, get_chat_model
 from chATLAS_Chains.prompt.starters import CHAT_PROMPT_TEMPLATE
 from chATLAS_Chains.query.query_rewriting import rewrite_query
@@ -55,6 +56,7 @@ def advanced_rag(
     pinecone_api_key: str | None = None,
     rrf_constant: float = 60.0,
     rrf_weights: dict[str, float] | None = None,
+    retry_config: RetryConfig | None = None,
 ) -> CompiledStateGraph:
     """
     Advanced Agentic RAG graph with optional query rewriting, dual-stage reranking and self-evaluation.
@@ -73,7 +75,8 @@ def advanced_rag(
     :param rerank_model: Name of the Pinecone reranker model to use.
     :param pinecone_api_key: Pinecone API key. If None, will use PINECONE_API_KEY environment variable.
     :param rrf_constant: Constant to use for RRF (Reciprocal Rank Fusion) calculations.
-    :rrf_weights: How to weight the RRF score for each retriever. Default is 1.0 for all.
+    :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
 
     :return: A compiled LangGraph with the chosen features
 
@@ -88,7 +91,7 @@ def advanced_rag(
     searcher = search_runnable(vectorstore)
 
     # Initialise models
-    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models)
+    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models, retry_config)
 
     if enable_query_rewriting:
         query_rewriting_model_instance = get_chat_model(
@@ -264,6 +267,11 @@ if __name__ == "__main__":
 
     twiki = get_vectorstore("twiki_prod")
 
+    retry_config = RetryConfig(
+        max_retries=1,
+        max_delay=120.0,
+    )
+
     # Create the hybrid graph
     graph = advanced_rag(
         vectorstore=[twiki],

{chatlas_chains-0.1.4 → chatlas_chains-0.1.5}/chATLAS_Chains/chains/basic.py

@@ -4,6 +4,7 @@ from typing import Optional
 from langchain_core.prompts import ChatPromptTemplate
 from langchain_core.runnables import RunnableSerializable
 
+from chATLAS_Chains.llm.groq import RetryConfig
 from chATLAS_Chains.llm.model_selection import get_chat_model
 from chATLAS_Chains.search.basic import search_runnable
 from chATLAS_Chains.utils.doc_utils import combine_documents
@@ -16,6 +17,8 @@ def basic_retrieval_chain(
     model_name: str,
     max_tokens: int | None = None,
     temperature: float | None = None,
+    use_preview_models: bool = False,
+    retry_config: RetryConfig | None = None,
 ) -> RunnableSerializable:
     """
     Baseline RAG retrieval chain. Searches one or several vectorstores in parallel, passes retrieved documents to the model
@@ -30,13 +33,17 @@ def basic_retrieval_chain(
     :type max_tokens: int | None
     :param temperature: The temperature to use for the model's response generation. Defaults to None.
     :type temperature: float | None
+    :param use_preview_models: Whether to allow the use of preview models from Groq. Defaults to False.
+    :type use_preview_models: bool
+    :param retry_config: Configuration for retrying requests to the model in case of failures. Defaults to None, which will use the default retry configuration.
+    :type retry_config: RetryConfig | None
 
     :return: A LangChain RunnableSerializable chain that performs retrieval and response generation.
     :rtype: RunnableSerializable
 
     """
     prompt_template = ChatPromptTemplate.from_template(prompt)
-    model = get_chat_model(model_name, max_tokens, temperature)
+    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models, retry_config)
 
     search = search_runnable(vectorstore)
 

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

@@ -26,8 +26,8 @@ GROQ_PREVIEW_MODELS = [
     "meta-llama/llama-prompt-guard-2-22m",
     "meta-llama/llama-prompt-guard-2-86m",
     "moonshotai/kimi-k2-instruct",
-    # "playai-tts", # text to audio model
-    # "playai-tts-arabic", # text to audio model
+    # "playai-tts",  # text to audio model
+    # "playai-tts-arabic",  # text to audio model
     "qwen/qwen3-32b",
     "gemma2-9b-it",
 ]

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.4
+Version: 0.1.5
 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,10 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.5
+
+Fix missing `retry_config` argument in `advanced_rag` caused by early PyPI upload
+
 #### 0.1.4
 
 Support for Groq-hosted models

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chATLAS_Chains"
-version = "0.1.4"
+version = "0.1.5"
 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.4 → chatlas_chains-0.1.5}/tests/test_chains.py

@@ -13,16 +13,19 @@ from pydantic import ValidationError
 
 from chATLAS_Chains.chains.advanced import advanced_rag
 from chATLAS_Chains.chains.basic import basic_retrieval_chain
+from chATLAS_Chains.llm.groq import RetryConfig
 from chATLAS_Chains.llm.model_selection import GROQ_PRODUCTION_MODELS
 from chATLAS_Chains.prompt.starters import CHAT_PROMPT_TEMPLATE
 from chATLAS_Chains.vectorstore import get_vectorstore
 
 
-def test_basic_retrieval_chain_returns_runnablesequence(twiki_vectorstore):
-    """
-    Test the search_runnable function with a populated vector store.
-    """
+@pytest.fixture
+def retry_config():
+    """Retry config to use in tests"""
+    return RetryConfig(max_retries=2, max_delay=120.0)
+
 
+def test_basic_retrieval_chain_returns_runnablesequence(twiki_vectorstore):
     chain = basic_retrieval_chain(prompt=CHAT_PROMPT_TEMPLATE, vectorstore=twiki_vectorstore, model_name="gpt-4o-mini")
 
     assert isinstance(chain, RunnableSequence)
@@ -30,7 +33,7 @@ def test_basic_retrieval_chain_returns_runnablesequence(twiki_vectorstore):
 
 def test_search_runnable_returns_runnablesequence_with_list_of_vectorstores(three_vectorstores):
     chain = basic_retrieval_chain(prompt=CHAT_PROMPT_TEMPLATE, vectorstore=three_vectorstores, model_name="gpt-4o-mini")
-    print(type(chain))
+
     assert isinstance(chain, RunnableSequence)
 
 
@@ -43,9 +46,12 @@ def test_search_runnable_error_on_invalid_input():
         )
 
 
-def test_basic_retrieval_chain_returns_docs_and_answer(twiki_vectorstore):
+def test_basic_retrieval_chain_returns_docs_and_answer(twiki_vectorstore, retry_config):
     chain = basic_retrieval_chain(
-        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=twiki_vectorstore, model_name=GROQ_PRODUCTION_MODELS[0]
+        prompt=CHAT_PROMPT_TEMPLATE,
+        vectorstore=twiki_vectorstore,
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        retry_config=retry_config,
     )
     output = chain.invoke("What is the Higgs boson?")
 
@@ -55,9 +61,12 @@ def test_basic_retrieval_chain_returns_docs_and_answer(twiki_vectorstore):
     assert isinstance(output["docs"], list)
 
 
-def test_basic_retrieval_chain_multiple_vectorstores(three_vectorstores):
+def test_basic_retrieval_chain_multiple_vectorstores(three_vectorstores, retry_config):
     chain = basic_retrieval_chain(
-        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=three_vectorstores, model_name=GROQ_PRODUCTION_MODELS[0]
+        prompt=CHAT_PROMPT_TEMPLATE,
+        vectorstore=three_vectorstores,
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        retry_config=retry_config,
     )
     output = chain.invoke("What is the Higgs boson?")
 
@@ -73,7 +82,7 @@ def test_basic_retrieval_chain_multiple_vectorstores(three_vectorstores):
     )
 
 
-def test_advanced_rag_chain(twiki_vectorstore):
+def test_advanced_rag_chain(twiki_vectorstore, retry_config):
     # TODO: this is required currently requried for search to work
     search_kwargs = {
         "k_text": 3,
@@ -83,7 +92,10 @@ def test_advanced_rag_chain(twiki_vectorstore):
     }
 
     chain = advanced_rag(
-        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=twiki_vectorstore, model_name=GROQ_PRODUCTION_MODELS[0]
+        prompt=CHAT_PROMPT_TEMPLATE,
+        vectorstore=twiki_vectorstore,
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        retry_config=retry_config,
     )
     output = chain.invoke(
         {"question": "What is the Higgs boson?", "search_kwargs": search_kwargs},
@@ -113,7 +125,7 @@ def test_advanced_rag_chain(twiki_vectorstore):
     assert isinstance(output["docs"][0], Document)
 
 
-def test_advanced_rag_chain_multiple_vectorstores(three_vectorstores):
+def test_advanced_rag_chain_multiple_vectorstores(three_vectorstores, retry_config):
     search_kwargs = {
         "k_text": 3,
         "k": 15,
@@ -122,7 +134,10 @@ def test_advanced_rag_chain_multiple_vectorstores(three_vectorstores):
     }
 
     chain = advanced_rag(
-        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=three_vectorstores, model_name=GROQ_PRODUCTION_MODELS[0]
+        prompt=CHAT_PROMPT_TEMPLATE,
+        vectorstore=three_vectorstores,
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        retry_config=retry_config,
     )
     output = chain.invoke(
         {"question": "What is the Higgs boson?", "search_kwargs": search_kwargs},