import os import json import sqlite3 import requests import numpy as np from tqdm import tqdm import pypdf # --- CONFIGURATION --- OLLAMA_URL = "http://localhost:11434/api/embed" EMBEDDING_MODEL = "qwen3-embedding:0.6b" DB_PATH = "/root/work/knowledge-base/knowledge_base.db" KB_JSON_PATH = "/root/work/knowledge-base/juniper-kb/juniper_kb_data_clean.json" OS_RELEASE_DIR = "/root/work/knowledge-base/juniper-os-release" BOOKS_DIR = "/root/work/knowledge-base/network-books" def init_db(db_path): """Initializes SQLite database and tables.""" conn = sqlite3.connect(db_path) cursor = conn.cursor() # 1. Table for Juniper KB cursor.execute(""" CREATE TABLE IF NOT EXISTS knowledge_base ( id INTEGER PRIMARY KEY AUTOINCREMENT, kb_id TEXT NOT NULL, salesforce_id TEXT NOT NULL, title TEXT, url TEXT, last_modified TEXT, products TEXT, categories TEXT, environment TEXT, symptoms TEXT, cause TEXT, description TEXT, solution TEXT, chunk_index INTEGER NOT NULL, text_content TEXT NOT NULL, prefix_content TEXT NOT NULL, embedding BLOB NOT NULL ) """) # 3. Table for Network Books cursor.execute(""" CREATE TABLE IF NOT EXISTS network_book ( id INTEGER PRIMARY KEY AUTOINCREMENT, file_name TEXT NOT NULL, page_num INTEGER NOT NULL, chunk_index INTEGER NOT NULL, text_content TEXT NOT NULL, prefix_content TEXT NOT NULL, embedding BLOB NOT NULL ) """) # 4. Progress tracker for resumable execution cursor.execute(""" CREATE TABLE IF NOT EXISTS indexing_progress ( source_type TEXT NOT NULL, source_name TEXT NOT NULL, status TEXT NOT NULL, PRIMARY KEY (source_type, source_name) ) """) # Create database indices cursor.execute("CREATE INDEX IF NOT EXISTS idx_kb_id ON knowledge_base (kb_id)") cursor.execute("CREATE INDEX IF NOT EXISTS idx_book_file ON network_book (file_name)") conn.commit() return conn def get_embeddings(texts): """Fetches embeddings from the local Ollama API for a batch of texts.""" try: response = requests.post(OLLAMA_URL, json={ "model": EMBEDDING_MODEL, "input": texts }, timeout=120) response.raise_for_status() data = response.json() return data.get("embeddings", []) except Exception as e: print(f"Error calling Ollama API: {e}") raise e def batch_list(lst, batch_size): """Helper to split a list into batches.""" for i in range(0, len(lst), batch_size): yield lst[i:i + batch_size] def chunk_text_words(text, chunk_size_words=300, overlap_words=30): """Splits text into chunks of given size using word boundaries.""" if not text: return [] words = text.split() if not words: return [] chunks = [] i = 0 while i < len(words): chunk_words = words[i:i+chunk_size_words] chunks.append(" ".join(chunk_words)) if i + chunk_size_words >= len(words): break i += chunk_size_words - overlap_words return chunks def split_markdown_by_headings(text): """Splits markdown content into sections based on H2 headings.""" import re parts = re.split(r'^##\s+(.*)$', text, flags=re.MULTILINE) sections = [] # Text before the first H2 heading intro = parts[0].strip() if intro: sections.append(("general", intro)) for i in range(1, len(parts), 2): heading = parts[i].strip() content = parts[i+1].strip() if i+1 < len(parts) else "" sections.append((heading, content)) return sections def process_kb(db_path, json_path, limit=None): """Processes and embeds Juniper KB articles from JSON and markdown files.""" print(f"\n[1/3] Processing Juniper Knowledge Base articles from {json_path}...") conn = sqlite3.connect(db_path) cursor = conn.cursor() if not os.path.exists(json_path): print(f"Warning: JSON file not found at {json_path}") conn.close() return with open(json_path, 'r', encoding='utf-8') as f: articles = json.load(f) count = 0 for article in tqdm(articles, desc="KB Articles"): salesforce_id = article.get("kb_id") title = article.get("title", "") # Extract human-readable KB/TN ID from title import re m = re.search(r'\b((?:KB|TN)\d+)\b', title) kb_id = m.group(1) if m else salesforce_id # Resumability check cursor.execute("SELECT status FROM indexing_progress WHERE source_type='kb' AND source_name=?", (salesforce_id,)) row = cursor.fetchone() if row and row[0] == 'completed': continue url = article.get("url", "") last_modified = article.get("last_modified", "") products = ", ".join(article.get("products", [])) if isinstance(article.get("products"), list) else str(article.get("products") or "") categories = ", ".join(article.get("categories", [])) if isinstance(article.get("categories"), list) else str(article.get("categories") or "") environment = article.get("environment", "") or "" # Read raw markdown file if available, to get headers file_path = f"/root/work/knowledge-base/juniper-kb/kb_markdown/{salesforce_id}.md" markdown_content = "" if os.path.exists(file_path): try: with open(file_path, 'r', encoding='utf-8') as f: markdown_content = f.read() except Exception: markdown_content = article.get("markdown_content", "") else: markdown_content = article.get("markdown_content", "") # Extract semantic sections sections = split_markdown_by_headings(markdown_content) desc_text = "" trigger_text = "" solution_text = "" fixed_text = "" symptoms_text = "" cause_text = "" for heading, content in sections: h_lower = heading.lower() if 'description' in h_lower: desc_text = content elif 'trigger' in h_lower: trigger_text = content elif 'solution' in h_lower or 'workaround' in h_lower or 'resolution' in h_lower: solution_text = (solution_text + "\n\n" + content) if solution_text else content elif 'fixed' in h_lower or 'upgrade' in h_lower or 'fix version' in h_lower or 'fix_version' in h_lower: fixed_text = (fixed_text + "\n\n" + content) if fixed_text else content elif 'symptom' in h_lower or 'issue' in h_lower: symptoms_text = (symptoms_text + "\n\n" + content) if symptoms_text else content elif 'cause' in h_lower: cause_text = content # Fallbacks to JSON fields if not extracted from markdown if not desc_text: desc_text = article.get("description", "") or "" if not solution_text: solution_text = article.get("solution", "") or "" if not symptoms_text: symptoms_text = article.get("symptoms", "") or "" if not cause_text: cause_text = article.get("cause", "") or "" # PRIORITIZE: Construct description that includes description, trigger, solution, and fixed info desc_parts = [] if desc_text: desc_parts.append(desc_text) if trigger_text: desc_parts.append(f"Trigger: {trigger_text}") if solution_text: desc_parts.append(f"Solution/Workaround: {solution_text}") if fixed_text: desc_parts.append(f"Fixed/Upgrade: {fixed_text}") prioritized_description = "\n\n".join(desc_parts) if desc_parts else desc_text # Build chunks based on sections all_kb_chunks = [] for heading, content in sections: if not content or len(content.strip()) < 5: continue # Split section content into chunks chunks = chunk_text_words(content, chunk_size_words=300, overlap_words=30) for idx, chunk in enumerate(chunks): prefix = f"[KB Article: {title} (ID: {kb_id}) - {heading}]" prefix_chunk = f"{prefix}\n{chunk}" all_kb_chunks.append((chunk, prefix_chunk)) if not all_kb_chunks: # Fallback if no sections or empty fallback_text = f"Title: {title}\nDescription: {prioritized_description}" all_kb_chunks = [(fallback_text, f"[KB Article: {title} (ID: {kb_id})]\n{fallback_text}")] try: # Get embeddings in batches embeddings = [] prefix_texts = [item[1] for item in all_kb_chunks] for batch in batch_list(prefix_texts, 32): embeddings.extend(get_embeddings(batch)) # Insert into database for idx, ((chunk, prefix_chunk), emb) in enumerate(zip(all_kb_chunks, embeddings)): emb_blob = np.array(emb, dtype=np.float32).tobytes() cursor.execute(""" INSERT INTO knowledge_base ( kb_id, salesforce_id, title, url, last_modified, products, categories, environment, symptoms, cause, description, solution, chunk_index, text_content, prefix_content, embedding ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """, ( kb_id, salesforce_id, title, url, last_modified, products, categories, environment, symptoms_text, cause_text, prioritized_description, solution_text, idx, chunk, prefix_chunk, emb_blob )) cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES ('kb', ?, 'completed')", (salesforce_id,)) conn.commit() count += 1 if limit and count >= limit: print(f"Reached KB test limit of {limit}") break except Exception as e: print(f"\nError processing KB {kb_id} (SF: {salesforce_id}): {e}") conn.rollback() cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES ('kb', ?, 'failed')", (salesforce_id,)) conn.commit() conn.close() def process_pdfs(db_path, dir_path, source_type, limit_files=None): """Processes PDF files (release notes or books) using optimized batching.""" print(f"\n[2/3 & 3/3] Processing PDF documents in {dir_path} ({source_type})...") conn = sqlite3.connect(db_path) cursor = conn.cursor() if not os.path.exists(dir_path): print(f"Warning: Directory not found at {dir_path}") conn.close() return files = [f for f in os.listdir(dir_path) if f.lower().endswith('.pdf')] files.sort() count_files = 0 for file in files: # Resumability check cursor.execute("SELECT status FROM indexing_progress WHERE source_type=? AND source_name=?", (source_type, file)) row = cursor.fetchone() if row and row[0] == 'completed': print(f"Skipping already processed PDF: {file}") continue file_path = os.path.join(dir_path, file) print(f"\nIndexing PDF file: {file}") try: import pdfplumber with pdfplumber.open(file_path) as pdf: num_pages = len(pdf.pages) # Step 1: Extract and chunk all pages all_chunks = [] for page_idx in tqdm(range(num_pages), desc=f"Reading pages of {file}"): page = pdf.pages[page_idx] page_num = page_idx + 1 try: text = page.extract_text(x_tolerance=1.5) or "" except Exception: text = page.extract_text() or "" text = text.strip() if not text: continue chunks = chunk_text_words(text, chunk_size_words=300, overlap_words=30) for idx, chunk in enumerate(chunks): prefix = f"[Doc: {file}, Page: {page_num}]" prefix_chunk = f"{prefix}\n{chunk}" all_chunks.append((page_num, idx, chunk, prefix_chunk)) if not all_chunks: print(f"No text extracted from PDF: {file}") cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'completed')", (source_type, file)) conn.commit() continue # Step 2: Batch embed print(f"Generating embeddings for {len(all_chunks)} chunks...") embeddings = [] prefix_texts = [item[3] for item in all_chunks] for batch in tqdm(batch_list(prefix_texts, 32), total=(len(prefix_texts) + 31) // 32, desc="Embedding batches"): embeddings.extend(get_embeddings(batch)) # Step 3: Insert into SQLite in a single transaction table_name = "os_release_note" if source_type == "release_note" else "network_book" for (page_num, idx, chunk, prefix_chunk), emb in zip(all_chunks, embeddings): emb_blob = np.array(emb, dtype=np.float32).tobytes() cursor.execute(f""" INSERT INTO {table_name} ( file_name, page_num, chunk_index, text_content, prefix_content, embedding ) VALUES (?, ?, ?, ?, ?, ?) """, ( file, page_num, idx, chunk, prefix_chunk, emb_blob )) cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'completed')", (source_type, file)) conn.commit() print(f"Successfully indexed PDF: {file}") count_files += 1 if limit_files and count_files >= limit_files: print(f"Reached {source_type} test file limit of {limit_files}") break except Exception as e: print(f"\nError processing PDF {file}: {e}") conn.rollback() cursor.execute("INSERT OR REPLACE INTO indexing_progress (source_type, source_name, status) VALUES (?, ?, 'failed')", (source_type, file)) conn.commit() conn.close() if __name__ == "__main__": import sys # Clean up previous database in test runs if clean DB is needed # (Optional: uncomment if you want a fresh DB for each run, but we keep it for resumability) # Initialize DB init_db(DB_PATH) # Quick CLI test flag is_test = len(sys.argv) > 1 and sys.argv[1] == "--test" kb_limit = 5 if is_test else None pdf_limit = 1 if is_test else None if is_test: print("!!! RUNNING IN TEST MODE (subset of data only) !!!") try: # Process 1: Juniper KB process_kb(DB_PATH, KB_JSON_PATH, limit=kb_limit) # Process 3: Network Books process_pdfs(DB_PATH, BOOKS_DIR, "network_book", limit_files=pdf_limit) print("\nAll done!") except KeyboardInterrupt: print("\nProcess interrupted by user. SQLite transactions rolled back safely.")