chATLAS_Chains 0.1.3__tar.gz → 0.1.4__tar.gz

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

@@ -1,22 +1,27 @@
 Metadata-Version: 2.4
 Name: chATLAS_Chains
-Version: 0.1.3
+Version: 0.1.4
 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
-Project-URL: Homepage, https://gitlab.cern.ch/belliot/chatlas-packages/
-Project-URL: Documentation, https://chatlas-packages.docs.cern.ch/chATLAS_Chain/
+Project-URL: Homepage, https://gitlab.cern.ch/atlasml/chatlas/chatlas-packages/
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License
 Classifier: Operating System :: OS Independent
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: chATLAS_Benchmark>=0.0.9
-Requires-Dist: chATLAS_Embed>=0.1.14
+Requires-Dist: chatlas-embed>=0.1.19
 Requires-Dist: langchain~=0.3.3
 Requires-Dist: langchain_core
 Requires-Dist: langchain_openai
+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
 Dynamic: license-file
 
 
@@ -26,10 +31,18 @@ This package implements and benchmarks various Retrieval Augmented Generation (R
 
 ## Installation
 
+### From PyPI
+
+```bash
+pip install chATLAS-Chains
+```
+
+### From source
+
+We recommend using [`uv`](https://docs.astral.sh/uv/)
 ```bash
-conda create -n venv chatlas_chains_env python=3.10
-conda activate chatlas_chains_env
-pip install chatlas-chains
+cd chATLAS_Chains
+uv sync
 ```
 
 ## Environment variables
@@ -38,30 +51,29 @@ These are required for the following use cases
 
 1. Using an OpenAI LLM
 ```bash
-export CHATLAS_OPENAI_KEY='your api key'
+export CHATLAS_OPENAI_KEY="your api key"
 ```
 
-2. Benchmarking, set the path to the question set
+2. Using LLMs via the Groq API
 ```bash
-export CHATLAS_BENCHMARK_QUESTIONS=/path/to/questions.josn
+export CHATLAS_GROQ_BASE_URL="http://cs-513-ml003:3000"
+export CHATLAS_GROQ_KEY="your groq api key"
 ```
 
-## Available Chains
-- chains.basic.basic_retrieval_chain
-- chains.basic_graph.basic_retrieval_graph
-
-## Benchmarking
-
-To benchmark e.g. the chains in `chATLAS_Chains.chains.basic` run this from the project root
+**note** The API address is local to the CERN network. If not at CERN, you can forward it like so:
 ```bash
-python benchmark/basic.py
+ssh -L 3000:cs-513-ml003:3000 <LXPLUS_USERNAME>@lxplus.cern.ch
+export CHATLAS_GROQ_BASE_URL="http://localhost:3000"
 ```
 
-## Testing
+## Available Chains
+- chains.basic.basic_retrieval_chain
+- chains.basic_graph.basic_retrieval_graph
+- chains.advanced.advanced_rag
 
-The tests require a running postgres server to work. If on lxplus you can modify `TEST_DB_CONFIG` in [tests/conftest.py](tests/conftest.py) to connect to the chATLAS server.
+## Postgres
 
-If you want to create a local dummy postgres server, you need to install `psql`. This can be done on macOS using [homebrew](https://brew.sh):
+If you want to create a local postgres server, you need to install `psql`. Some instructions to do this on macOS using [homebrew](https://brew.sh) are here:
 
 Software install
 ```bash
@@ -79,6 +91,19 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 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

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

@@ -5,10 +5,18 @@ This package implements and benchmarks various Retrieval Augmented Generation (R
 
 ## Installation
 
+### From PyPI
+
+```bash
+pip install chATLAS-Chains
+```
+
+### From source
+
+We recommend using [`uv`](https://docs.astral.sh/uv/)
 ```bash
-conda create -n venv chatlas_chains_env python=3.10
-conda activate chatlas_chains_env
-pip install chatlas-chains
+cd chATLAS_Chains
+uv sync
 ```
 
 ## Environment variables
@@ -17,30 +25,29 @@ These are required for the following use cases
 
 1. Using an OpenAI LLM
 ```bash
-export CHATLAS_OPENAI_KEY='your api key'
+export CHATLAS_OPENAI_KEY="your api key"
 ```
 
-2. Benchmarking, set the path to the question set
+2. Using LLMs via the Groq API
 ```bash
-export CHATLAS_BENCHMARK_QUESTIONS=/path/to/questions.josn
+export CHATLAS_GROQ_BASE_URL="http://cs-513-ml003:3000"
+export CHATLAS_GROQ_KEY="your groq api key"
 ```
 
-## Available Chains
-- chains.basic.basic_retrieval_chain
-- chains.basic_graph.basic_retrieval_graph
-
-## Benchmarking
-
-To benchmark e.g. the chains in `chATLAS_Chains.chains.basic` run this from the project root
+**note** The API address is local to the CERN network. If not at CERN, you can forward it like so:
 ```bash
-python benchmark/basic.py
+ssh -L 3000:cs-513-ml003:3000 <LXPLUS_USERNAME>@lxplus.cern.ch
+export CHATLAS_GROQ_BASE_URL="http://localhost:3000"
 ```
 
-## Testing
+## Available Chains
+- chains.basic.basic_retrieval_chain
+- chains.basic_graph.basic_retrieval_graph
+- chains.advanced.advanced_rag
 
-The tests require a running postgres server to work. If on lxplus you can modify `TEST_DB_CONFIG` in [tests/conftest.py](tests/conftest.py) to connect to the chATLAS server.
+## Postgres
 
-If you want to create a local dummy postgres server, you need to install `psql`. This can be done on macOS using [homebrew](https://brew.sh):
+If you want to create a local postgres server, you need to install `psql`. Some instructions to do this on macOS using [homebrew](https://brew.sh) are here:
 
 Software install
 ```bash
@@ -58,6 +65,19 @@ CREATE EXTENSION IF NOT EXISTS vector;
 ```
 ## CHANGELOG
 
+#### 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

chatlas_chains-0.1.4/chATLAS_Chains/chains/advanced.py

@@ -0,0 +1,299 @@
+"""
+More advanced RAG workflow with optional query rewriting, reciprocal rank fusion and reranking
+
+Stages:
+- (Optional) Query Rewriting - Correct typos and enhance query clarity using LLM
+- Retrieval - Retrieve documents using BM25 and vector search
+- (Optional) Reciprocal Rank Fusion - Combine results from retrieval modes (e.g. text and vector)
+upweighting results that appear in both
+- (Optional) Reranking - Rerank documents using Pinecone API cross-encoder model
+- Answer Generation - Generate answer using LLM with retrieved context
+"""
+
+import argparse
+import os
+import sys
+from typing import TypedDict
+
+from langchain_core.documents import Document
+from langchain_core.prompts import ChatPromptTemplate
+from langgraph.graph import END, StateGraph
+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.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
+from chATLAS_Chains.search.basic import search_runnable
+from chATLAS_Chains.utils.doc_utils import combine_documents
+from chATLAS_Embed.Base import VectorStore
+
+
+# Define TypedDict for the simplified state
+class HybridGraphState(TypedDict, total=False):
+    question: str
+    search_kwargs: dict
+    docs: list[Document]
+    answer: str
+
+
+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,
+    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,
+    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,
+) -> CompiledStateGraph:
+    """
+    Advanced Agentic RAG graph with optional query rewriting, dual-stage reranking and self-evaluation.
+
+    :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 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 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.
+
+    :return: A compiled LangGraph with the chosen features
+
+    """
+    if prompt is None:
+        prompt = CHAT_PROMPT_TEMPLATE
+
+    # Create prompt template
+    prompt_template = ChatPromptTemplate.from_template(prompt)
+
+    # Parallel searcher for vectorstore(s)
+    searcher = search_runnable(vectorstore)
+
+    # Initialise models
+    model = get_chat_model(model_name, max_tokens, temperature, use_preview_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,
+        )
+
+    # --------------- Define functions for the graph nodes ---------------
+
+    def query_rewriting(state: HybridGraphState) -> HybridGraphState:
+        """Rewrite the query to correct typos and enhance clarity."""
+        question = state["question"]
+
+        if not enable_query_rewriting:
+            print("[WARNING] Query rewriting disabled, but query_rewriting node was called.")
+            # return original state
+            return {**state}
+
+        rewritten_query = rewrite_query(question, model=query_rewriting_model_instance)
+
+        return {
+            **state,
+            "question": rewritten_query,
+            "unchanged_question": question,  # Keep original for reference
+        }
+
+    def retrieval(state: HybridGraphState) -> HybridGraphState:
+        """Call the search runnable to retrieve documents"""
+
+        question = state.get("question")
+        if question is None:
+            raise Exception("question field is None")
+
+        search_kwargs = state.get("search_kwargs", {})
+        if not search_kwargs:
+            print("[WARNING] No search_kwargs provided, using defaults.")
+            results = searcher.invoke(question)
+        else:
+            results = searcher.invoke(question, config={"metadata": {"search_kwargs": search_kwargs}})
+
+        if "docs" not in results:
+            raise Exception("Search results missing 'docs' field")
+
+        print(f"Retrieved {len(results['docs'])} documents")
+
+        return {
+            **state,
+            "docs": results["docs"],
+        }
+
+    # Define the document reranking function
+    def rerank(state: HybridGraphState) -> HybridGraphState:
+        """
+        Rerank the parent documents using the Pinecone API.
+        """
+
+        if not enable_reranking:
+            print("[WARNING] Reranking disabled, but rerank node was called.")
+            # return original state
+            return {**state}
+
+        docs = state.get("docs", [])
+
+        try:
+            reranked_docs = rerank_documents(
+                question=state["question"],
+                docs=docs,
+                reranker_model=rerank_model,
+                api_key=pinecone_api_key,
+                # num_return_docs = None # return everything
+            )
+
+        except Exception as e:
+            print(f"[WARNING] Reranking failed with exception {e}. Returning original documents.")
+            reranked_docs = docs
+
+        return {
+            **state,
+            "docs": reranked_docs,
+        }
+
+    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.")
+            # return original state
+            return {**state}
+
+        docs = state.get("docs", [])
+        if not docs:
+            print("[WARNING] No documents retrieved for RRF.")
+            return {**state}
+
+        try:
+            rrf_docs = reciprocal_rank_fusion(
+                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.")
+            rrf_docs = docs
+
+        return {
+            **state,
+            "docs": rrf_docs,
+        }
+
+    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)
+
+        response = model.invoke(final_prompt)
+        answer = response.content
+
+        return {
+            **state,
+            "answer": answer,
+        }
+
+    # --------------- 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 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
+        graph.add_edge("retrieval", "rrf")
+        graph.add_edge("rrf", "generate")
+
+    elif not enable_rrf and enable_reranking:
+        # retrieve → rerank → generate
+        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")
+
+    # generate at the end
+    graph.add_edge("generate", END)
+
+    # Compile the graph
+    return graph.compile()
+
+
+if __name__ == "__main__":
+    from chATLAS_Chains.llm.model_selection import GROQ_PRODUCTION_MODELS
+    from chATLAS_Chains.vectorstore import get_vectorstore
+
+    twiki = get_vectorstore("twiki_prod")
+
+    # Create the hybrid graph
+    graph = advanced_rag(
+        vectorstore=[twiki],
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        enable_query_rewriting=True,
+        enable_rrf=True,
+        enable_reranking=True,
+    )
+
+    # 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)

chatlas_chains-0.1.4/chATLAS_Chains/chains/basic.py

@@ -0,0 +1,82 @@
+from operator import itemgetter
+from typing import Optional
+
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnableSerializable
+
+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
+from chATLAS_Embed.Base import VectorStore
+
+
+def basic_retrieval_chain(
+    prompt: str,
+    vectorstore: VectorStore | list[VectorStore],
+    model_name: str,
+    max_tokens: int | None = None,
+    temperature: float | None = None,
+) -> RunnableSerializable:
+    """
+    Baseline RAG retrieval chain. Searches one or several vectorstores in parallel, passes retrieved documents to the model
+
+    :param prompt: The prompt template to use with the model.
+    :type prompt: str
+    :param vectorstore: The vectorstore or list of vectorstores to search over.
+    :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
+
+    :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)
+
+    search = search_runnable(vectorstore)
+
+    final_inputs = {
+        "context": lambda x: combine_documents(x["docs"]),
+        "question": itemgetter("question"),
+    }
+
+    answer = {
+        "answer": final_inputs | prompt_template | model,
+        "docs": lambda x: x["docs"],
+    }
+
+    chain = search | answer
+    return chain
+
+
+if __name__ == "__main__":
+    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
+
+    twiki_vectorstore = get_vectorstore("twiki_prod")
+    mkdocs_vectorstore = get_vectorstore("mkdocs_prod_v1")
+
+    chain = basic_retrieval_chain(
+        prompt=CHAT_PROMPT_TEMPLATE,
+        vectorstore=[twiki_vectorstore, mkdocs_vectorstore],
+        model_name=GROQ_PRODUCTION_MODELS[0],
+        # model_name="meta-llama/llama-4-maverick-17b-128e-instruct",
+        # model_name="gemma2-9b-it",
+        # model_name="qwen-qwq-32b",
+        # model_name="mistral-saba-24b",
+    )
+    SEARCH_HYPERPARAMS = {"k": 5, "k_text": 0, "date_filter": "01-01-2010"}
+
+    result = chain.invoke("What is the Higgs boson?", config={"metadata": {"search_kwargs": SEARCH_HYPERPARAMS}})
+
+    print(f"Answer: {result['answer'].content}")
+    print(f"Number of documents retrieved: {len(result['docs'])}")
+
+    for doc in result["docs"]:
+        print(doc.metadata.get("source"))

{chatlas_chains-0.1.3 → chatlas_chains-0.1.4}/chATLAS_Chains/chains/basic_graph.py

@@ -4,14 +4,15 @@ Example graph for running langgraph with this general setup and the postgres vec
 
 from typing import TypedDict
 
-import langgraph.graph as lg
 from langchain_core.documents import Document
 from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import END, StateGraph
+from langgraph.graph.state import CompiledStateGraph
 
+from chATLAS_Chains.llm.groq import RetryConfig
 from chATLAS_Chains.llm.model_selection import get_chat_model
-from chATLAS_Chains.utils.doc_utils import combine_documents
-from chATLAS_Embed import LangChainVectorStore
+from chATLAS_Chains.utils.doc_utils import combine_documents, truncate_to_context_window
+from chATLAS_Embed.LangChainVectorStore import LangChainVectorStore
 
 
 # Define TypedDict for the state
@@ -24,7 +25,15 @@ class GraphState(TypedDict, total=False):
     answer: str
 
 
-def basic_retrieval_graph(prompt: str, vectorstore, model_name: str) -> lg.Graph:
+def basic_retrieval_graph(
+    prompt: str,
+    vectorstore,
+    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,
     passes retrieved documents to the model, and returns the final answer.
@@ -38,7 +47,7 @@ def basic_retrieval_graph(prompt: str, vectorstore, model_name: str) -> lg.Graph
         A LangGraph graph that can be executed for RAG
     """
     # Initialize the model and prompt template
-    model = get_chat_model(model_name)
+    model = get_chat_model(model_name, use_preview_models=use_preview_models)
     prompt_template = ChatPromptTemplate.from_template(prompt)
 
     # Create a list of retrievers from the vectorstore(s)
@@ -84,8 +93,16 @@ def basic_retrieval_graph(prompt: str, vectorstore, model_name: str) -> lg.Graph
     # Define the document processing function
     def process_docs(state: GraphState) -> GraphState:
         """Process the merged documents into a context string."""
-        docs = state["merged_docs"]
-        context = combine_documents(docs)
+        initial_docs = state["merged_docs"]
+
+        truncated_docs = truncate_to_context_window(
+            docs=initial_docs, question=state["question"], prompt=prompt_template, model_name=model_name
+        )
+
+        if len(truncated_docs) < len(initial_docs):
+            print(f"Truncated from {len(initial_docs)} to {len(truncated_docs)} docs to fit context window.")
+
+        context = combine_documents(truncated_docs)
         return {
             "question": state["question"],
             "retrieved_docs": state["retrieved_docs"],
@@ -136,23 +153,26 @@ def basic_retrieval_graph(prompt: str, vectorstore, model_name: str) -> lg.Graph
 
 
 if __name__ == "__main__":
-    # Example of how to run the graph correctly
-    import os
+    # Example usage
+    from chATLAS_Chains.prompt.starters import CHAT_PROMPT_TEMPLATE
+    from chATLAS_Chains.vectorstore import get_vectorstore
 
-    os.environ["CHATLAS_EMBEDDING_MODEL_PATH"] = "<PATH TO YOUR EMBEDDING MODEL>"
-    os.environ["CHATLAS_OPENAI_KEY"] = "YOUR OPENAI API KEY"
-    os.environ["CHATLAS_DB_PASSWORD"] = "<>"
+    vs = get_vectorstore("twiki_prod")
 
-    from ..prompt.starters import CHAT_PROMPT_TEMPLATE
-    from ..vectorstore import vectorstore
+    # model_name = "llama-3.3-70b-versatile"
+    model_name = "deepseek-r1-distill-llama-70b"
+    # 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=vectorstore, model_name="gpt-4o-mini")
+    graph = basic_retrieval_graph(
+        prompt=CHAT_PROMPT_TEMPLATE, vectorstore=vs, model_name=model_name, use_preview_models=True
+    )
 
-    ans = graph.invoke(
+    result = graph.invoke(
         {
-            "question": "How many onions are in ATLAS",
+            "question": "What is the crack veto for electron reconstruction",
             "search_kwargs": {
-                "k_text": 3,
+                "k_text": 5,
                 "k": 10,
                 "date_filter": "01-01-2010",
                 "type": ["CDS", "twiki", "Indico"],
@@ -160,4 +180,5 @@ if __name__ == "__main__":
         }
     )
 
-    print(ans)
+    print("Num retrieved documents: ", len(result["merged_docs"]))
+    print("Answer: ", result["answer"])