trycua--cua/docs/scripts/generate_db.py

"""
Database generator for CUA documentation
Parses crawled JSON data and creates a LanceDB vector database for RAG
"""

import json
import re
from pathlib import Path
from typing import Optional

import lancedb
from lancedb.embeddings import get_registry
from lancedb.pydantic import LanceModel, Vector

# Configuration
CRAWLED_DATA_DIR = Path(__file__).parent.parent / "crawled_data"
DB_PATH = Path(__file__).parent.parent / "docs_db"
CHUNK_SIZE = 1000  # Characters per chunk
CHUNK_OVERLAP = 200  # Overlap between chunks

# Use sentence-transformers for embeddings
model = get_registry().get("sentence-transformers").create(name="all-MiniLM-L6-v2")


class DocChunk(LanceModel):
    """Schema for document chunks in the database"""

    text: str = model.SourceField()
    vector: Vector(model.ndims()) = model.VectorField()
    url: str
    title: str
    category: str
    subcategory: Optional[str]
    page: str
    chunk_index: int


def clean_markdown(markdown: str) -> str:
    """Clean markdown content for better chunking"""
    # Remove excessive whitespace
    text = re.sub(r"\n{3,}", "\n\n", markdown)
    # Remove image markdown
    text = re.sub(r"!\[.*?\]\(.*?\)", "", text)
    # Remove link URLs but keep text
    text = re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", text)
    # Remove HTML tags
    text = re.sub(r"<[^>]+>", "", text)
    # Clean up whitespace
    text = re.sub(r" {2,}", " ", text)
    return text.strip()


def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> list[str]:
    """Split text into overlapping chunks, respecting sentence boundaries"""
    if not text:
        return []

    # Split by paragraphs first
    paragraphs = text.split("\n\n")
    chunks = []
    current_chunk = ""

    for para in paragraphs:
        para = para.strip()
        if not para:
            continue

        # If adding this paragraph exceeds chunk size, save current and start new
        if len(current_chunk) + len(para) + 2 > chunk_size:
            if current_chunk:
                chunks.append(current_chunk.strip())
                # Start new chunk with overlap from previous
                if overlap > 0 and len(current_chunk) > overlap:
                    # Try to find a sentence boundary for overlap
                    overlap_text = current_chunk[-overlap:]
                    sentence_end = overlap_text.rfind(". ")
                    if sentence_end > 0:
                        overlap_text = overlap_text[sentence_end + 2 :]
                    current_chunk = overlap_text + "\n\n" + para
                else:
                    current_chunk = para
            else:
                # Single paragraph exceeds chunk size, split by sentences
                sentences = re.split(r"(?<=[.!?])\s+", para)
                for sentence in sentences:
                    if len(current_chunk) + len(sentence) + 1 > chunk_size:
                        if current_chunk:
                            chunks.append(current_chunk.strip())
                            # Start new chunk with overlap from previous, similar to paragraph logic
                            if overlap > 0 and len(current_chunk) > overlap:
                                overlap_text = current_chunk[-overlap:]
                                sentence_end = overlap_text.rfind(". ")
                                if sentence_end > 0:
                                    overlap_text = overlap_text[sentence_end + 2 :]
                                current_chunk = (overlap_text + " " + sentence).strip()
                            else:
                                current_chunk = sentence.strip()
                        else:
                            # No existing chunk; start with this sentence
                            current_chunk = sentence.strip()
                    else:
                        current_chunk = (current_chunk + " " + sentence).strip()
        else:
            current_chunk = (current_chunk + "\n\n" + para).strip()

    # Don't forget the last chunk
    if current_chunk:
        chunks.append(current_chunk.strip())

    return chunks


def load_crawled_data() -> list[dict]:
    """Load all crawled page data"""
    all_pages_file = CRAWLED_DATA_DIR / "_all_pages.json"

    if all_pages_file.exists():
        with open(all_pages_file, "r", encoding="utf-8") as f:
            return json.load(f)

    # Fallback: load individual files
    pages = []
    for json_file in CRAWLED_DATA_DIR.glob("*.json"):
        if json_file.name.startswith("_"):
            continue
        with open(json_file, "r", encoding="utf-8") as f:
            pages.append(json.load(f))

    return pages


def process_pages(pages: list[dict]) -> list[dict]:
    """Process pages into document chunks"""
    all_chunks = []

    for page in pages:
        markdown = page.get("markdown", "")
        if not markdown:
            continue

        # Clean the markdown
        cleaned_text = clean_markdown(markdown)
        if not cleaned_text or len(cleaned_text) < 50:
            continue

        # Get path info
        path_info = page.get("path_info", {})

        # Chunk the text
        text_chunks = chunk_text(cleaned_text)

        # Ensure non-null values for required fields
        url = page.get("url", "")
        title = page.get("title") or path_info.get("page", "") or "Untitled"
        category = path_info.get("category") or "unknown"
        page_name = path_info.get("page") or ""

        for i, chunk_text_content in enumerate(text_chunks):
            chunk = {
                "text": chunk_text_content,
                "url": url,
                "title": title,
                "category": category,
                "subcategory": path_info.get("subcategory"),
                "page": page_name,
                "chunk_index": i,
            }
            all_chunks.append(chunk)

    return all_chunks


def create_database(chunks: list[dict]):
    """Create LanceDB database from chunks"""
    # Remove existing database
    if DB_PATH.exists():
        import shutil

        shutil.rmtree(DB_PATH)

    # Create database
    db = lancedb.connect(DB_PATH)

    # Create table with schema
    table = db.create_table(
        "docs",
        schema=DocChunk,
        mode="overwrite",
    )

    # Add data in batches
    batch_size = 100
    for i in range(0, len(chunks), batch_size):
        batch = chunks[i : i + batch_size]
        print(f"Adding batch {i // batch_size + 1}/{(len(chunks) + batch_size - 1) // batch_size}")
        table.add(batch)

    print(f"Database created at: {DB_PATH}")
    print(f"Total chunks: {len(chunks)}")

    return db


def test_search(db: lancedb.DBConnection, query: str, limit: int = 5):
    """Test search functionality"""
    table = db.open_table("docs")

    print(f"\nSearching for: '{query}'")
    print("-" * 50)

    results = table.search(query).limit(limit).to_list()

    for i, result in enumerate(results):
        print(f"\n{i + 1}. [{result['category']}] {result['title']}")
        print(f"   URL: {result['url']}")
        print(f"   Score: {result.get('_distance', 'N/A'):.4f}")
        print(f"   Preview: {result['text'][:150]}...")


def main():
    print("Loading crawled data...")
    pages = load_crawled_data()
    print(f"Loaded {len(pages)} pages")

    if not pages:
        print("No crawled data found. Run crawl_docs.py first.")
        return

    print("\nProcessing pages into chunks...")
    chunks = process_pages(pages)
    print(f"Created {len(chunks)} chunks")

    if not chunks:
        print("No chunks created. Check your crawled data.")
        return

    print("\nCreating database...")
    db = create_database(chunks)

    # Test with sample queries
    print("\n" + "=" * 50)
    print("Testing search functionality")
    print("=" * 50)

    test_queries = [
        "how to install CUA",
        "computer use agent",
        "benchmark evaluation",
        "API reference",
    ]

    for query in test_queries:
        test_search(db, query)

    print("\n" + "=" * 50)
    print("Database generation complete!")
    print(f"Database location: {DB_PATH}")


if __name__ == "__main__":
    main()