Applying RAG for Wazuh Documentation: Complete Step-by-Step Guide#

1
# Use the base image with Python 3.9
2
FROM python:3.9
3

4
# Set the working directory
5
WORKDIR /app
6

7
# Copy the dependencies to the /tmp/requirements.txt folder
8
COPY wazuh-documentation/requirements.txt /tmp/requirements.txt
9

10
# Install the dependencies
11
RUN pip install -r /tmp/requirements.txt
12

13
CMD ["sleep", "infinity"]

1
services:
2
  wazuh-docs:
3
    build: .
4
    volumes:
5
      - ./wazuh-documentation:/app/wazuh-documentation

1
Running Sphinx v4.5.0
2
loading pickled environment... done
3
building [mo]: targets for 0 po files that are out of date
4
building [singlehtml]: all documents
5
updating environment: 0 added, 0 changed, 0 removed
6
looking for now-outdated files... none found
7
preparing documents... done
8
writing... done
9
writing additional files... done
10
copying static files... done
11
copying extra files... done
12
dumping object inventory... done
13
build succeeded.

1
NAME                    ID              SIZE      MODIFIED
2
llama3.2:latest         a80c4f17acd5    2.0 GB    2 hours ago
3
nomic-embed-text:latest 0a109f422b47    274 MB    2 hours ago

1
chromadb==0.6.3
2
unstructured==0.16.14
3
langchain==0.3.18
4
langchain-text-splitters==0.3.6
5
unstructured[all-docs]
6
langchain-community==0.3.14
7
langchain-ollama==0.2.2

1
#!/usr/bin/env python3
2
"""
3
Wazuh Documentation RAG Upload Script
4
Author: Anubhav Gain
5
Description: Uploads PDF documentation to vector database for RAG
6
"""
7

8
import argparse
9
import os
10
from pathlib import Path
11

12
from langchain_community.document_loaders import UnstructuredPDFLoader
13
from langchain_ollama import OllamaEmbeddings
14
from langchain_text_splitters import RecursiveCharacterTextSplitter
15
from langchain_community.vectorstores import Chroma
16

17

18
def upload(document_path, model_name, collection_name, ollama_base_url):
19
    """
20
    Upload PDF documentation to ChromaDB vector store
21

22
    Args:
23
        document_path (str): Path to PDF file
24
        model_name (str): Ollama embedding model name
25
        collection_name (str): ChromaDB collection name
26
        ollama_base_url (str): Ollama server URL
27

28
    Returns:
29
        Chroma: Vector store object
30
    """
31
    # Get current script directory
32
    current_path = os.path.dirname(os.path.realpath(__file__))
33
    chroma_persistent_directory = current_path + "/data"
34

35
    # Create data directory if it doesn't exist
36
    if not os.path.exists(chroma_persistent_directory):
37
        os.makedirs(chroma_persistent_directory, exist_ok=True)
38
        print(f"Created directory: {chroma_persistent_directory}")
39

40
    # Validate PDF file exists
41
    if not os.path.exists(document_path):
42
        raise FileNotFoundError(f"PDF file not found: {document_path}")
43

44
    print(f"Loading PDF: {document_path}")
45

46
    # Load PDF document
47
    loader = UnstructuredPDFLoader(file_path=document_path)
48
    data = loader.load()
49

50
    print(f"Loaded {len(data)} document(s)")
51

52
    # Split text into chunks
53
    # chunk_size: 7500 characters per chunk
54
    # chunk_overlap: 100 characters overlap between chunks
55
    text_splitter = RecursiveCharacterTextSplitter(
56
        chunk_size=7500,
57
        chunk_overlap=100
58
    )
59
    chunks = text_splitter.split_documents(data)
60

61
    print(f"Split into {len(chunks)} chunks")
62

63
    # Create vector embeddings and store in ChromaDB
64
    print(f"Creating embeddings using model: {model_name}")
65
    print(f"This may take several minutes...")
66

67
    vector = Chroma.from_documents(
68
        documents=chunks,
69
        embedding=OllamaEmbeddings(
70
            base_url=ollama_base_url,
71
            model=model_name,
72
            show_progress=True
73
        ),
74
        collection_name=collection_name,
75
        persist_directory=chroma_persistent_directory
76
    )
77

78
    print(f"✓ Successfully uploaded to collection: {collection_name}")
79
    print(f"✓ Vector store saved to: {chroma_persistent_directory}")
80

81
    return vector
82

83

84
if __name__ == '__main__':
85
    # Create argument parser
86
    parser = argparse.ArgumentParser(
87
        description='Upload PDF documentation to vector store for RAG'
88
    )
89

90
    # Add arguments
91
    parser.add_argument(
92
        '-p', '--path',
93
        type=str,
94
        help='Path to the PDF file',
95
        required=True
96
    )
97
    parser.add_argument(
98
        '-m', '--model',
99
        type=str,
100
        help='Name of the Ollama model for embedding',
101
        default='nomic-embed-text'
102
    )
103
    parser.add_argument(
104
        '-n', '--name',
105
        type=str,
106
        help='Collection name in ChromaDB',
107
        default='wazuh'
108
    )
109
    parser.add_argument(
110
        '-b', '--base-url',
111
        type=str,
112
        help='Base URL for the Ollama server',
113
        default='http://127.0.0.1:11434'
114
    )
115

116
    # Parse arguments
117
    args = parser.parse_args()
118

119
    print("=" * 60)
120
    print("WAZUH DOCUMENTATION RAG UPLOAD")
121
    print("=" * 60)
122
    print(f"PDF Path: {args.path}")
123
    print(f"Embedding Model: {args.model}")
124
    print(f"Collection Name: {args.name}")
125
    print(f"Ollama URL: {args.base_url}")
126
    print("=" * 60)
127

128
    # Call the upload function
129
    try:
130
        upload(
131
            document_path=args.path,
132
            model_name=args.model,
133
            collection_name=args.name,
134
            ollama_base_url=args.base_url
135
        )
136
        print("\n✓ Upload completed successfully!")
137
    except Exception as e:
138
        print(f"\n✗ Error during upload: {str(e)}")
139
        import traceback
140
        traceback.print_exc()
141
        exit(1)

1
current_path = os.path.dirname(os.path.realpath(__file__))

1
chroma_persistent_directory = current_path + "/data"
2
if not os.path.exists(chroma_persistent_directory):
3
    os.makedirs(chroma_persistent_directory, exist_ok=True)

1
loader = UnstructuredPDFLoader(file_path=document_path)
2
data = loader.load()

1
text_splitter = RecursiveCharacterTextSplitter(chunk_size=7500, chunk_overlap=100)
2
chunks = text_splitter.split_documents(data)

1
vector = Chroma.from_documents(
2
    documents=chunks,
3
    embedding=OllamaEmbeddings(
4
        base_url=ollama_base_url,
5
        model=model_name,
6
        show_progress=True
7
    ),
8
    collection_name=collection_name,
9
    persist_directory=chroma_persistent_directory
10
)

1
return vector

1
#!/usr/bin/env python3
2
"""
3
Wazuh Documentation RAG Query Script
4
Author: Anubhav Gain
5
Description: Query Wazuh documentation using RAG
6
"""
7

8
import argparse
9
import os
10

11
from langchain.retrievers import MultiQueryRetriever
12
import chromadb
13
from langchain_core.output_parsers import StrOutputParser
14
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
15
from langchain_core.runnables import RunnablePassthrough
16
from langchain_ollama import ChatOllama
17
from langchain_ollama import OllamaEmbeddings
18
from langchain_chroma import Chroma
19

20

21
def ask_ollama(question, collection_name='wazuh', embedding_model='nomic-embed-text',
22
               local_model='llama3.2'):
23
    """
24
    Query the RAG system with a question about Wazuh
25

26
    Args:
27
        question (str): User's question
28
        collection_name (str): ChromaDB collection name
29
        embedding_model (str): Embedding model name
30
        local_model (str): LLM model for generation
31

32
    Returns:
33
        str: Generated answer based on documentation
34
    """
35
    # Get current script path
36
    current_path = os.path.dirname(os.path.realpath(__file__))
37
    chroma_persistent_directory = current_path + "/data"
38

39
    # Check if data directory exists
40
    if not os.path.exists(chroma_persistent_directory):
41
        raise FileNotFoundError(
42
            f"Vector database not found at {chroma_persistent_directory}. "
43
            "Please run upload.py first to create the database."
44
        )
45

46
    print(f"Loading vector database from: {chroma_persistent_directory}")
47

48
    # Initialize embedding model
49
    embedding = OllamaEmbeddings(model=embedding_model)
50

51
    # Connect to persistent ChromaDB
52
    persistent_client = chromadb.PersistentClient(path=chroma_persistent_directory)
53

54
    # Load vector database
55
    vector_db = Chroma(
56
        client=persistent_client,
57
        collection_name=collection_name,
58
        embedding_function=embedding,
59
    )
60

61
    print(f"Loaded collection: {collection_name}")
62

63
    # Initialize LLM for generation
64
    load_ollama = ChatOllama(model=local_model)
65

66
    # Multi-query prompt template
67
    # This generates multiple versions of the question for better retrieval
68
    prompt_template = PromptTemplate(
69
        input_variables=["question"],
70
        template="""You are an AI language model assistant. Your task is to generate 2
71
different versions of the given user question to retrieve relevant documents from
72
a vector database. By generating multiple perspectives on the user question, your
73
goal is to help the user overcome some of the limitations of the distance-based
74
similarity search. Provide these alternative questions separated by newlines.
75
Original question: {question}""",
76
    )
77

78
    # Answer generation template
79
    template = """Answer the question based ONLY on the following context:
80
{context}
81

82
Question: {question}
83

84
Provide a detailed, accurate answer based on the Wazuh documentation.
85
If you cannot find the answer in the context, say so clearly."""
86

87
    # Create multi-query retriever
88
    # This improves retrieval by generating alternative phrasings
89
    retriever = MultiQueryRetriever.from_llm(
90
        vector_db.as_retriever(),
91
        load_ollama,
92
        prompt=prompt_template
93
    )
94

95
    print("Searching documentation...")
96

97
    # Create chat prompt
98
    prompt = ChatPromptTemplate.from_template(template)
99

100
    # Build RAG chain
101
    # This creates a pipeline: retrieve → format → generate → parse
102
    chain = (
103
        {"context": retriever, "question": RunnablePassthrough()}
104
        | prompt
105
        | load_ollama
106
        | StrOutputParser()
107
    )
108

109
    print("Generating answer...\n")
110

111
    # Invoke the chain with the question
112
    return chain.invoke(question)
113

114

115
if __name__ == '__main__':
116
    parser = argparse.ArgumentParser(
117
        description='Ask questions about Wazuh documentation'
118
    )
119

120
    parser.add_argument(
121
        '-q', '--question',
122
        type=str,
123
        help='The question to ask about Wazuh',
124
        required=True
125
    )
126
    parser.add_argument(
127
        '-c', '--collection',
128
        type=str,
129
        help='ChromaDB collection name',
130
        default='wazuh'
131
    )
132
    parser.add_argument(
133
        '-e', '--embedding',
134
        type=str,
135
        help='Embedding model name',
136
        default='nomic-embed-text'
137
    )
138
    parser.add_argument(
139
        '-m', '--model',
140
        type=str,
141
        help='LLM model name',
142
        default='llama3.2'
143
    )
144

145
    args = parser.parse_args()
146

147
    print("=" * 60)
148
    print("WAZUH DOCUMENTATION RAG QUERY")
149
    print("=" * 60)
150
    print(f"Question: {args.question}")
151
    print(f"Collection: {args.collection}")
152
    print(f"Embedding Model: {args.embedding}")
153
    print(f"LLM Model: {args.model}")
154
    print("=" * 60 + "\n")
155

156
    try:
157
        answer = ask_ollama(
158
            question=args.question,
159
            collection_name=args.collection,
160
            embedding_model=args.embedding,
161
            local_model=args.model
162
        )
163

164
        print("ANSWER:")
165
        print("=" * 60)
166
        print(answer)
167
        print("=" * 60)
168

169
    except Exception as e:
170
        print(f"\n✗ Error: {str(e)}")
171
        import traceback
172
        traceback.print_exc()
173
        exit(1)