chATLAS_Chains 0.1.5__tar.gz → 0.1.7__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.5
+Version: 0.1.7
 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
@@ -19,9 +19,9 @@ Requires-Dist: langgraph
 Requires-Dist: sentence-transformers>=3.0.0
 Requires-Dist: tiktoken
 Requires-Dist: pinecone
-Requires-Dist: torch==2.2.1
 Requires-Dist: chatlas-embed
 Requires-Dist: psycopg2-binary>=2.9.10
+Requires-Dist: httpx[socks]>=0.28.1
 Dynamic: license-file
 
 
@@ -62,14 +62,71 @@ export CHATLAS_GROQ_KEY="your groq api key"
 
 **note** The API address is local to the CERN network. If not at CERN, you can forward it like so:
 ```bash
-ssh -L 3000:cs-513-ml003:3000 <LXPLUS_USERNAME>@lxplus.cern.ch
+ssh -L 3000:cs-513-ml003:3000 $LXPLUS_USERNAME@lxplus.cern.ch
 export CHATLAS_GROQ_BASE_URL="http://localhost:3000"
 ```
 
-## Available Chains
-- chains.basic.basic_retrieval_chain
-- chains.basic_graph.basic_retrieval_graph
-- chains.advanced.advanced_rag
+3. Using LLMs via CERN's LiteLLM API, here is the [repo](https://gitlab.cern.ch/itgpt/litellm-okd/-/tree/main) and some [setup instructions](https://codimd.web.cern.ch/tQKiMa13Q4O-EJXWTO3N7w?view#Using-Your-Dedicated-API-Key-to-Access-LLMs) for reference.
+```bash
+export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"
+```
+
+## Supported Chains
+
+More details [here](chATLAS_Chains/chains/README.md)
+
+- `chains.basic.basic_retrieval_chain`
+- `chains.advanced.advanced_rag`
+
+### Model Configuration in Chains
+
+Supported chain constructors now accept a typed `chat_model_kwargs` argument for model options (for example:
+`temperature`, `max_tokens`, `service_provider`, `api_key`, `base_url`, `proxy`).
+
+```python
+from chATLAS_Chains.chains.basic import basic_retrieval_chain
+
+chain = basic_retrieval_chain(
+    prompt=...,
+    vectorstore=...,
+    model_name="gpt-4o-mini",
+    chat_model_kwargs={"temperature": 0.1, "max_tokens": 512},
+)
+```
+
+## Forwarding vectorstore connections
+
+If not on the CERN network, you can forward the connection to the postgres servers with:
+
+```bash
+ssh -N \         
+  -L 6624:dbod-chatlas.cern.ch:6624 \
+  -L 6606:dbod-chatlas-cds.cern.ch:6606 \
+  "$LXPLUS_USERNAME"@lxplus.cern.ch 
+export CHATLAS_PORT_FORWARDING=1
+```
+
+You can then the helper function [`get_vectorstore`](chATLAS_Chains/vectorstore.py)
+
+## Testing Environment Variables
+
+Some tests are DB-backed integration tests (`tests/test_chains.py`, `tests/test_conversational.py`, `tests/test_search.py`).
+If the DB/test environment is not configured, these tests are skipped by `tests/conftest.py`.
+
+`tests/conftest.py` now uses explicit controls:
+
+- `CHATLAS_PORT_FORWARDING`: enable localhost DB tunnels (`1`, `true`, `True`)
+- `CHATLAS_DB_PASSWORD`
+
+### Local Example (with DB tunnels)
+
+```bash
+export CHATLAS_DB_PASSWORD="..."
+export CHATLAS_PORT_FORWARDING=1
+unset GITLAB_PAT
+
+uv run pytest -q
+```
 
 ## Postgres
 
@@ -91,6 +148,20 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.7
+
+Support for CERN-hosted LiteLLM models
+
+Multi-turn conversational RAG with (local) conversation history 
+
+Bugfixes
+
+#### 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.7}/README.md

@@ -36,14 +36,71 @@ export CHATLAS_GROQ_KEY="your groq api key"
 
 **note** The API address is local to the CERN network. If not at CERN, you can forward it like so:
 ```bash
-ssh -L 3000:cs-513-ml003:3000 <LXPLUS_USERNAME>@lxplus.cern.ch
+ssh -L 3000:cs-513-ml003:3000 $LXPLUS_USERNAME@lxplus.cern.ch
 export CHATLAS_GROQ_BASE_URL="http://localhost:3000"
 ```
 
-## Available Chains
-- chains.basic.basic_retrieval_chain
-- chains.basic_graph.basic_retrieval_graph
-- chains.advanced.advanced_rag
+3. Using LLMs via CERN's LiteLLM API, here is the [repo](https://gitlab.cern.ch/itgpt/litellm-okd/-/tree/main) and some [setup instructions](https://codimd.web.cern.ch/tQKiMa13Q4O-EJXWTO3N7w?view#Using-Your-Dedicated-API-Key-to-Access-LLMs) for reference.
+```bash
+export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"
+```
+
+## Supported Chains
+
+More details [here](chATLAS_Chains/chains/README.md)
+
+- `chains.basic.basic_retrieval_chain`
+- `chains.advanced.advanced_rag`
+
+### Model Configuration in Chains
+
+Supported chain constructors now accept a typed `chat_model_kwargs` argument for model options (for example:
+`temperature`, `max_tokens`, `service_provider`, `api_key`, `base_url`, `proxy`).
+
+```python
+from chATLAS_Chains.chains.basic import basic_retrieval_chain
+
+chain = basic_retrieval_chain(
+    prompt=...,
+    vectorstore=...,
+    model_name="gpt-4o-mini",
+    chat_model_kwargs={"temperature": 0.1, "max_tokens": 512},
+)
+```
+
+## Forwarding vectorstore connections
+
+If not on the CERN network, you can forward the connection to the postgres servers with:
+
+```bash
+ssh -N \         
+  -L 6624:dbod-chatlas.cern.ch:6624 \
+  -L 6606:dbod-chatlas-cds.cern.ch:6606 \
+  "$LXPLUS_USERNAME"@lxplus.cern.ch 
+export CHATLAS_PORT_FORWARDING=1
+```
+
+You can then the helper function [`get_vectorstore`](chATLAS_Chains/vectorstore.py)
+
+## Testing Environment Variables
+
+Some tests are DB-backed integration tests (`tests/test_chains.py`, `tests/test_conversational.py`, `tests/test_search.py`).
+If the DB/test environment is not configured, these tests are skipped by `tests/conftest.py`.
+
+`tests/conftest.py` now uses explicit controls:
+
+- `CHATLAS_PORT_FORWARDING`: enable localhost DB tunnels (`1`, `true`, `True`)
+- `CHATLAS_DB_PASSWORD`
+
+### Local Example (with DB tunnels)
+
+```bash
+export CHATLAS_DB_PASSWORD="..."
+export CHATLAS_PORT_FORWARDING=1
+unset GITLAB_PAT
+
+uv run pytest -q
+```
 
 ## Postgres
 
@@ -65,6 +122,20 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 0.1.7
+
+Support for CERN-hosted LiteLLM models
+
+Multi-turn conversational RAG with (local) conversation history 
+
+Bugfixes
+
+#### 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
@@ -113,4 +184,4 @@ chATLAS_Benchmark is released under Apache v2.0 license.
 
 *For questions and support, please [contact](mailto:joseph.caimin.egan@cern.ch)*
 
-</div>
+</div>

{chatlas_chains-0.1.5 → chatlas_chains-0.1.7}/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
@@ -22,13 +23,19 @@ 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.llm.model_selection import (
+    GROQ_PRODUCTION_MODELS,
+    ChatModelKwargs,
+    get_chat_model,
+    sanitize_chat_model_kwargs,
+)
 from chATLAS_Chains.prompt.starters import CHAT_PROMPT_TEMPLATE
 from chATLAS_Chains.query.query_rewriting import rewrite_query
 from chATLAS_Chains.search.basic import search_runnable
 from chATLAS_Chains.utils.doc_utils import combine_documents
-from chATLAS_Embed.Base import VectorStore
+from chATLAS_Embed.VectorStores import VectorStore
+
+logger = logging.getLogger(__name__)
 
 
 # Define TypedDict for the simplified state
@@ -37,26 +44,24 @@ class HybridGraphState(TypedDict, total=False):
     search_kwargs: dict
     docs: list[Document]
     answer: str
+    chat_history: str  # Optional
 
 
 def advanced_rag(
     vectorstore: VectorStore | list[VectorStore],
     model_name: str,
     prompt: str | None = None,
-    max_tokens: int | None = None,
-    temperature: float = 0.1,
-    use_preview_models: bool = False,
+    chat_model_kwargs: ChatModelKwargs | None = None,
     enable_query_rewriting: bool = False,
     enable_rrf: bool = False,
     enable_reranking: bool = False,
     # enable_self_evaluation: bool = False,
     query_rewriting_model: str = GROQ_PRODUCTION_MODELS[0],
-    query_rewriting_temperature: float = 0.1,
+    query_rewriting_chat_model_kwargs: ChatModelKwargs | None = None,
     rerank_model: str = "cohere-rerank-3.5",
     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.
@@ -64,19 +69,17 @@ def advanced_rag(
     :param prompt: The prompt template to use for the language model. If None, uses chATLAS_Chains.prompt.starters.CHAT_PROMPT_TEMPLATE
     :param vectorstore: Single vectorstore instance or list of vectorstore instances to search
     :param model_name: The name of the language model to use for generating responses.
-    :param max_tokens: Maximum number of tokens to generate in the response. If None, uses the model's default value.
-    :param temperature: Temperature to use for the model.
-    :param use_preview_models: If True, allows the use of preview models from Groq.
+    :param chat_model_kwargs: Optional kwargs passed through to ``get_chat_model`` for the main model.
     :param enable_query_rewriting: Whether to enable LLM-powered query rewriting.
     :param enable_rrf: Whether to enable RRF (Reciprocal Rank Fusion) for combining results from multiple vectorstores.
     :param enable_reranking: Whether to rerank the retrieved results using the Pinecone API.
     :param query_rewriting_model: Name of the model to use for query rewriting.
-    :param query_rewriting_temperature: Temperature to use for query rewriting.
+    :param query_rewriting_chat_model_kwargs: Optional kwargs for the query rewriting model.
+        Inheritance order is: ``chat_model_kwargs`` -> ``{"temperature": 0.1}`` defaults -> explicit query rewriting kwargs.
     :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.
     :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
 
@@ -90,15 +93,18 @@ def advanced_rag(
     # Parallel searcher for vectorstore(s)
     searcher = search_runnable(vectorstore)
 
+    primary_model_kwargs = sanitize_chat_model_kwargs(chat_model_kwargs)
+
     # Initialise models
-    model = get_chat_model(model_name, max_tokens, temperature, use_preview_models, retry_config)
+    model = get_chat_model(model_name, **primary_model_kwargs)
 
     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,
-        )
+        query_model_kwargs: ChatModelKwargs = {
+            **primary_model_kwargs,
+            "temperature": 0.1,
+        }
+        query_model_kwargs.update(sanitize_chat_model_kwargs(query_rewriting_chat_model_kwargs))
+        query_rewriting_model_instance = get_chat_model(model_name=query_rewriting_model, **query_model_kwargs)
 
     # --------------- Define functions for the graph nodes ---------------
 
@@ -107,7 +113,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 +134,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 +142,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 +156,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 +172,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 +183,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 +197,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,9 +207,13 @@ 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"]}
+        prompt_input = {
+            "context": combine_documents(state["docs"]),
+            "question": state["question"],
+            "chat_history": state.get("chat_history", ""),
+        }
+
         final_prompt = prompt_template.format_messages(**prompt_input)
 
         response = model.invoke(final_prompt)
@@ -266,42 +276,28 @@ if __name__ == "__main__":
     from chATLAS_Chains.vectorstore import get_vectorstore
 
     twiki = get_vectorstore("twiki_prod")
-
-    retry_config = RetryConfig(
-        max_retries=1,
-        max_delay=120.0,
-    )
+    mkdocs = get_vectorstore("mkdocs_prod_v1")
 
     # 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,
+    )
+
+    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"],
+            },
+        }
     )
 
-    # 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
-
-        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.7}/chATLAS_Chains/chains/basic.py

@@ -1,24 +1,19 @@
 from operator import itemgetter
-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.llm.model_selection import ChatModelKwargs, get_chat_model, sanitize_chat_model_kwargs
 from chATLAS_Chains.search.basic import search_runnable
 from chATLAS_Chains.utils.doc_utils import combine_documents
-from chATLAS_Embed.Base import VectorStore
+from chATLAS_Embed.VectorStores import VectorStore
 
 
 def basic_retrieval_chain(
     prompt: str,
     vectorstore: VectorStore | list[VectorStore],
     model_name: str,
-    max_tokens: int | None = None,
-    temperature: float | None = None,
-    use_preview_models: bool = False,
-    retry_config: RetryConfig | None = None,
+    chat_model_kwargs: ChatModelKwargs | None = None,
 ) -> RunnableSerializable:
     """
     Baseline RAG retrieval chain. Searches one or several vectorstores in parallel, passes retrieved documents to the model
@@ -29,21 +24,16 @@ def basic_retrieval_chain(
     :type vectorstore: Any
     :param model_name: The name of the chat model to use for generating responses.
     :type model_name: str
-    :param max_tokens: The maximum number of tokens to generate in the response. Defaults to None.
-    :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
+    :param chat_model_kwargs: Optional kwargs passed through to ``get_chat_model``.
+    :type chat_model_kwargs: ChatModelKwargs | 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, use_preview_models, retry_config)
+    sanitized_kwargs = sanitize_chat_model_kwargs(chat_model_kwargs)
+    model = get_chat_model(model_name, **sanitized_kwargs)
 
     search = search_runnable(vectorstore)
 
@@ -72,7 +62,8 @@ if __name__ == "__main__":
     chain = basic_retrieval_chain(
         prompt=CHAT_PROMPT_TEMPLATE,
         vectorstore=[twiki_vectorstore, mkdocs_vectorstore],
-        model_name=GROQ_PRODUCTION_MODELS[0],
+        model_name="openai/gpt-oss-120b",
+        # chat_model_kwargs={"service_provider": "groq"},
         # model_name="meta-llama/llama-4-maverick-17b-128e-instruct",
         # model_name="gemma2-9b-it",
         # model_name="qwen-qwq-32b",

{chatlas_chains-0.1.5 → chatlas_chains-0.1.7}/chATLAS_Chains/chains/basic_graph.py

@@ -31,8 +31,6 @@ def basic_retrieval_graph(
     model_name: str,
     max_tokens: int | None = None,
     temperature: float | None = None,
-    use_preview_models: bool = False,
-    retry_config: RetryConfig | None = None,
 ) -> CompiledStateGraph:
     """
     Baseline RAG retrieval graph using LangGraph. Searches one or several vectorstores,
@@ -47,7 +45,7 @@ def basic_retrieval_graph(
         A LangGraph graph that can be executed for RAG
     """
     # Initialize the model and prompt template
-    model = get_chat_model(model_name, use_preview_models=use_preview_models)
+    model = get_chat_model(model_name)
     prompt_template = ChatPromptTemplate.from_template(prompt)
 
     # Create a list of retrievers from the vectorstore(s)
@@ -164,9 +162,7 @@ if __name__ == "__main__":
     # model_name="gemma2-9b-it" # smaller context window, to check the doc truncation
     # model_name = "gpt-4o-mini"
 
-    graph = basic_retrieval_graph(
-        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=vs, model_name=model_name, use_preview_models=True
-    )
+    graph = basic_retrieval_graph(prompt=CHAT_PROMPT_TEMPLATE, vectorstore=vs, model_name=model_name)
 
     result = graph.invoke(
         {