|
| 1 | +#!/usr/bin/env python3 |
| 2 | +""" |
| 3 | +Build vector database script extracted from the original rag.py logic. |
| 4 | +This script contains the specific vector database building functions that were previously in rag.py. |
| 5 | +
|
| 6 | +Usage: |
| 7 | + python3 build_rag_vector_db.py --source summarized_regulations --output regulations_db --strategy per_regulation |
| 8 | +""" |
| 9 | + |
| 10 | +import argparse |
| 11 | +from pathlib import Path |
| 12 | +from typing import List, Optional |
| 13 | + |
| 14 | +from langchain_ollama import OllamaEmbeddings |
| 15 | +from langchain_community.vectorstores import FAISS |
| 16 | +from langchain_core.documents import Document |
| 17 | +from langchain_text_splitters import RecursiveCharacterTextSplitter, MarkdownHeaderTextSplitter |
| 18 | + |
| 19 | +SUPPORTED_CONTENT_FILES = ["text.txt", "content.md"] # prefer plain text first |
| 20 | + |
| 21 | +def gather_local_regulations(source_dir: str = "filtered_regulations") -> List[Document]: |
| 22 | + """Collect one Document per regulation directory (already a doc-level split).""" |
| 23 | + base = Path(source_dir) |
| 24 | + if not base.exists(): |
| 25 | + raise FileNotFoundError( |
| 26 | + f"Source directory not found: {source_dir} (try running from test_warehouse_regulations dir)" |
| 27 | + ) |
| 28 | + docs: List[Document] = [] |
| 29 | + for item in sorted(base.iterdir()): |
| 30 | + if not item.is_dir() or not item.name.startswith("1910."): |
| 31 | + continue |
| 32 | + content_path: Optional[Path] = None |
| 33 | + for fname in SUPPORTED_CONTENT_FILES: |
| 34 | + candidate = item / fname |
| 35 | + if candidate.exists(): |
| 36 | + content_path = candidate |
| 37 | + break |
| 38 | + if not content_path: |
| 39 | + continue |
| 40 | + try: |
| 41 | + text = content_path.read_text(encoding="utf-8", errors="ignore") |
| 42 | + except Exception as e: |
| 43 | + print(f"WARN: failed reading {content_path}: {e}") |
| 44 | + continue |
| 45 | + docs.append(Document(page_content=text, metadata={"reg_dir": item.name, "source_file": content_path.name})) |
| 46 | + return docs |
| 47 | + |
| 48 | +def split_documents(base_docs: List[Document], strategy: str = "per_regulation", *, |
| 49 | + chunk_size: int = 1000, chunk_overlap: int = 200) -> List[Document]: |
| 50 | + """Return list of Documents according to chosen strategy. |
| 51 | +
|
| 52 | + Strategies: |
| 53 | + per_regulation: one chunk per regulation (no further splitting). |
| 54 | + recursive: standard RecursiveCharacterTextSplitter. |
| 55 | + markdown_headers: split by markdown headers (#, ##, etc.) then optionally merge small pieces. |
| 56 | + """ |
| 57 | + if strategy == "per_regulation": |
| 58 | + return base_docs |
| 59 | + if strategy == "recursive": |
| 60 | + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) |
| 61 | + return splitter.split_documents(base_docs) |
| 62 | + if strategy == "markdown_headers": |
| 63 | + header_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=[("#", "h1"), ("##", "h2"), ("###", "h3")]) |
| 64 | + out: List[Document] = [] |
| 65 | + for d in base_docs: |
| 66 | + parts = header_splitter.split_text(d.page_content) |
| 67 | + # propagate metadata |
| 68 | + for p in parts: |
| 69 | + p.metadata.update(d.metadata) |
| 70 | + out.extend(parts) |
| 71 | + # (Optional) re-chunk very small parts using recursive splitter for consistency |
| 72 | + normalized: List[Document] = [] |
| 73 | + tmp_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) |
| 74 | + for p in out: |
| 75 | + if len(p.page_content) < chunk_size // 3: |
| 76 | + normalized.extend(tmp_splitter.split_documents([p])) |
| 77 | + else: |
| 78 | + normalized.append(p) |
| 79 | + return normalized |
| 80 | + raise ValueError(f"Unknown split strategy: {strategy}") |
| 81 | + |
| 82 | +def build_local_index(source_dir: str = "filtered_regulations", |
| 83 | + strategy: str = "per_regulation", |
| 84 | + chunk_size: int = 1000, |
| 85 | + chunk_overlap: int = 200, |
| 86 | + embedding_model: str = "mxbai-embed-large") -> FAISS: |
| 87 | + """Build FAISS vector store from local regulation documents.""" |
| 88 | + print(f"Loading regulation documents from '{source_dir}' using '{strategy}' strategy ...") |
| 89 | + base_docs = gather_local_regulations(source_dir) |
| 90 | + print(f"Loaded {len(base_docs)} base regulation documents.") |
| 91 | + |
| 92 | + docs = split_documents(base_docs, strategy=strategy, chunk_size=chunk_size, chunk_overlap=chunk_overlap) |
| 93 | + print(f"Prepared {len(docs)} chunks (avg chars: {sum(len(d.page_content) for d in docs)//max(1,len(docs))}).") |
| 94 | + |
| 95 | + # Initialize embeddings |
| 96 | + embeddings = OllamaEmbeddings(model=embedding_model) |
| 97 | + |
| 98 | + # Build FAISS vector store |
| 99 | + vector_store = FAISS.from_documents(docs, embedding=embeddings) |
| 100 | + print("Initialized FAISS index (dimension inferred from first batch).") |
| 101 | + |
| 102 | + return vector_store |
| 103 | + |
| 104 | +def save_vector_store(vector_store: FAISS, output_path: str): |
| 105 | + """Save the FAISS vector store to disk.""" |
| 106 | + output_dir = Path(output_path) |
| 107 | + output_dir.mkdir(parents=True, exist_ok=True) |
| 108 | + vector_store.save_local(str(output_dir)) |
| 109 | + print(f"Saved FAISS vector store to: {output_dir}") |
| 110 | + |
| 111 | +def main(): |
| 112 | + parser = argparse.ArgumentParser( |
| 113 | + description="Build FAISS vector database from regulation documents using original rag.py logic" |
| 114 | + ) |
| 115 | + parser.add_argument( |
| 116 | + "--source", "-s", |
| 117 | + default="processed_regulations", |
| 118 | + help="Source directory containing regulation folders (default: processed_regulations)" |
| 119 | + ) |
| 120 | + parser.add_argument( |
| 121 | + "--output", "-o", |
| 122 | + default="regulations_db", |
| 123 | + help="Output directory for FAISS vector database (default: regulations_db)" |
| 124 | + ) |
| 125 | + parser.add_argument( |
| 126 | + "--strategy", |
| 127 | + choices=["per_regulation", "recursive", "markdown_headers"], |
| 128 | + default="per_regulation", |
| 129 | + help="Document splitting strategy (default: per_regulation)" |
| 130 | + ) |
| 131 | + parser.add_argument( |
| 132 | + "--chunk-size", |
| 133 | + type=int, |
| 134 | + default=1000, |
| 135 | + help="Chunk size for text splitting (default: 1000)" |
| 136 | + ) |
| 137 | + parser.add_argument( |
| 138 | + "--chunk-overlap", |
| 139 | + type=int, |
| 140 | + default=200, |
| 141 | + help="Chunk overlap for text splitting (default: 200)" |
| 142 | + ) |
| 143 | + parser.add_argument( |
| 144 | + "--embedding-model", |
| 145 | + default="mxbai-embed-large", |
| 146 | + help="Ollama embedding model to use (default: mxbai-embed-large)" |
| 147 | + ) |
| 148 | + parser.add_argument( |
| 149 | + "--test-query", |
| 150 | + help="Optional test query to run after building the database" |
| 151 | + ) |
| 152 | + |
| 153 | + args = parser.parse_args() |
| 154 | + |
| 155 | + # Build the vector store |
| 156 | + vector_store = build_local_index( |
| 157 | + source_dir=args.source, |
| 158 | + strategy=args.strategy, |
| 159 | + chunk_size=args.chunk_size, |
| 160 | + chunk_overlap=args.chunk_overlap, |
| 161 | + embedding_model=args.embedding_model |
| 162 | + ) |
| 163 | + |
| 164 | + # Save to disk |
| 165 | + save_vector_store(vector_store, args.output) |
| 166 | + |
| 167 | + # Optional test query |
| 168 | + if args.test_query: |
| 169 | + print(f"\nRunning test query: '{args.test_query}'") |
| 170 | + results = vector_store.similarity_search(args.test_query, k=3) |
| 171 | + for i, result in enumerate(results, 1): |
| 172 | + print(f"[{i}] {result.metadata.get('reg_dir', 'unknown')} ({len(result.page_content)} chars)") |
| 173 | + print(f" {result.page_content[:200]}..." if len(result.page_content) > 200 else f" {result.page_content}") |
| 174 | + print() |
| 175 | + |
| 176 | + print("Vector database build completed successfully!") |
| 177 | + |
| 178 | +if __name__ == "__main__": |
| 179 | + main() |
0 commit comments