chATLAS_Chains

This package implements and benchmarks various Retrieval Augmented Generation (RAG) chains for use in the chATLAS project.

Installation

From PyPI

pip install chATLAS-Chains

From source

We recommend using uv

cd chATLAS_Chains
uv sync

Environment variables

These are required for the following use cases

1. Using an OpenAI LLM

export CHATLAS_OPENAI_KEY="your api key"

2. Using LLMs via the Groq API

export CHATLAS_GROQ_BASE_URL="http://cs-513-ml003:3000"
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:

ssh -L 3000:cs-513-ml003:3000 $LXPLUS_USERNAME@lxplus.cern.ch
export CHATLAS_GROQ_BASE_URL="http://localhost:3000"

3. Using LLMs via CERN's LiteLLM API, here is the repo and some setup instructions for reference.

export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"

Supported Chains

More details here

• chains.basic.basic_retrieval_chain
• chains.advanced.advanced_rag

Prompt Routing

chATLAS_Chains.router.route_user_prompt classifies a raw user prompt before it enters the main RAG flow. The router defaults to CERN LiteLLM and returns a validated RouterDecision with:

• route: one of hep_atlas, person_lookup, dangerous, out_of_scope, or unknown
• confidence: score from 0 to 1
• normalized_query: cleaned prompt for downstream use
• search_kwargs: optional search parameters
• metadata: route-specific parameters

For implementation details and a function-by-function reference, see the frontend reviewer document: ../chATLAS_Frontend/docs/INTENT_ROUTER.md.

from chATLAS_Chains.router import route_user_prompt

decision = route_user_prompt(
    "Who is Jane Doe?",
    chat_model_kwargs={
        "service_provider": "litellm",
        "proxy": "socks5h://localhost:1080",  # optional when outside CERN
    },
)

if decision.route == "dangerous":
    print("Apply the safety response policy")
elif decision.route == "person_lookup":
    print("Use the person lookup flow")
elif decision.route == "hep_atlas":
    print("Use the HEP/ATLAS search flow")
elif decision.route == "out_of_scope":
    print("Explain the chATLAS scope")
else:
    print("Ask for clarification")

The dedicated mc_request route is currently disabled until the MC workflow is ready. MC production and job-option prompts are classified as hep_atlas and continue through the normal selected workflow.

Router context uses a hybrid policy: the current user prompt is authoritative, and recent conversation context is only a bounded disambiguation aid for follow-ups such as "What about Run 3?". The router keeps recent user prompts, includes only short assistant messages, and drops long assistant answers rather than summarizing them.

For live LiteLLM calls, set:

export CHATLAS_CHAINS_LITELLM_KEY="your litellm key"

If LiteLLM is unavailable or returns invalid JSON, the router falls back to deterministic rules by default.

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

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:

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

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)

export CHATLAS_DB_PASSWORD="..."
export CHATLAS_PORT_FORWARDING=1
unset GITLAB_PAT

uv run pytest -q

Postgres

If you want to create a local postgres server, you need to install psql. Some instructions to do this on macOS using homebrew are here:

Software install

brew install postgresql
brew services start postgresql
brew install pgvector
brew unlink pgvector && brew link pgvector

Create a user

psql -h localhost -U postgres
ALTER USER postgres WITH PASSWORD 'Set_your_password_here';
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

0.1.4

Support for Groq-hosted models

Some new functions that go beyond the "basic RAG" workflow:

• Reciprocal Rerank Fusion chATLAS_Chains.documents.rrf.reciprocal_rank_fusion
• Document reranking via the Pinecone API chATLAS_Chains.documents.rerank.rerank_documents
• Query rewriting step chATLAS_Chains.query.query_rewriting.rewrite_query

These are all usable via the new chain chATLAS_Chains.chains.advanced.advanced_rag

Added unit tests to gitlab CI/CD pipeline

0.1.3

Fixing imports

Changed output format of basic_retrieval_chain (docs key is now a list of Document objects, rather than a dict)

Unit tests for basic_retrieval_chain

0.1.2

Unit tests

First Langgraph chain

0.1.1

Initial Release

---

📄 License

chATLAS_Benchmark is released under Apache v2.0 license.

---