commit b64782b2b27d42ef02b37baafd25e508fa2da5b7 Author: thinhle Date: Sun Jun 14 23:11:03 2026 +0700 Initial commit: search script, indexing script, database and gitignore diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..733d6b0 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +# Exclude raw datasets and python cache +juniper-kb/ +network-books/ +__pycache__/ +*.log diff --git a/embed_docs.py b/embed_docs.py new file mode 100644 index 0000000..a5a8c2b --- /dev/null +++ b/embed_docs.py @@ -0,0 +1,405 @@ +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.") diff --git a/knowledge_base.db b/knowledge_base.db new file mode 100644 index 0000000..6587d6b Binary files /dev/null and b/knowledge_base.db differ diff --git a/search_kb.py b/search_kb.py new file mode 100644 index 0000000..a7cfb10 --- /dev/null +++ b/search_kb.py @@ -0,0 +1,233 @@ +import os +import sys +import sqlite3 +import requests +import numpy as np +from rich.console import Console +from rich.panel import Panel +from rich.text import Text + +# --- CONFIGURATION --- +OLLAMA_URL = "http://localhost:11434/api/embed" +EMBEDDING_MODEL = "qwen3-embedding:0.6b" +DB_PATH = "/root/work/knowledge-base/knowledge_base.db" + +console = Console() + +def get_query_embedding(query_text): + """Fetches embedding for the search query from Ollama.""" + try: + response = requests.post(OLLAMA_URL, json={ + "model": EMBEDDING_MODEL, + "input": query_text + }, timeout=30) + response.raise_for_status() + data = response.json() + return data.get("embeddings", [])[0] + except Exception as e: + console.print(f"[bold red]Error getting query embedding from Ollama:[/bold red] {e}") + console.print("[yellow]Make sure Ollama is running and the model qwen3-embedding:0.6b is pulled.[/yellow]") + sys.exit(1) + +def load_table_data(cursor, table_name, source_type): + """Loads text chunks and embeddings from a specific table.""" + if table_name == "knowledge_base": + cursor.execute("SELECT id, kb_id, salesforce_id, title, url, prefix_content, embedding FROM knowledge_base") + rows = cursor.fetchall() + data = [] + embeddings = [] + for row in rows: + doc_id, kb_id, salesforce_id, title, url, prefix_content, emb_bytes = row + emb = np.frombuffer(emb_bytes, dtype=np.float32) + if len(emb) == 1024: + embeddings.append(emb) + data.append({ + "id": doc_id, + "source_type": source_type, + "source_name": salesforce_id, + "title": f"{kb_id} : {title}" if not title.startswith(kb_id) else title, + "url": url, + "page_num": None, + "text": prefix_content + }) + return data, embeddings + else: + cursor.execute(f"SELECT id, file_name, page_num, prefix_content, embedding FROM {table_name}") + rows = cursor.fetchall() + data = [] + embeddings = [] + for row in rows: + doc_id, file_name, page_num, prefix_content, emb_bytes = row + emb = np.frombuffer(emb_bytes, dtype=np.float32) + if len(emb) == 1024: + embeddings.append(emb) + data.append({ + "id": doc_id, + "source_type": source_type, + "source_name": file_name, + "title": file_name, + "url": None, + "page_num": page_num, + "text": prefix_content + }) + return data, embeddings + +def get_full_content(item): + """Retrieves the full content of a KB article or PDF page.""" + if item['source_type'] == 'kb': + kb_id = item['source_name'] + file_path = f"/root/work/knowledge-base/juniper-kb/kb_markdown/{kb_id}.md" + if os.path.exists(file_path): + try: + with open(file_path, 'r', encoding='utf-8') as f: + return f.read() + except Exception as e: + console.print(f"[dim yellow]Warning: could not read file {file_path}: {e}[/dim yellow]") + + # Fallback to combining chunks from SQLite + conn = sqlite3.connect(DB_PATH) + c = conn.cursor() + c.execute("SELECT text_content FROM knowledge_base WHERE kb_id = ? ORDER BY chunk_index", (kb_id,)) + chunks = [row[0] for row in c.fetchall()] + conn.close() + return "\n\n".join(chunks) + else: + # Reconstruct page content + file_name = item['source_name'] + page_num = item['page_num'] + table_name = "os_release_note" if item['source_type'] == 'release_note' else "network_book" + + conn = sqlite3.connect(DB_PATH) + c = conn.cursor() + c.execute(f"SELECT text_content FROM {table_name} WHERE file_name = ? AND page_num = ? ORDER BY chunk_index", (file_name, page_num)) + chunks = [row[0] for row in c.fetchall()] + conn.close() + return "\n\n".join(chunks) + +def search(query, top_k=5, target_source=None, show_full=False): + """Performs cosine similarity search against the embedded documents.""" + if not os.path.exists(DB_PATH): + console.print(f"[bold red]Database not found at {DB_PATH}.[/bold red] Please run embed_docs.py first.") + return + + # 1. Embed query + console.print(f"[dim]Generating embedding for query...[/dim]") + query_emb = np.array(get_query_embedding(query), dtype=np.float32) + + # 2. Connect to database and load data + conn = sqlite3.connect(DB_PATH) + cursor = conn.cursor() + + all_data = [] + all_embeddings = [] + + tables_to_load = [ + ("knowledge_base", "kb"), + ("network_book", "book") + ] + + if target_source: + tables_to_load = [t for t in tables_to_load if t[1] == target_source] + + for table_name, source_type in tables_to_load: + try: + data, embs = load_table_data(cursor, table_name, source_type) + all_data.extend(data) + all_embeddings.extend(embs) + except sqlite3.OperationalError: + # Table might not exist yet if indexing has not run for this source + continue + + conn.close() + + if not all_embeddings: + console.print("[bold yellow]No matching tables or data found in the database. Run embed_docs.py to index documents.[/bold yellow]") + return + + # 3. Calculate cosine similarity + console.print(f"[dim]Searching {len(all_embeddings)} document chunks...[/dim]") + embs_matrix = np.array(all_embeddings, dtype=np.float32) + + q_norm = np.linalg.norm(query_emb) + m_norms = np.linalg.norm(embs_matrix, axis=1) + + # Avoid zero division + m_norms[m_norms == 0] = 1e-10 + if q_norm == 0: + q_norm = 1e-10 + + similarities = np.dot(embs_matrix, query_emb) / (m_norms * q_norm) + + # 4. Extract Top-K + top_indices = np.argsort(similarities)[::-1][:top_k] + + # 5. Display output + console.print(Panel( + f"[bold green]Query:[/bold green] '{query}'\n[dim]Model: {EMBEDDING_MODEL} | Top {top_k} results shown below[/dim]", + border_style="green", + expand=False + )) + + for rank, idx in enumerate(top_indices, 1): + item = all_data[idx] + score = similarities[idx] + + title_text = f"{rank}. {item['title']} (Score: [bold green]{score:.4f}[/bold green])" + meta_info = f"Type: [bold blue]{item['source_type'].upper()}[/bold blue]" + if item['page_num']: + meta_info += f" | Page: [bold yellow]{item['page_num']}[/bold yellow]" + if item['url']: + meta_info += f" | [link={item['url']}]URL[/link]" + + if not show_full: + meta_info += " | [dim yellow]Add --full to view complete doc/page[/dim yellow]" + content_to_show = Text(item['text']) + else: + if item['source_type'] == 'kb': + from rich.markdown import Markdown + try: + content_to_show = Markdown(get_full_content(item)) + except Exception: + content_to_show = Text(get_full_content(item)) + else: + content_to_show = Text(get_full_content(item)) + + console.print(Panel( + content_to_show, + title=title_text, + subtitle=meta_info, + border_style="cyan", + padding=(1, 2) + )) + +if __name__ == "__main__": + if len(sys.argv) < 2: + console.print("[bold yellow]Usage:[/bold yellow] python3 search_kb.py \"\" [--source kb|book] [--top-k ] [--full]") + sys.exit(0) + + query = sys.argv[1] + + # Parse args + target_source = None + top_k = 5 + show_full = "--full" in sys.argv + + if "--source" in sys.argv: + try: + idx = sys.argv.index("--source") + target_source = sys.argv[idx + 1] + if target_source not in ["kb", "book"]: + raise ValueError + except Exception: + console.print("[bold red]Error:[/bold red] --source must be one of: kb, book") + sys.exit(1) + + if "--top-k" in sys.argv: + try: + idx = sys.argv.index("--top-k") + top_k = int(sys.argv[idx + 1]) + except Exception: + console.print("[bold red]Error:[/bold red] --top-k must be an integer") + sys.exit(1) + + search(query, top_k=top_k, target_source=target_source, show_full=show_full)