Enable/Disable Knowledge Base

This guide will walk you through setting up a knowledge base toggle in the pipeline that allows users to dynamically choose whether to use your knowledge base or rely on the model's general knowledge.

This tutorial extends the Your First RAG Pipeline tutorial. Ensure you have followed the instructions to set up your repository and index your data.

Prerequisites

This tutorial specifically requires:

  • Completion of the Your First RAG Pipeline tutorial.

  • All setup steps listed on the Prerequisites page.

  • An Elastic Search vector data store that is already set up and available for use. Refer to supported-vector-data-store for tutorial.

Why Do We Need Knowledge Base Toggle?

This specific toggle allows users to dynamically control whether their queries should be processed using your knowledge base (RAG approach) or answered using the language model's general knowledge alone.

What is a Toggle?

The toggle functionality builds upon the concept of conditional steps introduced in the semantic routing tutorial. While conditional steps route queries to different handlers based on query content, a toggle step provides users with direct control over the pipeline behavior through a simple boolean parameter.

Key differences:

  • Conditional Steps: Automatic routing based on query analysis

  • Toggle Steps: User-controlled routing via explicit parameters

When setting up a toggle pipeline, ensure that any state dependent on the output of a toggled component(s) is defined beforehand to prevent errors. You can follow steps in Initialize States with Default Values section.

Installation

# you can use a Conda environment
pip install --extra-index-url "https://oauth2accesstoken:$(gcloud auth print-access-token)@glsdk.gdplabs.id/gen-ai-internal/simple/" gllm-rag gllm-core gllm-generation gllm-inference gllm-pipeline gllm-retrieval gllm-misc gllm-datastore

Set Up Your Project

We'll build upon the pipeline you created in the Your First RAG Pipeline tutorial. Make sure you have that working before proceeding.

Prepare your repository

Let’s prepare your workspace step by step.

1

Go to the repository you use for Your First RAG Pipeline:

cd my-rag-pipeline
2

Prepare your .env file:

Ensure you have a file named .env in your project directory with the following content:

CSV_DATA_PATH="data/imaginary_animals.csv"
ELASTICSEARCH_URL="http://localhost:9200/"
EMBEDDING_MODEL="text-embedding-3-small"
LANGUAGE_MODEL="gpt-4o-mini"
INDEX_NAME="first-quest"
OPENAI_API_KEY="<YOUR_OPENAI_API_KEY>"

This is an example .env file. You may adjust the variables according to your need.

Adjust Folder Structure

Extend your existing project structure to include the toggle pipeline:

my-rag-pipeline/
├── data/
│   ├── imaginary_animals.csv
├── modules/
│   ├── __init__.py
│   ├── retriever.py
│   ├── repacker.py
│   ├── response_synthesizer.py    # 👈 Will be modified
├── indexer.py
├── pipeline.py
├── toggle_pipeline.py             # 👈 New (alternatively you can modify the pipeline.py)
└── main.py                        # 👈 Will be modified

Index Your Data

Ensure you have your data indexed. If not, you should follow steps in Index Your Data before proceeding.

Build Core Components of Your Pipeline

Modify the Response Synthesizer

The response synthesizer needs to handle both knowledge-base and general knowledge modes.

1

Update the system prompt

Modify your modules/response_synthesizer.py to handle the toggle functionality:

def response_synthesizer_component() -> StuffResponseSynthesizer:
    """Response synthesizer component for generating answers from context.

    Returns:
        StuffResponseSynthesizer: An instance for generating responses from context.
    """
    SYSTEM_PROMPT = """
    You are a helpful assistant that recommends.

    Use knowledge base is {use_knowledge_base}.

    When user decided to use knowledge base, you should:
    - Use only the information provided in the context below to answer the user's question. You may infer simple, logical conclusions based on the context, but do not introduce new facts or external knowledge.
    - You may suggest or summarize the options listed in the context if the question asks for recommendations, ideas, or what to do.
    - If the context does not contain enough information to answer the user's question, respond with:
    "Sorry, I don't have enough information to answer that."

    When user decided to not use knowledge base, you should answer the user's question based on your general knowledge.

    Context:
    {context}
    """
    USER_PROMPT = "Question: {query}"

Key features:

  • Dynamic behavior: The prompt adapts based on the use_knowledge_base setting

  • Knowledge base mode: Restricts responses to context information only

  • General knowledge mode: Uses the model's training data

  • Fallback handling: Clear instructions for insufficient context

2

Keep the rest unchanged

The rest of the response synthesizer remains the same:

Build the Pipeline

Now let's create the pipeline that enables dynamic knowledge base switching.

1

Create the toggle pipeline file

Create toggle_pipeline.py with the necessary imports:

from gllm_pipeline.pipeline.states import RAGState
from gllm_pipeline.pipeline.pipeline import Pipeline
from gllm_pipeline.steps import bundle, step, toggle

from modules import (
    repacker_component,
    response_synthesizer_component,
    retriever_component,
)
2

Define the extended state

Create a custom state that includes the toggle setting:

class ToggleState(RAGState):
    use_knowledge_base: bool

This extends the default RAGState to include a boolean field that controls whether the knowledge base should be used.

3

Create component instances

Instantiate your existing components:

retriever = retriever_component()
repacker = repacker_component(mode="context")
response_synthesizer = response_synthesizer_component()

These are the same components from your original pipeline, ensuring consistency.

4

Define the individual pipeline steps

Create the standard pipeline steps that will be conditionally executed:

retriever_step = step(
    retriever,
    {"query": "user_query"},
    "chunks",
    {"top_k": "top_k"},
)

repacker_step = step(
    repacker,
    {"chunks": "chunks"},  # input variables
    "context",  # output variable
)

bundler_step = bundle(
    ["context", "use_knowledge_base"],
    "response_synthesis_bundle",
)

response_synthesizer_step = step(
    response_synthesizer,
    {"query": "user_query", "state_variables": "response_synthesis_bundle"},
    "response",
)

Key points:

  • Same steps as the original pipeline

  • The bundler now includes use_knowledge_base to pass the toggle state to the response synthesizer

5

Create the knowledge base toggle step

This is where the magic happens - the toggle step conditionally executes the retrieval and repacking:

knowledge_base_toggle_step = toggle(
    condition=lambda x: x["use_knowledge_base"],
    if_branch=[retriever_step, repacker_step],  # Both steps together
)

How it works:

  • condition: A lambda function that checks the use_knowledge_base field from the state

  • if_branch: A list of steps to execute when the condition is True

  • When False: The toggle step does nothing, leaving context empty (which we'll handle in the response synthesizer)

6

Compose the final pipeline

Connect all steps into the complete toggle pipeline:

e2e_pipeline_with_knowledge_base_toggle = knowledge_base_toggle_step | bundler_step | response_synthesizer_step

e2e_pipeline_with_knowledge_base_toggle.state_type = ToggleState

Pipeline flow:

  1. Toggle Step: Conditionally retrieves and repacks context

  2. Bundler Step: Packages context and toggle setting

  3. Response Synthesizer Step: Generates response based on available context and toggle setting

Modify the Application Code

Here we will update the main.py file. You can find the updated file below.

1KB
main.py

The breakdown of the modification is as follows.

Update Pipeline Import

1

Update the import in main.py

Add the import for the new toggle pipeline:

from toggle_pipeline import e2e_pipeline_with_knowledge_base_toggle
2

Update the pipeline execution

Modify the run_pipeline function to use the toggle pipeline:

async def run_pipeline(state: dict, config: dict):
   ...
    try:
        await event_emitter.emit("Starting pipeline")
        await e2e_pipeline_with_knowledge_base_toggle.invoke(state, config) # Change to new pipeline
    ...

Initialize States with Default Values

Since context becomes optional with the toggle functionality, we need to initialize it properly in our state.

1

Update the request handler

Modify the /stream endpoint to handle the toggle parameter and initialize context:

async def add_message(request: Request):
    ...
    state = {
        "user_query": user_query,
        "event_emitter": event_emitter,
        "use_knowledge_base": use_knowledge_base,
        "context": "",  # Initialize context as empty string
    }
    config = {"top_k": top_k, "debug": debug}

    asyncio.create_task(run_pipeline(state, config))
    return StreamingResponse(stream_handler.stream())

Key changes:

  • New parameter: use_knowledge_base with default value True

  • Context initialization: Start with empty context string

  • State enhancement: Include the toggle setting in the initial state

Run Your Application

Now let's test the toggle functionality with different configurations.

1

Start your server

Run your FastAPI server as before:

poetry run uvicorn main:app --reload
2

Test with knowledge base enabled

Try this query with the knowledge base enabled:

{
  "user_query": "Which animal lives in the forest?",
  "top_k": 5,
  "debug": true,
  "use_knowledge_base": true
}

Expected behavior:

  • The pipeline will retrieve information from your imaginary_animals.csv

  • You'll see retrieval and repacking steps in the debug logs

  • The response will be based on your knowledge base (imaginary animals)

3

Test with knowledge base disabled

Try the same query with the knowledge base disabled:

{
  "user_query": "Which animal lives in the forest?",
  "top_k": 5,
  "debug": true,
  "use_knowledge_base": false
}

Expected behavior:

  • The pipeline will skip retrieval and repacking steps

  • No chunks will be retrieved from your database

  • The response will be based on the model's general knowledge (real animals)

4

Compare the responses and debug output

You should notice clear differences:

With Knowledge Base (true):

Starting pipeline
[Start 'BasicVectorRetriever'] Processing input:
    - query: 'Which animal lives in the forest?'
    - top_k: 5
    - event_emitter: <gllm_core.event.event_emitter.EventEmitter object at ...xx>
[Finished 'BasicVectorRetriever'] Successfully retrieved 5 chunks.
...
[Start 'Repacker'] Repacking 5 chunks.
[Finished 'Repacker'] Successfully repacked chunks: ...
[Start 'StuffResponseSynthesizer'] Processing query: 'Which animal lives in the forest?'
[Finished 'StuffResponseSynthesizer'] Successfully synthesized response: "Based on the knowledge base, creatures like the Crystal Stag and Mystic Wolf live in enchanted forests..."
Finished pipeline

Without Knowledge Base (false):

Starting pipeline
[Start 'StuffResponseSynthesizer'] Processing query: 'Which animal lives in the forest?'
[Finished 'StuffResponseSynthesizer'] Successfully synthesized response: "Many real animals live in forests, including deer, bears, wolves, foxes, and various bird species..."
Finished pipeline
5

Verify debug output

With debug: true, you should see logs showing:

  • Whether the toggle step was executed or skipped

  • The presence or absence of retrieval operations

  • The context content (or lack thereof)

  • The different response styles

Understanding the Flow

Here's what happens in each mode:

Knowledge Base Enabled (use_knowledge_base: true)

  1. Toggle Check: Condition evaluates to true

  2. Retrieval: Searches your knowledge base for relevant information

  3. Repacking: Assembles retrieved chunks into context

  4. Bundling: Packages context and toggle setting

  5. Response: Generates answer using knowledge base information

Knowledge Base Disabled (use_knowledge_base: false)

  1. Toggle Check: Condition evaluates to false

  2. Skip Steps: Retrieval and repacking are bypassed

  3. Empty Context: Context remains as initialized empty string

  4. Bundling: Packages empty context and toggle setting

  5. Response: Generates answer using model's general knowledge

Extending the Toggle System

Adding Multiple Toggle Options

You can extend this pattern for multiple toggle options:

class ExtendedToggleState(RAGState):
    use_knowledge_base: bool
    use_web_search: bool
    use_technical_docs: bool

Conditional Toggle Chains

Create more complex conditional logic:

knowledge_base_toggle = toggle(
    condition=lambda x: x["use_knowledge_base"],
    if_branch=[retriever_step, repacker_step],
)

web_search_toggle = toggle(
    condition=lambda x: x["use_web_search"] and not x["use_knowledge_base"],
    if_branch=[web_search_step, web_repacker_step],
)

Dynamic Toggle Decisions

Make toggle decisions based on query analysis:

smart_toggle_step = toggle(
    condition=lambda x: should_use_knowledge_base(x["user_query"]),
    if_branch=[retriever_step, repacker_step],
)

Troubleshooting

Common Issues

  1. Context not being skipped properly:

    • Ensure context is initialized as empty string in the state

    • Verify the toggle condition is evaluating correctly

    • Check that the bundler includes both context and toggle setting

  2. Response synthesizer not adapting:

    • Confirm the system prompt includes the toggle logic

    • Verify use_knowledge_base is being passed to the response synthesizer

    • Check that the prompt template variables match your state keys

  3. Toggle parameter not being received:

    • Ensure the request JSON includes the use_knowledge_base field

    • Verify the FastAPI endpoint is extracting the parameter correctly

    • Check that the parameter is being added to the initial state

📂 Complete Tutorial Files

Coming soon!

Congratulations! You've successfully implemented a knowledge base toggle pipeline. This flexible system gives users control over whether to use your curated knowledge base or tap into the model's general knowledge, making your application more versatile and user-friendly.

PreviousYour First RAG PipelineNextImplement Semantic Routing

Last updated