406 lines
16 KiB
Python
406 lines
16 KiB
Python
|
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.")
|