chATLAS_Chains 0.1.6__tar.gz → 0.2.0__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.6
+Version: 0.2.0
 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,14 @@ 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

{chatlas_chains-0.1.6 → chatlas_chains-0.2.0}/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,14 @@ 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
@@ -119,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.6 → chatlas_chains-0.2.0}/chATLAS_Chains/chains/advanced.py

@@ -23,13 +23,17 @@ 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__)
 
@@ -40,27 +44,25 @@ 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,
-    fallback_models: list[str] | None = None,
+    skip_generation: bool = False,
 ) -> CompiledStateGraph:
     """
     Advanced Agentic RAG graph with optional query rewriting, dual-stage reranking and self-evaluation.
@@ -68,20 +70,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
-    :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
 
@@ -95,17 +94,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, fallback_models)
+    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,
-            retry_config=retry_config,
-            fallback_models=fallback_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 ---------------
 
@@ -209,7 +209,12 @@ 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)
@@ -223,50 +228,62 @@ def advanced_rag(
     # --------------- Build the graph ---------------
     graph = StateGraph(HybridGraphState)
 
-    # Add all the nodes, but don't link to them if not using
     graph.add_node("query_rewrite", query_rewriting)
     graph.add_node("retrieval", retrieval)
     graph.add_node("rrf", rrf)
     graph.add_node("rerank", rerank)
-    graph.add_node("generate", generate_answer)
-    # graph.add_node("assess", assess_answer)
-    # graph.add_node("refine", refine_answer)
+    if not skip_generation:
+        graph.add_node("generate", generate_answer)
 
     if enable_query_rewriting:
-        # rewrite the query first
         graph.add_edge("query_rewrite", "retrieval")
         graph.set_entry_point("query_rewrite")
     else:
-        # start with retrieval
         graph.set_entry_point("retrieval")
 
-    if not enable_rrf and not enable_reranking:
-        # no document processing, go straight to generation
-        graph.add_edge("retrieval", "generate")
-
-    elif enable_rrf and not enable_reranking:
-        # retrieve → rrf → generate
+    # Determine the last pre-generation node
+    if enable_rrf and enable_reranking:
+        last_node = "rerank"
         graph.add_edge("retrieval", "rrf")
-        graph.add_edge("rrf", "generate")
-
-    elif not enable_rrf and enable_reranking:
-        # retrieve → rerank → generate
+        graph.add_edge("rrf", "rerank")
+    elif enable_rrf:
+        last_node = "rrf"
+        graph.add_edge("retrieval", "rrf")
+    elif enable_reranking:
+        last_node = "rerank"
         graph.add_edge("retrieval", "rerank")
-        graph.add_edge("rerank", "generate")
-
     else:
-        # retrieve → rrf → rerank → generate
-        graph.add_edge("retrieval", "rrf")
-        graph.add_edge("rrf", "rerank")
-        graph.add_edge("rerank", "generate")
+        last_node = "retrieval"
 
-    # generate at the end
-    graph.add_edge("generate", END)
+    if skip_generation:
+        graph.add_edge(last_node, END)
+    else:
+        graph.add_edge(last_node, "generate")
+        graph.add_edge("generate", END)
 
-    # Compile the graph
     return graph.compile()
 
 
+def build_generation_prompt(prompt: str | None = None) -> ChatPromptTemplate:
+    """Build the prompt template used by the generation step."""
+    if prompt is None:
+        prompt = CHAT_PROMPT_TEMPLATE
+    return ChatPromptTemplate.from_template(prompt)
+
+
+def stream_generate_answer(state: HybridGraphState, model, prompt_template: ChatPromptTemplate):
+    """Generator that streams LLM tokens for the generation step."""
+    prompt_input = {
+        "context": combine_documents(state["docs"]),
+        "question": state["question"],
+        "chat_history": state.get("chat_history", ""),
+    }
+    final_prompt = prompt_template.format_messages(**prompt_input)
+    for chunk in model.stream(final_prompt):
+        if chunk.content:
+            yield chunk.content
+
+
 if __name__ == "__main__":
     from chATLAS_Chains.llm.model_selection import GROQ_PRODUCTION_MODELS
     from chATLAS_Chains.vectorstore import get_vectorstore
@@ -274,11 +291,6 @@ if __name__ == "__main__":
     twiki = get_vectorstore("twiki_prod")
     mkdocs = get_vectorstore("mkdocs_prod_v1")
 
-    retry_config = RetryConfig(
-        max_retries=1,
-        max_delay=120.0,
-    )
-
     # Create the hybrid graph
     graph = advanced_rag(
         vectorstore=[twiki, mkdocs],

{chatlas_chains-0.1.6 → chatlas_chains-0.2.0}/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.6 → chatlas_chains-0.2.0}/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(
         {