Build RAG using Milvus and Ollama

{
	"id": 2,
	"vector": array([ 1.92813469e-02,  ..... , -1.79081468e-02], 
	"text": "Debian is an operating system and a distribution of Free Software.", 
	"subject": "Linux"
}

mkdir -p rag/aaps-doc-repo
git clone https://github.com/openaps/AndroidAPSdocs.git rag/aaps-doc-repo

test@test:~/rag/aaps-doc-repo/docs$ du -hs *
32K	androidaps.jpg
252K	androidaps-logo.png
4.2G	CROWDIN
78M	EN # <============ WE JUST NEED TO FOCUS ON THIS DIRECTORY
16K	favicon.ico
16K	shared.conf.py
2.7M	_static
8.0K	_templates

curl -fsSL https://ollama.com/install.sh | sh

docker run -d -v ollama:/root/.ollama -p 11434:11434 --name ollama ollama/ollama

ollama pull phi4-mini:latest
ollama pull mxbai-embed-large

from langchain_milvus import Milvus
from langchain_text_splitters import MarkdownHeaderTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_ollama import OllamaEmbeddings
from langchain_core.documents import Document
from glob import glob
import re

# Vars definitions
collection_name = "aaps_docs_rag"
chunks = [] # will store all dataset chunks
milvus_data = [] # will store data to be inserted into Milvus
chunk_size = 800
chunk_overlap = 100

## Define Langchain Mardown splitter
headers_to_split_on = [
    ("#", "Header 1"),
    ("##", "Header 2"),
    ("###", "Header 3"),
    ("####", "Header 4"),
]

# Initialize Ollama embedding model
embeddings = OllamaEmbeddings(
    model="mxbai-embed-large",
)

# Initialize Milvus DB and create a collection
vector_store = Milvus(
    collection_name=collection_name,
    embedding_function=embeddings,
    connection_args={"uri": "./rag_demo.db"},
    index_params={"index_type": "IVF_FLAT", "metric_type": "COSINE"},
    auto_id=True,  # Enable automatic ID generation
    drop_old=True,  # Drop existing collection if it exists. Use this for test only
)

# Define text splitters
markdown_header_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=headers_to_split_on, strip_headers=False
)

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=chunk_size,
    chunk_overlap=chunk_overlap,
    separators=[
        "\n### ", "\n## ", "\n# ",
        "\n- ", "\n* ", "\n\n",
        ". ", " ", ""
    ]
)

# Define function to clean up AAPS docs
def clean_markdown(text: str) -> str:
    # Remove Sphinx/Myst anchors like (anchor-name)=
    text = re.sub(r'^\([^)]+\)=\s*$', '', text, flags=re.MULTILINE)
    # Remove images
    text = re.sub(r'!\[.*?\]\(.*?\)', '', text)
    # Consolidate and improve code block removal (matches 3+ backticks)
    text = re.sub(r'`{3,}[\s\S]*?`{3,}', '[CODE BLOCK REMOVED]', text, flags=re.DOTALL)
    # Remove HTML tags
    text = re.sub(r'<[^>]+>', '', text)
    # Normalize multiple newlines
    text = re.sub(r'\n{2,}', '\n', text)
    return text.strip()


# Add chunks to the list which will be used to collect all vectors to stores in Milvus
def add_chunks(sub_chunks, metadata):
    for chunk in sub_chunks:
        if len(metadata.items()) > 1:
            items = list(metadata.items())
            chunk_metadata = {
                "section": items[0][1],
                "sub_section": items[1][1],
                "source_file": file_path
            }
        elif len(metadata.items()) == 1:
            key, value = list(metadata.items())[0]
            chunk_metadata = {
                "section": value,
                "sub_section": "",
                "source_file": file_path
            }
        else:
            chunk_metadata = {
                "section": "",
                "sub_section": "",
                "source_file": file_path
            }

        chunks.append(Document(
            page_content=chunk,
            metadata=chunk_metadata
        ))


# Collect and chunk Markdown files
for file_path in glob("aaps-doc-repo/docs/EN/**/*.md", recursive=True):
    print(f"Processing file: {file_path}")
    loader = TextLoader(file_path, encoding="utf-8")
    documents = loader.load()
    # Extract markdown text content and remove unwanted parts
    md_doc = clean_markdown(documents[0].page_content)
    # Split respecting markdown hierarchy
    md_header_splits = markdown_header_splitter.split_text(md_doc)
    for header_chunk in md_header_splits:
        # Further split each header chunk into smaller chunks, for long sentences
        if len(header_chunk.page_content) > chunk_size:
            sub_chunks = text_splitter.split_text(header_chunk.page_content)
            add_chunks(sub_chunks, header_chunk.metadata)
        else:
            add_chunks([header_chunk.page_content], header_chunk.metadata)


print(f"Total chunks created: {len(chunks)}")

# Store each document chunk in the DB, after generating embedding
print("Inserting documents into Milvus collection...")
batch_size = 500
for i in range(0, len(chunks), batch_size):
    batch = chunks[i:i + batch_size]
    vector_store.add_documents(documents=batch)
embedding = vector_store.add_documents(documents=chunks)

# Print Finished message
print(f"Inserted {len(chunks)} documents into Milvus collection '{collection_name}'")

cd rag/
python3 -m venv .venv
source .venv/bin/activate
pip3 install langchain langchain-milvus langchain-text-splitters langchain-ollama langchain-core langchain-community pymilvus[milvus-lite]
python3 gen_embeddings.py

(.venv) test@test:~/rag$ python3 gen_embeddings.py 
Processing file: aaps-doc-repo/docs/EN/index.md
[truncated for readability]
Total chunks created: 2343
Inserting documents into Milvus collection...
Inserted 2343 documents into Milvus collection 'aaps_docs_rag'

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_milvus import Milvus

# Load local Ollama model
llm = ChatOllama(
    model="phi4-mini:latest",
    temperature=0.2,
)

# Connect to existing Milvus Lite DB
vectorstore = Milvus(
    collection_name="aaps_docs_rag",
    connection_args={"uri": "./rag_demo.db"},
    embedding_function=OllamaEmbeddings(model="mxbai-embed-large"),
)

# Create a data retriever
retriever = vectorstore.as_retriever(
    search_type="similarity", # using similarity search
    search_kwargs={
        "k": 5, # number of documents to return
    }
)

question = "How can I restore my AAPS settings backup on a new phone?"

search_res = retriever.invoke(question)

# Display the retrieved documents
for i, doc in enumerate(search_res):
    print(f"Document {i+1}:")
    print(f"Content: {doc.page_content}")
    print(f"Metadata: {doc.metadata}")
    print("-" * 50)

$ python3 milvus_read.py 
Document 1:
Content: ### Reinstalling AAPS
When uninstalling **AAPS** you will lose all your settings, objectives and the current Pod session. **To restore them make sure you have a recent exported settings file available!**
When on an active Pod, make sure that you have an export for the current pod session or you will lose the currently active pod when importing older settings.
1. Export your settings and store a copy in a safe place (e.g Google Drive).
2. Uninstall **AAPS** and restart your phone.
3. Install the new version of **AAPS**.
4. Import your settings.
5. Verify all preferences (optionally import settings again).
6. Activate a new pod.
7. When done: Export current settings.
---
Metadata: {'pk': 462052157073658950, 'section': 'Omnipod DASH', 'source_file': 'aaps-doc-repo/docs/EN/CompatiblePumps/OmnipodDASH.md', 'sub_section': 'Troubleshooting'}
--------------------------------------------------
Document 2:
Content: ### Reinstalling AAPS
When uninstalling **AAPS** you will lose all your settings, objectives and the current Pod session. **To restore them make sure you have a recent exported settings file available!**
When on an active Pod, make sure that you have an export for the current pod session or you will lose the currently active pod when importing older settings.
1. Export your settings and store a copy in a safe place (e.g Google Drive).
2. Uninstall **AAPS** and restart your phone.
3. Install the new version of **AAPS**.
4. Import your settings.
5. Verify all preferences (optionally import settings again).
6. Activate a new pod.
7. When done: Export current settings.
---
Metadata: {'pk': 462051399138804716, 'section': 'Omnipod DASH', 'source_file': 'aaps-doc-repo/docs/EN/CompatiblePumps/OmnipodDASH.md', 'sub_section': 'Troubleshooting'}
--------------------------------------------------
Document 3:
Content: # Creating and restoring back-ups
When installing AAPS on your phone it becomes a "medical device" you rely on daily. It is highly recommended to have an
emergency backup plan for when your phone gets defective, stolen or lost. Therefore, it is essential to prepare by asking yourself, "What if?
To restore your AAPS setup to an existing or new phone, it's important to keep following items in a secure location (read: not on your phone).
Best practice is to keep at least two separate backups: on a local hard drive, USB stick and (preferred) on Cloud storage like Google Drive or
Microsoft 365 OneDrive
Metadata: {'pk': 462052157073659440, 'section': 'Creating and restoring back-ups', 'source_file': 'aaps-doc-repo/docs/EN/Maintenance/ExportImportSettings.md', 'sub_section': ''}
--------------------------------------------------
Document 4:
Content: # Creating and restoring back-ups
When installing AAPS on your phone it becomes a "medical device" you rely on daily. It is highly recommended to have an
emergency backup plan for when your phone gets defective, stolen or lost. Therefore, it is essential to prepare by asking yourself, "What if?
To restore your AAPS setup to an existing or new phone, it's important to keep following items in a secure location (read: not on your phone).
Best practice is to keep at least two separate backups: on a local hard drive, USB stick and (preferred) on Cloud storage like Google Drive or
Microsoft 365 OneDrive
Metadata: {'pk': 462051399138805206, 'section': 'Creating and restoring back-ups', 'source_file': 'aaps-doc-repo/docs/EN/Maintenance/ExportImportSettings.md', 'sub_section': ''}
--------------------------------------------------
Document 5:
Content: ## Export settings
- AAPS 2.7 uses a new encrypted backup format.
- You must [export your settings](ExportImportSettings.md) after updating to version 2.7.
- Settings files from previous versions can only be imported in AAPS 2.7. Export will be in new format.
- Make sure to store your exported settings not only on your phone but also in at least one safe place (your pc, cloud storage...).
- If you build AAPS 2.7 apk with the same keystore than in previous versions you can install new version without deleting the previous version.
- All settings as well as finished objectives will remain as they were in the previous version.
Metadata: {'pk': 462052157073659412, 'section': 'Necessary checks after update coming from AAPS 2.6', 'source_file': 'aaps-doc-repo/docs/EN/Maintenance/Update2_7.md', 'sub_section': 'Export settings'}
--------------------------------------------------

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_milvus import Milvus
import time
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

# Load local Ollama model
llm = ChatOllama(
    model="phi4-mini:latest",
    temperature=0.2,
)

# Connect to existing Milvus Lite DB
vectorstore = Milvus(
    collection_name="aaps_docs_rag",
    connection_args={"uri": "./rag_demo.db"},
    embedding_function=OllamaEmbeddings(model="mxbai-embed-large"),
)

# Create a data retriever
retriever = vectorstore.as_retriever(
    search_type="similarity", # using cosine similarity
    search_kwargs={
        "k": 7  # <-- Get the Top 7 most similar docs
    }
)

# Define the prompt for our RAG
RAG_PROMPT_TEMPLATE = ChatPromptTemplate.from_template("""
You are a helpful assistant specialized in Android APS (a.k.a. AAPS) documentation.
Use ONLY the following retrieved context to answer the user's question.
If the context does not contain the answer, state that you don't have enough information.
Cite the source file for your answer using (Source: <filename>).

Context:
{context}

Question:
{question}
""")

def format_docs(docs):
    """Prepares the retrieved documents for the LLM."""
    return "\n\n".join(
        (f"Source: {doc.metadata.get('source_file', 'Unknown')}\nContent: {doc.page_content}")
        for doc in docs
    )

# Create a RAG chain using LangChain Expression Language (LCEL)
rag_chain = (
    {"context": retriever 
    | RunnableLambda(format_docs), "question": RunnablePassthrough()}
    | RAG_PROMPT_TEMPLATE
    | llm
    | StrOutputParser()
)

# Run the chain with the query and stream the output
query = "How can I restore my AAPS settings backup on a new phone?"

print(f"Query: {query}\n")
print("Streaming answer:\n")

start_time = time.time()
full_response = ""
for chunk in rag_chain.stream(query):
    print(chunk, end="", flush=True)
    full_response += chunk

end_time = time.time()
print(f"\n\n--- Replied in {end_time - start_time:.2f} seconds ---")

python3 RAG.py

test@test:~/rag$ python3 RAG.py 
Query: How can I restore my AAPS settings backup on a new phone?

Streaming answer:

To restore your AAPS settings backup on a new phone, follow these instructions:

1. Install Android APK Manager and the latest version of AAPS.
2. Import all exported files from Google Drive or Microsoft 365 OneDrive (or any other cloud storage you used) to your computer using an app like MobaHoaxer for Windows PCs.

Source: aaps-doc-repo/docs/EN/Maintenance/ExportImportSettings.md

--- Replied in 76.67 seconds ---

from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_milvus import Milvus
from langchain_core.runnables import RunnableLambda
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
import gradio as gr

# Load local Ollama model
llm = ChatOllama(
    model="phi4-mini:latest",
    temperature=0.2,
)

# Connect to existing Milvus Lite DB
vectorstore = Milvus(
    collection_name="aaps_docs_rag",
    connection_args={"uri": "./rag_demo.db"},
    embedding_function=OllamaEmbeddings(model="mxbai-embed-large"),
)

# Create a data retriever
retriever = vectorstore.as_retriever(
    search_type="similarity", # using cosine similarity
    search_kwargs={
        "k": 7  # <-- Get the Top 7 most similar docs
    }
)

# Define the prompt for our RAG
RAG_PROMPT_TEMPLATE = ChatPromptTemplate.from_template("""
You are a helpful assistant specialized in Android APS (a.k.a. AAPS) documentation.
Use the following retrieved context to answer the user's question.
If the context does not contain the answer, state that you don't have enough information.
Cite the source file for your answer using (Source: <filename>).

Here is the conversation history:
{chat_history}                                                       

Context:
{context}

Question:
{question}
""")

def format_docs(docs):
    """Prepares the retrieved documents for the LLM."""
    return "\n\n".join(
        (f"Source: {doc.metadata.get('source_file', 'Unknown')}\nContent: {doc.page_content}")
        for doc in docs
    )

def format_history(history):
    if not history:
        return "No previous conversation."
    
    formatted = []
    for user_msg, bot_msg in history:
        formatted.append(f"User: {user_msg}")
        formatted.append(f"Assistant: {bot_msg}")
    return "\n".join(formatted)

# Create a RAG chain using LangChain Expression Language (LCEL)
rag_chain = (
    {
        "context": RunnableLambda(lambda inputs: inputs["question"]) | retriever | RunnableLambda(format_docs),
        "question": RunnableLambda(lambda inputs: inputs["question"]),
        "chat_history": RunnableLambda(lambda inputs: inputs["chat_history"])
    }
    | RAG_PROMPT_TEMPLATE
    | llm
    | StrOutputParser()
)

# function called by Gradio interface
def get_rag_response(message, history):
    
    #print(f"User query: {message}")

    # Format the history to pass it to the LLM
    formatted_chat_history = format_history(history)
    
    # We now pass a dictionary to the chain
    inputs = {
        "question": message,
        "chat_history": formatted_chat_history
    }
    
    # Use .stream() to get a streaming response
    response_stream = rag_chain.stream(inputs)
    #if len(inputs["chat_history"]) > 0:
    #    print(f"Chat history: {inputs['chat_history']}")
    
    full_response = ""
    # Yield each chunk as it comes in, accumulating the full response
    for chunk in response_stream:
        full_response += chunk
        yield full_response

gui_iface = gr.ChatInterface(
    fn=get_rag_response,
    title="Android APS (AAPS) RAG Chatbot",
    description="Ask questions about AAPS documentation. The assistant will answer using retrieved information.",
)

gui_iface.launch()

pip3 install gradio
python3 RAG.py