RAG System Design Document: Markdown-Based Knowledge Retrieval¶

from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler
import os
class MarkdownWatcher(FileSystemEventHandler):
    def init(self, processor):
        self.processor = processor
    def on_created(self, event):
        if event.is_directory or not event.src_path.endswith('.md'):
            return
        self.processor.process_file(event.src_path)
    def on_modified(self, event):
        if event.is_directory or not event.src_path.endswith('.md'):
            return
        self.processor.process_file(event.src_path)
    def on_deleted(self, event):
        if event.is_directory or not event.src_path.endswith('.md'):
            return
        self.processor.delete_file(event.src_path)
class DirectoryWatcherService:
    def init(self, directory_path, processor):
        self.observer = Observer()
        self.directory_path = directory_path
        self.event_handler = MarkdownWatcher(processor)
    def start(self):
        self.observer.schedule(self.event_handler, self.directory_path, recursive=True)
        self.observer.start()
    def stop(self):
        self.observer.stop()
        self.observer.join()

from langchain_community.document_loaders import TextLoader
from langchain.text_splitter import MarkdownTextSplitter, RecursiveCharacterTextSplitter
import frontmatter
import hashlib
class DocumentProcessor:
    def init(self, embedding_service, vectorstore_service, metadata_store):
        self.embedding_service = embedding_service
        self.vectorstore_service = vectorstore_service
        self.metadata_store = metadata_store
        self.text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
            chunk_size=512,
            chunk_overlap=50,
            separators=["\n## ", "\n### ", "\n#### ", "\n", " ", ""]
        )
    def process_file(self, file_path):
        # Generate document hash for tracking changes
        file_hash = self._hash_file(file_path)
        # Check if document has changed
        existing_hash = self.metadata_store.get_document_hash(file_path)
        if existing_hash == file_hash:
            print(f"No changes detected for {file_path}")
            return
        # Extract content and metadata
        with open(file_path, 'r') as f:
            content = f.read()
        try:
            post = frontmatter.loads(content)
            metadata = post.metadata
            content_text = post.content
        except:
            metadata = {}
            content_text = content
        # Add file path and hash to metadata
        metadata['source'] = file_path
        metadata['file_hash'] = file_hash
        # Split the document
        docs = self.text_splitter.create_documents([content_text], [metadata])
        # Delete old vectors if they exist
        if existing_hash:
            self.vectorstore_service.delete_document(file_path)
        # Generate embeddings and store
        self.embedding_service.embed_documents(docs)
        self.vectorstore_service.add_documents(docs)
        # Update metadata store
        self.metadata_store.update_document_metadata(file_path, metadata, file_hash)
    def delete_file(self, file_path):
        self.vectorstore_service.delete_document(file_path)
        self.metadata_store.delete_document(file_path)
    def process_all_files(self, directory_path):
        for root, , files in os.walk(directorypath):
            for file in files:
                if file.endswith('.md'):
                    file_path = os.path.join(root, file)
                    self.process_file(file_path)
    def hashfile(self, file_path):
        with open(file_path, 'rb') as f:
            return hashlib.md5(f.read()).hexdigest()

from langchain_ollama import OllamaEmbeddings
import numpy as np
class EmbeddingService:
    def init(self, model_name="mxbai-embed-large", batch_size=32):
        self.embeddings_model = OllamaEmbeddings(
            model=model_name,
            base_url="http://localhost:11434"
        )
        self.batch_size = batch_size
    def embed_documents(self, documents):
        """Generate embeddings for a list of documents"""
        texts = [doc.page_content for doc in documents]
        # Process in batches to avoid memory issues
        all_embeddings = []
        for i in range(0, len(texts), self.batch_size):
            batch_texts = texts[i:i+self.batch_size]
            batch_embeddings = self.embeddings_model.embed_documents(batch_texts)
            all_embeddings.extend(batch_embeddings)
        # Add embeddings to document objects
        for i, doc in enumerate(documents):
            doc.embedding = all_embeddings[i]
        return documents
    def embed_query(self, query):
        """Generate embedding for a query string"""
        return self.embeddings_model.embed_query(query)

import sqlite3
import json
class SQLiteMetadataStore:
    def init(self, db_path="metadata.db"):
        self.db_path = db_path
        self._create_tables()
    def createtables(self):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        # Create documents table
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS documents (
            id TEXT PRIMARY KEY,
            title TEXT,
            hash TEXT,
            metadata TEXT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        )
        ''')
        conn.commit()
        conn.close()
    def update_document_metadata(self, document_id, metadata, document_hash):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        # Convert metadata to JSON string
        metadata_json = json.dumps(metadata)
        # Check if document exists
        cursor.execute("SELECT id FROM documents WHERE id = ?", (document_id,))
        exists = cursor.fetchone()
        if exists:
            cursor.execute('''
            UPDATE documents SET
                title = ?,
                hash = ?,
                metadata = ?,
                updated_at = CURRENT_TIMESTAMP
            WHERE id = ?
            ''', (metadata.get('title', document_id), document_hash, metadata_json, document_id))
        else:
            cursor.execute('''
            INSERT INTO documents (id, title, hash, metadata)
            VALUES (?, ?, ?, ?)
            ''', (document_id, metadata.get('title', document_id), document_hash, metadata_json))
        conn.commit()
        conn.close()
    def get_document_hash(self, document_id):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT hash FROM documents WHERE id = ?", (document_id,))
        result = cursor.fetchone()
        conn.close()
        if result:
            return result[0]
        return None
    def get_document_metadata(self, document_id):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT metadata FROM documents WHERE id = ?", (document_id,))
        result = cursor.fetchone()
        conn.close()
        if result:
            return json.loads(result[0])
        return None
    def delete_document(self, document_id):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("DELETE FROM documents WHERE id = ?", (document_id,))
        conn.commit()
        conn.close()
    def list_all_documents(self):
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        cursor.execute("SELECT id, title, updated_at FROM documents ORDER BY updated_at DESC")
        results = cursor.fetchall()
        conn.close()
        return [{"id": r[0], "title": r[1], "updated_at": r[2]} for r in results]

from langchain.chains import RetrievalQAWithSourcesChain
from langchain.prompts import PromptTemplate
from langchain_ollama import Ollama
from typing import List, Dict, Any
class RAGPipeline:
    def init(
        self,
        embedding_service,
        vectorstore_service,
        metadata_store,
        model_name="llama3",
        temperature=0.1,
        top_k=5
    ):
        self.embedding_service = embedding_service
        self.vectorstore_service = vectorstore_service
        self.metadata_store = metadata_store
        self.model_name = model_name
        self.temperature = temperature
        self.top_k = top_k
        # Initialize LLM
        self.llm = Ollama(
            model=model_name,
            temperature=temperature
        )
        # Create QA prompt
        self.qa_prompt = PromptTemplate(
            template="""You are an assistant that helps users find information in a collection of markdown documents.
            Answer the question based on the following context:
            {context}
            Question: {question}
            Provide a comprehensive answer. If the answer cannot be found in the context, say so clearly.
            Include relevant source information when possible.
            """,
            input_variables=["context", "question"]
        )
    def query(self, query_text, filter_metadata=None):
        # Embed the query
        query_embedding = self.embedding_service.embed_query(query_text)
        # Construct filter expression if needed
        filter_expr = None
        if filter_metadata:
            filter_parts = []
            for key, value in filter_metadata.items():
                filter_parts.append(f'metadata["{key}"] == "{value}"')
            filter_expr = " && ".join(filter_parts)
        # Retrieve relevant documents
        results = self.vectorstore_service.search(
            query_embedding=query_embedding,
            top_k=self.top_k,
            filter=filter_expr
        )
        # Construct context from results
        context_texts = []
        sources = []
        for result in results:
            context_texts.append(result["text"])
            if result["source"]:
                sources.append(result["source"])
        context = "\n\n".join(context_texts)
        # Generate response with LLM
        prompt = self.qa_prompt.format(context=context, question=query_text)
        response = self.llm.invoke(prompt)
        return {
            "query": query_text,
            "response": response,
            "sources": list(set(sources)),
            "results": results
        }

from fastapi import FastAPI, HTTPException, Body
from pydantic import BaseModel
from typing import Dict, List, Optional, Any
import uvicorn
import os
app = FastAPI(title="Markdown RAG API")
# Initialize services
embedding_service = EmbeddingService(model_name="mxbai-embed-large")
metadata_store = SQLiteMetadataStore(db_path="metadata.db")
vectorstore_service = MilvusVectorStore(
    host=os.environ.get("MILVUS_HOST", "localhost"),
    port=os.environ.get("MILVUS_PORT", "19530")
)
document_processor = DocumentProcessor(
    embedding_service=embedding_service,
    vectorstore_service=vectorstore_service,
    metadata_store=metadata_store
)
rag_pipeline = RAGPipeline(
    embedding_service=embedding_service,
    vectorstore_service=vectorstore_service,
    metadata_store=metadata_store,
    model_name=os.environ.get("OLLAMA_MODEL", "llama3")
)
# Initialize directory watcher
watcher_service = DirectoryWatcherService(
    directory_path=os.environ.get("VAULT_DIR", "./vault"),
    processor=document_processor
)
# Start directory watcher on startup
@app.on_event("startup")
async def startup_event():
    # Process all existing files first
    print(f"Processing existing files in {os.environ.get('VAULT_DIR', './vault')}")
    document_processor.process_all_files(os.environ.get("VAULT_DIR", "./vault"))
    # Start watching for changes
    print("Starting directory watcher")
    watcher_service.start()
# Shutdown directory watcher on shutdown
@app.on_event("shutdown")
async def shutdown_event():
    print("Stopping directory watcher")
    watcher_service.stop()
# Define API models
class QueryRequest(BaseModel):
    query: str
    filter_metadata: Optional[Dict[str, Any]] = None
    model: Optional[str] = None
    temperature: Optional[float] = None
    top_k: Optional[int] = None
class QueryResponse(BaseModel):
    query: str
    response: str
    sources: List[str]
# Define API endpoints
@app.post("/api/query", response_model=QueryResponse)
async def query(request: QueryRequest = Body(...)):
    try:
        # Use provided model params or defaults
        model = request.model or rag_pipeline.model_name
        temperature = request.temperature or rag_pipeline.temperature
        top_k = request.top_k or rag_pipeline.top_k
        # Create a customized pipeline if needed
        if model != rag_pipeline.model_name or temperature != rag_pipeline.temperature or top_k != rag_pipeline.top_k:
            custom_pipeline = RAGPipeline(
                embedding_service=embedding_service,
                vectorstore_service=vectorstore_service,
                metadata_store=metadata_store,
                model_name=model,
                temperature=temperature,
                top_k=top_k
            )
            result = custom_pipeline.query(request.query, request.filter_metadata)
        else:
            result = rag_pipeline.query(request.query, request.filter_metadata)
        return {
            "query": result["query"],
            "response": result["response"],
            "sources": result["sources"]
        }
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/documents")
async def list_documents():
    try:
        documents = metadata_store.list_all_documents()
        return documents
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/reindex")
async def reindex():
    try:
        document_processor.process_all_files(os.environ.get("VAULT_DIR", "./vault"))
        return {"status": "success", "message": "Reindexing complete"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
# Run the API server
if name == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)