rag_pipeline.py

"""
rag_pipeline.py

This module is responsible for:
1. Retrieving relevant documents from ChromaDB.
2. Creating context from the retrieved documents.
3. Preparing the prompt for the LLM.
4. Sending the prompt to the LLM.
5. Returning the final response.
"""

from langchain_core.prompts import PromptTemplate
from langchain_ollama import ChatOllama
from config import LLM_MODEL


def generate_response(retriever, question):
    """
    Generate a response using the RAG pipeline.

    Args:
        retriever: LangChain retriever object
        question (str): User question

    Returns:
        str: Response generated by the LLM
    """

    # Step 1: Retrieve relevant documents from ChromaDB
    retrieved_docs = retriever.invoke(question)

    print("\n" + "=" * 60)
    print(f"Retrieved Documents: {len(retrieved_docs)}")

    for i, doc in enumerate(retrieved_docs):
        print(f"\nDocument {i+1}")
        print("Source:", doc.metadata.get("source", "Unknown"))
        print("Preview:")
        print(doc.page_content[:300])

    print("=" * 60)

    # Step 2: Combine all retrieved document text into one context
    context = "\n\n".join(
        doc.page_content for doc in retrieved_docs
    )

    # Step 3: Create the Prompt Template
    prompt = PromptTemplate(
        input_variables=["context", "question"],
        template="""
You are an expert DevOps Engineer specializing in AWS, Kubernetes,
Docker, Linux, Terraform, Jenkins, CI/CD and Python.

Use ONLY the information provided in the given context to answer the user's question.

If the answer is not available in the context, reply exactly:

"I don't have enough information in the knowledge base."

Do not make up any information.

Context:
{context}

Question:
{question}

Answer:
"""
    )

    # Step 4: Replace placeholders with actual values
    formatted_prompt = prompt.format(
        context=context,
        question=question
    )

    print("\nPrompt Preview")
    print("=" * 60)
    print(formatted_prompt[:1000])
    print("=" * 60)

    # Step 5: Initialize the LLM
    llm = ChatOllama(model=LLM_MODEL)

    # Step 6: Generate the response
    response = llm.invoke(formatted_prompt)

    # Step 7: Return only the response text
    return response.content