Initial commit: search script, indexing script, database and gitignore

.gitignore

@@ -0,0 +1,5 @@
+# Exclude raw datasets and python cache
+juniper-kb/
+network-books/
+__pycache__/
+*.log

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.")

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 \"<search query>\" [--source kb|book] [--top-k <number>] [--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)